Dream to Control: Learning Behaviors by Latent Imagination

Danijar Hafner, Timothy Lillicrap, Jimmy Ba, Mohammad Norouzi

Seminal December 2019

TL;DR Dreamer learns a latent world model from pixels and trains an actor and a value model purely on imagined rollouts, backpropagating value gradients through the learned dynamics. On 20 DeepMind Control tasks it averages 823 after 5 million steps, above D4PG's 786 after 100 million.

What they do

Model-based agents that maximize imagined rewards over a fixed horizon tend to be shortsighted, and many rely on derivative-free planning. Dreamer learns a world model from images and then learns behavior purely inside that model’s compact latent space: an actor and a value model are trained on imagined rollouts, and the actor is updated with analytic gradients through the learned dynamics.

How it works

World model. PlaNet’s recurrent state space model (RSSM) encodes 64x64 images and actions into latent states and predicts rewards, trained by image reconstruction with a KL regularizer. Contrastive and reward-only objectives are also tested.
Imagination. Starting from latent states of replayed sequences, the model rolls forward H=15 steps with actions sampled from a tanh-Gaussian policy.
Actor critic. The value model regresses λ-returns, a weighted mix of k-step returns bootstrapped with the value model, so rewards beyond the horizon count. The actor maximizes these values by backpropagating through value, reward and transition predictions (reparameterization; straight-through gradients for discrete actions).
Interaction. Training alternates 100 gradient steps with collecting one episode; one hyperparameter set covers all continuous tasks.

Results

20 DeepMind Control tasks from pixels: average score 823 after 5×10^6 steps, versus 332 for PlaNet and 786 for D4PG after 10^8 steps.
About 3 h per 10^6 environment steps on one V100, versus 11 h for PlaNet’s online planning and the 24 h D4PG uses to reach similar performance.
At horizon 20, Dreamer beats an action model without values and PlaNet on 16 of 20 tasks, with 4 ties.
Pixel reconstruction outperforms contrastive estimation on most tasks; reward prediction alone was not sufficient.

Limitations

Everything is evaluated in simulation. Dreamer trails D4PG on some tasks (Finger Spin 499 vs 986). On Atari and DeepMind Lab it learns successful behaviors only on a subset of tasks, and the authors note that agents learning purely through world models are not yet competitive there. Results depend on the representation objective, which the authors name as the path to visually more complex environments.

Figures

From Figure 1: Dreamer learns a world model from past experience, then learns farsighted behaviors in its latent spac... — From Figure 1: Dreamer learns a world model from past experience, then learns farsighted behaviors in its latent space by backpropagating value estimates through imagined trajectories.

From Figure 3: (b) Learning behavior in imagination. From the latent state of a real image, Dreamer predicts values, ... — From Figure 3: (b) Learning behavior in imagination. From the latent state of a real image, Dreamer predicts values, rewards and actions along an imagined trajectory and propagates gradients back through it.

From Figure 4: Final performance versus imagination horizon for Dreamer, an action model without a value model, and P... — From Figure 4: Final performance versus imagination horizon for Dreamer, an action model without a value model, and PlaNet online planning. The learned value model makes Dreamer more robust to the horizon length.

From Figure 6: Scores on the 20 control tasks. After 5 million environment steps Dreamer averages 823, compared to Pl... — From Figure 6: Scores on the 20 control tasks. After 5 million environment steps Dreamer averages 823, compared to PlaNet at 332 and the model-free D4PG at 786 after 100 million steps; averages over 5 seeds.

Abstract

Learned world models summarize an agent's experience to facilitate learning complex behaviors. While learning world models from high-dimensional sensory inputs is becoming feasible through deep learning, there are many potential ways for deriving behaviors from them. We present Dreamer, a reinforcement learning agent that solves long-horizon tasks from images purely by latent imagination. We efficiently learn behaviors by propagating analytic gradients of learned state values back through trajectories imagined in the compact state space of a learned world model. On 20 challenging visual control tasks, Dreamer exceeds existing approaches in data-efficiency, computation time, and final performance.

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