Artificial Neural Networks and Deep Learning · Insper

World Models

Learn the world, then plan in it

Start with the cost of a lesson

A robot arm must learn to pick up a cup. Each real attempt takes a minute and may break something.
What would you want before trying?

  • A way to try actions in its head: predict what each one does, rank them, then act.
  • That is a world model. PlaNet, one of the first that works from pixels, trained for 1,000 episodes on each of its six control tasks and came close to the model-free D4PG, which used 100,000.

"If the organism carries a 'small-scale model' of external reality and of its own possible actions within its head, it is able to try out various alternatives…" — Craik, 1943

What a world model is

symbol meaning
observation at time (a frame)
encoder: frame → latent state
action taken at time
, transition: the model of
the hat = something the model imagined

Optional heads on the latent: decoder (pixels), reward, value, policy. Which heads exist, and what the loss compares, separates the families.

Why an agent wants one

  1. Plan — search action sequences in the model, execute the best. The goal can change at test time.
  2. Imagine — generate synthetic experience and train a policy on it (Sutton's Dyna, 1991; every Dreamer).
  3. Save real interaction — the expensive resource.

The price: every decision now passes through something that can be wrong.

This class is about where it is wrong and what that does to a planner.

Diagram · five families, one pipeline

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

part network params
V VAE: 64×64 → 4,348,547
M MDN-RNN: LSTM 256, 5 Gaussians 422,368
C one linear layer 867

CarRacing: 906 ± 21 (solved = 900). VizDoom: C trained entirely inside M's dream.

"Cheating the world model": in the dream, C found a policy where the monsters never shoot a fireball.

M's temperature dream real game
0.10 2,086 193
1.15 918 1,092

Ha & Schmidhuber (2018). The score is the number of time steps survived.

The RSSM · PlaNet and Dreamer

  • : deterministic memory (a GRU). : stochastic part. The paper found both crucial.
  • Loss = reconstruct and + — the VAE recipe stretched over time.
  • The KL is what makes the prior, the only thing available in imagination, agree with what the observation revealed.

PlaNet plans with CEM: horizon 12, 10 iterations, 1,000 candidates, best 100 kept.

Dreamer · learn behaviour in imagination

An actor and a critic trained only on imagined rollouts of steps; the value's gradient flows back through the model.

The critic regresses a λ-return ( = discount, = how far to trust imagined rewards):

  • : trust the critic after one step.
  • : trust the imagined rewards to the horizon — and every error of the model with them.
  • Dreamer: , .

Exercise · a λ-return by hand

, , . Imagined rewards , . Critic values , .

Compute , and .

A λ-return by hand · answer

step computation value
3
2.97
3.892
  • of the weight goes to the critic at each step, to the rest of the imagined return.
  • The target mixes imagined rewards and critic estimates: the critic caps how far a wrong model can push it.

The Dreamer line

what it added
DreamerV2 (2021) 32 categorical variables × 32 classes; KL balancing (). Human level on 55 Atari games, one GPU
DreamerV3 (2023; Nature 2025) one set of hyperparameters for 150+ tasks: , free bits (no KL below 1 nat). First to collect diamonds in Minecraft from scratch
Dreamer 4 (2025) a 2B-parameter video world model, 21 frames/s on one GPU. Trained on 2,541 h of recorded play; the agent learns only inside it and still gets diamonds (0.7% of episodes)

MuZero · a model of what matters

  • : representation · : dynamics · : prediction ( = policy, = value).
  • Loss: reward + value + policy, along unrolled steps. No reconstruction.
  • Planned in with Monte Carlo tree search: 800 simulations per move (board games), 50 (Atari).

A model judged only by the decisions it supports — free to ignore the background, the texture of the board. Matched AlphaZero at Go, chess and shogi without being given the rules.

Latent predictive models · no pixels at all

Predict the next embedding, not the next frame — the JEPA idea, with actions.

model encoder planning
TD-MPC2 (2024) learned, + reward and value MPPI · one 317M agent, 80 tasks
DINO-WM (2025) frozen DINOv2 patches CEM · MSE to the goal-image embedding
V-JEPA 2-AC (2025) video encoder, 1M+ hours of video; < 62 h of robot video CEM · L1 to the goal-image embedding · zero-shot on Franka arms

The planning cost is a distance between embeddings. Remember this — it comes back.

Planning in the model · MPC with CEM

model: imagine N planssearch H actions in the modelworldexecute the first actionworldobserve the real stateMPC: repeat every step

CEM, times:

  1. draw plans
  2. roll each out in the model, cost
  3. refit to the cheapest (the elites)

MPPI: soft weights instead of top-.

No gradient anywhere: the model only runs forward.

Simulator · one MPC decision, line by line

world: x′ = x + a, a ∈ [−1, 1] · start x₀ = 0 · goal g = 2

Exercise · one CEM iteration

, goal , cost , model , .
Which plans are the elites, and what are the new and ?

plan
1 1.000 −1.000 0.336
2 −0.456 −0.360 −0.176
3 −1.000 −0.184 −0.696
4 1.000 0.184 −0.280
5 −0.224 −0.536 −0.848
6 −0.312 0.384 −0.192

σ is the standard deviation of the elites + 0.05.

One CEM iteration · answer

plan imagined cost
1 1, 0, 0.336
4 1, 1.184, 0.904
others — 8.264, 10.064, 8.592, 6.360
  • Elites: plans 4 and 1.
  • — e.g. .
  • — e.g. .
  • After a second iteration: , cost 2.846. The optimum costs 1.0 — six samples are a crude search; PlaNet draws 1,000.

Simulator · open loop or MPC

Failure 1 · compounding error

Trained on real inputs, used on its own outputs: .

True , model , :

error
1 0.02
10 0.219
30 0.811

Measured in the lab (dot position, room 2 wide):
trained 1 step 0.0146 0.188
trained 10 steps 0.0192 0.068

The objective decides what the model is good at.

Remedies: train on the model's own rollouts · short rollouts from real states (MBPO) · a learned value at the end (TD-MPC, Dreamer) · re-plan every step.

Failure 2 · the planner exploits the model

plans, all with true cost . The model adds an independent error . Which predicted cost does the chosen plan have?

  • The planner takes . The minimum of 1,000 standard normals is about .
  • Predicted: . Real: . It found the plan the model is most wrong about, in its favour.
  • More search makes it worse.

In the lab, the wall appears in 1.75% of the random-play data. With the right cost, the planner in the learned model solves 2 of 10 wall tasks; 8 of 10 imagined plans go through the wall (promised 0.68 of path left, real 0.94).

Fixing exploitation

make the model honest

  • Ensembles (PETS: 5 probabilistic networks): where members disagree, the plan looks risky instead of attractive.
  • Ha & Schmidhuber's temperature is a cruder version.

go and get the data

  • Close the loop: act with the planner, record what happened, retrain — Dyna, PlaNet, Dreamer, MuZero.
  • Lab: 600 planner transitions, 72% of them hitting the wall.
  • Wall tasks 2 → 7 of 10; imagined crossings 8 → 2.

Failure 3 · which goal is closer?

In a room with a wall up to , the agent is at and the goal at , on the other side.
Each step moves at most 0.1 on each axis. How far is the goal — for the eye, and for the agent?

  • Two dots on either side of a wall make similar images. What will the encoder say?

Failure 3 · the train–plan gap

sggapwall (up to y = 0.3)straight line: 1.2 · 12 stepsaround the wall: 2.163 · 18 stepss = (−0.6, −0.6) · g = (0.6, −0.6)2√(0.6² + 0.9²) = 2.163

  • Straight line 1.2 (12 steps); around the tip 2.163 (18 steps).
  • Lab, 1,124 pairs with the wall between them — rank correlation of the latent distance: 0.93 with Euclidean, 0.41 with path length.
  • Perfect simulator, Euclidean cost: 1/10. Same simulator, path-length cost: 10/10.

The model was never the problem there. The cost was never trained to mean progress.

The gap the next class closes

"JEPA-style training optimizes short-horizon latent prediction, whereas planning requires a multi-step ranking of imagined futures by goal progress." — Temporal-Distance JEPA (2026)

  • Lambert et al. (2020) named the general form objective mismatch: train for one-step likelihood, hope for control performance.
  • TD-JEPA's remedy: learn, from reward-free trajectories, a cost that measures temporal progress — steps along the same trajectory as positives, other trajectories as negatives — and plan with it.
  • The next class, TDWM, fixes the representation instead: frames close in time are trained to be close in the latent.

Video world models, 2024–2026

model numbers that are public
GameNGen (DOOM) 20 frames/s on one TPU · PSNR 29.4 · raters 58–60% right
Genie 11B params · 30,000 h of 2-D platformer video · 8 latent actions · ~1 frame/s
Oasis ViT autoencoder + DiT · 20 frames/s · 500M weights released
Genie 2 / 3 consistent "up to a minute" / 720p, 24 frames/s, "a few minutes" · Genie 3 size: not disclosed
Cosmos ~100M clips from 20M h of video · diffusion 7B/14B, autoregressive 4B/12B
Dreamer 4 2B params · 21 frames/s on one H100 · an agent trained inside it

Pixels or embeddings?

generative (pixels) latent predictive
predicts the next frame the next embedding
cost per imagined step a denoising pass on a frame a small network on a vector
errors look like visual drift nothing visible: geometry
hard part consistency over minutes; physics that only looks right a cost that means progress

V-JEPA 2-AC vs. Cosmos, same GPU: 16 s vs. 4 min per robot action; pick-and-place 80% / 50% (cup / box) vs. 0% / 0%.
For a planner, embeddings suffice; for a person who watches or plays, pixels are the point.

Honest limits

  • Long horizons are still the frontier: PlaNet and Dreamer imagine 12–15 steps; V-JEPA 2-AC, a few seconds of motion.
  • "Looks right" ≠ "is right": raters barely telling GameNGen from DOOM says nothing about physics. No agreed benchmark.
  • Exploitation never goes away — reduced by uncertainty and closed-loop data, never removed.
  • Goal-image planning needs a goal image; language or rewards as a planning cost is open.
  • Much is not disclosed: for Genie 3 and the closed video models, architecture, size and data are not public.

Lab · learn a world, then plan in it

Colab notebook, ~2 h, CPU only, nothing to download: a dot in a 2-D room with a wall, seen as a 16×16 image.

  1. the world and its two distances
  2. random play — the wall in 1.75%
  3. encoder + decoder + transition
  4. compounding error over 30 steps
  5. training on 10-step rollouts
  1. MPC with CEM: random, true sim, model
  2. the train–plan gap, measured
  3. exploitation: promise vs. reality
  4. closing the loop
  5. your turn: learn the cost, add pessimism or tune the planner

open in Colab →

Takeaways

  1. A world model learns in a latent space, to plan, to imagine and to save real interaction.
  2. Families differ in what is predicted: pixels (V–M–C, RSSM, video), decisions (MuZero), embeddings (TD-MPC2, DINO-WM, V-JEPA 2-AC).
  3. MPC plans, executes one action, re-plans; CEM/MPPI search without gradients.
  4. Compounding error: trained on real inputs, used on its own outputs. Train on rollouts; keep them short.
  5. Exploitation: the planner picks the model's favourable mistakes. Ensembles and closed-loop data.
  6. Train–plan gap: the cost was never trained to mean progress — next class.
  7. Pixels for people to watch and play; embeddings for planners.

Outline: (1) what a world model is and why an agent wants one; (2) the families: V–M–C, RSSM and Dreamer, MuZero, latent predictive models; (3) planning in the model with MPC and CEM; (4) the three ways it goes wrong: compounding error, exploitation, the train–plan gap; (5) video world models and honest limits. Closes with the Colab lab.

The PlaNet numbers are Table 1 of the paper: 100× fewer episodes than D4PG. The paper's "200×" is a different measure: how many more episodes D4PG needs to reach PlaNet's final score, estimated per task (40× to 500×) and averaged.

Step through the tabs. Ask at each one: what does the loss compare? Pixels (V–M–C, RSSM, video), decisions (MuZero), embeddings (JEPA-style).

Walk the 14 steps with the true model. Then switch to the model 25% too fast and jump to line 8: promised 2.057, real 2.846.

Default: 15° of direction error. Blind plan: 0.43 from the goal, the gap grows 0.13 every 5 steps. MPC: reaches it in 17 steps. Try 30° and −30% speed.