A world model implemented from scratch in PyTorch, following the Vision–Memory–Controller decomposition of Ha & Schmidhuber's World Models (2018). A convolutional VAE compresses game frames into a 32-dimensional latent, and an MDN-LSTM learns the environment's dynamics in that latent space — so the model can imagine driving that never happened, generating coherent video with no environment at all.
Every frame after the first is a hallucination. The model is given one real latent, then feeds its own predictions back into itself for 100 steps — the road persists, curves plausibly, and the car stays on it. No simulator is running.
A world model factors an agent into three parts:
- V (Vision) — a convolutional VAE that compresses each 64×64×3 frame into a 32-dimensional latent. Turns a high-dimensional observation into a small, structured representation.
- M (Memory) — an MDN-LSTM that models how the latent evolves. Given the current latent, the action taken, and its hidden state (a running summary of history), it predicts the next latent. This is the learned physics of the environment, operating entirely in compressed latent space.
- C (Controller) — a small policy acting on the latent and the LSTM's hidden state. (Not implemented in this repo — the V and M components and the dream rollout are the focus.)
The M component doesn't regress to a single next latent. The environment is stochastic — given the current state and an action, several different next states are genuinely possible — and a point estimate trained with MSE would learn the average of those possibilities, which is often none of them.
Instead the MDN outputs a mixture of Gaussians over the next latent: several possible modes, each with a mean, variance, and mixture weight. Sampling from that mixture gives one plausible next state, and the predicted variance doubles as an uncertainty signal — when the model is out of its depth, the modes spread out.
The loss is the negative log-likelihood of the true next latent under the predicted mixture: how probable was what actually happened, according to my distribution?
Top row: real frames. Bottom row: VAE reconstructions. Slightly soft — the characteristic blur of the reconstruction-versus-KL tradeoff — but the road, grass, car, and HUD are all preserved, which is what the dynamics model needs.
The three stages are strictly sequential; each consumes the previous stage's output.
1. collect.py random rollouts in CarRacing -> frames + actions
2. train_vae.py conv-VAE on the frames -> vae.pth
3. encode_rollouts.py encode frames through the VAE -> latent sequences
4. train_mdn.py MDN-LSTM on (latent, action) -> mdn.pth
5. dream_gif.py roll M forward, decode via V -> dream.gif
collect.py # collect rollouts with a temporally-correlated random policy
vae_model.py # ConvVAE architecture and the ELBO loss
train_vae.py # trains V, saves reconstructions
encode_rollouts.py # encodes rollouts into latent sequences (VAE frozen)
mdn_model.py # MDN-LSTM, mixture NLL loss, and mixture sampling
train_mdn.py # trains M on latent dynamics
dream_gif.py # generates the imagined trajectory
pip install -r requirements.txt
python collect.py # ~300 rollouts (takes a while)
python train_vae.py # trains the VAE
python encode_rollouts.py # encodes rollouts to latents
python train_mdn.py # trains the dynamics model
python dream_gif.py # dreamThe collection policy matters more than it looks. A uniformly random policy drives the car off the track within seconds, so most frames end up as empty grass and the VAE learns nothing about roads. Temporally correlated actions (held for several steps) with gentle steering keep the car near the track long enough for the data to be useful.
The VAE is frozen when training M. The dynamics are learned within a fixed latent space, not co-trained with it. This staged training is the structural core of the V-M-C decomposition.
Temperature controls the dreams. Scaling the mixture's variances at sampling time trades stability for variety — low temperature gives conservative, stable rollouts; high temperature gives chaotic ones that dissolve faster.
Ha & Schmidhuber — World Models (2018).
MIT