World Models
Paper: Ha, D. & Schmidhuber, J. (2018). World Models. arXiv:1803.10122.
Authors: David Ha (Google Brain Tokyo), Jürgen Schmidhuber (NNAISENSE / IDSIA)
Published: NIPS 2018
arXiv: https://arxiv.org/abs/1803.10122
Interactive demo: https://worldmodels.github.io/
Summary
World Models proposes that an RL agent can learn to solve tasks by first building an internal generative model of its environment — a “world model” — and then training its policy not in the real environment, but inside the dream that the world model generates. The system has three components trained in sequence: (1) a Variational Autoencoder (VAE) that compresses each observation frame into a compact latent vector; (2) a Mixture-Density-Network–Recurrent Neural Network (MDN-RNN) that predicts the distribution of the next latent state given the current latent state, action, and a recurrent hidden state; and (3) a tiny Controller — a linear policy that maps (latent, hidden) → action — trained with evolution strategies (CMA-ES). On the CarRacing-v0 environment, the controller trained inside the VAE+MDN-RNN “dream” transfers back to the real environment and achieves competitive scores. The paper’s central insight is that most of the difficulty in RL comes from learning in the raw observation space; if you first learn a good internal representation and a predictive model, the policy itself can be trivially small.
Core idea: Separate perception (VAE), memory/prediction (MDN-RNN), and action (linear controller) into three modules, train them unsupervised first, then train the controller inside the learned dream.
Key method details for the code:
- VAE (“V”): A convolutional encoder–decoder trained to reconstruct 64×64 pixel frames. The encoder outputs the mean and log-variance of a Gaussian; the latent vector z is sampled via the reparameterization trick. We use a latent dimension of 32.
- MDN-RNN (“M”): An LSTM that takes (z_t, a_t) as input and predicts the distribution of z_{t+1} as a mixture of K Gaussians (mixture-density network). The LSTM hidden state h_t carries temporal context.
- Controller (“C”): A linear map
(z_t, h_t) → a_t. Only a few hundred parameters. Trained with CMA-ES on the dream — rollout episodes are generated by the VAE decoder + MDN-RNN, not the real simulator. - Dream training: After the VAE and MDN-RNN are trained on real rollouts, the controller is optimized entirely inside the hallucinated environment: the VAE decoder renders frames, the MDN-RNN steps the latent dynamics, and the controller’s reward comes from the dream.
- Simplifications: We use a 2D toy “point-mass navigation” environment (instead of CarRacing-v0) to keep everything trainable on CPU in minutes. The architecture (VAE → MDN-RNN → linear controller → dream training) is faithful to the paper.
Influence: World Models has been cited over 3,000 times and influenced subsequent work on model-based RL, dreamer-style world models (Dreamer, DreamerV2/V3), and the idea that agents can learn in their own simulated imagination. The “can agents learn inside their own dreams?” framing became iconic in the model-based RL community.
Kaggle:
What problem does it solve
Imagine you’re learning to drive a car for the first time. The real world is overwhelming — you see millions of pixels of road, trees, other cars, and you have to figure out steering and gas from all that raw visual mess. That’s really hard.
This paper says: what if you first built a mental model of the road in your head — a little simulation of how the world works — and then practiced driving inside your imagination instead of on the real road? You could try things a thousand times faster, make mistakes safely, and when you finally drive the real car, your brain already knows what to do.
That’s exactly what World Models does for an AI agent. It has three parts: (1) a Vision module that compresses what it sees into a small summary, (2) a Memory module that predicts what happens next based on past experience, and (3) a tiny Controller that decides what action to take. The agent first watches the real world to learn its vision and memory, then closes its eyes and practices inside its own dream. When it wakes up, the controller works on the real thing. This is how humans instinctively predict a baseball’s path and swing before consciously thinking — your internal model does the work.
Resources
Kaggle Notebook:
| File | Description |
|---|---|
| README.md | This overview — a 5-minute read covering the paper’s summary, core idea, key method details, influence, and a plain-English explanation of the problem it solves |
| solution.ipynb | Complete from-scratch PyTorch implementation of the paper’s architecture, with training loops, visualizations, and worked examples |
| CODE_ARCHITECTURE.md | Section-by-section breakdown of the notebook — what each cell does, data flow, tensor shapes, and key design decisions |
| TALK.md | Press coverage, blog posts, interview Q&A, citations, and common misconceptions about the paper |