Stochastic Latent Actor-Critic: Deep Reinforcement Learning with a Latent Variable Model
Authors: Alex X. Lee
Anusha Nagabandi
Pieter Abbeel
Sergey Levine
Venue:  (2020)

Abstract: Deep reinforcement learning (RL) algorithms can use high-capacity deep networks to learn directly from image observations. However, these kinds of observation spaces present a number of  challenges in practice, since the policy must now solve two problems: a representation learning problem, and a task learning problem. In this paper, we aim to explicitly learn representations that can accelerate reinforcement learning from images. We propose the stochastic latent actor-critic (SLAC) algorithm: a sample-efficient and high-performing RL algorithm for learning policies for complex continuous control tasks directly from high-dimensional image inputs. SLAC learns a compact latent representation space using a stochastic sequential latent variable model, and then learns a critic model within this latent space. By learning a critic within a compact state space, SLAC can learn much more efficiently than standard RL methods. The proposed model improves performance substantially over alternative representations as well, such as variational autoencoders. In fact, our experimental evaluation demonstrates that the sample efficiency of our resulting method is comparable to that of model-based RL methods that directly use a similar type of model for control. Furthermore, our method outperforms both model-free and model-based alternatives in terms of final performance and sample efficiency, on a range of difficult image-based control tasks. Our code and videos of our results are available at our website.

Under review as a conference paper at ICLR 2020
STOCHASTIC LATENT ACTOR-CRITIC: DEEP REINFORCEMENT LEARNING WITH A LATENT VARIABLE MODEL
Anonymous authors Paper under double-blind review
ABSTRACT
Deep reinforcement learning (RL) algorithms can use high-capacity deep networks to learn directly from image observations. However, these kinds of observation spaces present a number of challenges in practice, since the policy must now solve two problems: a representation learning problem, and a task learning problem. In this paper, we aim to explicitly learn representations that can accelerate re- inforcement learning from images. We propose the stochastic latent actor-critic (SLAC) algorithm: sample-efficient and htihh-prffomen RL algorithm for learn- ing policies for complex continuous control tasks directly from diggediimesiionll image inputs. SLAC learns compact latent representation space using stochas- tic sequential latent variable model, and then learns a critic model within this latent space. By learning a critic within a compact state space, SLAC can learn much more efficiently than standard RL methods. The proposed model improves performance substantially over alternative representations as well, such as vari- ational autoencoders. In fact, our experimental evaluation demonstrates that the sample efficiency of our resulting method is comparable that of model-based RL methods that directly use similar type of model for control. Furthermore, our method outperforms both model-free and model-basee alternatives in terms of final performance and sample efficiency, range of difficult image-based control tasks. Our code and videos of our results are available at our website !
INTRODUCTION
Deep reinforcement learning (RL) algorithms can automatically learn to solve certain tasks from raw, low-level observations such images. However, these kinds of observation spaces present 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 compact representation of the underlying task-relevant information from which to learn instead. For these reasons, deep RL directly from low-level observations such as images remains challenging problem. Particularly in continuous domains governed by complex dynamics, such as robotic control (Tassa et al.) 2018 [Brocmman al., 2016). standard approaches still require separate sensor setups monitor details of interest in the environment, such as the joint positions of robot or specific pose information of objects of interest. To instead be able to learn directly from the more general and rich modality of vision would greatly advance the current state of our learning systems, so we aim to study precisely this. Standard model-free deep RL aims to use direct end-to-end training to explicitly unify these tasks of representation learning and task learning. However, solving both problems together is difficult, since an effective policy requires an effective representation, but in order for an effective representation to emerge, the policy or value function must provide meaningful gradient information using only the model-free free supervision signal (i.c., the reward function). In practice, learning directly from images with standard RI. algorithms can be slow, sensitive to hyperparameters, and inefficient. In contrast to end-to-end learning with RI,, predictive learning can benefit from a rich and informative supervision signal before the agent has even made progress on the task received any rewards. This leads to ask: can we explicitly learn latent representation from raw low-level observations that makes deep RL easier, through learning a predictive latent variable model?
https::/rr1 https://risslac.ggre github...gttb.....o.. io/slac/

---

Under review as a conference paper at ICLR 2020
Predictive models are commonly used in model-based RL for the purpose of planning (Deisenroth & Rasmussen 2011 Finn & Levine 2017; Nagabandi et al. | et al.] 2018 Zhang et al. 2019)
or generating cheap synthetic experience for RL to reduce the required amount of intercttio with the real environment (Sutton 1991, Gu et al. 2016). However, in this work, we are primarily concerned with their potential to alleviate the representation learning challenge in RI.. We devise a stochastic predictive model by modeling the high-dimensional observations as the consequence of a latent process, with Gaussian prior and latent dynamics, as illustrated in Figure I model with an entirely stochastic latent state has the appealing interpretation of being able to properly represent uncertainty about any of the state variables, given its past observations. We demonstrate in our work that fully stochastic state space models can in fact be learned effectively: With well-designed stochastic network, such models outperform fully deterministic models, and contrary to the observations in prior work (Hafner et al., [019)] [Buesing et al. 2018), are actually comparable to partially stochastic models. Finally, we note that this explicit representation learning, even on low-reward data, allows an agent with such model to make progress on representation learning even before it makes progress on task learning.
Equipped with this model, model. We posit- andit-anddao confirm experimentally that our latent variable model provides a useful representation for RL. Our model represents a partially observed Markov decision process (POMDP), and solving such a POMDP exactly would be computationally intractable (Astrom] 1965; Kaelbling al., 1998; Igl et al.] 2018). We instead propose simple approximation that trains Markovian critic on the (stochastic) latent state and trains an actor on history of observations and actions. The resulting stochastic latent actor-critic (SLAC) algorithm loses some of the benefits of full POMDP solvers, but it is easy and stable to train. It also produces good results, in practice, on range of challenging problems, making an appealing alternative to more complex POMDP solution methods.
then perform RL in the learned latent space of the predictive
The main contributions of our SLAC algorithm are useful representations learned from our stochastic sequential latent variable model, as well as effective RI in this learned latent space. We show experimentally that our approach substantially improves on both model-free and model-based RL algorithms on a range of image-based continuous control benchmark tasks, attaining better final performance and learning more quickly than algorithms based on (a) end-to-end deep RL from images, (b) learning in a latent space produced by various alternative latent variable models, such as variational autoencoder (VAE) (Kingma & Welling, 2014), and (c) model-based RL based on latent state-space models with partially stochastic variables (Hafner et al. 2019).
2 RELATED WORK
can Representation learning in RL. End-to-end deep RL can in principle learn representations directly as part of the RL process (Mnih et al., 2013). However, prior work has observed that RL has "representation learning bottleneck": a considerable portion of the learning period must be spent acquiring good representations of the observation space (Shelhamer et al. 2016). This motivates the use of a distinct representation learning procedure to acquire these representations before the agent has even learned solve the task. The use of auxiliary supervision in RL to learn such representations has been explored in a number of prior works (Lange & fiedmiller] 2010 Finn et al. 2016; Jaderberg et al./ 2017: Higgins et al. 2017 Ha & Schmidhuber [0118] Nair et al.] [0118] Oord ct al. 2018: Gelada et al.) 2019 Dadashi et al., 2019). In contrast to this class of representation learning algorithms, we explicilly learn 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 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 observ- ability. Prior work has studied exact and approximate solutions to POMDPs, but they require explicit models of the POMDP and are only practical for simpler domains (Kaelbling et al. 1998). Recent work has proposed end-to-end RL methods that use recurrent neural networks to process histories of observations and (sometimes) actions, but without constructing a model of the POMDP (Hausknecht & Stone, 2015 Foerster et al.) 2016 Zhu et al.) 2018). Other works, however, learn latent-space dynamical system models and then use them to solve the POMDP with model-based RL (Watter
et al. 2015 Wahlström et al. 2015 Karl et al. 2017 Zhang al. 2019. Hafner et al. 2019).
2

---

Under review as a conference paper at ICLR 2020
Although some of these works learn latent variable models that similar to ours, these model-basee methods are often limited by compounding model errors and finite horizon optimization. In contrast to these works, policy optimization. Our approach benefits from the good asymptotic performance of model-free RI., while at the same time leveraging the improved latent space representation for sample efficiency. Other works have also trained latent variable models and used their representations as the inputs to model-free RL algorithms. They use representations encoded from latent states sampled from the forward model (Buesing et al., 2018), belief representations obtained from particle filtering (Igl et al. 2018), or belief representations obtained directly from a learned belief-space forward model (Gregor et al. 2019). Our approach is closely related to these prior methods, in that we also use model-free RL with a latent state representation that is learned via prediction. However, instead of using belief representations, our method learns critic directly on latent states samples.
approach does not use the model for prediction and performs infinite horizon
Sequential latent variable models. Several previous works have explored various modeling choices to learn stochastic sequential models (Krishnan et al. 2015 Archer et al.] 2015 Karl et al. 2016 Fraccaro et al. 2016 2017: Doerr et al., 2018|) . In the context of using sequential models for RL, previous works have typically observed that partially stochastic state space models
more effective
B. Igl et al. [2018) Hafner et al., 2019). In these models, the state of the underlying MDP is modeled with the deterministic state of a recurrent network (e.g.. LSTM (Hochreiter & Schmidhuber, 1997) or GRU (Cho et al., 2014)), and optionally with some stochastic random variables. As mentioned earlier, model with latent state that entirely stochastic has the appealing interpretation of learning a representation that can properly represent uncertainty about any of the state variables, given past observations. We demonstrate in our work that fully stochastic state space models can in fact be learned effectively and, with well-designed stochastic network, such models perform on par to partially stochastic models and outperform fully deterministic models.
than fully stochastic ones (Buesing et al. 2018)
REINFORCEMENT LEARNING AND MODELING
This work addresses the problem of learning maximum entropy policies from high-dimensional observations in POMDPs, by simultaneously learning latent representation of the underlying MDP state using variational inference and learning the policy in a maximum entropy RL framework. In
this section, we describe maximum entropy RL (Ziebart 2010 Haarnoja et al. 2018a Levine 2018)
in fully observable MDPs, as well as variational methods for training latent state space models for POMDPs.
3.1 MAXIMUM ENTROPY RI. IN FULLY OBSERVABLE MDPs
In a Markov decision process (MDP), an agent at time t takes an action E A from state St E S and reaches the next state St+1 E S according to some stochastic transition dynamics p(s [s at). The initial state $ comes from a distribution p(sı), and the agent receives a reward on each of the transitions. Standard RL aims to leam the parameters & of some policy T (s:) such that the expected sum of rewards is maximized under the induced trajectory distribution Pr. This objective can be modified to incorporate an entropy term, such that the policy also aims to maximize the expected entropy H(T (-(s))) under the induced trajectory distribution Pz. This formulation has close connection to variational inference (Ziebart, 2010 [Harroja et al. [0][2008a] Levine 2018), and we build on this in our work. The resulting maximum entropy objective
T
arg max
Σ
E (-
[r(St, + aFH(To(-(ss))
(1)
where r is the reward function, and is temperature parameter that controls the trade-off between optimizing for the reward and for the entropy (i.c., stochasticity) of the policy. Soft actor-critic (SAC) (Haarnoja et al. 2018a) uses this maximum entropy RL framework to derive soft policy iteration, which alternates between policy evaluation and policy improvement within the described maximum entropy framework. SAC then extends this soft policy iteration to handle continuous action spaces by using parameterized function approximators to represent both the Q-function Qa (critic) and the policy no (actor). The soft Q-function parameters 0 are optimized to minimize the
3

---

Under review as a conference paper at ICLR 2020
soft Bellman residual,
JØ(0)=
E
2 (Qo(se, ar) - - (re + + (Vo(5++1)))
(2)
(0, (e.
(g(+++1)= = E 16:1~TT
(Qa( art1) - a log T (a++1 [8++)]]
(3)
where D is the replay buffer, y is the discount factor, and @ are delayed parameters. The policy parameters @ are optimized to update the policy towards the exponential the soft Q-function,
(Ø)
=
(a (olo BLTEE
(s(ss,a)))
S,
Results of this stochastic, entropy maximizing RL framework demonstrate improved robustness and stability. SAC also shows the sample efficiency benefits an off-policy learning algorithm, in conjunction with the high performance benefits of long-horizon planning algorithm. Precisely for these reasons, we choose to extend the SAC algorithm in this work to formulate our SLAC algorithm.
3.2 SEQUENTIAL LATENT VARIABLE MODELS AND AMORTIZED VARIATIONAL INFERENCE IN POMDPS
To learn representations for RL, we use latent variable models trained with amortized variational inference. The learned model must be able to process a large number of pixels that present in the entangled image x, and it must tease out the relevant information into a compact and disentangled representation To learn such model, can consider maximizing the probability of each observed datapoint x from some training set 'D under the entire generative process p(x) ) p(x z p(())) dz. This objective is intractable compute in general due to the marginalization of the latent variables In amortized variational inference, we utilize the following bound on the log-likelihood (Kingma & Welling 2014).
EX [log ((x))( 2 EX [Fa-g [log [z]] - DEL (g(zlx) p(z))]-
We can maximize the probability of the observed datapoints (i.e., the left hand side of Equation (5) by learning an encoder g( (x) and a decoder p( xx (x) and then directly performing gradient ascent on the right hand side of the equation. In this setup, the distributions interest are the prior p(z), the observation model p(x 1z. and the posterior q(z (x)
Although such generative models have been shown to successfully model various types of complex distributions (Kingma& & Welling 2014) by embedding knowledge of the distribution into informa- tive latent space, they do not have built-in mechanism for the use of temporal information when performing inference. In the case of partially observable environments, representative latent state 2, corresponding to a given non-Markovian observation x, needs to be informed by past observations.
discuss below, the
Consider a partially observable MDP (POMDP), where an action a, E A from latent state z, E 2 results in latent state 2 Z and emits a corresponding observation X X. We make an explicit distinction between an observation
and the underlying latent state 2,, to emphasize that the latter unobserved and the distribution not known priori. Analogous to the fully observable MDP, the initial state distribution is p(z1), the transition probability distribution is p(zz-1 az). and the reward rz. In addition, the observation model given by p(x, (2:).
As in the case for VAES, a generative model of these observations X; can be learned by maximizing the ligeliiklooodd. the POMDP setting, however, we note that x, alone does not provide all necessary information infer 2,, and thus, prior temporal information must be taken into account. This brings us to the discussion of sequential latent variable models. The distributions of interest are the priors p(z1) and p(z++1 (2, a,), the observation model p(x, (2z), and the approximate posteriors g(z, Xx) and q(zz-1 Xt+11 Zt, a, ).). The likehood of the observations can then be bounded, similarly to the VAE bound in Eqqutioo (5)
log p(x (
Σ log p(xx (z,) DKL (g(21(x1) || p(z1))
DKL (9(2)-1|xee1) an) |l p(z (2, ai))
(6)
4

---

Under review as a conference paper at ICLR 2020
Prior work (Hafner et al. 2019 Buesing et al. 2018 Doerr et al. 2018b) 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, enable the effective training of a fully stochastic sequential latent variable model, and bring it together with maximum entropy actor-critic RI. algorithm create SLACC 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
Our method aims to learn maximum entropy policies from high-dimensional, non-Markovian observations in POMDP, while also learning model of that POMDP. The model alleviates the representation learning problem, which in turn helps with the policy learning problem. We formulate the control problem as inference in a proba- bilistic graphical model with latent variables, as shown in Figure
Or
ar+1
PT
For fully observable MDP, the control problem can be embedded into graphical model by introducing bi- nary random variable O₁, which indicates time step t is optimal. When its distribution is chosen to be p(O, = 1|s,, az) = exp(r(s;, at)). then maximization of p(O1.)) via approximate inference in that model yields the optimal policy for the maximum entropy objective (Levine 2018).
Figure I:
Graphical model of POMDP
with optimality variables for 2 T 1.
In POMDP setting, the distribution can analogously be given by p(O, = 1|z, a.) exp(r(z,, at)). Instead of maximizing the likelihood of the optimality variables alone, jointly model the observa- tions (including the observed rewards of the past time steps) and learn maximum entropy policies by maximizing the marginal likelihood p(x (-+11|a1:77 This objective represents both the likelihood of the observed data from the past steps, well as the optimality of the agent's is actions for future steps. We factorize our variational distribution into a product of recognition terms g(z) xx1
and q(z++1 1x Ze, a dynamics terms p(z++1 [z, ax), and policy terms т (a [z)::
((ZL.T,8410) 11 a1:r)
T-1 = { x1) II a ((zz+1 2g, a.) II p(z 22, ax)
T II (=++1
T (z))
(7)
The variational distribution uses the dynamics for future time steps to prevent the agent from controlling the transitions and from choosing optimistic actions (Levine 2018). The posterior over the actions represents the agent's policy æ. Although this derivation uses policy that conditioned on the latent state, algorithm, which will be described in the next section, learns parametric policy that is directly conditioned on observations and actions. This approximation allows us to directly execute the policy without having perform inference the latent state at time.
We use the posterior from Eqquation (7) to obtain the evidence lower bound (ELBO) of the marginal likelihood,
log p(x O++1:1 a1:))
E
log p(x)
(R 1:3 )m
- DKL (g(2((x)) p(z1)) -
DEL |1 p(ze41/2,, at))
Σ (r(z,, + log log -log-loooooooo ((al)2))
(8)
where p(a.) is the action prior. The full derivation of the ELBO is given in Appendix A This derivation assumes that the reward function, which determines p(O, 2 a,), ). is known. However, in many RL problems, this is not the case. In that situation, we can simply append the reward to the
5

---

Under review as a conference paper at ICLR 2020
observation, and learn the reward along with p(x; (z.). This requires modification other than changing the observation space, and we use this approach in all of our experiments. We do this to learn latent representations that are more relevant to the task, but we do not use predictions from it. Instead, the RI. objective uses rewards from the agent's experience, as in model-free RI..
method
5 STOCHASTIC LATENT ACTOR CRITIC
We now describe our stochastic latent actor critic (SLAC) algorithm, which approximately maximizes the ELBO using function approximators to model the prior and posterior distributions. The ELBO objective in [Equation can be split into a model objective and a maximum entropy RL. objective. The model objective can directly be optimized, while the maximum entropy RL objective can be solved via message passing. We can learn Q-functions for the messages, and then we can rewrite the RL objective to express it in terms of these messages. Additional details of the derivation of the SLAC objectives are given in Appendix A
The first part of the ELBO corresponds to training the latent variable model to maximize the likelihood of the observations, analogous to the BLOO in Eqqatiin (6) for the sequential latent variable model. The distributions of the latent variable model are diagonal Gaussian distributions, where the means and variances are outputs of neural networks. The distribution parameters z of this model are optimized maximize the first part of the ELBO, The model loss
Latent Variable Model:
(M) =
E
E M1 Tug
log Pvp (xe|zz)
(X AND
Σ DEL (qw (2 (x+1, Zz, ac) || pa (Z++ 12c, az))
-KLL (z)) (l))
(9)
We use the reparameterization trick to sample from the filtering distribution q: (Z X1:++1+ a..)).
Critic and Actor: The second part of the ELBO corresponds to the maximum entropy RL objective. As in the fully observable case from[Section [Scction I and as described by Levine (2018), this optimization can be solved via message passing of soft Q-values, except that we use the latent states z rather than the true states For continuous state and action spaces, this message passing is approximated by minimizing the soft Bellman residual, which we use to train our soft Q-function parameters 0,
1 (Qo(2r)a+)
1
(o)
E (I Nr)
E Mir mo
)| (10)
(0(2)+1,a44)) |
47410
where 0 are delayed parameters, obtained as exponential moving averages of Notice that the latents z, and Z,+1, which are used in the Bellman backup, sampled from the same joint, i.e. Z,+1 ~ (2 27, a,). The RL objective, which corresponds to the second part of the ELBO, can be rewritten in terms of the soft Q-function. The policy parameters are optimized to
maximize this objective, analogously to soft actor-critic (Haarnoja et al.) 2018a) The policy loss is then
(=(
E Bir+i~gg
E art1~
la ( a1r))
a am))
(X (1 DND
(11)
We assume uniform action prior, so p(a:) is constant term that we omit from the policy loss. We use the reparameterization trick to sample from the policy, and the policy loss only uses the last sample 2 of the sequence for the critic. Although the policy used in our derivation conditioned in the latent state, our learned parametric policy is conditioned directly on the past observations and actions, so 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 (10), rather than sampling latent states z ~ Z, we sample latent states from the filtering distribution qe [21:++1 [1:r+1,, a1.7). ). This design choice allows to minimize the critic loss for samples that are most relevant for Q, while also allowing the critic loss to use the Q-function in the
same way implied by the policy loss in Eqqatton (11)

---

Under review as a conference paper at ICLR 2020
SLAC is outlined in Algorithm The actor-critic component follows prior work, with automatic tuning of the temperature a and two Q-functions to mitigate underestimation (Fujimoto et al. [0188) Haarojj et al. 20188 b). SLAC can be viewed as variant of SAC (Haarnoja et al. 201 8a) where the
critic is trained on the stochastic latent state of our sequential latent variable model. T'ee hackup for the
critic is performed on tuple (2,, ar, rr; 2,+ sampled from the posterior g( 2, X1:r+1, a1:r). The critic can, in principle, take advantage of the perfect knowledge of the state Z,, which makes learning easier. However, the parametric policy does have access to Zt, and must make decisions based on a history of observations and actions. SLAC is not a model-basee 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.
Algorithm Stochastic Latent Actor-Critic (SLAC)
Require: F., 1/, 01, Oz ob X1 ~ Breser()
Environment and initial parameters for model, actor, and critic
> Sample initial observation from the environment Initialize replay buffer with initial observation
D (x1)
for each iteration do for each environment step do
a. T (a, a1:t-1)
> Sample action from the policy Sample transition from the environment > Store the transition in the replay buffer
Tr. X Exeep (a))
+ DU( (ax; rz; X(+1)
for each gradient step do s t v/ - Aas V JN (bb)
> Update model weights b Update the Q-function weights D Update policy weights D Update target critic network weights
0, + 0, Ag V 0, (0.) for [1,2]
& < Φb A V (
0, + vl, + (1 v)0; for i [1,2]
6 LATENT VARIABLE MODEL
We briefly summarize our full model architecture here, with full details in Appendix B Motivated by the recent success of autoregressive latent variables in VAEs (Razavi et al. 2019 Maaloe et al. 2019), we factorize the latent variable z, into two stochastic layers, zi and z as shown in Figure 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 produce high quality reconstructions and samples, and utilize it in all of our experiments. The generative model and the inference model g are given by
Pe (21) = pp(zllz()p(
Pv (Zc-1 [2,, ar) (z |2i+1 , 2 a (e)pw(zh (z) (2,,ar), ar),
Iv (21 (x) (2 ( | g0 ( (x1);
qu (Z++1 1 Ze, ar) - (x 12441 2 a y (2 (2 (x+11 at).
ar).
Note that we choose the variational distribution ! 2? to be the same the model p. Thus, the KL divergence in JM simplifies to the divergence be- tween y and p over z, We normal distribution for p(z)), since is not condi- tioned on any variables, i.c. z (N I). The conditional distributions of our model are diagonal Gaussian, with means and variances given by neural networks. Unlike models from prior work (Hafner ct al. 2019 Buesing et al., 2018 Doerr et al. 8b). which have deterministic and stochastic paths and use recurrent neural networks, ours is fully stochastic, i.e. latent state is Markovian latent random variable formed by the concatenation of z and z,. Further details discussed in Appendix
multivariate standard
#
Figure 2: Diagram of our full model. Solid arrows show the generative model, dashed arrows show the inference model. Rewards not shown for clarity.
7

---

Under review as a conference paper at ICLR 2020
7 EXPERIMENTAL EVALUATION
We evaluate SLAC on numerous image-based continuous control tasks from both the DeepMind Control Suite (Tassa et al. 2018) and OpenAl Gym (Brockman et al. (0166), as illustrated Figure Full details of SLAC's s network architecture are described in Appendix B Aside from the value of action repeats (i.e. control frequency) for the tasks, we kept all SLAC's hyperparameters constant across all tasks in all domains. Training and evaluation details are given in Appendix C. and image samples from our model for all tasks are shown in Appendix D Additionally, visualizations of our results and code are available on the project website
7.1 COMPARATIVE EVALUATION ON CONTINUOUS CONTROL BENCHMARK TASKS
To provide 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 (Tassa et al. 2018), and four tasks (cheetah, walker, ant, hopper) from OpenAI Gym (Brockman et al. (2166]. 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 the following stae-offthhe-ar model-based and model-free algorithms:
SAC (Haarnoj al. 2018a): This is off-policy actor-critic algorithm, which represents comparison to state-of- model-free learning. We include experiments showing the performance state (as an upper bound on performance) as well as directly from raw images.
of SAC based
MPO (Abdolmaleki et al. 201 8b a): This is an off-policy actor-critic algorithm that performs an expectation maximization form of policy iteration, learning directly from raw images.
D4PG (Barth- Maron et al. 2018): 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 training steps, as stated in the benchmarks from (Tasaa et al. 2018).
PlaNet (Hafner al. 2019): This 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 optimized with the cross entropy method (CEM).
DVRL (gi et al. 2018): This an on-policy model-free RL algorithm that also 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 history of actions and observations.
Our experiments on the DeepMind Control Suite in Figure 4 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-base 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 off-policy RL algorithm, which can better take advantage of the learned representation.
We also evaluate SLAAC on continuous control benchmark tasks from OpenAl Gym in Figure 5] We notice that these tasks are much more challenging than the DeepMind Control Suite tasks, because the rewards are not as shaped and not bounded between 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
Example image observations for our continuous control benchmark tasks: DeepMind Control's cheetah run, walker walk, hall-in-cup catch, and finger spin, and OpenAl Gym's 'ss's half cheetah, walker, hopper, and ant (left to right). Images are rendered a resolution of 64 64 pixels.
Figure 3:
htttps:/rri
io/lac//
8

---

Under review as a conference paper at ICLR 2020
rus
1230
Walken walk
cup, caiads
Finges. spin
Retrun
310
620
230
15
20
25
62
65
08
322
6.1
4.5
80
62
04
36
Ensirmmmnn Stepps (Siillsons)
Eas Ruvirmmnnaa Steps (Millions)
Environmen Stars (Millinn)
Envirommen Stepp (Millloo))
(state)
D4PG (10" stape))
MPO (10P stepe)
MPO
SAC
PlaSet
DVRL
SLAC (ouss)
Figure 4: Experiments on the DeepMind Control Suite from images (unless otherwise labeled as "state"). SLAC (ours) converges similar better final performance than the other methods, while almost always achieving reward as high as the upper bound SAC baseline that learns from true state. Note that for these experiments, 1000 environments steps corresponds to I episode.
Walerriie2
Hurgweee
L00
Averng
85
Enstronmen Steps (Millson)
Factrumment Steps (Millionn)
Eaveinmment Staps (Millions)
Emvrroaeet Stege (Million)
steped
MPO
PlaNoi
(sus)
Figure 5: Experiments on the OpenAl Gym benchmark tasks from images. SLAC (ours) converges to higher performance than both PlaNet and SAC all four of these tasks. The number of environments steps in each episode is variable, depending on the termination.
the last three tasks, while for the cheetah task, it learns 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 algorithmss, 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 PlaNet on three of the tasks, and matches its performance on the walker task. However, on the harder image-based OpenAl Gym tasks, SLAC outperforms PlaNet by a large margin. In both domains, SLAC substantially outperforms all prior model-free methods. We note that the prior methods that we tested generally performed poorly on the image-based OpenAl Gym tasks, despite considerable hyperparameter tuning.
7.2 EVALUATING THE LATENT VARIABLE MODEL
We next study the tradeoffs between different design choices for the latent variable model. We compare fully stochastic model, as described in Section 6 to a standard non-sequential VAE model (Kingma & Welling, 2014), which has been used in multi- ple prior works for representation learning in RL (Higgins et al. 2017, Ha & Schmidhuber, 2018: Nair et al.] 2018). the partially stochastic model used by PlaNet (Hafner et al. 2019), as well as three variants of our model: simple filtering model that does not factorize the latent variable into two layers of stochastic units, a fully deterministic model that removes all stochasticity from the hidden state dynamics, and partially stochastic model that has both deterministic and stochastic transitions, similar to the PlaNet model, but with our architecture. Both the fully deterministic and partially stochastic models use the same architecture as our fully stochastic model, including the same two-level factorization of the latent variable. In all cases, we use the RL framework of SLAC and only vary the choice of model for representation learning. As shown in the comparison in FFiuure 6) our fully stochastic model outperforms prior models well as the deterministic and simple variants
Cheetah
Avesagee Moturn
Environmen Steps (Millions)
Figure 6: Comparison of differ- ent design choices for the latent variable model.

---

Under review as a conference paper at ICLR 2020
Ground
Trah
I
K
Posterior Sample
Cheetah run
{
K
K
7
Conditional
Prior Sample
K
K
K
x
Prior Sample
Figure 7: Example image sequence seen for the chectah task (first row), corresponding posterior sample (reconstruction) from model (sccond row), and generated prediction from the generative model (last two rows). The second last row is conditioned on the first frame (i.e., the posterior model is used for the first time step while the prior model is used for all subsequent steps), whereas the last row not conditioned on any ground truth images. Note that all of these sampled sequences are conditioned on the same action sequence, and that our model produces highly realistic samples, even when predicting via the generative model.
of our own model. The partially stochastic variant of our model matches the performance our fully stochastic model but, contrary to the conclusions in prior work (Hafner et al. 2019; Buesing et al.) 2018). the fully stochastic model performs on par, 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 these latent samples are less sensitive to the propagation of the latent states through time.
7.3 QUALITATIVE PREDICTIONS FROM THE LATENT VARIABLE MODEL.
We show example image samples from our learned sequential latent variable model for the cheetah task Figure 7, and we include the other tasks in Appendix D Samples from the posterior show the images X, as constructed by the decoder P (x, k,), using sequence of latents z, that are encoded and sampled from the posteriors, qu [z [=1xx) and gu (Z Xx++1, Zg, at). Samples from the prior, on the other hand, use sequence of latents where 1 is sampled from p(z:) and all remaining latents z; are from the propagation of the previous latent state through the latent dynamics pw [Z++1 (z, a,). Note that these prior samples do not use any image frames inputs, and thus they do not correspond to any ground truth sequence. We also show samples from the conditional prior, which conditioned on the first image from the true sequence: for this, the sampling procedure is the same the prior, except that 2 is encoded and sampled from the posterior qu (z [x1), rather than being sampled from p(z)). We notice that the generated images samples can be sharper and more realistic by using smaller variance for Pv (x, (z)) 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.
8 DISCUSSION
We presented SLAC, an efficient RL algorithm for learning from high-dimensional image inputs that combines efficient off- off-colicy model-free RL with representation learning via 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 RI. 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 bandel-based/mmdelll. and performance. More broadly, the 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.
method could further improve sample efficiency
10

---

Under review as a conference paper at ICLR 2020
REFERENCES
Abbas Abdolmaleki, Jost Tobias Springenberg. Jonas Degrave, Steven Bohez, Yuval Tassa, Dan Belov, Nicolas Heess, and Martin A. Riedmiller. Relative entropy regularized policy iteration. arXiv preprint ar1:::
2018a.
Abbas Abdolmalcki, Jost Tobias Springenberg, Yuval Tassa, Rémi Munos, Nicolas Heess, and Martin A. Riedmiller. Maximum posteriori policy optimisation. In International Conference Learning Representations (ICLR), 2018b.
Evan Archer, Il Memming Park, Lars Buesing, John Cunningham, and Liam Paninski. Black box variational inference for state space models. arXiv preprint 1rXiv: 2015.
Karl Astrom. Optimal control of markov processes with incomplete state information. Journal of mathematical analysis and applications, 1965.
Gabriel Barth-Maron, Matthew W Hoffman, David Budden, Will Dabney, Dan Horgan, Alistair Muldal, Nicolas Heess, and Timothy Lillicrap. Distributed distributional deterministic policy gradients. In International Conference on Learning Representations (ICLR), 2018.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. OpenAI Gym. arXiv preprint arXiv: /60i: 01500015440 2016.
Lars Buesing, Theophane Weber, Sébastien Racanière, S. M. Ali Eslami, Danilo Jimenez Rezende, David Reichert, Fabio Viola, Frederic Besse, Karol Gregor, Demis Hassabis, and Daan Wier- stra. Learning and querying fast generative models for reinforcement learning. arXiv preprint 1802. 2018.
Kyunghyun Cho, Bart van Merrienboer, Çaglar Gülçehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation. In Conference Empirical Methods in Natural Language Processing (EMNLP), 2014.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. In Neural Information Processing Systems (NeurIPS), 2018.
Robert Dadashi, Adrien Ali Taiga, Nicolas Le Roux, Dale Schuurmans, and Mare Bellemare. The value function polytope in reinforcement learning. In International Conference Machine Learning (ICML), 2019.
Marc Deisenroth and Carl E Rasmussen. PILCO: A model-based and data-efficient approach to policy search. In International Conference on Machine Learning (ICML), 2011.
Andreas Doem, Christian Daniel, Martin Schiegg, Nguyen-Tuong Duy, Stefan Schaal, Mare Toussaint, and Trimpe Sebastian. Probabilistic recurrent state-space models. In International Conference on Machine Learning (ICML), 2018a.
Andreas DOeIT, Christian Daniel, Martin Schiegg, Duy Nguyen- Tuong, Stefan Schaal, Marc Toussaint, and Sebastian Trimpe. Probabilistic recurrent sate- models. In International Conference on Machine Learning (ICML), 2018b.
Chelsea Finn and Sergey Levine. Deep visual foresight for planning robot motion. In International Conference Robotics and Automation (ICRA), 2017.
Chelsea Finn, Xin Yu Tan, Yan Duan, Trevor Darrell, Sergey Levine, and Pieter Abbeel. Deep spatial autoencoders for visuomotor learning. In International Conference on Robotics and Automation (ICRA), 2016.
Jakob Foerster, Ioannis Alexandros Assael, Nando de Freitas, and Shimon Whiteson. Learning to communicate with deep multi-agent reinforcement learning. In Neural Information Processing Systems (NIPS), 2016.
11

---

Under review as a conference paper at ICLR 2020
Marco Fraccaro, Soren Kaae Sonderby. UIrich Paquet, and Ole Winther. Sequential neural models with stochastic layers. In Neural Information Processing Systems (NIPS), 2016.
Marco Fraccaro, Simon Kamronn, Ulrich Paquet, and Ole Winther. A disentangled recognition and nonlinear dynamics model for unsupervised learning. In Neural Information Processing Systems (NIPS), 2017.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International Conference Machine Learning (ICML), 2018.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International Conference on Machine Learning (ICMI., 2019.
Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aaron den Oord. Shaping belief states with generative environment models for rl. In Neural Information Processing Systems (NeurIPS), 2019.
Shixiang Gu, Timothy Lillicrap, Ilya Sutskever, and Sergey Levine. Continuous deep q-learning with masel- acceleration. In International Conference on Machine Learning (ICML), 2016.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint 10XXv:8803.11122 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International Conference on Machine Learning (ICML), 2018a
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Schoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, and Sergey Levine. Soft actor-critic algorithms and applications. arXiv preprint 1812.0591 2018b.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning (ICML), 2019.
Matthew Hausknecht and Peter Stone. Deep recurrent Q-learning for partially observable MDPs, In AAAI Fall Symposium Sequential Decision Making for Intelligent Agents, 2015.
Irina Higgins, Arka Pal, Andrei Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. DARLA: Improving zero-shot transfer in reinforcement learning. In International Conference Machine Learning (ICMI), 2017.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 1997.
Maximilian Igl, Luisa Zintgraf. Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational reinforcement learning for POMDPs. In International Conference on Machine Learning (ICML), 2018.
Max Jaderberg, Volodymyr Mnih, Wojeiech Marian Czarnecki, Tom Schaul, Joel Z Leibo. David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In International Conference Learning Representations (ICI.R) 2017.
Leslie Pack Kaelbling, Michael L Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, (01(19))999134, 1998.
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick van der Smagt. Deep variational bayes filters: Unsupervised learning of state space models from raw data. In International Conference Learning Representations (ICLR), 2016.
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick van der Smagt. Deep variational bayes filters: Unsupervised learning of state space models from raw data. In International Conference on Learning Representations (ICLR), 2017.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, In International Conference Learning Representations (ICLR), 2015.
12

---

Under review as a conference paper at ICLR 2020
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Representations (ICLR), 2014.
Rahul G. Krishnan, Uri Shalit, and David Sontag. Deep kalman filters. arXiv:1511005121, 2015.
arXiv preprint
Sascha Lange and Martin Riedmiller. Deep auto-encoder neural networks in reinforcement learning. In International Joint Conference on Neural Networks (IJCNN), 2010.
Sergey Levine. Reinforcement learning and control probabilistic inference: Tutorial and review. arXiv preprint 1rXiv:1805.00900 00909, 2018.
Lars Maaloe, Marco Fraccaro, Valentin Liévin, and Ole Winther. Biva: A very deep hierarchy of latent variables for generative modeling. In Neural Information Processing Systems (NeurIPS), 2019.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin A. Riedmiller. Playing Atari with deep reinforcement learning. In NIPS Deep Learning Workshop, 2013.
Anusha Nagabandi, Gregory Kahn, Ronald S Fearing, and Sergey Levine. Neural network dynam- ics for masel-based deep reinforcement learning with model-free fine-tuning. In International Conference Robotics and Automation (ICRA), 2018.
Ashvin V Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Visual reinforcement learning with imagined goals. In Neural Information Processing Systems (NeurIPS), 2018.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiiv1807.03748 2018.
Ali Razavi, Aaron den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. 2019.
Evan Shelhamer, Parsa Mahmoudich, Max Argus, and Trevor Darrell. Loss is its own reward: Self-supervision for reinforcement learning. arXiv preprint arXiv:1612.07300 2016.
Richard S Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM SIGART Bulletin, 160-166,,, 1991.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy Lillicrap, and Martin Riedmiller. DeepMind control suite. arXiv preprint arXiv: 2018.
Niklas Wahlström, Thomas Schön, and Mare Peter Deisenroth. From pixels to torques: Policy learning with deep dynamical models. arXiv preprint arXiv: 10Xi::/502.0225/, 2015.
Manuel Watter, Jost Springenherg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. In Neural Information Processing Systems (NIPS), 2015.
Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Mattew J Johnson, and Sergey Levine. SOLAR: Deep structured latent representations for model-based reinforcement learning. In International Conference Machine Learning (ICML), 2019.
Pengfei Zhu, Xin Li, Pascal Poupart, and Guanghui Miao. On improving deep reinforcement learning for POMDPs. arXiv preprint arXiv:1804.06300 2018.
Brian D Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy. 2010.
13

---

Under review as a conference paper at ICLR 2020
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 the control as inference framework for reinforcement learning (Levine) 2018 In this framework, the problem of finding the optimal policy is cast as an inference problem, conditioned on the evidence that the agent behaving optimally. While[Leving (2018) derives this in the fully observed case, we present derivation in the POMDP setting.
We aim to maximize the marginal likelihood +(X (+r1/aa)))) where T the number of steps that the agent has already taken. This likelihood reflects that the agent cannot modify the past actions and they might have not been optimal, but it can choose the future actions up to the end 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 dynamics model of the state would not change the model-free equations derived from the maximum entropy RI objective.
For reference, we restate the factorization of our variational distribution:
q(Zi:) 4741:7 X1:++1, a1:7)
T-1 II
= [((zi(x))
II
[ X Zr, ax)
p Z (,, ar) II (atlz))) (12)
As discussed by Levine use the dynamiic p(z[+1 Z,, ar) for t 2 1 in the variational distribution in order to prevent the agent from choosing optimistic actions.
2018), , the agent does not have control over the stochastic dynamics, so we
The joint likelihood is
p(Xir+1, Or41. ZI:T. Art17|/a17
T+1
T II t=++1
=((z))
p(z, Be a,)
p(x
p(O, 2, a.) II p(a.). (13)
1=7+1
We use the posterior from Equation (12) and Jensen's 's inequality to obtain the ELBO of the marginal likelihood,
log p(x1:+111 O++1:1 [1:r))
= log /
/
p(xi++11, O++1:T 2i:r; a.417 (a.r) d21. dar+1:1
81:T Br:1:1
Σ
log pl (xe
DKI. (x(z) Il p(zi))
Σ DKL. (g(z)+1 (X++1 Zr, a,) p(z,+ (z,, a,))
Σ +ogg(a))))) log Soga(tn) m(a)lzz))
log
(14)
Notice that the dynamics terms log p(z) a.) for t 2 T + 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 (9)
The second part of the ELBO corresponds to the maximum entropy RL objective. We assume a uniform action prior, so the p(a)) term is constant term that be omitted when optimizing
14

---

Under review as a conference paper at ICLR 2020
this objective. We message passing to optimize this objective, with messages defined as
Q(ze, at) r(zz; ac) +
[V((+1]]
(15)
V(zi) log exp(Q(zze, ar))
(16)
Then, the maximum entropy RL objective can be expressed in terms of the messages as
(r(z,, an) - log m(au/zz)) -log
E
(x++1:7 Prr
(((T+1, - (ar+ [mtti]]
E
E
(17)
2,41~9(/
(
exp ( (, ar+1))
E
DRT. ((ar+1/2r+4)) exp (V (V++1)))
+ v(ZE+i)
Zr410ig(1 2..
(18)
where the first equality is obtained from dynamic programming (see Levine (2088) for details), the second equality holds from the definition of KL divergence, and exp (V(z,) is the normalization
factor for exp (2, ar)) with respect a.. Since the KL divergence term minimized when its two arguments represent the same distribution, the optimal policy is given by
π aaa ( exp (Q(z, a.) V(z.)).
(19)
Noting that the KL divergence term is zero for the optimal action, the equality from Equation (18) can be used Equation (15) to obtain
Q(2,, a²) = r(z; ax) +
E
E
Q(Z4+1, a411) - log л (at+1 24+1
#e41mof (..1
(20)
This equation corresponds to the standard Bellman backup with soft maximization for the value function.
As mentioned in [Section 5| our algorithm conditions the parametric policy in the history of observa- tions and actions, which allows us to directly execute the policy without having to perform inference on the latent state run time. When using the parametric function approximators, the negative of the maximum entropy RL objective, written as in Equation (17) corresponds to the policy loss in Equation (11) Lastly, the Bellman backup of Equation (20) corresposs the Bellman residual in Equation (10) 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 terms of messages of Q-functions, which turn are learned by minimizing the Bellman residual. These objectives lead to the model, policy, and critic losses.
B NETWORK ARCHITECTURES
Recall that our full sequential latent variable model has two layers latent variables, which we denote z and z². We found this design to provide good balance between ease of training and expressivity, producing good reconstructions and generations and, crucially, providing good representations for rein- forcement learning. For reference, we reproduce the model diagram from the main paper in Figure 8 Note that this diagram represents the Bayes correspond- ing to our full model. However, since all the la- tent variables are stochastic, this visualization also
Figure 8: Diagram of our full model, reproduced from the main paper. Solid arrows show the gen- erative model, dashed arrows show the inference model. Rewards shown for clarity.
15

---

Under review as a conference paper at ICLR 2020
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 part of the generative model and all dashed arrows as part of approximate posterior. The generative model consists of the following probability distributions, described in the main paper:
of ~ p(z))
zî ~ Pa z z
2(41 ~ pv [2 Zi at)
z4t p (2 | ar)
X, px k,)
rr ~ (ra m2, 2i, ac, 2 2(+1)
The initial distribution p(z}) a multivariate standard normal distribution NV 1). All of the other distributions are conditional and parameterized by neural networks with parameters tb. The networks for Pv (2]]zz)) (z)z ar), Pel (z7_, |z,,,, zi; az), and p (ralz), 27 ar, z z14 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 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-linear.ty. The observation model Puo (x,|z), 2,) consists of 5 transposed convolutional layers (256 4 x 4, 128 3 x 3. 64 3 x 3, 32 3 x 3, and 3 5 x 5 filters, respectively, stride 2 each, except for the first layer). The output variance for each image pixel is fixed to 0.1.
The variational distribution q, also referred as the inference model or the posterior, is represented by the following factorization:
z qu (z (x1)
xî p (z) 2
2 ~ qu [ X141 27, at)
z ~ (Z 2 27 ar).
Note that the variational distribution over z and z is intentionally chosen to exactly match the generative model such that this term does not appear in the KL-divergence within the ELBO, and a separate variational distribution is only learned over zi and z This intentional design decision simplifies the inference process. The networks representing the distributions qu (z) |x:) and tv (z/41 [x+1,, n2; a,) both consist of convolutional layers (32 5 5, 64 3 3, 128 x x. 256 3 and 256 4 x 4 filters, respectively, stride 2 each, except for the last layer). 2 fully connected layers (256 units each), and Gaussian output layer. The parameters of the convolution layers are shared among both distributions.
The latent variables have 32 and 256 dimensions, respectively, i.e. z R and z 256 For the image observations, X, E [0, 1] 1x64x44x3 All the layers, except for the output layers, use leaky ReLU -oneeiirasss Note that there are no deterministic recurrent connections in the network all networks are feedforward, and the temporal dependencies all flow through the stochastic units 2 and 2
For the reinforcement learning process, we use a critic network Qo consisting of 2 fully connected layers (256 units each) and a linear output layer. The actor network F& consists of 5 convolutional layers, 2 fully connected layers (256 units each), Gaussian layer, and tanh bijector, which constrains the actions to be in the bounded action space of [-1, The convolutional layers are the same as the ones from the latent variable model. but the parameters of these layers are not updated by the actor objective. The same exact network architecture is used for every one of the experiments in the paper.
16

---

Under review as a conference paper at ICLR 2020
Action repeat
Original control
Effective control time step
Benchmark
Task
time step
cheetah
4
0.01
0.04
walker walk
2
0.025
0.05
DeepMind Control Suite
ball-in-cup catch
4
0.02
0.08
finger spin
2
0.02
0.04
Half HalfCheetah-v2
0.05
0.05
Walker2d-v2
4
0.008
0.032
OpenAI Gym
Hopper-v2
2
0.008
0.016
Ant-v2
4
0.05
0.2
Table 1: Action repeats and the corresponding agent's control time step used in our experiments.
C TRAINING AND EVALUATION DETAILS
The control portion of our algorithm uses the same hyperparameters as SAC (Haarnjja al. 2018 8a except for smaller replay buffer size of 100000 environment steps (instead of a million) due to the high memory usage of image observations. All of the parameters optimizer (Kingma & Ba 2015), and we perform one gradient step per environment step. The Q-function and policy parameters are trained with learning rate of 0.0003 and batch size of 256. The model parameters are trained with a learning rate of 0.0001 and batch size of 32. We use for all the tasks. Note that the sequence length can be less than 7 for the first l steps (l < -) of each episode.
trained with the Adam
sequences of length
We use action repeats for all the methods, except for D4PG for which we use the reported results from prior work TTasaa et al. 2018). 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 fraction of the environment steps reported in plots. For example, 3 million environment steps of the cheetah task correspond to 750000 samples when using an action repeat The action repeats used in experiments are given in Table
Unlike in prior work (Haamojj et al. 2018a b), we use the same stochastic policy as both the behavioral and evaluation policy since found the deterministic greedy policy to be comparable worse than the stochastic policy.
D ADDITIONAL PREDICTIONS FROM THE LATENT VARIABLE MODEL
We show additional samples from our model in Figure 9 and Figure 10) Samples from the posterior show the images X; constructed by the decoder P (x (2)), using sequence of latents z, that are encoded and sampled from the posteriors, qu Ixi) and qu (Z++1 |X++ Zr, az). Samples from the prior, the other hand, use a sequence of latents where 2 is sampled from p(z1) and all remaining latents are from the propagation of the previous latent state through the latent dynamics Pv (Z141 Zg, a.). These samples do 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 2 is encoded and sampled from the posterior qu (z) (x1). ). rather than being sampled from p(z)).
17

---

Under review as a conference paper at ICLR 2020
Ground
Posterior Sample
Walker walk
Conditional Sample
Prior Sample
U
Ground
Ball-in-cupp catch
Postrior Sample
Conditional Sample
Prior Sumple
Ground
Poseriar Sample
Finger spin
Conditional Prior Sumple
Prior Sample
Figure 9: Example image sequences. along with generated image samples, for three of the DM Control tasks that we used in our experiments. See Figure 7 for more details and for image samples from the cheetah task.
18

---

Under review as a conference paper at ICLR 2020
Ground
Trih
HalfCheetahhv2
Posterior Sample
Conditional Prior Sample
Phior Sample
Ground Truth
Poeterioe Sample
Walker2d-v2
Conditional Price Sumple
Prior Sample
Ground
Posterior Sample
Hopper-v2
Conditional Sample
Phior Sample
Ground
Posterior Sample
Ant-v2
Conditional
Sample
Prioe Sample
Figure 10: Example image sequences, along with generated image samples, for the four OpenAl Gym tasks that we used our experiments.
19


Reviews:
[{"review_id": "SJe5StvycH", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose SLAC, an important extension of the recently introduced soft actor critic (SAC) algorthm, which operates on a learned latent state, rather than an observed one, and therefore aims to jointly learn to represent high dimensional inputs and execute continuous control based on this representation. \n\nSLAC is based on a dynamic, non-linear Markov generative model and incorporates structured variational inference to learn the latent state distribution. The generative model and amortized variational inference support the learning of strong latent expected future reward estimates (Q functions that condition on the latent state), which the policy, which conditions directly on the observations (i.e. image) is distilled against for fast inference. The paper demonstrates solid gains over existing techniques, brings together recent work under a rigorous framework, and is a pleasure to read.\n\nStrengths:\n-Novel formulation, SOTA results, well written.\n\nLimitations:\n-While the most important ablation, the role of making the primary latent variable stochastic, is investigated, a deeper investigation of what makes the model more effective than existing techniques would be insightful, and further strengthen the paper.\n-Related, the approach seems closest to PlaNet in structure, but rather than being used for planning, is executed directly as an off-policy actor-critic algorithm, generalizing SAC. A discussion, and possibly some additional experiments to explain the differences and understand the tradeoffs would strengthen the paper. The authors mention \", contrary to the conclusions in prior work (Hafner et al., 2019; Buesing et al., 2018), the fully stochastic model performs on par or better.\" Why?\n\nMinor:\n-Figure 6 partially stochastic in figure, mixed in text.\n\nOverall:\nA strong paper, that brings together and generalizes existing work, with strong experimentation and SOTA results. Definite accept.\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "BJgddhA3YH", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work proposed a fully stochastic RL method and demonstrated significantly improved performance on multiple tasks.\n\nPros:\n1. The presentation is very clear and easy to read. \n2. The proposed method is sensible \n3, The experimental evaluation shows great practical gain\n\nCons:\n1. The method itself is incremental. As discussed in the related work, this work can be seen as complementary to many related works such as Igl 18, but the novelty of the idea is rather limited. \n2. The claims and the real-benefit of the method may not be consistent. (My biggest concern)\nThe paper claims that full stochasticity contributed to the practical gain but in the experiment Figure 6, we can see the simple filtering does not perform well. \nIt seems that the benefit of the method is rather from such particular latent space design rather than the stochastic vs deterministic. \n3. Continue with the previous point, Figure 2 is not very well motivated and I believe that from Figure 1 to figure 2 design was the most important part of the performance gain. Such important designed was very briefly described without any motivation. \n4. With the previous point, the experiments may be unfair, because, another partially stochastic method can easily utilize such design and further improve the performance. \n5. The related work should add a discussion about stochastic sequential models such as Kalman VAE etc.  paragraph 3 motivates your contribution as VAE does not model sequential information. But there are many works such as the KVAE that are stochastic and models sequential information. \n\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I made a quick assessment of this paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "B1l4nYmjqH", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes an actor-critic method that tries to aid learning good policies via learning a good representation of the state space (via a latent variable model). In actor-critic methods, the critic is learnt to evaluate policies in the latent space, which further helps with efficient policy optimization. The proposed method is evaluated on image-based control tasks, with baseline evaluations against both model-based and model-free methods in terms of sample efficiency. \u000b\n\n- The key argument is that learning policies in the latent space is more efficient, as it is possible to learn good representations in the latent space. There are quite a few recent works (e.g DeepMDP, Gelada et al., 2019;  Dadashi et al., 2019) that talks about representation learning in RL, and yet the paper makes no relations or references to previous works. I find it surprising that none of the past related works are mentioned in the paper. \u000b\n\n- I find the arguments on solving a POMDP instead of a MDP a bit vague in this context. I understand that the goal is to solve image based control tasks - for which learning good representations via a latent variable model might be useful, but it does not explicitly require references to a POMDP? In most ALE tasks, we have pixel based observations too, which makes the ALE environments a POMDP in some sense, but we use approximations to it to make it equivalent to a MDP with sufficient history. The arguments on POMDP seems rather an additional mention, with no necessary significance to it?\u000b\n\n- The paper mentions solving RL in the learned latent space, which is empirically proposed to be a good approach without theoretical justifications. There are several recent works that tries to understand the representation learning in RL problem from a theoretical perspective too - it would be useful to see where this approach stands in light of those theoretical results? Otherwise, the contribution seems rather limited : solving RL in latent space is useful, but there are no justifications to it? Why should this approach even be adapted or what is the significance of it?\u000b\n\n- The proposed actor-critic method in the latent space is built on top of Soft Actor-Critic (SAC). I understand this is a design/implementation approach building from previous works - but it would have been useful to add more context as to what it means to learn a critic in the latent space. If the critic evaluates a policy in the latent space - then is this a good policy evaluation for actor-critic itself? Why or why not? I do not understand why the critic evaluation in the latent space is even a good approach? \u000b\n\n- My first impression was that the paper proposes a separate auxilliary objective for learning good representations based on which actor-critic algorithms can be made more efficient. However, this does not seem to be the case directly? Following on previous point - I find the argument of solving a critic in the latent space rather vague. \u000b\n\n- The sequential latent variable model proposed is based on existing literature. This can be any latent variable model (e.g VAEs), but I understand, as mentioned in the paper, the design choice of using sequential models to capture the temporal aspect. \u000b\n\n- The proposed algorithm is in fact a combination of SAC and sequential latent variable models, both of which are well-known in the literature. The SLAC algorithm combines these to solve image-based control tasks. As per equation 10, which is the regular policy optimization objective with max entropy - the only difference is that the critic is evaluated in the latent space. This appears to me as more of an engineeering choice, and experimentally one that perhaps give good results - but the lack of justifications of why equation 10 is even the right objective to solve makes the paper rather less appealing. \u000b\n\n- I think overall the contribution of the paper is rather limited. It is more of an experimental design and engineering approach that combines previous known techniques. The paper mentions learning good representations for RL, without any references or justifications - and it appears that overall there are bold claims made in the paper but it lacks significant scientific contribution. \u000b\n\n- Experimental evaluations are made on image based control tasks. Experimental results are compared to few baselines - but it is not clear whether these are even the right baselines. For example, it would have been good to include analysis of the proposed model with different latent variable models (including VAE) to perhaps justify the choice of the latent variable model. Results in figure 5 appear a bit concerning to me - these are mostly the standard Mujoco tasks from the OpenAI suite. Are these all image based benchmarks too, or the standard baselines? It is not clear from the text. Assuming they are standard baselines, the comparisons made are rather unfair (for example : SAC and MDP performs much better on tasks like HalfCheetah-v2). Why are the baselines performing so poorly in the results?\u000b\n\n- Overall, I think the paper needs more work in terms of writing and justifying the choice of the approach. There are significant references missing in the paper. Most importantly, there are quite a few claims made in the paper which are not properly justified, that makes the overall contribution and novelty of the paper rather limited. I would tend for a rejection of this paper, as it requires more work - both in terms of theoretical justifications (including references) and experimetnal ablation studies and more simpler benchmarks explaining the choice of the approach. \n\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly."}, "extra_text": {}}]

Author rebuttal:

Meta-review: An actor-critic method is introduced that explicitly aims to learn a good representation using a stochastic latent variable model. There is disagreement among the reviewers regarding the significance of this paper. Two of the three reviewers argue that several strong claims made in the paper that are not properly backed up by evidence. In particular, it is not sufficiently clear to what degree the shown performance improvement is due to the stochastic nature of the model used, one of the key points of the paper. I recommend that the authors provide more empirical evidence to back up their claims and then resubmit.
Decision: Reject