[Paper Review] World Models

David Ha, Jürgen Schmidhuber|arXiv (Cornell University)|Mar 27, 2018

Reinforcement Learning in Robotics40 references179 citations

TL;DR

This paper presents a large RNN-based world model (VAE + MDN-RNN) that compresses observations and predicts future latent states, enabling a small controller to perform tasks, including training inside a learned dream and transferring policies to the real environment.

ABSTRACT

We explore building generative neural network models of popular reinforcement learning environments. Our <em>world model</em> 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 article is available at worldmodels.github.io.

Motivation & Objective

• Motivate and construct a predictive world model inspired by human cognitive prediction mechanisms.
• Demonstrate how a large, unsupervised world model can compress and predict spatial-temporal representations from raw observations.
• Show that a tiny controller can solve tasks when guided by features from the world model.
• Explore training policies inside a hallucinated dream environment and transfer them back to the real environment.
• Discuss iterative training and potential curiosity-driven enhancements for more complex environments.

Proposed method

• Use a Variational Autoencoder (VAE) to compress each frame into a latent vector z.
• Use a Mixture Density Network (MDN) with an RNN to model the distribution P(z_{t+1} | a_t, z_t, h_t) as a mixture of Gaussians.
• Implement a minimal linear Controller C that maps [z_t, h_t] to actions, keeping C lightweight for evolution strategies (CMA-ES) optimization.
• Train V and MDN-RNN separately on data collected from random rollouts, then optimize C using CMA-ES.
• Optionally run the controller in a dream-like world generated by M, or transfer policies learned in dream to the real environment.
• Discuss iterative training, curiosity mechanisms, and robustness against model imperfectness.

Experimental results

Research questions

• RQ1 Can a large unsupervised world model learn compact spatial-temporal representations from high-dimensional observations?
• RQ2 To what extent can a simple controller leverage world-model features to solve control tasks?
• RQ3 Is it possible to train policies inside a learned dream and transfer them to the actual environment?
• RQ4 How does incorporating uncertainty via a MDN-RNN affect policy learning and transfer?
• RQ5 What are the vulnerabilities of model-based policies when the dynamics model is imperfect, and how can these be mitigated?

Key findings

• V model + MDN-RNN enable a compact controller to achieve strong performance on CarRacing-v0, surpassing several prior methods.
• Using the full world model (V + M) yields higher average scores than using V alone or a simple controller.
• On CarRacing-v0, the Full World Model achieves 906 ± 21, outperforming DQN, A3C variants, and prior leaderboard results.
• Training inside a dream environment (VizDoom Take Cover) can produce policies that transfer to the real VizDoom environment, with notable survival time gains.
• Adjusting the MDN-RNN temperature τ modulates dream difficulty and transferability, balancing realism and exploitability.
• The paper discusses potential adversarial exploitation of the world model and the need for uncertainty to mitigate such issues.