SOLAR: Deep Structured Representations for Model-Based Reinforcement Learning

1 INTRODUCTION . Model-based reinforcement learning ( RL ) methods use learned models in a variety of ways , such as planning ( Levine & Abbeel , 2014 ; Deisenroth et al. , 2014 ) and generating synthetic experience ( Sutton , 1990 ) . We can categorize model-based algorithms as either global model methods , where models are used for planning and trained to give accurate predictions for a wide range of states , or local model methods , where simple models provide gradient directions that are used for policy improvement . On simple , low-dimensional tasks , model-based approaches have demonstrated remarkable data efficiency , learning policies for systems like cart-pole swing-up with under 30 seconds of experience ( Deisenroth et al. , 2014 ; Moldovan et al. , 2015 ) . However , for more complex systems , one of the main difficulties in applying model-based methods is model bias : local models will often underfit complex systems , but may still be preferred over global models which tend to overfit in the low-data regime and may be difficult to incorporate into control methods ( Deisenroth et al. , 2014 ) . Most global model methods use the model to make forward predictions and then backpropagate through those predictions . However , this places a heavy burden on the dynamics model , and forward prediction often suffers from significant drift over longer trajectories . In contrast , local models are typically only used to provide gradient directions for local policy improvement ( Levine & Abbeel , 2014 ) , and thus a common choice for local model methods is to use linear models , which can themselves be interpreted as gradients . As illustrated in Figure 1 , in our work , we present a method that automatically encourages learning representations where linear models better fit the data . From this , we devise an efficient local model method based on the linear-quadratic regulator ( LQR ) ( Camacho & Bordons , 1997 ; Todorov & Li , 2005 ; Levine & Abbeel , 2014 ) that utilizes linear models for gradient directions for policy improvement . Our motivation is similar to that of Watter et al . ( 2015 ) ; Finn et al . ( 2016 ) ; however , as discussed in section 5 , our representation learning method specifically allows us to construct a local model method that performs inference in the latent space in order to improve the policy , rather than focusing on forward prediction and planning . Our main contribution is a representation learning and model-based RL procedure , which we term stochastic optimal control with latent representations ( SOLAR ) , which jointly optimizes a latent representation and model such that inference produces local linear models that provide good gradient directions for policy improvement . We demonstrate empirically in section 6 that SOLAR is able to learn policies directly from raw , high-dimensional observations in several robotic environments including a simulated nonholonomic car , a simulated two degree-of-freedom ( DoF ) arm , and a real 7-DoF Sawyer arm , all of which are learned directly from image pixels . We compare to existing state-of-the-art RL methods and show that SOLAR , while significantly more data efficient than model-free methods , exhibits superior performance compared to other model-based methods . 2 PRELIMINARIES . We first formalize our problem setting as a Markov decision process ( MDP ) M = ( S , A , p , C , ρ , T ) , where the state space S , action space A , and horizon T are known , but the dynamics function p ( st+1|st , at ) , cost function C ( st , at ) , and initial state distribution ρ ( s0 ) are unknown . The goal of reinforcement learning is to optimize a policy π ( at|st ) to minimize the expected sum of costs η [ π ] = Eπ , p , ρ [ ∑T t=0 C ( st , at ) ] under the distribution induced by the initial state distribution , dynamics function , and policy . Model-based methods decompose this problem into policy and model optimization subproblems , and we discuss each subproblem as it relates to our approach . 2.1 MODEL-BASED POLICY SEARCH . Policy search methods directly optimize parameterized policies with respect to η ( θ ) , η [ πθ ] where the parameters θ may be , for example , weights in a neural network or matrices for a linear policy . Model-based policy search methods typically build models ( ρ̂ , p̂ , Ĉ ) of the unknown quantities and compute the gradient of η̂ ( θ ) , Eπθ , p̂ , ρ̂ [ ∑T t=0 Ĉ ( st , at ) ] with this model . One particularly tractable model is the linear-quadratic system ( LQS ) , which models the initial state distribution as Gaussian , the dynamics as time-varying linear-Gaussian ( TVLG ) , and the cost as quadratic , i.e. , p̂ ( st+1|st , at ) = N ( st+1 ∣∣∣∣ Ft [ stat ] , Σt ) , Ĉ ( st , at ) = 1 2 [ st at ] > C [ st at ] + c > [ st at ] . ( 1 ) Any deterministic policy operating in an environment with smooth dynamics can be locally modeled with a time-varying LQS ( Boyd & Vandenberghe , 2004 ) , while low-entropy stochastic policies are modeled approximately . This makes the time-varying LQS a reasonable local model for many dynamical systems . Furthermore , the optimal policy at any time step given the model is a linear function of the state and the optimal maximum-entropy policy is linear-Gaussian ( Tassa et al. , 2012 ; Levine & Koltun , 2013 ) . As shown in Jacobson & Mayne ( 1970 ) ; Todorov & Li ( 2005 ) , these optimal policies can be computed in closed form using dynamic programming by computing the first and second derivatives of the Q ( cost-to-go ) and value functions : Qs̃ , t = cs̃ , t + F > s̃ , tVs , t+1 , Qs̃s̃ , t = Cs̃s̃ , t + F > s̃s̃ , tVss , t+1Fs̃s̃ , t , Vs , t = Qs , t −Qsa , tQ−1aa , tQa , t , Vss , t = Qss , t −Qsa , tQ−1aa , tQas , t . Here , similar to Tassa et al . ( 2012 ) , we use subscripts to denote derivatives , and we use s̃ to abbreviate [ s a ] . Once these values are computed , the optimal maximum-entropy policy is TVLG , i.e. , πθ ( at|st ) = N ( Ktst + kt , St ) , where Kt = −Q−1aa , tQas , t , kt = −Q−1aa , tQa , t , St = −Q−1aa , t . We refer the reader to Appendix A and Levine & Abbeel ( 2014 ) for further details . Prior work assumes access to a compact , low-dimensional state representation ( Deisenroth et al. , 2014 ; Levine & Abbeel , 2014 ; Nagabandi et al. , 2018 ) , and as we show in section 6 , this precludes these local model methods from operating on complex observations such as images . In subsection 2.2 and section 3 , we describe a probabilistic latent variable model and variational inference procedure that , conditioned on a full trajectory of observations , produces local models that can be used for policy improvement , enabling us to utilize this local model method in image-based domains . 2.2 LEARNING LATENT DYNAMICS MODELS . The local model-based method described above requires us to learn both a quadratic cost function as well as a linear dynamical system ( LDS ) . We utilize the Bayesian LDS model , which is given by µρ̂ , Σρ̂ ∼ NIW ( Ψ , ν , µ0 , κ ) , Ft , Σt ∼MNIW ( Ψ , ν , M0 , V ) for t ∈ [ 0 , . . . , T − 1 ] , s0 | µρ̂ , Σρ̂ ∼ N ( µρ̂ , Σρ̂ ) , st+1 | st , at ∼ N ( Ft [ st at ] , Σt ) for t ∈ [ 0 , . . . , T − 1 ] , Where NIW is the normal-inverse-Wishart distribution andMNIW is the matrix normal-inverseWishart ( MNIW ) distribution . This probabilistic graphical model ( PGM ) allows for tractable approximate inference , i.e. , Bayesian linear regression , and also captures uncertainty in the form of a posterior distribution over the initial state and dynamics . However , for dynamical systems with complex non-linear dynamics , this model still suffers from significant bias . Even when the system is poorly modeled by an LDS in the state space , we might be able to find a latent embedding and model the system as approximately linear in that latent space , which may allow us to find a better-performing policy that operates in the learned latent space . This shifts our problem setting to that of a partially observed MDP , as we do not observe the latent state . In particular , our modeling assumption is that we receive an observation as generated from an underlying unobserved state , and as discussed in section 3 , we address this by training a recognition model to infer the latent state . In our experiments in section 6 , we provide several observations to our recognition model in order to infer information that can not be observed from a single observation , such as velocity . We can jointly train an embedding and model using the SVAE framework ( Johnson et al. , 2016 ) , which allows us to combine arbitrary embedding functions , such as neural networks , with PGMs . The model we build off of is a version of the LDS SVAE presented in Johnson et al . ( 2016 ) and is given by µρ̂ , Σρ̂ ∼ NIW ( Ψ , ν , µ0 , κ ) , Ft , Σt ∼MNIW ( Ψ , ν , M0 , V ) for t ∈ [ 0 , . . . , T − 1 ] , ( 2 ) z0 | µρ̂ , Σρ̂ ∼ N ( µρ̂ , Σρ̂ ) , zt+1 | zt , at ∼ N ( Ft [ zt at ] , Σt ) for t ∈ [ 0 , . . . , T − 1 ] , ( 3 ) st | zt ∼ fγ ( zt ) for t ∈ [ 0 , . . . , T ] , ( 4 ) Where fγ ( z ) is an observation model , parameterized by neural network weights γ , that outputs a distribution over s , e.g. , Gaussian or Bernoulli , depending on the nature of the data . This is very similar to the Bayesian LDS , except we are learning the PGM in the latent space . Though this model does not admit the same efficient approximate inference algorithms when fγ is nonlinear , an efficient variational inference algorithm has previously been derived by Johnson et al . ( 2016 ) . We describe the relevant aspects of this algorithm in the next section . 3 LEARNING AND MODELING THE LATENT SPACE . In this section , we describe how we extend the LDS SVAE for model-based RL , such that we learn an action-conditioned LQS model in the latent space . This then enables a local model method that can leverage the LQS to infer the dynamics of sampled trajectories . In this way , our model-based RL algorithm circumvents the need for forward prediction , in contrast to model-based RL methods that use model-based rollouts or planning ( Nagabandi et al. , 2018 ; Deisenroth et al. , 2014 ) . In section 4 , we describe how these components are combined into our final method , SOLAR . Our goal with this model is to learn a latent representation of the state and a prior over the dynamics in this latent representation that is suitable for fitting local dynamics models via posterior inference . Specifically , we are interested in the setting where we have access to trajectories of the form [ s0 , a0 , c0 , . . . , sT−1 , aT−1 , cT−1 , sT ] , sampled from the system using our current policy and set of previous policies . Our aim is to infer local linear dynamics in the neighborhood of these trajectories , and we learn a model that makes this fitting process more accurate for the observed trajectories , thus enabling our local model method to find good directions for policy improvement . We build upon the variational inference algorithm presented in Johnson et al . ( 2016 ) , such that we are maximizing , with respect to both the PGM and neural network parameters , the variational lower bound ( ELBO ) of our observed data . This algorithm requires variational factors of the form q ( zt | st ) = N ( eφ ( st ) ) , q ( Ft , Σt ) =MNIW ( Ψ′t , ν′t , M ′0t , V ′t ) for t ∈ [ 0 , . . . , T − 1 ] . ( 5 ) eφ ( s ) is a recognition model , parameterized by neural network weights φ , that outputs the mean and diagonal covariance of a Gaussian distribution over z . This recognition model is identical to that used in Kingma & Welling ( 2014 ) ; Rezende et al . ( 2014 ) ; Gao et al . ( 2016 ) , however , as with prior work in the LDS SVAE , we also have variational factors of the form q ( Ft , Σt ) , which represent our posterior belief about the system dynamics after observing the collected data . We also model this distribution as MNIW but with updated parameters compared to the prior from Equation 2 . Given this , we can formulate the variational lower bound ( ELBO ) which is given by L = Eq [ log p ( { F , Σ } T−1t=0 , { st } Tt=0 , { at } T−1 t=0 , zt } Tt=0 ) q ( { Ft , Σt } T−1t=0 , { zt } Tt=0| { st } Tt=0 ) ] = Eq [ log ( T∏ t=0 pγ ( st|zt ) ) ] − T−1∑ t=0 KL ( q ( Ft , Σt ) ‖p ( F , Σ ) ) − T∑ t=1 Eq [ KL ( qφ ( zt|st ) ‖p ( zt|zt−1 , at−1 , Ft , Σt ) ] . Prior work has shown that , for conjugate exponential models such as the Bayesian LDS , the variational model parameters can be updated using natural gradients , which can be computed in closed form using the variational message passing framework ( Winn & Bishop , 2005 ) . Specifically , letting λ denote the MNIW parameters of the variational factors on { Ft , Σt } t , the natural gradient update is ∇̃λL = λ0 +BEq [ tF , Σ ( F , Σ ) ] − λ , ( 6 ) Where B is the number of minibatches in the dataset , λ0 is the parameter for the prior distribution p ( F , Σ ) , and tF , Σ ( F , Σ ) is the sufficient statistic function for p ( F , Σ ) . Thus , we can use this equation to compute the natural gradient update for λ , whereas for γ and φ we use stochastic gradient updates on Monte Carlo estimates of the ELBO , specifically using the Adam optimization scheme ( Kingma & Ba , 2015 ) . This leads to two simultaneous optimizations for the PGM parameters and the neural network parameters , and their learning rates are treated as separate hyperparameters . We have found 10−3 and 10−4 to be generally suitable for the natural gradient and Adam updates , respectively . Algorithm 1 SOLAR 1 : Hyperparameters : # iterations K , # trajectories N , model training buffer size B 2 : Initialize policy π ( 0 ) θ , modelM ( 0 ) 3 : for iteration k ∈ { 1 , . . . , K } do 4 : Collect rollouts from the real world D ( k ) = { ( s ( i ) 0 , a ( i ) 0 , . . . , s ( i ) T ) } Ni=1 5 : M ( k ) ← MODELUPDATE ( M ( k−1 ) , { D ( i ) } ki=k−B ) ( section 3 ) 6 : π̃ ( k−1 ) θ ← LINEARIZEPOLICY ( D ( k ) , M ( k ) ) ( Appendix C ) 7 : { F ( k ) t , Σ ( k ) t } t ← INFERDYNAMICS ( D ( k ) , M ( k ) ) ( subsection 4.1 ) 8 : π ( k ) θ ← POLICYUPDATE ( π̃ ( k−1 ) θ , { F ( k ) t , Σ ( k ) t } t , M ( k ) ) ( subsection 4.2 ) 9 : end for Figure 2 details the graphical model presented in Equation 2-Equation 4 along with the variational family described above . Since we are interested in control and RL , there is the added notion of observed costs from the environment , and there are many ways we could model these additional observations . A natural choice is to model costs as a quadratic function of the latent state and action , such that we arrive at the LQS presented in Equation 1 except in the learned latent space . Specifically , given trajectories of the form [ s0 , a0 , c0 , . . . , sT−1 , aT−1 , cT−1 , sT ] , we first embed the observations { st } using the mean of our recognition model µ ( eφ ( s ) ) to obtain a set of latent states { zt } . We then model our cost samples as ct = 12z > t LL > zt + c > zt + α‖at‖22 + b , where we assume that the action-dependent part of the cost is known and we learn L , c , and b by minimizing the mean-squared error of the observed costs with stochastic gradient descent . L is a lower-triangular matrix with strictly positive diagonal entries , and thus by constructing our cost matrix as C = LL > we guarantee that the learned cost matrix is positive definite , which improves the conditioning of the policy update .

Summary: In this work the authors propose an end to end approach for model based reinforcement learning from images, where the main building blocks are locally-linear dynamical systems and variational auto-encoders (VAE). Specifically, it is assumed that the input features (i.e., the images) are generated from a low dimensional latent representation mapped through parametric random functions; the latter are modeled via neural networks. A recognition model based on convolutional neural networks operates on the reverse way and is responsible for projecting the input features to the latent space, in order to proceed with the reinforcement learning task. The variational framework is employed in order to jointly learn the VAE and the linear dynamics on the latent state. As a final step, once the model is fitted a linear quadratic system (LQS) is solved in order to learn the cost function and the optimal policy.