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: The paper proposes SOLAR a model based RL algorithm that learns a low dimensional embedding such that the dynamics within the latent space are linear. Within this latent space the linear dynamics  are learned using a Bayesian regression. In addition, a quadratic cost function is approximated. The learned dynamics and the cost function are used to update the policy, while simultaneously bounding the change in policy by a KL-bound.  In contrast to other model-based RL algorithms the learned dynamics are not used for planning or imaginary roll-outs and are only used to improve the policy.