CarRacing World Model
A world model in PyTorch. A conv-VAE compresses CarRacing frames into a latent space where an MDN-LSTM learns the dynamics, then dreams coherent driving that never happened.
Every frame after the first is a hallucination. The model gets one real latent, then feeds its own predictions back into itself for 100 steps.
Problem
An agent learning from pixels solves two problems at once: understanding what it sees, and predicting what happens next. Ha and Schmidhuber’s World Models splits them apart, but the paper only makes sense once you build it.
How it was solved
A convolutional VAE compresses each frame into a 32-dimensional latent. An MDN-LSTM then models how that latent evolves given the action taken: the learned physics of the game.
The dynamics head is a mixture density network, not a regressor, because the environment is stochastic. Several next states are genuinely possible, and MSE would learn their average, which is often none of them.
Results
The dream above: 100 steps of self-fed prediction producing a persistent, plausibly curving road with no simulator running.
Two lessons the paper omits. The collection policy matters more than it looks: random actions leave the track in seconds, so the VAE never learns roads. And the VAE stays frozen while the dynamics train, so physics are learned inside a fixed latent space.