SOLAR: D EEP S TRUCTURED R EPRESENTATIONS FOR M ODEL -B ASED R EINFORCEMENT L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Model-based reinforcement learning (RL) methods can be broadly categorized as global model methods, which depend on learning models that provide sensible predictions in a wide range of states, or local model methods, which iteratively reﬁt simple models that are used for policy improvement.
While predicting future states that will result from the current actions is difﬁcult, local model methods only attempt to understand system dynamics in the neighborhood of the current policy, making it possible to produce local improvements without ever learning to predict accurately far into the future.
The main idea in this paper is that we can learn representations that make it easy to retrospectively infer simple dynamics given the data from the current policy, thus enabling local models to be used for policy learning in complex systems.
We evaluate our approach against other model-based and model-free RL methods on a suite of robotics tasks, including a manipulation task on a real Sawyer robotic arm directly from camera images.
I NTRODUCTION
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 efﬁciency, 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 difﬁculties in applying model-based methods is model bias: local models will often underﬁt complex systems, but may still be preferred over global models which tend to overﬁt in the low-data regime and may be difﬁcult 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 signiﬁcant 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 ﬁt the data.
From this, we devise an efﬁcient 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.
Finn et al. (2016); however, as discussed in section 5, our representation learning method speciﬁcally 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
[Figure]
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 signiﬁcantly more data efﬁcient than model-free methods, exhibits superior performance compared to other model-based methods.
P RELIMINARIES
We ﬁrst 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 ( s t +1 | s t , a t ) , cost function C
( s t , a t ) , and initial state distribution ρ ( s 0 ) are unknown.
The goal of reinforcement learning is to optimize a policy π ( a t
| s t ) to minimize the expected sum of costs η [ π ] =
E π,p,ρ (cid:104)(cid:80) Tt =
C ( s t , a t ) (cid:105) 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. M ODEL - BASED POLICY SEARCH
Policy search methods directly optimize parameterized policies with respect to η ( θ ) (cid:44) η
[ π θ ] 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 (cid:16) ˆ ρ, ˆ p, ˆ C (cid:17) of the unknown quantities and compute the gradient of ˆ η ( θ ) (cid:44)
E π θ , ˆ p, ˆ ρ (cid:104)(cid:80) Tt =
ˆ C ( s t , a t ) (cid:105) 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.,
[Equation]
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 ﬁrst and second derivatives of the Q (cost-to-go) and value functions:
[Equation]
Here, similar to Tassa et al.
(2012), we use subscripts to denote derivatives, and we use ˜ s to abbreviate (cid:104) sa (cid:105) .
Once these values are computed, the optimal maximum-entropy policy is TVLG, i.e., π θ ( a t
| s t )
=
N ( K t s t + k t , S t ) , where K t = − Q − 1 aa ,t Q as ,t , k t = − Q − 1 aa ,t Q a ,t , S t = − Q − 1 aa ,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. L EARNING 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
[Equation]
) for t ∈ [
[Equation] [Equation]
Where NIW is the normal-inverse-Wishart distribution and MNIW 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 signiﬁcant bias.
Even when the system is poorly modeled by an LDS in the state space, we might be able to ﬁnd a latent embedding and model the system as approximately linear in that latent space, which may allow us to ﬁnd 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 cannot 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.
and is given by
[Equation]
t [Equation] 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 efﬁcient approximate inference algorithms when f γ is nonlinear, an efﬁcient variational inference algorithm has previously been derived by Johnson et al.
We describe the relevant aspects of this algorithm in the next section.
L EARNING AND M ODELING THE L ATENT S PACE
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 ﬁnal 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 ﬁtting local dynamics models via posterior inference.
Speciﬁcally, we are interested in the setting where we have access to trajectories of the form [ s 0 , a 0 , c 0 , . .
, s T − 1 , a T − 1 , c T − 1 , s T ] , 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 ﬁtting process more accurate for the observed trajectories, thus enabling our local model method to ﬁnd 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 (
[Equation]
MNIW
[Equation]
t ) for t ∈ [0 , . . . [Equation] 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.
; Gao et al.
(2016), however, as with prior work in the LDS SVAE, we also have variational factors of the form q ( F t , Σ 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 (see Appendix B) which is optimized using stochastic gradient updates, for the neural network parameters γ and φ , and natural gradient updates and the variational message passing (VMP) framework (Winn & Bishop, 2005), for the graphical model parameters.
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.
Speciﬁcally, we choose to ﬁt C and c to the data collected by the current policy by encoding the observed states from our policy and regressing the latent states and actions to their observed costs.
We refer readers to Appendix B for more details about the model training and cost ﬁtting procedures.
P OLICY L EARNING IN THE L ATENT S PACE
While we could use a variety of model-based policy learning methods in the learned latent space, the ability to infer local time-varying linear dynamics lends itself naturally to the particular analytic local solution to the policy described in subsection 2.1.
This approach yields a policy that is TVLG in the latent space, which in general corresponds to a class of nonlinear policies in the original space formed by the composition of the nonlinear neural network embedding and the TVLG policy.
As discussed in the following sections, we can use the PGM in the previous section to formulate local model ﬁtting as probabilistic inference, in order to obtain a dynamics estimate that we can then use to improve the policy.
Note that this use of the model is quite different from how dynamics models are typically used in standard model-based RL algorithms: instead of using the model to predict into the future, we only use the model to infer local linear dynamics conditioned on real-world trajectory samples.
While local models are not burdened by forward prediction compared to global forward
1: Hyperparameters: # iterations K , # trajectories N , model training buffer size B 2: Initialize policy π (0) θ , model M (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 = 5:
M ( k ) ← M ODEL
U PDATE ( M ( k − 1) , {D ( i ) } ki = k − B ) (section 3) 6: ˜ π ( k − 1) θ
← L INEARIZE P OLICY ( D ( k ) , M ( k ) )
(Appendix D) 7: { F ( k ) t , Σ ( k ) t } t ← I NFER D YNAMICS
( D ( k ) , M ( k ) )
(subsection 4.1) 8: π ( k ) θ ← P OLICY U PDATE (˜ π ( k − 1) θ , { F ( k ) t , Σ ( k ) t } t , M ( k ) )
(subsection 4.2) 9: end for models, the simplicity of linear local models prevents accurate modeling of complex systems, and our method mitigates this through a latent representation that is optimized for local linear model ﬁtting.
Our overall algorithm, SOLAR, is presented in algorithm 1.
At every iteration, we collect N rollouts from the real world (line 4). Then, we update our model using data from the last B iterations (line 5), we linearize our policy given the updated model (line 6, see Appendix D for details), we perform inference within our model to get the dynamics estimates (line 7), and we update our policy using the rollouts from our current iteration and our updated model (line 8).
The following subsections detail the modules of our method that are involved in policy learning and improvement.
4. D YNAMICS INFERENCE UNDER THE MODEL
To obtain a TVLG dynamics model, we could directly use linear regression to ﬁt F t and Σ t to the observed latent trajectories τ
= [ z 0 , a 0 , . .
, z T − 1 , a T − 1 , z T ] .
However, this may be poorly conditioned in the low-data regime.
Instead, we can perform inference within our model to obtain dynamics estimates for policy improvement.
As described in section 3, our model provides us with variational approximations to the posterior over dynamics models, i.e., { q ( F t , Σ t ) } T − 1 t =0 , which are
MNIW.
We can use these as a prior and condition on the data to obtain new variational posteriors { q ( F t , Σ t |{ τ } Ni =0 ) } T − 1 t =0 , which are also MNIW.
Writing the parameters of these posteriors – for which the closed form solutions are given in Appendix C– as { Ψ t , M 0 t , V t , ν t } t , we compute a maximum a posteriori estimate of the dynamics parameters at time step t as: F t
= M 0 t , Σ t =
Ψ t ν t .
This inference procedure corresponds to Bayesian linear regression and can be interpreted as resolving the uncertainty in the global dynamics model conditioned on a real-world rollout.
In essence, { q ( F t , Σ t ) } T − 1 t =0 captures uncertainty over the latent system dynamics by acting as a global model over all observed data, but in order to accurately model the system within the local region around the current policy, we condition on trajectories collected from the policy in order to resolve the uncertainty and obtain dynamics estimates { F t , Σ t } T − 1 t =0 that allow us to improve the policy.
4. P OLICY UPDATE
As described in subsection 2.1, once we have our TVLG dynamics estimates { F t , Σ t } t and quadratic cost ﬁt C , c , we can use dynamic programming on the Q and value functions to compute the optimal policy in closed form.
However, doing so is typically undesirable as the resulting policy will overﬁt to the model and likely will not perform well in the real environment.
Since our modeling assumption is not that our model will be globally valid, but rather that our model will be valid close to the data distribution of the previous policy, we utilize a constrained policy update such that our new policy does not drastically change the induced trajectory distribution.
Speciﬁcally, similar to prior work, we impose a KL-divergence constraint on the policy update such that the shift in the induced trajectory distributions before and after the update, which we denote as ¯ p ( τ ) and p ( τ ) , respectively, is bounded by a step size (cid:15) (Levine & Abbeel, 2014).
This leads to a constrained optimization of the form max θ ˆ η ( θ ) s.t. D KL ( p ( τ ) (cid:107)
¯ p ( τ )) ≤ (cid:15) .
As shown in Levine & Abbeel (2014), this constrained optimization can be solved by augmenting the cost function to penalize the deviation from the previous policy π ¯ θ , i.e., ˜ C ( z t , a t )
= 1 λ C ( z t , a t ) − log π ¯ θ ( a t | z t ) .
Note that this augmented cost function is still quadratic, since the policy is TVLG, and thus we can still compute the optimal policy under this cost function in closed form using the procedure described in subsection 2.1.
λ is a dual variable that trades off between optimizing the cost function and staying close in distribution to the previous policy, and the weight of this term can be determined through a dual gradient descent procedure.
Combined with the model learning from section 3, we arrive at the SOLAR algorithm.
R ELATED W ORK
Model-based RL methods have achieved signiﬁcant efﬁciency beneﬁts compared to model-free RL methods (Chebotar et al., 2017; Nagabandi et al., 2018; Deisenroth et al., 2014).
Many of these prior methods learn global models of the system that are then used for planning, generating synthetic experience, or policy search (Atkeson & Schaal, 1997; Peters et al., 2010).
These methods require an accurate and reliable model and will typically suffer from modeling bias, hence these models are still limited to short horizon prediction in more complex domains (Mishra et al., 2017; Nagabandi et al., 2018; Gu et al., 2016; V.Feinberg et al., 2018).
Another class of model-based methods rely only on local system models to compute the gradient for a policy update (An et al., 1988;
Kolter &
Ng, 2005; Heess et al., 2015; Levine & Abbeel, 2014; Bansal et al., 2017).
These methods do not use models for long-term forward prediction, allowing for the use of simple models that enable policy improvement (Montgomery et al., 2017; Levine et al., 2016).
As we show in section 6, modeling bias for prior methods can be severely limiting in systems with complex observations such as images, whereas we are able to learn representations that mitigate the effects of modeling bias.
Utilizing representation learning within model-based RL has been studied in a number of previous works (Lesort et al., 2018), including using embeddings for state aggregation (Singh et al., 1994), dimensionality reduction (Nouri & Littman, 2010), self-organizing maps (Smith, 2002), value prediction (Oh et al., 2017), and deep auto-encoders (Lange & Riedmiller, 2010; Finn et al., 2016;
Watter et al., 2015; Higgins et al., 2017).
Within these works, deep spatial auto-encoders (DSAE)
(Finn et al., 2016) and embed to control (E2C) (Watter et al., 2015; Banijamali et al., 2017) are the most closely related to our work in that they consider local model methods combined with representation learning.
The key difference in our work is that, rather than using a learning objective for reconstruction and forward prediction, we formulate a Bayesian latent variable model such that inference corresponds to ﬁtting local models within the learned representation.
As such, our objective enables local model methods by directly encouraging learning representations where ﬁtting local models accurately explains the observed data.
We also do not assume a known cost function, goal state, or access to the underlying system state as in DSAE and E2C, thus SOLAR is applicable even when the underlying states and cost function are unknown.
1 We ﬁnd that our approach tends to produce better results on a number of complex image-based tasks, as we discuss in the next section.
E XPERIMENTS
We aim to answer the following questions through our experiments: (1) How does SOLAR compare to state-of-the-art model-free and model-based RL algorithms?
(2) How do local and global model methods compare when operating in our learned representations?
(3) How much beneﬁt do we derive from our particular representation learning method?
To answer (1), we compare SOLAR to trust region policy optimization (TRPO)
(Schulman et al., 2015) and proximal policy optimization (PPO) (Schulman et al., 2017), two state-of-the-art model-free methods, and LQR with ﬁtted linear models (LQR-FLM) (Levine & Abbeel, 2014), a state-of-the-art model-based method.
To answer (2), we test an ablation of our method where we learn a neural network dynamics model with which we perform model-predictive control (MPC) in the latent space.
We refer to this as the “global model ablation”.
To answer (3), we replace our LDS SVAE model with a variational auto-encoder (VAE) (Kingma & Welling, 2014; Rezende et al., 2014) and with the robust locally-linear controllable embedding (RCE) model (Banijamali et al., 2017), an improved version of the E2C model (Watter et al., 2015).
We refer to these as the “VAE ablation” and “E2C-like ablation”, respectively.
We additionally compare to a pixel space model similar to Finn & Levine (2017) that utilizes no [Figure] representation learning and instead learns both a dynamics and cost model on images in order to run
MPC in pixel space.
Videos of the learned policies are available on the project website.
6. E XPERIMENTAL T ASKS
We set up simulated image-based robotic domains for a 2-dimensional navigation task, a nonholonomic car, and a 2-DoF arm, as shown in Figure 3a.
We also learn a block stacking task directly from camera images on a real Sawyer robotic arm, as shown in Figure 3b.
Details regarding experimental setup and training hyperparameters are provided in Appendix E.
2D navigation.
We consider a 2-dimensional navigation task similar to Watter et al.
; Banijamali et al.
(2017) except we move the goal every episode rather than ﬁxing it to the bottom right.
Observations consist of two 32-by-32 images indicating the positions of the agent and goal.
Nonholonomic car.
The nonholonomic car starts in the bottom right of the 2-dimensional space and controls its acceleration and steering velocity in order to reach the target in the top left.
We use a sliding window of four 64-by-64 images as the observation to capture velocity information.
Reacher.
We experiment with the reacher environment from OpenAI Gym (Brockman et al., 2016), where a 2-DoF arm has to reach a target denoted by a red dot, which we specify to be in the bottom left.
For observations, we directly use 64-by-64-by-3 images of the rendered environment, which provides a top-down view of the reacher and target, and we use a sliding window of four images.
Sawyer Lego block stacking.
To demonstrate a challenging task in the real world, we use our method to learn Lego block stacking with a real 7-DoF Sawyer robotic arm, as depicted in Figure 3b.
The observations used are raw 84-by-84-by-3 images from a camera pointed at the robot, and the controller only receives images as the observation, without joint angles or other information.
6. S IMULATION R ESULTS
Figure 4 details our results on the simulated image-based experimental domains, where each method is tested on three random seeds and the mean and standard deviation of the performance is reported.
For the 2D navigation and car tasks from images, we plot the average ﬁnal distance to the goal as a function of the number of episodes, so lower is better.
On the reacher task, we plot the reward function as deﬁned by Gym since this is the standard metric used to evaluate performance on this task, [Figure] requires over more perform worse than but need episodes learn whereas performs noticeably and as shown by the videos on our project website, achieving high Gym reward correlates strongly with solving the task in terms of distance to the goal.
On 2D navigation from images, our method, the VAE ablation, and the global model ablation are all able to learn very quickly, converging to high-performing policies within 200 episodes.
LQR-FLM struggles to learn the task, likely because the images are too complex for local linear model ﬁtting, and makes no progress at all.
In fact, LQR-FLM fails to learn on all of the simulated tasks, and we note that this precludes the guided policy search (GPS) method from solving these tasks (Levine et al., 2016), as GPS uses LQR-FLM as a subroutine.
For the sake of clarity in the plots, we omit the
LQR-FLM results, which are qualitatively similar to the E2C-like ablation results.
PPO eventually learns a successful policy, as indicated by the dashed line depicting this method’s ﬁnal performance, but this requires roughly three orders of magnitude more samples than our method.
We present log-scale plots that illustrate the full learning progress of model-free methods in Appendix F.
Despite using code directly from the authors of RCE, we were unable to get the E2C-like ablation to learn a good model for this task, and thus the learned policy does not improve over the initial policy.
In fact, we were unable to learn successful policies for any of the simulated tasks, though in Appendix F, we demonstrate that this ablation can learn a more successful policy on the 2D navigation domain used by Watter et al.
; Banijamali et al.
(2017), where the target is ﬁxed to the bottom right.
This highlights the difﬁculty of the tasks we consider.
On the image-based car, our method is able to learn a good policy with about 1500 episodes of experience.
The global model ablation is competitive with our method, however, we obtained this result by biasing the mean of the MPC random action selection to be positive, effectively encoding prior information that the car should move forward.
We also noticed that, even with more data, the variance of the MPC performance remained higher than the policy learned by our method.
These observations indicate that forward prediction using the learned global models may be inaccurate, leading to inconsistent control performance.
In contrast, our method does not heavily rely on an accurate model and can achieve consistently good behavior on this task.
The VAE ablation is able to solve this task for some random seeds, however this method’s performance is less consistent compared to our method.
PPO eventually learns a successful policy for this task that performs better than our method, however it uses over 25 times more data than our method.
Finally, on the image-based reacher task, our method achieves worse ﬁnal policy performance than PPO, though we do so with about 40 times fewer episodes, i.e., we use under 700 episodes whereas PPO uses about 30000.
This gain in data efﬁciency compared to model-free methods is typical of model-based methods, however, SOLAR is able to handle this domain directly from raw image observations, which is challenging for other model-based methods.
The VAE ablation also makes progress toward the goal, however, the performance is noticeably worse compared to our method.
The global model ablation makes very little improvement over its initial behavior, which is better than the other methods as it learns both a dynamics and cost model from the pretraining data and uses these models right away for planning.
This performance drop compared to the previous tasks indicates the difﬁculty in forward prediction for this domain, coupled with the failure of short-horizon control for this task as greedily minimizing distance to the goal often simply leads to collapsing the arm. As it is also less intuitive to encode prior information into this task compared to biasing the actions in the car domain to drive forward, we could not get this ablation to succeed on this task.
Figure 5 details performance on the Lego block stacking tasks in terms of the average ﬁnal distance in centimeters to the goal, where we test on three random seeds and report the mean and standard deviation of the performance.
We deﬁne the goal position of the end effector such that reaching the goal leads to successful stacking of the block
Not only is our method able to solve this task directly from raw, high-dimensional camera images within 200 episodes, corresponding to about half an hour of interaction time, our method is also successful at handling the complex, contact-rich dynamics of block stacking. As seen in the video on our project website, our method learns a policy that can react to slightly different contacts, due to the bottom block shifting between episodes, and is ultimately successful in stacking the block in most episodes.
We compare to the VAE and global model ablations, as these proved to be the most successful and data efﬁcient baselines in simulation.
These ablations are competitive with our method for this real world task, though our method still achieves a better ﬁnal policy that is able to more consistently stack the block.
The pixel space model is signiﬁcantly worse than the other methods that learn a latent representation, and we suspect that this method would need more data in order to learn this task.
D ISCUSSION AND F UTURE W ORK
We presented SOLAR, a model-based RL algorithm that is capable of learning policies in a dataefﬁcient manner directly from raw high-dimensional observations.
The key insights in SOLAR involve learning latent representations where simple models are more accurate and utilizing PGM structure to infer dynamics from data conditioned on entire real-world trajectories.
Our experimental results demonstrate that SOLAR is competitive in sample efﬁciency, while exhibiting superior ﬁnal policy performance, compared to other model-based methods.
Furthermore, SOLAR is signiﬁcantly more data-efﬁcient compared to state-of-the-art model-free RL methods.
There are several interesting directions for future work.
First, the ability to learn representations lends itself naturally to multi-task and transfer settings, where new tasks could potentially be learned much more quickly by starting from a latent embedding that has been learned from previous tasks.
We can also in principle share dynamics models, where the PGM we learn from solving previous tasks can be used as a global prior when inferring local dynamics ﬁts for a new task.
Second, our model is designed for and tested on continuous action domains as we focus on robotic applications. Extending our model to discrete actions would necessitate some type of continuous relaxation or learned action representation, and we believe that this is another interesting direction for future work.
A P OLICY L EARNING D ETAILS
Given a TVLG dynamics model and quadratic cost approximation, we can approximate our Q and value functions to second order with the following dynamic programming updates, which proceed from the last time step t =
T to the ﬁrst step t = 1 :
[Equation]
It can be shown (e
[Equation]
g., by Tassa et al.
that the action a t that minimizes the second-order approximation of the Q-function at every time step t is given by
[Equation]
This action is a linear function of the state s t , thus we can construct an optimal linear policy by setting
K t = − Q − 1 aa ,t Q as ,t and k t = − Q − 1 aa ,t Q a ,t .
We can also show that the maximum-entropy policy that minimizes the approximate Q-function is given by
[Equation]
Furthermore, as in Levine & Abbeel (2014)
[Equation]
we can impose a constraint on the total KL-divergence between the old and new trajectory distributions induced by the policies through an augmented cost function ¯ c ( s
[Equation] [Equation]
s t , a t ) − log π ( i − 1) ( a t | s t ) , where solving for λ via dual gradient descent can yield an exact solution to a KL-constrained LQR problem.
B M ODEL L EARNING D ETAILS
To derive the variational lower bound (ELBO) of the observed data for our model, we ﬁrst formulate our full variational posterior, given the variational factors presented in the paper, as
[Equation]
Where we use λ to denote the matrix normal-inverse-Wishart (MNIW) parameters of the variational factors on { F t , Σ t } t .
Thus, the ELBO is given by
[Equation]
Prior work has shown that, for conjugate exponential models such as the Bayesian LDS, the parameters can be updated using natural gradients, which can be computed in closed form using the variational message passing framework (Winn & Bishop, 2005).
Speciﬁcally, for the parameters of the MNIW λ , the natural gradient update is
[Equation]
Where B is the number of minibatches in the dataset, λ 0 is the parameter for the prior distribution p ( F , Σ) , and t F , Σ ( F , Σ) is the sufﬁcient 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, speciﬁcally 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 respective learning rates are treated as separate hyperparameters.
We have found 10 − 3 to be generally suitable for the natural gradient updates and 10 − 4 to be a good default value for Adam.
To perform our policy update, we ﬁt a quadratic cost function to observed data in the learned latent space as follows.
Given trajectories of the form [ s 0 , a 0 , c 0 , . . .
, s T − 1 , a T − 1 , c T − 1 , s T ] , we ﬁrst embed the observations { s t } using the mean of our recognition model µ ( e φ ( s )) to obtain a set of latent states { z t } .
We then model our cost samples as c t = 12 z (cid:62) t LL (cid:62) z t + c (cid:62) z t + α (cid:107) a t (cid:107)
+ 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 (cid:62) we guarantee that the learned cost matrix is positive deﬁnite, which improves the conditioning of the policy update.
C D YNAMICS I NFERENCE
Here we provide the closed form parameter computations for the posteriors of our dynamics given observed trajectories, as described in Section 4.1 of the main paper.
Given variational factors from our model of the form
[Equation]
We can condition on observed trajectories τ to obtain new variational posteriors {
[Equation]
t , Σ t |{ τ } Ni =0 ) } T − 1 t =0 .
These posteriors are also MNIW, and the parameters of these posteriors can be computed in closed form as
[Equation]
Then, a maximum a posteriori estimate gives us the TVLG dynamics parameters as described in the main paper.
D P OLICY L INEARIZATION
The policy update described in Section 4.2 of the main paper requires us to compute the KL-divergence between the trajectory distributions before and after the policy update, denoted as ¯ p ( τ ) and p ( τ ) , respectively.
We compute p ( τ ) =
ˆ ρ ( z 0 ) (cid:81) T − 1 t =
π θ ( a t | z t )ˆ p ( z t +1 | z t , a t ) , and analogously for ¯ p ( τ ) with the previous policy, and we are able to compute these analytically because the policies and dynamics model are TVLG, thus the induced trajectory distributions are also Gaussian.
However, this operates under the assumption that z is ﬁxed, which does not hold since the model update changes the latent representation.
Since our overall policy is a combination of the model embedding, given by e φ ( s ) , and the TVLG policy π θ ( a t | z t ) , training e φ ( s ) will change the behavior of the policy even if π θ ( a t | z t ) stays ﬁxed.
In some cases, this may lead to a policy with worse performance, and constraining against this policy for the policy update may lead to poor results.
In fact, what we want to do is to account for the model update by changing π θ ( a t | z t ) accordingly, so that the overall policy does not change in its distribution.
Thus, using ( s t , a t ) pairs from the previous data collection phase, we embed z t = µ ( e φ ( s t )) with our updated model and use linear regression to ﬁnd the TVLG policy ˜
π θ ( a t | z t ) that best explains the data collected from the policy This is line 6 of the SOLAR algorithm presented in the main paper, and after this, we can perform the policy update constrained against the trajectory distribution induced by ˜
π
θ ( a t | z t ) .
[Figure]
E E XPERIMENT S ETUP
Image-based 2D navigation.
Our recognition model architecture for the 2D navigation domain consists of two convolution layers with 2-by-2 ﬁlters and 32 channels each, with no pooling layers and ReLU non-linearities, followed by another convolution with 2-by-2 ﬁlters and 2 channels.
The output of the last convolution layer is fed into a spatial softmax layer (Finn et al., 2016), which then outputs a Gaussian distribution with a ﬁxed diagonal covariance of 10 − 4 for the latent distribution.
Our observation model consists of two fully-connected (FC) hidden layers with 256 ReLU activations, and the last layer outputs a categorical distribution over pixels.
We initially collect 200 episodes which we use to train our model, and for every subsequent iteration we collect 20 episodes to ﬁne tune our model.
The cost function we use is the sum of the L 2 -norm squared of the distance to the target and the commanded action, with weights of 1 and 0.001, respectively.
Image-based nonholonomic car.
The image-based car domain consists of 64-by-64 image observations.
We include a window of the 3 previous 64-by-64 images in our observation to preserve velocity information.
Our recognition model is a convolutional neural network that operates on each image in the sliding window independently.
Its architecture is four convolutional layers with 4-by-4 ﬁlters with 4 channels each, and the ﬁrst two convolution layers are followed by a ReLU non-linearity.
The output of the last convolutional layer is fed into three FC ReLU layers of width 2048, 512, and 128, respectively.
Our ﬁnal layer outputs a Gaussian distribution with dimension 8.
This leads to a ﬁnal latent dimension of 32.
Our observation model consists of four FC ReLU layers of width 256, 512, 1024, and 2048, respectively, followed by a Bernoulli distribution layer that models the image.
Like the recognition model, the observation model only operates on each section of the latent representation corresponding to the image window independently.
For this domain, we collectepisodes initially to train our model, and we collect 100 episodes per iteration after this.
The cost function we use is the sum of the L 2 -norm squared of the distance from the center of the car to the target and the commanded action, with weights of 1 and 0.001, respectively.
Reacher.
The reacher domain consists of 64-by-64-by-3 image observations.
Similar to the car, we include a window of the 3 previous 64-by-64-by-3 images in our observation.
Our recognition model is a convolutional neural network that again operates on each image in the sliding window independently.
Its architecture is three convolutional layers with 2-by-2 ﬁlters with 64, 32 andchannels respectively.
Each layer has a ReLU non-linearity followed by a 2-by-2 max-pooling.
The output of the last convolutional layer is fed into an FC ReLU layer of width 200, followed by another
FC ReLU layer of width 200.
Our ﬁnal layer outputs a Gaussian distribution with dimension 10, leading to a ﬁnal latent dimension of 40.
Our observation model consists of three FC ReLU layers of width 256, followed by a Bernoulli distribution layer and separately models each image in the sliding window.
We collect 200 episodes initially to train our model, and we collect 100 episodes per iteration after this.
The cost function we use is the sum of the L 2 -norm of the distance from the ﬁngertip to the target and the L 2 -norm squared of the commanded action.
Sawyer Lego block stacking.
The image-based Sawyer block-stacking domain consists of 84-by-84-by-3 image observations.
The policy outputs velocities on the end effector in order to control the robot.
Our recognition model is a convolutional neural network with the following architecture: a 5-by-5 ﬁlter convolutional layer with 16 channels followed by two convolutional layers using 5-by-5 ﬁlters with 32 channels each.
The ﬁrst two convolutional layers are followed by ReLU activations and the last by a FC ReLU layer of width 256 leading to a 16 dimensional Gaussian distribution layer.
Our observation model consists of a FC ReLU layer of width 128 feeding into three deconvolutional layers, the ﬁrst with 5-by-5 ﬁlters with 32 channels and the last two of 6-byﬁlters with 16 and 3 channels respectively.
These are followed by a ﬁnal Bernoulli distribution layer.
For this domain, we collect 50 episodes initially to train our model, 20 episodes per iteration for the ﬁrst 5 iterations, then 10 episodes per iteration for the remainder.
The cost function is the sum of the L 1 -norm of a weighted displacement vector between the end-effector and the target in 3D-space (weighted 1, 2, 1 for x , y , z ), the L 2 -norm in the same space, and the angle of rotation required to reach a valid wrist orientation, with weights of 1, .1, and .15, respectively.
F A DDITIONAL E XPERIMENTS F. E2C- LIKE ABLATION ON SIMPLIFIED 2D NAVIGATION
As mentioned in Section 6, our E2C-like ablation was unable to make progress for the 2D navigation task, though we were able to get more successful results by ﬁxing the position of the goal to the bottom right as is done in the image-based 2D navigation task considered in
E2C (Watter et al., 2015) and RCE (Banijamali et al., 2017).
Figure 7 details this experiment, which we ran for three random seeds and report the mean and standard deviation of the average ﬁnal distance to the goal as a function of the number of training episodes.
It is clear that the policy is improving, and two of the seeds are able to make substantial progress, though the ﬁnal seed is less successful and signiﬁcantly worsens the average performance of the method.
This indicates that the latent representation learned through RCE is less suitable for local model ﬁtting, as accurate local model ﬁtting is not explicitly encouraged by their representation learning objective.
F. M ODEL - BASED COMPARISONS ON STATE - BASED NONHOLONOMIC CAR
that we maintain the sample model based
To provide a point of comparison to modelbased RL methods, we consider the car domain where the underlying state is observed.
The states for the car domain include the position of the center of mass, orientation, forward and angular velocity of the car, and the position of the target, making for a 9-dimensional system.
Since this observation is already quite simple, we use a single linear layer for our recognition and observation models that output Gaussian distributions, and we use the same dimensionality for our latent representation as the state dimensionality.
We plot the performances of our method, LQRFLM (Levine & Abbeel, 2014), and Nagabandi et al.
(2018), which we refer to as model- [Figure] predictive control with neural networks (MPCNN), again based on the average ﬁnal distance to the target, in Figure 8.
In this setting, our method is competitive with LQR-FLM, learning a policy with similar performance in 200 episodes.
MPC-NN performs the best for this task, learning a policy that consistently reaches the target in just 20 episodes, though it is given the true cost function whereas our method and LQR-FLM are not.
For this simple setup where modeling bias is not an issue, we expect model-based methods to perform very well and learn efﬁciently.
However, when we make the problem more challenging by using image observations, model-based methods will fail quickly: LQR-FLM is unable to ﬁt complex pixel transitions using local linear models, as shown through the 2D navigation experiment, and MPC-NN has never been used with images, as forward video prediction and deﬁning a cost function on images are both very difﬁcult.
We extend MPC-NN to the image-based task, and we term this the “global model ablation” of our method – as shown in the paper, this approach is able to make progress toward the goal, though our method is still signiﬁcantly better at solving this difﬁcult task.
F. F ULL PERFORMANCE OF TRPO ON 2D NAVIGATION AND REACHER
In Figure 9 we include the plots for the simulated tasks comparing SOLAR, PPO, and TRPO.
Note that the x-axis is on a log scale, i.e., though our method is sometimes worse in ﬁnal policy performance to PPO and TRPO, we do so with one to three orders of magnitude fewer samples.
This demonstrates our method’s sample efﬁciency compared to model-free methods, while being able to solve complex image-based domains that are difﬁcult for model-based methods.