Dream and Search to Control: Latent Space Planning for Continuous Control

1 INTRODUCTION . Deep reinforcement learning ( RL ) has been effective at solving sequential decision-making problems with varying levels of difficulty . The solutions generally fall into one of two categories : model-free and model-based methods . Model-free methods ( Haarnoja et al. , 2018 ; Silver et al. , 2014 ; Lillicrap et al. , 2015 ; Fujimoto et al. , 2018 ; Schulman et al. , 2017 ) directly learn a policy or action-values . However , they are usually considered to be sample-inefficient . Model-based methods ( Lee et al. , 2019 ; Gregor et al. , 2019 ; Zhang et al. , 2018 ; Ha & Schmidhuber , 2018 ) learn the environment ’ s dynamics . Common model-based approaches involve sampling trajectories from the learned dynamics to train using RL or applying a planning algorithm directly on the learned dynamics ( Ha & Schmidhuber , 2018 ; Hafner et al. , 2019a ; b ) . However , learning multi-step dynamics in the raw observation space is challenging . This is primarily because it involves the reconstruction of high-dimensional features ( e.g. , pixels ) in order to roll-out trajectories , which is an error-prone process . Instead , recent work has focused on learning latent space models ( Hafner et al. , 2019a ; b ) . This has been shown to improve robustness and sample efficiency by eliminating the need for high-dimensional reconstruction during inference . Learning on latent dynamics : Once the dynamics have been learned , a classic approach is to sample trajectories from it to learn a policy using RL . This approach is usually motivated by sample efficiency . Dreamer ( Hafner et al. , 2019a ) took this approach and demonstrated state-of-the-art performance on continuous control by performing gradient-based RL on learned latent dynamics . Another approach is to perform a look-ahead search - where the dynamics are used for a multi-step rollout to determine an optimal action . This could be accompanied by a value estimate and/or a policy that produces state-action mappings to narrow the search space or reduce the search ’ s depth . MuZero ( Schrittwieser et al. , 2019 ) took this approach and applied tree-based search on latent dynamics - however , it was restricted to discrete action spaces only . The role of look-ahead search using learned latent dynamics has not been explored sufficiently for continuous action spaces . Our contribution : In this paper , we extend the idea of performing look-ahead search using learned latent dynamics to continuous action spaces . Our high level approach is shown in Fig 1 . We build on top of Dreamer and modify how actions are sampled during online planning . Instead of sampling actions from the current policy , we search over a set of actions sampled from a mix of distributions . For our search mechanism , we implement MCTS but also investigate a simple rollout algorithm that trades off performance for compute complexity . We observed that look-ahead search results in optimal actions early on . This , in turn , leads to faster convergence of model estimates and of optimal policies . However , these benefits come at the cost of computational time since deeper and iterative look-ahead search takes more time . 2 RELATED WORK . Model-Based RL involves learning a dynamical model and utilizing it to sample trajectories to train a policy or performing look-ahead search for policy improvement . World Model ( Ha & Schmidhuber , 2018 ) learns a latent representation of observations and learns a dynamical model over this latent representation . It then utilizes the learned model to optimize a linear controller via evolution . The latent representations help to plan in a low dimensional abstract space . PlaNet ( Hafner et al. , 2019b ) does end-to-end training of the latent observation encoding and dynamics and does online planning using CEM . Dreamer , the current state-of-the-art , ( Hafner et al. , 2019a ) builds on top of PlaNet and uses analytic gradients to efficiently learn long-horizon behaviors for visual control purely by latent imagination . We built our work on top of Dreamer . Planning using look-ahead search : Given a dynamics model , one can search over its state-action space to determine an optimal action . If one has a perfect model but a sub-optimal action-value function , then in fact , a deeper look-ahead search will usually yield better policies . This look-ahead search could be done both during online decision making and during policy improvement . AlphaGo ( Silver et al. , 2017 ; Lowrey et al. , 2018 ) combined MCTS with a value estimator to plan and explore with known ground-truth dynamics . MuZero ( Schrittwieser et al. , 2019 ) learned a latent dynamic model and performed MCTS over it . They conditioned the latent representation on value and multistep reward prediction rather than observation reconstruction . These works are limited to discrete action space due to large action space in continuous action space tasks . A variation on MCTS is Progressive Widening ( Coulom , 2006 ; Chaslot et al. , 2008 ; Couëtoux et al. , 2011 ) where the the child-actions of a node are expanded based on its visitation count . This has been exploited in continuous action spaces by adding sampled actions from a proposal distribution to create a discrete set of child-actions . AOC ( Moerland et al. , 2018 ) utilized this approach to apply MCTS in continuous action spaces over true dynamics.However , they showed their results only for the Pendulum-v0 task in OpenAI Gym . A common thread across all these prior works is the availability of known ground-truth dynamics . Another approach that can be used to reduce the size of the look-ahead search space is to utilize hierarchical optimistic optimization ( HOO ) , which splits the action space and gradually increases the resolution of actions ( Mansley et al. , 2011 ; Bubeck et al. , 2009 ) . 3 DREAMER : LEARNING LATENT DYNAMICS AND POLICY . MBRL approaches usually involve collecting data from the environment , dynamics learning , and behavior learning by utilizing data collected from environment-agent interaction and/or trajectories generated via learned dynamics . We adopt Dreamer ’ s dynamics and behavior learning approach . And , unlike Dreamer , we perform decision-time planning via look-ahead search during exploration to determine an action , instead of just using learned behavior policy . In this section , we review Dreamer ’ s dynamics learning component in subsection ( 3.1 ) and the behavior learning component in subsection ( 3.2 ) . After that , we describe our investigation of look-ahead search for decision-time planning in section ( 4 ) . 3.1 DYNAMICS LEARNING . This comprises of a representation model that encodes the history of observations and actions to create a Markovian belief of the current state . This belief state and an action are consumed by a transition model that predicts the model ’ s future latent state . The transition model is conditioned on an observation model and a reward model . The observation model predicts the observation corresponding to the current latent state and the reward model predicts the corresponding reward . The observation model is used only to condition the transition model and does not play a role during inference . These models are summarized n Eq . ( 1 ) where ( st , at , ot , rt ) represent latent state , action , observation , and reward at discrete-time step t. Representation model : pθ ( st|st−1 , at−1 , ot ) Observation model : qθ ( ot|st ) Reward model : qθ ( rt|st ) Transition model : qθ ( st|st−1 , at−1 ) . ( 1 ) The transition model is implemented as a recurrent state-space model ( RSSM ) ( Hafner et al. , 2019b ) . The representation , observation model and reward model are dense networks . Observation features are given to the representation model , which are reconstructed by the observation model . This is unlike Dreamer which trains the model over raw pixels . The combined parameter vector θ in Eq . 1 is updated by stochastic backpropagation ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) . It optimizes θ using a joint reconstruction loss ( JREC ) comprising of an observation reconstruction loss ( JO ) , a reward reconstruction loss ( JR ) and a KL regularizer ( JD ) , as described in Eq . ( 2 ) . JREC . = Ep ( ∑ t ( J tO + J tR + J tD ) ) + const J tO . = ln q ( ot|st ) J tR . = ln q ( rt|st ) J tD . = −βKL ( p ( st|st−1 , at−1 , ot ) || ( st|st−1 , at−1 ) ) . ( 2 ) 3.2 BEHAVIOR LEARNING . Given a transition dynamics , the action and value models are learnt from the imagined trajectories { sτ , aτ , rτ } t+Hτ=t of finite-horizon ( H ) over latent state space . Dreamer utilizes these trajectories for learning behaviour via an actor-critic approach . 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τ , where τ is the discrete-time index during imagination . Action model : aτ ∼ qφ ( aτ |sτ ) Value model : vψ ( sτ ) ≈ Eq ( ·|sτ ) ∑t+H τ=t γ τ−trτ . ( 3 ) As in Eq . ( 3 ) , action and value models are implemented as dense neural network with parameters φ and ψ , respectively . As shown in Eq . 4 , the action model outputs a tanh-transformed Gaussian ( Haarnoja et al. , 2018 ) . This allows for re-parameterized sampling ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) that views sampled actions as deterministically dependent on the neural network output , allowing Dreamer to backpropagate analytic gradients through the sampling operation . aτ = tanh ( µφ ( sτ ) + σφ ( sτ ) ) , ∼ Normal ( 0 , I ) . ( 4 ) Dreamer uses Vλ ( Eq . 8 ) ( Sutton & Barto , 2018 ) , an exponentially weighted average of the estimates for different k-step returns for the imagination rollouts . The k varies from 1 . . . H . It helps to balance bias and variance of target for value estimates by acting as an intermediary between the 1-step and the Monte-Carlo return . The objective for the action model qφ ( aτ |sτ ) is to predict actions that result in state trajectories with high value estimates as shown in Eq . 5 . It uses analytic gradients through the learned dynamics to maximize the value estimates . The objective for the value model vψ ( sτ ) , in turn , is to regress the value estimates as shown in Eq . ( 6 ) . Also , Dreamer freezes the world model while learning behaviors . max φ Eqθ , qφ t+H∑ τ=t Vλ ( sτ ) , ( 5 ) min ψ Eqθ , qφ t+H∑ τ=t 1 2 ∥∥∥vψ ( sτ ) −Vλ ( sτ ) ) ∥∥∥2 . ( 6 ) VkN ( sτ ) . = Eqθ , qφ h−1∑ n=τ γn−τrn + γ h−τvψ ( sh ) with h = min ( τ + k , t+H ) , ( 7 ) Vλ ( sτ ) . = ( 1− λ ) H−1∑ n=1 λn−1VnN ( sτ ) + λ H−1VHN ( sτ ) , ( 8 )

Summary: This paper proposes to integrate planning into Dreamer. The main idea is to apply a planning module on top of Dreamer to improve the quality of action selection. The planning via MCTS is on the learnt latent dynamics and the policy learnt by Dreamer. One of the challenges addressed in the paper is to perform planning on continuous action spaces. The proposed method is evaluated on 20 control tasks from the DeepMind Control suite, and compared against the original Dreamer algorithm, and a baseline planning method that does only rollout simulations.