Recurrent World Models Facilitate Policy Evolution

lastmod 2026-10-28created 2026-10-283 edits

World Models is a paper published in 2018 by David Ha and Jürgen Schmidhuber. It builds a generative world model trained in an unsupervised manner on reinforcement learning environments, and shows that an agent can be trained entirely inside a "dream" produced by that model and that the policy can then be transferred to the actual environment.

• Subtitle: "Can agents learn inside of their own dreams?"
• Authors: David Ha (Google Brain), Jürgen Schmidhuber (NNAISENSE, IDSIA)
• Released: 2018-03-27 (arXiv:1803.10122)
• Conference: NeurIPS 2018 (presented as Recurrent World Models Facilitate Policy Evolution)
• Format: interactive web paper that runs the models directly in the browser (uses p5.js, deeplearn.js)

Categories

• Model Works
• Reinforcement learning

Abstract

Builds generative neural network models (World Models) of popular reinforcement learning environments. The World Model can be trained quickly in an unsupervised manner to learn a compressed spatial and temporal representation of the environment. Using features extracted from the World Model as inputs to an agent, a very compact and simple policy can solve the task. Furthermore, it shows that an agent can be trained entirely inside a hallucinated dream generated by the World Model, and that the policy still works when transferred back to the actual environment.

Key ideas

• Humans build an internal model (mental model) of the world from information obtained through limited senses, and make decisions and act based on it. The brain perceives while predicting the future; a baseball batter reflexively swinging the bat before consciously planning is also thanks to this predictive model.
• Large, expressive RNNs are advantageous, but model-free RL has difficulty training large networks because of the credit assignment problem, so it usually uses small networks.
• So the agent is split in two.
  • Large World Model (V, M): compresses and predicts the environment through unsupervised learning.
  • Small Controller (C): learns actions through reinforcement learning. Its small search space lets it focus on credit assignment.
• The Controller can be trained entirely in a virtual environment generated by the World Model instead of the actual environment. The problem of the agent exploiting flaws in the model is mitigated by raising uncertainty with the temperature τ.

Agent Model

The agent consists of three components.

V: VAE

• A Convolutional Variational Autoencoder that compresses the 2D image frame at each time step into a low-dimensional latent vector \(z\) and reconstructs it.
• The input is a 64×64×3 RGB image. It uses a 4-layer convolution encoder and a 4-layer deconvolution decoder; all layers use stride 2 and, except for the output layer, ReLU.
• The latent dimension is 32 for Car Racing and 64 for Doom.
• The loss function is the sum of the L2 reconstruction loss and the KL loss. The Gaussian prior limits information capacity, but makes it robust to unrealistic \(z\) produced by M.

M: MDN-RNN

• Compresses what V sees over time. Instead of predicting the next \(z\) deterministically, it outputs a probability density \(p(z)\).
• An LSTM with a Mixture Density Network output layer attached. It mixes 5 Gaussians and uses a diagonal covariance.
• The number of hidden units is 256 for Car Racing and 512 for Doom.
• Trained for 20 epochs with teacher forcing. Overfitting is reduced by resampling \(z\) for each batch from the stored \(\mu, \sigma\).
• When sampling, the temperature τ controls the model's uncertainty.
• Works the same way as earlier MDN-RNN-family sequence generation models such as SketchRNN.

C: Controller

A single linear layer trained to maximize the expected cumulative reward over a rollout.

\(a_t = W_c \; [z_t \; h_t] \; + \; b_c\)

• Because \(h_t\) already contains the information needed to predict the future, the predicted \(z_{t+1}\) is not fed as input.
• Most of the model complexity is placed in V and M, keeping C small. This allows it to be optimized with CMA-ES (an evolution strategy) instead of backpropagation. CMA-ES works well in spaces of up to a few thousand parameters and is easy to parallelize.
  • The population size is 64, and each individual's fitness is the average cumulative reward over 16 rollouts with different seeds.
• The output is fitted to the action range with tanh.

Combining V, M, and C

def rollout(controller):
  ''' env, rnn, vae are '''
  ''' global variables  '''
  obs = env.reset()
  h = rnn.initial_state()
  done = False
  cumulative_reward = 0
  while not done:
    z = vae.encode(obs)
    a = controller.action([z, h])
    obs, reward, done = env.step(a)
    cumulative_reward += reward
    h = rnn.forward([a, z, h])
  return cumulative_reward

Car Racing experiment

Drives in OpenAI Gym's CarRacing-v0 from pixel input alone. The track is randomly generated every time, and an average of 900 or more over 100 consecutive trials counts as solved. Actions are continuous values for steering (−1 to 1), acceleration (0 to 1), and brake (0 to 1).

Training procedure

1. Collect 10,000 rollouts with a random policy.
2. Train V (VAE) to encode each frame into \(z \in \mathbb{R}^{32}\).
3. Train M (MDN-RNN) to model \(P(z_{t+1} \mid a_t, z_t, h_t)\).
4. Evolve C with CMA-ES to maximize the expected cumulative reward.

The World Model (V, M) never sees the reward signal and only learns to compress and predict frames. Only C has access to the reward.

Results

• V only (input \(z_t\) only): driving is wobbly and the car leaves the track on sharp corners. A linear C scores 632 ± 251, and even with an added hidden layer (40 tanh units, 1,443 parameters) it scores 788 ± 141.
• V and M (input \(z_t\), \(h_t\)): driving is stable and sharp corners are handled well. It does not explicitly plan the future, but reacts 'instinctively' using the information about the future probability distribution contained in \(h_t\).
• After 1,800 generations it reached an average of 900.46 over 1,024 random rollouts, and the result over 100 trials is 906 ± 21.

Earlier Deep RL methods needed preprocessing such as edge detection and frame stacking. World Models learned spatial and temporal representations directly from raw RGB pixels and, to the authors' knowledge, was the first to reach the solving criterion for this task.

Car Racing Dreams

Since M can predict the future, a sampled \(z_{t+1}\) can be fed back in as if it were a real observation to generate imaginary driving scenes. The trained C can be placed in this dream environment to drive, and the web paper lets you try it yourself while adjusting τ.

VizDoom experiment: learning inside a dream

In VizDoom's DoomTakeCover-v0, the agent must survive as long as possible by dodging fireballs shot by monsters on the other side of the room. Survival time is the cumulative reward. One episode lasts at most 2,100 time steps (about 60 seconds), and an average of 750 time steps (about 20 seconds) or more over 100 consecutive trials counts as solved.

Unlike Car Racing, M predicts not only the next \(z\) but also whether the agent dies, \(d\). So M can be wrapped in a Gym environment interface and used as a full virtual environment. This virtual environment does not render pixels and operates only in latent space. If the probability of death exceeds 50%, done is set to True.

Training procedure

1. Collect 10,000 rollouts with a random policy.
2. Train V (VAE) to encode each frame into \(z \in \mathbb{R}^{64}\).
3. Train M (MDN-RNN) to model \(P(z_{t+1}, d_{t+1} \mid a_t, z_t, h_t)\).
4. Evolve C with CMA-ES to maximize survival time inside the virtual environment.
5. Apply the learned policy to the actual environment.

Training inside the dream and policy transfer

• From raw images alone, M imitates the game logic, monster behavior, physics, and 3D graphics. It also learns being blocked by walls and the movement of multiple fireballs.
• About 900 time steps were achieved in the virtual environment. When this policy was transferred to the actual environment, it recorded an average of about 1,100 time steps over 100 trials. This is well above the solving criterion (750), and higher even than the score in the virtual environment, which was set to be harder.
• Even though V cannot accurately reproduce details such as the number of monsters, the agent makes good use of the policy learned in the noisy virtual environment in the clean actual environment.

Cheating the World Model

• An early version of the agent found an adversarial policy that made the monsters never shoot fireballs at all.
• Because M is a probabilistic model approximating the actual environment, it sometimes generates trajectories that do not follow the laws of reality. C also has access to M's hidden state, so it easily exploits these flaws. Deterministic models are even more vulnerable.
• The countermeasure is to model M as a probability distribution (a mixture of Gaussians) and raise the temperature τ to make the virtual environment more uncertain and harder.
  • When τ is very low (e.g., 0.1), mode collapse occurs and the monsters do not shoot fireballs. The agent then gets a near-perfect score in the virtual environment but does worse than a random policy in the actual environment.
  • A mixture of Gaussians is used instead of a single Gaussian to represent discrete events such as "whether the monster shoots or not".

• Raising τ makes adversarial policies harder to find, but raising it too far leaves the agent unable to learn anything. τ is therefore a hyperparameter that must be tuned.
• The best actual environment score was obtained at τ=1.15. At τ=1.30, a safer policy with a lower mean but smaller variance emerged.

Iterative Training Procedure

Data collected with a random policy alone cannot sufficiently explore complex environments. So the paper proposes repeating the following procedure.

1. Initialize M and C with random parameters.
2. Roll out N times in the actual environment, saving all actions \(a_t\) and observations \(x_t\).
3. Train M to model \(P(x_{t+1}, r_{t+1}, a_{t+1}, d_{t+1} \mid x_t, a_t, h_t)\), and train C to optimize the expected reward inside M.
4. If the task is not finished, go back to step 2.

• One iteration is enough for simple tasks, but complex tasks require multiple iterations.
• Using M's loss as C's reward yields artificial curiosity. The agent collects new data by seeking out situations that M does not predict well.
• If M also predicts the action \(a_{t+1}\), it can absorb C's behavior. Low-level skills such as walking can then be left to M while C focuses on high-level skills.
• Neuroscientifically, this connects to the phenomenon in which the hippocampus replays recent experiences to consolidate memories while an animal rests or sleeps.

Related Work

• PILCO: a probabilistic model-based method that learns environment dynamics with Gaussian processes. It is efficient only with low-dimensional observations.
• Prior work introduced includes dynamics models based on Bayesian neural networks (BNNs) and research that learned dynamics in an autoencoder's latent space to control a pendulum from pixel input.
• Guzdial et al.'s CNN-based forward simulation model of games, Graves's "Hallucination with RNNs" lecture, and Finn's Model-Based RL lecture are mentioned.
• Combining an RNN-based World Model with a Controller was already explored in Schmidhuber's 1990 work "Making the World Differentiable" and others, and "Learning to Think" generalized it.
• The evolution strategies (ES) used to train the Controller need only the final cumulative reward and are easy to parallelize.

Discussion

• Pros
  • Instead of the heavy rendering and physics computation of a game engine, the agent can train as much as it wants in a lightweight virtual environment. Distributed training on GPUs is also easy.
  • Since the World Model is differentiable, the policy in the dream can also be fine-tuned directly with backpropagation.
• Limitations
  • Training the VAE separately can make it reproduce task-irrelevant parts such as brick patterns while missing important parts such as road tiles. Training it together with an M that predicts rewards focuses it on task-relevant features, but makes it harder to reuse for other tasks.
  • An LSTM-based M has limited capacity and suffers from catastrophic forgetting when trained with backpropagation. Larger models or external memory are needed.
  • It only simulates the future time step by time step, without hierarchical planning or abstract reasoning. It can be extended in a way where C queries M as a subroutine, as in "Learning to Think".
• Merging C and M into a single network (One Big Net) is also suggested as future work.

Citation

@incollection{ha2018worldmodels,
  title = {Recurrent World Models Facilitate Policy Evolution},
  author = {Ha, David and Schmidhuber, J{\"u}rgen},
  booktitle = {Advances in Neural Information Processing Systems 31},
  pages = {2451--2463},
  year = {2018},
  publisher = {Curran Associates, Inc.},
  url = {https://papers.nips.cc/paper/7512-recurrent-world-models-facilitate-policy-evolution},
  note = "\url{https://worldmodels.github.io}",
}

Diagrams and text are licensed under CC-BY 4.0. However, figures taken from other sources are excluded.

Documentation

World Models

https://arxiv.org/abs/1803.10122

Recurrent World Models Facilitate Policy Evolution (NeurIPS 2018)

https://papers.nips.cc/paper/7512-recurrent-world-models-facilitate-policy-evolution

See also

• World Models - the general concept of world models
• World Foundation Model
• Reinforcement learning
• Autoencoder
• VAE
• RNN
• LSTM
• Mixture Density Network
• Gaussian mixture model
• Physical AI

Favorite site

• World Models - Can agents learn inside of their own dreams?
• Github - worldmodels/worldmodels.github.io
• World Models Experiments - guide to reproducing the experiments (author's blog)
• World Models talk video (YouTube)