Stochastic Latent Actor-Critic: Deep Reinforcement Learning with a Latent Variable Model

1 INTRODUCTION . Deep reinforcement learning ( RL ) algorithms can automatically learn to solve certain tasks from raw , low-level observations such as images . However , these kinds of observation spaces present a number of challenges in practice : on one hand , it is difficult to directly learn from these high-dimensional inputs , but on the other hand , it is also difficult to tease out a compact representation of the underlying task-relevant information from which to learn instead . For these reasons , deep RL directly from low-level observations such as images remains a challenging problem . Particularly in continuous domains governed by complex dynamics , such as robotic control ( Tassa et al. , 2018 ; Brockman et al. , 2016 ) , standard approaches still require separate sensor setups to monitor details of interest in the environment , such as the joint positions of a robot or specific pose information of objects of interest . To instead be able to learn directly from the more general and rich modality of vision would greatly advance the current state of our learning systems , so we aim to study precisely this . Standard model-free deep RL aims to use direct end-to-end training to explicitly unify these tasks of representation learning and task learning . However , solving both problems together is difficult , since an effective policy requires an effective representation , but in order for an effective representation to emerge , the policy or value function must provide meaningful gradient information using only the model-free supervision signal ( i.e. , the reward function ) . In practice , learning directly from images with standard RL algorithms can be slow , sensitive to hyperparameters , and inefficient . In contrast to end-to-end learning with RL , predictive learning can benefit from a rich and informative supervision signal before the agent has even made progress on the task or received any rewards . This leads us to ask : can we explicitly learn a latent representation from raw low-level observations that makes deep RL easier , through learning a predictive latent variable model ? 1https : //rl-slac.github.io/slac/ Predictive models are commonly used in model-based RL for the purpose of planning ( Deisenroth & Rasmussen , 2011 ; Finn & Levine , 2017 ; Nagabandi et al. , 2018 ; Chua et al. , 2018 ; Zhang et al. , 2019 ) or generating cheap synthetic experience for RL to reduce the required amount of interaction with the real environment ( Sutton , 1991 ; Gu et al. , 2016 ) . However , in this work , we are primarily concerned with their potential to alleviate the representation learning challenge in RL . We devise a stochastic predictive model by modeling the high-dimensional observations as the consequence of a latent process , with a Gaussian prior and latent dynamics , as illustrated in Figure 1 . A model with an entirely stochastic latent state has the appealing interpretation of being able to properly represent uncertainty about any of the state variables , given its past observations . We demonstrate in our work that fully stochastic state space models can in fact be learned effectively : With a well-designed stochastic network , such models outperform fully deterministic models , and contrary to the observations in prior work ( Hafner et al. , 2019 ; Buesing et al. , 2018 ) , are actually comparable to partially stochastic models . Finally , we note that this explicit representation learning , even on low-reward data , allows an agent with such a model to make progress on representation learning even before it makes progress on task learning . Equipped with this model , we can then perform RL in the learned latent space of the predictive model . We posit—and confirm experimentally—that our latent variable model provides a useful representation for RL . Our model represents a partially observed Markov decision process ( POMDP ) , and solving such a POMDP exactly would be computationally intractable ( Astrom , 1965 ; Kaelbling et al. , 1998 ; Igl et al. , 2018 ) . We instead propose a simple approximation that trains a Markovian critic on the ( stochastic ) latent state and trains an actor on a history of observations and actions . The resulting stochastic latent actor-critic ( SLAC ) algorithm loses some of the benefits of full POMDP solvers , but it is easy and stable to train . It also produces good results , in practice , on a range of challenging problems , making it an appealing alternative to more complex POMDP solution methods . The main contributions of our SLAC algorithm are useful representations learned from our stochastic sequential latent variable model , as well as effective RL in this learned latent space . We show experimentally that our approach substantially improves on both model-free and model-based RL algorithms on a range of image-based continuous control benchmark tasks , attaining better final performance and learning more quickly than algorithms based on ( a ) end-to-end deep RL from images , ( b ) learning in a latent space produced by various alternative latent variable models , such as a variational autoencoder ( VAE ) ( Kingma & Welling , 2014 ) , and ( c ) model-based RL based on latent state-space models with partially stochastic variables ( Hafner et al. , 2019 ) . 2 RELATED WORK . Representation learning in RL . End-to-end deep RL can in principle learn representations directly as part of the RL process ( Mnih et al. , 2013 ) . However , prior work has observed that RL has a “ representation learning bottleneck ” : a considerable portion of the learning period must be spent acquiring good representations of the observation space ( Shelhamer et al. , 2016 ) . This motivates the use of a distinct representation learning procedure to acquire these representations before the agent has even learned to solve the task . The use of auxiliary supervision in RL to learn such representations has been explored in a number of prior works ( Lange & Riedmiller , 2010 ; Finn et al. , 2016 ; Jaderberg et al. , 2017 ; Higgins et al. , 2017 ; Ha & Schmidhuber , 2018 ; Nair et al. , 2018 ; Oord et al. , 2018 ; Gelada et al. , 2019 ; Dadashi et al. , 2019 ) . In contrast to this class of representation learning algorithms , we explicitly learn a latent variable model of the POMDP , in which the latent representation and latent-space dynamics are jointly learned . By modeling covariances between consecutive latent states , we make it feasible for our proposed algorithm to perform Bellman backups directly in the latent space of the learned model . Partial observability in RL . Our work is also related to prior research on RL under partial observability . Prior work has studied exact and approximate solutions to POMDPs , but they require explicit models of the POMDP and are only practical for simpler domains ( Kaelbling et al. , 1998 ) . Recent work has proposed end-to-end RL methods that use recurrent neural networks to process histories of observations and ( sometimes ) actions , but without constructing a model of the POMDP ( Hausknecht & Stone , 2015 ; Foerster et al. , 2016 ; Zhu et al. , 2018 ) . Other works , however , learn latent-space dynamical system models and then use them to solve the POMDP with model-based RL ( Watter et al. , 2015 ; Wahlström et al. , 2015 ; Karl et al. , 2017 ; Zhang et al. , 2019 ; Hafner et al. , 2019 ) . Although some of these works learn latent variable models that are similar to ours , these model-based methods are often limited by compounding model errors and finite horizon optimization . In contrast to these works , our approach does not use the model for prediction and performs infinite horizon policy optimization . Our approach benefits from the good asymptotic performance of model-free RL , while at the same time leveraging the improved latent space representation for sample efficiency . Other works have also trained latent variable models and used their representations as the inputs to model-free RL algorithms . They use representations encoded from latent states sampled from the forward model ( Buesing et al. , 2018 ) , belief representations obtained from particle filtering ( Igl et al. , 2018 ) , or belief representations obtained directly from a learned belief-space forward model ( Gregor et al. , 2019 ) . Our approach is closely related to these prior methods , in that we also use model-free RL with a latent state representation that is learned via prediction . However , instead of using belief representations , our method learns a critic directly on latent states samples . Sequential latent variable models . Several previous works have explored various modeling choices to learn stochastic sequential models ( Krishnan et al. , 2015 ; Archer et al. , 2015 ; Karl et al. , 2016 ; Fraccaro et al. , 2016 ; 2017 ; Doerr et al. , 2018a ) . In the context of using sequential models for RL , previous works have typically observed that partially stochastic state space models are more effective than fully stochastic ones ( Buesing et al. , 2018 ; Igl et al. , 2018 ; Hafner et al. , 2019 ) . In these models , the state of the underlying MDP is modeled with the deterministic state of a recurrent network ( e.g. , LSTM ( Hochreiter & Schmidhuber , 1997 ) or GRU ( Cho et al. , 2014 ) ) , and optionally with some stochastic random variables . As mentioned earlier , a model with a latent state that is entirely stochastic has the appealing interpretation of learning a representation that can properly represent uncertainty about any of the state variables , given past observations . We demonstrate in our work that fully stochastic state space models can in fact be learned effectively and , with a well-designed stochastic network , such models perform on par to partially stochastic models and outperform fully deterministic models . 3 REINFORCEMENT LEARNING AND MODELING . This work addresses the problem of learning maximum entropy policies from high-dimensional observations in POMDPs , by simultaneously learning a latent representation of the underlying MDP state using variational inference and learning the policy in a maximum entropy RL framework . In this section , we describe maximum entropy RL ( Ziebart , 2010 ; Haarnoja et al. , 2018a ; Levine , 2018 ) in fully observable MDPs , as well as variational methods for training latent state space models for POMDPs . 3.1 MAXIMUM ENTROPY RL IN FULLY OBSERVABLE MDPS . In a Markov decision process ( MDP ) , an agent at time t takes an action at ∈ A from state st ∈ S and reaches the next state st+1 ∈ S according to some stochastic transition dynamics p ( st+1|st , at ) . The initial state s1 comes from a distribution p ( s1 ) , and the agent receives a reward rt on each of the transitions . Standard RL aims to learn the parameters φ of some policy πφ ( at|st ) such that the expected sum of rewards is maximized under the induced trajectory distribution ρπ . This objective can be modified to incorporate an entropy term , such that the policy also aims to maximize the expected entropyH ( πφ ( ·|st ) ) under the induced trajectory distribution ρπ . This formulation has a close connection to variational inference ( Ziebart , 2010 ; Haarnoja et al. , 2018a ; Levine , 2018 ) , and we build on this in our work . The resulting maximum entropy objective is φ∗ = arg max φ T∑ t=1 E ( st , at ) ∼ρπ [ r ( st , at ) + αH ( πφ ( ·|st ) ) ] , ( 1 ) where r is the reward function , and α is a temperature parameter that controls the trade-off between optimizing for the reward and for the entropy ( i.e. , stochasticity ) of the policy . Soft actor-critic ( SAC ) ( Haarnoja et al. , 2018a ) uses this maximum entropy RL framework to derive soft policy iteration , which alternates between policy evaluation and policy improvement within the described maximum entropy framework . SAC then extends this soft policy iteration to handle continuous action spaces by using parameterized function approximators to represent both the Q-function Qθ ( critic ) and the policy πφ ( actor ) . The soft Q-function parameters θ are optimized to minimize the soft Bellman residual , JQ ( θ ) = E ( st , at , rt , st+1 ) ∼D [ 1 2 ( Qθ ( st , at ) − ( rt + γVθ̄ ( st+1 ) ) ) 2 ] , ( 2 ) Vθ̄ ( st+1 ) = E at+1∼πφ [ Qθ̄ ( st+1 , at+1 ) − α log πφ ( at+1|st+1 ) ] , ( 3 ) where D is the replay buffer , γ is the discount factor , and θ̄ are delayed parameters . The policy parameters φ are optimized to update the policy towards the exponential of the soft Q-function , Jπ ( φ ) = E st∼D [ E at∼πφ [ α log ( πφ ( at|st ) ) −Qθ ( st , at ) ] ] . ( 4 ) Results of this stochastic , entropy maximizing RL framework demonstrate improved robustness and stability . SAC also shows the sample efficiency benefits of an off-policy learning algorithm , in conjunction with the high performance benefits of a long-horizon planning algorithm . Precisely for these reasons , we choose to extend the SAC algorithm in this work to formulate our SLAC algorithm .

Summary: This paper proposes an actor-critic method that tries to aid learning good policies via learning a good representation of the state space (via a latent variable model). In actor-critic methods, the critic is learnt to evaluate policies in the latent space, which further helps with efficient policy optimization. The proposed method is evaluated on image-based control tasks, with baseline evaluations against both model-based and model-free methods in terms of sample efficiency.