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

Introduction

Deep reinforcement learning (RL) algorithms can learn to solve tasks directly from raw, low-level observations such as images. However, such high-dimensional 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. Standard model-free deep RL aims to unify these challenges of representation learning and task learning into a single end-to-end training procedure. However, solving both problems together is difficult, since an effective policy requires an effective representation, and an effective representation requires meaningful gradient information to come from the policy or value function, while using only the model-free supervision signal (i.e., the reward function). As a result, learning directly from images with standard end-to-end RL algorithms can in practice be slow, sensitive to hyperparameters, and inefficient.

Instead, we propose to separate representation learning and task learning, by relying on predictive model learning to explicitly acquire a latent representation, and training the RL agent in that learned latent space. This alleviates the representation learning challenge because predictive learning benefits from a rich and informative supervision signal even before the agent has made any progress on the task, and thus results in improved sample efficiency of the overall learning process. In this work, our predictive model serves to accelerate task learning by separately addressing representation learning, in contrast to existing model-based RL approaches, which use predictive models either for generating cheap synthetic experience or for planning into the future .

Our proposed stochastic sequential model (Figure 1) models the high-dimensional observations as the consequence of a latent process, with a Gaussian prior and latent dynamics. This model represents a partially observed Markov decision process (POMDP), where the stochastic latent state enables the model to represent uncertainty about any of the state variables, given the past observations. Solving such a POMDP exactly would be computationally intractable, since it amounts to solving the decision problem in the space of beliefs . Recent works approximate the belief as encodings of latent samples from forward rollouts or particle filtering , or as learned belief representations in a belief-state forward model . We instead propose a simple approximation, which we derive from the control as inference framework, that trains a Markovian critic on latent state samples and trains an actor on a history of observations and actions, resulting in our stochastic latent actor-critic (SLAC) algorithm. Although this approximation loses some of the benefits of full POMDP solvers (e.g. reducing uncertainty), it is easy and stable to train in practice, achieving competitive results on a range of challenging problems.

The main contribution of this work is a novel and principled approach that integrates learning stochastic sequential models and RL into a single method, performing RL in the model’s learned latent space. By formalizing the problem as a control as inference problem within a POMDP, we show that variational inference leads to the objective of our SLAC algorithm. We empirically show that SLAC benefits from the good asymptotic performance of model-free RL while also leveraging the improved latent space representation for sample efficiency, by demonstrating that SLAC substantially outperforms both prior model-free and model-based RL algorithms on a range of image-based continuous control benchmark tasks.

Related Work

Representation learning in RL. End-to-end deep RL can in principle learn representations implicitly as part of the RL process . 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 . This motivates the use of a distinct representation learning procedure to acquire these representations before the agent has even learned to solve the task. A number of prior works have explored the use of auxiliary supervision in RL to learn such representations . 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 . 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 . Other works, however, learn latent-space dynamical system models and then use them to solve the POMDP with model-based RL . Although some of these works learn latent variable models that are similar to ours, these 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 , belief representations obtained from particle filtering , or belief representations obtained directly from a learned belief-space forward model . 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 state samples, which more tractably enables scaling to more complex tasks. Concurrent to our work, Hafner et al. proposed to integrate model-free learning with representations from sequence models, as proposed in this paper, with model-based rollouts, further improving on the performance of prior model-based approaches.

Sequential latent variable models. Several previous works have explored various modeling choices to learn stochastic sequential models . They vary in the factorization of the generative and inference models, their network architectures, and the objectives used in their training procedures. Our approach is compatible with any of these sequential latent variable models, with the only requirement being that they provide a mechanism to sample latent states from the belief of the learned Markovian latent space.

Preliminaries

This work addresses the problem of learning policies from high-dimensional observations in POMDPs, by simultaneously learning a latent representation of the underlying MDP state using variational inference, as well as learning a policy in a maximum entropy RL framework. In this section, we describe maximum entropy RL in fully observable MDPs, as well as variational methods for training latent state space models for POMDPs.

where $\gamma$ is the discount factor, and $\bar{\theta}$ are delayed parameters. The policy parameters $\phi$ are optimized to update the policy towards the exponential of the soft Q-function, resulting in the policy loss

SLAC builds on top of this maximum entropy RL framework, by further integrating explicit representation learning and handling partial observability.

2 Sequential Latent Variable Models and Amortized Variational Inference in POMDPs

We can maximize the probability of the observed datapoints (i.e., the left hand side of Equation (missing) 3) by learning an encoder $q(\bm{\mathbf{z}}|\bm{\mathbf{x}})$ and a decoder $p(\bm{\mathbf{x}}|\bm{\mathbf{z}})$, and then directly performing gradient ascent on the right hand side of the equation. In this setup, the distributions of interest are the prior $p(\bm{\mathbf{z}})$, the observation model $p(\bm{\mathbf{x}}|\bm{\mathbf{z}})$, and the variational approximate posterior $q(\bm{\mathbf{z}}|\bm{\mathbf{x}})$.

In order to extend such models to sequential decision making settings, we must incorporate actions and impose temporal structure on the latent state. Consider a partially observable MDP (POMDP), with latent states $\bm{\mathbf{z}}_{t}\in\mathcal{Z}$ and its corresponding observations $\bm{\mathbf{x}}_{t}\in\mathcal{X}$. We make an explicit distinction between an observation $\bm{\mathbf{x}}_{t}$ and the underlying latent state $\bm{\mathbf{z}}_{t}$, to emphasize that the latter is unobserved and its distribution is unknown. Analogous to the MDP, the initial and transition distributions are $p(\bm{\mathbf{z}}_{1})$ and $p(\bm{\mathbf{z}}_{t+1}|\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t})$, and the reward is $r_{t}$. In addition, the observation model is given by $p(\bm{\mathbf{x}}_{t}|\bm{\mathbf{z}}_{t})$.

As in the case for VAEs, a generative model of these observations $\bm{\mathbf{x}}_{t}$ can be learned by maximizing the log-likelihood. In the POMDP setting, however, we note that $\bm{\mathbf{x}}_{t}$ alone does not provide all necessary information to infer $\bm{\mathbf{z}}_{t}$, and prior observations must be taken into account during inference. This brings us to the discussion of sequential latent variable models. The distributions of interest are $p(\bm{\mathbf{z}}_{1})$ and $p(\bm{\mathbf{z}}_{t+1}|\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t})$, the observation model $p(\bm{\mathbf{x}}_{t}|\bm{\mathbf{z}}_{t})$, and the approximate variational posteriors $q(\bm{\mathbf{z}}_{1}|\bm{\mathbf{x}}_{1})$ and $q(\bm{\mathbf{z}}_{t+1}|\bm{\mathbf{x}}_{t+1},\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t})$. The log-likelihood of the observations can then be bounded,

For notational convenience, we define $q(\bm{\mathbf{z}}_{1}|\bm{\mathbf{x}}_{1},\bm{\mathbf{z}}_{0},\bm{\mathbf{a}}_{0})\coloneqq q(\bm{\mathbf{z}}_{1}|\bm{\mathbf{x}}_{1})$ and $p(\bm{\mathbf{z}}_{1}|\bm{\mathbf{z}}_{0},\bm{\mathbf{a}}_{0})\coloneqq p(\bm{\mathbf{z}}_{1})$. Prior work has explored modeling such non-Markovian observation sequences, using methods such as recurrent neural networks with deterministic hidden state, as well as probabilistic state-space models. In this work, we enable the effective training of a fully stochastic sequential latent variable model, and bring it together with a maximum entropy actor-critic RL algorithm to create SLAC: a sample-efficient and high-performing RL algorithm for learning policies for complex continuous control tasks directly from high-dimensional image inputs.

Joint Modeling and Control as Inference

For a fully observable MDP, the control problem can be embedded into a graphical model by introducing a binary random variable $\mathcal{O}_{t}$, which indicates if time step $t$ is optimal. When its distribution is chosen to be ${p(\mathcal{O}_{t}=1|\bm{\mathbf{s}}_{t},\bm{\mathbf{a}}_{t})=\exp(r(\bm{\mathbf{s}}_{t},\bm{\mathbf{a}}_{t}))}$, then maximization of $p(\mathcal{O}_{1:T})$ via approximate inference in that model yields the optimal policy for the maximum entropy objective .

In this paper, we extend this idea to the POMDP setting, where the probabilistic graphical model includes latent variables, as shown in Figure 1, and the distribution can analogously be given by $p(\mathcal{O}_{t}=1|\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t})=\exp(r(\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t}))$. Instead of maximizing the likelihood of the optimality variables alone, we jointly model the observations (including the observed rewards of the past time steps) and learn maximum entropy policies by maximizing the marginal likelihood $p(\bm{\mathbf{x}}_{1:\tau+1},\mathcal{O}_{\tau+1:T}|\bm{\mathbf{a}}_{1:\tau})$. This objective represents both the likelihood of the observed data from the past $\tau+1$ steps, as well as the optimality of the agent’s actions for future steps, effectively combining both representation learning and control into a single graphical model. We factorize our variational distribution into a product of recognition terms $q(\bm{\mathbf{z}}_{t+1}|\bm{\mathbf{x}}_{t+1},\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t})$, dynamics terms $p(\bm{\mathbf{z}}_{t+1}|\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t})$, and policy terms $\pi(\bm{\mathbf{a}}_{t}|\bm{\mathbf{x}}_{1:t},\bm{\mathbf{a}}_{1:t-1})$:

The variational distribution uses the dynamics for future time steps to prevent the agent from controlling the transitions and from choosing optimistic actions, analogously to the fully observed MDP setting described by Levine . The posterior over the actions represents the policy $\pi$.

We use the posterior from Equation (missing) 5 to obtain the evidence lower bound (ELBO) of the likelihood,

where $r(\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t})=\log p(\mathcal{O}_{t}=1|\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t})$ by construction and $p(\bm{\mathbf{a}}_{t})$ is the action prior. The full derivation of the ELBO is given in Appendix A.

Stochastic Latent Actor Critic

We now describe our stochastic latent actor critic (SLAC) algorithm, which maximizes the ELBO using function approximators to model the prior and posterior distributions. The ELBO objective in Equation (missing) 8 can be split into a model objective and a maximum entropy RL objective. The model objective can be optimized directly, while the maximum entropy RL objective can be optimized via approximate message passing, with messages corresponding to the Q-function. We can rewrite the RL objective to express it in terms of these messages, yielding an actor-critic algorithm analogous to SAC. Additional details of the derivation of the SLAC objectives are given in Appendix A.

Latent variable model. The first part of the ELBO corresponds to training the latent variable model to maximize the likelihood of the observations, analogous to the ELBO in Equation (missing) 4 for the sequential latent variable model. The generative model is given by $p_{\psi}(\bm{\mathbf{z}}_{1})$, $p_{\psi}(\bm{\mathbf{z}}_{t+1}|\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t})$, and $p_{\psi}(\bm{\mathbf{x}}_{t}|\bm{\mathbf{z}}_{t})$, and the inference model is given by $q_{\psi}(\bm{\mathbf{z}}_{1}|\bm{\mathbf{x}}_{1})$ and $q_{\psi}(\bm{\mathbf{z}}_{t+1}|\bm{\mathbf{x}}_{t+1},\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t})$. These distributions are diagonal Gaussian, where the means and variances are given by outputs of neural networks. Further details of our specific model architecture are given in Appendix B. The distribution parameters $\psi$ are optimized with respect to the ELBO in Equation (missing) 8, where the only terms that depend on $\psi$, and therefore constitute the model objective, are given by

where we define $q_{\psi}(\bm{\mathbf{z}}_{1}|\bm{\mathbf{x}}_{1},\bm{\mathbf{z}}_{0},\bm{\mathbf{a}}_{0})\coloneqq q_{\psi}(\bm{\mathbf{z}}_{1}|\bm{\mathbf{x}}_{1})$ and $p_{\psi}(\bm{\mathbf{z}}_{1}|\bm{\mathbf{z}}_{0},\bm{\mathbf{a}}_{0})\coloneqq p_{\psi}(\bm{\mathbf{z}}_{1})$. We use the reparameterization trick to sample from the filtering distribution $q_{\psi}(\bm{\mathbf{z}}_{1:\tau+1}|\bm{\mathbf{x}}_{1:\tau+1},\bm{\mathbf{a}}_{1:\tau})$.

Actor and critic. The second part of the ELBO corresponds to the maximum entropy RL objective. As in the fully observable case from Section 3.1 and as described by Levine , this optimization can be solved via message passing of soft Q-values. However, in our method, we must use the latent states $\bm{\mathbf{z}}$, since the true state is unknown. The messages are approximated by minimizing the soft Bellman residual, which we use to train our soft Q-function parameters $\theta$,

where $V_{\theta}$ is the soft state value function and $\bar{\theta}$ are delayed target network parameters, obtained as exponential moving averages of $\theta$. Notice that the latents $\bm{\mathbf{z}}_{\tau}$ and $\bm{\mathbf{z}}_{\tau+1}$, which are used in the Bellman backup, are sampled from the same filtering distribution, i.e. $\bm{\mathbf{z}}_{\tau+1}\sim q_{\psi}(\bm{\mathbf{z}}_{\tau+1}|\bm{\mathbf{x}}_{\tau+1},\bm{\mathbf{z}}_{\tau},\bm{\mathbf{a}}_{\tau})$. The RL objective, which corresponds to the second part of the ELBO, can then be rewritten in terms of the soft Q-function. The policy parameters $\phi$ are optimized to maximize this objective, resulting in a policy loss analogous to soft actor-critic :

We assume a uniform action prior, so $\log p(\bm{\mathbf{a}}_{t})$ is a constant term that we omit from the policy loss. This loss only uses the last sample $\bm{\mathbf{z}}_{\tau+1}$ of the sequence for the critic, and we use the reparameterization trick to sample from the policy. Note that the policy is not conditioned on the latent state, as this can lead to over-optimistic behavior since the algorithm would learn Q-values for policies that have perfect access to the latent state. Instead, the learned policy in our algorithm is conditioned directly on the past observations and actions. This has the additional benefit that the learned policy can be executed at run time without requiring inference of the latent state. Finally, we note that for the expectation over latent states in the Bellman residual in Equation (missing) 11, rather than sampling latent states for all $\bm{\mathbf{z}}\sim\mathcal{Z}$, we sample latent states from the filtering distribution $q_{\psi}(\bm{\mathbf{z}}_{1:\tau+1}|\bm{\mathbf{x}}_{1:\tau+1},\bm{\mathbf{a}}_{1:\tau})$. This design choice allows us to minimize the critic loss for samples that are most relevant for $Q_{\theta}$, while also allowing the critic loss to use the Q-function in the same way as implied by the policy loss in Equation (missing) 12.

SLAC is outlined in 1. The actor-critic component follows prior work, with automatic tuning of the temperature $\alpha$ and two Q-functions to mitigate overestimation . SLAC can be viewed as a variant of SAC where the critic is trained on the stochastic latent state of our sequential latent variable model. The backup for the critic is performed on a tuple $(\bm{\mathbf{z}}_{\tau},\bm{\mathbf{a}}_{\tau},r_{\tau},\bm{\mathbf{z}}_{\tau+1})$, sampled from the filtering distribution $q_{\psi}(\bm{\mathbf{z}}_{\tau+1},\bm{\mathbf{z}}_{\tau}|\bm{\mathbf{x}}_{1:\tau+1},\bm{\mathbf{a}}_{1:\tau})$. The critic can, in principle, take advantage of the perfect knowledge of the state $\bm{\mathbf{z}}_{t}$, which makes learning easier. However, the policy does not have access to $\bm{\mathbf{z}}_{t}$, and must make decisions based on a history of observations and actions. SLAC is not a model-based algorithm, in that in does not use the model for prediction, but we see in our experiments that SLAC can achieve similar sample efficiency as a model-based algorithm.

Experimental Evaluation

We evaluate SLAC on multiple image-based continuous control tasks from both the DeepMind Control Suite and OpenAI Gym , as illustrated in Figure 2. Full details of SLAC’s network architecture are described in Appendix B. Training and evaluation details are given in Appendix C, and image samples from our model for all tasks are shown in Appendix E. Additionally, visualizations of our results and code are available on the project website.https://alexlee-gk.github.io/slac/

To provide a comparative evaluation against prior methods, we evaluate SLAC on four tasks (cheetah run, walker walk, ball-in-cup catch, finger spin) from the DeepMind Control Suite , and four tasks (cheetah, walker, ant, hopper) from OpenAI Gym . Note that the Gym tasks are typically used with low-dimensional state observations, while we evaluate on them with raw image observations. We compare our method to the following state-of-the-art model-based and model-free algorithms:

SAC : This is an off-policy actor-critic algorithm, which represents a comparison to state-of-the-art model-free learning. We include experiments showing the performance of SAC based on true state (as an upper bound on performance) as well as directly from raw images.

D4PG : This is also an off-policy actor-critic algorithm, learning directly from raw images. The results reported in the plots below are the performance after $10^{8}$ training steps, as stated in the benchmarks from Tassa et al. .

MPO : This is an off-policy actor-critic algorithm that performs an expectation maximization form of policy iteration, learning directly from raw images.

DVRL : This is an on-policy model-free RL algorithm that trains a partially stochastic latent-variable POMDP model. DVRL uses the full belief over the latent state as input into both the actor and critic, as opposed to our method, which trains the critic with the latent state and the actor with a history of actions and observations.

PlaNet : This is a model-based RL method for learning from images, which uses a partially stochastic sequential latent variable model, but without explicit policy learning. Instead, the model is used for planning with model predictive control (MPC), where each plan is optimized with the cross entropy method (CEM).

DrQ : This is the same as the SAC algorithm, but combined with data augmentation on the image inputs.

Our experiments on the DeepMind Control Suite in Figure 3 show that the sample efficiency of SLAC is comparable or better than both model-based and model-free alternatives. This indicates that overcoming the representation learning bottleneck, coupled with efficient off-policy RL, provides for fast learning similar to model-based methods, while attaining final performance comparable to fully model-free techniques that learn from state. SLAC also substantially outperforms DVRL. This difference can be explained in part by the use of an efficient off-policy RL algorithm, which can better take advantage of the learned representation. SLAC achieves comparable or slightly better performance than subsequent work DrQ, which also uses the efficient off-policy SAC algorithm.

We also evaluate SLAC on continuous control benchmark tasks from OpenAI Gym in Figure 4. We notice that these tasks are more challenging than the DeepMind Control Suite tasks, because the rewards are not as shaped and not bounded between 0 and 1, the dynamics are different, and the episodes terminate on failure (e.g., when the hopper or walker falls over). PlaNet is unable to solve the last three tasks, while for the cheetah task, it learns a suboptimal policy that involves flipping the cheetah over and pushing forward while on its back. To better understand the performance of fixed-horizon MPC on these tasks, we also evaluated with the ground truth dynamics (i.e., the true simulator), and found that even in this case, MPC did not achieve good final performance, suggesting that infinite horizon policy optimization, of the sort performed by SLAC and model-free algorithms, is important to attain good results on these tasks.

Our experiments show that SLAC successfully learns complex continuous control benchmark tasks from raw image inputs. On the DeepMind Control Suite, SLAC exceeds the performance of prior work PlaNet on the four tasks, and SLAC achieves comparable or slightly better performance than subsequence work DrQ. However, on the harder image-based OpenAI Gym tasks, SLAC outperforms PlaNet by a large margin. We note that the prior methods that we tested generally performed poorly on the image-based OpenAI Gym tasks, despite considerable hyperparameter tuning.

2 Ablation Experiments

We investigate how SLAC is affected by the choice of latent variable model, the inputs given to the actor and critic, the model pretraining, and the number of training updates relative to the number of agent interactions. Additional results are given in Appendix D, including experiments that compare the effect of the decoder output variance and using random cropping for data augmentation.

Latent variable model. We study the tradeoffs between different design choices for the latent variable model in 5(a) and Appendix C. We compare our fully stochastic model to a standard non-sequential VAE model , which has been used in multiple prior works for representation learning in RL , and a non-sequential factored VAE model, which uses our autoregressive two-variable factorization but without any temporal dependencies. We also compare to a sequential filtering model that uses temporal dependencies but without the two-variable factorization, the partially stochastic model used by PlaNet , as well as two additional variants of our model: a fully deterministic model that removes all stochasticity from the hidden state dynamics, and a partially stochastic model that adds deterministic paths in the transitions, similar to the PlaNet model, but with our latent factorization and architecture. All the models, except for the PlaNet model, are variants of our model that use the same architecture as our fully stochastic model, with minimal differences in the transitions or the latent variable factorization. In all cases, we use the RL framework of SLAC and only vary the choice of model for representation learning.

Our fully stochastic model outperforms all the other models. Contrary to the conclusions in prior work , the fully stochastic model slightly outperforms the partially stochastic model, while retaining the appealing interpretation of a stochastic state space model. We hypothesize that these prior works benefit from the deterministic paths (realized as an LSTM or GRU) because they use multi-step samples from the prior. In contrast, our method uses samples from the posterior, which are conditioned on same-step observations, and thus it is less sensitive to the propagation of the latent states through time. The sequential variants of our model (including ours) outperform the non-sequential VAE models. The models with the two-variable factorization perform similarly or better than their respective equivalents among the non-sequential VAE models and among the sequential stochastic models. Overall, including temporal dependencies results in the largest improvement in performance, followed by the autoregressive latent variable factorization and using a fully stochastic model.

Actor and critic inputs. We next investigate alternative choices for the actor and critic inputs as either the observation-action history or the latent sample. In SLAC, the actor is conditioned on the observation-action history and the critic is conditioned on individual latent samples. The images in the history are first compressed with the model’s convolutional network before they are given to the networks. However, the actor and critic losses do not propagate any gradient signal into the model nor its convolutional layers, i.e. the convolutional layers used for the observation-action history are only trained by the model loss.

5(b) and Appendix C show that, in general, the performance is significantly worse when the critic input is the history instead of the latent sample, and indifferent to the choice for the actor input. This is consistent with our derivation—the critic should be given latent samples, but the actor can be conditioned on anything (since the policy is the variational posterior). However, we note that a latent-conditioned actor could lead to overconfident behaviors in uncertain environments. For generality, we choose to give the raw history directly to the actor.

Model pretraining. We next study the effect of pretraining the model before the agent starts learning on the task. In our experiments, the agent first collects a small amount of data by executing random actions, and then the model is pretrained with that data. The model is pretrained for 50000 iterations on the DeepMind Control Suite experiments, unless otherwise specified. 5(c) and Appendix C show that little or no pretraining results in slower learning and, in some cases, worse asymptotic performance. There is almost no difference in performance when using 100000 instead of 50000 iterations, although the former resulted in higher variance across trials in some of the tasks. Overall, these results show that the agent benefits from the supervision signal of the model even before the agent has made any progress on the task.

Training updates per iteration. We next investigate the effect of the number of training updates per iteration, or equivalently, the number of training updates per environment step (we use 1 environment step per iteration in all of our experiments). 5(d) and Appendix C show that, in general, more training updates per iteration speeds up learning slightly, but too many updates per iteration causes higher variance across trials and slightly worse asymptotic performance in some tasks. Nevertheless, this drop in asymptotic performance (if any) is small, which indicates that our method is less susceptible to overfitting compared to methods in prior work. We hypothesize that using stochastic latent samples to train the critic provides some randomization, which limits overfitting. The best tradeoff is achieved when using 2 training updates per iteration, however, in line with other works, we use 1 training update per iteration in all the other experiments.

Conclusion

We presented SLAC, an efficient RL algorithm for learning from high-dimensional image inputs that combines efficient off-policy model-free RL with representation learning via a sequential stochastic state space model. Through representation learning in conjunction with effective task learning in the learned latent space, our method achieves improved sample efficiency and final task performance as compared to both prior model-based and model-free RL methods.

While our current SLAC algorithm is fully model-free, in that predictions from the model are not utilized to speed up training, a natural extension of our approach would be to use the model predictions themselves to generate synthetic samples. Incorporating this additional synthetic model-based data into a mixed model-based and model-free method could further improve sample efficiency and performance. More broadly, the use of explicit representation learning with RL has the potential to not only accelerate training time and increase the complexity of achievable tasks, but also enable reuse and transfer of our learned representation across tasks.

Broader Impact

Despite the existence of automated robotic systems in controlled environments such as factories or labs, standard approaches to controlling systems still require precise and expensive sensor setups to monitor the relevant details of interest in the environment, such as the joint positions of a robot or pose information of all objects in the area. To instead be able to learn directly from the more ubiquitous and rich modality of vision would greatly advance the current state of our learning systems. Not only would this ability to learn directly from images preclude expensive real-world setups, but it would also remove the expensive need for human-engineering efforts in state estimation. While it would indeed be very beneficial for our learning systems to be able to learn directly from raw image observations, this introduces algorithm challenges of dealing with high-dimensional as well as partially observable inputs. In this paper, we study the use of explicitly learning latent representations to assist model-free reinforcement learning directly from raw, high-dimensional images.

Standard end-to-end RL methods try to solve both representation learning and task learning together, and in practice, this leads to brittle solutions which are sensitive to hyperparameters but are also slow and inefficient. These challenges illustrate the predominant use of simulation in the deep RL community; we hope that with more efficient, stable, easy-to-use, and easy-to-train deep RL algorithms such as the one we propose in this work, we can help the field of deep RL to transition to more widespread use in real-world setups such as robotics.

From a broader perspective, there are numerous use cases and areas of application where autonomous decision making agents can have positive effects in our society, from automating dangerous and undesirable tasks, to accelerating automation and economic efficiency of society. That being said, however, automated decision making systems do introduce safety concerns, further exacerbated by the lack of explainability when they do make mistakes. Although this work does not explicitly address safety concerns, we feel that it can be used in conjunction with levels of safety controllers to minimize negative impacts, while drawing on its powerful deep reinforcement learning roots to enable automated and robust tasks in the real world.

Acknowledgments and Disclosure of Funding

We thank Marvin Zhang, Abhishek Gupta, and Chelsea Finn for useful discussions and feedback, Danijar Hafner for providing timely assistance with PlaNet, and Maximilian Igl for providing timely assistance with DVRL. This research was supported by the National Science Foundation through IIS-1651843 and IIS-1700697, as well as ARL DCIST CRA W911NF-17-2-0181 and the Office of Naval Research. Compute support was provided by NVIDIA.

References

Appendix A Derivation of the Evidence Lower Bound and SLAC Objectives

In this appendix, we discuss how the SLAC objectives can be derived from applying a variational inference scheme to the control as inference framework for reinforcement learning . In this framework, the problem of finding the optimal policy is cast as an inference problem, conditioned on the evidence that the agent is behaving optimally. While Levine derives this in the fully observed case, we present a derivation in the POMDP setting. For reference, we reproduce the probabilistic graphical model in Figure 6.

We aim to maximize the marginal likelihood $p(\bm{\mathbf{x}}_{1:\tau+1},\mathcal{O}_{\tau+1:T}|\bm{\mathbf{a}}_{1:\tau})$, where $\tau$ is the number of steps that the agent has already taken. This likelihood reflects that the agent cannot modify the past $\tau$ actions and they might have not been optimal, but it can choose the future actions up to the end of the episode, such that the chosen future actions are optimal. Notice that unlike the standard control as inference framework, in this work we not only maximize the likelihood of the optimality variables but also the likelihood of the observations, which provides additional supervision for the latent representation. This does not come up in the MDP setting since the state representation is fixed and learning a dynamics model of the state would not change the model-free equations derived from the maximum entropy RL objective.

For reference, we restate the factorization of our variational distribution:

As discussed by Levine , the agent does not have control over the stochastic dynamics, so we use the dynamics $p(\bm{\mathbf{z}}_{t+1}|\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t})$ for $t\geq\tau+1$ in the variational distribution in order to prevent the agent from choosing optimistic actions.

We use the posterior from Equation (missing) 13, the likelihood from Equation (missing) 14, and Jensen’s inequality to obtain the ELBO of the marginal likelihood,

We are interested in the likelihood of optimal trajectories, so we use $\mathcal{O}_{t}=1$ for $t\geq\tau+1$, and its distribution is given by $p(\mathcal{O}_{t}=1|\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t})=\exp(r(\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t}))$ in the control as inference framework. Notice that the dynamics terms $\log p(\bm{\mathbf{z}}_{t+1}|\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t})$ for $t\geq\tau+1$ from the posterior and the prior cancel each other out in the ELBO.

The first part of the ELBO corresponds to the model objective. When using the parametric function approximators, the negative of it corresponds directly to the model loss in Equation (missing) 9.

The second part of the ELBO corresponds to the maximum entropy RL objective. We assume a uniform action prior, so the $\log p(\bm{\mathbf{a}}_{t})$ term is a constant term that can be omitted when optimizing this objective. We use message passing to optimize this objective, with messages defined as

Then, the maximum entropy RL objective can be expressed in terms of the messages as

That is, the optimal policy is optimal with respect to the expectation over the belief of the Q value of the learned MDP. This is equivalent to the Q-MDP heuristic, which amounts to assuming that any uncertainty in the belief is gone after the next action .

Noting that the KL divergence term is zero for the optimal action, the equality from Equation (missing) 22 and Equation (missing) 24 can be used in Equation (missing) 20 to obtain

This equation corresponds to the Bellman backup with a soft maximization for the value function.

As mentioned in Section 5, our algorithm conditions the parametric policy in the history of observations and actions, which allows us to directly execute the policy without having to perform inference on the latent state at run time. When using the parametric function approximators, the negative of the maximum entropy RL objective, written as in Equation (missing) 22, corresponds to the policy loss in Equation (missing) 12. Lastly, the Bellman backup of Equation (missing) 26 corresponds to the Bellman residual in Equation (missing) 11 when approximated by a regression objective.

We showed that the SLAC objectives can be derived from applying variational inference in the control as inference framework in the POMDP setting. This leads to the joint likelihood of the past observations and future optimality variables, which we aim to optimize by maximizing the ELBO of the log-likelihood. We decompose the ELBO into the model objective and the maximum entropy RL objective. We express the latter in terms of messages of Q-functions, which in turn are learned by minimizing the Bellman residual. These objectives lead to the model, policy, and critic losses.

Appendix B Latent Variable Factorization and Network Architectures

In this section, we describe the architecture of our sequential latent variable model. Motivated by the recent success of autoregressive latent variables in VAEs , we factorize the latent variable $\bm{\mathbf{z}}_{t}$ into two stochastic variables, $\bm{\mathbf{z}}_{t}^{1}$ and $\bm{\mathbf{z}}_{t}^{2}$, as shown in Figure 7. This factorization results in latent distributions that are more expressive, and it allows for some parts of the prior and posterior distributions to be shared. We found this design to provide a good balance between ease of training and expressivity, producing good reconstructions and generations and, crucially, providing good representations for reinforcement learning. Note that the diagram in Figure 7 represents the Bayes net corresponding to our full model. However, since all of the latent variables are stochastic, this visualization also presents the design of the computation graph. Inference over the latent variables is performed using amortized variational inference, with all training done via reparameterization. Hence, the computation graph can be deduced from the diagram by treating all solid arrows as part of the generative model and all dashed arrows as part of approximate posterior.

The generative model consists of the following probability distributions:

The initial distribution $p(\bm{\mathbf{z}}_{1}^{1})$ is a multivariate standard normal distribution $\mathcal{N}(\bm{\mathbf{0}},\bm{I})$. All of the other distributions are conditional and parameterized by neural networks with parameters $\psi{}$. The networks for $p_{\psi}(\bm{\mathbf{z}}_{1}^{2}|\bm{\mathbf{z}}_{1}^{1})$, $p_{\psi}(\bm{\mathbf{z}}_{t+1}^{1}|\bm{\mathbf{z}}_{t}^{2},\bm{\mathbf{a}}_{t})$, $p_{\psi}(\bm{\mathbf{z}}_{t+1}^{2}|\bm{\mathbf{z}}_{t+1}^{1},\bm{\mathbf{z}}_{t}^{2},\bm{\mathbf{a}}_{t})$, and $p_{\psi}(r_{t}|\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t},\bm{\mathbf{z}}_{t+1})$ consist of two fully connected layers, each with 256 hidden units, and a Gaussian output layer. The Gaussian layer is defined such that it outputs a multivariate normal distribution with diagonal variance, where the mean is the output of a linear layer and the diagonal standard deviation is the output of a fully connected layer with softplus non-linearity. The pre-transformed standard deviation right before the softplus non-linearity is gradient clipped element-wise by value to within $$during the backward pass. The observation model$p_{\psi}(\bm{\mathbf{x}}_{t}|\bm{\mathbf{z}}_{t})$consists of 5 transposed convolutional layers (256$4\times 4$, 128$3\times 3$, 64$3\times 3$, 32$3\times 3$, and 3$5\times 5$filters, respectively, stride 2 each, except for the first layer). The output variance for each image pixel is fixed to a constant$\sigma^{2}$, which is a hyperparameter$\sigma^{2}\in\{0.04,0.1,0.4\}$on DeepMind Control Suite and$\sigma^{2}=0.1$ on OpenAI Gym.

The variational distribution $q$, also referred to as the inference model or the posterior, is represented by the following factorization:

The networks representing the distributions $q_{\psi}(\bm{\mathbf{z}}_{1}^{1}|\bm{\mathbf{x}}_{1})$ and $q_{\psi}(\bm{\mathbf{z}}_{t+1}^{1}|\bm{\mathbf{x}}_{t+1},\bm{\mathbf{z}}_{t}^{2},\bm{\mathbf{a}}_{t})$ both consist of 5 convolutional layers (32 $5\times 5$, 64 $3\times 3$, 128 $3\times 3$, 256 $3\times 3$, and 256 $4\times 4$ filters, respectively, stride 2 each, except for the last layer), 2 fully connected layers (256 units each), and a Gaussian output layer. The parameters of the convolution layers are shared among both distributions.

Note that the variational distribution over $\bm{\mathbf{z}}_{1}^{2}$ and $\bm{\mathbf{z}}_{t+1}^{2}$ is intentionally chosen to exactly match the generative model $p$, such that this term does not appear in the KL-divergence within the ELBO, and a separate variational distribution is only learned over $\bm{\mathbf{z}}_{1}^{1}$ and $\bm{\mathbf{z}}_{t+1}^{1}$. In particular, the KL-divergence over $\bm{\mathbf{z}}_{t+1}$ simplifies to the KL-divergence over $\bm{\mathbf{z}}_{t+1}^{1}$:

This intentional design decision simplifies the training process.

For the reinforcement learning process, we use a critic network $Q_{\theta}{}$ consisting of 2 fully connected layers (256 units each) and a linear output layer. The actor network $\pi_{\phi}{}$ consists of 5 convolutional layers, 2 fully connected layers (256 units each), a Gaussian layer, and a tanh bijector, which constrains the actions to be in the bounded action space of $$. The convolutional layers are shared with the ones from the latent variable model, but the parameters of these layers are only updated by the model objective and not by the actor objective.

Appendix C Training and Evaluation Details

The control portion of our algorithm uses the same hyperparameters as SAC , except for a smaller replay buffer size of 100000 environment steps (instead of a million) due to the high memory usage of image observations.

The network parameters are initialized using the default initialization distributions. In the case of the DeepMind Control Suite experiments, the scale of the policy’s pre-transformed Gaussian distribution is scaled by 2. This, as well as the initial tanh-transformed Gaussian policy, contributes to trajectories with larger actions (i.e. closer to $-1$ and $1$) at the beginning of training. This didn’t make a difference for the DeepMind Control Suite tasks except for the walker task, where we observed that this initialization resulted in less variance across trials and avoided trials that would otherwise get stuck in local optima early in training.

All of the parameters are trained with the Adam optimizer , and we perform 1 gradient step per environment step for DeepMind Control Suite and 3 gradient steps per environment step for OpenAI Gym. The Q-function and policy parameters are trained with a learning rate of 0.0003 and a batch size of 256. The model parameters are trained with a learning rate of 0.0001 and a batch size of 32. We use fixed-length sequences of length 8, rather than all the past observations and actions within the episode.

We use action repeats for all the methods, except for D4PG for which we use the reported results from prior work . The number of environment steps reported in our plots correspond to the unmodified steps of the benchmarks. Note that the methods that use action repeats only use a fraction of the environment steps reported in our plots. For example, 1 million environment steps of the cheetah task correspond to 250000 samples when using an action repeat of 4. The action repeats used in our experiments are given in Appendix C.

Unlike in prior work , we use the same stochastic policy as both the behavioral and evaluation policy since we found the deterministic greedy policy to be comparable or worse than the stochastic policy.

Our plots show results over multiple trials (i.e. seeds), and each trial computes average returns from 10 evaluation episodes. We used 10 trials for the DeepMind Control Suite experiments and 5 trials for the OpenAI Gym experiments. In the case of the DeepMind Control Suite experiments, we sweep over $\sigma^{2}\in\{0.04,0.1,0.4\}$ and plot the results corresponding to the hyperparameter $\sigma^{2}$ that achieves the best per-task average return across trials averaged over the first half a million environment steps. In Figure 3, the best $\sigma^{2}$ values are 0.1, 0.4, 0.04, and 0.1 for the cheetah run, walker walk, ball-in-cup catch, and finger spin tasks, respectively.

Appendix E Predictions from the Latent Variable Model

We show example image samples from our learned sequential latent variable model in Figure 14 and Figure 15. Samples from the posterior show the images $\bm{\mathbf{x}}_{t}$ as constructed by the decoder $p_{\psi}(\bm{\mathbf{x}}_{t}|\bm{\mathbf{z}}_{t})$, using a sequence of latents $\bm{\mathbf{z}}_{t}$ that are encoded and sampled from the posteriors, $q_{\psi}(\bm{\mathbf{z}}_{1}|\bm{\mathbf{x}}_{1})$ and $q_{\psi}(\bm{\mathbf{z}}_{t+1}|\bm{\mathbf{x}}_{t+1},\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t})$. Samples from the prior, on the other hand, use a sequence of latents where $\bm{\mathbf{z}}_{1}$ is sampled from $p_{\psi}(\bm{\mathbf{z}}_{1})$ and all remaining latents $\bm{\mathbf{z}}_{t}$ are from the propagation of the previous latent state through the latent dynamics $p_{\psi}(\bm{\mathbf{z}}_{t+1}|\bm{\mathbf{z}}_{t},\bm{\mathbf{a}}_{t})$. Note that these prior samples do not use any image frames as inputs, and thus they do not correspond to any ground truth sequence. We also show samples from the conditional prior, which is conditioned on the first image from the true sequence: for this, the sampling procedure is the same as the prior, except that $\bm{\mathbf{z}}_{1}$ is encoded and sampled from the posterior $q_{\psi}(\bm{\mathbf{z}}_{1}|\bm{\mathbf{x}}_{1})$, rather than being sampled from $p_{\psi}(\bm{\mathbf{z}}_{1})$. We notice that the generated images samples can be sharper and more realistic by using a smaller variance for $p_{\psi}(\bm{\mathbf{x}}_{t}|\bm{\mathbf{z}}_{t})$ when training the model, but at the expense of a representation that leads to lower returns. Finally, note that we do not actually use the samples from the prior for training.