MiniPong: building V and M from scratch

Module 3 · 7 min read

Building V: squeezing a frame into 12 numbers

The 2018 World Models paper by Ha & Schmidhuber split the problem into three small networks, and that split is still the clearest way in:

V (vision) compresses each frame. M (memory) remembers and predicts. C (controller) picks the action. I’ll write c for the compressed code and m for the memory.

The test bed: MiniPong

Instead of car racing, the course builds V and M on a tiny game: a 32×32 screen, a teal paddle at the bottom with three actions (left, stay, right), and a glowing yellow ball that bounces off the walls and the paddle. The data is 200 episodes of random play, 24,000 frames in total. Random play is fine because the goal here is to learn how the world works, not how to win.

V turns each frame (32×32×3 = 3,072 numbers) into 12 numbers, and a decoder redraws the frame from those 12. Two convolution layers shrink the image 32 → 16 → 8, and a linear layer maps that to the code. The decoder mirrors it.

Four small decisions that do most of the work

The architecture is ordinary. What surprised me is that almost all the quality comes from four lines in the training loop, and each one exists because a simpler version failed.

1. Make the ball count. The ball is a handful of pixels out of 1,024. Under a plain pixel loss, forgetting the ball entirely costs almost nothing, and that is exactly what the first version did: a perfect paddle and no ball. The fix is to up-weight red pixels, since only the ball is red.

ball weight

1×

2. Keep the code exact. A textbook VAE (see the appendix) gives each frame a small cloud of codes and samples from it. Here, the codes handed to M are the cloud’s centre μ, with no noise. M’s job is to predict the next code, and it can only be as exact as its answer key. If the same frame gave a slightly different code every time, M would chase a moving target forever.

3. Teach the decoder to read slightly wrong codes. M’s predictions will land near real codes, never exactly on them. So during training, the decoder alone sees the code plus a little noise and still has to draw the right frame. Think of it as learning to read messy handwriting.

4. Neighbouring frames get neighbouring codes. The ball moves about a pixel per step, so the code should move a little too. Without a push in that direction, V is free to map frame 5 to one corner of code space and frame 6 to the opposite corner. Both would decode perfectly, but M could never learn the jump. So the loader takes pairs of consecutive frames and penalises the distance between their codes.

Tricks 1, 3 and 4 in the loss:

z    = mu + 0.1 * torch.randn_like(mu)          # 3: noise for the decoder only
xhat = vnet.decode(z)
w    = 1.0 + 40.0 * x[..., :1]                  # 1: red pixel (ball) weighs 41
recon  = (w * F.binary_cross_entropy(xhat, x, reduction="none")).sum() / len(x)
smooth = ((mu - mu_next) ** 2).sum(1).mean()    # 4: nearby frames, nearby codes
loss = recon + 1.0 * smooth + 0.01 * (mu ** 2).sum(1).mean()

Here is what that buys. Top row: real frames. Bottom row: the same frames squeezed to 12 numbers and redrawn from those 12 alone.

If I had to keep one sentence from this module: a model keeps only what the loss pays for.

Module 4 · 7 min read

Building M, and why dreams drift

One frame isn’t enough, so give it a memory

Module 1’s problem comes back here. From one code, M can’t know which way the ball is heading. So M carries a memory, a vector of 128 numbers that persists from step to step, updated by a GRU (a recurrent layer whose gates learn what to keep and what to overwrite):

mₜ₊₁ = GRU(mₜ, cₜ, aₜ)          update the memory
ĉₜ₊₁ = cₜ + head(mₜ₊₁)      predict the next code

Nobody gives the network velocity. After it has seen two ball positions, the memory can hold their difference, and that difference is the velocity. The memory derives it because prediction is impossible without it.

You can watch this happen. Here is M’s prediction error at each step of an episode, averaged over many episodes:

Three more small decisions

• Rescale every code dimension. The busy paddle dominates some of the 12 numbers, while the ball’s numbers barely move. Without rescaling, the loss obsesses over the paddle. Subtracting the mean and dividing by the standard deviation per dimension makes them count equally.
• Predict the change, not the code. When a network trained with squared error is unsure, it drifts toward the average answer, and the average MiniPong frame has a smeared, invisible ball. If it predicts a change instead, its unsure answer is “nothing moves”, which at least keeps the ball where it was.
• Train on slightly noisy inputs. In the dream, M eats its own imperfect guesses. Adding a little noise to its inputs during training teaches it to correct small errors instead of amplifying them.

z_in = z[:, :-1] + 0.05 * torch.randn_like(z[:, :-1])   # noisy "now"
pred, _ = mnet(z_in, a[:, :-1])                         # predicted change
loss = F.mse_loss(z_in + pred, z[:, 1:])                # now + change vs real next

Because M works on 12 numbers instead of 3,072 pixels, it trains in seconds. That speed is what compression bought us.

First, the easy test. At every step, show M the real current frame and the key, and ask it to paint the next frame before the game shows it:

That looks great. The catch is that the model was never left on its own.

The dream: switching the game off

Now the real test. Warm the model up on three real frames, switch the game engine off, and loop:

1. a key is pressed;
2. M updates its memory and predicts the next code;
3. V’s decoder paints that code, and that’s what you see;
4. the predicted code is fed back in as if it were real.

c = first_real_code
for a in actions:
    change, m = M(c, a, m)
    c = c + change              # the model's own guess becomes its next input
    frames.append(V.decode(c))

Step 4 is where things go wrong. Every guess is slightly off, and the next step starts from that slightly-off guess, so the errors pile up. It’s a photocopy of a photocopy. The paddle holds up, because the key press tells M exactly where it goes. The ball has to be inferred, so it fades.

Real game

Dream

error per dream step

0.08

← ■ → Restart

And here is the real thing, from the trained network on my run. Same keys in both rows; the bottom row never sees the game after the warm-up:

In the course’s own run, the ball stays sharp for the first few steps, then fades and re-forms. Mine was worse: the ball was gone by step 3, and even the paddle didn’t follow the keys (Module 6 has more). This isn’t a bug in the notebook. Keeping a long dream from drifting is the open problem in this field, and most later work (Dreamer, IRIS, DIAMOND) is a better answer to it.

Why the paper’s M outputs several futures

MiniPong is deterministic, so one guess per step is enough. Car racing isn’t: the track can be driven both ways, so from one frame the future really has two branches. The paper’s M therefore outputs a mixture of Gaussians: several candidate next codes, each with a probability, and the dream samples one.

• Why not average the candidates? The average of “turn left” and “turn right” is “drive into the wall”, a future that never happens.
• Why not always take the most likely one? Then the dream never shows rarer but real possibilities.

The course shows this nicely: with one frame of history, the dreamed cars split about 54/46 between clockwise and anticlockwise. With six frames of history, 90% drive the way the real car was going, because the memory now knows the direction.