Ever since I first read the World Models paper, there has been a nagging irresistible thought at the back of my mind: “I want to try that out on Deathchase". And, to be fair, I have over the years spent a few Sunday nights trying to make it work. The RNN or the VAE code wasn’t the problem; I have working code on my old 1080Ti PC tower that trains well enough on the example toy racing world. The problem was somehow hooking it up to a Spectrum. I tried a lot of different approaches, including some major surgery to MAME, and I got tantalizingly close, to the point of hand-disassembling Z80 code to work out where information like lives and scores were being kept but never enough to really make it work in a fashion I could run a full training loop without it keeling over and dying a few iterations in.
(you know how this is going to go)
It popped up in my head again recently, and this time I asked Claude. It laughed, saying: “Why do you not just use zx, as it’s basically built exactly for what you want to do?” because I didn’t know it existed, you smug little SkyNet! And then we set to work.
The World Model paper is actually three different models. Firstly, a variational autoencoder (VAE) takes the pixels from the screen (resized in the paper to 64x64, in my code to 84x84) and compresses them down to a latent vector. Just 64 numbers to describe everything going on in every single frame fed into the model. This is then chained into a LSTM network, which takes the latent vectors and learns to predict what comes next. These two models together form the world model; eventually, you can feed a starting frame into the autoencoder and then ‘dream’ the entire sequence of the game inside the LSTM. And this is the trick that seemed magic in 2018; it’s annoyingly hard to do reinforcement learning or any sort of gradient descent work when you have to keep going back to the game state, perhaps in an awkwardly coded emulator. There’s a lot to go wrong. Here, the controller model trains directly on the dream world and then you transfer the trained controller model back to the real world…and it works! Mostly.
Fig 1 · V–M–C pipeline
mermaid source
flowchart LR
o["obs oₜ<br/>64×64 frame"]
subgraph WM["world model · runs with no environment (the dream)"]
direction LR
V["V · VAE encoder<br/>compress → zₜ ∈ ℝ³²"]
M["M · MDN-RNN<br/>predict P(zₜ₊₁ | zₜ, aₜ, hₜ)"]
V -->|"zₜ"| M
end
o --> V
M -->|"zₜ, hₜ"| C["C · controller<br/>1 linear layer · ~900 params"]
C -->|"aₜ"| E(["real environment"])
E -->|"next frame"| o
classDef ha fill:#2a2214,stroke:#eaa63f,stroke-width:2px,color:#f3e7d4;
classDef ctl fill:#1b2130,stroke:#cfd6e2,stroke-width:2px,color:#eef2f8;
classDef io fill:#141a26,stroke:#6b7688,stroke-width:1.5px,color:#c7d0de;
class V,M ha; class C ctl; class o,E io;
style WM fill:transparent,stroke:#eaa63f,stroke-dasharray:5 5,color:#c69a52;After spending an afternoon with Opus a couple of weeks ago, the first discovery was that Ha’s world model…just didn’t work very well for Deathchase. It did learn a grand plan of surviving…which was to slowly go through the first day, shoot at a few things, and simply sit tight when night falls. Which is a strategy, but one that would have got you slapped if you tried it back in the day round your friend’s house. Also, it does have a tendency to hit the trees a lot. But then I have that problem as well.
Despite that, looking at the dream sequences of the VAE model is spooky; you can see exactly what the model is seeing and, yep, well, it’s certainly Deathchase.
Anyhow, I was not to be stopped by that little failure. Because I kept reading these sorts of papers, and I specifically remembered DreamerV3 from DeepMind. Back in 2023 it pretty much set the new state of the art for autonomously playing Atari games. So Claude, myself, and the DGX Spark worked for a day or two to train a new model based on this new approach.
The DreamerV3 framework is a lot more complicated (see the image, yikes!), but I’d say that the too key points are: a CNN model to better capture what’s going on versus the older autoencoder, and a training regime that keeps the world model anchored to reality, with joint training focused on reconstructing the frame, optimising the reward, and whether the model has left the game in a playable state. And this works a lot better.
Fig 2 · the DreamerV3 world model, one timestep
mermaid source
flowchart LR
prev["hₜ₋₁ · zₜ₋₁ · aₜ₋₁"] --> GRU["sequence model<br/>GRU"]
GRU --> H["hₜ<br/>deterministic memory"]
X["obs xₜ"] --> ENC["encoder<br/>CNN"]
ENC --> Q["representation q<br/>posterior zₜ | hₜ, xₜ"]
H --> Q
H -.-> P["dynamics p<br/>prior ẑₜ | hₜ · no obs"]
Q -. "KL: pull prior → posterior" .-> P
H --> S["state sₜ<br/>(hₜ, zₜ)"]
Q --> S
S --> DEC["decoder → x̂ₜ"]
S --> REW["reward → r̂ₜ"]
S --> CON["continue → ĉₜ"]
classDef dr fill:#123039,stroke:#43c0d0,stroke-width:2px,color:#dff3f6;
classDef prior fill:#123039,stroke:#43c0d0,stroke-dasharray:5 4,color:#dff3f6;
classDef neut fill:#1b2130,stroke:#8b97ac,stroke-width:1.5px,color:#e9edf6;
class GRU,H,Q,S dr; class P prior; class prev,X,ENC,DEC,REW,CON neut;This is Dreamer dreaming of trees and motorbikes, learning how to ride in just 12,000 iterations.
And this is Dreamer playing the actual game, having learnt to ride safely and shoot things with 300k steps played in the imagined worlds1
I think we can say it’s mostly solved at this point…but if you don’t believe me, then here’s five minutes of it playing without losing a life.
Of course, now that I had Deathchase sorted, I started thinking about other games. We’ve got another post coming that stays firmly in the 16K era, but is a touch more complicated…and another classic. We’ll be going to Ashby-de-la-Zouch…
-
Although I will admit that the model is definitely cheesing things by not hitting the full acceleration. ↩︎