Building a First World Model › Implementation: The Full Loop
World Models HandbookAll chapters Building a First World Model › Implementation: The Full LoopOne section · Sections in this chapter6 / 9 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Now we assemble everything. The code below is organized as the complete pipeline: environment, data collection, model, training, diagnostics, planner, and closed-loop evaluation. Each block is small and self-contained, and blocks build on earlier ones.
As you read the code, keep the conceptual map from the previous sections in view. The environment implements the physics we are pretending not to know. Data collection stands in for the experience an embodied agent would accumulate. The model class implements the four-component anatomy, with the filtering and imagination steps exposed as separate methods so their different roles remain visible. Training implements the alternating objective. Diagnostics implement the two evaluation protocols. The planner implements CEM over imagined returns. And the final loop wires filtering, planning, and acting together into the closed agent. Nothing in the code introduces a concept the prose has not already prepared you for.
Collecting Experience
First the environment. The true state holds position and velocity; the observation exposes position plus Gaussian noise; the reward is negative distance to goal minus a small force penalty, evaluated at the resulting state so that it is exactly predictable from the predicted next latent and the action taken.
import numpy as np
class PointMassEnv:
"""A 2D point mass with friction. Observations reveal position only."""
def __init__(self, dt=0.1, friction=0.6, goal=(1.0, 1.0), obs_noise=0.02):
self.dt = dt
self.friction = friction
self.goal = np.array(goal)
self.obs_noise = obs_noise
def reset(self, rng):
self.state = rng.uniform(-1.0, 1.0, size=4) # [x, y, vx, vy]
return self._observe(rng)
def _observe(self, rng):
return self.state[:2] + rng.normal(0.0, self.obs_noise, size=2)
def step(self, action, rng):
force = np.clip(action, -1.0, 1.0)
accel = force - self.friction * self.state[2:]
self.state[2:] = self.state[2:] + self.dt * accel
self.state[:2] = self.state[:2] + self.dt * self.state[2:]
obs = self._observe(rng)
reward = -np.linalg.norm(self.state[:2] - self.goal) - 0.1 * np.sum(
force**2
)
return obs, rewardTwo details in this implementation repay attention. The semi-implicit Euler integration, updating velocity first and then position with the new velocity, is more stable than fully explicit Euler for damped systems, which keeps our simulated physics well-behaved over long episodes. And the reward's dependence on the post-step state means that when the model predicts the next latent correctly, the reward head has all the information it needs, so reward-prediction error becomes a clean proxy for transition-model error. Had we evaluated the reward at the pre-step state instead, part of the reward signal would be predictable without any dynamics knowledge at all, diluting the training signal.
Next we collect a dataset with a uniformly random policy. Random exploration wanders aimlessly, which is realistic: most practical world models are trained on data gathered by policies far less competent than the one they eventually enable.
SEED = 7
rng = np.random.default_rng(SEED)
env = PointMassEnv()
N_EPISODES, EP_LEN = 150, 100
ACT_DIM, OBS_DIM = 2, 2
episodes = []
for _ in range(N_EPISODES):
obs = env.reset(rng)
obs_list, act_list, rew_list = [obs], [], []
for _ in range(EP_LEN):
action = rng.uniform(-1.0, 1.0, size=ACT_DIM)
obs, reward = env.step(action, rng)
obs_list.append(obs)
act_list.append(action)
rew_list.append(reward)
episodes.append(
{
"obs": np.array(obs_list), # [T+1, 2]
"act": np.array(act_list), # [T, 2]
"rew": np.array(rew_list), # [T,]
}
)Collected 150 episodes of 100 steps Random-policy return: mean -192.4, range [-455.7, -43.7]
The random policy scores terribly, as expected: it never approaches the goal on purpose. Improving on this number is the concrete target for the rest of the chapter.
Note also what the dataset contains: fifteen thousand transitions, each holding two positions, two force values, and one reward. The total information content is a few hundred kilobytes. That a competent controller can be extracted from so little is a consequence of the environment's simplicity, and it is worth remembering when we scale up: the fundamental currency of world-model learning is transitions, and the cost of acquiring informative ones dominates most real applications.
The figure shows eight sample trajectories wandering without any preference for the goal region. Look closely and you can see the effect of friction: paths curve and slow, occasionally lingering in one area before a burst of force sends the puck elsewhere. This aimlessness is not a defect of the data but a feature of the setup. Because the random policy covers the arena broadly, the model gets to see the dynamics in many regions of state space, which is exactly what a dynamics learner needs, even though the data contains almost no information about the task.
Defining the Model
The model class mirrors the four-component anatomy exactly. One design detail worth noticing: filter_step accepts a precomputed embedding, so the same method serves both the training loop (which embeds whole windows at once) and the closed-loop agent (which embeds one observation at a time).
import torch
import torch.nn as nn
class LatentDynamicsModel(nn.Module):
def __init__(
self, obs_dim=OBS_DIM, act_dim=ACT_DIM, hidden=64, recurrent=True
):
super().__init__()
self.recurrent = recurrent
self.encoder = nn.Sequential(
nn.Linear(obs_dim, 64),
nn.ReLU(),
nn.Linear(64, hidden),
)
self.transition = nn.Sequential(
nn.Linear(hidden + act_dim, 64),
nn.ReLU(),
nn.Linear(64, hidden),
)
self.gru = nn.GRUCell(hidden + act_dim, hidden) if recurrent else None
self.decoder = nn.Sequential(
nn.Linear(hidden, 64),
nn.ReLU(),
nn.Linear(64, obs_dim),
)
self.reward_head = nn.Sequential(
nn.Linear(hidden + act_dim, 64),
nn.ReLU(),
nn.Linear(64, 1),
)
def filter_step(self, h, emb, prev_action):
"""Posterior update: fold a fresh observation embedding into the latent."""
if self.gru is None:
return emb
return self.gru(torch.cat([emb, prev_action], dim=-1), h)
def imagine_step(self, h, action):
"""Prior step: predict the next latent from the current one and an action."""
return self.transition(torch.cat([h, action], dim=-1))Setting recurrent=False removes the GRU, leaving a memoryless model whose latent is just the current observation's embedding. This is our experimental control: identical data, identical loss, identical planner, with the memory removed.
The ablation is deliberately surgical. By removing only the GRU and rerouting filter_step to pass the raw embedding through, we ensure that any performance difference between the two models is attributable to memory and nothing else. Architectural capacity differs slightly, since the recurrent model has additional parameters, but the difference is minor relative to the shared components, and the direction of the result will confirm that capacity is not the operative variable. This style of minimal paired comparison, changing one capability while freezing everything else, is the workhorse of empirical analysis in this field.
Training
Training windows are sampled uniformly from the collected episodes. Within each window we alternate prior predictions (which generate losses) and filtering updates (which consume fresh observations), exactly as the objective specifies.
WINDOW = 25 # transitions per training window
def sample_batch(batch_size, rng):
idx_ep = rng.integers(0, len(episodes), size=batch_size)
idx_t = rng.integers(0, EP_LEN - WINDOW + 1, size=batch_size)
obs_b, act_b, rew_b = [], [], []
for ie, it in zip(idx_ep, idx_t):
ep = episodes[ie]
obs_b.append(ep["obs"][it : it + WINDOW + 1])
act_b.append(ep["act"][it : it + WINDOW])
rew_b.append(ep["rew"][it : it + WINDOW])
to_tensor = lambda x: torch.tensor(np.array(x), dtype=torch.float32)
return to_tensor(obs_b), to_tensor(act_b), to_tensor(rew_b)
def compute_losses(model, obs, act, rew):
emb = model.encoder(obs) # [B, WINDOW+1, H]
h = emb[:, 0]
recon_total, reward_total = 0.0, 0.0
for t in range(act.shape[1]):
h_pred = model.imagine_step(h, act[:, t])
obs_pred = model.decoder(h_pred)
rew_in = torch.cat([h_pred, act[:, t]], dim=-1)
rew_pred = model.reward_head(rew_in).squeeze(-1)
recon_total = (
recon_total + ((obs_pred - obs[:, t + 1]) ** 2).sum(dim=-1).mean()
)
reward_total = reward_total + ((rew_pred - rew[:, t]) ** 2).mean()
h = model.filter_step(h, emb[:, t + 1], act[:, t])
return recon_total / act.shape[1], reward_total / act.shape[1]
def train_model(recurrent, epochs=30, lr=3e-3, batch_size=256, seed=0):
torch.manual_seed(seed)
model = LatentDynamicsModel(recurrent=recurrent)
opt = torch.optim.Adam(model.parameters(), lr=lr)
rng_np = np.random.default_rng(seed)
history = []
for _ in range(epochs):
sums, n = {"recon": 0.0, "reward": 0.0}, 0
for _ in range(len(episodes) * EP_LEN // batch_size):
obs, act, rew = sample_batch(batch_size, rng_np)
recon_loss, reward_loss = compute_losses(model, obs, act, rew)
loss = recon_loss + reward_loss
opt.zero_grad()
loss.backward()
opt.step()
b = obs.shape[0]
sums["recon"] += recon_loss.item() * b
sums["reward"] += reward_loss.item() * b
n += b
history.append({k: v / n for k, v in sums.items()})
return model, historyTrace the flow inside compute_losses once, slowly, because it is the objective made executable. Embeddings for the whole window are computed in one batched pass, an efficiency choice that exploits the encoder's feedforward structure. The initial latent is simply the first embedding, since there is no prior history to filter. Then, for each step: imagine forward with the action, decode and score the prediction against the actual next observation, score the reward head against the actual received reward, and only afterward fold the actual next embedding into the latent via the filter step. The ordering is the whole point. Prediction always precedes correction, so the loss measures the model's foresight, not its hindsight.
We train both variants and compare their final losses.
model_rnn, hist_rnn = train_model(recurrent=True, seed=1)
model_ff, hist_ff = train_model(recurrent=False, seed=1)recurrent: recon 0.00250 reward 0.00088 feedforward: recon 0.00307 reward 0.00112
Both models drive their losses down, but look closely at the reconstruction values: the recurrent model reaches a markedly lower floor. That gap is the hidden velocity making itself visible. A memoryless encoder looking at a single noisy position simply cannot know how fast the puck moves, so its best possible one-step prediction is systematically wrong whenever the puck is in motion. The recurrent model infers velocity from the recent observation history and closes most of that gap.
It is worth articulating why the gap cannot be closed by the feedforward model no matter how long it trains. The observation at time is a noisy draw centered on the true position; conditioned on that single draw, the velocity is statistically independent of the observation, so no function of the observation alone can predict the next position better than the conditional mean allows. The recurrent model escapes this ceiling because its input is not one observation but a trajectory of them, and consecutive noisy positions jointly constrain velocity far more tightly than any single one does. The loss floor difference is thus not a training artifact but an information-theoretic necessity, which is what makes this comparison a genuine test of memory rather than of optimization luck.
On the log-scale plot, notice that the reward losses of the two models sit closer together than their reconstruction losses. This is partly because the reward is dominated by distance to goal, which depends mostly on position, a quantity both models represent well. The residual advantage of the recurrent model on reward prediction comes from situations where velocity matters for the next distance, for example when the puck is moving rapidly toward or away from the goal. Even in this tiny problem, the pattern is instructive: task-relevant losses can mask representational deficits that reconstruction exposes clearly, which is one reason to monitor multiple losses rather than a single aggregate.
Measuring Rollout Quality
Now the diagnostic that separates genuine simulators from lucky one-step predictors. On held-out episodes we track prediction error under two protocols: teacher-forced, where the latent is corrected with the true observation every step, and recursive, where the model consumes only its own predictions.
rng_eval = np.random.default_rng(123)
eval_eps = []
for _ in range(20):
obs = env.reset(rng_eval)
obs_list, act_list = [obs], []
for _ in range(EP_LEN):
a = rng_eval.uniform(-1.0, 1.0, size=ACT_DIM)
obs, _ = env.step(a, rng_eval)
obs_list.append(obs)
act_list.append(a)
eval_eps.append({"obs": np.array(obs_list), "act": np.array(act_list)})
MAX_HORIZON = 50
@torch.no_grad()
def rollout_errors(model, eps):
one_step = np.zeros(MAX_HORIZON)
recursive = np.zeros(MAX_HORIZON)
for ep in eps:
obs = torch.tensor(ep["obs"], dtype=torch.float32)
act = torch.tensor(ep["act"], dtype=torch.float32)
emb = model.encoder(obs)
h_true, h_free = emb[0], emb[0]
for t in range(MAX_HORIZON):
p_true = model.imagine_step(h_true, act[t])
p_free = model.imagine_step(h_free, act[t])
one_step[t] += (model.decoder(p_true) - obs[t + 1]).norm().item()
recursive[t] += (model.decoder(p_free) - obs[t + 1]).norm().item()
h_true = model.filter_step(h_true, emb[t + 1], act[t])
h_free = p_free # no correction: pure imagination
return one_step / len(eps), recursive / len(eps)
err_rnn_1step, err_rnn_rec = rollout_errors(model_rnn, eval_eps)
_, err_ff_rec = rollout_errors(model_ff, eval_eps)The evaluation episodes are freshly generated with a different seed from training, so none of the measured errors reflect memorization. Inside rollout_errors, two latents evolve in parallel from the same starting point: h_true receives a corrective filter update at every step, implementing the teacher-forced protocol, while h_free is overwritten with its own prediction, implementing the recursive protocol. Both are scored against the same ground-truth observations, so the only difference between the two curves is whether reality was allowed to intervene. That isolation is what makes the comparison interpretable.
horizon 1: one-step 0.073 recursive RNN 0.073 recursive FF 0.071 horizon 5: one-step 0.073 recursive RNN 0.935 recursive FF 0.862 horizon 10: one-step 0.046 recursive RNN 1.006 recursive FF 0.886 horizon 25: one-step 0.066 recursive RNN 1.261 recursive FF 1.146 horizon 50: one-step 0.069 recursive RNN 1.563 recursive FF 1.461
Three phenomena are visible, and each teaches something general.
First, the teacher-forced curve stays essentially flat. Correcting the latent with real observations every step keeps errors from accumulating, confirming that the one-step transition model is genuinely accurate.
Second, the recurrent recursive curve jumps sharply and then climbs steadily. Nothing went wrong in execution; this is the train-test mismatch discussed above, visible in miniature. The model was optimized with corrected latent states, so it has no incentive to remain stable when its own predictions are fed back recursively.
Third, the feedforward recursive curve is slightly below the recurrent one in this particular run. That result is not evidence that memory is useless. It shows that representational richness and rollout stability are different properties: the recurrent latent can encode velocity, as the probe below demonstrates, while its learned transition can still amplify off-manifold errors more strongly than the simpler feedforward system. A fair world-model evaluation must measure both what a latent contains and how its dynamics behave under self-feedback.
A fourth observation hides in the shape of the curves rather than their levels. The recursive curves rise steeply at first and then flatten somewhat, reflecting the fact that the arena is bounded: once the imagined puck has drifted far from the truth, both trajectories wander within the same confined box, and the error saturates near the scale of the arena itself. In unbounded systems, compounding error grows without limit; in bounded ones it plateaus at a level that can still be useless for planning, since a prediction that is wrong by half the arena's width cannot distinguish good actions from bad ones. Always read error curves with the geometry of the domain in mind.
To verify directly that the recurrent latent really does encode the hidden velocity, we can run a simple linear probe. We generate fresh episodes where we also record the true velocities, filter the latent at every step, and fit a linear map from latent to velocity on half the data. If the GRU integrates observation history into an implicit velocity estimate, even a plain linear readout should recover it well; the memoryless model, whose latent sees only one noisy position, should fail.
rng_probe = np.random.default_rng(321)
probe_eps = []
for _ in range(20):
obs = env.reset(rng_probe)
vel_list, obs_list, act_list = [env.state[2:].copy()], [obs], []
for _ in range(EP_LEN):
a = rng_probe.uniform(-1.0, 1.0, size=ACT_DIM)
obs, _ = env.step(a, rng_probe)
obs_list.append(obs)
act_list.append(a)
vel_list.append(env.state[2:].copy())
probe_eps.append(
{
"obs": np.array(obs_list),
"act": np.array(act_list),
"vel": np.array(vel_list), # [T+1, 2], ground-truth velocities
}
)
@torch.no_grad()
def collect_latents(model, eps):
hs, vels = [], []
for ep in eps:
obs = torch.tensor(ep["obs"], dtype=torch.float32)
act = torch.tensor(ep["act"], dtype=torch.float32)
emb = model.encoder(obs)
h = emb[0]
hs.append(h.numpy())
vels.append(ep["vel"][0])
for t in range(len(act)):
h = model.filter_step(h, emb[t + 1], act[t])
hs.append(h.numpy())
vels.append(ep["vel"][t + 1])
return np.array(hs), np.array(vels)
H_rnn, V_true = collect_latents(model_rnn, probe_eps)
H_ff, _ = collect_latents(model_ff, probe_eps)
n_train = len(H_rnn) // 2
W_rnn, *_ = np.linalg.lstsq(H_rnn[:n_train], V_true[:n_train], rcond=None)
W_ff, *_ = np.linalg.lstsq(H_ff[:n_train], V_true[:n_train], rcond=None)
pred_rnn = H_rnn[n_train:] @ W_rnn
pred_ff = H_ff[n_train:] @ W_ff
def r2(y_true, y_pred):
ss_res = ((y_true - y_pred) ** 2).sum()
ss_tot = ((y_true - y_true.mean(axis=0)) ** 2).sum()
return 1.0 - ss_res / ss_tot
print(
f"velocity probe R^2: recurrent {r2(V_true[n_train:], pred_rnn):.3f} "
f"feedforward {r2(V_true[n_train:], pred_ff):.3f}"
)
The probe clarifies the representation story from the inside. A single linear layer fitted to the recurrent latents recovers both velocity components with high accuracy on unseen episodes, meaning the GRU's belief state dedicates linearly readable structure to the hidden variable. The feedforward latents, by contrast, support almost no velocity readout: their information content is exhausted by the current noisy position. This is direct evidence that the recurrent model remembers useful hidden state, while the rollout diagnostic reminds us that possessing that information does not automatically produce stable long-horizon predictions.
Imagining a Trajectory
Before trusting the model with control, let us look at what an open-loop rollout actually looks like in observation space. We take one held-out episode, filter the latent at the first step, then replay the episode's real actions through pure imagination and decode positions throughout.
@torch.no_grad()
def open_loop_rollout(model, ep):
obs = torch.tensor(ep["obs"], dtype=torch.float32)
act = torch.tensor(ep["act"], dtype=torch.float32)
h = model.encoder(obs[0:1])
preds = []
for t in range(len(act)):
h = model.imagine_step(h, act[t : t + 1])
preds.append(model.decoder(h)[0])
return torch.stack(preds).numpy()
pred_traj = open_loop_rollout(model_rnn, eval_eps[0])
true_traj = eval_eps[0]["obs"][1:]
The shortened view is the honest face of this learned world model. Even inside the planning-relevant window, the imagined and true paths begin to separate. Receding-horizon control limits the damage by executing only the first proposed action and then filtering a real observation before planning again, but replanning does not make an inaccurate model accurate. It merely prevents one long hallucinated rollout from controlling the whole episode.
Look at the divergence pattern rather than just its final magnitude. The imagined path remains smooth and physically plausible even as it loses alignment with the realized path. The model's failure is therefore one of calibration and accumulated state error rather than a sequence of impossible jumps. Such a model may still rank nearby action sequences usefully over short horizons, but the rollout-error plot warns us not to equate plausibility with accuracy.
The Planner
The planner is a compact, batched CEM. All candidate sequences in a population are rolled out simultaneously through the transition model, which keeps the per-step planning cost negligible.
PLAN_HORIZON, POP, ELITES, ITERS = 10, 100, 25, 3
@torch.no_grad()
def imagine_returns(model, h0, actions, gamma=1.0):
"""Predicted returns for a batch of action sequences, rolled out in parallel."""
h = h0
total = torch.zeros(h0.shape[0])
disc = 1.0
for t in range(actions.shape[1]):
h = model.imagine_step(h, actions[:, t])
r_in = torch.cat([h, actions[:, t]], dim=-1)
total = total + disc * model.reward_head(r_in).squeeze(-1)
disc *= gamma
return total
@torch.no_grad()
def cem_plan(model, h0, act_limit=1.0):
mu = torch.zeros(PLAN_HORIZON, ACT_DIM)
sigma = torch.full((PLAN_HORIZON, ACT_DIM), act_limit)
for _ in range(ITERS):
noise = torch.randn(POP, PLAN_HORIZON, ACT_DIM) * sigma
candidates = (mu + noise).clamp(-act_limit, act_limit)
returns = imagine_returns(model, h0.expand(POP, -1), candidates)
elites = candidates[returns.topk(ELITES).indices]
mu = elites.mean(dim=0)
sigma = elites.std(dim=0) + 1e-3
return mu[0].numpy()Notice how the imagined rollout in imagine_returns mirrors the recursive protocol from the diagnostics section exactly: one filter at the start, then chained prior steps with no corrective observations. The planner therefore inherits precisely the error profile we measured, and our choice of a ten-step horizon is informed by the rollout-error plot, staying comfortably inside the region where the recurrent model's recursive predictions remained reliable. This is good practice in general: planning horizons should be chosen with reference to measured rollout fidelity, not set arbitrarily. Also note the expansion of h0 across the population dimension, which lets a single starting belief seed all hundred candidate futures simultaneously, a batching trick that turns per-candidate simulation cost into a single wide matrix product.
It is instructive to watch CEM's search distribution contract over its few iterations. We instrument one planning call from a real filtered belief and record the Gaussian's mean and standard deviation after each refitting.
@torch.no_grad()
def cem_trace(model, h0, act_limit=1.0):
mu = torch.zeros(PLAN_HORIZON, ACT_DIM)
sigma = torch.full((PLAN_HORIZON, ACT_DIM), act_limit)
trace = [(mu.numpy().copy(), sigma.numpy().copy())]
for _ in range(ITERS):
noise = torch.randn(POP, PLAN_HORIZON, ACT_DIM) * sigma
candidates = (mu + noise).clamp(-act_limit, act_limit)
returns = imagine_returns(model, h0.expand(POP, -1), candidates)
elites = candidates[returns.topk(ELITES).indices]
mu = elites.mean(dim=0)
sigma = elites.std(dim=0) + 1e-3
trace.append((mu.numpy().copy(), sigma.numpy().copy()))
return trace
_obs0 = torch.tensor(eval_eps[0]["obs"], dtype=torch.float32)
_h0 = model_rnn.encoder(_obs0[0:1])
cem_iters = cem_trace(model_rnn, _h0)
plan_axis = np.arange(PLAN_HORIZON)
The contraction is visible but incomplete. At iteration zero the search distribution spans nearly the entire legal force range, so the sampled plans are essentially noise. Elite refitting narrows the band and produces a differentiated mean plan, yet substantial uncertainty remains after the final iteration. That residual spread matters: with only three iterations and an approximate return model, CEM has found a better region rather than a uniquely trustworthy optimum. This is still useful for receding-horizon control, where the plan will be discarded and recomputed after one real action.
Closing the Loop
Finally, the complete agent: filter, imagine, plan, act, repeat. We evaluate three controllers on fresh episodes: the random policy, MPC driven by the feedforward model, and MPC driven by the recurrent model.
@torch.no_grad()
def run_episode(model, rng, plan=True):
obs = env.reset(rng)
h = model.encoder(torch.tensor(obs, dtype=torch.float32).unsqueeze(0))
total, positions = 0.0, [obs.copy()]
for _ in range(EP_LEN):
action = (
cem_plan(model, h) if plan else rng.uniform(-1, 1, size=ACT_DIM)
)
obs, reward = env.step(action, rng)
total += reward
positions.append(obs.copy())
emb = model.encoder(torch.tensor(obs, dtype=torch.float32).unsqueeze(0))
a_t = torch.tensor(action, dtype=torch.float32).unsqueeze(0)
h = model.filter_step(h, emb, a_t)
return total, np.array(positions)
N_EVAL = 12
configs = [
("Random", model_rnn, False),
("Feedforward MPC", model_ff, True),
("Recurrent MPC", model_rnn, True),
]
results, demo_trajs = {}, {}
for name, mdl, plan in configs:
rets, trajs = [], []
for i in range(N_EVAL):
ret, traj = run_episode(mdl, rng_eval, plan=plan)
rets.append(ret)
if plan and i < 5:
trajs.append(traj)
results[name] = np.array(rets)
if trajs:
demo_trajs[name] = trajsThe inner loop of run_episode is the receding-horizon cycle rendered in five lines: plan from the current belief, execute one action in the real environment, then filter the resulting observation into the belief before the next iteration. Everything the chapter has built converges here. The encoder supplies the sensory contribution, the GRU maintains the belief across steps, CEM searches imagined futures seeded by that belief, and the environment delivers the corrective observation that keeps the whole loop honest. It is worth appreciating that the agent's competence emerges entirely from this cycle; there is no separate policy network, no supervised demonstration, and no task-specific engineering beyond the reward definition.
Random: return -268.3 +/- 34.2 Feedforward MPC: return -195.2 +/- 93.4 Recurrent MPC: return -153.8 +/- 64.2
Examine the marked starts and endpoints rather than reading every curved path as a success. The controller often bends trajectories toward the goal and improves average return, but several episodes finish far away or loop through high-error regions. This is exactly the distinction between better than a baseline and reliable control. The recurrent belief contains useful momentum information, yet inaccurate recursive dynamics and imperfect return predictions still mislead the planner. The figure therefore previews two themes developed later in the book: uncertainty-aware planning and protection against model exploitation.
Read these results as a compact demonstration with a warning attached. Both planners improve mean return over the random policy, so decisions made inside a learned model can transfer to the real environment. The recurrent planner also improves on the feedforward planner in this run, consistent with the value of remembering hidden velocity. But the error bars overlap and the trajectory plot shows conspicuous failures, so this small experiment does not establish reliable control or isolate memory as the only cause. It establishes a working loop and gives us diagnostics for where that loop breaks.
It is worth pausing on what did not happen here. Nobody labeled velocity. Nobody provided equations of motion. Nobody gave the planner a curriculum. A generic recurrent network, squeezed by a reconstruction and reward objective, discovered that tracking an unobservable quantity was the cheapest way to predict the future, and a generic search procedure exploited those predictions to behave competently. That recipe, scaled up and refined, is the skeleton of essentially every modern world-model agent.
There is also a quieter lesson in the comparison worth naming explicitly. The feedforward model is not useless; it beats random by a clear margin, because position information alone supports short-horizon planning when velocities happen to be small. Its failure is graded, not absolute, and it concentrates exactly where the task demands momentum reasoning: fast approaches, sharp turns, final braking. Real systems exhibit the same graded structure of competence, and diagnosing where a model-based agent fails, by correlating failures with the hidden variables its latent cannot express, is one of the most practically valuable skills this book aims to teach.