Learning to Plan Optimistically: Uncertainty-Guided Deep Exploration via Latent Model Ensembles

1 INTRODUCTION . The ability to learn complex behaviors through interaction will enable the autonomous deployment of various robotic systems in the real world . Reinforcement learning ( RL ) provides a key framework for realizing these capabilities , but efficiency of the learning process remains a prevalent concern . Reallife applications yield complex planning problems due to high-dimensional environment states , which are further exacerbated by the agent ’ s continuous actions space . For RL to enable real-world autonomy , it therefore becomes crucial to determine efficient representations of the underlying planning problem , while formulating interaction strategies capable of exploring the resulting representation efficiently . In traditional controls , planning problems are commonly formulated based on the underlying statespace representation . This may inhibit efficient learning when the environment states are highdimensional or their dynamics are susceptible to non-smooth events such as singularities and discontinuities ( Schrittwieser et al. , 2019 ; Hwangbo et al. , 2019 ; Yang et al. , 2019 ) . It may then be desirable for the agent to abstract a latent representation that facilitates efficient learning ( Ha & Schmidhuber , 2018 ; Zhang et al. , 2019 ; Lee et al. , 2019 ) . The latent representation may then be leveraged either in a model-free or model-based setting . Model-free techniques estimate state-values directly from observed data to distill a policy mapping . Model-based techniques learn an explicit representation of the environment that is leveraged in generating fictitious interactions and enable policy learning in imagination ( Hafner et al. , 2019a ) . While the former reduces potential sources of bias , the latter offers a structured representation encoding deeper insights into underlying environment behavior . The agent should leverage the chosen representation to efficiently identify and explore informative interactions . We provide a motivational one-dimensional example of a potential action-value mapping in Figure 1 ( left ) . The true function and its samples are visualized in red with the true maximum denoted by the green dot . Relying only on the predicted mean can bias policy learning towards local optima ( orange dot ; Sutton & Barto ( 2018 ) ) , while added stochasticity can waste samples on un-informative interactions . Auxiliary information-gain objectives integrate predicted uncertainty , however , uncertain environment behavior does not equate to potential for improvement ( pink dot ) . It is desirable to focus exploration on interactions that harbor potential for improving overall performance . Combining mean performance estimates with the associated uncertainty into an upper confidence bound ( UCB ) objective provides a concise method of doing so ( blue dot ; Auer et al . ( 2002 ) ; Krause & Ong ( 2011 ) ) . The underlying uncertainty can be explicitly represented by maintaining an ensemble of hypothesis on environment behavior ( Osband et al. , 2016 ; Lakshminarayanan et al. , 2017 ) . Figure 1 ( right ) demonstrates this selective uncertainty reduction by showcasing forward predictions of a ensembled model on two motion patterns of a walker agent . The expected high-reward walking behavior has been sufficiently explored and model hypotheses strongly agree ( top ) , while little effort has been extended to reduce uncertainty over the expected low-reward falling behavior ( bottom ) . This paper demonstrates that exploring interactions through imagined positive futures can yield information-dense sampling and data-efficient learning . We present latent optimistic value exploration ( LOVE ) , an algorithm that leverages optimism in the face of uncertain long-term rewards in guiding exploration . Potential futures are imagined by an ensemble of latent variable models and their predicted infinite-horizon performance is obtained in combination with associated value function estimates . Training on a UCB objective over imagined futures yields a policy that behaves inherently optimistic and focuses on interactions with the potential to improve performance under the current world model . This provides a concise , differentiable framework for driving deep exploration while not relying on stochasticity . We present empirical results on challenging visual control tasks that highlight the necessity for deep exploration in scenarios with sparse reward signals , and demonstrate improved sample-efficiency on a selection of benchmarking environments from the DeepMind Control Suite ( Tassa et al. , 2018 ) . We compare to both Dreamer ( Hafner et al. , 2019a ) , the current state-of-the-art model-based agent , and DrQ ( Kostrikov et al. , 2020 ) , a concurrent model-free approach . 2 RELATED WORK . Problem representation Model-free approaches learn a policy by directly estimating performance from interaction data . While their asymptotic performance previously came at the cost of sample complexity ( Lillicrap et al. , 2015 ; Fujimoto et al. , 2018 ; Haarnoja et al. , 2018 ) , recent advances in representation learning through contrastive methods and data augmentation have improved their efficiency ( Srinivas et al. , 2020 ; Laskin et al. , 2020 ; Kostrikov et al. , 2020 ) . However , their implicit representation of the world can make generalization of learned behaviors under changing task specifications difficult . Model-based techniques leverage a structured representation of their environment that enables them to imagine potential interactions . The nature of the problem hereby dictates model complexity , ranging from linear ( Levine & Abbeel , 2014 ; Kumar et al. , 2016 ) , over Gaussian process models ( Deisenroth & Rasmussen , 2011 ; Kamthe & Deisenroth , 2018 ) to deep neural networks ( Chua et al. , 2018 ; Clavera et al. , 2018 ) . In high-dimensional environments , latent variable models can provide concise representations that improve efficiency of the learning process ( Watter et al. , 2015 ; Ha & Schmidhuber , 2018 ; Lee et al. , 2019 ; Hafner et al. , 2019a ) . Planning interactions Model-based approaches leverage their representation of the world in predicting the performance of action sequences . The agent may then either solve a model predictive control-style optimization ( Nagabandi et al. , 2018 ; Chua et al. , 2018 ; Hafner et al. , 2019b ) or train a policy in simulation ( Kurutach et al. , 2018 ; Clavera et al. , 2018 ) . The resulting finite-horizon formulations can be extended by value function estimates to approximate an infinite-horizon planning problem ( Lowrey et al. , 2019 ; Hafner et al. , 2019a ; Seyde et al. , 2020 ) . When considering learned models , ensembling the model predictions may further be leveraged in debiasing the actor strategy during training ( Kurutach et al. , 2018 ; Chua et al. , 2018 ; Clavera et al. , 2018 ; Seyde et al. , 2020 ) . Both explicit and implicit model rollouts in combination with value estimates can also be utilized to accelerate model-free learning ( Oh et al. , 2017 ; Feinberg et al. , 2018 ; Buckman et al. , 2018 ) . Directed exploration Directed exploration can improve over random exploration by focusing on information-dense interactions ( Schmidhuber , 2010 ) . These methods are commonly driven by uncertainty estimates . Information gain techniques define an auxiliary objective that encourages exploration of unexpected environment behavior or model disagreement and have been applied in discrete ( Stadie et al. , 2015 ; Ostrovski et al. , 2017 ; Pathak et al. , 2017 ) and continuous actions spaces ( Still & Precup , 2012 ; Houthooft et al. , 2016 ; Henaff , 2019 ) . When driving interactions to improve knowledge of the environment dynamics , effort may be wasted on uncertain interactions that are tangential to the objective . Alternatively , exploration can be driven by uncertainty over expected performance as encoded by value functions ( Osband et al. , 2016 ; 2017 ; Chen et al. , 2017 ; O ’ Donoghue et al. , 2018 ; Lee et al. , 2020 ) , multi-step imagined returns ( Depeweg et al. , 2018 ; Henaff et al. , 2019 ) or their combination ( Lowrey et al. , 2019 ; Schrittwieser et al. , 2019 ; Seyde et al. , 2020 ) . Model-ensemble agents Related work on ensemble agents has demonstrated impressive results . We note key differences to our approach . ME-TRPO ( Kurutach et al. , 2018 ) leverages a dynamics ensemble to debias policy optimization on finite-horizon returns under a known reward function and random exploration . MAX ( Shyam et al. , 2019 ) and Amos et al . ( 2018 ) explore via finite-horizon uncertainty in a state and latent dynamics ensemble , respectively . RP1 ( Ball et al. , 2020 ) explores in reward space via finite-horizon returns , but assumes access to the nominal reward function and full proprioceptive feedback . Seyde et al . ( 2020 ) also leverage full proprioception and embed optimism into the value function , which prohibits adjustment of the exploration trade-off during policy learning and limits transferability . Exploring uncertain dynamics samples interactions orthogonal to task completion and finite-horizon objectives limit exploration locally , while full-observability and access to the reward function are strong assumptions . We learn latent dynamics , reward and value functions in partially observable settings to explore uncertainty over infinite-horizon returns . This enables backpropagation through imagined trajectories to recover analytic policy gradients , while offering a flexible framework to guide exploration based on expected potential for long-term improvement . 3 PRELIMINARIES . In the following , we first introduce the general problem definition and the associated nomenclature . We then provide an overview of the representation learning framework used to abstract environment behavior , which leverages the recurrent state space model ( RSSM ) proposed by Hafner et al . ( 2019a ) . 3.1 PROBLEM FORMULATION . We formulate the underlying optimization problem as a partially observable Markov decision process ( POMDP ) defined by the tuple { X , A , T , R , Ω , O , γ } , where X , A , O denote the state , action and observation space , respectively , T : X ×A → X signifies the transition mapping , R : X ×A → R the reward mapping , Ω : X → O the observation mapping , and γ ∈ [ 0 , 1 ) is the discount factor . We define xt and at to be the state and action at time t , respectively , and use the notation rt = R ( xt , at ) . Let πφ ( at|ot ) denote a policy parameterized by φ and define the discounted infinite horizon return Gt = ∑∞ τ=t γ τ−tR ( xτ , aτ ) , where xt+1 ∼ T ( xt+1|xt , at ) and at ∼ πφ ( at|ot ) . The goal is then to learn the optimal policy maximizing Gt under unknown nominal dynamics and reward mappings . 3.2 PLANNING FROM PIXELS . We build on the world model introduced in Hafner et al . ( 2019b ) and refined in Hafner et al . ( 2019a ) . High-dimensional image observations are first embedded into a low-dimensional latent space using a neural network encoder . An RSSM then serves as a probabilistic transition model and defines the model state s. Together , the encoder and RSSM define the representation model . The agent therefore abstracts observation ot of environment state xt into model state st , which is leveraged for planning . Consistency of the learned representations is enforced by minimizing the reconstruction error of a decoder network in the observation model and the ability to predict rewards of the reward model . For details , we refer the reader to Hafner et al . ( 2019a ) , and provide their definitions of the models as Representation model : pθ ( st|st−1 , at−1 , ot ) Transition model : qθ ( st|st−1 , at−1 ) Observation model : qθ ( ot|st ) Reward model : qθ ( rt|st ) , ( 1 ) where p and q denote distributions in latent space , with θ as their joint parameterization . The action model πφ ( at|st ) is then trained to optimize the predicted return of imagined world model rollouts . The world model is only rolled-out over a finite horizon H , but complemented by predictions from the value model vψ ( st ) at the terminal state st+H to approximate the infinite horizon return . The policy and value function are trained jointly using policy iteration on the objective functions maxφEqθ , qφ ( ∑t+H τ=t Vλ ( sτ ) ) , minψ Eqθ , qφ ( ∑t+H τ=t 1 2‖vψ ( sτ ) − Vλ ( sτ ) ‖ 2 ) , ( 2 ) respectively . Here , Vλ ( sτ ) represents an exponentially recency-weighted average of the k-step value estimates V kN ( sτ ) along the trajectory to stabilize the learning ( Sutton & Barto , 2018 ) , such that Vλ ( sτ ) . = ( 1− λ ) H−1∑ n=1 λn−1V nN ( sτ ) + λ H−1V HN ( sτ ) , V kN ( sτ ) . = Eqθ , qφ ( h−1∑ n=τ γn−τrn + γ h−τvψ ( sh ) ) , ( 3 ) with h = min ( τ + k , t+H ) .

Summary: The authors proposed latent optimistic value exploration (LOVE) as a mechanism to leverage optimistic exploration for continuous visual control. The main idea is to use a small (~5) ensemble of latent models with shared encoders (and therefore shared learned latent space) but different transition, reward and value models. The variance of predictions from this ensemble can be used as uncertainty estimates of each action sequence while the mean is the typical policy learning objective. Then LOVE puts more weighs on the states with high variance during exploration to enforce the agent to visit (optimistically) uncertain states. This is inherently optimistic because there is a positive bias (beta > 0 and var > 0) addition to the expected reward.