Source-linked AI summary
Dream to Control: Learning Behaviors by Latent Imagination
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, Mohammad Norouzi
TL;DR
Dreamer addresses how to derive long-horizon behaviors from learned world models for visual control. It learns an actor-critic policy by propagating analytic gradients of multi-step values through imagined latent trajectories, and outperforms previous methods in data-efficiency, computation time, and final performance across challenging image-based tasks.
Problem
Learning behaviors from high-dimensional sensory inputs and learned world models remains open because many ways exist to derive behaviors from those models.
Method
Dreamer learns an actor-critic policy by propagating analytic gradients of multi-step values through learned latent dynamics and imagined trajectories.
Results
Dreamer outperforms previous methods in data-efficiency, computation time, and final performance across challenging continuous-control tasks with image inputs.
Takeaways & Limitations
Latent imagination with learned world models can learn long-horizon behaviors and outperform top experience-replay methods on visual control tasks.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Dreamer uses a learned world model to represent experience in compact latent states, then learns farsighted behaviors by imagining trajectories and backpropagating value estimates. It is evaluated on challenging visual control tasks where prior world-model methods had limitations.
- Motivation: Dreamer represents observations in compact latent spaces that support prediction and generalization to novel situations.Latent states use less memory than image-space predictions, enabling many imagined trajectories in parallel.
- Evaluation: Dreamer is evaluated on 20 visual control tasks spanning contact dynamics, sparse rewards, many degrees of freedom, and 3D environments.Several tasks had previously not been solved through world models.
- Contributions: The agent learns long-horizon behaviors by propagating analytic value gradients through trajectories imagined in latent space.This addresses shortsightedness from finite imagination horizons by predicting both actions and state values.
- Results: Using the same hyperparameters across tasks, Dreamer exceeds previous model-based and model-free agents in data-efficiency, computation time, and final performance.The evaluation uses image inputs from the DeepMind Control Suite.
2 CONTROL WITH WORLD MODELS
Dreamer combines latent dynamics learning, behavior learning, and environment interaction. Its world model predicts compact latent states and rewards, while imagined trajectories support value-based action learning before the policy acts in the environment.
- Environment interaction: The agent executes the learned action model in the environment to collect new experience and grow its dataset.This interaction cycle is interleaved with or performed in parallel to model and behavior learning.
- Dynamics learning: Dreamer learns a latent dynamics model from past experience to predict future rewards from actions and past observations.The world model can incorporate different learning objectives.
- Behavior learning: Dreamer learns action and value models from predicted latent trajectories, updating actions by propagating value gradients through neural-network dynamics.The value model enforces Bellman consistency for imagined rewards.
- Latent dynamics: The representation, transition, and reward models encode observations, predict latent states, and predict rewards in compact state space.The transition model predicts ahead without observing or imagining corresponding images.
3 LEARNING BEHAVIORS BY LATENT IMAGINATION
Dreamer learns long-horizon behaviors by imagining trajectories in a learned compact latent space and propagating value gradients through the latent dynamics. A learned value model estimates rewards beyond the imagination horizon, improving robustness to horizon length.
- Imagination environment: Dreamer starts imagined trajectories from model states and predicts states, rewards, and actions through learned latent dynamics.The imagination environment uses transition, reward, and policy models to generate hypothetical trajectories from past experience.
- Imagination horizons: Learning a state value model makes Dreamer more robust to imagination horizon length than action-only learning or online planning with PlaNet.The comparison uses pixel reconstruction for representation learning and action repeat R = 2.
- Action and value models: The action and value models are trained cooperatively: the action model maximizes estimated value, while the value model regresses toward changing value targets.Both models operate on imagined trajectories in the latent world model.
- Value estimation: The value model estimates imagined rewards beyond the finite horizon, while Vλ combines estimates across different horizons to balance bias and variance.Without value prediction, truncated returns ignore rewards beyond the imagination horizon.
- World-model predictions: Accurate 45-step latent predictions from five observed images support learning successful behaviors in a compact latent space.The predictions use only the actions after encoding the initial images from hold-out trajectories.
- Learning objective: The action model is optimized by propagating analytic gradients of value estimates backward through imagined actions, states, rewards, and values.Dreamer uses stochastic backpropagation with reparameterization for continuous actions and latent states, and straight-through gradients for discrete actions.
4 LEARNING LATENT DYNAMICS
Dreamer evaluates three representation-learning objectives for latent dynamics, using learned models to support imagination-based control. Pixel reconstruction performs best on most tasks, while contrastive learning solves about half and reward prediction alone is insufficient in these experiments.
- Dreamer reviews reward prediction, pixel reconstruction, and contrastive estimation as representation-learning objectives for latent dynamics.
- Reconstruction: The reconstruction-based world model jointly optimizes observation and reward reconstruction terms with a KL regularizer.
- Reconstruction: The reconstruction implementation combines an RSSM, CNN-based representation model, transposed CNN observation model, and dense reward model.
- Contrastive estimation: Contrastive estimation predicts latent states from images and uses other observations in the batch to prevent representation collapse.
- Comparison: Pixel reconstruction performs best for most tasks, contrastive learning solves about half, and reward prediction alone is insufficient in these experiments.
- Control implications: Dreamer succeeds on visual control tasks requiring long-horizon credit assignment, including acrobot and hopper.
5 RELATED WORK
Dreamer combines learned latent dynamics with analytic value gradients to learn long-horizon visual-control behaviors purely through latent imagination. It differs from prior approaches that use derivative-free learning, online planning, multi-step Q-learning, or gradients from shorter model predictions.
- Prior latent-dynamics methods use derivative-free policy learning, online planning, guided policy search, or other strategies for deriving behaviors.
- Dreamer uses analytic gradients through learned latent dynamics to learn long-horizon visual-control behaviors purely by latent imagination.
- Several methods learn multi-step returns from imagined dynamics for Q-learning, while AlphaGo and POLO assume access to true dynamics.
- Actor-critic methods such as DPG, DDPG, and SAC use action-value gradients, while SVG applies analytic gradients to one-step model predictions.
- Representation-learning objectives remain an orthogonal design choice whose task performance is compared in Dreamer.
6 EXPERIMENTS
Dreamer is evaluated on 20 challenging image-based control tasks spanning sparse rewards, contact dynamics, 3D scenes, and varied action settings. It achieves strong cross-task performance while reducing training time relative to PlaNet and D4PG.
- Experimental design: The experiments compare Dreamer with current methods, test long horizons and action types, examine early termination, and vary world-model learning objectives.
- Control tasks: Dreamer is evaluated on 20 DeepMind Control Suite visual tasks with sparse rewards, contact dynamics, 3D scenes, and 64 × 64 × 3 image observations.
- Computation time: Dreamer training takes about 3 hours per 10^6 environment steps, compared with 11 hours for PlaNet's online planning and 24 hours for D4PG to reach similar performance.
- Baselines: The baseline comparison includes D4PG with pixel inputs, A3C with state inputs, and PlaNet using the same learned world model with online action selection.
- Performance: Dreamer reaches an average score of 823 after 5 × 10^6 environment steps, versus D4PG's 786 within 10^8 steps.
- Performance: Dreamer inherits PlaNet's data-efficiency while exceeding the asymptotic performance of the best model-free agents.
7 CONCLUSION
Dreamer learns long-horizon behaviors purely by latent imagination and uses analytic gradients of multi-step values through learned latent dynamics. It outperforms previous methods in data-efficiency, computation time, and final performance across challenging visual-control settings.
- Dreamer learns long-horizon behaviors purely by latent imagination using an actor-critic method with analytic gradients through learned latent dynamics.
- Dreamer outperforms previous methods in data-efficiency, computation time, and final performance on challenging continuous image-based control tasks.
- Dreamer also applies to tasks with discrete actions and early episode termination.
A HYPER PARAMETERS
Dreamer uses compact neural components and a 30-dimensional diagonal-Gaussian latent space for its model distributions.
- Model components: Dreamer implements functions with three dense layers of size 300 using ELU activations.The encoder and decoder are convolutional networks, while other functions use the dense-layer configuration.
- Model components: Dreamer represents latent-space distributions as 30-dimensional diagonal Gaussians.
- Model components: The action model outputs a tanh-scaled mean and a softplus standard deviation for a Normal distribution transformed by tanh.The tanh mean is scaled by a factor of 5.
B DERIVATIONS
The information bottleneck objective encourages latent states to predict observations and rewards while limiting extracted information. Its variational formulation yields a tractable lower-bounded objective using reconstruction terms and contrastive estimation.
- Information bottleneck: The information bottleneck objective maximizes predictive information about observations and rewards while penalizing information extracted from dataset indices.The penalty is weighted by the scalar β.
- Information bottleneck: Maximizing the objective encourages reconstructing each image from preceding information whenever possible and accessing additional current information only as needed.
- Variational formulation: The variational objective includes observation and reward log-likelihood terms conditioned on latent states.
- Contrastive objective: The contrastive objective uses Bayes rule and the InfoNCE mini-batch bound after subtracting the data's constant marginal probability.
- Variational formulation: Non-negativity of the KL divergence provides an upper bound, producing a lower bound on the objective.
C DISCRETE CONTROL
Dreamer is evaluated on visually complex discrete-action tasks from Atari and DeepMind Lab, where successful behavior is demonstrated on a subset but broader competitiveness remains future work.
- Evaluation setting: Dreamer is evaluated on Atari and DeepMind Lab tasks with discrete actions, visual complexity, sparse rewards, and early termination.Agents observe 64 × 64 × 3 images and choose between 3 and 18 actions.
- Evaluation setting: The evaluation uses sticky actions for Atari and refers to Figure 9 for the experiments.
- Observed performance: Dreamer learns successful behaviors on a subset of Atari games and the object collection level of DeepMind Lab.
- Future direction: The paper identifies representation learning as future work toward competitive performance across all Atari games and DeepMind Lab levels.
D BEHAVIOR LEARNING
The experiments compare Dreamer's action-selection and representation-learning choices, test control-frequency robustness, and report aggregate scores under differing input modalities and environment-step budgets.
- Action selection: Dreamer is compared with action-only imagination, PlaNet's online planning, and the model-free algorithm D4PG on continuous pixel-control tasks.The figure reports mean scores and standard deviations across 5 seeds.
- Representation learning: Image reconstruction provides the best representation-learning signal across most tasks, followed by the contrastive objective.Learning purely from rewards was not sufficient in these experiments.
- Control frequency: R = 2 works best across tasks when Dreamer is trained with different action-repeat values.The experiment uses a previous hyperparameter setting and reports one standard deviation across 2 seeds.
- Aggregate scores: The reported averages are 243.70, 786.32, 332.97, and 823.39 across the listed input-modality and environment-step conditions.
- Evaluation protocol: PlaNet is rerun with fixed action repeat R = 2 rather than tuning this value separately for each of the 20 tasks.The paper notes that these scores therefore differ from those reported by Hafner et al. (2018).