Dream to Control: Learning Behaviors by Latent Imagination

1 INTRODUCTION 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 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 ( Watter et al. , 2015 ; Oh et al. , 2017 ; Gregor et al. , 2019 ) . Compared to predictions in image space , latent states have a small memory footprint that enables imagining thousands of trajectories in parallel . Learning effective latent dynamics models is becoming feasible through advances in deep learning and latent variable models ( Krishnan et al. , 2015 ; Karl et al. , 2016 ; Doerr et al. , 2018 ; Buesing et al. , 2018 ) . Behaviors can be derived from dynamics models in many ways . Often , imagined rewards are maximized with a parametric policy ( Sutton , 1991 ; Ha and Schmidhuber , 2018 ; Zhang et al. , 2019 ) or by online planning ( Chua et al. , 2018 ; Hafner et al. , 2018 ) . However , considering only rewards within a fixed imagination horizon results in shortsighted behaviors ( Wang et al. , 2019 ) . Moreover , prior work commonly resorts to derivative-free optimization for robustness to model errors ( Ebert et al. , 2017 ; Chua et al. , 2018 ; Parmas et al. , 2019 ) , rather than leveraging analytic gradients offered by neural network dynamics ( Henaff et al. , 2019 ; Srinivas et al. , 2018 ) . 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 imagination 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 ( Lillicrap et al. , 2015 ; Mnih et al. , 2016 ; Schulman et al. , 2017 ; Haarnoja et al. , 2018 ; Lee et al. , 2019 ) , world models can interpolate past experience and offer analytic gradients of multi-step returns for efficient policy optimization . ∗Correspondence to : Danijar Hafner < mail @ danijar.com > . The key contributions of this paper are summarized as follows : • Learning long-horizon behaviors by latent imagination Model-based agents can be short- sighted if they use a finite imagination horizon . We approach this limitation by predicting both actions and state values . Training purely by imagination in a latent space lets us efficiently learn the policy by propagating analytic value gradients back through the latent dynamics . • Empirical performance for visual control We pair Dreamer with existing representation learning methods and evaluate it on the DeepMind Control Suite with image inputs , illustrated in Figure 2 . Using the same hyper parameters for all tasks , Dreamer exceeds previous model-based and model-free agents in terms of data-efficiency , computation time , and final performance . 2 CONTROL WITH WORLD MODELS . 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 at ∼ p ( at | o≤t , a < t ) generated by the agent , and high-dimensional observations and scalar rewards ot , rt ∼ p ( ot , rt | o < t , a < t ) generated by the unknown environment . The goal is to develop an agent that maximizes the expected sum of rewards Ep ( ∑T t=1 rt ) . 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 ( Sutton , 1991 ) . 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 interleaved or in parallel : • Learning the latent dynamics model from the dataset of past experience to predict future re- wards 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 . • Learning 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 . • Executing 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 st with Markovian transitions ( Watter et al. , 2015 ; Zhang et al. , 2019 ; Hafner et al. , 2018 ) . The transition model predicts future model states without seeing the corresponding observations that will later cause them . The reward model predicts the rewards given the model states , Representation model : p ( st | st−1 , at−1 , ot ) Transition model : q ( st | st−1 , at−1 ) Reward model : q ( rt | st ) . ( 1 ) We use p for distributions that generate samples in the real environment and q for their approximations that enable latent imagination . Specifically , the transition model lets us predict ahead in the compact latent space without having to observe or imagine the corresponding images . This results in a low memory footprint and fast predictions of thousands of imagined trajectories in parallel . The model mimics a non-linear Kalman filter ( Kalman , 1960 ) , 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 . observations and actions into compact latent states ( ) , for example via reconstruction , and predicts environment rewards ( ) . ( b ) In the compact latent space , Dreamer predicts state values ( ) and actions ( ) that maximize future value predictions by propagating gradients back through imagined trajectories . ( c ) The agent encodes the history of the episode to compute the current model state and predict the next action to execute in the environment . See Algorithm 1 for pseudo code of the agent . 3 LEARNING BEHAVIORS BY LATENT IMAGINATION . Dreamer learns long-horizon behaviors in the compact latent space of a learned world model by efficiently leveraging the neural network latent dynamics . 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 the main contribution of our paper . Imagination environment The latent dynamics define a Markov decision process ( MDP ; Sutton , 1991 ) that is fully observed because the compact model states st are Markovian . We denote imagined quantities with τ as the time index . Imagined trajectories start at the true model states st of observation sequences drawn from the agent ’ s past experience . They follow predictions of the transition model sτ ∼ q ( sτ | sτ−1 , aτ−1 ) , reward model rτ ∼ q ( rτ | sτ ) , and a policy aτ ∼ q ( aτ | sτ ) . The objective is to maximize expected imagined rewards Eq ( ∑∞ τ=t γ τ−trτ ) with respect to the policy . Algorithm 1 : Dreamer Initialize dataset D with S random seed episodes . Initialize neural network parameters θ , φ , ψ randomly . while not converged do for update step c = 1 .. C do // Dynamics learning Draw B data sequences { ( at , ot , rt ) } k+Lt=k ∼ D. Compute model states st ∼ pθ ( st | st−1 , at−1 , ot ) . Update θ using representation learning . // Behavior learning Imagine trajectories { ( sτ , aτ ) } t+Hτ=t from each st . Predict rewards E ( qθ ( rτ | sτ ) ) and values vψ ( sτ ) . Compute value estimates Vλ ( sτ ) via Equation 6 . Update φ← φ+ α∇φ ∑t+H τ=t Vλ ( sτ ) . Update ψ← ψ− α∇ψ ∑t+H τ=t 1 2 ∥∥vψ ( sτ ) 9Vλ ( sτ ) ∥∥2 . // Environment interaction o1 ← env.reset ( ) for time step t = 1 .. T do Compute st ∼ pθ ( st | st−1 , at−1 , ot ) from history . Compute at ∼ qφ ( at | st ) with the action model . Add exploration noise to action . rt , ot+1 ← env.step ( at ) . Add experience to dataset D ← D ∪ { ( ot , at , rt ) Tt=1 } . Model components Representation pθ ( st | st-1 , at-1 , ot ) Transition qθ ( st | st-1 , at-1 ) Reward qθ ( rt | st ) Action qφ ( at | st ) Value vψ ( st ) Hyper parameters Seed episodes S Collect interval C Batch size B Sequence length L Imagination horizon H Learning rate α 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 for this . The action model implements the policy and aims to predict actions that solve the imagination environment . The value model estimates the expected imagined rewards that the action model achieves from each state sτ , Action model : aτ ∼ qφ ( aτ | sτ ) Value model : vψ ( sτ ) ≈ Eq ( ·|sτ ) ( ∑t+H τ=t γ τ−trτ ) . ( 2 ) 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 value models with parameters φ and ψ , respectively . The action model outputs a tanh-transformed Gaussian ( Haarnoja et al. , 2018 ) with sufficient statistics predicted by the neural network . This allows for reparameterized sampling ( Kingma and Welling , 2013 ; Rezende et al. , 2014 ) that views sampled actions as deterministically dependent on the neural network output , allowing us to backpropagate analytic gradients through the sampling operation , aτ = tanh ( µφ ( sτ ) + σφ ( sτ ) ) , ∼ Normal ( 0 , I ) . ( 3 ) 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 st 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 ways that trade off bias and variance ( Sutton and Barto , 2018 ) , VR ( sτ ) . = Eqθ , qφ ( t+H∑ n=τ rn ) , ( 4 ) VkN ( sτ ) . = Eqθ , qφ ( h−1∑ n=τ γn−τrn + γ h−τvψ ( sh ) ) with h = min ( τ + k , t+H ) , ( 5 ) Vλ ( sτ ) . = ( 1− λ ) H−1∑ n=1 λn−1VnN ( sτ ) + λ H−1VHN ( sτ ) , ( 6 ) where the expectations are estimated under the imagined trajectories . VR simply sums the rewards from τ until the horizon and ignores rewards beyond it . This allows learning the action model without a value model , an ablation we compare to in our experiments . VkN estimates rewards beyond k steps with the learned value model . Dreamer uses Vλ , an exponentially-weighted average of the estimates for different k to balance bias and variance . Figure 4 shows that learning a value model in imagination enables Dreamer to solve long-horizon tasks while being robust to the imagination horizon . The experimental details and results on all tasks are described in Section 6 . Learning objective To update the action and value models , we first compute the value estimates Vλ ( sτ ) for all states sτ along the imagined trajectories . The objective for the action model qφ ( aτ | sτ ) is to predict actions that result in state trajectories with high value estimates . The objective for the value model vψ ( sτ ) , in turn , is to regress the value estimates , max φ Eqθ , qφ ( t+H∑ τ=t Vλ ( sτ ) ) , ( 7 ) min ψ Eqθ , qφ ( t+H∑ τ=t 1 2 ∥∥∥vψ ( sτ ) −Vλ ( sτ ) ) ∥∥∥2 ) . ( 8 ) The value model is updated to regress the targets , around which we stop the gradient as typical ( Sutton and Barto , 2018 ) . 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 all steps are implemented as neural networks , we analytically compute ∇φEqθ , qφ ( ∑t+H τ=t Vλ ( sτ ) ) by stochastic backpropagation ( Kingma and Welling , 2013 ; Rezende et al. , 2014 ) . We use reparameterization for continuous actions and latent states and straight-through gradients ( Bengio et al. , 2013 ) for discrete actions . The world model is fixed while learning behaviors . In tasks with early termination , the world model also predicts the discount factor from each latent state to weigh the time steps in Equations 7 and 8 by the cumulative product of the predicted discount factors , so terms are weighted down based on how likely the imagined trajectory would have ended . Comparison to actor critic methods Agents using Reinforce gradients ( Williams , 1992 ) , such as A3C and PPO ( Mnih et al. , 2016 ; Schulman et al. , 2017 ) , employ value baselines to reduce gradient variance , while Dreamer backpropagates through the value model . This is similar to deterministic or reparameterized actor critics ( Silver et al. , 2014 ) , such as DDPG and SAC ( Lillicrap et al. , 2015 ; Haarnoja et al. , 2018 ) . However , these do not leverage gradients through transitions and only maximize immediate Q-values . MVE and STEVE ( Feinberg et al. , 2018 ; Buckman et al. , 2018 ) extend them to multi-step Q-learning with learned dynamics to provide more accurate Q-value targets . We predict state values , which is sufficient for policy optimization since we backpropagate through the dynamics . Refer to Section 5 for a more detailed comparison to related work .

Summary: This paper introduced a latent space model for reinforcement learning in vision-based control tasks. It first learns a latent dynamics model, in which the transition model and the reward model can be learned on the latent state representations. Using the learned latent state representations, it used an actor-critic model to learn a reactive policy to optimize the agent's behaviors in long-horizon continuous control tasks. The method is applied to vision-based continuous control in 20 tasks in the Deepmind control suite.