World Models

David Ha, Jürgen Schmidhuber

Seminal March 2018

The CarRacing agent sees both z and the RNN hidden state h and drives through sharp corners without wobbling.

TL;DR A VAE compresses each frame, an MDN-RNN predicts the next latent, and a linear controller with under 1,100 parameters is evolved on top. It scores 906 ± 21 on CarRacing-v0, and a VizDoom policy trained only inside the learned model survives 1092 ± 556 steps in the real game (750 counts as solved).

What they do

The agent is split into a large world model, which learns without rewards to compress game frames and predict how they evolve, and a tiny linear controller evolved on its features. In VizDoom the world model replaces the game: the policy is trained inside the model’s “dream” and then run in the real game.

How it works

V (vision): a convolutional VAE encodes each 64x64 frame into z (32-D for CarRacing, 64-D for Doom). It is trained on 10,000 rollouts of a random policy.
M (memory): an LSTM with a 5-component mixture-density output models P(z_t+1 a_t, z_t, h_t). For Doom it also predicts whether the agent dies. A temperature τ sets how random the samples are.
C (controller): a_t = W_c [z_t h_t] + b_c (867 parameters for CarRacing). It is optimized with CMA-ES: population 64, 16 rollouts per candidate.
Dream training: M is wrapped as a Gym environment that runs entirely in latent space. C is evolved there and then transferred to the real game unchanged.

Results

CarRacing-v0: 906 ± 21 over 100 random tracks. The best leaderboard entry scored 838 ± 11, and earlier deep-RL methods scored 591-652. With z only, the controller scores 632 ± 251, or 788 ± 141 with an extra hidden layer.
VizDoom Take Cover: the dream-trained policy (τ = 1.15) survives 1092 ± 556 steps in the real game. Solving the task takes 750, and the Gym leader scores 820 ± 58.
Temperature matters. At τ = 0.1 the dream policy scores 2086 in the model but only 193 ± 58 in the real game, below a random policy (210 ± 108). At τ = 1.30 it scores 753 ± 139, a lower mean with less variance.

Limitations

The controller exploits model errors (it learns to make dream fireballs vanish), so τ has to be tuned. The VAE ignores the task: it reproduces irrelevant wall bricks in Doom but misses CarRacing’s road tiles. Random-policy data suffices only for these simple games; the iterative exploration loop for harder tasks is proposed but untested. The LSTM has limited capacity and plans step by step, with no hierarchy.

Videos

A policy trained only inside the dream environment, deployed in the actual VizDoom Take Cover game.

Ablation: a controller that sees only z wobbles and misses sharp corners.

Figures

From Figure 3: Frames from generative world models of CarRacing (left) and VizDoom (right). The models are trained on... — From Figure 3: Frames from generative world models of CarRacing (left) and VizDoom (right). The models are trained on observations recorded from the actual games and can then stand in for them when training agents.

From Figure 8: The agent. V encodes each observation into z; C maps z and M's hidden state h to an action; M takes z ... — From Figure 8: The agent. V encodes each observation into z; C maps z and M's hidden state h to an action; M takes z and the action to update h for the next step.

From Figure 6: The M model, an RNN with a mixture density output layer that gives a mixture of Gaussians over the nex... — From Figure 6: The M model, an RNN with a mixture density output layer that gives a mixture of Gaussians over the next latent vector z.

From Figure 29: Time steps survived in the actual VizDoom environment over 100 trials by the policy trained in the dr... — From Figure 29: Time steps survived in the actual VizDoom environment over 100 trials by the policy trained in the dream; score 1092 ± 556.

Abstract

We explore building generative neural network models of popular reinforcement learning environments. Our world model can be trained quickly in an unsupervised manner to learn a compressed spatial and temporal representation of the environment. By using features extracted from the world model as inputs to an agent, we can train a very compact and simple policy that can solve the required task. We can even train our agent entirely inside of its own hallucinated dream generated by its world model, and transfer this policy back into the actual environment. An interactive version of this paper is available at https://worldmodels.github.io/

Summaries are written automatically from each paper's text. Figures and videos come from the paper and its project page; all credit goes to the authors.

Matt Strong