Practice · Fundamentals
How a World Model
Learns, Dreams, and Plans.
DE Auf Deutsch lesen
A world model is a program that predicts how an environment changes after an action. I built one in a single Python file, without PyTorch and without a graphics card. This piece explains what four parts something like this is made of, and what the effort buys you.
The problem a world model solves
Anyone who has written game AI before knows the setup. You need a function
next_state(state, move), use it to play out possible sequences of moves, score the
results, and execute the best first move. For chess or tic-tac-toe you write that function
yourself, because the rules are known.
As soon as the input consists of sensor data, nobody can do that anymore. A robot sees a camera
image, not a coordinate system. For the question "What will the image look like if I turn
right?" there's no rule you can write down. So you learn it from recorded data instead. This
learned next_state is a world model.
The approach comes out of research on agents that learn from images. David Ha and Jürgen Schmidhuber trained a controller in 2018 partly inside a learned simulation instead of the environment itself (World Models). The Dreamer work by Danijar Hafner and colleagues took the same path further and learned behavior from imagined rollouts (Dream to Control). My experiment is a few sizes smaller: a grid world, two small networks, a planner that rolls dice.
The world consists of 49 pixels
The environment is a room of 7×7 tiles. # is a wall, Z the goal,
o the agent:
#######
#..#..#
#..#..#
#..o..#
#.##..#
#....Z#
#######
Four actions, a reward of 1 on the goal tile, nothing otherwise. What matters is what the agent does not get: no coordinates, no map, no list of walls. Its observation consists of 49 grayscale values, one per tile. A wall is 1.0, the goal 0.6, itself 0.3, empty tiles 0. Everything it knows about its world, it has to pull out of these 49 numbers.
That's the stripped-down version of the robot problem. Small enough that a full run takes two minutes on a laptop, and large enough that the typical failure modes still show up.
Step 1: wander aimlessly and log everything
A model needs training data, and here the only way to get it is by trying things out. The agent walks randomly through the room for 3,000 steps. Every step gets logged as a row: image before, action, image after, reward, episode ended yes or no. That's a log file, nothing more.
In my run that came out to 94 episodes, in which chance landed on the goal 38 times. 1.3 percent of the 3,000 rows contain a reward at all. Everything downstream has to work with that yield. Anyone who has worked with imbalanced datasets can guess the problem: a model that stubbornly predicts "reward 0" is right 98.7 percent of the time and is worthless for planning.
Step 2: compressing 49 numbers into 8
Planning directly on pixels is expensive. So first comes an autoencoder. That's a network shaped like an hourglass: it squeezes the 49 values through a bottleneck of 8 numbers and then has to reconstruct the original image from just those 8 numbers. It trains against its own input; it needs no labels.
The trick is in the bottleneck. Not much fits through 8 numbers, so the network has to choose what matters. Walls and the goal sit in the same spot in every image and therefore carry no information. The only thing that's new is where the agent stands. If the compression works, those 8 numbers essentially encode its position, derived without anyone telling it that something like a position exists. The technical term for this is latent state.
x and y. The network
distributes the image information across that space on its own.
This is where my first attempt went wrong, and the mistake is the most instructive part of the project. The reconstruction error dropped to 0.0015 and stayed there. That looks like a cleanly converged network. In fact it was exactly the error you get when you paint the walls and goal perfectly and leave the agent out entirely: one wrong pixel out of 49 doesn't show up in the average. The network had memorized the map and ignored the one thing that actually changes.
Step 3: the dream
Now comes the actual world model. A second network takes the 8 numbers and an action and predicts: the 8 numbers afterward, the reward, and whether the episode ends. It never sees a single pixel again. This split is the reason for the work in step 2. A planning step now costs eight numbers instead of forty-nine.
Because the model can turn its own output back into an input, it runs in a loop: one observation at the start, then nothing but actions from there on. The program calls this a dream. The decoder from step 2 turns the 8 numbers back into an image so you can watch. Left is reality, right is the imagined version, from a run made today:
Welt Traum
####### #######
#..#..# #..#..#
#..#..# #..#..#
#.....# #.....#
#.##o.# #.##o.#
#....Z# #....Z#
####### #######
1. rechts Belohnung echt 0 / getraeumt 0.00
#.##.o# #.##.o#
2. rechts Belohnung echt 0 / getraeumt 0.00
#.##.o# #.##..#
3. runter Belohnung echt 1 / getraeumt 0.00
#....o# #..o.Z#
Step 1 is correct. In step 2 the dream loses the agent, in step 3 it puts it in the wrong spot and misses the reward entirely. That's what decay looks like when a model keeps feeding its own errors back into itself.
Six dreams at once
All runs start from the same image and use different action sequences. The gold dot shows the model's prediction, the cyan outline the position after the same action in the real grid world. Tiles marked red stand for predicted rewards.
The colored dots sit on the shared starting tile.
How far the dream carries
This decay can be measured. The model gets one image and after that nothing but actions. We count whether the dreamed agent ends up on the correct tile. Alongside it runs a baseline model that simply claims nothing changes at all. That's the bar a world model has to clear at minimum.
| Dreamed steps | Position correct | Static model |
|---|---|---|
| 1 | 95.0 % | 35.0 % |
| 3 | 92.7 % | 27.1 % |
| 5 | 83.6 % | 16.9 % |
| 10 | 71.3 % | 16.0 % |
The dream decays slowly and predictably. From that follows the operating rule for the next step: the agent only plans a few moves ahead and then reads reality again.
Step 4: planning means rolling dice and filtering
The planner is the simplest part of the program and learns nothing at all. Before every real move it rolls 100 random action sequences of six moves each. It plays out each sequence inside the dream and sums the predicted rewards. From the best sequence it only executes the first move, looks at the real world, and rolls again. The procedure is called random shooting with model predictive control, and it does roughly what a navigation device does at every junction.
For comparison, the same planner also runs with the real game rules instead of the learned model — called "Oracle" in the program — and one that just moves randomly. 20 episodes, cutoff after 40 steps, planning horizon 6:
| Planner | Goal reached | Average steps |
|---|---|---|
| World model | 20/20 | 4.5 |
| Oracle (real rules) | 19/20 | 11.3 |
| Random | 9/20 | 33.1 |
The dreamer beats reality, because it's wrong
A learned model can't be more accurate than the original it imitates. And yet there it is: 20/20 against 19/20, with the dreamer needing a third of the steps. I looked into where that lead comes from.
It comes from starting tiles more than six moves from the goal. There the oracle doesn't find a single one of its 100 rolled sequences that reaches the goal. All 100 score 0, there's nothing to compare, and the planner essentially picks at random. The world model, by contrast, predicts small rewards between 0.00003 and 0.77 in places where there are none. These errors accumulate along paths that head toward the goal. Out of a reward that exists on exactly one tile, this creates a gradient across the entire space.
| Planner | Move shortens the path | Measurement scope |
|---|---|---|
| World model | 82.2 % | 90 moves |
| Oracle (real rules) | 46.7 % | 90 moves |
So the lead was borrowed. At horizon 12, once the oracle finds the reward on its own, it catches up at 20 out of 20. The blur cuts both ways: at horizon 3 the world model predicts rewards in places that lead nowhere and drops to 8 out of 20 — worse than pure random rolling at 9 out of 20.
What sticks for everyday development
- A good average can hide the one case that matters. A reconstruction error of 0.0015 looked like convergence while the network was ignoring the agent entirely. A second metric measuring exactly the thing that mattered functionally cost twelve lines.
- A cheap baseline belongs in the project from day one. The static "nothing changes" model is right 35 percent of the time after one step. Without that number I wouldn't know whether 95 percent is a lot.
- A good end result doesn't prove the parts are correct. The planner succeeded because the model was wrong. Anyone who only measures the success rate mistakes a broken model for a good one.
- Twenty episodes and a single random seed aren't statistics. 20/20 against 19/20 is a difference of one episode. I verified the explanation behind it by measurement; I wouldn't defend the number itself without multiple seeds.
Run it yourself
The program needs Python 3.6 or later and not a single pip package. A full run takes about two minutes, most of it spent training the autoencoder.
python weltmodell.py zeigen # view the environment
python weltmodell.py sammeln # collect 3000 transitions
python weltmodell.py trainieren # autoencoder and dynamics net
python weltmodell.py traeumen # dream and reality side by side
python weltmodell.py planen 6 # world model vs. oracle vs. random
python weltmodell.py warum # where the lead comes from
The numbers in this post come from a run made today on a Windows laptop with Python 3.8.10. They match the logged values from the original post, because all random seeds are fixed. They say nothing about different seeds, maps, or machines.
Agents that do more than answer?
I build LLM and agent systems for production use, and in training sessions I use small, traceable examples like this one to show what's actually happening inside them. 30-minute intro call, free.