To select effective actions in complex environments, intelligent agents need to generalize from past experience. World models can represent knowledge about the environment to facilitate such generalization. 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 purely by latent imagination. We efficiently learn behaviors by backpropagating analytic gradients of learned state values 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. Intelligent agents can achieve goals in complex environments even though they never encounter the exact same situation twice. This ability requires building representations of the world from past experience that enable generalization to novel situations. World models offer an explicit way to represent an agent's knowledge about the world in a parametric model learned from experience that can make predictions about the future. When the sensory inputs are high-dimensional images, latent dynamics models can abstract observations to predict forward in compact state spaces (; ;). Compared to predictions in image space, latent states have a small memory footprint and enable imagining thousands of trajectories in parallel. Learning effective latent dynamics models is becoming feasible through advances in deep learning and latent variable models (; ; ;). Behaviors can be derived from learned dynamics models in many ways. Often, imagined rewards are maximized by learning a parametric policy (; ;) or by online planning . However, considering only rewards within a fixed imagination horizon in shortsighted behaviors. Moreover, prior work commonly resorts to derivative-free optimization for robustness to model errors (; ;), rather than leveraging the analytic gradients offered by neural network dynamics models . We present Dreamer, an agent that learns long-horizon behaviors from images purely by latent imagination. A novel actor critic algorithm accounts for rewards beyond the planning horizon while making efficient use of the neural network dynamics. For this, we predict state values and actions in the learned latent space as summarized in Figure 1. The values optimize Bellman consistency for imagined rewards and the policy maximizes the values by propagating their analytic gradients back through the dynamics. In comparison to actor critic algorithms that learn online or by experience replay;;, world models enable interpolating between past experience and offer analytic gradients of multi-step returns for efficient policy optimization. Figure 2: Agent observations for 5 of the 20 control tasks used in our experiments. These pose a variety of challenges including contact dynamics, sparse rewards, many degrees of freedom, and 3D environments that exceed the difficult to tasks previously solved through world models. The agent observes the images as 64 × 64 × 3 pixel arrays. The key contributions of this paper are summarized as follows: • Learning long-horizon behaviors in imagination Purely model-based agents can be shortsighted due to finite imagination horizons. We approach this limitation in latenby predicting both actions and state values. Training purely by latent imagination lets us efficiently learn the policy by propagating analytic gradients of the value function back through latent state transitions. • Empirical performance for visual control We pair Dreamer with three representation learning objectives to evaluate it on the DeepMind Control Suite with image inputs, shown in Figure 2. Using the same hyper parameters for all tasks, Dreamer exceeds existing model-based and model-free agents in terms of data-efficiency, computation time, and final performance. Reinforcement learning We formulate visual control as a partially observable Markov decision process (POMDP) with discrete time step t ∈ [1; T], continuous vector-valued actions a t ∼ p(a t | o ≤t, a <t) generated by the agent, and high-dimensional observations and scalar rewards o t, r t ∼ p(o t, r t | o <t, a <t) generated by the unknown environment. The goal is to develop an agent that maximizes the expected sum of rewards E p T t=1 r t. Figure 2 shows a selection of our tasks. Agent components The classical components of agents that learn in imagination are dynamics learning, behavior learning, and environment interaction . In the case of Dreamer, the behavior is learned by predicting hypothetical trajectories in the compact latent space of the world model. As outlined in Figure 3 and detailed in Algorithm 1, Dreamer performs the following operations throughout the agent's life time, either sequentially interleaved or in parallel: • Learn the latent dynamics model from the dataset of past experience to predict future rewards from actions and past observations. Any learning objective for the world model can be incorporated with Dreamer. We review existing methods for learning latent dynamics in Section 4. • Learn action and value models from predicted latent trajectories, as described in Section 3. The value model optimizes Bellman consistency for imagined rewards and the action model is updated by propagating gradients of value estimates back through the neural network dynamics. • Execute the learned action model in the world to collect new experience for growing the dataset. Latent dynamics Dreamer uses a latent dynamics model that consists of three components. The representation model encodes observations and actions to create continuous vector-valued model states s t with Markovian transitions (; ;). The transition model predicts future model states without seeing the corresponding observations that will cause them. The reward model predicts the rewards given the model states, The model mimics a non-linear Kalman filter , latent state space model, or HMM with real-valued states. However, it is conditioned on actions and predicts rewards, allowing the agent to imagine the outcomes of potential action sequences without executing them in the environment. Dreamer learns long-horizon behaviors in the compact latent space of a learned world model. For this, we propagate stochastic gradients of multi-step returns through neural network predictions of actions, states, rewards, and values using reparameterization. This section describes our core contribution. The latent dynamics define a Markov decision process (MDP;) that is fully observed since the compact model states s t are Markovian. We denote imagined quantities with τ as the time index. Imagined trajectories start at the true model states s t of observation sequences drawn from the agent's past experience. They follow predictions of the transition model, and a policy a τ ∼ q(a τ | s τ). The objective is to maximize expected imagined rewards E q ∞ τ =t γ τ −t r τ with respect to the policy. Action and value models Consider imagined trajectories with a finite horizon H. Dreamer uses an actor critic approach to learn behaviors that consider rewards beyond the horizon. We learn an action model and a value model in the latent space of the world model. The action model implements the policy and aims to predict actions that solve the imagination environment. The value model estimates the state values V(s τ) E q(·|sτ) t+H τ =t γ τ −t r τ for the action model, the expected sum of imagined rewards that it achieves in each state s τ, Action model: The action and value models are trained cooperatively as typical in policy iteration: the action model aims to maximize an estimate of the value, while the value model aims to match an estimate of the value that changes as the action model changes. We use dense neural networks for the action and the value model with parameters φ and ξ, respectively. The action model outputs a tanh-transformed Gaussian with sufficient statistics predicted by the neural network. This allows for reparameterized sampling that lets sampled actions depend deterministically on the neural network output, allowing to backpropagate analytic gradients through the sampling operation, Value estimation To learn the action and value models, we need to estimate the state values of imagined trajectories {s τ, a τ, r τ} t+H τ =t. These trajectories branch off of the model states s t of sequence batches drawn from the agent's dataset of experience and predict forward for the imagination horizon H using actions sampled from the action model. State values can be estimated in multiple Figure 4: Imagination horizons. We compare the final performance of Dreamer to learning an action model without value prediction and to online planning using PlaNet. Learning a state value model to estimate rewards beyond the imagination horizon makes Dreamer more robust to the horizon length. The agents use reconstruction for representation learning and an action repeat of R = 2. ways that trade off bias and variance , where the expectations are estimated with the imagined trajectories. V R simply sums the rewards from τ until the horizon and ignores rewards beyond it. This allows learning the action model without value model, an ablation we compare to in our experiments. V k N estimates rewards beyond k steps with the learned value model. Dreamer uses V λ, which computes an exponentially-weighted average of the estimates for different k to balance bias and variance. Figure 4 shows that learning a value function in imagination enables Dreamer to solve long-horizon tasks while being robust to the imagination horizon. The experimental details and on all tasks are described in Section 5. Learning objective To update the action and value models, we first compute the value estimates V λ (s τ) for states s τ along the imagined trajectories. The objective for the action model q φ (a τ | s τ) is to output actions that in state trajectories with high value estimates. The objective for the value model v ξ (s τ), in turn, is to regress the value estimates, The value model is simply updated to regress the targets, around which we stop the gradient as typical in temporal difference learning . The action model uses analytic gradients through the learned dynamics to maximize the value estimates. To understand this, we note that the value estimates depend on the reward and value predictions, which depend on the imagined states, which in turn depend on the imagined actions. Since these steps are all implemented as neural networks with reparameterized sampling, we analytically compute . The world model is fixed while learning the action and value models. Comparison to actor critic methods Agents using Reinforce gradients employ a value baseline to reduce gradient variance, such as A3C and PPO , while Dreamer backpropagates through the value model. This is similar to analytic actor critics , such as DDPG and SAC . However, these do not leverage gradients through the state transitions and only maximize immediate Q-values. MVE and STEVE extend these to multi-step Q-learning using learned dynamics to help rewards propagate faster into the value estimates. We simply predict state values, which is sufficient for policy optimization since we backpropagate through the dynamics. We empirically compare learning action and value models from V λ, learning the action model from V R which does not require a value model, and online planning in our experiments in Learning behaviors in imagination requires a world model that generalizes well. We focus on latent dynamics models that predict forward in a compact latent space, facilitating long-term predictions and allowing to imagine thousands of trajectories in parallel. Several objectives for learning representations for control have been proposed (; ; ; . We review three approaches for learning representations to use with Dreamer: image reconstruction, contrastive estimation, and reward prediction. Reward prediction Latent imagination requires a representation model p(s t | s t−1, a t−1, o t), transition model q(s t | s t−1, a t−1,), and reward model q(r t | s t), as described in Section 2. In principle, this could be achieved by simply learning to predict future rewards given actions and past observations . Given a large and diverse dataset, such representations should be sufficient for solving a given control problem. However, while the agent is still exploring and especially when the reward signal is limited, additionally learning about observations is likely to improve the world model . Representation learning The world model is learned from sequences {(o t, a t, r t)} T t=1 drawn from the agent's dataset of experience. To learn representations that generalize, the model states s 1:T should be predictive of observations o 1:T and rewards r 1:T while not overfitting to individual examples in the dataset. At a high level, this is formalized by an information bottleneck , max The first terms encourages mutual information between the model states and the observations and rewards. The second term penalizes information between model states and dataset indices i 1:T by an amount 0 ≤ β ≤ 1. The dataset indices relate to the images by a Dirac delta p(o t | i t) as in. The information bottleneck poses the representation learning problem in a generic way and provides a common view on pixel reconstruction and contrastive estimation. While the two information terms are difficult to estimate, they are easy to bound an optimize . Reconstruction We first describe the world model used by PlaNet , shown in Figure 3a. It bounds the objective by predicting observations and rewards from the model states, where the expectation samples sequences from the dataset and states from the representation model. The bound includes a reconstruction term, a reward prediction term, and a KL regularizer. We refer to Appendix C for the derivation. The bound uses four distributions that we implement as neural networks and optimize jointly to increase the bound, Representation model: Episode Return n/a n/a n/a n/a Dreamer (2e6) PlaNet (2e6) D4PG (1e9 steps) A3C (1e9 steps, proprio) Figure 6: Performance comparison to existing methods. Dreamer exhibits the data-efficiency of PlaNet while exceeding the asymptotic performance of the best model-free agents. After 2 * 10 6 environment steps, Dreamer reaches an average performance of 802 across tasks, compared to PlaNet at 312 and the top model-free D4PG agent at 786 after 10 9 steps. Results are averages over 3 seeds. We implement the transition model as recurrent state space model (RSSM;), the representation model by combining the RSSM with a convolutional neural network (CNN;) applied to the image observation, the observation model as a transposed CNN, and the reward model as dense network. The combined parameter vector θ is updated by reparameterization gradients . Contrastive estimation Accurately predicting pixels in visually complex environments can be a challenging task. We can avoid reconstruction by instead predicting model states . While the observation marginal above was a constant, we now face the state marginal. Using the InfoNCE bound (Gutmann and Hyvärinen, 2010;) as described in Appendix C, where o q(s t | o) estimates the marginal by summing over observations o of the current sequence batch. Intuitively, q(s t | o t) makes the state predictable from the current image and ln o q(s t | o) keeps it diverse to prevent collapse. Instead of the observation model, the bound uses a state model, State model: We implement the state model as a CNN and again optimize the bound with respect to the combined parameter vector θ using reparameterization gradients. While avoiding pixel prediction, the amount of information this bound can extract efficiently is limited . We empirically compare reconstruction, contrastive, and reward objectives in our experiments in Figure 8. Visual control tasks We evaluate Dreamer on 20 continuous control tasks with image observations of the DeepMind Control Suite , illustrated in Figure 2. These tasks pose a variety of challenges, including partial observability, sparse rewards, contact dynamics, and 3D environments. We selected the tasks on which report non-zero performance from image inputs. Agent observations are images of shape 64 × 64 × 3, actions range from 1 to 12 dimensions, rewards are between 0 and 1, episodes contain 1000 steps, and initial states are randomized. Visualizations of our agent are available at https://dreamrl.github.io. Implementation All experiments used a single Nvidia V100 GPU and 10 CPU cores per training run. Our implementation uses TensorFlow Probability and will be open sourced. The training time for our implementation of Dreamer is 10 hours per 10 6 environment steps without parallelization, compared to 17 hours for online planning using PlaNet, and 24 hours for D4PG. We use the same hyper parameters across all tasks including a fixed action repeat of R = 2, as detailed in Appendix B. The world models are learned by reconstruction unless noted otherwise. Baseline methods We compare Dreamer to several baselines: The current best reported performance on the considered tasks is by D4PG , an improved variant of DDPG that uses distributed experience collection, distributional Q-learning, multi-step returns, and prioritized experience replay. PlaNet learns the same world model as Dreamer and selects actions using online planning instead of learning an action model. We include the numbers for D4PG from and re-run PlaNet with R = 2 for a fair comparison. Performance comparison To evaluate the performance of Dreamer, we compare with state-of-theart reinforcement learning agents. The are summarized in Figure 6. With an average score of 802 across tasks after 2 * 10 6 environment steps, Dreamer exceeds the performance of the strong model-free D4PG agent that achieves an average of 786 within 10 9 environment steps. At the same time, Dreamer inherits the data-efficiency of PlaNet, confirming that the learned world model can help to generalize from small amounts of experience. The empirical success of Dreamer shows that learning behaviors by latent imagination can outperform top methods based on experience replay. Long-horizon behavior To investigate its ability to learn long-horizon behaviors, we compare Dreamer to alternatives for deriving behaviors from the world model at various horizon lengths. For this, we learn an action model to maximize imagined rewards without value model and compare to online planning using PlaNet. Figure 4 shows the final performance for different imagination horizons, confirming that the value model makes Dreamer more robust to the horizon and in high performance even for short horizons. Performance curves for all 20 tasks with horizon of 20 are shown in Appendix D, where Dreamer outperforms the alternatives on 15 of 20 tasks and 3 ties. Representation learning Dreamer can be used with any dynamics model that predicts future rewards given actions and past observations. Since the representation learning objective is orthogonal to our algorithm, we compare three natural choices described in Section 4: pixel reconstruction, contrastive estimation, and pure reward prediction. Figure 8 shows clear differences in task performance for different representation learning approaches, with pixel reconstruction outperforming contrastive estimation on most tasks. This suggests that future improvements in representation learning are likely to transfer over to task performance with Dreamer. Reward prediction along was not sufficient to solve any of the tasks in our experiments. Further ablations are included in the appendix of the paper. Prior work learns latent dynamics for visual control by derivative-free policy learning or online planning, augments model-free agents with multi-step predictions, or uses analytic gradients of Qvalues or multi-step rewards, often for low-dimensional tasks. In comparison, Dreamer uses analytic gradients to efficiently learn long-horizon behaviors for visual control purely by latent imagination. Control with latent dynamics E2C and RCE embed images to predict forward in a compact space to solve simple tasks. World Models learn latent dynamics in a two-stage process to evolve linear controllers in imagination. PlaNet learns them jointly and solves visual locomotion tasks by latent online planning. Similarly, SOLAR latent space. I2A hands imagined trajectories to a model-free policy, while and learn belief representations to accelerate model-free agents. Imagined multi-step returns VPN , MVE , and STEVE learn dynamics for multi-step Q-learning from a replay buffer. AlphaGo combines predictions of actions and state values with planning, assuming access to the true dynamics. Also assuming access to the dynamics, POLO plans to explore by learning a value ensemble. PETS , VisualMPC , and PlaNet plan online using derivative-free optimization, and POPLIN improves online planning by self-imitation. Planning with neural network gradients was shown on small problems but has been challenging to scale . Analytic value gradients DPG , DDPG, and SAC leverage gradients of learned immediate action values to learn a policy by experience replay. SVG reduces the variance of model-free on-policy algorithms by analytic value gradients of one-step model predictions. ME-TRPO accelerates learning of a model-free agent via gradients of predicted rewards for proprioceptive inputs. DistGBP directly uses model gradients for online planning in simple tasks. We present Dreamer, an agent that learns long-horizon behaviors purely by latent imagination. For this, we propose a novel actor critic method that optimizes a parametric policy by propagating analytic gradients of multi-step values back through latent neural network dynamics. Dreamer outperforms previous approaches in data-efficiency, computation time, and final performance on a variety of challenging continuous control tasks from image inputs. While our approach compares favourably on these tasks, future research on learning representations is likely needed to scale latent imagination to visually more complex environments. A DETAILED ALGORITHM Update θ to predict rewards using representation learning. Imagine trajectories {(s τ, a τ)} t+H τ =t from each s t using action model. Predict rewards E q θ (r τ | s τ) and values v ξ (s τ). Compute value estimates V λ (s τ) via Equation 6. Update φ according to maximize t+H τ =t V λ (s τ) by gradient ascent. Update ξ to minimize Compute action a t ∼ q φ (a t | s t) with the action model. Add exploration noise to action. Model components We use the convolutional encoder and decoder networks from , the RSSM of , and implement all other functions as three dense layers of size 300 with ELU activations . Distributions in latent space are 30-dimensional diagonal Gaussians. The action model outputs an unconstrained mean and softplus standard deviation for the Normal distribution that is then transformed using tanh. Learning updates We draw batches of 50 sequences of length 50 to train the world model, value model, and action model models using Adam with learning rates 10 −3, 3 * 10 −4, 3 * 10 −4, respectively. We do not scale the KL regularizers (β = 1) but clip them below 3 free nats as in PlaNet. The imagination horizon is H = 20 and the same trajectories are used to update both action and value models. We use a slow moving value network that is updated every 100 gradient steps to compute the V λ value estimates with γ = 0.99 and λ = 0.95. Environment interaction The dataset is initialized with C = 5 episodes collected using random actions. We iterate between 100 training steps and collecting 1 episode by executing the predicted mode action with Normal(0, 0.3) exploration noise. Instead of manually selecting the action repeat for each environment as in and, we fix it to 2 for all environments. See Figure 11 for an assessment of the robustness to different action repeat values. The information bottleneck objective defined in Equation 9 for latent dynamics models is, For the generative objective, we lower bound the first term using the non-negativity of the KL divergence and drop the marginal data probability as it does not depend on the representation model, For the contrastive objective, we subtract the constant marginal probability of the data under the variational encoder, apply Bayes rule, and use the InfoNCE mini-batch bound , E ln q(o t | s t) + ln q(r t | s t) + = E ln q(o t | s t) − ln q(o t) + ln q(r t | s t) = E ln q(s t | o t) − ln q(s t) + ln q(r t | s t) ≥ E ln q(s t | o t) − ln o q(s t | o) + ln q(r t | s t). For the second term, we use the non-negativity of the KL divergence to obtain an upper bound, I(s 1:T ; i 1:T | a 1:T) = E p(o 1:T,r 1:T,s 1:T,a 1:T,i 1:T) t ln p(s t | s t−1, a t−1, i t) − ln p(s t | s t−1, a t−1) = E t ln p(s t | s t−1, a t−1, o t) − ln p(s t | s t−1, a t−1) ≤ E t ln p(s t | s t−1, a t−1, o t) − ln q(s t | s t−1, a t−1) = E t KL p(s t | s t−1, a t−1, o t) q(s t | s t−1, a t−1). This lower bounds the objective. Comparison of action selection schemes on the continuous control tasks of the DeepMind Control Suite from pixel inputs. The lines show mean scores over environment steps and the shaded areas show the standard deviation across 3 seeds. We compare Dreamer that learns both actions and values in imagination, to only learning actions in imagination, and to online planning using CEM without policy learning. The baselines include the top model-free algorithm D4PG, the common A3C agent, and the hybrid SLAC agent.
We present Dreamer, an agent that learns long-horizon behaviors purely by latent imagination using analytic value gradients.