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: This paper proposes LOVE (Latent Optimistic Value Exploration), a model-based exploration algorithm for POMDP or pixel-based control systems. The method builds upon Dreamer (Hafner et al. 2019) for learning latent models, and the variance of value estimates (one transition/reward/value model per particle) estimates the epistemic uncertainty of the world, which can be used as a UCB-style exploration reward. LOVE can achieve a comparable or better performance on a simple pointmass environment with a trap and standard DeepMind continuous control suites.