PaperlayerView exact revision on arXiv

Source-linked AI summary

World Models

David Ha, Jürgen Schmidhuber

arXiv:1803.10122v4cs.LGstat.ML

TL;DR

The paper addresses how to train expressive agents for reinforcement-learning environments despite credit-assignment difficulties in large policies. It learns a compressed, predictive world model separately from a small controller, enabling compact policies and successful task performance, including car racing. The discussion notes limitations in task relevance and memory capacity.

  • Problem

    Credit-assignment difficulties make it hard for traditional model-free RL methods to train large neural networks, motivating more efficient training of expressive agents.

  • Method

    The method trains a large world model unsupervised to represent and predict environments, then trains a small controller using those representations, including through efficient optimization methods.

  • Results

    The approach learns a highly compact policy and achieves the score required to solve the CarRacing-v0 task.

  • Takeaways & Limitations

    Separating world-model complexity from controller learning provides a simplified framework for applying RNN-based world models and controllers across RL environments.

  • Takeaways & Limitations

    A standalone VAE may encode task-irrelevant details, while the LSTM-based world model has limited capacity and can suffer catastrophic forgetting.

Abstract

from arXiv · show

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/

1. Introduction

The paper frames reinforcement-learning agents as relying on compact internal models that represent and predict environments. It proposes separating a large world model from a small controller to retain expressive representations while easing policy learning.

  • Motivation: Humans use abstract spatial and temporal representations and predictions of future sensory information to perceive and act in the world.The introduction connects this predictive modeling principle to fast, largely instinctive behavior.
  • Motivation: RL agents benefit from representing past and present states and predicting future states, potentially with expressive recurrent neural networks.The paper identifies predictive representations as useful for reinforcement-learning problems.
  • Problem: Credit assignment makes it difficult for model-free RL methods to train millions of parameters, so practical agents often use smaller networks.Smaller networks can iterate faster toward a good policy during training.
  • Approach: The proposed framework divides the agent into a large, unsupervised world model and a small controller trained to perform the task.This division lets the controller search a small space without sacrificing the larger world model’s capacity and expressiveness.
  • Contribution: The resulting agent can learn a highly compact policy while retaining the expressive capacity of a large neural network.The framework is presented as a simplified way to demonstrate concepts from earlier work on RNN-based world models and controllers.

2. Agent Model

The agent combines compressed visual observations, predictive memory, and a small controller. Vision produces latent codes, memory predicts future codes, and the controller maps current and predicted information to actions.

  • 2. Agent Model: The agent consists of Vision (V), Memory (M), and Controller (C), which respectively encode observations, predict future representations, and choose actions.The controller acts only on representations produced by the vision and memory components.
  • 2.1. VAE (V) Model: V uses a Variational Autoencoder to compress each high-dimensional image frame into a small latent vector z.The Gaussian prior limits latent information capacity and improves robustness to unrealistic vectors generated by M.
  • 2.2. MDN-RNN (M) Model: M uses an MDN-RNN to model the probability distribution of the next latent vector zt+1 from current and past information, including action at, latent state zt, and hidden state ht.A temperature parameter τ controls uncertainty during sampling.
  • 2.3. Controller (C) Model: C is a single-layer linear model that maps the concatenated vector [zt ht] directly to the action at.Its parameters are Wc and bc, and it is deliberately kept small so most model complexity remains in V and M.
  • 2.4. Putting V, M, and C Together: At each time step, V encodes the observation, C uses zt and ht to produce an action, and M updates its hidden state from zt and at.The controller is optimized separately with CMA-ES using parallel environment rollouts.

3. Car Racing Experiment

The CarRacing-v0 experiment tests whether world-model features support compact continuous control from pixels. The agent reaches the task’s solving threshold and exceeds prior reported scores while using separately trained visual and predictive models.

  • 3. Car Racing Experiment: The agent uses a predictive world model to extract spatial and temporal features for a compact controller that drives from pixel inputs in CarRacing-v0.The environment randomly generates tracks and provides steering, acceleration, and braking actions.
  • Training: 10,000 random rollouts provide the data for training V to encode frames and M to model P(zt+1 | at, zt, ht).V is trained on reconstructed frames, and M models the latent-state transition as a mixture of Gaussians.
  • Results: CarRacing-v0 defines solving as an average reward of 900 over 100 consecutive trials, allowing only a few driving mistakes.The experiment describes the reported agent as the first known solution to achieve the required score.
  • Training: The world model excludes environmental reward signals, while the 867-parameter controller alone accesses rewards and is optimized by an evolutionary algorithm.The V and M models only compress and predict image-frame sequences.
  • Training: The 32-dimensional latent code z ∈R32 is used with M’s hidden state as input to the linear controller.C is optimized with CMA-ES to maximize expected cumulative reward.

V Model Only

Using only the visual representation z_t makes driving unstable, while adding the memory state h_t provides predictive information that substantially improves control. The combined representation enables stable cornering and solves the racing task.

  • V Model Only: The controller initially receives only the visual representation z_t, defined as a linear policy input.This setting tests driving with momentary visual features without the memory state.
  • V Model Only: 632 ± 251 average score over 100 random trials matches other leaderboard and traditional Deep RL agents, but the agent wobbles and misses sharper corners.Adding a hidden layer raises performance to 788 ± 141, still insufficient to solve the environment.
  • V Model Only: The visual representation z_t captures the current moment, whereas h_t summarizes predictive information learned by the memory model.Combining z_t with h_t therefore gives the controller both current-observation and future-predictive features.
  • V Model Only: Access to both z_t and h_t greatly improves driving stability and lets the agent attack sharp corners effectively.The controller can query the memory state’s future distribution for reflexive action decisions without rolling out hypothetical scenarios.
  • Full World Model (V and M): 906 ± 21 average score over 100 random trials solves CarRacing-v0 and exceeds the best reported leaderboard score of 838 ± 11.The world model learns directly from raw RGB streams, unlike cited methods using frame preprocessing and recent-frame stacking.

4. VizDoom Experiment

The VizDoom experiment trains a controller inside a learned, latent-space world model and transfers it to the actual environment. Despite model inaccuracies and exploitable dynamics, the transferred policy substantially exceeds the task threshold.

  • Setup: The experiment asks whether a policy can be trained entirely inside a world model’s hallucinated VizDoom environment and transferred back to reality.The virtual environment is built to mimic VizDoom, and its interface matches the real environment for deployment.
  • Setup: The pipeline collects 10,000 random rollouts, trains a VAE and MDN-RNN, defines a controller from latent and hidden states, then optimizes it with CMA-ES.The MDN-RNN predicts the next latent frame and death event, allowing it to serve as a complete virtual environment.
  • Virtual-environment result: ∼900 time steps: the controller learns to navigate the dream environment and escape fireballs launched by monsters generated by the world model.The agent is trained using survival time as its objective.
  • Transfer result: The world model learns essential game properties from raw images, including game logic, enemy behaviour, physics, and 3D graphics rendering.The learned representation is sufficient for the controller to navigate the real environment even when visual details are imperfectly reproduced.
  • Transfer result: ∼1100 time steps over 100 random consecutive trials: the transferred policy exceeds the required 750-step survival score in the original VizDoom scenario.Its real-environment score is also higher than the score obtained inside the more difficult virtual environment.
  • World-model exploitation: Low sampling temperature causes mode collapse: the dream omits fireballs, producing policies that score 2100 virtually but fail in the actual world.Higher temperature makes adversarial policies harder to find, but excessive temperature makes the virtual environment too difficult to learn.

5. Iterative Training Procedure

For more difficult environments, the paper proposes iteratively improving the world model with new observations gathered through active exploration. Curiosity can be encouraged by rewarding regions where the model predicts poorly.

  • Motivation: More sophisticated environments require iterative training because the agent must explore strategically to reveal new parts of the world.The world model is refined over time as the agent collects additional observations.
  • Procedure: The procedure initializes M and C, rolls out the actual environment, trains M on observed transitions, trains C inside M, and repeats until completion.The iterative model predicts observations, rewards, actions, and done states for the next time step.
  • Scope: One iteration was enough to solve the paper’s simple tasks, whereas harder tasks require exploration that improves the world model.The authors identify intrinsic motivation and information seeking as possible ways to encourage novel exploration.
  • Curiosity: Poor MDN-RNN prediction loss identifies unfamiliar parts of the environment that can guide curiosity-driven exploration.Flipping the sign of the model’s loss in the actual environment rewards collecting data from unfamiliar regions, potentially improving M.

6. Related Work

Related work spans dynamics-model learning, visual prediction, recurrent world models, and evolution-based policy optimization. This paper situates its approach among methods that learn from high-dimensional visual observations and train controllers using learned dynamics.

  • Dynamics models: Model-based RL methods learn dynamics models and use them to train policies, including PILCO’s Gaussian-process trajectory sampling for control tasks.The paper contrasts its visual, recurrent setting with approaches using lower-dimensional, well-defined states.
  • Visual dynamics: The paper focuses on modeling dynamics from sequences of raw pixel frames, a challenging setting for robotic control and visual RL.Earlier active-vision work predicted video frames and used them to train a fovea-shifting control network.
  • Sequence modeling: Feed-forward models have predicted future video-game frames, while RNNs capture longer-term temporal dependencies and can hallucinate game levels.The cited RNN work demonstrates probabilistic modeling of Atari environments.
  • Recurrent world models: RNN-based internal world models were explored in early work and later developed into frameworks for reasoning about the future and generating long sequences.Subsequent research used recurrent models both to generate future frames and to support future-oriented reasoning.
  • Policy optimization: Evolution strategies are useful for controller training because they require only final cumulative reward and parallelize across many rollouts.The paper uses evolution strategies to optimize the compact controller while training the world model with backpropagation.

7. Discussion

The world model enables training agents inside simulated latent-space environments, while its design introduces task-relevance and capacity limitations. Its recurrent, differentiable structure also supports accelerated simulation and potential direct policy optimization.

  • Practical benefits: Training entirely inside a simulated latent-space dream can reduce reliance on computationally expensive real-environment rendering and physics.The authors also suggest incremental world models may support transferring policies back to the real world.
  • Practical benefits: A fully differentiable recurrent world model may allow direct backpropagation through the dream to fine-tune a policy.GPU-based distributed simulation can further accelerate world-model rollouts.
  • Limitations: Standalone unsupervised VAE training can encode irrelevant visual details while missing task-relevant features.Joint reward prediction may improve task focus, but can reduce reuse across new tasks without retraining.
  • Limitations: The world model’s limited capacity and catastrophic forgetting constrain how much historical information its recurrent weights can retain.The authors propose higher-capacity recurrent models or external memory for more complicated worlds.
  • Limitations: The approach simulates futures step by step without human-like hierarchical planning or abstract reasoning that can ignore irrelevant details.The broader Learning To Think approach is described as allowing recurrent controllers to address and reuse world-model subroutines.

A. Appendix

This appendix introduces the models and training methods used in the work.

  • The appendix describes the models and training methods used in this work.

A.1. Variational Autoencoder

The V Model uses a convolutional variational autoencoder to compress image observations into low-dimensional latent representations. Its Gaussian prior constrains information capacity and improves robustness to unrealistic latent vectors generated by the temporal model.

  • Model: The ConvVAE serves as the V Model and compresses each image frame into a low-dimensional latent vector.The model uses a Gaussian prior over z to limit per-frame information capacity.
  • Architecture: 64x64x3 image inputs pass through four convolutional layers that encode vectors µ and σ, each of size Nz.The latent vector z is sampled from the encoded distribution.
  • Architecture: Convolutional and deconvolutional layers use stride 2, with ReLU activations except at the output layer.The output activation keeps reconstructed pixel values between 0 and 1.
  • Training: The model was trained for 1 epoch on random-policy data using reconstruction L2 loss together with KL loss.

A.2. Recurrent Neural Network

The M Model is an LSTM-based Mixture Density Network that predicts distributions over future latent vectors rather than deterministic next states. Sampling these distributions generates hallucinated environments, with an additional death-probability mechanism for Doom.

  • Model: The M Model combines an LSTM with a Mixture Density Network to model the probability distribution of the next latent vector.It predicts a mixture of Gaussian distributions rather than a deterministic latent state.
  • Generation: The MDN-RNN samples future latent vectors from its predicted distribution at each time step to generate hallucinated environments.The model conditions predictions on the current action, latent vector, and recurrent hidden state.
  • Doom termination: In Doom, the MDN-RNN also predicts whether the agent has died and sets done to true when that probability exceeds 50%.The cutoff is reported as more stable than sampling from a Bernoulli distribution for this low-probability event.
  • Training: The MDN-RNNs were trained for 20 epochs on random-policy data, using 256 LSTM hidden units for Car Racing and 512 for Doom.Both tasks used five Gaussian mixtures without modeling the correlation ρ parameter.
  • Training: Teacher-forced training samples z from pre-computed frame-wise distributions to prevent overfitting to a specific sampled latent vector.

A.3. Controller

The controller uses bounded continuous actions and CMA-ES to evolve policy weights across randomized rollouts. In CarRacing-v0, the evolved agent reached the required average score threshold after 1800 generations, while the figures also report results for CarRacing and VizDoom evaluations.

  • Action space: Tanh nonlinearities bound CarRacing steering, acceleration, and braking actions to their task-specific ranges.Steering spans -1 to 1, while acceleration and braking span 0 to 1.
  • Optimization: CMA-ES evolved the C Model using 64-agent populations, with fitness equal to average cumulative reward over 16 random-seed rollouts.The population’s best, worst, and mean fitness were tracked during evolution.
  • CarRacing-v0: 900.46 average score was achieved in CarRacing-v0 after 1800 generations, evaluated across 1024 random rollouts.The task requirement was an average score above 900 over 100 random rollouts.
  • CarRacing-v0: 906 ± 21 cumulative reward was reported for the CarRacing-v0 evaluation.Additional CarRacing figures report scores of 632 ± 251 using only z_t and 788 ± 141 using only z_t with a hidden layer.
  • VizDoom: 1092 ± 556 time steps survived over 100 consecutive trials in the actual VizDoom environment.The surrounding figure sequence also documents DoomRNN training.
Loading 1803.10122v4…