World Models

An agent that can only learn by acting pays for every lesson in the world: a crashed car, a dropped cup, a robot arm that spent the night trying. A world model moves part of that cost inside the agent's head. It is a network that has learned how the environment responds to actions — given where things are now and what I do, what comes next — so that the agent can try actions in imagination, rank them, and only then act.

The idea is old: in 1943 Kenneth Craik argued that an organism carrying "a 'small-scale model' of external reality and of its own possible actions within its head" could "try out various alternatives" before committing to one¹. What changed between 2018 and 2026 is that the model can now be learned from pixels with the machinery of the earlier chapters: the encoder of a VAE, the attention blocks of Transformers, the denoisers of diffusion and DiT, the tokens of autoregressive generation. Learned world models now collect diamonds in Minecraft, match AlphaZero at Go without being told the rules, plan pick-and-place with robot arms in labs whose data they never saw, and run as playable games.

This chapter is about the part that is easy to underrate: a world model is only as good as what a planner does with it. A model with excellent one-step predictions can still lead a planner into a wall, for three distinct reasons, and the lab measures each of them.

Slides

Click the slides and use ← → to move through them; F goes full screen and P opens the presenter view, with the notes. The simulators of this page run inside the slides. Open the slides in a new tab

What a world model is

Take an environment with a state \(s_t\), an action \(a_t\) chosen by the agent, and an observation \(o_t\) the agent receives (an image, usually; the state itself is rarely visible). The environment has dynamics: a distribution over what comes next, \(p(s_{t+1} \mid s_t, a_t)\). A world model is a learned approximation of it, almost always in a latent space:

Optional heads hang off the latent: a decoder \(d_\psi(z) \approx o\) (to train by reconstruction, or to look at the dreams), a reward model \(\hat r(z)\), a value \(v(z)\), a policy \(\pi(a \mid z)\). Which heads exist, and what the loss compares, is what separates the families below.

Why an agent wants one:

• Planning. Search over action sequences in the model and execute the best. Nothing has to be learned about behaviour beforehand; the goal can change at test time.
• Imagination. Generate synthetic experience from the model and train a policy on it, as in Sutton's Dyna (1991)² and every Dreamer.
• Sample efficiency. Real interaction is the expensive resource. On its six control tasks PlaNet trained for 1 000 episodes, against 100 000 for the model-free D4PG, and came close to D4PG's final scores; the paper estimates that D4PG needs, on average, about 200× more episodes than PlaNet to reach the same score⁴.

The price is that every decision now passes through something that can be wrong. The rest of the chapter is about where it is wrong and what that does.

Five families, one pipeline

Every world model in this chapter fits the same diagram: observation, encoder, latent state, a predictor conditioned on the action, and something trained on top of the prediction. Step through the five families and watch which boxes change.

V–M–C: the first agent trained in its own dream

Ha & Schmidhuber's World Models (2018)³ split the agent in three:

The asymmetry is the design: almost all the capacity is in the world model, and C is so small it can be trained by an evolution strategy (CMA-ES) instead of by gradients. On CarRacing the agent scored 906 ± 21 over 100 trials, above the 900 that counts as solving the task. On VizDoom, C was trained entirely inside M's dream and then transferred to the real game.

The same paper found the problem this chapter is organized around. In the dream the controller "discovered an adversarial policy to move around in such a way so that the monsters in this virtual environment governed by M never shoots a single fireball". The authors' diagnosis: the policy "visits states where the model is wrong because they are away from the training distribution". Raising the temperature \(\tau\) of M's samples made the dream harder to exploit: at \(\tau = 0.1\) the agent scored 2 086 in the dream and 193 in the real game; at \(\tau = 1.15\), 918 in the dream and 1 092 in reality.

The RSSM: PlaNet and Dreamer

PlaNet (2019)⁴ introduced the recurrent state-space model, the core of every Dreamer since. The state has a deterministic part \(h_t\) (a GRU's memory) and a stochastic part \(s_t\):

Training maximizes an evidence lower bound, exactly the VAE's recipe stretched over time: reconstruct \(o_t\) and \(r_t\) from the posterior state, and pay a KL term \(\mathrm{KL}\big(q(s_t \mid h_t, o_t)\,\|\,p(s_t \mid h_t)\big)\) that forces the prior — the only thing available when imagining — to agree with what the observation revealed. The paper found "having both components to be crucial for high planning performance": a deterministic path alone cannot represent several possible futures, and a stochastic one alone has trouble carrying information over many steps.

PlaNet used the model only for planning: CEM (next section) with horizon 12, 10 iterations, 1 000 candidate sequences and the best 100 kept. Dreamer (2020)⁵ replaced the planner with a policy: an actor and a critic trained entirely on imagined rollouts of 15 steps, with gradients of the value flowing back through the learned dynamics. Its critic regresses a λ-return, which blends the imagined rewards with the critic's own estimates. With \(\gamma\) the discount and \(\lambda\) the mixing weight,

With Dreamer's \(\gamma = 0.99\), \(\lambda = 0.95\), a 2-step imagined rollout with rewards \(r_0 = 1\), \(r_1 = 0\) and critic values \(v(s_1) = 2\), \(v(s_2) = 3\): \(R^\lambda_2 = 3\); \(R^\lambda_1 = 0 + 0.99\,(0.05 \cdot 3 + 0.95 \cdot 3) = 2.97\); \(R^\lambda_0 = 1 + 0.99\,(0.05 \cdot 2 + 0.95 \cdot 2.97) = 1 + 0.99 \cdot 2.9215 = 3.892\). \(\lambda = 0\) would trust the critic after one step; \(\lambda = 1\) would trust the imagined rewards all the way to the horizon, and with them every error of the model.

The line continued with three papers worth knowing by what each one added:

• DreamerV2 (2021)⁶: the stochastic state became 32 categorical variables with 32 classes each, and "KL balancing" trains the prior faster than the posterior (\(\alpha = 0.8\)). First agent to reach human level on the 55-game Atari benchmark by learning inside a world model, on one GPU.
• DreamerV3 (2023; Nature, April 2025)⁷: one fixed set of hyperparameters across more than 150 tasks in 8 domains, made possible by normalizations such as \(\mathrm{symlog}(x) = \mathrm{sign}(x)\ln(|x|+1)\) for targets of unknown scale and "free bits" (no KL penalty below 1 nat). It was "the first algorithm to collect diamonds in Minecraft from scratch without human data or curricula".
• Dreamer 4 (2025)⁸: the world model became a 2-billion-parameter video model trained on 2 541 hours of recorded Minecraft play, fast enough to run in real time on one GPU. The agent learned only inside it, with no environment interaction, and still obtained diamonds — in 0.7% of episodes, a small number for a task that takes a long chain of correct mouse and keyboard actions.

MuZero: a model of what matters

MuZero (2020)⁹ took the opposite position on reconstruction. Its three functions are a representation \(s^0 = h_\theta(o_1, \dots, o_t)\), a dynamics \(r^k, s^k = g_\theta(s^{k-1}, a^k)\) and a prediction \(p^k, v^k = f_\theta(s^k)\) (policy and value). The loss compares only those predictions — reward, value and policy, along \(K = 5\) unrolled steps — with the targets from real play. In the authors' words, "there is no direct constraint or requirement for the hidden state to capture all information necessary to reconstruct the original observation".

A model judged only by the decisions it supports is free to ignore everything else: the colour of the background, the texture of the board. It is planned in with Monte Carlo tree search (800 simulations per move in board games, 50 in Atari) and matched AlphaZero at Go, chess and shogi without being given the rules.

Latent predictive models: no pixels at all

A fourth family keeps a learned latent dynamics model and drops the decoder, predicting the next embedding instead of the next frame — the JEPA idea of the previous chapter applied to actions. TD-MPC2 (2024)¹¹ trains it with "joint-embedding prediction …, reward prediction, and TD-learning … without decoding observations" and plans with MPPI; one 317M-parameter agent handles 80 tasks. DINO-WM (2025)¹⁴ predicts the future patch features of a frozen DINOv2 encoder and plans by minimizing the squared distance to the embedding of a goal image. V-JEPA 2-AC (2025)¹⁵ does the same on top of a video encoder pretrained on over a million hours of internet video, post-trained on fewer than 62 hours of unlabeled robot video, and is deployed zero-shot on Franka arms in two labs.

Planning in the model: MPC

The planner that PlaNet, TD-MPC2, DINO-WM and V-JEPA 2-AC all use is model-predictive control:

1. from the current state, search for the action sequence \(a_{t:t+H-1}\) that minimizes a cost computed in the model;
2. execute only its first action;
3. observe what really happened, and repeat from there.

The search most often used is the cross-entropy method (CEM). Keep a Gaussian over plans, with a mean \(\mu \in \mathbb R^{H \times d_a}\) and a standard deviation \(\sigma\) per entry. Then, \(I\) times:

with \(\hat z^{(i)}_{t+h}\) the latent the model imagines after the first \(h\) actions of plan \(i\), and \(c\) a per-step cost (a distance to the goal, or a negative predicted reward). MPPI¹⁰ replaces the hard top-\(K_e\) selection with soft weights \(w_i \propto \exp\!\big(-(C^{(i)} - \min_j C^{(j)})/\lambda\big)\) and sets \(\mu \leftarrow \sum_i w_i a^{(i)}\). Neither uses a gradient: the model only ever runs forward, which is why both work with any model, including a non-differentiable simulator.

The panel runs that loop on a problem small enough to check every number: a point on a line, start \(x_0 = 0\), goal \(g = 2\), actions in \([-1, 1]\), horizon \(H = 3\), \(N = 6\) plans, \(K_e = 2\) elites, two iterations, cost \(\sum_h |x_h - g|\).

The worked example, with the model equal to the world (\(x' = x + a\)). The first plan drawn is \(a = (1, -1, 0.336)\) — the first draw, \(0.8 \cdot 2.04 = 1.632\), was clipped to 1 — so it imagines \(x = (1, 0, 0.336)\) and costs \(|1-2| + |0-2| + |0.336-2| = 4.664\). The six costs are \(4.664,\ 8.264,\ 10.064,\ 2.912,\ 8.592,\ 6.360\); the elites are plans 4 and 1, so the new mean is \(\mu = (1,\ -0.408,\ 0.028)\) and the new spread \(\sigma = (0.05,\ 0.642,\ 0.358)\) — e.g. \(\sigma_2 = |0.184 - (-1)|/2 + 0.05\). After the second iteration the plan is \(\mu = (0.996,\ -0.058,\ 0.284)\), with cost 2.846. The optimum, \(a = (1, 1, 0)\), costs 1.0: six samples are a crude search, which is why PlaNet draws 1 000.

Now switch the panel to the learned model that is 25% too fast (\(x' = x + 1.25\,a\)). The same plan is chosen, the model promises a cost of 2.057, and the world delivers 2.846. MPC survives this: after executing \(a = 0.996\) the world is at \(x = 0.996\), not the promised \(1.244\), and the next decision starts from the truth. The panel below makes that difference visible over a whole episode.

With a model that misjudges directions by 15°, the plan executed blindly ends 0.43 from the goal — the gap grows by 0.13 every 5 steps — while re-planning at every step reaches it in the same 17 steps the model promised. Feedback is what makes an imperfect model usable at all. It is also why the lab measures each failure under MPC, the way the model is actually used.

Three ways it goes wrong

1. Compounding error

A model is trained on real inputs: \(z_t = e(o_t)\), then one step. It is used on its own outputs: \(\hat z_{t+2} = f(\hat z_{t+1}, a_{t+1})\), and so on for \(H\) steps. Each prediction becomes the input of the next, so an error made at step 1 is carried — and transformed by the model's own dynamics — into step 2. Talvitie put it plainly: "a flawed sample is given as input to the model, which causes more errors, and so on"¹². Venkatraman et al. pointed at the cause: the errors "change the input distribution for future prediction steps, breaking the train-test i.i.d assumption"¹³.

The simplest case already shows the growth. Let the true dynamics be \(x_{t+1} = x_t\) (nothing moves) and the model's be \(\hat x_{t+1} = 1.02\,\hat x_t\) — a 2% error per step. From \(x_0 = 1\), the error after \(h\) imagined steps is \(1.02^h - 1\): 0.02 after one step, \(1.02^{10} - 1 = 0.219\) after ten and \(1.02^{30} - 1 = 0.811\) after thirty. When the model's error is multiplied by its own Jacobian at every step, it compounds like interest.

The lab measures it on a learned model. Trained on one-step prediction, it locates a moving dot to within 0.0146 after one imagined step (the room is 2 wide; a step moves at most 0.14) and misses by 0.188 after thirty — 12.9 times more. Trained with a loss over 10 unrolled steps instead, it is slightly worse at one step (0.0192) and almost three times better at thirty (0.068). The objective decides what the model is good at.

2. Model exploitation

The planner is an optimizer, and it optimizes the model. Suppose \(N\) candidate plans all have the same true cost \(C\), and the model adds an independent error \(\varepsilon_i \sim \mathcal N(0, \sigma^2)\) to each. The planner picks the plan with the lowest predicted cost, \(C + \min_i \varepsilon_i\). For \(N = 1\,000\) the expected minimum of 1 000 standard normals is about \(-3.24\), so with \(\sigma = 0.1\) the chosen plan is predicted to cost \(C - 0.324\) while really costing \(C\). The planner has not found a better plan; it has found the plan on which the model is most wrong in its favour. More search makes it worse, not better.

In practice errors are not independent: they concentrate where the data were thin, and the planner goes there. In the lab, the wall of a small maze appears in 1.75% of the random-play transitions. Given the right cost, a planner in the learned model succeeds on 2 of 10 tasks that require going around the wall; 8 of its 10 first plans go through the wall in imagination, promising 0.68 of path left and delivering 0.94. Ha & Schmidhuber's fireball-free Doom is the same phenomenon.

The fixes either make the model honest about what it does not know, or go and get the data:

• Uncertainty. PETS¹⁷ trains an ensemble of 5 probabilistic networks and propagates particles through them, so that regions where the members disagree look risky instead of attractive; Ha & Schmidhuber's temperature is a cruder version of the same idea.
• Close the loop. Collect new data with the planner itself and retrain — the loop of Dyna, PlaNet, Dreamer and MuZero. The planner goes where the model is wrong, and the loop turns each mistake into a training example. In the lab, 600 transitions collected by the planner (72% of them hitting the wall) raise the success from 2 to 7 out of 10, and the imagined crossings drop from 8 to 2.

3. The train–plan gap

The third failure is not in the model at all. A planner ranks imagined futures with a cost, and in a latent world model that cost is usually a distance between the imagined embedding and the embedding of a goal image — L1 for V-JEPA 2-AC, squared L2 for DINO-WM. But nothing in the training of the encoder asked that distance to mean how far the goal is for an agent that has to get there.

The lab's maze makes the gap concrete. From \(s = (-0.6, -0.6)\) to \(g = (0.6, -0.6)\), with a wall between them up to height 0.3, the straight-line distance is \(1.2\) and the shortest path, bent at the top of the wall \((0, 0.3)\), is \(2\sqrt{0.6^2 + 0.9^2} = 2.163\). With actions that move at most 0.1 on each axis per step, that is 12 steps if the wall were not there and 18 around it. Two dots on either side of a wall make similar images, so their embeddings are close too: on 1 124 random pairs of states with the wall in between, the rank correlation of the latent distance is 0.93 with the Euclidean distance and only 0.41 with the length of the path.

A planner with a Euclidean cost and a 15-step horizon sees every plan that climbs toward the gap as getting worse, and presses against the wall. In the lab this fails even in the perfect simulator (1 task in 10); the same planner in the same simulator with the path-length cost solves 10 in 10. The model was never the problem there. Lambert et al. named the general version objective mismatch: "training the forward dynamics model w.r.t. the likelihood of the one-step ahead prediction" while the goal is "improving performance on a downstream control task"¹⁸.

Video world models, 2024–2026

A different line of work treats the world model as a video generator that takes actions: the predicted frame is the product, and you can play it.

These are the models behind most "world model" headlines, and the right comparison with the latent models above is a trade-off, not a ranking:

V-JEPA 2-AC measured the cost side directly. Planning one robot action took 16 seconds with the latent model and 4 minutes with Cosmos on the same GPU; on pick-and-place the latent planner succeeded in 80% (cup) and 50% (box) of the trials in that lab, the video model in none¹⁵. For planning, predicting the embedding is enough and pixels are a heavy luxury; for a person who wants to watch or play the world, pixels are the point.

Honest limits

• Long horizons are still the frontier. Every result above plans or imagines over seconds — 12 to 15 steps for PlaNet and Dreamer, a few seconds of robot motion for V-JEPA 2-AC. Tasks that need minutes of coherent behaviour rely on a learned value at the end of the rollout, or on hierarchy, and neither is solved.
• "Looks right" is not "is right". A video world model can be visually convincing and physically wrong; human raters barely telling GameNGen from DOOM says nothing about whether a robot trained in it would transfer. There is no agreed benchmark for physical correctness.
• Exploitation never goes away. It can be reduced with uncertainty and closed-loop data, never removed: an optimizer pointed at an imperfect model will find its imperfections.
• Goal-image planning needs a goal image. The latent planners above are given a picture of the finished task. Turning language or rewards into a planning cost is open.
• Much is not disclosed. For the largest systems (Genie 3, the closed video models) the architecture, size and training data are not public, so the claims about them are claims about demos.

Lab: learn a world, then plan in it

Every number on this page marked as measured comes from a Colab notebook that builds a world model end to end on a toy small enough to see whole: a dot in a 2-D room with a wall, observed as a 16 × 16 image. It runs on CPU in a few minutes, with nothing to download. Ten sections, each ending in a checkpoint:

1. the world, its true simulator and the two distances (straight line, around the wall);
2. random play, and how rarely it meets the wall;
3. a latent world model — encoder, decoder, residual transition — trained one step at a time;
4. compounding error, measured over 30 imagined steps;
5. the same model trained on 10-step rollouts;
6. MPC with CEM in the true simulator and in the learned model, against random actions;
7. the train–plan gap: latent distance against path length, and the planner with the right cost;
8. model exploitation: what the planner was promised and what it got;
9. closing the loop with data collected by the planner;
10. your turn: learn the cost, add pessimism, or tune the planner.

Open the lab in Colab — save a copy to your Drive (File → Save a copy in Drive) to run and edit it.

Key takeaways

1. A world model learns \(p(s_{t+1} \mid s_t, a_t)\), almost always in a latent space: an encoder, an action-conditioned transition and optional heads (pixels, reward, value). Agents use it to plan, to imagine training data, and to spend less real interaction.
2. The families differ in what the model must predict: pixels (V–M–C, RSSM, video models), only reward, value and policy (MuZero), or only the next embedding (TD-MPC2, DINO-WM, V-JEPA 2-AC).
3. MPC plans in the model, executes the first action and re-plans. CEM and MPPI search without gradients. Re-planning is what makes an imperfect model usable.
4. Compounding error: a model trained one step at a time is used on its own outputs. Errors carry over and grow; training on multi-step rollouts, short horizons and terminal values contain it.
5. Exploitation: the planner chooses the plans the model is most wrong about, in its favour. More search makes it worse. Uncertainty (ensembles) and closed-loop data collection are the remedies.
6. The train–plan gap: the planning cost — a distance between embeddings — was never trained to mean progress toward the goal. A perfect model with the wrong cost still fails; the next chapter trains the latent so that its distance measures progress.
7. Video world models (Genie 3, Cosmos, GameNGen, Dreamer 4) predict pixels and can be played; latent models predict embeddings and are far cheaper to plan with. Which is right depends on whether the output is for a person or for a planner.