Latent State Marginalization as a Low-cost Approach for Improving Exploration
Authors: Dinghuai Zhang
Aaron Courville
Yoshua Bengio
Qinqing Zheng
Amy Zhang
Ricky T. Q. Chen
Venue: ICLR 2023 poster (2023)
Keywords: MaxEnt RL
latent variable modeling
world models

Abstract: While the maximum entropy (MaxEnt) reinforcement learning (RL) framework -- often touted for its exploration and robustness capabilities -- is usually motivated from a probabilistic perspective, the use of deep probabilistic models have not gained much traction in practice due to their inherent complexity. In this work, we propose the adoption of latent variable policies within the MaxEnt framework, which we can provably approximate any policy distribution, and additionally, naturally emerges under the use of world models with a latent belief state. 
We discuss why latent variable policies are difficult to train, how naive approaches can fail, and subsequently introduce a series of improvements centered around low-cost marginalization of the latent state, allowing us to make full use of the latent state at minimal additional cost. We instantiate our method under the actor-critic framework, marginalizing both the actor and critic. The resulting algorithm, referred to as Stochastic Marginal Actor-Critic (SMAC), is simple yet effective. We experimentally validate our method on continuous control tasks, showing that effective marginalization can lead to better exploration and more robust training. Our implementation is open sourced at https://github.com/zdhNarsil/Stochastic-Marginal-Actor-Critic.

Published as a conference paper at ICLR 2023
LATENT STATE MARGINALIZATION AS A LOW-COST APPROACH FOR IMPROVING EXPLORATION
Dinghuai Zhang", Aaron Courville, Yoshua Bengio Mila, University de Montreal
Qinqing Zheng, Amy Zhang, Ricky T. Q. Chen Meta AI (FAIR)
ABSTRACT
While the maximum entropy (MaxEnt) reinforcement learning (RL) framework- often touted for its exploration and robustness capabilities-is usually usualll motivated from a probabilistic perspective, the use of deep probabilistic models has not gained much traction in practice due to their inherent complexity. In this work, we propose the adoption of latent variable policies within the MaxEnt framework, which we show can provably approximate any policy distribution, and additionally, naturally emerges under the use of world models with latent belief state. We discuss why latent variable policies are difficult to train, how naïve approaches can fail, then subsequently introduce series of improvements centered around low-cost marginalization of the latent state, allowing us to make full use of the latent state at minimal additional cost. We instantiate our method under the actor- critic framework, marginalizing both the actor and critic. The resulting algorithm, referred to as Stochastic Marginal Actor-Critie (SMAC), is simple yet effective. We experimentally validate our method on continuous control tasks, showing that effective marginalization Our implementation is open sourced https: / //gittuut cem/ 2dhNarsil/ Stochast MetissiccMMaai ----arginnn----- Actor- AAtor-Critic..
lead better exploration and more robust training.
INTRODUCTION
A fundamental goal of machine learning is to develop methods capable of sequential decision making, where reinforcement learning (RL) has achieved great success in recent decades, One of the core problems in RL exploration, the process by which an agent learns to interact with its environment. To this end, a useful paradigm is the principle of maximum entropy, which de- fines the optimal solution to one with the highest amount of randomness that solves the task at hand. While the maximum entropy (Max- Ent) RL framework (Todorov, 2006; Rawlik et al., 2012) is often motivated for learning com- plex multi-modal behaviors through stochas- tic agent, algorithms that are most often used in practice rely on simple agents that only make local perturbations around a single action. Part of this is due to the need to compute the entropy of the agent and use it as part of the training objective.
Figure 1: The world model infers latent states (5)) from observation inputs (I). While most ex- isting methods only take one sample the mean from this latent belief distribution, the agent (+) of the proposed SMAC algorithm marginalizes out the latent state for improving exploration. Icons adapted from Mendonca al. (2021).
Meanwhile, the use of more expressive models have not gained nearly as much traction in the community. While there exist works that have increased the fiexibility of their agents by making use of more complex distributions such as energy- based models (Haarnoja et al., 2017), normal- izing flows (Haarnoja et al., 2018 8a; Ward et al., 2019), mixture-of-expers (Ren et al., 2021), and
Work done during an internship at Meta AI. Correspondence to: <di <dinghuai. zhang@hiinn. la quebec>. By multi-modality, we're referring distributions with different and diverse modes.

---

Published as a conference paper at ICLR 2023
autoregressive models (Zhang al., 2021b), these constructions often result in complicate training procedures and are inefficient in practice.
Instead, note that relatively simple approach to increasing expressiveness is to make of latent variables, providing the agent with its own inference procedure for modeling stochasticity in the observations, environment, and unseen rewards. Introducing latent variables into the policy makes it possible to capture a diverse set of scenarios that are compatible with the history of observations. In particular, a majority of approaches for handling partial observability make use of world models (Hafner et al., 2019; 2020), which already result in a latent variable policy, but existing training algorithms do make use of the latent belief state to its fullest extent. This is due in part to the fact that latent variable policies do not admit a simple expression for its entropy, and we show that naïvely estimating the entropy can lead to catastrophic failures during policy optimization. Furthermore, high-variance stochastic updates for maximizing entropy do not immediately distinguish between local random perturbations and multi-modal exploration. We propose remedies these aforementioned downsides of latent variable policies, making use of recent advances in stochastic estimation and variance reduction. When instantiated in the actor-critic framework, the result is a simple yet effective policy optimization algorithm that can perform better exploration and lead to more robust training in both fully-observed and partially-obsered settings.
Our contributions can be summarized as follows:
We motivate the use of latent variable policies for improving exploration and robustness to partial observations, encompassing policies trained on world models as a special instance.
We discuss the difficulties in applying latent variable policies within the MaxEnt RL paradigm. We then propose several stochastic estimation methods centered around cost- efficiency and variance reduction.
When applied to the actor-critic framework, this yields an algorithm (SMAC; Figure 1) that is simple, effective, and adds minimal costs.
We show through experiments that SMAC is more sample efficient and can more robustly find optimal solutions than competing actor-critic methods in both fully-observed and partially continuous control tasks.
BACKGROUND
2.1 MAXIMUM ENTROPY REINFORCEMENT LEARNING
We first consider a standard Markov decision process (MDP) setting. We denote states x E S and actions a, E A, for timesteps E N. There exists an initial state distribution p(x1), stochastic transition distribution pxx,|x--1, a)) and deterministic reward function r, S x A + R. We can then learn a policy π (( (x:) such that the expected sum rewards maximized under trajectories
T A (x), a....., XT, ar) sampled from the policy and the transition distributions.
While it is known that the fully-observed MDP setting has at least one deterministic policy solution (Sutton & Barto, 2018; Puterman, 1990), efficiently searching for an optimal policy generally requires exploring sufficiently large part of the state space and keeping track of frontier of current best solutions. As such, many works focus on the use of stochastic policies, often in conjunction with the maximum entropy (MaxEnt) framework,
max E Ep(r)
y (rx(x,, a:) a(H(((((x)))) where H(T(-Ixx)) = En, - (xxx)) [- log I (Ta,) [x)]]
where p(T) is the trajectory distribution with policy n, H((-) is entropy and y is discount factor.
The MaxEnt RL objective has appeared many times the literature (e.g. Todorov (2006); Rawlik et al. (2012); Nachum et al. (2017)), and is recognized for its exploration (Hazan et al., 2019) and robustness (Eysenbach & Levine, 2022) capabilities. can be equivalently interpreted as variational inference from probabilistic modeling perspective (Norouzi et al., 2016; Levine, 2018; Lee et al., 2020a). Intuitively. MaxEnt RL encourages the policy to obtain sufficiently high reward while acting randomly possible, capturing the largest possible set of optimal actions. Furthermore, it also
2

---

Published as a conference paper at ICLR 2023
optimizes policies to reach future states where has high entropy (Haarnoja al., 2017), resulting in improved exploration.
Soft Actor-Critic popular algorithm for solving MaxEnt RL is Soft Actor-Critic (SAC; Haarnoja et al. (2018b)), which we directly build on in this work due to its reasonable good performance and relative simplicity. Briefly, SAC alternates between learning a soft Q-function Q(x, ar) that satisfies the soft Bellman equation,
(x,, ar) = rx(x,, a.) + + (at (Xe41) (x...) ((xe+1,a44)) + aH([((Xx+)))) (2)
and learning a policy with the maximum entropy objective,
max Ex (x.lx)).) (Q(x,, a+) o²H(=(-(x))))
where states are sampled from a replay buffer D during training. In practice, SAC is often restricted to the use of policies where the entropy can be computed efficiently, e.g. factorized Gaussian policy for continuous control environments. This allows random movements occur as noise added independently for each action dimension. Our proposed approach, on the other hand, introduces structure in the exploration noise level.
2.2 WORLD MODELS FOR PARTIALLY-OBSERRY ENVIRONMENTS
In many practically motivated settings, the agents only have access certain observations, e.g. partial states, and the complete states must be inferred through observations. This can modelled through the partially observed MDP (POMDP) graphical model, which encompasses wide range of problem settings involving uncertainty. POMDP can be used to model uncertainty in the state, reward, or even the transition model itself (Áström, 1964). Here, the optimal policy must take into account these uncertainties, naturally becoming stochastic and may exhibit multi-modal behaviors (Todorov, 2006). Notationally, we only have access to observations F, X' with incomplete information, while the latent state S; S is unobserved, leading to latent state transition distribution p(s; a(-1)) observation distribution p(x; (s)) and reward function re as).
In order to tackle this regime. people have resorted to learning world models (Deisenroth & Ras- mussen, 10 Ha & Schmidhuber, 2018) that attempt to learn a belief state conditioned on the history of observations and actions, typically viewed as performing variational inference on the POMDP, i.e. the world model is then responsible for tracking a belief state S,, which update based on new observations through inference model q(s, 1, Xt). The POMDP and the inference model are often jointly trained by maximizing a variational bound on the likelihood of observations,
(og((xi:Taa)) 2 Ey logp(xe(se) -
a (:)ll((s av-)))
(4)
The world model is then typically paired with a policy that makes use of the belief state to take actions, i.e. x(ay)s:) with S, (ast, Xge), as the assumption is that the posterior distribution over St contains all the information we have so far regarding the current state.
3 STOCHASTIC MARGINAL ACTOR-CRITEL (SMAC)
We now discuss the use of latent variables for parameterizing policy distributions, and how these appear naturally under the use of a world model. We discuss the difficulties in handling latent variable policies in reinforcement learning, and derive cost-efficient low- low-variance stochastic estimators for marginalizing the latent state. Finally, we put all together an actor-critic framework.
3.1 LATENT VARIABLE POLICIES
We advocate the use of latent variables for constructing policy distributions as an effective yet simple way of increasing flexibility. This generally adds minimal changes to existing stochastic policy algorithms. Starting with the MDP setting, a latent variable policy (LVP) can be expressed as
raa:(x)) := ((a sy (s(ss(x)))))))))))))))))){)))( ds
(5)
3

---

Published as a conference paper at ICLR 2023
where Sy is latent variable conditioned the current observation. In the MDP setting, the introduction of a latent ((st)x/) mainly increases the expressiveness of the policy. This thus allows the policy to better capture wider frontier of optimal actions, which can be especially helpful during the initial exploration when we lack information regarding future rewards. We discuss extensions to POMDPs shortly in the following section, where the policy is conditioned on history of observations.
For parameterization, we use factorized Gaussian distributions for both π ((az)s)) and g(s)|x,). Firstly, this results in latent variable policy that is computationally efficient: sampling and density eval- uations both remain cheap. Furthermore, this allows us to build upon existing stochastic policy algorithms and architectures that have been used with a single Gaussian distribution, by simply adding new stochastic node S.. Secondly. we can show that this is also a sufficient parameterization: with standard neural network architectures, latent variable policy can universally approximate any distribution if given sufficient capacity. Intuitively, it is known that mixture of factorized Gaussian is universal as the number of mixture components increases, and we can roughly view a latent variable т and q as an infinite mixture of Gaussian distributions.
model with Gaussian-distribus
Proposition 1. For any d-dimensional continuous distribution p" (x), there exist sequence of two-level latent variable model pa pa(x)) = f Pr. (zl2))p, (z) dz, E 4 that converge to it, where both (x])) and (2) are factorized Gaussian distributions with mean and variance parameterized by neural networks.
Proof can be found in Appendix D.I.
3.1.1 WORLD MODELS INDUCE LATENT VARIABLE POLICIES
Perhaps unsurprisingly, latent variables already exist part of many reinforcement learning works. In particular, in the construction of prohabilistic world models, used when the environment is highly complex or only partially observable. Some works only use intermediate components of the world model as a deterministic input to their policy distribution (e.g. Lee et al. (2020a)), disregarding the distributional aspect, while other approaches use iterative methods for producing an action (e.g. Hafner et al. (2020)). We instead simply view the world model for what it is a latent state inference model- which naturally induces a latent variable policy,
T Taa acc, xsr)= f π (s) Xe)) ds;.
f
(6)
This follows the form of Equation 5 where the context includes the entire history, i.e. h, (act, Xgx). Note that π aa, ss)) conditions only the current latent state due to a Markov assumption typically used in existing world models (see Figure 2), though our algorithms easily extend non- Markov settings as well. Furthermore, this marginalizes over the full latent history due to the recurrence,
a+
SY
a1-1, / / -t-.--//---- 1
g( x(s)))))
which when recursively applied, we can see that the belief state sn-and hence the policy-marginaliann over the entire history of belief states. A more thorough discussion is in Appendix B.2.
Figure 2: Graphi-
cal model of POMDP (solid), world model, and induced latent vari- able policy (dashed).
Our approaches for handling latent variables are agnostic to what y conditions
on, to unify and simplify notation, we use the shorthand h, & (ace, Xx() to denote history information. This subsumes the MDP setting, where (S,, (h)) is equivalent to q(s/ xx) due to Markovian conditional independence.
3.2 MAXENT RL IN THE PRESENCE OF LATENT VARIABLES
The presence of latent variables makes training with the maximum entropy objective (equations I and 3) difficult. Firstly, it requires an accurate estimation of the entropy term, and the entropy of a latent variable model is notoriously hard to estimate due to the intractability of marginalization (Paninski, 2003; Lim et al., 2020). Secondly, the use of latent variables results in an increase in gradient variance, which we remedy with variance reduction methods at a negligible cost. Finally, the appearance of
4

---

Published as a conference paper at ICLR 2023
latent variables can also be used within the Q-function better aggregate uncertainty. For each, derive principled methods for handling latent variables, while the end result is actually fairly simple and only adds minimal amount of extra cost compared to non-latent variable policies.
3.2.1 ESTIMATING THE MARGINAL ENTROPY
An immediate consequence of using latent variables is that the entropy, or marginal entropy, becomes intractable, due the proobababylltyy being intractable, i.e.
H(T((-Ib)) = (a.)) [ log /
π ((s ds;
(8)
Failure cases of naïve entropy estimation Applying methods developed for amortized variational inference (Kingma & Welling, 2013; Burda et al., 2016) can result in a bound on the entropy that is in the wrong direction. For instance, the standard evidence lower bound (ELBO) results in an entropy estimator,
(9) where q is any variational distribution, for example setting ((s:Ja, he) = q(s; he). Adopting this naive estimator will result in maximizing upper bound on the MaxEnt RL objective, which we can see by writing out the error.
Hnavee (hr) A Ex(a, (h.) Er. ha.) [- log ग (aa) [s:)assss))))))))> + log ((s ar, h₂) log ( hh)]],
Hmuive (h,) = H(+(-Iba)) + En(aa/he)) DKL aa he)llgss//a,, h,)]
(10)
Therefore, replacing the entropy in the MaxEnt RL objective (Equa- tion 1) with Hnaive will lead to maximizing the error-i.e. the KL the variational distribution to be as far as it can from the true posterior (s, ar, h₂). Furthermore, this error is unbounded so it may become arbitrarily large without actually affecting the true entropy we want to be maximizing, H (h,)). which leads to serious numerical instability issues. In Figure 3, we show the results from a preliminary experiment where this approach to entropy estimation during policy optimization led to extremely large values (scale of 1028), significantly overestimating the true entropy, and resulted in policies that did not learn. More details are in Appendix C. To overcome this overestimation issue, we propose the following method for achieving accurate estimation.
Naive BO]
Naive (MAE)
Entropy estimate
Ground truth
1000 2000 4000 5000
3000
Figure 3: Training with naïve entropy estimators results in extremely loose upper bounds.
Lower bounding the marginal entropy with nested estimator To be amenable to entropy max- imization, we must construct lower bound estimator of the marginal entropy. For this, inspired by advances in hierarchical inference (Yin Zhou, 2018; Soboley & Vetrov, 2019), a method to estimate the marginal entropy (Equation 8) via a lower bound can be obtained. Specifically, for any K E N, we define
)
Hx Hx(h:) A E Ea E. ~p( B (g()((b))) ((.)))(e)))b))))
log
Σ
γ (acls/ (k)
K+
(11)
where p(s: la; h:) is the (unknown) posterior the policy distribution; however, we can easily sample from this by first sampling s/ then sample a, conditioned on ( This results in nested estimator where we effectively sample K + 1 times from q(s) ht), use only the first latent variable s for sampling the action, while using all the latent variables to estimate the marginal entropy. Note that this is not equivalent to replacing the expectation inside the logarithm with independent samples, which would correspond to an IWAE estimator (Burda et al., 2016). Equation 11 results in a nested estimator that is monotonically increasing in K, which in the limit, becomes an unbiased estimator of the marginal entropy, i.e. (h,)) < H [=([) Hx (h,) & H (h;) [h:)]. Thus, replacing the marginal entropy with HK results in
and lim Hx (h,) maximizing a tight lower bound on the MaxEnt RL objective, and is much more numerically stable in practice. Proofs for these results are in Appendix D. In practice, find that using reasonable values for does not increase computation times since sampling multiple times is easily done in parallel, and the evaluation of n rr (s.) is cheap relative to other components such the world model.
H [=
5

---

Published as a conference paper at ICLR 2023
3.2.2 VARIANCE REDUCTION WITH ANTITHETIC MULTI-LEVEL MONTE CARLO
While latent variable policies can optimize for the MaxEnt RL objective better in expectation, its reliance on stochastic estimation techniques introduces additional gradient variance. This higher variance actually results in poorer sample efficiency, negating any gains obtained from using more flexible distribution. In particular, it has been shown that multi-sample estimators like Equation 11 can result more noise than signal as K increases (Rainforth et al., 20) To remedy this, we adopt a simple yet reliable variance reduction method referred to as antithetic multi-level Monte Carlo (MLMC). While this method has been used in simulations of stochastic differential equations (Giles, 2008; Giles & Szpruch, 2014) and more recently, in variational inference (Ishikawa & Goda, 2021; Shi & Cornish, 2021), it has not yet seen uses in the context of reinforcement learning.
Applying MLMC to the estimator in Equation 11, we have
[K)]
H2 H21 - (FL 1 H
if = O.
MLMMC KK
Σ
Azze
where AN2
(12)
otherwise.
At the f-th level, after we have generated 2 ii.d. samples, we use half compute Hg 1. the other half to compute 7/ and all of the samples to compute Hge. This antithetic sampling scheme is key ingredient in reducing variance, and can achieve the optimal computational complexity for a given accuracy (Ishikawa & Goda, 2021). We compute all the A72z terms in parallel in our implementation, so there is an almost negligible additional cost compared to H.. The only consideration involved for using H HMLMC is that K should be power of two.
3.2.3 ESTIMATING THE MARGINAL Q-FUNCTION
Under the POMDP setting, we aim to build Q-function upon the inferred belief states S, as these contain the relevant dynamics and reward information. However, while most existing approaches such as Lee et al. (2020a); Hafner et al. (2020) only take one sample of the inferred latents as input of the Q-function, we propose marginalizing out the latent distribution in the critic calculation. This can be seen by interpreting the Q-function through the probabilistic inference framework in Levine (2018), where the reward function is viewed as the log-likelihood of observing binary optimality random variable O, i.e. satisfying p(O, 1|s,, ar) exp(r(s;, (s),))))).. a)) As a result, the Q-function is equivalent to Q(s;, at) = log p(OrT|ss, az). Since our latent belief state represents the uncertainty regarding the system in the context of POMDPs, including the current state and unseen rewards, we propose marginalizing the value function over the belief state. Hence, through this probabilistic interpretation, the marginal Q-function is related to the Q-function over latent states through
Q(h,, ar) = log f p(Ou.Tis; (s,, a (ss hh)) ds; = log / exp (Q(s,, (q(s (he) ds..
(13)
Given this, we hence propose the following estimator to be used during policy optimization,
K
Q(by ax) = Qx at) & log
Σ exp (
(K) a)))) si (I:K) ~ q(s hb),
(14)
K
where Q(si, a.) trained to satisfy the soft Bellman equation in Equation 2. A closely related approach is from Lee et al. (2020a) who similarly trains Q-function on latent states; however, they directly use Q(s,, ax) during policy optimization, which is a special case with K = 1, whereas using K > results in closer approximation to the marginal Q-function. We found this construction for marginalizing the Q-function to be useful mainly when used in conjunction with a world model.
3.3 STOCHASTIC MARGINAL ACTOR-CRITIC (SMAC)
While the above methods can each be applied to general MaxEnt RL algorithms, we instantiate concrete algorithm termed Stochastic Marginal Actor- (SMAC). SMAC is characterized by the use of a latent variable policy and maximizes a lower bound marginal MaxEnt RL objective. Specifically, we use the same method as SAC to train Q(s,, a,) on latent states, but we train the policy using a low-variance debiased objective for taking into account latent state marginalization,
max EAi~P.
@x (be, ax) + a MMLMC (h:)
(15)

---

Published as a conference paper at ICLR 2023
We train the inference model ((s) with standard amortized variational inference (Equation 4), and we train only π (ax))s) using the objective in Equation 15. When not used with a world model, we train both 5T (s,))s and ((s,) h,) using Equation 15. See Algorithms and 2 for a summary of the training procedures, and more details regarding implementation for SMAC in Appendix B.
4 RELATED WORK
Maximum entropy reinforcement learning Prior works have demonstrated multiple benefits of MaxEnt RL, including improved exploration (Han & Sung, 2021), regularized behaviors (Neu et al., 2017; Vieillard et al., 2020a), better optimization property (Ahmed et al., 2018), and stronger robustness (Eysenbach & Levine, 2022). Generally, policies optimizing the MaxEnt RL objective sample actions that are proportional to the exponentiated reward, and alternatively can be viewed as a noise injection procedure for better exploration (Attias, 2003; Ziebart, 2010; Haarnoja et al., 2017; Nachum et al., 2017; Levine, 2018; Abdolmaleki et al., 2018; Haarnoja et al., 201 8b; Vieillard et al., 2020b; Pan et al., 2022; 2023; Lahlou et al., 2023). However, this noise injection is commonly done directly in action space, leading to only local perturbations, whereas we inject noise through nonlinear mapping.
Latent variable modeling The usage of latent variable models originates from graphical mod- els (Dayan et al., 1995; Hinton et al., 2006) and has been recently popularized in generative model- ing (Kingma & Welling, 2013; Rezende et al., 2014; Zhang et al., 202 2022b;a). The estimation of the log marginal probability and marginal entropy has long been a central problem in Bayesian statistics and variational inference (Newton, 1994; Murray & Salakhutdinov, 2008; Nowozin, 2018; Ishikawa & Goda, 2021; Malkin al., 2022). However, most of these works consider a lower bound on the log marginal probability for variational inference, which is not directly applicable to maximum entropy as discussed in Section 3.2.1. few works have proposed upper bounds (Sobolev & Vetrov, 2019; Dieng et al., 2017) or even unbiased estimators (Luo et al., 2020), and while we initially experimented with a couple of these estimators, we found that many results in high gradient variance and ultimately identified an approach based on hierarchical inference technique for its efficiency and suitability in RI.
Latent structures in POMDPs and world models Settings with only partial observations are natural applications for probabilistic inference methods, which help learn latent belief states from observational data. As such, variational inference has been adopted for learning sequential latent variable models (Ghahramani & Hinton, 2000; Krishnan et al., 2015; Fraccaro et al., 2016; Karl et al., 2017; Singh et al., 2021). One paradigm is to use the learned recurrent model to help model-free RL. algorithms (Wahlstrom et al., 2015; T'schiatschek et al., 2018: 8; Buesing et al., 2018: 8; Igl et al., 2018; Gregor et al., 2019; Han et al., 2020). Another approach is to use world models for solving POMDP and building model-based RL agents (Deisenroth & Rasmussen, 2011; Hausknecht & Stone, 2015; Watter et al., 2015; Zhang et al., 2019; Hafner et al., 2019; 2020; 2021; Nguyen et al., 2021; Chen et al., 2022) due to their planning capabilities. It is also sometimes the case that the world model is treated mainly as a representation, without much regard for the measure of uncertainty (Ha & Schmidhuber, 2018; Schrittwieser al., 2020; Amos et al., 2021; Hansen al., 2022).
EXPERIMENTS
We evaluate SMAC on a series of diverse continuous control tasks from DeepMind Control Suite (DMC; Tassa et al. (2018)). These tasks include challenging cases in the sense of having sparse rewards, high dimensional action space, or pixel observations. We also perform a preliminary unit test with a
reward in Appendix B.
5.1 STATE- BASED CONTINUOUS CONTROL ENVIRONMENTS
Setting We first compare SMAC with SAC and TD3 (Fujimoto et al., 2018) baselines on variety of state-based environments to demonstrate the advantage of latent variable policies. We show eight environments in Figure 4, and leave more results in Appendix due space limitations.
Results We find that even in the simple MDP setting, we improve upon SAC by simply introduc- ing a latent variable. Specifically, our method is almost never worse than SAC, implying that the extra gradient variance from the entropy estimation is not incurring a penalty in terms of sample efficiency.
7

---

Published as a conference paper at ICLR 2023
cheetah,
caripole wwwngup.ppppocco... spaars
finger,
hopper stand
800
800
8:00
Episode Return
Epicade Buturn
600
600
600
400
Algorithm
300
Episode
400
Episode
6:00
SMAC
200
SAC
200
200
TD3
0.0
0.2
0.4
0.6
0.8
1,0
0.0
0.2
0.4
0.4
0.0
1.0
0.0
0.2
0.4
0.5
0.8
1.0
0.0
0.2
0.4
366
0.8
1.0
Step
1c6
Step
1e6
Step
166
Step
155
heopperr hop
nguu
quadruped
1030
250
800
Eplode Return
200
Epicode Ruturn
600
MI
800
8.00
150
800
620
400
100
400
400
200
200
200
0
1.0 0
0.2
0.6
1.0
0.2
0.6
1.0
0.4 0.5
1.0
0.0
0.4
2.6
Step
Step
Step
Step
Figure 4: Experiments on eight DMC environments where agents are given state-based inputs. The SMAC approach improves upon SAC with better exploration and more robust training.
By being able to track a wider frontier of optimal action trajectories, SMAC can be more robust to find optimal policies, particularly when the reward sparse (e.g., cartpolee.singgp..pp...
Comparison with other probabilistic modeling approaches We further conduct extensive pirical comparisons with other probabilistic policy modeling methods including normalizing flow and mixture-of-experts (Ren et al., 2021) based SAC methods in Figure 11. Our proposed SMAC generally achieves the best sample efficiency on almost all environments. Due to limited space, we defer related discussion to Appendix C.
Marginalization Marginalizing over the latent state has significant effect on training, though this often exhibits a diminishing rate, suggest- ing that using reasonable number of particles is sufficient. Figure shows this behavior for the quadruped wall k task.
Hortiiees
Variance-reduced updates We find that the
Stne
use of MLMC as a variance reduction tool is
(a) Marginalization
(b) MLMC
crucial for the latent variable policy to perform well in some difficult environments such as the quadruped. wwtrruppeescccee escape task. Figure 5b shows that using MLMC clearly reduces variance and makes training much more robust, whereas using only the nested estimator performs closer to the baseline SAC (see comparison in Figure 10).
Figure 5: Ablation experiments.
5.2 PIXEL-BASED CONTINUOUS CONTROL ENVIRONMENTS
Setting We next compare different algorithms series of DMC environments with pixel-based input. Since this task is much more challenging, and pixels only provide partial observability. we make use of a world model (as described in Section 2.2) to supplement our algorithm. We use the recurrent state-space model (RSSM) architecture from Hafner et al. (2019) as the world model. We refer to this baseline as "Latent-SAAC S and follow the practice in Wang et al. (2022), which samples from the belief distribution g x<t) and directly trains using SAC on top of the belief state. A closely related work, SLAC (Lee et al., 2020a), only uses S, as input to learned Q-function, while the policy does not use S; and instead uses intermediate layers of the world model input. Finally, RL (MBRL) algorithm that performs rollouts on the
we also compare to Dreamer, dynamics model (Hafner et al., 2020). This iterative procedure results in higher computational cost as it requires iteratively sampling from the belief state and differentiating through the rollouts. In contrast, proposed SMAC aggregates samples from the current belief state and does require
8

---

Published as a conference paper at ICLR 2023
ball in cup
cartpoli swingup
finger spin
1000
800
Episodl Return
300
Episode Return
Episode Return
Epicode Raturrn
600
Algenthim
SMAC
409
Latent SAC
200
SLAC
200
MBAL
0.0
0.2
0.4
3.4
0.8
1,0
0.5
0.2
0.4
0.5
0.8
1.0
0.0
0.2
0.4
3.6
0.8
1.0
0.0
0.2
0.4
0.5
0.8
1,0
Step
166
Stop
166
Step
166
Step
165
happer
easy
quadruped
quadrupeed
50
Eptole Return
800
Episade Ruturn
800
Episode Return
800
40
600
600
20
400
4.00
20
200
200
10
0.2
0.8
0.8
0.0
0.4 0.8 1.0
0.0
0.4
1.0
Step
5tep
Step
1c6
Step
Figure Experiments on eight DMC environments where agents are given pixel-based inputs. The proposed SMAC approach achieves better expected return than similar Latent-SAC. and SLAC baselines. Notice that our method does not involve any planning, but still achieves comparable (sometimes even better) performance to the model-based RL algorithm.
Table 1: Experiments with noisy observations. We experiment with two different perturbation level for two kinds noise. GAUSSIAN PERTURBATION adds independent white noise to the pixel images while SENSOR MISSING randomly turns portion of the pixels to black. SMAC is trained with world model while L-SAC denotes a SAC baseline trained with a world model.
GAUSSIAN PERTURBATION
SENSOR MISSING
NOISE
SMALL
LARGE
SMALL
LARGE
METHOD
L-SAC
SMAC
L-SAC
SMAC
L-SAC
SMAC
L-SAC
SMAC
FINGER
912:138
959: 30
880:1109
924 43
933 63
955 36
921 44
921
HOPPER
671::111
731:1:077
516:1385
702 1391
721 16
872 02
703
866 61
REACHER
869 28
928 10
782
925
883
937 89
854
925 50
differentiating through the dynamics model. For training, we follow Hafner et al. (2020) and repeat each action 2 times. We show the comparison on eight tasks in Figure 6, and again relegate more results to the Appendix due to space constraints.
Results Comparing SMAC to the Latent-SAC baseline, we again find that we can often find an optimal policy with fewer environment interactions. We find that SLAC and Latent-SAC are roughly on par, while SLAC can also some times perform worse, as their policy does not condition on the latent state. The model-based approach has widely-variable performance when compared to the actor-critic approaches. Interestingly, in most of the environments where the model-based approach is performing well, find that SMAC can often achieve comparable performance, though it does not make use of planning. Overall, we find that our method improves upon actor-critic approaches and bridges the gap planning- approaches.
Robustness to noisy observations While the pixel-based setting already provides partial observa- tions, we test the robustness of our approach in settings with higher noise levels. In Table 1 we report sinnenrsppii
the episodic rewards of both SMAC and SAC baseline on three environments hopper.ste asaan, and rea reacher.ed checceesry) under noisy perturbations and missing pixels (Meng et al., 2021). We find that SMAC behaves more robustly than the baseline across almost all settings.
Efficiency Despite the extra estimation procedures, SMAC does not incur significant computational costs as we can compute all terms in the estimators in parallel. Tested with an NVIDIA Quadro GV 100 on the pixel-based environments, SMAC implementation does 60 frames per second (FPS)

---

Published as a conference paper at ICLR 2023
average, almost the same training speed compared to Latent-SA4 SAC (63 FPS), whereas differentiating through a single rollout over the dynamics model already reduces to 51 FPS (roughly 20% slower).
6 CONCLUSION
We propose methods for better handling of latent variable policies under the MaxEnt RI. framework, centered around cost-efficient computation and low-variance estimation, resulting in tractable algorithm SMAC when instantiated in the actor-critic framework. We find that SMAC can better make use the belief state than competing actor-critic methods and can more robustly find optimal policies, while adding only very minimal amounts of extra compute time.
ACKNOWLEDGEMENT
The authors would like thank Zhixuan Lin. Tianwei Ni, Chinwei Huang, Brandon Amos, Ling Pan. Max Schwarzer, Yuandong Tian, Tianjun Zhang, Shixiang Gu, and anonymous reviewers for helpful discussions. Dinghuai also expresses gratitude towards his fellow interns FAIR for creating a lot of joyful memories during the summer in New York City.
REFERENCES
Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Rémi Munos, Nicolas Manfred Otto Heess, and Martin A. Riedmiller. Maximum a posteriori policy optimisation. ArXiv, abs/1 abs/1800.,6990, 2018.
Zafarali Ahmed, Nicolas Le Roux, Mohammad Norouzi, and Dale Schuurmans. Understanding the impact of entropy policy learning. 2018.
Brandon Amos, Samuel Stanton, Denis Yarats, and Andrew Gordon Wilson. On the model-based stochastic value gradient for continuous reinforcement learning. In Learning for Dynamics and Control, pp. 6-20. PMLR, 2021.
Karl Johan Aström. Optimal control of markov processes with incomplete state information. Journal of Mathematical Analysis and Applications, 10:174-205, 1964.
Hagai Attias. Planning by probabilistic inference. In AISTATS, 2003.
Lars Buesing, Théophane Weber, Sébastien Racanière, S. M. Ali Eslami, Danilo Jimenez Rezende, David P. Reichert, Fabio Viola, Frederic Besse, Karol Gregor, Demis Hassabis, and Daan Wierstra. Learning and querying fast generative models for reinforcement learning. ArXiv, abs/1 00/11002000000 2018.
Yuri Burda, Roger Baker Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. CoRR,
2016.
Ricky T. Q. Chen, Jens Behrmann, David Kristjanson Duvenaud, and föm--Henil Jacobsen. Residual 2019.
flows for invertible generative modeling. ArXiv,
Xiaoyu Chen, Yao Mu, Ping Luo, Sheng Li, and Jianyu Chen. Flow-based recurrent belief state learning for pomdps. In ICML, 2022.
Peter Dayan, Geoffrey E. Hinton, Radford M. Neal, and Richard S. Zemel. The helmholtz machine. Neural Computation, 7:889-904, 904, 1995.
Mare Peter Deisenroth and Carl Edward Rasmussen. Pilco: A model-based and data-efficient approach to policy search. In /CML, 2011.
Adji Bousso Dieng, Dustin Tran, Rajesh Ranganath, John Paisley and David Blei. Variational inference via X upper bound minimization. Advances in Neural Information Processing Systems, 30, 2017.
10

---

Published as a conference paper at ICLR 2023
Laurent Dinh, Jascha Narain ArXiv, 0as/16055088833 2017.
and Samy Bengio. Density estimation using real nyp.
Benjamin Eysenbach and Sergey Levine. Maximum entropy (provably) solves some robust rl problems. ArXiv, abs/2103.06257, 2022.
Marco Fraccaro, Søren Kaac Sønderby, UIrich Paquet, and Ole Winther. Sequential neural models with stochastic layers. ArXiv, abs/ leds.0077,77 2016.
Scott Fujimoto. Herke van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. ArXiv, abs/1802.09947, 2018.
Zoubin Ghahramani and Geoffrey Hinton. Variational learning for switching state-space models. 2000.
Neural Compulation,
Michael B Giles. Multilevel monte carlo path simulation. Operations research, 56(3):607-617, 2008.
Michael B Giles and Lukasz Szpruch. Antithetic multilevel monte carlo estimation for multi- dimensional sdes without lévy area simulation. The Annals of Applied Probability, 24(4): 1620, 2014.
Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aäron den Oord. Shaping belief states with generative environment models for rl. ArXiv, abs/1906 1906.09237., 2019.
David R Ha and Jürgen Schmidhuber. World models. ArXiv, abs/ htts//883.100222. 2018
Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International conference machine learning, pp. 1352-1361. PMLR, 2017.
Tuomas Haarnoja, Kristian Hartikainen, Picter Abbeel, and Sergey Levine. Latent space policies for hierarchical reinforcement learning. In International Conference Machine Learning, pp- 1851-1860. PMLR, 2018a.
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, al. Soft actor-critic algorithms and applications. arXiv preprint arXiv:
2018b.
Danijar Hafner, Timothy P. Lillicrap, Ian S. Fischer, Ruben Villegas, David R Ha. Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. ArXiv, abs/181 004551, 2019.
Danijar Hafner, Timothy P. Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. ArXiv, abs/1912.01 12361163,,, 2020.
Danijar Hafner, Timothy P. Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. ArXiv, abs/2010.02193,, 2021.
Dongqi Han, Kenji Doya, and Jun Tani. Variational recurrent models for solving partially observable control tasks. ArXiv,
2020.
Seungyul Han and Youngchul Sung. A max-min entropy framework for reinforcement learning. In NeurlPS, 2021.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. In /CML. 2022.
Matthew Hausknecht and Peter Stone. Deep recurrent q- q-learning for partially observable mdps. In AAAI Fall Symposia, 2015.
Elad Hazan, Sham Kakade, Karan Singh, and Abby Van Soest. Provably efficient maximum entropy exploration. International Conference on Machine Learning, pp. 2681-2691. PMLR, 2019.
Geoffrey E. Hinton, Simon Osindero, and Yee Whye Teh. A fast learning algorithm for deep belief' nets. Neural Computation, 18:1527-1554, 2006.

---

Published as a conference paper at ICLR 2023
Maximilian Igl, Luisa M. Zintgraf. Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational reinforcement learning for pomdps. In ICML, 2018.
Kei Ishikawa and Takashi Goda. Efficient debiased evidence estimation by multilevel Monte Carlo sampling. In Cassio de Campos and Marloes H. Maathuis (eds.), Proceedings of the -hirty-Seventh Conference on Uncertainty in Artificial Intelligence. volume 161 of Proceedings of Machine Learning Research, pp. 34 43. PMLR, 27-30 Jul 2021.
Herman Kahn. Use of different Monte Carlo sampling techniques. Rand Corporation, 1955.
Maximilian Karl, Maximilian Sölch, Justin Bayer, and Patrick van der Smagt. Deep variational bayes filters: Unsupervised learning of state space models from raw data. ArXiv, abs/1600.00632, 2017.
Diederik P Kingma and Max Welling.
Auto-encoding variational bayes.
arXiv preprint
arXiv:1312.6114 14,66111400 2013.
Rahul G. Krishnan, Uri Shalit, and David A. Sontag. Deep kalman filters. ArXiv, abs/151 2015.
Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hern's andez- Gare'ia. ii, L'ena N'chale Ezzine, Yoshua Bengio, and Nikolay Malkin. A theory of continuous generative flow networks. ArXiv, abs/2301.12594, 2023.
Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine, Stochastic latent actor-critic: Deep reinforcement learning with latent variable model. Advances in Neural Information Processing Systems, 33:741-752, 2020a.
Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic. https: /ttp::://githht wtt.://...ommma.e.ae...eeeeeeeee /a/ slac 2020b.
Sergey Levine. Reinforcement learning and control probabilistic inference: Tutorial and review.
arXiv preprint 1800v: 00909, 2018
Jae Hyun Lim, Aaron Courville, Christopher Pal, and Chin-Wei Huang. AR-DAE: Towards unbiased neural entropy gradient estimation. In Hal Daumé III and Aarti Singh (eds.). Proceedings the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 6061-6071 PMLR, 13-18 Jul 2020.
Zhixuan Lin. Dreamer-pytorch implementation. https: www..////.hhiiiiinn. dreemee-pytooe 2022.
Yucen Luo, Alex Beatson, Mohammad Norouzi, Jun Zhu, David Kristjanson Duvenaud, Ryan P. Adams, and Ricky T. Q. Chen. Sumo: Unbiased estimation of log marginal probability for latent variable models. ArXiv, abs/2004.00353, 2020.
Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward J. Hu, Katic Elizabeth Everett, Dinghuai Zhang, and Yoshua Bengio. Gflownets and variational inference. ArXiv, abs/2210.00580, 2022.
Russell Mendonea, Oleh Rybkin, Kostas Danilidis, Danijar Hafner, and Deepak Pathak. Discovering and achieving goals via world models. ArXiv, abs/2 I 10. 14, 2021.
Lingheng Meng, Robert B. Gorbet, and Dana Kuli' Memory-based deep reinforcement learning for pomdps. 2021 JEEE/RSJ International Conference Intelligent Robots and Systems (IROS), pp. 5619-5626, 2021.
Laurence Illing Midgley, Vincent Stimper, Gregor NC Simm, Bernhard Schölkopf, and José Miguel Hemández-Lobato. Flow annealed importance sampling bootstrap. arXiv preprint arXiv:2208.01899 2022.
lain Murray and Ruslan Salakhutdinov. Evaluating probabilities under high-dimensional latent variable models. In NIPS, 2008.
Ofir Nachum, Mohammad Norouzi, Kelvin Xu, and Dale Schuurmans. Bridging the gap between value and policy based reinforcement learning. Advances neural information processing systems, 30, 2017.
12

---

Published as a conference paper at ICLR 2023
Gergely Neu, Anders Jonsson, and V. Gómez. A unified view of entropy-regularizg markov decision 2017.
processes. ArXiv,
Michael A. Newton. Approximate bayesian-inference with the weighted likelihood bootstrap. Journal of the royal statistical society series b-methodological, 56:3-26, 1994.
Tung D. Nguyen, Rui Shu, Tu Pham, Hung Hai Bui, and Stefano Ermon. Temporal predictive coding for model-based planning in latent space. ArXiv, 00os/106.07156, 2021.
Mohammad Norouzi, Samy Bengio, Navdcep Jaitly, Mike Schuster, Yonghui Wu, Dale Schuurmans, et al. Reward augmented maximum likelihood for neural structured prediction. Advances In Neural Information Processing Systems, 29, 2016.
Sebastian Nowozin. Debiasing evidence approximations: On importance--weighed autoencoders and jackknife variational inference. In ICLR, 2018.
L. Pan. Dinghuai Zhang, Aaron C. Courville, Longbo Huang, and Yoshua Bengio. Generative augmented flow networks. ArXiv, obs/2210.003388, 2022.
L. Pan, Nikolay Malkin, Dinghuai Zhang, and Yoshua Bengio. Better training of gflownets with local credit and incomplete trajectories. ArXiv, abs/2302.01687, 2023.
Liam Paninski. Estimation of entropy and mutual information. Neural computation, 2003.
Dmytro Perekrestenko, Stephan Müller, and Helmut Böleskei. Constructive universal high- dimensional distribution generation through deep relu networks. ArXiv, abs/2006. 16664, 2020.
Martin L Puterman. Markov decision processes. Handbooks in operations research and management science, 2:331-434, 1990.
Tom Rainforth, Adam Kosiorek, Tuan Anh Le, Chris Maddison, Maximilian Igl, Frank Wood, and Yee Whye Teh. Tighter variational bounds are not necessarily better. In International Conference on Machine Learning, pp. 4277-4225 PMLR, 2018a.
Tom Rainforth, Yuan Zhou, Xiaoyu Lu, Yee Whye Teh, Frank Wood, Hongseok Yang, and Jan- Willem van de Meent. Inference trees: Adaptive inference with exploration. arXiv preprint /7X.v:8806.09500 2018b.
Konrad Rawlik, Mare Toussaint, and Sethu Vijayakumar. On stochastic optimal control and rein- forcement learning by approximate inference. Proceedings of Roboties: Science and Systems VIII. 2012.
Jie Ren, Yewen Li, Zihan Ding, Wei Pan, and Hao Dong. Probabilistic mixture-of-experts for efficient deep reinforcement learning. ArXiv, abs/2104.09122, 2021.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In JCML, 2014.
Julian Schrittwieser, Joannis Antonoglou, Thomas Hubert, Karen Simonyan, L.. Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy P. Lillicrap, and David Silver. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588 7839: 604-609, 2020.
Yuyang Shi and Rob Cornish. On multilevel monte carlo unbiased gradient estimation for deep latent variable models. In Arindam Banerjee and Kenji Fukumizu (eds.), Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pp. 3925-3933, PMLR, 13-15 Apr 2021.
Gautam Singh, Skand Vishwanath Peri, Junghyun Kim, Hyunseok Kim, and Sungjin Ahn. Structured world belief for reinforcement learning in pomdp. In ICML, 2021.
Artem Sobolev and Dmitry P. Vetrov. Importance weighted hierarchical variational inference. In NeurIPS, 2019.
13

---

Published as a conference paper at ICLR 2023
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Pranjal Tandon. Pytorch implementation of soft actor-critic. https: http:/twww........gggg./oo.....a Sytorch-soft--eee
critic, 2020.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefraneq. Timothy P. Lillicrap, and Martin Riedmiller. Deepmind control suite. ArXiv, abs/180) 801.00690, 2018.
Emanuel Todorov. sinearly-solvabll markov decision problems. Advances in neural information processing systems, 19, 2006.
Sebastian Tschiatschek, Kai Arulkumaran, Jan Stühmer, and Katja Hofmann. Variational inference for data-efficient model learning in pomdps. ArXiv, abs/1 80500921, 2018.
Nino Vieillard, Tadashi Kozuno, Bruno Scherrer, Olivier Pietquin. Rémi Munos, and Matthieu Geist. Leverage the average: an analysis of regularization in rl. ArXiv, abs/2003. 14089, 2020a.
Nino Vieillard, Olivier Pietquin, and Matthieu Geist. Munchausen reinforcement learning, ArXiv, abs/2007. 14430, 2020b.
Niklas Wahlstrom, Thomas Bo Schön, and Mare Peter Deisenroth, From pixels to torques: Policy learning with deep dynamical models. ArXiv, abs/1502.02251, 2015.
Tongzhou Wang. Simon Shaolei Du, Antonio Torralba, Phillip Isola, Amy Zhang, and Yuandong Tian. Denoised mdps: Learning world models better than the world itself. ArXiv, abs/2206.154777 2022.
Patrick Nadeem Ward, Ariella Smofsky, and Avishek Joey Bose. Improving exploration in soft-actor- critic with normalizing flows policies. arXiv preprint arXhv:1906.02771, 2019.
Manuel Watter, Jost Tobias Springenberg. Joschka Boedecker, and Martin A. Riedmiller. Em- bed to control: A locally linear latent dynamics model for control from raw images. ArXiv, 00111506.77666. 2015.
Kai Xu, Akash Srivastava, and Charles Sutton. Variational russian roulette for deep bayesian nonparametrics. In ICML, 2019.
Mingzhang Yin and Mingyuan Zhou. Semi-implicit variational inference. In International Conference on Machine Learning, pp. 5660-5669. PMLR, 2018.
Dinghuai Zhang, Jie Fu, Yoshua Bengio, and Aaron C. Courville. Unifying likelihood-free inference with black-box optimization and beyond. In International Conference Learning Representations, 2021a.
Dinghuai Zhang, Ricky T. Q. Chen, Nikolay Malkin, and Yoshua Bengio. Unifying generative models with gflownets. 2022a.
Dinghuai Zhang, Nikolay Malkin, Liu, Alexandra Volokhova, Aaron C. Courville, and Yoshua Bengio. Generative flow networks for discrete probabilistic modeling. In International Conference on Machine Learning. 2022b.
Marvin Zhang, Sharad Vikram, Laura Smith, P. Abbeel, Matthew J. Johnson, and Sergey Levine Solar: Deep structured representations for model-based reinforcement learning. In ICML, 2019.
Michael Zhang, Tom Le Paine, Ofir Nachum, Cosmin Paduraru, George Tucker, Ziyu Wang, and Mohammad Norouzi. Autoregressive dynamics models for offline policy evaluation and optimization. arXiv preprint urXiv:2104. 13877, 2021b.
Brian D. Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy. 2010.
14

---

Published as a conference paper at ICLR 2023
A
NOTATIONS
Symbol
Description
X₁
Environment state (human designed feature or pixel image) at time
a,
Action at time &
h,
History up to time t. defined as (a.c,
Reward
Policy (distribution over action)
Q
Q-function
Q
Q-function estimator
Discount factor
O
Optimality binary random variable as defined in Levine (2018)
8
Domain of state
A
Domain of action
H()
Entropy
H())
Entropy estimator
Temperature for MaxEnt RL
St
Latent variable inferred belief state
((s xx)
Belief distribution over inferred states from observation X
(ss act xgi)
Belief distribution over inferred states from past data Unifies notation for the above two
(5 h
n (a, St
Policy conditioned on latent variable St
T (a Xx.
Latent variable policy, equals f T(a (s)x) ds
T a actx
xg1)
Latent variable policy, equals / IT is (s))g(s (s xx.) ds;
[a; h:)
Unifies notation for the above
p(X, (8.)
Observation model in the world model Reward model in the world model Transition model, also the prior of a world model Inferred posterior dynamics model of learned world model
p(r:
p(S St, ar)
( a a1
B ADDITIONAL DETAILS REGARDING METHODOLOGY
B. MULTI-MODALITY OF LATENT VARIABLE POLICIES
While a latent variable policy can theoretically model any distribution (see Proposition 1), training this policy can still be difficult, especially if the true reward is actually multi-modal. Here, we test in a controlled setting, whether our method can truly recover a multi-modal policy.
A standard interpretation of the MaxEnt RL objective is reverse KL objective (Levine, 2018), [r(x, a) log π (a[x))
max Eax
(16)
r(x, a)
< max (pix)x(alx))
log ग (alx)
(17)
a
a
r(xx, a)
max m (((x)(()))) *
logexp exp min Ep(x) DKL. (TTaa/))lp** [aax))))
log 77 (all))
(18)
a
(19)
where p" (alx) x exp ((x.a)) }. i.e., target distribution defined by the exponentiated reward function and annealed with temperature
Despite the ubiquity of the reverse KL objective-such as appearing standard posterior inference- the training of latent variable models for this objective is still relatively under-explored due to the difficulty in estimating this objective properly. Luo et al. (2020) showed that using improper bounds on the objective can lead to catastrophic failure, but only showed successful training for unimodal target distribution, while Sobolev & Vetrov (2019) discussed proper bounds but did not perform such experiment.
15

---

Published as a conference paper at ICLR 2023
(a) Target [xxp
{
rixaa)
(b) Latent variable policy
(c) Gaussian policy
Figure 7: Optimizing latent variable policy for one-step multi-modal MaxEnt RL objective.
We experiment by setting reward function that has multiple optimal actions. Using sufficiently large a creates a target distribution with four modes (Figure 7a). In Figure 7b, we show that we can successfully learn multi-mode modal distribution with latent variable policy using the methods discussed in Section 3.2.1 and 3.2.2. On the other hand, a Gaussian policy can only capture one out of four modes (Figure 7c), with the exact mode depending on the random initialization.
B.2 WORLD MODEL LEARNING
In Figure 8 we visualize the graphical model for the RSSM similarly with Hafner et al. (2019)
described Section 2.2. We use solid arrows denote the generative machinery (p the following cquations) and dotted arrows to denote the inference machinery (g in the following equations). A
variational bound for the likelihood on observed trajectory could be written as follows,
log pl (OxTTT (srlagr) 2
BR Exgrra [x rgT, Sgrlaggi
log ((<r X agr))
(20)
Eg
logp(xx(s:) + logp(r:(s:) log (s 1) log 21
(21)
-yy
a log (xg logp(r:(s:) DKL (4
(22)
The world model RSSM is then learned by maximizing Equation 22 with regard to parameters of p(x)s), p(r(s), ((s/|st-1, a,_,, and p(s//s, a4-1)) Note that in Section 2.2 omit the reward modeling part for simplicity. Due to the Markovian assumption on latent dynamics and the shorthand of h, P (ac;)x()) we could also use q(s, h,) to denote q(s, S at-1,
B.3 SMAC ALGORITHM
In this section, present the algorithmic details of SMAC with and without world model in Algorithm and Algo- rithm 2 respectively. Both of the two algorithms follow the commonly adopted off-policy actor-critic style and utilize a replay buffer to save data for the update of both the actor and critic networks (the dependency of buffer D is omitted in the algorithms). Our SMAC based on SAC algorithm, whose critic is trained by minimizing TD error,
a
aa
Bry
Sy
x
X
To (Q(x,a)= (r (r+70(x 1xx (F(((((1x)))
(23)
where (x, x.) D. a' x(-x'), ) Q denotes a stop gradient operator and H an estimate for the policy en- tropy. In our case, estimate the entropy of latent variable
Figure Graphical model of a POMDP
(solid) and a world model (dashed).
16

---

Published as a conference paper at ICLR 2023
policy with Equation 12 as discussed in Section 3.2. What's more, the actor is updated via minimizing
Jx = Q(x a) - (H((T(-))))
(24)
where X ~ D and a x(-xx), which is equivalent to a ~ T (-|ss) s ~ g(s In the algorithm box we omit the moving average of the critic network for simplicity, which is adopted as common-sense. We remark that SMAC has not much difference with SAC in the sense of RL algorithmic details, but mainly achieve improvement with the structured exploration behavior achieved from latent variable modeling. For SMAC in conjunction with world model, we learn the critic network by minimizing TD error on the latent level,
Jo (Q(s.)) - -(r (r+nQ(s', + (+(())))))
(25)
whose terms can be seen as one sample estimate to each term in Equation 23. We also try directly train the critic network on the observation level, but the empirical difference is negligible, As result, we keep the latent level TD learning for simplicity's sake.
We next state a method to encourage exploration through conditional entropy minimization. Entropy in latent variable policy (Equation 5) can be increased by either dispersing the probability density to other modes, or by increasing the entropy at single mode. The latter corresponds to increasing the entropy of the conditional distribution x (a (st), which can end up as a shorteut to increasing entropy and can result in spurious local minima during training. This is in fact a well-known issue that sampling-basel objectives run into (Rainforth et al., 20; Midgley et al., 2022), resulting in a policy that explores the space of action trajectories at slower pace. On the other hand, notice that in the proof of Proposition I in Section D.1 we require the decoder variance to be sufficiently small to be expressive, thus we propose remedying this issue for MaxEnt RL by adding conditional entropy term to the objective:
max E Ep(r)
Y (F((x)) + aHH(T())))) - (ss(s)))))) [H(T(( (s))
(26)
The conditional entropy H((T(-s:)) represents the entropy around single mode. By minimizing this in conjunction with maximizing the marginal entropy, we incentivize the policy to disperse its density to other modes, encouraging it to explore and find all regions with high reward. This allows the latent variable policy to make better use of its source of randomness St, and encourages nonlinear mapping between S, and the action space. This is in stark contrast to entropy maximization with a Gaussian policy, where random noise is simply added linearly and independently to the action. This technique only useful for a few experiments (hopper-hog humanoid run humanoid stand quadruped .eeeaee quadruped run reacher eaac walker_rum run) on state-based model-free experiments. In the pixel-basee experiments where SMAC leverages a world model, the distribution of the latent variable is learned within the world model, thus there is no need to further involve such a regularizer.
Unadopted techniques Our proposed technique is universal and could be applied to any MaxEnt RL algorithm. For example, we also try to combine Dreamer with MaxEnt in the policy optimization part, and wish to further improve the performance with our method. Nonetheless, as stated in the Appendix of Hafner et al. (2020), there is much positive effect of introducing MaxEnt principle into Dreamer implementation. Indeed, we find that doing MaxEnt on the actor of Dreamer will in fact lower down the sample efficiency. Therefore, we think it is not meaningful to put the experiments of this part into this work.
Another unused technique lies in latent variable modeling. MLMC with finite samples (i.e., K < oo) still gives a biased estimator. On the other hand, Russian Roulette estimator (Kahn, 1955) enables an unbiased estimate of infinite series of summation with the help of randomized truncations, together with corresponding term upweighting operation. This technique is also used in many modern machine learning problems (Xu et al., 2019; Chen al., 2019). As result, we also try to introduce Russian Roulette calculation into our MLMC estimator. However, we do not find much evident improvement in our RI. experiments, thus we do not take this technique into the final SMAC algorithm.
C ADDITIONAL DETAILS REGARDING EXPERIMENTS
For all episode return curves, we report the mean over 5 seeds and 95% confidence intervals.
17

---

Published as a conference paper at ICLR 2023
Algorithm 2 SMAC (with a world model)
1:
for each step do
2:
/ //Environment interaction
Algorithm SMAC (without world model)
3:
for t ........TT do
for each step do
4:
at ~ I (t). 8 1(s
2: //Env interaction
5:
Tt, Tt,N t- env. step env.step (a,)
(als)) (({)x) 4: x, - env env.step (a)
6:
end for
7:
D (X (()
D (x,
8:
6:
9:
//World model learning
//Critic learning
10:
Update world model maximize Eq. 22
8:
Update Q to minimize Eq. 23
11:
9:
12:
//Critic learning
10:
//Actor learning
Update minimize Eq. 25
11:
Calculate HM via Eq. 12
14:
12: Update minimize Eq. 24
//Actor learning Calculate HMLMC via Eq. 12
13:
end for
17: Calculate QK via Eq. 14
18: Update policy to maximize Qx + MMMMC
end for
Regarding Figure We run SMAC on DMC finger spin environment with different marginal log probability estimators. We obtain the ground truth value via expensive Monte Carlo estimation with 1 x 10 samples. From the figure, we can see that there little hope of using a naïve upper bound of entropy for MaxEnt RL, where reasonable scale of the entropy term is 10". For IWAE (Burda al., 2016) implementation, we set the number particles to 32. Other values for the number of particles give similar results.
Regarding Section 5.1 For model-free experi- ments, the agents are fed with state-based inputs. We follow the Py Torch model-free SAC imple- mentation of Tandon (2020) for this part. The actor is parametrized by a tanh-Gaussian distri- bution. For our method, we additionally use two layer MLP to parametrize the latent distribu- tion q(s)x). We set the neural network width of the baselines and SMAC to 400 and 256 respec- tively keep comparable number of parameters. For the entropy coefficients, we use the same autotuning approach from SAC (Haarnoja et al., 2018b). We follow Fujimoto al. (2018) for the TD3 implementation details, except that we do not take its 1 10-3 learning rate. This is because we found the original learning rate gives very poor performance in our experiments, so instead we set its learning rate x 10-1 which empirically much better and also consistent with two other algorithms. We conduct ablation study for the number of particles (i.e., K in Equation 12) used Figure This indicates that the number of level particles used in the estimation has an
quadruped
reacher hard
Epiadde Belurr
Spaand Teturm
Particess
Stop
Figure 9: Ablation study of the number of particles for SMAC on quadruped wa lk and reacher hard environments. The effect is evident some but not all environments.
effect on some of the environments. We choose the best hyperparameters (number of particles in (8, 16, 32), dimension of the latent in {8, 16, 32)) for each environment. We show the full set of experimental results in Figure 10.
Regarding other probabilistic policy modeling methods We further compare with normalizing flow based policy and mixture-of-experts. For the normalizing flow based method, we follow the practice Haarnoja et al. (2018a); Ward et al. (2019) to use RealNVP (Dinh et al., 2017) architecture for the policy distribution π =(ais), (a)s). The neural network is also followed by tanh transformation
18

---

Published as a conference paper at ICLR 2023
ball in cup catch
cartpele balance
cartpole belance sparse
cartpolc swingup
1000
1000
1000
1000
Eptods Retur
200
Epicode fatum
800
Eptcode Retum
200
Episdd Retum
500
600
600
400
Algorithe
400
400
400
200
200
200
TD3
0.2
0.8
0.4
0.8
0.0
0.1
0.8
1.0
0.0
0.4
0.5
1,0
Step
Step
166
Step
1c6
Step
cartpole swingup sparse
cheatah run
finger spin
Finger tumn easy
800
1000
Eps:de Return
500
Epiodd Rutum
Eptode Return
800
600
Epiiode
400
4.00
200
0.0
0.2
0.4
3.6
0.0
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.5
0.8
1.0
Step
106
Stop
166
Step
165
Step
166
haspper
horpor
800
250
Egocdd Rztum
Ep Return
Epiiode Retur
200
6000
4.00
4.00
100
200
50
0.2
0.8
0.8
0.0
0.4
1.0
24b
1se
1mb
Step
1eb
humanoid_ ka waak
hum armidadid_ru
pendulumm_swngg swwingu
reacher_easy
500
1.50
800
1000
tp.00e Wecumm
Eploode Retum
122
Erpindd Retur
Episoda falum
800
200
1.00
6000
200
400
300
200
200
0.2
0.8
0.8
0.0
0.4
0.8
0.0
0.4
0.5
1.0
Step
Step
Step
Step
reacher, hard
walker stand
walker walk
walker run
1000
1000
1000
Eplode Retum
800
Epicode Ratumn
800
Episode Retur
800
Ratum
500
600
600
4.00
400
400
400
200
200
200
2.0
0.2
0.4
3.6
0.8
1.0
0.0
0.2
0.4
0.5
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.5
0.8
1.0
Step
166
Step
166
Step
1e6
Step
166
quadruped walk
quadruped
quaciruped
quadruped fetch
3000
Epiods Retum
800
M
Episade Roturn
800
Episade Return
250
600
Epoade
4.00
200
50
0.2
3.6
0.8
0.4
0.8
0.0
0.4
0.8
0.0
0.4
0,8
1.0
166
166
1c6
185
Figure 10: Experiment results on different DMC environments with basee observations.
19

---

Published as a conference paper at ICLR 2023
ball in cup catch
cartpole balance
cartpolo balance Sparso
cartpole swinguu
1000
1000
8:0
Ep-s0dl Recrm
800
Algorithm
Epooodd Nuturn
800
Episode Return
Epiood Reture
Return
600
SMAC
600
SAC
Flow?
400
6:00
Fload
200
MCES
200
MDE19
0.2
0.4
0.6
0.6
1.0
0.0
0.2
0.4
0.5
0.8
1.0
0.0
0.2
0.4
0.6
0.8
10
0.0
0.2
0.4
0.5
0.8
1.0
1ef
1ed
Step
1e5
Step
1+6
caripole_ swingup. Sparss
cheatah_run
finger_spin
finger finger_tum _tm__aay
800
1000
1000
800
spisode Retur
Eploode Raturn
Eploode Return
800
800
400
400
4.00
400
Spicode
200
Step
1e6
Step
166
Step
166
ungerrturn__eee hhard
hopper_stand
hopper _hop
humanniid_stame
Eplodd Retumn
800
Edisode Return
Epiiode Return
200
Epicode Return
600
250
400
200
200
50
200
0
0.2
0.4
3.6
0.8
0.0
0.2
0.4
0.8
0.0
0.4
0.8
1.0
0.0
0.4
0.6
5.
1.8
Step
106
Stap
166
Step
1c6
Step
humanold, walk
hum Lumanoid_run run
pendulum, swingup
reacher caaay
1000
500
1.50
200
Eploode NotoleRetur Retur
400
Edisodo Reburn
122
Episode Episnde-Return Reburn
600
Edisoda Ruturnn
8.00
1000
300
200
75
400
1//
4.00
300
200
0.0
0,2
0.4
3.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Stop
245
Stap
165
Stop
20b
Stap
14b
hard
walker
walker
800
Episodle Return
Spocodle Roturn
Ep-sode Return
Return
800
400
400
200
200
0.2
0.8
0.4
1ed
Step
1+6
qua dww.ddduddddeco walk
quadruped run
quaciruped ascape
quadruped fwich
Epiocds Return
Epioode Raturn
Eprsode Iteurn
200
Episode Return
600
600
150
400
4.00
150
200
50
50
0.2
0.8
0.2
0.8
0.0
0.4
0.8
0.0
0.4
0.5
1.0
Step
50ep
Step
Step
Figure 11: Experiment results about comparison with other probabilistic policy modeling methods on different DMC environments with state-based observations. "Flow2" and "Flow4" refer to normalizing flow based policy with two or four RealNVP blocks as backend, while "MOE5" and "MOE10" refer to probabilistic mixture-of ofatuure-of-expeee ex-eeersss policy with five or ten mixtures. MOE, SAC and SMAC share similar number of parameters, while flow methods have about two or four times more parameters. Our proposed SMAC general achieves the best sample efficiency on majority of environments.
20

---

Published as a conference paper at ICLR 2023
Ball in cup, caach
cartpolle balance
Cartpole balance sparse
Cartpolle Swingup
1000
1000
Eptode Retumn
Episode flatum
Eptode Retum
Retum
600
Algerithen
600
400
Latent SAC
400
400
4.00
SLAC
230
MBRL
0.0
0.2
0.4
2.6
0.8
1.0
0.0
0.2
0.4
0.5
0.8
1.0
0.0
0.2
0.4
0.6
0.8
100
0.0
0.2
0.4
0.5
0.8
1.0
Step
1e6
Step
166
Step
1e6
Step
166
cartpole
cheetah_run
finger_spin
Finger
800
Epiodd Return
500
400
Edsode Return
800
Edsode Return
800
600
300
4:00
200
200
200
0.8
Step
Step
Stop
Step
fenger turn, hard
Pepppor stand
hopper hop
pendulum swinggup
Eptodd Retum
600
Episode Rabum
Epiode Retum
250
8000
409
150
4.00
200
200
200
50
0.2
0.8
0.4
0.8
0.0
0.4
0.8
0.5
0.4
0.5
1.0
166
166
1e6
Step
166
easy
reacher
walker
walloor walk
1000
Epio0ds Return
800
Episode Ratur
Raturn
Episode Return
Raturn
8:0
600
600
600
400
1.00
3.0
0.2
0.4
3.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
10
0.0
0.2
0.4
0.6
0.8
1.0
Step
Step
lec
Step
1mB
Step
1e6
walker
quadruped walk
quacruped run
quadrupedd escape
500
Ep:s0le Retum
600
Episade Return
Epiode Return
8.00
400
200
20
200
200
30
166
Step
166
Figure 12: Experiment results on different DMC environments with pixel-based observations.
21

---

Published as a conference paper at ICLR 2023
other methods. We experiment with two-block and four-block RealNVVSs, which makes their number of parameters to be upproximately twice and four times as the number of SAC and SMAC methods (these two share approximate the same number of parameters). We show their performance in Figure 11 as "Flow2" and "Flow4", respectively. A About the probabilistic mixture-of-experts (MOE) method, we follow the practice of Ren et al. (2021) to use five Gaussian mixture and ten Gaussian mixture. We show their performance in Figure 11 as "MOES" and "MOE10", respectively. This probabilistic MOE shares roughly around the same number of parameters and compute cost as SAC and SMAC. To conclude, SMAC achieves more favorable performance across a majority of the control tasks with less or equal number of parameters (notice that SMAC is the red curve in the figure). Normalizing flow based polices outperform SMAC on only one environment, while MOE requires the use of biased gradient estimators which often required extra hyperparameter tuning or could lead to worse performance.
Regarding Section 5.2 For this part, the agents are fed with pixel-based inputs. We mainly follow the Py'Torch world model implementation of Lin (2022) for this part. We implement the Latent-SAC algorithm according to instructions and hyperparameters in Wang et al. (2022). We implement the SLAAC (Leee et al., 2020a) algorithm following its original github repo (Lee et al., 2020b). Note that Latent-SAC and SLAC are two different algorithms, in the sense of actor modeling and world model design, although they both build on a SAC backend. We select the hyperparameters in the same way as the above part. We show the full version of experimental results in Figure 12. We do not plot the results of the humanoid domain as well as quadruped fetch, as all methods could not obtain meaningful results within 1 million frames.
For the robustness experiments, we add two kinds of noises in the pixel space, namely Gaussian perturbation and sensor missing perturbation. For Gaussian perturbation, we add isotropic Gaussian noise with scale 0.01 and 0.05. For sensor missing, we randomly drop the pixel value for each dimen- sion to zero according to a Bernoulli distribution (with parameter 0.01 and 0.05) in an independent way. For results in Table we report the best episodic reward across training iterations with standard deviation estimated from 5 seeds.
D THEORETICAL DERIVATIONS
D.1 UNIVERSALITY OF LATENT VARIABLE MODELS
We would first need the help of the following Lemma for the proof.
Lemma 2. For any continuous d-dimensional distribution p" (x) and Vc O. there exists neural network V IR- + i with finite depth and width that satisfies W [1 #., 1)Ile S e. Here # the push-forward operator, W(u, v) == inf inf J 1x y di (I, y) is the Wasserstein metric. and N 1) is the standard Gaussian distribution.
Proof. The Theorem 5.1 of Perekrestenko et al. (2020) shows that for any p" (x) and > there exists nonlinear ReLU neural network with finite size R + Rd that satisfies W #UO, 1)]l=(()) S c, where U00, 1] is the uniform distribution on [0, 1]. On the other hand, it is well known that the cumulative distribution function (cdf) of standard Gaussian of.)) : R R.. could map the standard Gaussian to the uniform distribution U00 1]. thus we have the construction of V = v o o.
We use this result to prove Proposition 1.
Proposition 1. For any d-dimensional continuous distribution p" (x), there exist sequence of two-level latent variable model pm( = (x)x))))) (2) dz,n n N4 that converge to it, where both (z)z)) and 2z) are factorized Gaussian distributions with mean and variance parameterized by neural networks.
Proof of Proposition 1. We let 2 E R and set Pn = - N.. From the above Lemma 2, know that Vn N+, AV ,, R Rd s.t. W (V. (+) S where V, is finite size
22

---

Published as a conference paper at ICLR 2023
neural network. We then set (x]2)) = (x; (x),),, 2²)) Note that this falls into the category of factorized Gaussian.
Let To be coupling between (x) and Ln #n (2), where no (z, x²) is the joint distribution (x,z") and x' = x C/n, { ~ [O, 10,,,),z ~ Vr (z)...) We thus have W (2)] #n (2)] S
I x' d 2
Since the Wasserstein metric satisfies the triangle inequality,
have W (x) (z)) & W PP (r) #P. + W M (z)llp (x)) "
D.2 TIGIT LOWER BOUND ON THE MARGINAL ENTROPY
Proposition 3 (Lower bound marginal entropy). For a latent variable policy r(a/h) (<(as))((s))) dz with prior ((s(h) and likelihood π (ais), s), consider
K (als(h)
)
Σ π
(k) Ea Egt (p (0 -
log
(27)
K+1
k=0
where p(sla, h) (sa)s) is the posterior and K is any positive integer, then the following holds:
(1) Hv(hh) < H(T(-Ih)) SA lo ((als)a(s(h) da.
(2) (k) - (h))
(3) limg-+ H (h) H(((+Ibb).
The following proofs roughly follow the derivations from Sobolev & Vetrov (2019). We describe them here for completeness.
For (1):
Proof. We write
K p(s(i) la, h)
H((T(- h)) - Hx (h) Be
log
K
+ 1
Am ((s(k)(b)
p(sgla a, ())(s(1:k)
Egg Eppsp/sso Ma h)
log
w(s(0:K) a, h)
DKI. p(sola, (b)l a, h)
where (s((0:)) ((s(0:K//a,h)) h) = (0) a ():(( h)) We only need to show that ((s((:K)) a, h) is normal-
x+1 Σ* ((o(t) pp)
ized density function.
Consider such generation process:
1. sample + samples ((A) q(([h), O,...,
2. set weight for each sample we (a))(()))
3. sample a categorical random variable ~ Cal - We
Wig
WK
4. put the /-th sample to the first: So = g(k), s(1:k) = 8(\))
It is easy to see the joint probability of this generation process is
(( (0:K), h)) ((s (0:K) h)
5(50 (()))((((1 -((b))
ΣK
k=
Wk
23

---

Published as a conference paper at ICLR 2023
Then the marginal of s (0:K) is
Σ
((((0:K(1h)
hh)
w/
8(80 - s((N)) ds(o:k)
k
k=
wk
(K 11 0(s(0:K)]
0(s(0) - 300)) - -((:k)) d5(0:
k=
wk
((s(1:K)
h))(s (0)
a. h)
=(K 1) Σx=o
a h).
(alls())
k
pa bm
K-I
2k=0
Thus (as(0:))|a, a, h) a normalized density function.
For (2).
Proof. We write
K+T Σ k= T(als(A))
Hk-1(h) Hx (h) Beg
log
K+2 k (Tals(K))
p(s(0) a, h)g(s(1:k+1) h)
Bsgnpissa,h) B Bexx
log
(1)
=KLL p(s haggs (1:+11
[h)]lu(s
(0:K la.
K4a ΣK+1 ((als()))
where (as(0, = p(s(0) a, b)q(s((:*+t(b) ((s(0:K+1)aa a, h) is normalized density function similarly (1).
We could then show that
Σi Ii (aa)s())))
For (3).
Proof. For the estimator, we have
Ax
Bx
K K Σ (als(())
K
K + T(als(0)) +
Σ
(H))(())
K + 1 K
R=1
By law of large numbers, we have Aς + 0, O, + 1, C, + p(x). Thus the limit of the left-hand side p(z).
24


Reviews:
[{"review_id": "r1IilhKrKEs", "reviewer": "Reviewer_fjCR", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "The paper proposes to marginalize (or rather, draw multiple samples from) latent variable policies as a means of improving estimates of policy entropy and Q-value estimates. The authors point out that a naïve ELBO estimate can result in a loose bound (i.e., a biased estimate), so they suggest using a nested estimator and multi-level Monte Carlo. They perform experiments using continuous control environments from DeepMind control suite, both in the setting of joint-based and pixel-based observations. They find that their method outperforms SAC and TD3 in the joint-based setting and outperforms various world-model-based setups in the pixel-based setting.", "summary_of_the_review": "The paper provides a technique for improving entropy and values estimates in latent variable policies using multiple samples. While I believe that this is a useful idea, the somewhat limited novelty, combined with the marginal performance boost (over inappropriate baselines) leads me to feel that this paper is somewhat below the acceptance threshold. The paper could be improved by improving the baselines to illuminate the various aspects (i.e., hierarchy vs. multi-sample vs. variance) of the proposed method. The authors may also consider whether this method allows one to scale up to larger hierarchical policies to tackle more challenging tasks where performance gains may be more significant.\n\n\n**UPDATE:**\nI appreciate that the authors have corrected various issues in the text that I raised. The normalizing flows experiments also add a useful multi-modal baseline to the results. I have increased my score from 5 --> 6 to acknowledge the effort on the part of the authors. However, I still feel that the complexity of the method and its limited novelty do not warrant the somewhat minimal boost in performance as currently demonstrated in the paper. For instance, in light of the normalizing flows experiments (Figure 11), it appears SMAC has statistically significant performance benefits on only 1 (or maybe 2) of the 24 DM control suite environments. I do not find this to be a compelling argument in favor of using SMAC.", "strength_and_weaknesses": "**Strengths**\n\n- **Proposes a principled approach to improving entropy and Q-value estimates.** To the best of my understanding, the method proposed here appears principled, building off of previous methods for entropy estimation. While many previous works have generally ignored the fact that world model latent variables induce hierarchical policies, few previous works have taken advantage of this fact to improve entropy and value estimators.\n\n- **Builds off of popular existing methods.** The paper builds off of several established algorithms in the literature, namely SAC and world models. These algorithms have fairly widespread adoption in the academic literature, so there’s a higher likelihood that methods built upon these algorithms will be more widely adopted, particularly if they are not difficult to implement.\n\n- **Uses standard benchmark setups/environments.** The authors evaluate SMAC using the DeepMind control suite, a standard set of continuous control environments built on MuJoCo. These environments are fairly commonly used in the literature, so demonstrating results here will likely be compelling to the deep RL sub-field.\n\n**Weaknesses**\n\n- **Fairly complex method for a fairly minimal performance boost.** Early in the paper, the authors lament that “the use of more expressive models have not gained nearly as much traction in the community…these constructions often result in complicate[d] training procedures and are inefficient in practice.” Two thoughts here: First, I would disagree with this characterization and analysis. Normalizing flow-based policies do not yield a significantly more complicated or inefficient procedure. Indeed, many modern neural network libraries (PyTorch, Tensorflow, and JAX/Distrax) make it trivial to implement these policies and sample / evaluate log-densities. It’s not that these more expressive policies aren’t useful. Rather, the limited boost in performance in our fairly simple benchmark environments hasn’t sufficiently justified their widespread adoption. Second, the method presented here strikes me as more complicated than, for example, normalizing flow-based policies. I find it difficult to believe that SMAC would receive widespread adoption when normalizing flow-based policies have not.\n\n- **Unclear if the baselines are really appropriate/fair comparisons.** The authors, at times, appear to confuse the distinction between several proposals: 1) using hierarchical policies, 2) using more sophisticated/multi-sample estimators, and 3) their particular approach for estimation, which itself consists of multiple components (entropy and Q-value estimators). For instance, in some of their experiments, the authors only compare with SAC and TD3, claiming to show the benefits of hierarchical policies. This is hardly a novel finding, as it’s well-known that more expressive policy distributions generally outperform less expressive ones. At the very least, I would have expected comparisons with normalizing-flow-based policies, which, having personally implemented these in one afternoon, are not difficult to use (and these do, indeed, outperform Gaussian policies in SAC). Likewise, in the only other set of experiments, the authors compare with “Latent-SAC”, SLAC, and Dreamer, but it’s not at all clear what these baseline comparisons are meant to accomplish. Is the boost in performance coming from using multiple samples, or is it the estimator itself? Is it from having better value estimates, entropy estimates, or do they both contribute? Likewise, as far as I understand, this estimator could also be used with a model-based value estimator, so it’s unclear whether the comparison with Dreamer tells us much about the setup proposed here.\n\n- **Somewhat minimal contribution.** While I’m not discounting the insight and effort required to develop the main techniques (nested estimator with MLMC) in the paper, the main contribution of the paper reduces to just drawing more samples for estimation. Assuming we have a reasonable estimator, it’s hardly surprising that drawing additional samples should yield improved entropy and value estimates, and likely, improved performance. The proposal of the paper is effectively to use a well-known knob within the existing RL framework to improve performance, which I view as a somewhat minimal contribution.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "SoxETlRu0R", "reviewer": "Reviewer_y62x", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "This paper considers the problem of MaxEnt reinforcement learning, with the goal of increasing the expressiveness of the policy while still allowing for practical entropy maximization. The proposed method relies on latent variables to avoid complexities with using EBMs. \n\nThe challenge with directly using latent variables is that the entropy is difficult to estimate, therefore entropy maximization is difficult. Typically, we'd want to optimize latent variable models via variational approaches with the ELBO, but this gives a lower bound of likelihood, hence an upper bound of the entropy term and the maxent rl objective. Maximizing the upper bound of course leads to practical issues.\n\nInstead the paper proposes to use a technique in hierarchical inference (Sobolev & Vetrov 2019), which gives a lower bound estimate to the entropy term. The rough intuition is use multiple samples, but unlike IWAE, do so in a way that reuses samples, and this can be shown to flip the direction of the bound.\n\nWith this new estimator, the paper optimizes the maxent rl objective with variational methods on a latent variable model and shows improvements over SAC.\n", "summary_of_the_review": "In summary the paper proposes a method for optimizing latent variable policies with the maxent RL objective, which is difficult to do with naive techniques.", "strength_and_weaknesses": "Strengths\nThe general idea makes sense, and the methodology is well explained. The technical contribution is not novel, but applying it to the RL setting appears to be novel.\n\nWeaknesses\nThe experiments don't seem to compare against energy based model and autoregressive model variants. The introduction says these \"complicate training procedures and are inefficient in practice\", but it seems important to provide experimental evidence.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "HmNOU0FefEJ", "reviewer": "Reviewer_cPcG", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "The paper proposes to use the latent variable policy for MaxEnt framework, which is testified to improve the sample efficiency and stability of RL. It uses some techniques to estimate the marginal entropy, as well as reducing the variances in the gradient. Experiments on DM control suite show the performance improvement of the proposed SMAC algorithm.\n", "summary_of_the_review": "Overall, the paper is written well. But the novelty of the proposed algorithm is not high enough, and the experiments show marginally improved performance in most environments. Some important baselines are missed, and more experiments should be added to clarify the problems I asked in the Weakness section.", "strength_and_weaknesses": "Strength: \n\nThe paper has good writing and clear delivery of the idea. Two techniques for solving the entropy estimation and variance reduction are effective. Experiments show performance improvement over previous algorithms almost consistently. Necessary ablation studies are conducted.\n\nWeakness:\n\nUnclear about Eq. (11). In Eq. (11), only the first sampling $s^{(0)}$ is used for sampling the action, while the other K $s$ are used for estimating the marginal entropy. Why is it? Why not use all $s$ for sampling action as well as estimating the entropy? Please explain in the paragraph.\n\nExplanations about notations. I’m not clear about the upper scripts $(a), (b)$ in Eq.(12). Please explain this in the paragraph.\n\nMissing references and baselines. There is a related paper (Probabilistic mixture-of-experts for efficient deep reinforcement learning. Jie et al. 2021) which also adopts the Gaussian mixture model for policy function approximation. Also I think Dreamer-v2 should also serve as an important baseline method. The paper should compare the proposed method with these.\n\nThe variance reduction in Sec. 3.2.2 seems to be a computationally heavy one. Please compare the computational efficiency in the alation study about whether this variance reduction is applied or not.\n\nHow about the on-policy algorithms? The proposed SMAC algorithm seems to only be compared with off-policy algorithms like SAC and TD3. How about on-policy algorithms like PPO? Is SMAC possible to be applied on PPO with an on-policy update as well?\n\nIn sec 5.1, is this model-free setting? If so, please write it explicitly in the paragraph.\n\n“Reasonable number of particles”. In Sec.5.1, the paper suggests to use a reasonable number of particles. What specific number is it, 8? What’s the number of the particles in each experiment?\n\nPerformance improvement in most of the experiments over existing algorithms seems to be marginal. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster