Mahir Shahriar Tamim

AI Researcher · Explainable & Responsible AI · Vision, Language & Agentic Systems

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:

  1. make a population of slightly different copies of C;
  2. let each one play a full episode and record its score;
  3. keep the best few, and centre the next population around them;
  4. repeat.

A pretend controller with just 2 parameters, so we can draw it. Brighter means a higher episode score. Grey dots are this generation’s copies, filled dots are the best 6, and the ring is where the next generation will be centred. Each Reset starts from a random controller: some runs find the tall hill, others settle on the small one and can’t leave.

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:

  1. Drive the car randomly for 10,000 episodes and record every frame and action.
  2. Train V, a VAE that squeezes each 64×64 frame into 32 numbers.
  3. Train M, an LSTM with 256 memory units that predicts the next code as a mixture of 5 Gaussians.
  4. 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 seeScore
V’s code only632 ± 251
V’s code, through an extra hidden layer788 ± 141
V’s code + M’s memory906 ± 21
for comparison: DQN / A3C343 ± 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 dreamScore in the real game
0.102086 ± 140193 ± 58
0.502060 ± 277196 ± 50
1.001145 ± 690868 ± 511
1.15918 ± 5461092 ± 556
1.30732 ± 269753 ± 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):

Runstep 1step 3step 10V’s redrawThe dream
baseline0.0940.0420.022goodball gone by step 3, faint ghosts later
no_exact0.1120.0900.057goodball partly comes back around steps 9–11
no_fuzz0.1050.0450.023goodball gone, paddle slightly stronger
no_smooth0.1050.0460.023good, looks identicalball gone, paddle fades fastest
baseline
0.022 no_exact
0.057 no_fuzz
0.023 no_smooth
0.023
M’s prediction error at step 10. Removing exact codes costs about 2.5×; the other two barely register.

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.

Top row: noise ε drawn once from N(0, 1). Bottom row: the same noise pushed through c = μ + σε. Moving the sliders slides and stretches the samples smoothly, and that smoothness is exactly what makes them differentiable. The grey curve is the prior N(0, 1).

How MiniPong’s V differs

Textbook VAEMiniPong’s V
Code used downstreama sample, μ + σεthe centre μ only
What keeps codes tidyKL to N(0, I)the smoothness penalty and a small pull toward 0
Where the noise goesinto the code itselfonly 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).