Part of World Models Handbook
Build a minimal world model in PyTorch with encoder, transition model, decoder, and reward head. Topics include filtering, imagination, CEM planning, and MPC.
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
Building a First World Model
You have probably had the experience of catching a ball without consciously computing anything. Your brain did not solve the projectile equations symbolically. Instead, somewhere inside you, there is a machinery that takes what you currently see, combines it with what happened a moment ago, and produces a prediction: "if I keep moving my hand like this, the ball will land here." That machinery is a world model. It is an internal simulator that turns observations and actions into predictions about the future, and it is the engine behind planning, motor skill, and much of what we call common sense.
Notice how much work that single sentence does. A world model is not merely a memory of the past, because memories alone do not answer questions about things that have not happened yet. Nor is it merely a predictor of observations, because predictions are only useful if they can be steered by hypothetical choices. The essential trick is conditioning: the model accepts an action as input and produces a different future for each action you might take. That conditional structure is what separates a world model from a passive forecaster, and it is what makes search over futures possible at all. When you catch the ball, your nervous system is not replaying remembered catches; it is running forward simulations of candidate hand trajectories and selecting among them before committing your muscles.
This chapter builds the smallest complete version of that idea. By the end, you will have trained a neural network that watches a partially observable environment, maintains an internal state it never observes directly, predicts the consequences of actions it has not yet taken, and uses those imagined futures to choose good actions. Everything fits in a few hundred lines of PyTorch and runs on a laptop CPU.
We deliberately start with a working system rather than definitions. Formal questions, like what exactly qualifies as a world model and how the many proposed variants relate to each other, come in the following chapters. But having built and interrogated a concrete system first, those questions will have something to attach to. There is a practical pedagogical reason for this ordering. Definitions in this field are contested, and different research communities use overlapping vocabulary with subtly different meanings. If we began with taxonomy, you would have no way to judge which distinctions matter. If we begin with a system small enough to hold entirely in your head, every later distinction can be tested against something concrete: does this refinement change what the latent contains, how errors compound, or how plans behave? Those three questions will organize much of what follows.
A Task to Think With
Abstract descriptions of "learning to predict the consequences of actions" are hard to reason about, so let us pin everything to one small environment for the whole chapter.
Imagine a puck sliding on a frictionful table. You control the puck by applying forces, and your goal is to bring it to a target location. Here is the twist that makes the problem interesting: you can only see where the puck is, never how fast it is moving. Velocity is invisible. Yet velocity determines the future, because momentum carries the puck after you stop pushing.
This is a deliberately engineered gap between what you observe and what you need to know. A controller that reacts only to the current position will overshoot the goal again and again, because it cannot anticipate that the puck will keep drifting. To succeed, the agent must somehow remember: recent positions imply recent motion, and recent motion implies future positions. Memory of the right kind is not a luxury here. It is the whole problem.
It is worth spelling out precisely why reactive control fails, because the failure mode is instructive. Suppose the puck sits directly on the goal but is moving quickly away from it. A position-only controller sees zero error and applies no force, or perhaps a tiny corrective nudge. Either way, momentum carries the puck off target, and only then does the controller react, chasing a puck that is already moving. The result is the familiar oscillating pursuit you may have seen in badly tuned thermostats and camera autofocus systems. The information needed to avoid this trap, namely velocity, exists in the world and even in the observation history, but not in any single observation. The agent's only route to success is to reconstruct the hidden variable from the sequence of things it has seen. This reconstruction-from-history requirement is the defining feature of partially observable problems, and our little puck is the cleanest possible instance of it.
A world model is a learned computational object that, given the history of an agent's interactions, predicts how the environment will evolve under hypothetical future actions. It separates what the agent has seen from what the agent believes, and it supports asking "what would happen if?" without touching the real environment.
The environment itself is simple enough to define in a few lines. The true physical state is a four-dimensional vector: position and velocity in the plane. Actions are two-dimensional forces, clipped to a bounded range. Observations are the two position coordinates plus a little sensor noise. Rewards encourage reaching the goal while penalizing large forces.
Each of these design choices earns its keep. The four-dimensional state is small enough that we can later inspect what the model has and has not learned, yet rich enough that one dimension of it is genuinely hidden. The bounded forces prevent the trivial degenerate solution of applying infinite force, which would make dynamics modeling unnecessary. Sensor noise matters more than its small magnitude suggests: it guarantees that the agent cannot recover velocity perfectly from two consecutive positions, so any useful velocity estimate must be an integrated, smoothed inference over many steps rather than a finite-difference calculation. Finally, evaluating the reward at the resulting state, a detail we return to during implementation, ensures the reward is predictable from quantities the model actually represents. Small design decisions like these determine whether a toy problem teaches the right lesson or an accidental one.
The Shape of a Minimal World Model
This section covers the architecture of our model: what it stores, how it updates that storage when new observations arrive, and how it predicts forward when they do not. We build intuition first and write down the mathematics alongside.
Why a Latent State?
The environment's true state has four numbers, but the agent never sees two of them. So the agent must construct its own state, and we call this constructed state the latent state. It is "latent" because it is never observed directly. It exists only inside the model, inferred from the stream of observations and actions.
Why not skip latents entirely and train a network to map raw observation histories directly to predictions? You could, for short histories. But a fixed-size latent state that is recursively updated gives you something a growing raw-history window does not: a bounded, persistent belief that can be maintained indefinitely with constant computation per step. When we formalize this in Part II: Inference, Decisions, and Control, this idea becomes the belief state of a partially observable decision process. Here we simply build one and watch it work.
There is a second, equally important reason for latents. The quantities we ultimately care about, like "where will the puck be in ten steps" and "how much reward will I collect," are functions of the physical state, not of any single observation. A latent state gives the network a place to accumulate exactly the information those predictions require, and nothing forces it to waste capacity replaying irrelevant pixel-level detail.
A third reason deserves mention even at this early stage, because it shapes everything downstream. A fixed-size latent is a bottleneck, and bottlenecks force compression, and compression forces the network to decide what matters. Given unlimited capacity, a model can hedge by memorizing everything about its inputs; given a tight latent, it must prioritize the features that actually reduce prediction error. In our puck environment this prioritization is easy to see: the single most predictive unobserved quantity is velocity, so pressure from the loss pushes the latent toward representing it. In richer environments the same pressure operates on messier trade-offs, deciding whether surface texture or object identity better serves future prediction, and those trade-offs become a central design question rather than an accident. The bottleneck is thus not merely an efficiency measure. It is the mechanism by which the objective's preferences get written into the representation.
The Four Components
Our minimal world model has four neural networks, each with one job:
- Encoder : compresses a raw observation into an embedding, the sensory contribution to the latent state.
- Transition model : the heart of the system. Given the current latent state and a candidate action, it predicts the next latent state. This is the "physics engine" of the model, learned rather than programmed.
- Decoder : maps a latent state back to observation space, so the model can be trained by reconstructing what actually happened.
- Reward head : maps a latent state and action to a predicted reward, connecting the model to the task rather than only to the sensations.
The subscripts denote the parameters of each network. All four train jointly, end to end, from nothing but recorded interaction data.
It helps to think of these components as answering four distinct questions. The encoder answers "what did I just see?", translating raw sensation into a form the rest of the system can consume. The transition model answers "what happens next if I act?", and it is the only component consulted during imagination. The decoder answers "what would I see if my belief were true?", providing the training signal that anchors latents to reality. The reward head answers "how good is this situation for my task?", tying the whole apparatus to the objective the agent actually cares about. Notice the asymmetry between them: the encoder, decoder, and reward head are ordinary feedforward networks, while the transition model is the piece whose quality under repeated self-application determines whether planning works. Most of the failure modes we discuss later live in that one component.
One might ask why the decoder is needed at all if the goal is control. The answer lies in how gradients reach the latent state. Without a decoder, the only pressure shaping the latent comes from the reward head, and reward signals are sparse and weak in most environments. Reconstruction provides a dense, per-step learning signal that forces the latent to retain enough information to describe the world, and the reward head then builds task structure on top of that foundation. The division of labor is deliberate: reconstruction keeps the latent honest, reward prediction keeps it relevant.
Two Kinds of Transition: Filtering and Imagination
Here is the conceptual core of the chapter, and the distinction most often muddled in casual discussions. The transition machinery plays two different roles depending on whether a fresh observation is available.
When the agent acts and then sees what happened, it performs a filtering update (also called the posterior step). The latent state is revised using both the action just taken and the observation just received:
where denotes concatenation and the GRU cell is a small recurrent network. Because the update sees the newest observation, errors in get corrected. This is the model operating online, in contact with reality.
When the agent wants to imagine, no future observation exists. It applies the prior transition (the imagination step), which relies on the action alone:
From this predicted latent, the decoder produces a predicted observation and the reward head produces a predicted reward . Crucially, the output of the prior step can be fed back into itself: , and so on for as many steps as we like. This recursion is what turns a one-step predictor into a simulator.
The word "filtering" is borrowed from estimation theory, where a filter maintains a running estimate of a hidden state by combining predictions with measurements. Our GRU plays exactly that role, though it learns the combination rule rather than deriving it from assumptions about noise statistics. The intuition worth carrying forward is that filtering is a correction operation: whatever nonsense accumulated in the latent last step, the fresh observation offers a chance to repair it. Imagination, by contrast, is a commitment operation: the model stakes its entire picture of the future on the accuracy of its current belief and its transition function, with no external check until reality eventually weighs in. Every design tension in this book, from multi-step losses to uncertainty penalties to replanning schedules, traces back to managing the gap between these two modes of operation.
Filtering folds a fresh observation into the latent state and corrects past mistakes. Imagination predicts forward from the latent state using actions only, with no corrective signal. Every world model that plans needs both: filtering to stay anchored to reality, imagination to evaluate futures that have not happened. In later chapters these two steps acquire probabilistic formal dress as the posterior and prior of a state-space model.
Notice what the latent state must contain for either step to work. Since observations reveal only position, the encoder output at a single time step says nothing about velocity. But the recurrent filtering update receives a sequence of embeddings, and a sufficiently trained GRU can integrate them into an implicit velocity estimate. We will verify empirically that this happens, by comparing against a model stripped of its recurrence.
This verification plan is worth pausing on, because it illustrates a general experimental strategy. We cannot directly inspect the latent and read off "velocity lives here," since distributed representations do not label their dimensions. What we can do is construct a controlled comparison: two models identical in every respect except the presence of memory, trained identically, evaluated identically. If the recurrent variant systematically outperforms the memoryless one on precisely the tasks that require velocity knowledge, we have strong indirect evidence that the recurrence is being used to encode velocity. Controlled ablation of this kind is often more informative than inspection, and we will rely on it repeatedly.
The Training Objective
Training data comes from an exploratory policy, here just random forces, collected as episodes of transitions . We cut these episodes into overlapping windows of length and train on the following objective. Starting from the first observation of a window, alternate prior predictions and filtering updates:
Read the summand carefully, because every symbol earns its place. At step of the window we hold the filtered latent , which has digested observations up to and including . We apply the prior transition with action , obtaining a predicted next latent. The decoder must reconstruct the actually observed from that predicted latent, and the reward head must predict the actually received . Only then do we perform the filtering update with the real and move on. The gradients therefore flow through chains of prior transitions, teaching the transition model to be accurate over the exact kind of recursive application it will face at planning time.
That last sentence deserves emphasis, because it is easy to misread the objective as pure one-step training. Within a window of length , the loss at step depends on the latent , which was itself shaped by the filtering updates of earlier steps, and gradient descent adjusts earlier transitions based on how well later predictions turn out. The alternation of prior steps and filtering steps means the model is always asked to predict from beliefs that were built up through its own prior machinery, punctuated by corrections. This is a middle ground between pure one-step training and full open-loop rollouts, and the window length controls where that middle ground sits: short windows approximate one-step training, long windows approach open-loop training with all of its stability challenges. Choosing is therefore not an arbitrary hyperparameter but a statement about how far into the future you want the model to be accountable.
A few explicit assumptions are baked into this setup, and naming them now will save confusion later:
- Markov property in the latent: we assume the filtered latent summarizes everything relevant for predicting the future. If important information falls outside the GRU's capacity, predictions degrade silently.
- Approximate determinism: our transition model outputs a single predicted latent, not a distribution. This is adequate when the environment is nearly deterministic apart from small sensor noise, and inadequate when outcomes are genuinely multimodal. Probabilistic transitions arrive in Part V: World-Model Architectures.
- Stationarity: the environment's dynamics do not change over time, so one model fit on old data remains valid.
- Off-policy data: training uses random actions, while evaluation uses planned ones. The model is never explicitly warned about this mismatch, a theme that will recur throughout the book.
Each assumption hides a potential failure. The Markov assumption fails gracefully here, since a modest GRU easily integrates enough history to estimate velocity, but in environments with long-range dependencies, say a door opened twenty minutes ago that still affects airflow, a fixed-size latent may simply be too small, and the degradation appears only as mysteriously poor predictions in rare situations. Approximate determinism fails whenever the same latent and action legitimately lead to different outcomes: a point-estimate model averages those outcomes, producing a prediction of a future that never actually occurs, and a planner acting on that average can be systematically misled. Stationarity fails in any adaptive environment, including environments containing other learning agents. And the off-policy mismatch is subtler than it looks: random actions visit a broad but shallow coverage of the state space, while planned actions drive the system along narrow corridors of high-value states, and the model's accuracy along those corridors is largely untested at training time. None of these assumptions is wrong for our toy problem, but each defines a boundary beyond which the simple recipe breaks, and the parts of this handbook that follow are largely a tour of what lies beyond those boundaries.
From One-Step Predictions to Rolled-Out Futures
This section covers how to evaluate a world model's predictions, and why the evaluation protocol matters as much as the model itself.
Teacher-Forced Versus Recursive Prediction
There are two ways to ask a model about the future, and they produce very different answers.
Teacher-forced (one-step) prediction always conditions on ground truth. At every step, we filter the latent with the real observation, apply one prior transition, and score the prediction. Errors never accumulate because reality resets the latent each step. This measures the quality of a single transition.
Recursive (open-loop) rollout conditions on the model's own outputs. We filter once at the start, then chain prior transitions indefinitely, decoding observations along the way. Each prediction becomes the input to the next. Whatever error enters at step one is amplified and distorted by every subsequent step.
The names come from language modeling, where "teacher forcing" refers to training and evaluating a sequence model while feeding it ground-truth tokens rather than its own generations. The analogy is exact. A language model evaluated one token at a time with correct context can look fluent, while free-running generation drifts into repetition or incoherence; the mechanism of failure, consuming your own imperfect outputs as inputs, is identical. Keeping this analogy in mind is useful because the remedies also transfer: scheduled sampling, exposure to model-generated contexts during training, and explicit multi-step objectives all appear in both fields.
The gap between these two curves is arguably the single most important diagnostic in world-model research. A model can look excellent under teacher forcing and fall apart within twenty imagined steps. Since planning always involves recursive rollouts, teacher-forced accuracy alone tells you surprisingly little about whether the model is useful for decisions. We will measure both, and you should internalize the habit of distrusting one-step metrics.
To make the distrust concrete: imagine reading a paper reporting that a world model achieves 0.01 mean squared error on next-frame prediction. That number tells you the transition function is locally accurate. It says nothing about whether a fifty-step rollout stays in the vicinity of reality, and it says even less about whether returns computed over such rollouts rank action sequences correctly. A planner does not need the model to be right everywhere; it needs the model's relative ordering of candidate futures to survive contact with reality. These are related but distinct demands, and conflating them leads to models optimized for the wrong property.
Compounding Error
Why does recursive error grow? Three mechanisms stack on top of each other:
- Error propagation: a slightly wrong latent state is fed back in, producing a more wrong latent state. For smooth systems the error often grows roughly geometrically at first.
- Distribution shift: the model was trained on latents produced by filtering real observations. During a recursive rollout it visits latents produced by its own predictions, a region of latent space it has rarely seen. Its behavior there is untested.
- Exploitation: when a planner searches over action sequences, it gravitates toward whichever regions of the model are optimistically wrong. We return to this failure mode shortly; it deserves its own treatment in Part XII: Reliable World Models.
It is worth separating the first two mechanisms carefully, because they demand different remedies. Error propagation is a property of the dynamical system interacting with an imperfect transition function: even a perfect model of a chaotic system diverges from reality given an imperfect initial condition, so some growth is unavoidable and the remedy is frequent correction, which is precisely what receding-horizon control provides. Distribution shift, by contrast, is a property of the training procedure: the model has simply never been asked to predict from self-generated latents, so its behavior there reflects extrapolation rather than fitting. That failure is addressable at the source, by training on rollouts, mixing imagined and real contexts, or regularizing the latent space so that self-generated latents stay near filtered ones. Confusing the two mechanisms leads to wasted effort, for example by building ensembles to fight a problem that was really a training-distribution problem.
None of these mechanisms is a bug in our implementation. They are structural properties of learning dynamics models from finite data, and every serious world-model system contains machinery designed specifically to contain them: multi-step training objectives, ensembles, uncertainty penalties, and frequent replanning.
Reconstruction Is Not Enough
One more caution before we build. Our model is trained to reconstruct observations and predict rewards, and it is tempting to equate low reconstruction error with a good world model. Resist that temptation. Reconstruction measures how well the latent supports describing the past and present. Decision-making needs the latent to support predicting futures under interventions. These overlap but are not identical. A latent could achieve beautiful reconstructions while discarding exactly the dynamical information a planner needs, or vice versa. The reward head partially bridges the gap by tying the representation to task-relevant structure, and Chapter 3 develops a proper qualification test for world models that goes well beyond reconstruction scores.
A concrete thought experiment makes the distinction vivid. Suppose the encoder learned to represent position with exquisite precision but packed velocity into a single noisy dimension, since velocity contributes less to immediate reconstruction than position does. One-step reconstruction scores would look excellent, yet every multi-step prediction would inherit the corrupted velocity estimate, and plans built on those predictions would misjudge momentum exactly as our feedforward baseline does below. Conversely, a latent that discarded fine visual detail irrelevant to dynamics, say the exact shade of the table surface, would reconstruct slightly worse but plan better. Which latent is the "better" world model? The question answers itself once you fix the purpose: model quality is not one number but several, covering predictive fidelity, representation adequacy, calibration of uncertainty, and usefulness for decisions, and different applications weight them differently. Holding these notions apart, rather than collapsing them into a single score, is a discipline we will maintain throughout the handbook.
Planning with Imagined Trajectories
This section covers the payoff: using the trained model to choose actions without consulting the real environment.
Scoring Action Sequences
Given the current filtered latent and a candidate action sequence , the model can unroll a purely imaginary trajectory:
and summarize it with a predicted return:
where is a discount factor. Planning is now a search problem over action sequences: find the sequence maximizing . Note what has been achieved. The search touches the real environment zero times. All evaluation of candidate futures happens inside the model, at the speed of matrix multiplication.
The discount factor plays a quiet but important role here. Setting treats all steps within the horizon equally, which suits our task since reaching the goal at any point within ten steps is equally valuable. Values below one emphasize near-term rewards, which has a side benefit for imperfect models: distant imagined rewards are computed from deeply compounded predictions, the least trustworthy numbers the model produces, so discounting them is a crude form of epistemic humility baked into the objective. In general, shorter horizons and stronger discounts both reduce exposure to compounding error, at the cost of myopia about long-term consequences. The choice of horizon and discount is thus another place where model reliability and task requirements must be balanced explicitly.

With a differentiable model we could even backpropagate the return gradient to the actions themselves, a technique explored in Part VII: Planning and Agency. Here we use a simpler and very robust derivative-free method.
The Cross-Entropy Method
The cross-entropy method (CEM) maintains a Gaussian distribution over action sequences and iteratively sharpens it around high scorers:
- Sample a population of action sequences from the current Gaussian.
- Roll out each candidate in imagination and compute its predicted return.
- Keep the top few sequences, the elites.
- Refit the Gaussian's mean and standard deviation to the elites.
- Repeat for a handful of iterations, then return the mean's first action.
CEM is not clever, and that is its virtue. It makes no assumptions about the return landscape beyond local smoothness, it parallelizes trivially, and it degrades gracefully when the model is imperfect. For a first world model it is exactly the right partner.
A few implementation details deserve explanation, since they matter in practice. Sampling the initial Gaussian with a standard deviation equal to the full action range means the first iteration explores broadly across the space of possible plans, and each subsequent refitting concentrates the distribution around whatever worked. The small constant added to the refitted standard deviation prevents premature collapse, the situation where the population becomes so concentrated that all candidates are effectively identical and the search stalls. Taking only the first action of the best plan, rather than executing the whole plan, is deferred to the next subsection. And because every candidate rollout consists of the same number of matrix multiplications through the same networks, the entire population can be scored in a single batched pass, which is why CEM feels instantaneous even on a laptop despite evaluating hundreds of imagined futures per decision.
Receding-Horizon Control
Executing a whole planned sequence open-loop would be reckless: any model error compounds over the horizon. Instead we use receding-horizon control, also known as model-predictive control (MPC):
- Every environment step, filter the newest observation into the latent state.
- Plan a fresh -step action sequence by CEM in imagination.
- Execute only the first action.
- Observe what actually happened, and repeat.
Frequent replanning converts the model from an oracle into a advisor whose suggestions are constantly fact-checked by reality. Short model horizons can suffice because each cycle only needs to identify a good first move. The price is computation per step, which is why the efficiency of imagined rollouts matters in practice.
There is a deeper way to understand why this works, connecting to classical control theory. MPC predates learned world models by decades, and in its classical form it uses analytically specified dynamics, typically linear equations fit by system identification. The loop structure, filter, optimize, apply first move, re-filter, is unchanged whether the dynamics come from least squares on a linear model or from a recurrent network trained on pixels. What learned world models contribute is scope: the same loop now applies to systems whose dynamics are nonlinear, high-dimensional, or unknown in analytic form. What classical MPC contributes to us is a mature understanding of the loop's robustness properties, including the fact that closed-loop feedback can stabilize systems far beyond the horizon over which the internal predictions remain accurate. Our puck controller is, quite literally, a modern instance of a fifty-year-old idea, with a neural simulator substituted for the identified linear model.
Implementation: The Full Loop
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.
Limitations and Impact
It is equally important to be clear-eyed about what this minimal system cannot do, because each limitation defines a branch of the research tree we explore in the rest of this handbook.
The most immediate weakness is compounding rollout error, which we measured directly. Our model is trained exclusively on one-step predictions, so nothing in its objective encourages stability under the recursive self-feedback that planning requires. Modern systems address this with multi-step training objectives, latent overshooting, and scheduled sampling, topics taken up in Part V: World-Model Architectures. Relatedly, our model is deterministic. Real environments are often genuinely stochastic, and a point-estimate transition model collapses multiple possible futures into one average future, which can be a future that never occurs. Representing outcome distributions, rather than averages, is a central theme of probabilistic world models.
The second cluster of weaknesses concerns the interface between the model and the planner. CEM optimizes predicted return, and predicted return is computed by a model with errors. A search procedure will happily find action sequences that exploit pockets where the model is optimistically wrong, achieving magnificent imagined returns that evaporate on contact with reality. This model exploitation problem grows more severe as the planner gets stronger relative to the model, and managing it, through uncertainty estimation, ensembles, pessimism, and careful evaluation, occupies much of Part XII: Reliable World Models. Our frequent replanning is a crude but effective patch: reality audits the plan every single step. But an agent that must replan at every step pays a heavy computational tax, and scaling world models to long-horizon reasoning requires planners that can trust their imaginations for longer stretches.
Third, our treatment of partial observability is minimal. The GRU accumulates information, but it maintains a single point-estimate belief with no notion of confidence. When two distinct situations produce the same observation history ambiguity, our model silently commits to one interpretation. Principled belief representations, where the agent knows what it does not know, are the subject of Bayesian filtering and belief-state methods in Part II: Inference, Decisions, and Control. And finally, our entire pipeline assumes a fixed, stationary environment and a fixed task reward. Environments that change, tasks specified in language, and models that must be updated online raise questions that Part VI: Learning World Models at Scale and Part IX: Foundation and World-Action Models begin to answer.
Against these limitations stands a remarkable record of impact. The recipe demonstrated here, encode, recurrently filter, predict forward, decode, and plan in imagination, is the direct ancestor of a family of systems that learned to master visual control tasks from pixels alone, culminating in agents that imagine entire episodes of experience internally and learn behaviors almost entirely inside their own dreams. Beyond robotics and games, the same template now underlies video prediction models, model-based offline reinforcement learning, and the emerging field of foundation world models trained on internet-scale video. The vocabulary you exercised in this chapter, latent state, prior and posterior transitions, rollout, imagined return, receding-horizon planning, is the shared language of all of them.
It is worth noticing how each limitation doubles as a research program with a recognizable lineage. Compounding error motivated the shift from purely one-step objectives to multi-step and rollout-based training. Determinism motivated probabilistic and ensemble transition models. Model exploitation motivated uncertainty-aware planning and pessimism. Point-estimate beliefs motivated variational and Bayesian state-space formulations. And the stationarity assumption motivates continual and online adaptation. Read in this light, the history of world-model research is not a sequence of disconnected tricks but a systematic response to the specific gaps this chapter's minimal system exhibits. You have now seen the gaps firsthand, which means the solutions ahead have something concrete to solve.
Summary
- A world model predicts environmental consequences of actions from an internal latent state, enabling search over futures without touching the real environment.
- The minimal architecture has four parts: an encoder, a transition model, a decoder, and a reward head, trained jointly by observation reconstruction and reward prediction.
- Filtering folds fresh observations into the latent state and corrects errors; imagination rolls the transition model forward on actions alone. Planning needs both.
- Teacher-forced accuracy and recursive rollout accuracy are different properties. Only the second matters for planning, and it is strictly harder, because errors compound and the model visits its own outputs.
- Latent memory is decisive under partial observability. In our experiment, removing the recurrent update left the model unable to represent velocity, visibly degrading both its predictions and its control performance.
- Receding-horizon MPC contains model imperfection by replanning every step from freshly filtered beliefs, trading computation for robustness.
- The remaining gaps, compounding error, stochasticity, model exploitation, and uncertain beliefs, are not defects of this particular implementation but structural research problems that organize the rest of this handbook.
With a concrete system in hand, we can now ask the question this chapter deliberately postponed: what, precisely, makes an object a world model, and what should we demand of one? That is the business of the next chapter.
Quiz
Ready to test your understanding? Take this quick quiz to reinforce what you've learned about building a first world model.
Building a First World Model
Reference
Citation details
Cite or share this article.
Continue with the full handbook
This chapter is part of World Models Handbook. Use the handbook page to browse the complete table of contents and continue reading in sequence.
Explore World Models HandbookStay up to date
Get articles, book updates, and news delivered to your inbox.
No spam, unsubscribe anytime.
Join the community
Sign in to remove popups, track your reading progress, and join the discussion.

Comments
No comments yet. Be the first to share your thoughts!