What a world model is, and the tools to build one

Module 1 · 6 min read

What a world model actually is

A system that predicts the result of an action

Push a cup near its edge and you already know it will tip. Push it in the middle and you know it will slide. You didn’t compute any physics; you ran a tiny simulation in your head. The psychologist Kenneth Craik described this in 1943: a mind that carries a “small-scale model” of reality can try out actions before committing to them. So the idea isn’t new. What’s new is that since 2018 we can actually learn such models with neural networks.

The agent and the environment

Almost every setup in this field looks like the same loop. An agent sees something and picks an action. The environment takes the action, changes, and reports back what the agent can now see (and sometimes a reward).

So where does the world model go: on the agent’s side or the environment’s? This question took me a while. The answer is the environment’s. A world model is a learned replacement for the environment box. If it’s good, the agent can’t tell the difference, and that’s exactly the point.

Three jobs, not one

Most of my early confusion came from people using “world model” for three different jobs. Sorting a paper into one of these clears things up immediately:

The simulator is the core. A renderer is a nice-looking layer on top of one, and a planner needs something to plan in. That also settles a common mix-up: world models and RL are not the same thing. RL is one way to build the planner. DINO-WM plans with no RL at all.

State versus observation

This is the idea that clicked most for me. The state is everything that decides what happens next: every position, every velocity, every hidden card. The observation is what your sensors actually give you, which is usually a thin slice of the state.

Take a single photo of a thrown ball. You can see where the ball is. You can’t see where it’s going. Position is in the observation; velocity is in the state but invisible. Chess is fully observable (the board is the state). Poker isn’t (your opponent’s cards are part of the state). The real world is always poker.

1 frame 3 frames Reveal state Pause

There are classic fixes for this. You can stack the last few frames, as the demo does, so velocity appears as a difference between frames. You can filter noisy sensor readings into a state estimate; the Kalman filter did this on Apollo. Or you can give the model a memory, which is what world models do. We’ll build exactly that in Module 4.

Four dials for reading any paper

Once a paper is placed on the map, four more questions describe it almost completely:

1. Representation: does it think in raw pixels, discrete tokens, or continuous latent vectors?
2. Training signal: does it learn by redrawing pixels, or by predicting latents (JEPA never touches pixels)?
3. Action source: are the actions logged in the data, or inferred from video alone (Genie)?
4. Behaviour: are actions chosen by a trained policy, or by searching at decision time (MPC)?

Module 2 · 6 min read

The toolkit: compression, return, value

Before building anything, four tools come up in every world-model paper. Most papers assume you already know them.

Tool 1: compress, then predict

A 64×64 colour frame is 12,288 numbers. A world model has to predict the next frame thousands of times to imagine one episode. Doing that in pixels wastes almost all the effort on textures that no decision depends on.

So we first train an autoencoder. An encoder squeezes the image down to a short vector (the latent), and a decoder tries to redraw the image from that vector alone. Comparing the redraw with the original is the teacher, so no labels are needed. If the redraw is good, the short vector must have kept what matters.

There’s a bonus. In a good latent space, similar situations end up close together, so “predict the next frame” turns into “take a small step in a small space”. That sentence is the recipe for every system in this series.

Tool 2: judge an action by everything it leads to

Studying tonight costs an evening; the payoff arrives months later. A good action often looks bad right now, so we score it by the total reward that follows, with far-away rewards shrunk a little per step:

Gₜ = rₜ + γ·rₜ₊₁ + γ²·rₜ₊₂ + …      with 0 ≤ γ < 1

The discount γ has two jobs. It keeps the sum finite in endless games, and it encodes how much the agent cares about the future.

γ

0.990

Computing it is easiest backwards, because each step’s return is its reward plus the discounted return of the next step:

def discounted_return(rewards, gamma=0.99):
    G = 0.0
    for r in reversed(rewards):   # G_t = r_t + gamma * G_{t+1}
        G = r + gamma * G
    return G

discounted_return([0, 0, 0, 1])   # 0.970299: the late +1 still counts, slightly shrunk

The hard part isn’t the formula. It’s credit assignment: you get one “+1, you won” at move 40. Which of the 40 moves earned it?

Tool 3: a price tag on a situation

A value function fills that silence. V(s) estimates the total future reward from situation s, the way a chess player says “White is winning” long before any piece is taken. Q(s, a) does the same for a specific action in that situation.

World models need this for a specific reason: imagined futures get unreliable after a handful of steps. So an agent dreams a short distance in detail, then lets V summarise everything past the horizon in a single number.

Tool 4: two ways to learn behaviour