World Models · Chapter 1

Why would an agent dream?

Before RSSMs, KL balancing and two-hot critics, one question: what does a model of the world give you that trial and error never can?

~35 minreading + playing

6 problems+ 1 challenge

AssumesMDPs, Q-learning

DestinationDreamerV3

Cold openWhere will it be?

A ball is flying around a box. You get to watch for a moment. Then the lights go out, and the ball keeps moving in the dark. Click where you think it is when the clock runs out. Four rounds, and each one asks you to look a little further ahead.

Figure 1.1.The solid trail is what you saw. When the lights go out, the ball keeps moving for 0.3 s in round 1, 0.6 s in round 2 and 0.9 s in round 3. The violet trail, shown after you click, is what really happened in the dark.

Act I · 1The price of trying

A child touches a hot stove once. An RL agent might need to touch it ten thousand times. Both are "learning from experience", so where does the gap come from? It isn't just that the network is dumber.

One way to see it

The child doesn't only learn "stove → pain". They update a model: hot things hurt, this kind of glow means hot, the kettle is probably hot as well. One experience changes predictions about thousands of situations they never tried. A model-free agent turns each experience into a small nudge to the value of the exact state it was in. A model generalises experience into structure, and structure can be reused.

This matters most when experience is expensive. A simulator can hand out a million frames in minutes. A robot arm, a car or a chemistry lab can't. So how much experience does classical deep RL actually use?

Your guess: how many hours of game time did DQN (2015) train on per Atari game?

Figure 1.2.Hours of real game time at 60 frames/s, on a log scale. The human reference is the professional tester in the DQN paper, who practised each game for about two hours. Atari100k is a benchmark that caps agents at 100k decisions (400k frames). It was introduced for model-based agents, and DreamerV3 reports results on it.

Numerical●●●

DQN trained for 50 million agent steps per game. Suppose you ran that budget on a real robot arm instead. It acts at 10 Hz, episodes last 20 s, and after each one a human needs 10 s to reset the scene. Assume it runs day and night with no breaks. How many days does training take?

Solution

50M steps at 10 steps/s is $5\times10^6$ s of acting. Each 20 s episode costs 30 s of wall time, so multiply by $30/20 = 1.5$: $7.5\times 10^6$ s. Divide by 86,400 s/day to get ≈ 86.8 days, almost three months of continuous use for one skill, and that assumes nothing breaks. That's why robotics was one of the first places model-based RL had to work.

Act I · 2Two ways to use one experience

An agent takes action $a$ in state $s$, sees reward $r$ and lands in $s'$. That one tuple $(s,a,r,s')$ is all experience ever is. What can you do with it?

Model-free: push it straight into the value function and move on. Q-learning:

$$Q(s,a) \;\leftarrow\; Q(s,a) + \alpha\Big[\,\textcolor{#15a34a}{r} + \gamma \max_{a'} Q(\textcolor{#1c1d21}{s'},a') - Q(s,a)\Big]$$

Model-based: also use it to learn a model of the world, $\hat{T}(s,a) \to (\hat{r}, \hat{s}')$. From then on, the model can generate as many imagined tuples as you have compute for, and you apply exactly the same update to them:

$$Q(s,a) \;\leftarrow\; Q(s,a) + \alpha\Big[\,\textcolor{#7046e0}{\hat r} + \gamma \max_{a'} Q(\textcolor{#7046e0}{\hat s'},a') - Q(s,a)\Big], \qquad (\textcolor{#7046e0}{\hat r},\textcolor{#7046e0}{\hat s'}) = \hat T(s,a)$$

Figure 1.3.The real loop is slow and expensive. Once a model exists, a imagination loop runs several times per real step, and each real step also improves the model. This is Sutton's Dyna architecture (1990), the ancestor of every algorithm in this series.

"Dyna" is from Sutton's 1990 paper Integrated architectures for learning, planning, and reacting. Dreamer's "learning behaviours in imagination" is the same loop with neural networks in place of the table. That's the whole trick. In Dyna-Q the model is just a lookup table: "last time I did $a$ in $s$, I got $r$ and went to $s'$." After every real step, the agent replays $n$ random remembered transitions through the Q-update. Here $n$ is the number of planning steps: how much it dreams per real step.

Act I · 3The race

Three agents in the same maze, Sutton & Barto's classic 6×9 grid. The reward is +1 on reaching the goal and 0 everywhere else, with γ = 0.95, ε = 0.1 and α = 0.1. All three take exactly the same number of real steps per tick. The only difference is how much they dream: $n = 0$ (plain Q-learning), $n = 5$ and $n = 50$.

The shortest path is 14 steps. With ε-exploration, a well-trained agent averages about 16–18.

Predict first: after how many episodes will each agent consistently solve the maze in ≤ 20 steps?

n = 0 n = 5 n = 50

speed

n = 0 (Q-learning) n = 5 n = 50 faint: average of 30 runs

Figure 1.4.Cell shading is each state's value $\max_a Q(s,a)$. Violet flashes are imagined updates: cells the agent is dreaming about this tick. The orange dot is the agent's real position. The plot shows steps per episode on a log scale.

MCQ●●●

Watch the very first episode of the $n=0$ agent. When it finally stumbles into the goal, how many state–action pairs in its Q-table have a non-zero value?

Exactly one. Every reward before the goal is 0 and every Q starts at 0, so every target $r + \gamma\max Q(s',\cdot)$ is 0 until the goal is reached. The final transition then gets $\alpha\cdot 1 = 0.1$, and that's all. Value creeps backwards one step per real visit. That's why the $n=0$ agent's shading spreads out from the goal so slowly.

The $n=50$ agent re-runs that single rewarding transition, and its neighbours, dozens of times in its head straight away, so value floods back through the maze after one success. A model lets you propagate one experience much further than the place it happened.

Numerical●●●

Is the dream worth its cost? Imagination isn't free: each planning step costs compute. Suppose that on some task, plain Q-learning needs $N_0 = 200{,}000$ real steps. Dyna-Q with $n=50$ needs only $N_{50} = 15{,}000$, but every planning update costs $c_p = 0.2$ ms. Let $c_r$ be the wall-clock cost of one real step.

(a) Real robot, $c_r = 0.5$ s. Hours for each method?
(b) Fast emulator, $c_r = 0.1$ ms. Seconds for each method?
(c) Above what $c_r$ (in milliseconds) does Dyna win on wall-clock time?

Check your answer to (c):

Solution

Wall time is $N_0 c_r$ for Q-learning and $N_{50}(c_r + 50\,c_p)$ for Dyna. Note $50 c_p = 10$ ms.

(a) Q: $200\text{k}\times 0.5 = 100{,}000$ s ≈ 27.8 h. Dyna: $15\text{k}\times 0.51 = 7{,}650$ s ≈ 2.1 h. The model wins by 13×.

(b) Q: $200\text{k}\times 10^{-4} = $ 20 s. Dyna: $15\text{k}\times 0.0101 = $ 151.5 s. The model loses by 7.5×.

The lesson: world models pay off when reality is expensive relative to compute. Robots, the physical world, slow simulators and human feedback all qualify. That's also why the world-model literature cares so much about sample efficiency benchmarks like Atari100k rather than wall-clock time.

Act I · 4What a model buys, and what it costs

Sample efficiency is the headline, but it isn't the only thing a model gives you. Here are four, and one bill:

Gift 1

Reuse every experience

One real transition gets replayed and propagated many times, as in the race above.

Gift 2

Think before acting

Roll candidate action sequences forward and pick the best before committing. This is planning (PlaNet, MuZero).

Gift 3

Counterfactuals

"What if I had braked?" The model can answer questions about actions you never took.

Gift 4

Fail safely

Crash the car a thousand times in the dream. Crash it zero times on the road.

The bill

The model is wrong

The agent optimises against the model, so it will find, and exploit, every mistake the model makes. Remember round 4.

MCQ●●●

You're handed a perfect model of an environment's dynamics and rewards. Which of these do you still not get for free?

C. A perfect model gives you a perfect simulator, not an answer. You still have to search or optimise inside it, and that can be enormously expensive (think chess: the rules are a perfect model). Every Dreamer paper is really two problems: learn the model, and then learn to act inside it efficiently. Chapter by chapter we'll attack both.

Numerical●●●

Dyna-Q's model is a table keyed by state. Atari frames are 210×160 pixels, and each pixel takes one of 128 palette colours. Ignoring game logic, how many distinct frames are possible? Give the answer as a power of 10, $10^{x}$, and enter $x$.

Solution

There are $210\times160 = 33{,}600$ pixels with 128 choices each, giving $128^{33600}$ frames. $\log_{10}128 \approx 2.1072$, so this is about $10^{70{,}802}$. There are roughly $10^{80}$ atoms in the observable universe.

You'll essentially never see the same frame twice, so a table that memorises "what happened last time in exactly this state" is useless. The model has to generalise: map a frame to a compact latent state where similar situations are close together, and learn dynamics there. That one sentence is the reason latent world models exist. We'll build one soon.

Socratic●●●

Models are wrong in two different ways. Think about each one before opening the answer.

Optimistic error: the model believes there's a gap in a wall where there's actually a wall.
Pessimistic error: the model believes there's a wall where there's actually a gap.

A Dyna agent plans with each wrong model. Which error corrects itself, and which can last forever? Why?

Answer

Optimism corrects itself. The agent plans a route through the imaginary gap, walks there for real, hits the wall, and the model updates. Being wrong in a hopeful direction sends you to exactly the place where you'll find out.

Pessimism can last forever. If the model says "wall", the planned route never goes there, so the agent never collects the data that would prove the model wrong. Apart from rare random ε-steps, the error is self-protecting.

This asymmetry is about how a learner gathers its own data, and it'll come back in every Dreamer. It's also exactly the trap in this chapter's challenge.

Chapter 1 challenge

The Shifting Maze →

The world changes under a Dyna agent's feet. First predict what happens, then design a fix and test it in a live simulator. Work it on your iPad and send me your answers as photos or text.