World Models
David Ha, Jürgen Schmidhuber
cs.LG, stat.ML
2018-03-27
A VAE plus MDN-RNN world model lets a 867-parameter linear controller score 906 on CarRacing and transfer a policy trained entirely in its dream to VizDoom.
Model-free RL has to push credit assignment through the whole net, so in practice policy networks stay in the 10^3–10^6 parameter range. Reflexive human action does not unroll the future step by step; it queries a compressed internal model. Schmidhuber wrote the RNN-world-model-plus-controller story across the 1990s and 2015. What was missing was a stripped-down experiment that actually runs on pixel environments. This paper is that experiment: train a large world model unsupervised, and leave credit assignment to a tiny controller.
The agent has three parts. Vision V is a convolutional VAE that compresses 64×64 RGB frames into a diagonal-Gaussian latent z. Memory M is an RNN with a mixture-density head (MDN-RNN) that models P(z{t+1}|at,zt,ht); the mixture is there for stochastic discrete events. Controller C is a linear map at=Wc[zt;ht]+bc. V and M train with backprop on a GPU; C is evolved with CMA-ES on CPUs, because the search space is a few hundred to a thousand dimensions.
CarRacing: 10,000 random rollouts; V encodes frames to z∈R^{32}; M learns the mixture for the next z; CMA-ES maximizes real-environment return. Parameter counts are about 4.35M (V), 422k (M), 867 (C). The world model never sees reward; only C does.
For VizDoom Take Cover, M also predicts a death bit dt, so M itself is a Gym interface. C is evolved entirely in the latent dream and then copied into the real game. Sampling temperature τ makes the dream noisier than the real env, which makes it harder for C to exploit holes in M.
CarRacing-v0 is solved at a mean of 900 over 100 trials. Giving C only z, not h, scores 632±251, in the A3C ballpark. A hidden layer on C reaches 788±141, still short. z and h together score 906±21. Baselines: DQN 343±18, continuous A3C 591±45, discrete A3C 652±10, Gym leaderboard 838±11. The authors report this as the first solution of the task, from a raw RGB stream, with no edge detection and no frame stacking.
Take Cover is solved at 750 steps alive (about 20 seconds) averaged over 100 trials, capped at 2100. The dream scores about 900; transfer to the real env scores about 1100. Temperature scan:
| τ | dream | real |
| 0.10 | 2086±140 | 193±58 |
| 1.00 | 1145±690 | 868±511 |
| 1.15 | 918±546 | 1092±556 |
| 1.30 | 732±269 | 753±139 |
| random | — | 210±108 |
| Gym leader | — | 820±58 |
At τ=0.1 the dream is almost a deterministic LSTM, monsters never shoot, C gets a perfect dream score and transfers worse than random. τ=1.15 is best on the real env, with large variance; τ=1.30 is lower and more stable.
C also finds an adversarial dream policy that extinguishes fireballs, because M generates off-distribution trajectories. MDN stochasticity and higher τ are the patch. The patch is incomplete.
Once “learn the world” and “learn the policy” are split, the large model can take differentiable unsupervised training and the small controller can take evolution. Credit assignment is no longer stuck on a million weights. Training in the dream means fewer rendering and physics cycles; the policy rolls in latent space. Before the sim2real slogan took over, this paper already ran “learn in the dream, deploy in reality” on two Gym tasks.
The caveats are equally plain. Random-policy data is enough only for simple tasks; harder worlds need iterative exploration. The VAE encodes irrelevant brick textures and drops useful track tiles. The controller can exploit the world model, which is the generic failure of replacing the real env with a learned one.
CarRacing “solved” hangs on the 900/100-trial line; a standard deviation of 21 still means occasional crashes. VizDoom’s 1100 steps come with ±556, so the policy is unstable. The world model is trained on random-policy data and fabricates unseen states. τ is yet another hyperparameter; too low and modes collapse. C sees all of M’s hidden state, equivalent to handing the player the game engine’s memory, which inflates exploitability. Iterative training, curiosity, and predicting actions and rewards are listed as future work, with no experiment. VAE and MDN are trained separately, not end-to-end.