Lesson 18.2 · 28 min

World Models: Teaching AI to Simulate Its Environment

Before reaching for a hot pan, you already “see” what would happen if you grabbed it bare-handed. What would it take for an AI to imagine consequences before acting?

In short: A world model is a learned simulator: given the current state and an action, it predicts the next state (and often the reward). It is trained from recorded experience, usually in a compressed latent state rather than raw pixels. Once learned, an agent can roll out imagined futures inside the model to plan or to train its policy, needing far fewer real-world trials. Dreamer-style agents, MuZero, video world models such as Genie, and driving and robotics simulators all build on this idea, with compounding prediction error as the main limit.

Environment, state and action

We need three words from reinforcement learning (RL). The environment is the world the agent lives in: a video game, a robot's room, a road. The state is a description of the environment at one moment: a car's position and speed, or the pixels on the game screen. An action is what the agent can do: press left, turn the wheel, push with some force. After each action the environment moves to a new state and may give a reward, a number saying how good that was.

Running example: a small cart on a track that must reach position 5. Its state is (position, velocity); its action is a push between −1 and +1; friction slows it down. The agent does not know the physics in advance.

What is a world model? The human analogy

A world model is the agent's own learned approximation of the environment's dynamics: a function f̂ that, given a state and an action, predicts the next state (and often the reward). It is an internal simulator of how the world responds.

Imagining a move before making it A chess player thinks “if I move my knight there, they will take it with the pawn”. A driver thinks “if I brake now, the car behind might hit me”. We run little simulations in our heads, choose the best imagined outcome, and only then act. We learned the simulator from years of experience; nobody gave us the equations.

Two ways an agent can use a world model: planning (search through imagined action sequences at decision time and pick the best first action) and learning in imagination (train a policy on imagined experience instead of, or in addition to, real experience). Agents that use a model are called model-based; agents that learn directly from real trials without one are model-free.

Why we need a world model

• Sample efficiency: real trials are slow, expensive or dangerous (a robot arm, a car). Imagined trials are cheap: thousands per second on a GPU.
• Safety: we can test risky actions in imagination first. Crashing a simulated car costs nothing.
• Planning and foresight: looking several steps ahead helps in tasks where a good move only pays off later.
• Generalisation: a model of how the world works can be reused for new goals; a model-free policy is tied to the reward it was trained on.
• Understanding: predicting the future forces the system to learn objects, motion and cause and effect.

How a world model learns: predicting the next state

Learning a world model is supervised learning on recorded experience. We collect transitions (s, a, s′, r) by acting in the environment (even randomly at first). Then we train a model to map (s, a) to s′ and r, minimising the prediction error. The data labels itself: the next state is simply what happened.

The basic recipe

1. Collect: Act in the real environment and store many transitions: state, action, next state, reward.
2. Fit: Train f̂(s, a) ≈ s′ (and R̂(s, a) ≈ r) by minimising squared error or a likelihood, like any regression model.
3. Use: Plan or train a policy using f̂ instead of the real environment.
4. Act and refresh: Act in the real world with the plan, record the new transitions, and refit the model, so it improves where the agent actually goes.

Real environments are often stochastic (the same action can lead to different outcomes), so good world models predict a distribution over next states, not just one guess. For our cart, a simple linear model is enough; for video games and robots, deep neural networks are used.

The latent state: compressing what we see

When the observation is an image, predicting the next image pixel by pixel is expensive and mostly wasted on irrelevant detail (as we saw in the JEPA lesson). So most modern world models first compress each observation into a small latent state z, a vector of, say, a few hundred numbers, and learn dynamics in that latent space.

The influential 2018 paper World Models by David Ha and Jürgen Schmidhuber made this concrete with three parts: V (vision), a variational autoencoder that compresses each game frame into a small z; M (memory), a recurrent network that predicts the next z given the current z and action; and C (controller), a tiny policy that acts from z and M's hidden state. They even trained the controller entirely inside M's “dream” for a VizDoom task and then ran it in the real game.

Imagining the future: rollouts without touching the real world

A rollout is a simulated trajectory: start from the current state, pick an action, ask the model for the next state, pick another action, and so on for H steps (the horizon). By scoring many rollouts we can pick good actions. A simple planning method, random shooting with model predictive control (MPC), tries many random action sequences in the model, executes only the first action of the best one, observes the real result, and plans again.

world_model_cart.py

import numpy as np
rng = np.random.default_rng(1)
def real_env(s, a):                   # hidden physics: s = [position, velocity]
pos, vel = s
vel = 0.9 * vel + 0.5 * a         # friction + push
return np.array([pos + vel, vel])
# 1) Collect 200 random transitions (s, a, s') from the real world
S, A, S2 = [], [], []
s = np.zeros(2)
for _ in range(200):
a = rng.uniform(-1, 1)
s2 = real_env(s, a) + rng.normal(0, 0.01, 2)   # small sensor noise
S.append(s); A.append(a); S2.append(s2)
s = s2 if abs(s2[0]) < 10 else np.zeros(2)
X = np.column_stack([np.array(S), np.array(A)])     # inputs [pos, vel, a]
W, *_ = np.linalg.lstsq(X, np.array(S2), rcond=None)  # 2) learn s' ≈ [s, a] @ W
model = lambda s, a: np.array([*s, a]) @ W
print("learned vel' =", np.round(W[:, 1], 2), "(true 0, 0.9, 0.5)")
# 3) Plan by imagination: try 300 random 8-step action plans inside the model
def plan(s, goal=5.0, n=300, horizon=8):
best, best_cost = None, np.inf
for _ in range(n):
acts, x, cost = rng.uniform(-1, 1, horizon), s.copy(), 0.0
for a in acts:
x = model(x, a); cost += (x[0] - goal) ** 2   # imagined, no real step
if cost < best_cost: best, best_cost = acts, cost
return best[0]                    # do only the first action, then re-plan
s = np.zeros(2)
for t in range(12):                   # 4) act in the real world, re-planning each step
s = real_env(s, plan(s))
print(f"position after 12 real steps: {s[0]:.2f} (goal 5.0)")
print("imagined steps used:", 12 * 300 * 8, "| real steps used:", 200 + 12)

Output:

learned vel' = [0.  0.9 0.5] (true 0, 0.9, 0.5)
position after 12 real steps: 4.85 (goal 5.0)
imagined steps used: 28800 | real steps used: 212

The danger of imagination is compounding error. A small one-step mistake feeds into the next prediction, which adds its own mistake, and so on. Over long horizons, imagined trajectories can drift far from reality, and a planner may even exploit model errors (finding “shortcuts” that only exist in the flawed model).

Pause and think: Why does the cart code execute only the first action of the best plan and then plan again?

Because the model is imperfect and errors compound along the horizon. Re-planning from the real observed state after every step corrects drift and keeps the agent grounded in reality, which is the core idea of model predictive control.

Dreamer-style agents that plan inside the model

The Dreamer family by Danijar Hafner and colleagues (Dreamer in 2019, DreamerV2 in 2020, DreamerV3 in 2023) is the best-known line of agents that learn behaviour inside a latent world model. Instead of searching at decision time like our cart planner, Dreamer trains an actor (policy) and a critic (value estimator) on imagined trajectories.

Dreamer's three interleaved loops

1. Learn the world model: From replayed real experience, learn a recurrent latent model (the recurrent state-space model, RSSM, introduced in the earlier PlaNet work) that encodes observations, predicts next latent states, rewards and episode ends, and can reconstruct observations as a training signal.
2. Learn behaviour in imagination: Start from latent states seen in real data, roll out the actor for a short horizon (around 15 steps) inside the model, and train the actor to maximise the critic's predicted returns; train the critic to predict those returns.
3. Act in the real environment: Use the actor to collect new real experience, add it to the replay buffer, and repeat.

DreamerV3 was notable for working across many different domains with one fixed set of hyper-parameters, and for being reported as the first algorithm to collect diamonds in Minecraft from scratch, without human demonstrations or a hand-made curriculum. A related landmark, DeepMind's MuZero (2019–2020), learned a model that predicts only what planning needs (reward, value and policy) rather than observations, and used it with tree search to master Go, chess, shogi and Atari without being told the rules.

World models and predicting the future

Since 2024 the term “world model” has also been used for large video world models: generative models trained on huge amounts of video that predict future frames, sometimes conditioned on actions. OpenAI described its Sora video model as a step toward “world simulators”. Google DeepMind's Genie models learn interactive environments from video, with Genie 3 (2025) generating explorable worlds in real time from a text prompt. NVIDIA's Cosmos platform offers world foundation models aimed at robotics and autonomous driving. Meta's V-JEPA 2 (previous lesson) takes the non-generative route, predicting in embedding space.

Milestones in world models

1. World Models (Ha & Schmidhuber): VAE + recurrent model + small controller; a policy trained inside the model's dream.
2. PlaNet and Dreamer: Recurrent latent state-space models; planning and then actor-critic learning in imagination.
3. MuZero: A learned model of value, policy and reward used with tree search to master board games and Atari without the rules.
4. DreamerV3: One configuration across many domains; Minecraft diamonds from scratch.
5. Video world models: Large generative and predictive video models (Sora, Genie 2 and 3, Cosmos, V-JEPA 2) used as simulators and for planning.

Whether a video generator really “understands” physics is debated: such models can produce plausible-looking clips that break physical laws, and looking right is not the same as predicting correctly under new actions. For control, what matters is accuracy on the consequences of actions, not visual quality.

World models in the real world

Where they are used Robotics: learning manipulation with fewer physical trials, and planning toward goal images. Autonomous driving: generating rare or dangerous scenarios (a child running into the road) to train and test driving systems; Wayve's GAIA models are an example. Games and simulation: generating interactive environments for training agents. Industry: model predictive control with learned dynamics for processes such as heating and cooling systems.

Limits and common mistakes Compounding error over long horizons; policies that exploit model flaws; models that are accurate only where data was collected and fail on new situations (distribution shift); and confusing pretty video with correct dynamics. Mitigations: short horizons with frequent re-planning, uncertainty estimates (for example ensembles of models), continual data collection, and always validating in the real world before trusting a plan.

When not to use one: if a fast, accurate simulator already exists (many board games, some physics tasks), we can use it directly; and if real data is cheap and plentiful, a model-free method may be simpler and reach higher final performance.

Worked example, step by step

The cart code tried 300 random plans of 8 steps. Let us shrink that to something we can do on paper: three plans of two steps. The cart starts at position 0 with velocity 0, and the goal is position 5. We use the model the agent learned: new velocity = 0.9 × velocity + 0.5 × push, and new position = position + new velocity. The cost of a plan is the sum of (position − 5)² over its imagined steps.

From imagined plans to one real action

1. Imagine plan A: Step 1: velocity = 0.9 × 0 + 0.5 × 1 = 0.5, position = 0.5. Step 2: velocity = 0.9 × 0.5 + 0.5 = 0.95, position = 1.45. Cost = (0.5 − 5)² + (1.45 − 5)² ≈ 32.85.
2. Imagine plans B and C: The same arithmetic gives costs of about 40.95 and 59.95. Plan C starts by pushing the wrong way and never recovers within two steps.
3. Pick the cheapest: Plan A wins. No real step has been taken yet; all six transitions happened inside the model.
4. Act once: We execute only A's first push, +1, in the real world, and read the real new state.
5. Throw the rest away and re-plan: The second push of plan A is discarded. From the real new state we imagine fresh plans. This is model predictive control.

This tiny example also shows a limit of short horizons. With only two imagined steps, the planner sees no reason to brake: pushing hard always looks best because the goal is far away. Near the goal, a cart moving at high speed will overshoot, and a planner can only avoid that if its horizon is long enough to see the overshoot. Too short a horizon gives greedy, short-sighted plans. Too long a horizon lets model errors pile up. Choosing the horizon is a balance between the two.

How to check a learned model before trusting its plans Hold back some real transitions the model never trained on. Compare its one-step predictions with what really happened, then do the same for 5-step and 10-step rollouts. If the error grows quickly with the number of steps, keep the planning horizon short and re-plan often.

Practice: try it yourself

We will measure compounding error directly. The real cart has friction 0.9. Our pretend learned model has it slightly wrong: 0.93. We push the same 20-step plan through both and compare positions. Then we let the model look at reality every few steps and see how much that helps.

practice_compounding_error.py

import numpy as np
def real_step(pos, vel, a):          # the true physics (hidden from the agent)
vel = 0.9 * vel + 0.5 * a
return pos + vel, vel
def model_step(pos, vel, a):         # the learned model: friction slightly wrong
vel = 0.93 * vel + 0.5 * a
return pos + vel, vel
rng = np.random.default_rng(4)
actions = rng.uniform(-1, 1, 20) + 0.3    # one fixed 20-step plan, mostly pushing right
def imagine(resync_every):
"""Roll the plan through both worlds. Every `resync_every` steps the model
is reset to the real state, as if the agent had looked at the world again."""
real = model = (0.0, 0.0)
errors = []
for t, a in enumerate(actions, start=1):
real = real_step(*real, a)
model = model_step(*model, a)
errors.append(abs(model[0] - real[0]))    # position error at step t
if t % resync_every == 0:
model = real                          # look at reality, start again
return errors
open_loop = imagine(resync_every=99)              # never looks: pure imagination
print("steps ahead:       ", "  ".join(f"{t:5d}" for t in (1, 2, 5, 10, 20)))
print("open-loop error:   ", "  ".join(f"{open_loop[t - 1]:5.2f}" for t in (1, 2, 5, 10, 20)))
for k in (5, 1):
e = imagine(resync_every=k)
print(f"resync every {k}: worst error over 20 steps = {max(e):.2f}")

Output:

steps ahead:            1      2      5     10     20
open-loop error:     0.00   0.02   0.22   1.04   5.18
resync every 5: worst error over 20 steps = 1.01
resync every 1: worst error over 20 steps = 0.09

Now change it:

• Make the model worse: change 0.93 to 0.99. Predict whether the open-loop error at 20 steps grows by a little or by a lot.
• Give the model a different flaw: restore 0.93 to 0.9 and change the model's push strength from 0.5 to 0.45. Predict whether the error at step 1 is still 0.00.
• Remove the + 0.3 from actions so the pushes average out to zero. Predict whether the open-loop error grows faster or slower, and explain it using the cart's speed.

Pause and think: The open-loop error is 0.00 after one step, even though the model's friction is wrong. Why does the flaw not show up straight away?

The wrong number multiplies the velocity, and the cart starts with velocity 0. So on the first step both worlds compute the same thing. The flaw only appears once the cart is moving. A model's errors show up only in states that exercise them, which is why a model can look perfect on the data it was tested on and still fail in situations it has not seen.

Pause and think: From 10 to 20 steps the horizon doubles, but the error grows about five times (1.04 to 5.18). Why is the growth faster than a straight line, and what does resyncing change?

Each step's velocity error is carried into the next step and added to a new one, and position then adds up all those velocity errors. Errors stack on errors, so the gap accelerates. Resyncing throws away the accumulated gap and restarts from the true state: with a resync every step the worst gap is only 0.09. This is the numerical reason model predictive control re-plans after every real action.

Key takeaways

• A world model is a learned simulator: (state, action) → next state and reward.
• It is trained by supervised learning on recorded transitions, usually in a compressed latent state.
• Agents use it to plan (search imagined futures) or to learn a policy in imagination, saving real trials.
• Compounding error is the main limit; short horizons, re-planning and uncertainty estimates help.
• Dreamer and MuZero are landmark model-based agents; video world models extend the idea to rich visual worlds.
• Use world models when real experience is costly; use model-free methods when simulation is free.

Key terms

• Environment: The world an agent acts in, which returns new states and rewards after actions.
• State: A description of the environment at one moment.
• World model: A learned model that predicts the next state (and often reward) from the current state and an action.
• Latent state: A compressed vector representation of an observation in which dynamics are learned.
• Rollout: A simulated trajectory produced by repeatedly applying a model to chosen actions.
• Model predictive control (MPC): Planning a sequence of actions with a model, executing only the first, then re-planning from the new real state.

← 18.1 JEPA: LeCun's Vision for World Model AI · 18.3 Recursive Self-Improvement: Can AI Improve Itself Indefinitely? →