Part 3 of Learning World Models
Learning inside a dream, and what breaks without the tricks
Module 5 · 10 min read
C: learning to act inside a dream
A deliberately tiny controller
The MiniPong code stops at V and M; it never trains a controller. This module follows the 2018 paper, where C reads what V sees and what M remembers and outputs an action. There, C is a single linear layer:
action = W · [cₜ, mₜ] + b (32 code + 256 memory) inputs × 3 outputs + 3 biases = 867 parameters
That’s tiny next to V and M, and on purpose. V and M already did the hard work of understanding the scene, so C only has to answer one question: given this, which way do I steer?
No gradients: evolution instead
C gets feedback as one number per episode: the total reward. In car racing that’s a reward for every new piece of track covered minus a little per frame; in the paper’s Doom task it’s simply how long you survive. With so few parameters, you don’t need backpropagation at all. The paper used CMA-ES, an evolution strategy:
- make a population of slightly different copies of C;
- let each one play a full episode and record its score;
- keep the best few, and centre the next population around them;
- repeat.
The same loop in a few lines of Python:
import numpy as np
def evolve(score_fn, n_params, pop=64, elite=8, sigma=0.1, gens=100):
mean = np.zeros(n_params)
for _ in range(gens):
cands = mean + sigma * np.random.randn(pop, n_params) # nudged copies
scores = np.array([score_fn(w) for w in cands]) # one episode each
best = cands[np.argsort(scores)[-elite:]] # keep the top few
mean = best.mean(axis=0) # breed around them
return mean
Here score_fn(w) plays one episode with controller weights w, either in the real game or inside M’s dream. Inside a dream, M also has to predict the score, for example a “did I just die?” flag.
The paper’s full recipe
Here is how the 2018 system was trained, start to finish, for car racing:
- Drive the car randomly for 10,000 episodes and record every frame and action.
- Train V, a VAE that squeezes each 64×64 frame into 32 numbers.
- Train M, an LSTM with 256 memory units that predicts the next code as a mixture of 5 Gaussians.
- Freeze V and M, and evolve C with CMA-ES: 64 candidate controllers per generation, each scored as its average reward over 16 episodes with different random seeds.
Averaging over 16 episodes matters. One lucky track can make a bad controller look good; sixteen rarely do.
Car racing: memory is what made it work
The paper ran the same controller with and without M, scored over 100 random tracks (solved means averaging above 900):
| What C gets to see | Score |
|---|---|
| V’s code only | 632 ± 251 |
| V’s code, through an extra hidden layer | 788 ± 141 |
| V’s code + M’s memory | 906 ± 21 |
| for comparison: DQN / A3C | 343 ± 18 / 591–652 |
With only the current frame, the car wobbles and sometimes misses sharp corners: the huge ±251 says some tracks went fine and some went badly. Adding M’s memory both raised the score and made it consistent, because the memory carries what a single frame can’t show, such as where the road is heading. It’s the same lesson as MiniPong’s error-by-step chart in Module 4: state is more than one observation. Note that in this experiment C was still trained in the real game; V and M were its eyes and memory, not its training ground.
Doom: training entirely inside the dream
The famous result is the second task, VizDoom’s Take Cover: monsters throw fireballs and the agent survives by dodging, for up to 2,100 steps. “Solved” means surviving 750 steps on average. Here M (an LSTM with 512 units, reading 64-number codes) also predicts a done flag, the probability that the agent just died, so the dream can end an episode on its own. C, now 1,088 parameters, was trained only inside M’s dream and then dropped into the real game.
Cheating the dream, and the temperature fix
A controller that trains in a dream will exploit every flaw in it. In some dreams, the agent found moves that made the monsters never fire at all; when a fireball started to form, it moved in a way that made the fireball fizzle out. That’s a perfect score in a world that doesn’t exist. It can even steer M’s memory directly, since C reads that memory as an input.
The fix is a temperature τ on M’s sampling: higher τ makes the dream more random and harder to game. The paper’s table is the clearest picture I know of the gap between a model and reality:
| Temperature τ | Score in the dream | Score in the real game |
|---|---|---|
| 0.10 | 2086 ± 140 | 193 ± 58 |
| 0.50 | 2060 ± 277 | 196 ± 50 |
| 1.00 | 1145 ± 690 | 868 ± 511 |
| 1.15 | 918 ± 546 | 1092 ± 556 |
| 1.30 | 732 ± 269 | 753 ± 139 |
| random policy | — | 210 ± 108 |
Read the two score columns against each other. At τ = 0.1 the dream is nearly deterministic and the monsters never fire, so the agent scores almost the maximum in the dream and worse than random in reality. As the dream gets harder, the dream score falls and the real score rises, peaking at τ = 1.15. A world that’s slightly harder than reality trains skills that transfer; a world that’s easier trains tricks that don’t.
What the paper admits doesn’t work yet
- V doesn’t know what matters. It’s trained only to redraw frames, so it faithfully kept Doom’s brick-wall texture but lost road tiles that mattered for driving. That’s the same problem as MiniPong’s ball, solved there with a hand-tuned weight.
- The tasks are simple. Random play is enough to see everything in car racing and Take Cover. Harder worlds need the agent to explore, retrain the model on what it finds, and repeat.
- Memory is limited. One LSTM can’t hold a large world, and it forgets old experience when trained on new.
Each of these is a thread later work picks up: Dreamer trains the controller with gradients through the dream, IRIS and DIAMOND make the dream sharper, and JEPA stops redrawing pixels altogether.
Module 6 · 5 min read
My experiment: what breaks without the tricks
Module 3 listed four small decisions in V. The notebook claims each one matters. I wanted to see it, so I removed them one at a time and retrained everything.
Setup
The course’s MiniPong script, CPU only, fixed seed. Four runs, each identical except for one changed line:
# no_exact: M's training codes get noise
- zs.append(mu.cpu())
+ zs.append((mu + 0.3 * torch.randn_like(mu)).cpu())
# no_fuzz: the decoder only ever sees exact codes
- z = mu + 0.1 * torch.randn_like(mu)
+ z = mu
# no_smooth: no penalty for consecutive codes being far apart
- smooth = ((mu - mun) ** 2).sum(1).mean()
+ smooth = torch.tensor(0.0)
Results
The number is M’s error when predicting the next code, measured at different points in an episode (lower is better):
| Run | step 1 | step 3 | step 10 | V’s redraw | The dream |
|---|---|---|---|---|---|
| baseline | 0.094 | 0.042 | 0.022 | good | ball gone by step 3, faint ghosts later |
| no_exact | 0.112 | 0.090 | 0.057 | good | ball partly comes back around steps 9–11 |
| no_fuzz | 0.105 | 0.045 | 0.023 | good | ball gone, paddle slightly stronger |
| no_smooth | 0.105 | 0.046 | 0.023 | good, looks identical | ball gone, paddle fades fastest |
The baseline’s dream strip is the one in Module 4.
What I learned
V looked fine every time. The redrawn frames were indistinguishable across all four runs. These tricks aren’t there to help V draw; they’re there to help M predict. A code can be perfectly decodable and still hard to predict.
Exact codes mattered most, at about 2.5× the error by step 10. Part of that is by construction, since that run’s answer key is noisy too.
The other two barely mattered on my run, around 5%, which is much less than the notebook suggests.
The dream was bad even in the baseline. The ball disappeared within about three steps and the paddle didn’t really follow the keys, well short of the notebook’s “the paddle obeys every key”. My guess is that the author’s numbers came from different hardware, and this dream is fragile enough that small differences in training show up. Compounding error dominated every other effect.
The number and the picture disagreed. no_exact had the worst error but the best-looking dream, with the ball reappearing near the right place. If I’d only looked at one of them, I’d have drawn the wrong conclusion.
One run per setting is not evidence. Differences like 0.022 against 0.023 could easily be noise; a fair version needs several seeds per setting, reported as mean ± spread.
Appendix · 3 min read
A quick look at the VAE
MiniPong’s V borrows from the variational autoencoder. This is just enough to see the connection; VAEs will get a proper treatment in a later series on diffusion models.
Points versus clouds
A plain autoencoder maps each frame to one point in code space. A VAE maps it to a small cloud: a centre μ and a spread σ. Its loss has two parts:
loss = reconstruction error ( can you redraw the frame from a code in the cloud? )
+ KL to N(0, 1) ( keep every cloud close to a standard bell curve )
The second term keeps the code space free of holes, so any point near a real code still decodes to something sensible.
The reparameterisation trick
Training needs codes sampled from each cloud, and “draw a random sample” has no gradient, so backpropagation would stop there. The fix is to move the randomness outside the network:
c = μ + σ · ε, ε ~ N(0, 1)
Now c is an ordinary function of μ and σ, and ε is just another input. Move the sliders below: the same noise slides and stretches smoothly, and that smoothness is what lets gradients flow.
How MiniPong’s V differs
| Textbook VAE | MiniPong’s V | |
|---|---|---|
| Code used downstream | a sample, μ + σε | the centre μ only |
| What keeps codes tidy | KL to N(0, I) | the smoothness penalty and a small pull toward 0 |
| Where the noise goes | into the code itself | only into the decoder’s training input (fixed size 0.1) |
MiniPong drops the sampling on purpose: M needs the same frame to give the same code every time (Module 3, decision 2).