World Models

01The trade you are making

Look at the two loops first. Every decision further down this page is a consequence of choosing the second one.

Model-free RL, the subject of page 15, learns a policy directly from reward. It is simple and it works. The cost is that every bit of improvement is paid for in real interactions.

That bill is trivial in some places and ruinous in others:

• An Atari emulator runs faster than real time and costs nothing but electricity, so you can spend 200 million frames.
• A physics simulator is nearly as cheap, but it only simulates what somebody wrote down.
• A robot arm runs at one times real time, wears out, and breaks things. Ten million steps would take about a decade, which is not a workable budget.
• A power grid or a live recommender cannot be explored at all. Bad actions have consequences nobody can take back.

Model-based RL changes what the real interactions buy. Rather than spending them on trial and error, spend them on fitting a model of the dynamics. Then do the trial and error inside the model, where a step costs a matrix multiply.

02What a world model is

Strip the neural networks away and three functions are left. Everything from Dyna in 1991 to Genie is a choice about how to represent them.

• Transition. Given a state and an action, predict the next state.
• Reward. Given a state, predict the reward. Usually learned, occasionally known in closed form.
• Continuation. Predict whether the episode ends here. Easy to forget, and imagined rollouts run straight off the end of the world without it.

A world model is one learned function standing in for the environment's step method.

Dyna, from Sutton in 1991, is the minimal version. Act in the real environment, store each transition, and update the value function on both real transitions and transitions replayed from the learned model. The model manufactures extra experience at no interaction cost.

The payoff is sample efficiency, and it is easiest to see on something small enough to reason about completely.

import numpy as np
rng = np.random.default_rng(0)
N = 20                       # a 20-state chain, reward at the far end

def step(s, a):              # a=1 moves right, a=0 moves left
    s2 = min(N - 1, s + 1) if a else max(0, s - 1)
    return s2, float(s2 == N - 1)

Q = np.zeros((N, 2))                # model-free: one backup per real step
T, R = np.full((N, 2), -1), np.zeros((N, 2))   # model-based: keep s', r
s = mf = mb = 0
for t in range(1, 100000):          # the SAME random data feeds both
    a = int(rng.integers(2))
    s2, r = step(s, a)
    Q[s, a] += .5 * (r + .95 * Q[s2].max() - Q[s, a])
    T[s, a], R[s, a] = s2, r
    if not mf and (Q[:-1].argmax(1) == 1).all(): mf = t
    if not mb and (T[:-1] >= 0).all():
        V = np.zeros(N)
        for _ in range(100): V = (R + .95 * V[T]).max(1)  # costs 0 steps
        if ((R + .95 * V[T]).argmax(1)[:-1] == 1).all(): mb = t
    s = 0 if r else s2
    if mf and mb: break
print("Q", Q.shape, " learned model T", T.shape, "R", R.shape)
print("model-free  policy correct after", mf, "real steps")
print("model-based policy correct after", mb, "real steps")

The model-based learner needed no extra exploration, no better algorithm and no bigger network. It only had to stop discarding each transition after using it once.

03Compounding error, and where the bill arrives

A model that is 99% right per step is nowhere near 99% right over a rollout, and this failure mode shapes most of the field.

Rolling a model out means feeding it its own predictions. At step two the input is already slightly wrong, so the output is wronger. By step three the input is wronger still. The trajectory walks off the distribution the model was fitted on.

It is the same mechanism that makes behavioural cloning drift on page 21, with the dynamics model in the policy's place.

import numpy as np
rng = np.random.default_rng(0)

# Ground truth: a gentle rotation. The model never sees this matrix.
A = np.array([[0.995, 0.100], [-0.100, 0.995]])
X = rng.normal(size=(50, 2))                  # 50 states, short rollouts
Y = X @ A.T + 0.02 * rng.normal(size=(50, 2))   # one noisy step each
Ahat = np.linalg.lstsq(X, Y, rcond=None)[0].T   # the entire world model

print("pairs", X.shape, "-> model", Ahat.shape)
print("worst one-step entry error %.4f" % np.abs(Ahat - A).max())

true = model = np.array([1.0, 0.0])
for h in range(1, 201):
    true = A @ true
    model = Ahat @ model        # imagine on, never corrected by data
    if h in (1, 5, 15, 50, 100, 200):
        drift = np.linalg.norm(model - true)
        print("horizon %3d   drift %.3f   %3.0f%% of true state"
              % (h, drift, 100 * drift / np.linalg.norm(true)))

For deterministic dynamics with Lipschitz constant $L$ and a model whose one-step error never exceeds $\varepsilon$, the gap after $H$ steps obeys a standard bound:

When $L \approx 1$ this collapses to $\varepsilon H$, linear and survivable, which is the benign case above where drift grew almost exactly linearly, at about 0.0048 per step. Any expansion in the dynamics pushes $L$ above 1, the sum turns geometric, and the usable horizon collapses.

The whole Dreamer line follows this rule: imagine briefly and constantly, and re-anchor on fresh real data all the time.

04Latent dynamics: predict in a space built for predicting

If the observation is a 64×64 colour image, predicting the next one means predicting 12,288 numbers, almost none of which matter for control.

The latent dynamics idea splits the problem in two. An encoder maps an observation to a compact state. The dynamics are learned inside that compact space. A decoder exists mainly to supply a training signal, and planning never calls it.

Three things improve at once. The model has far fewer parameters, so it fits from less data. A rollout becomes a small matrix multiply and stops being a generated image. And the latent is trained to be predictable, which raw pixels never were.

import numpy as np
rng = np.random.default_rng(0)
D, K, N = 64, 2, 100         # 64-dim observations of a 2-dim world
A = np.array([[0.99, 0.12], [-0.12, 0.99]])   # true latent dynamics
W = rng.normal(size=(K, D)) / np.sqrt(K)      # latent -> observation
Z = rng.normal(size=(N, K))
O = Z @ W + 0.05 * rng.normal(size=(N, D))
On = (Z @ A.T) @ W + 0.05 * rng.normal(size=(N, D))

# (a) predict the observation itself: D*D weights from N pairs
Mpix = np.linalg.lstsq(O, On, rcond=None)[0].T
# (b) find a K-dim latent with an SVD, then predict inside it
E = np.linalg.svd(O, full_matrices=False)[2][:K]
Mlat = np.linalg.lstsq(O @ E.T, On @ E.T, rcond=None)[0].T
print("pixel model", Mpix.shape, " latent model", Mlat.shape)

z = rng.normal(size=K)               # a fresh start, unseen in training
pix, lat = z @ W, (z @ W) @ E.T
for h in range(1, 21):
    z, pix, lat = A @ z, Mpix @ pix, Mlat @ lat
    if h in (1, 5, 20):
        t = z @ W
        print("h=%2d  pixel err %.3f   latent err %.3f"
              % (h, np.linalg.norm(pix - t) / np.linalg.norm(t),
                 np.linalg.norm(lat @ E - t) / np.linalg.norm(t)))

Real systems learn the encoder instead of taking an SVD, and the latent carries a hundred or so dimensions instead of two, but the argument is the same: a learned bottleneck makes long-horizon prediction tractable.

05PlaNet, the RSSM, and the Dreamer line

One architecture family has carried this idea for six years. It helps to know what each version added and why.

PlaNet (2018) learned latent dynamics from pixels and planned directly in the latent space. Its lasting contribution is the recurrent state-space model (RSSM), which keeps two things in the state.

• A deterministic path, a recurrent hidden state that carries information forward reliably. Without it the model forgets what it just saw.
• A stochastic path, a sampled latent that represents genuine uncertainty. Without it the model cannot say "the ball might bounce either way" and averages the two futures into mush.

• h_t is the deterministic state. It uses plain recurrence with no sampling, so information survives many steps.
• q is the posterior. It sees the actual observation $o_t$, so it is only available while training.
• p is the prior. It must guess $z_t$ from $h_t$ alone. Training pulls the prior toward the posterior.

The last line is the trick. Once the prior can predict the posterior without seeing the next frame, you can run the model forward with no observations at all, which is what imagination means here.

Dreamer (2019) changed what you do with the model. PlaNet ran a fresh search at every step. Dreamer instead trains an actor and a critic on imagined latent trajectories, backpropagating value gradients through the learned dynamics. An expensive online search is replaced by a network that has already done it.

• Dreamer (2019) learns an actor-critic purely inside latent imagination, on visual control tasks.
• DreamerV2 (2020) uses discrete categorical latents in place of Gaussian ones, which handle multimodal futures better. It was the first agent to reach human-level on the 55-game Atari benchmark by learning behaviours inside a separately trained world model.
• DreamerV3 (2023) adds normalisation and transformation fixes that make a single hyperparameter setting work across more than 150 tasks. Applied out of the box, it was the first system to collect diamonds in Minecraft from scratch, with no human data and no curriculum.
• DayDreamer (2022) runs the same algorithm online on physical robots. A quadruped learned to roll off its back, stand and walk in one hour of real experience, without a simulator and without resets.

DayDreamer is the cleanest evidence for the claim at the top of this page. One hour of wall-clock experience is the kind of budget a physical robot has.

06Planning inside the model

A world model hands you a cheap simulator. What you do with it is a separate choice, and there are three standard answers.

• Sample and score. Draw hundreds of action sequences, roll each through the model, keep the best few, refit the sampling distribution, repeat. This is the cross-entropy method and needs no policy network. PlaNet and PETS work this way.
• Train a policy in imagination. This is Dreamer's answer: amortise the search into an actor and a critic so that acting is one forward pass.
• Search with a learned value. Monte Carlo tree search over a learned model, which is MuZero.

MuZero is the interesting case, because it declines to model the world. Its learned dynamics predict three quantities and no others: reward, value and policy. It never reconstructs a frame or a board position.

This is value equivalence: the model has to be right about the quantities the search consumes, and any detail that cannot change a value is free to be wrong.

The same machinery gives you counterfactual prediction. The model answers "what would have happened had I turned left" without turning left, which is how an agent compares options it never took. Generating counterfactuals is cheap, but being right about them is hard.

07Video prediction as a world model

A model that predicts the next frame of a video is a world model with no actions and no reward. Adding the actions turns it into something a policy can use.

Action-conditioned video prediction conditions the next frame on the past frames and the action taken. Finn and Levine did this for robot arms in 2016 under the name visual foresight, then planned pushes by searching over action sequences and scoring the predicted frames.

The modern versions are generative video models, and they are far better at frames:

• GameNGen (2024) trained a diffusion model to predict the next DOOM frame from past frames and inputs. It runs at 20 frames per second on a single TPU, and human raters are barely better than chance at telling short clips from the real game.
• Genie (2024) learned a latent action space from unlabelled internet video. It can turn a single image or sketch into a controllable environment, with no action labels anywhere in training.
• GAIA-1 (2023) and NVIDIA's Cosmos platform do the same for driving and physical AI, generating futures conditioned on actions and text.

The live question through 2025 and 2026 is whether the video route reaches control. Frame quality and real-time interactivity have improved enormously. Long-horizon consistency and physical reasoning have improved much less, and those are exactly the properties a planner leans on. Page 26 tracks where that has got to.

08JEPA: predict representations, not pixels

There is a third position worth taking seriously: predicting pixels is the wrong objective in the first place.

A pixel loss spends most of its capacity on things that are both unpredictable and irrelevant. Leaf texture, sensor noise, the exact shape of a shadow. A model that gets those wrong is penalised as hard as one that misses the oncoming car.

A joint-embedding predictive architecture (JEPA) predicts in representation space instead. Encode the context, encode the target, and train a predictor to match the target's embedding. There is no decoder and no reconstruction.

The obvious objection is collapse. If the encoder emits a constant, prediction is perfect and nothing has been learned. JEPA variants avoid this with asymmetry: a stop-gradient, a slowly-updated target encoder, or both. Page 03 covers those tricks in their original setting.

V-JEPA 2 (2025) carried this into control. Pre-train on over a million hours of internet video with no action labels, then post-train an action-conditioned latent model on under 62 hours of unlabelled robot video, and plan with it zero-shot on a real Franka arm.

The bet is to learn the structure of the world from video nobody labelled, then attach actions using a small amount of interaction data.

09Evaluating a world model

The metric you reach for first is the one most likely to mislead you.

A world model has two very different jobs, and they are measured in incompatible ways:

• Prediction quality. One-step and multi-step error, PSNR, FVD for video. Cheap, reproducible, and only loosely related to usefulness.
• Downstream control utility. Does a policy trained or planned inside the model perform in the real environment? Expensive, noisy, and the only thing that settles the question.

Three specific things to probe, because an aggregate error number hides all of them:

• Long-horizon consistency. Roll out for hundreds of steps and check the scene is still the same scene. Does the room still have the same number of doors?
• Object permanence and physical reasoning. Occlude something, wait, reveal it. Benchmarks such as IntPhys and Physion probe this directly, and models that look excellent frame to frame fail it routinely.
• Counterfactual accuracy. Take one start state and two different action sequences. The model has to predict different futures, and not a blurred average of both.

One more catches teams late: error where the policy goes. A model fitted on data from the old policy is accurate on the old policy's states. Training a new policy inside it moves the state distribution to exactly the region the model was never fitted on.

10Build this

The horizon trade-off on this page is a curve with a peak in it. You can plot that curve in an afternoon, on a laptop, in numpy.

11What breaks

World-model failures are rarely a bad architecture. They are a model being trusted somewhere it was never fitted.

• The policy finds a hole in the model. Imagined return climbs steadily while real return is flat or falling. The agent has located a state where the model promises reward the environment does not pay.
• One-step error looks superb and 50-step rollouts are garbage. Validating on one-step transitions is the default, and it certifies nothing about rollouts.
• The model is accurate everywhere the old policy went. Then the new policy goes somewhere else, and the model has no data there. Each round of improvement invalidates part of the model.
• Imagined episodes never end. The continuation head was skipped, so the agent imagines collecting reward past a terminal state and learns to run into walls.
• The stochastic latent collapses. The posterior stops carrying information, the model turns deterministic, and every uncertain future is predicted as its own blurry average.
• Long rollouts lose the scene. Objects change colour, doors vanish, counts drift. A video world model can be excellent frame to frame and no longer the same world after thirty seconds.
• Reward is the hardest head to fit and the one that matters most. With sparse reward almost every imagined transition has reward zero, so the head predicts zero everywhere and imagination becomes worthless.
• Weeks go into the model's own metrics. Reconstruction loss falls steadily and control performance does not move, because nobody was running the policy in the real environment.

12Where you meet this in the wild

Same idea, four very different reasons for wanting it.

13Interview questions

14Go deeper

●Now write it yourself

Reading the derivation and being able to produce it are different skills. These are Deep-ML problems that exercise what this page covers — each one is checked against real test cases, not multiple choice.

Matched to this page from Deep-ML's catalogue of 1,380 problems. More at deep-ml.com, and Where to practise covers the other platforms and what each one trains.

15Related pages

Other techniques for this problem

A scoped slice of the full Technique Map — every technique this page covers, grouped by what it solves.