PaperlayerView exact revision on arXiv

Source-linked AI summary

Recurrent World Models Facilitate Policy Evolution

David Ha, Jürgen Schmidhuber

arXiv:1809.01999v1cs.LGstat.ML

TL;DR

The paper addresses whether compact policies can exploit learned predictive representations instead of relying on extensive direct interaction with reinforcement-learning environments. It combines unsupervised visual and temporal modeling with evolutionary controller training, including training inside generated environments, and reports state-of-the-art task results. The approach’s feature learning has task-relevance limitations, especially when the visual model is trained standalone.

  • Problem

    Reinforcement-learning agents may benefit from predictive models, but existing model-based approaches still train their policies in the actual environment.

  • Method

    The agent uses a VAE-based visual compressor V, an MDN-RNN world model M, and a compact controller C trained by evolution from their representations.

  • Results

    The method solves the CarRacing-v0 task from raw RGB pixels, achieving an average score of 906 ± 21 and new state-of-the-art results.

  • Takeaways & Limitations

    Controllers can be trained inside environments generated by their own world models and transferred back to actual environments.

  • Takeaways & Limitations

    Standalone unsupervised visual learning can encode irrelevant details while missing task-relevant features, limiting reuse across tasks without retraining.

Abstract

from arXiv · show

A generative recurrent neural network is quickly trained in an unsupervised manner to model popular reinforcement learning environments through compressed spatio-temporal representations. The world model's extracted features are fed into compact and simple policies trained by evolution, achieving state of the art results in various environments. We also train our agent entirely inside of an environment generated by its own internal world model, and transfer this policy back into the actual environment. Interactive version of paper at https://worldmodels.github.io

1 Introduction

The paper builds recurrent world models that learn predictive representations of environments and uses compact controllers to act from those representations. It additionally trains controllers inside generated environments and addresses model imperfections through uncertainty.

  • Predictive world models: A recurrent model M learns to predict future observations from past information in partially observable reinforcement-learning environments.M uses recurrent memory to improve predictions based on previous observation sequences.
  • Predictive world models: The world model combines unsupervised prediction with a controller C that learns task behavior from M’s internal representations.The controller is intentionally small, limiting its credit-assignment search space while retaining the capacity of the larger model.
  • Contribution: The approach combines recurrent world models, probabilistic modeling, and modern reinforcement-learning environments.The authors report solving a challenging pixel-based race-car navigation task not previously solved by more traditional methods.
  • Contribution: Unlike standard model-based reinforcement learning, the controller can be trained entirely inside an environment generated by its own world model and then transferred to the actual environment.Most existing approaches learn an environment model but continue training in the actual environment.
  • Robustness: Training the controller in a noisier generated environment, controlled by M’s temperature parameter, helps prevent exploitation of model imperfections.The temperature adjusts uncertainty in the generated environment.

2 Agent Model

The agent compresses visual inputs, predicts their temporal evolution probabilistically, and uses a compact controller to select actions. The components are trained with specialized methods so most model complexity remains in the world model.

  • Agent architecture: The agent consists of a visual compressor V, predictive memory M, and decision-making controller C.V compresses observations, M predicts future codes, and C chooses actions from their representations.
  • Agent architecture: A variational autoencoder V compresses each image frame into a low-dimensional latent vector z.The compressed representation is intended to capture the agent’s observed input in a smaller form.
  • Agent architecture: The recurrent model M predicts a probability density over future latent vectors because environments can be stochastic.Its output models P(z_t+1 | a_t, z_t, h_t).
  • Predictive memory: M uses a mixture of Gaussians to model the next latent vector, with temperature τ controlling uncertainty during sampling.This combination is called a Mixture Density Network with an RNN, or MDN-RNN.
  • Controller: The controller C maps the current latent code and recurrent hidden state directly to actions through a simple linear model.The policy parameters are optimized with CMA-ES across parallel environment rollouts.
  • Controller: Keeping C small leaves most parameters in V and M and enables evolutionary optimization over a comparatively small search space.The design also supports training large differentiable components with backpropagation while evolving the controller separately.

3 Car Racing Experiment: World Model for Feature Extraction

The car-racing experiment uses compressed spatial and temporal features from V and M as inputs to a compact controller. The training procedure separates unsupervised representation learning from reward-driven controller optimization.

  • Task and representation: The authors describe the agent as the first known solution to the car-racing task under the stated evaluation criterion.The task requires an average reward above 900 over 100 consecutive trials.
  • Task and representation: The experiment applies V and M features to continuous control from pixels in the randomly generated CarRacing-v0 environment.The controller steers, accelerates, and brakes while being rewarded for visiting many track tiles quickly.
  • Training procedure: The training pipeline collects 10,000 random rollouts, trains a VAE for frame encoding, trains an MDN-RNN for latent prediction, and evolves C for reward.The stages are performed in sequence rather than optimizing the full agent jointly.
  • Training procedure: V and M learn only to compress and predict image sequences, while C alone receives reward information and is optimized by evolution.The linear controller has 867 parameters, making CMA-ES suitable for its optimization.
  • Training procedure: Training V and M separately is described as more practical and sufficiently effective than end-to-end training for this experiment.The authors also report that M can train on long latent-vector sequences without images.

V without M

Using only the visual representation z_t, the controller can navigate but struggles with sharp corners and does not solve CarRacing-v0. Adding model-derived temporal information later improves this limitation.

  • V without M: The controller receives only the VAE representation z_t through a linear policy, excluding the predictive model M.Its policy is defined as a_t = W_c z_t + b_c.
  • V without M: 632 ± 251 average reward leaves the agent able to navigate but wobbling and missing sharper corners.This performance is comparable to other leaderboard agents and traditional Deep RL methods such as A3C.

4 VizDoom Experiment: Learning Inside of a Generated Environment

The authors train a controller entirely inside a VizDoom environment generated by an RNN world model, then transfer it to the actual game. Varying model uncertainty helps reduce exploitation of modeling errors, but excessive uncertainty makes learning harder.

  • World-model environment: The controller is trained in a virtual DoomTakeCover-v0 environment whose model predicts both the next latent observation and whether the agent dies.This prediction enables an OpenAI Gym interface over the learned model, replacing real pixel frames during controller training.
  • World-model environment: The world model learns to simulate essential game aspects from raw image data, including game logic, enemy behaviour, physics, and 3D rendering.The generated environment can be used as a playable substitute for the human-designed game.
  • Uncertainty control: Increasing the sampling temperature τ makes the generated environment more uncertain and difficult, discouraging policies that exploit predictable model imperfections.The temperature changes the stochasticity of generated fireball trajectories and the strategies discovered by the controller.
  • Transfer results: The visual encoder misses details such as the exact number of monsters, but controllers trained under noisier settings generally perform better in the normal setting.Increasing τ helps prevent the controller from taking advantage of imperfections in the world model.
  • Uncertainty control: At τ = 0.1, mode collapse prevents generated monsters from shooting fireballs, allowing policies trained there to exploit an unrealistically easy environment.Temperature therefore must be tuned: too little uncertainty enables cheating, while too much makes learning difficult; the best score was 1092 ± 556 at τ = 1.15, versus 820 ± 58 previously reported.

5 Related Work

Related work develops learned dynamics models and policies across control, visual prediction, and recurrent world-model settings. This paper uses evolution strategies to train a controller from learned representations, targeting high-dimensional visual environments.

  • Dynamics models: PILCO learns system dynamics with a Gaussian process and samples trajectories to train controllers for difficult low-dimensional control problems.Gaussian processes work well with small, low-dimensional datasets but are difficult to scale to long histories of high-dimensional observations.
  • Visual representations: Prior visual-control work compresses video frames with neural networks so reinforcement learning can operate on low-dimensional feature vectors.Learning dynamics in a compressed latent space can make reinforcement learning more data-efficient.
  • Sequential prediction: Earlier video-game models used feed-forward convolutional networks to predict future states and support action selection.Recurrent networks are considered better suited to capturing longer-term temporal dependencies in sequential observations.
  • Controller optimization: Evolution strategies train the controller using only final cumulative reward and can evaluate many rollout solutions in parallel.This provides an alternative to traditional deep reinforcement learning methods for optimizing compact policies.

6 Discussion

The paper demonstrates training agents entirely inside simulated latent-space worlds, while identifying limitations involving task relevance, exploration, model capacity, and planning.

  • Training controllers entirely inside generated latent environments can reduce the need to render game states or calculate irrelevant physics.The authors describe this as a practical benefit of simulated latent-space training.
  • Standalone unsupervised VAE training may encode irrelevant observations and miss task-relevant visual features.The paper reports detailed brick patterns reproduced in Doom but task-relevant road tiles missed in Car Racing.
  • More sophisticated environments require iterative training because random-policy data may not expose strategically important parts of the world.Future work proposes active exploration and intrinsic motivation to improve the world model over time.
  • The LSTM-based world model has limited capacity and can suffer catastrophic forgetting as it stores more historical information.The authors suggest higher-capacity models or external memory for more complicated worlds.
  • The system simulates futures step by step without hierarchical planning or abstract reasoning that could ignore irrelevant spatio-temporal details.The paper leaves more general recurrent controller–model approaches for future work.

A Supplementary Materials

The supplementary materials introduce the models and training methods used in the work.

  • The section provides additional details about the models used in the work.
  • The section describes the training methods used in the work.
  • The supplementary material expands on the paper’s model and training descriptions.

A.1 Comparing V, M, C Model Sizes

The section presents a parameter-count table for the CarRacing-v0 environment.

  • The section concerns comparing model sizes for CarRacing-v0.
  • Table 3 is titled “CarRacing-v0 Parameter Count.”
  • The table identifies parameter count as the comparison dimension.

A.2 Variational Autoencoder

The variational autoencoder compresses image observations into latent representations using a convolutional architecture, with tensor shapes described layer by layer.

  • Figure 3 presents the tensor shapes for the ConvVAE and the architecture of the MDN-RNN.
  • The ConvVAE is the visual component used to encode observations into compressed latent representations.It receives image observations and produces latent vectors for the world model.
  • The latent vector z is sampled from the Gaussian prior N(µ, σI).The Gaussian prior limits frame information capacity and makes the world model more robust to unrealistic latent vectors.
  • The ConvVAE applies four convolutional layers to a 64x64x3 RGB input and encodes it into vectors µ and σ of size Nz.Nz is 32 for Car Racing and 64 for Doom.
  • Each convolutional and deconvolutional layer uses stride 2, with ReLU activations except at the output layer.The output layer keeps reconstructed values between 0 and 1.

A.3 Mixture Density Network + Recurrent Neural Network

The MDN-RNN models the distribution of future latent vectors rather than making deterministic predictions, then samples those vectors to generate environments. It uses an LSTM with a Mixture Density Network output layer and is trained with teacher forcing on randomized latent inputs.

  • An LSTM recurrent network with a Mixture Density Network output layer represents the next z as a mixture of Gaussian distributions.
  • The MDN-RNN models the probability distribution of the next latent vector z instead of the next pen stroke or a deterministic prediction.At each time step, the model samples from this distribution to generate environments.
  • The MDN-RNNs were trained for 20 epochs on data collected from a random-policy agent, using 5 Gaussian mixtures without modeling the correlation ρ parameter.The Car Racing LSTM used 256 hidden units, while the Doom LSTM used 512.
  • During teacher-forced training, precomputed µ and σ values were used to sample z inputs for each batch, reducing overfitting to a particular sampled latent vector.

A.4 Controller

The controller maps world-model features to bounded actions and is optimized with CMA-ES across parallel rollouts. Evolution produced strong CarRacing and Doom results, including transfer from the generated DoomRNN environment to actual VizDoom.

  • A.5 Evolution Strategies: The controller’s actions are bounded with tanh nonlinearities, matching Car Racing’s steering, acceleration, and brake ranges.
  • A.5 Evolution Strategies: CMA-ES evolves controller weights using a population of 64 agents, each evaluated over 16 random-seed rollouts.Fitness is the average cumulative reward across those 16 rollouts.
  • A.5 Evolution Strategies: After 1800 generations, CarRacing reached an average score of 900.46 over 1024 random rollouts.The environment’s stated requirement is an average score above 900 over 100 random rollouts.
  • A.5 Evolution Strategies: A controller using only z_t scored 632 ± 251, while adding a hidden layer increased the score to 788 ± 141.
  • A.6 DoomRNN: DoomRNN was trained without the actual game engine, operating in latent space before deployment to the actual VizDoom environment.
  • A.6 DoomRNN: In DoomRNN, training at τ = 1.15 yielded 959 average score over 1024 rollouts, while the transferred controller achieved 1092 ± 556 over 100 actual DoomTakeCover-v0 rollouts.
Loading 1809.01999v1…