Simplifying Model-based RL: Learning Representations, Latent-space Models, and Policies with One Objective
Authors: Raj Ghugare
Homanga Bharadhwaj
Benjamin Eysenbach
Sergey Levine
Russ Salakhutdinov
Venue: ICLR 2023 poster (2023)
Keywords: Latent-space models
objective mismatch
model based RL

Abstract: While reinforcement learning (RL) methods that learn an internal model of the environment have the potential to be more sample efficient than their model-free counterparts, learning to model raw observations from high dimensional sensors can be challenging.
Prior work has addressed this challenge by learning low-dimensional representation of observations through auxiliary objectives, such as reconstruction or value prediction. However, the alignment between these auxiliary objectives and the RL objective is often unclear.
In this work, we propose a single objective which jointly optimizes a latent-space model and policy to achieve high returns while remaining self-consistent. This objective is a lower bound on expected returns. Unlike prior bounds for model-based RL on policy exploration or model guarantees, our bound is directly on the overall RL objective. We demonstrate that the resulting algorithm matches or improves the sample-efficiency of the best prior model-based and model-free RL methods. While sample efficient methods typically are computationally demanding, our method attains the performance of SAC in about 50\% less wall-clock time.


Published as a conference paper at ICLR 2023
SIMPLIFYING MODEL-BASED RL: LEARNING REPRESENTATIONS,
LATENT-SPACE MODELS, AND POLICIES WITH ONE OBJECTIVE
Raj Ghugare
Homanga Bharadhwaj
Benjamin Eysenbach
Sergey Levine
Ruslan Salakhutdinov
VNIT Nagpur
*Cannegie Mellon University
UC Berkeley
rajl9@students, vnit. ac.in, hbharadhees. emu. edu, beysenbalcs. cmu.
ABSTRACT
While reinforcement learning (RI.) methods that learn an internal model of the environment have the potential to be more sample efficient than their model-free counterparts, learning to model raw observations from high dimensional sensors can be challenging. Prior work has addressed this challenge by learning low- dimensional representation of observations through auxiliary objectives, such as reconstruction value prediction. However, the alignment between these auxiliary objectives and the RL objective is often unclear. In this work, we propose single objective which jointly optimizes latent-space model and policy achieve high returns while remaining self-consistent. This objective is a lower bound on expected returns. Unlike prior bounds for model-based RL on policy exploration model guarantees, our bound is directly on the overall RL objective. We demonstrate that the resulting algorithm matches or improves the sample-effiienn prior model-bass typically are computationally demanding, our method attains the performance SAC in about 50% less wall-clock time
of the best and model-free RL methods. While sample efficient methods
INTRODUCTION
While RL algorithms that learn internal model of the world can learn more quickly than their model-free counterparts (Hafner et al.] 2018 Janner et al.] 2019), figuring out exactly what these models should predict has remained open problem: the real world and even realistic simulators too complex to model accurately. Although model errors may be rare under the training distribution, learned RI agent will often seek out the states where otherwise accurate model makes mis- takes (Jafferjee et al.] 2020). Simply training the model with maximum likelihood will not, in general, produce model that good for model-based RL (MBRL). The discrepancy between the policy objective and the model objective is called the objective mismatch problem (Lambert et al. 2020). and remains an active area of research. The objective mismatch problem especially important in settings with high- -igh--imensional observations, which are challenging to predict with high fidelity.
Prior model-based methods have coped with the difficulty to model high-dimensional observa- tions by learning the dynamics of a compact representation of observations, rather than the dynamics of the raw observations. Depending on their learning objective, these representations might still be hard to predict or might not contain task relevant information. Besides, the accuracy of prediction depends not just on the model's parameters, but also the states visited by the policy. Hence, another way of reducing predic- tion errors is to optimize the policy to avoid transitions where the model is inaccurate, while achieving high returns. In the end, we want to train the model, representations, and policy to be self-consistent: the policy should only visit states where the model is accurate, the representation should encode information that is task-relevant and predictable. Can we design model- based RL algorithm that automatically learns compact yet sufficient representations for -based reasoning?
police
pc cylio
by
Inee braad
Haaruatiaie boa
mevdaadeeeeebbrrdd slyothme
Figure 1: (left) Most model- based RL methods learn the representations, latent-space model, and policy using three different objectives. (Right) We derive a single objective for all three components, which is a lower bound on expected returns. Based on this objective, we develop a practical deep RI. algorithm.
Project website with code: https::///a/ gned1 atentmodel s.gi thub. ic/

---

Published as a conference paper at ICLR 2023
In this paper, we present simple yet principled answer this question by devising single objective that jointly optimizes the three components of the model- based algorithm: the representation, the model, and the policy. As shown in Fig. | this is in contrast to prior methods, which use three separate objectíves. We build upon prior work that views model-based RI. as a latent-variable problem: the objective is to maximize the returns (the likelihood), which is an expectation over trajectories (the unobserved latent variable) (Botvinick & Toussaint 2012 Attias 2003 Eysenbach al. 202 This different from prior work that maximizes the likelihood of observed data, independent of the reward function (Hafner et al. 2019 Lee et al. 2020). This perspective suggests that model- based RL algorithms should resemble inference algorithms, sampling trajectories (the latent variable) and then maximizing returns (the likelihood) on those trajectories. However, sampling trajectories is challenging when observations are high-dimensional. The key to our work is to infer both the trajectories (observations, actions) and the representations of the observations. Crucially, we show how to maximize the expected returns under this inferred distribution by sampling only the representations, without the need to sample high-dimensional observations.
The main contribution of this paper is Aligned Latent Models (ALM), an MBRL algorithm that jointly optimizes the observation representations, model that predicts those representations, and a policy that acts based on those representations. To the best of our knowledge, this objective is the first lower bound for model-based RL method with a latent-space model. Across a range of continuous control tasks, we demonstrate that ALM achieves higher sample efficiency than prior model-based and model-free RL methods, including on tasks that stymie prior MBRL methods. Because ALM does not require ensembles (Chua et al. [0188] Janner et al.) (2019) or decision- planning (Deisenroth & Rasmussen 201T Sikchi et al. 2020, Morgan et al. 202T). our open-source implementation performs updates I0X and 6x faster than MBPO (Janner et al.) (219)) and REDQ (Chen et al. 2021
respectively, and achieves near-optimal returns about 50% less time than SAC.
RELATED WORK
Prior model-based RL methods use models in many ways, using it search for optimal action
sequences (Garcia et al.
1989; Springenberg et al. 2020 Hafner et al. 2018 Chua et al. 2018
Hatner et al.
[ et al. 22200) generate synthetic data (Sutton 199 Luo al. 2018 Hatner
et al. [2199] Janner et al. [019]] Shen et al. (020). to better estimate the value function Deisenroth &
Rasmussen 2011 Chua et al. 2018 [ et al. | Feinberg et al.] 2018)) or some combination
thereof (Schrittwieser et al.) 2020 Hamrick et al. 2020 Hansen et al. 2022). Similar to prior work
on stochastic value gradients (Heess et al. 2015| Hafner et aal.2011) Clavera et al., 2020 AAms et al., 2020), our approach uses model rollouts to estimate the value function for policy gradient. Prior work find that taking gradients through learned dynamics model can be unstable (Metz et al. 2021 Parmas al. 2019). However, similar to dreamer Hatner et al. (2019 2020), we found that BPTT can work successfully if done with a latent-space model with appropriately regularized representations. Unlike these prior works, the precise form of the regularization emerges from our principled objective.
Because learning a model of nigh-dimensional observations is challenging, many prior model-based methods first learn compact representation using a representation learning objective (e.g., image reconstruction (Kaiser et al. 2019 Oh et al. 2015 Buesing et al. 2018 Ha & Schmidhuber 2018) Hafner al. 2018 2019 2020). value and action prediction (Oh et 2017 SScrittwieserr et al.
2020: Grimm et al. 2020), planning performance (Tamar et al.] 2016 Racanière el al. 2017|
al. 2017), or sell sell-supervised learning (Deng et al. [02]]] Nguyen al.) 2021 Okada Taniguchi
2020). These methods then learn the dynamics of these representations (not of the raw observations). and use the model for RL. The success of these methods depends on the representation pppresentation/Auuugaaan [Arumugam & Roy (2022): the representations should be compact (i.e., easy to predict) while retaining task-relevant information. However, prior work does not optimize for this criterion, but instead optimizes the representation using some auxiliary objective.
The standard RL objective is to to maximize the expected returns, but models are typically learned via a different objective (maximum likelihood) and representations are learned via third objective (e.g., image reconstruction). To solve this objective mismatch (Lambert et al., 2020 Joseph et al. 2013: Grimm et al.) 2020), prior work study decision aware loss functions which optimize the model
to minimize the difference between true and imagined next step values (Farahmand et al., 2017]
Farahmand 2018 D'Oro et al. 2020 Abachi et al. 2020
Voeleker et al. 2022) directly optimize
2

---

Published as a conference paper at ICLR 2023
the model to produce high-return policies (Eysenbach et al. 2021a Amos al. 2018 Nikishin et al. 2021). However, effectively addressing the objective mismatch problem for latent-space models remains open problem. Our method makes progress on this problem by proposing single objective to be used for jointly optimizing the model, policy, and representation. Because all components optimized with the same objective, updates to the representations make the policy better (on this objective), as do updates to the model. While prior theoretical work in latent space models has proposed bounds on exploratory behavior of the policy (Misra et al., 2020), and on learning compressed latent representations (Efroni et al. 2021), our analysis lifts some of the assumptions (e.g.. removing the block-MDP assumption), and bounds the overall RL objective model- bodel--bseed setting.
A UNIFIED OBJECTIVE FOR LATENT-SPACE MODEL-BASSE RL
We first introduce notation, then provide high-level outline of the objective, and then derive our objective. Sec. 4 will discuss a practical algorithm based on this objective.
3.1 PRELIMINARIES
The agent interacts with Markov decision process (MDP) defined by states 8t, actions a,, an initial state distribution Ps(s), dynamics function p(s st, az), a positive reward function r(st, a.) 2 0 and a discount factor y E (0, 1). The RL objective is to learn a policy <a, | 8e) that maximizes the discounted sum of expected rewards within an infinite-horizon episode:
(1 (1-2) Σ Err((8,,a)))
max E
The factor of (1 11-) 7) does not change the optimal policy, but simplifies the analysis (Janner et al. 2020; Zahavy et al. 2021). We consider policies that are factored into two parts: an observation encoder (2 {,) and paaresentation-oondiioon
policy a 2z). Our analysis considers infitite-length trajectories т, which include the actions at, observations $t, and the corresponding
observation representations Ze: T & (So: ag, 20, $1, 01, 21. To simplify notation, we write the discounted sum of rewards as R(T)) (1 7)) Ei= (rr(st, at). Lastly, we define the Q-function of policy parameterized by d, Q ae) E Ey (R(T) 8c: orl-
3.2
METHOD OVERVIEW
Our method consists of three components, The first component an encoder c,(2; sr), which takes as input
shown Fig. 2
high-dimensional observation 8; and produces compact representation 2t. This representa- tion should be as compact as possible, while retaining the bits for selecting good actions and for predicting the Q-function. The second com- ponent is dynamics model of representations, m²(2++1 Ze, at), which takes as input the rep- resentation of the current observation and the ac- tion and predicts the representation of the next observation. The third component policy (a 21), which takes representations as in- puts and chooses an action. This policy will be optimized to select actions to maximize rewards, while also keeping the agent in states where the dynamics model is accurate.
enooder
palicy
space
(111)) 11501--122-11 7 (leqasa ex
Figure 2: Aligned Latent Models (ALM) performs model-based RL by jointly optimizing the policy, the latent-space model, and the representations produced by the encoder using the same objective: maximize predicted rewards while minimizing the errors the predicted representations. This objective corresponds RI. with an augmented reward function f. ALM esti- mates this objective without predicting high-dimensional observations $1411
3.3 DERIVING THE OBJECTIVE
To derive our objective, we build on prior work (Toussaint) 2009 Kappen et al., 2012) and view the RL objective vataalevariable problem, where the return I((7) is the likelihood and the trajectory
3

---

Published as a conference paper at ICLR 2023
is the latent variable. Different from prior work, include representations in this trajectory, in addition to the raw states, a difference which allows our method to learn good representations for MBRL. We can write the RL objective (Eq. I in terms of trajectories as (p(>)[[(TT) by defining the distribution over trajectories p(T)
Ps (T( _ (s()
II p(se+1 8t, (o) 2))e st).
Estimating and optimizing this objective directly is challenging because drawing samples from p(T) requires interacting with the environment, an expensive operation. What we would like to do instead estimate this same expectation via trajectories sampled from different distribution, g(T). We can estimate a lower bound on the expected return objective using samples from this different objective, by using the standard evidence lower bound (Jordan al. 1999):
log Ep(r) [R(T)] 2 Eey(r) [log R(r) + (()) log ((T))
This lower bound resolves first problem, allowing
estimate (a bound on) the expected return by drawing samples from the learned model, rather than from the true environment. However, learning a distribution over trajectories is difficult due to challenges in modeling high-dimensional observations, and potential compounding errors during sampling that can cause the policy to incorrectly predict high returns. We resolve these issues by only sampling compact representations of observations, instead of the observations themselves. By learning to predict the rewards and Q-values as function of these representations, we are able estimate this lower bound without sampling high-dimensional observations. Further, we carefully parameterize the learned distribution q(T) to support an arbitrary length of model rollouts (K), which allows us to estimate the lower bound accurately:
K (T) =P0(80)) sn) 20) II p(s. 18-11, 1)-1) Zt-1, at- (eee-)TT(ee zc). (4) 1
While it may seem strange that the future representations sampled from o (r) are independent of states, this is an important design choice. It allows us to estimate the lower bound for any policy, using only samples from the latent-space model, without access high dimensional states from the environment's 's dynamics function.
Combining the lower bound (Eq. 3) with this choice of parameterization, we obtain the following objective for model-based RL:
5 ( (r)
Σ
("r( aes st-1)) + nx log Q(xxQ (x))
(5)
where 7(se,ar, (1 7) logr(si, ai) log ca(441 I 8441) - log ms (2441 21, a1,)s)
(6)
(n)
(6)
This objective is an evidence lower bound on the RL objective (see Proof Appendix A.2
Theorem 3.1. For any representation eg(2 st), Tolay 24) and K E N, the ALM objective C corresponds to lower bound on the expected return objective:
model m 24, a.), policy
1 I- exp(Cx) Epg(T) En'r(s,
Here, we provide intuition for our objective and relate it to objectives in prior work. We start by looking at the augmented reward. The first term (a: extrinsic term) in this augmented reward function is the log of true rewards, which is analogous to maximizing the true reward function in the real environment, albeit on different scale. The second term (b: intrinsic term), i.e., the negative KL divergence between the latent-space model and the encoder, is reminiscent of the prior methods (Goyal et al., 2019 Eysenbach et al., 20216 Bharadhwaj et al. 2021 Rakelly et al.) 2021) that regularize the encoder against prior, to limit the number of bits used from the observations. Taken together, all the components are trained to make the model self-coonsisent with the policy and representation.
4

---

we find that the Published as a conference paper at ICLR 2023
Algorithm The ALM objective can be optimized with any RI. algorithm. We present an implemen- tation based on DDPG (Lillicrap et al., 2015).
1: Initialize the encoder EØ(2) st), model m²(=++1 Re, az), policy Ta(a, 2x), classifier Co(Z++1, C1, 2t), reward ro(21, at), Q-function Qo(2,, at), replay buffer B.
2: forn= 1, , do do
3: Select action On TT(- (s()) using the current policy.
4:
Execute action On and observe reward r, and next state $n+1-
5:
Store transition (sn, I': Sn+1) in B.
6:
Sample length-KK sequences from the replay buffer (s;, dz 8141 Compute the objective L (($;, ai, 814 1== using the sampled sequence.
1 1
B
7:
> Eq.
8:
Update encoder and model by gradient ascent on L Leg.ms.
9: Compute the objective - ((si=)) using on-policy model-based trajectories.
> Eq.
Update policy by gradient ascent on - Update classifier, Q-function and reward using losses LC LQus Lro.
Eq. 9
The last part of our objective the length of rollouts (K), that the model is used for. Our objective is directly applicable to all prior model-based RL algorithms which use the model only for fixed
number rollouts rather than the entire horizon, like SVG style updates (Amos et al. [2200] Heess et al. 2015) and trajectory optimization (Tedrake 2022). Larger values of K correspond to looser bounds (see Appendix A.6). Although this suggests that a model-free estimate is the tightest, a Q-function learned using function approximation with TD-estimates (Thrun & Schwartz 1993) is
biased and difficult to learn. A larger value of K decreases this bias by reducing the dependency on a learned O-function. While the lower bound in Theorem 3.1 is not tight, discount factor such that becomes tight (see Appendix A.5) In our experiments 4 objective in Theorem 3.1 still good estimate of true expected returns.
include a learnable
In Appendix A.4 we derive a closed loop form for the optimal latent- latent-dynamics and show that they biased towards high-return trajectories: they reweight the true probabilities of trajectories with their rewards. We also derive lower bound for the model-based offline RL setting, obtaining similar objective with an additional behavior eloning term (Appendix A.7).
A PRACTICAL ALGORITHM
We now describe practical method to jointly train the policy, model, and encoder using the lower bound (CK). ). We call the resulting algorithm Aligned Latent Models (ALM) because joint optimization means that the objectives for the model, policy, and encoder are the same; they aligned. For training the encoder and model (latent-space learning phase), g K (r) is unrolled using actions from replay buffer, whereas for training the policy (planning phase), 4 (r) is unrolled using actions imagined from the latest policy. To estimate our objective using just representations, method also learns to predict the reward and Q Q-function from the learned representations 2; using real data only (see Appendix for details). Algorithm provides pseudocode.
Maximizing the objective with respect to the encoder and latent-space model. To train the encoder and the latent-space model, we estimate the objective Cg using K-length sequences of 144K] sampled from the replay buffer:
transitions (si, ai, 8841 1
Cggmm (186, as 8141) 14
y Qo(ex,
cg(simelee) malasseltt
u4K-1
1 (ro(2), as) KL(mp(zu41 aelle (2411 ((41))) ). (7)
Σ
To optimize this objective, we sample an initial representation from the encoder and roll out the latent-space model using action sequences taken in the real environment (see Fig, We find that using a target encoder Exp,, (Zz sx) to calculate the KL consistency term leads to stable learning.
In code, we do not the 7 discounting for Equation 7
5

---

Published as a conference paper at ICLR 2023
Maximizing the objective with respect to the policy. The latent-space model allows us to evaluate the objective for the current policy by generating on on-policy trajectories. Starting from a sampled state s, the latent-space model is recurrently unrolled using actions from the current policy to generate K- length trajectory of representations and actions (ZE++K dx=++).. k)). Calculating the intrinsic term in the augmented reward (log ea( 22++1 St+1) log mp((++1 ze, (t)) for on-policy actions is challenging.
as we do not have access to the next high dimensional state 8t-1. Following prior work (Eysenbach)
al. 2020 Eysenbach et al. 2021a), we learn a classifier to differentiate between representations sampled from the encoder c 21++ S(+1) and the latent- model mø(#++1 24, Q1) to estimate this term (see Appendix details).
We train the latent policy by recurrently backpropagating stochastic gradients (Heess et al. 2015 Amos et al. 2020: Hafner et al. 2019) of our objective evaluated on this trajectory:
K+
x-1 (r((z))a)) + t.loog log Co( av,
C (s1) B (sLR) bee)
E i=
+ Qo(Ex,
Co C0(2441,
(8)
During policy training (Bq. {) estimating the intrinsic reward, log eo (2
(++1) log m (2++1
2t, at), is challenging because do not have access the next high dimensional state (8++1). Following prior work (Eysenbach et al.) 2020: Eysenbach et al.] (021a), we note that a learned classifier between representations sampled from the encoder (2 8r+1) versus the latent-space model m²(#+1 zz, 01) can also be used to estimate the difference between log-likelihoods under them, which is exactly equal to the augmented reward function:
Se+1) log ms (2241 2t, a+) log Co(2++1, at, 21)
log cœ(z141
Co 01, (t)
Here, C2(+111, at 21) E [0, is learned classifier's prediction of the probability that 2441 is sampled from the encoder conditioned on the next state after starting at (æ,, az) pair. The classifier is trained via the standard cross entropy loss:
LC es(: 8c), ze+1; i log(Co(Z441
zt; a:)) + log(1 Co(zeti, 24, dt)).
where R++1 c.(.1 8t-1) is the real next representation and že+1 mp( (- Ze, az) is the imagined next representation, both from the same starting pair az).
Differences between theory and experiments. While our theory suggests a coefficient use c = 0.1 in our experiments because it slightly improves the results. We do provide an ablation which shows that ALM performs well across different values of e Figure T6] We also omit the log of the true rewards in both Eq. [ and 7 We show that this change is equivalent to the first one for all practical purposes (see Appendix A.8) ). Nevertheless, these changes mean that the objective we use in practice is not guaranteed to be lower bound.
we
5 EXPERIMENTS
Our experiments focus on whether jointly optimizing the model, representation, and policy yields benefits relative to prior methods that use different objectives for different components. We use SAC-SVG (Amos al. 2020) the main baseline, as it structurally resembles method but uses different objectives and architectures. While design choices like ensembling (MBPO, REDQ) orthogonal to our paper's contribution, we nonetheless show that ALM achieves similar sample efficiency MBPO and REDQ without requiring ensembling; as a consequence, it achieves good performance in 6x less wall-clock time. Additional experiments analyze the Q-values, ablate components of our method, and visualize the learned representations and model. All plots and tables show the mean and standard deviation across five random seeds. Where possible, we used hyperparameters from prior works; for example, our network dimensions were taken from Amos et al. (2020). Additional implementation details and hyperparameters are Appendix [ and a summary of successful and failed experiments are in Appendix G. We have released the code
https://al gnedl ieedlatennt atentmode1 w..githoe thub. 10/

---

Published as a conference paper at ICLR 2023
Table 1: On the model-based benchmark from Wang et al. 2019), ALM outperforms model-based and model- free methods 4/s tasks, often by wide margin. We report mean and std. dev. across random seeds. We T-Humanoid-v2 and T-Ant-v2 to refer to the respective truncated environments from Wang 2019).
T- T-Humanoid-v2
T-Ant-v2
HalfCheetah--2
Walker2d-v2
Hopper-v2
ALM(3)
5306
4887 = 1027
10789
3006 = 1183
2546 1074
SAC-SVG(2)
Amos ot
501
4473
8752 178
448 + 2852
SAC-SVG(3)
472
3833 + 1418
9220 1431
878 # 1533
2024 #: 1981
SVG(1)
Teess et al
2015
811.8 241
185 4 141
336 387
252 =
435 +: 163
SLBO
et
2018
1377
200 1:40
1097 + 166
207
805 + 142
TD3
Fujimoto
147
870 283
3016 1:
-516
1817
SAC
Jhamuja et
2018
1470 =: 794
548 +: 146
3460 1326
166 :: 1318
788 1: 977
ALM(3)
MBPO
SAC
REDQ
Humanoid-v2
Ant-v2
HalfCheetah-v2
average returns
1.2
3.0
0.8
0.6
0.0
0.0
2.5
5.0
0.0
2.5
5.0
steps
165
env steps
1e5
env steps
1e5
ALM(3)
MBPO
REDQ
165
Humanoid-v2
1:4
HalfCheetah-v2
Ant-v2
average retrrs
1.0
seconds
0.5
0.4
MBPO
REDO
ALM33
time (hours)
16]
time (hours)
1n1
time (hours)
1n1
Figure 3: Good performance without ensembles. Our method (ALM) can (Top) match the sample complexity of ensembling based methods (MBPO, REDQ) while (Bottom) requiring less runtime. Compared to MBPO, ALM takes 10x less time per environment step. See Appendix Fig. for results on other environments.
Baselines. We provide a quick conceptual comparison to baselines Table 2
In Sec. 5. we
will compare with the most similar prior methods, SAC-SVG (Amos et al. 2020) and SVG (Heess et al. 2013). Like ALM, these methods use a learned model to perform S SVG-style actor updates. SAC-SVG SVG also maintains hidden representation, using GRU, make recurrent predictions in the observation space. While SAC-SVG learns the model and representations using reconstruction objective, ALM trains these components with the same objective the policy. The dynamics model from SAC-SVG bears resemblance to Dreamer-v2 (Hafner et al. (0020) and other prior work that targets image-based tasks (our experiments target state-based tasks) SAC-SVG reports better results
than many prior model- based methods (POPLIN-P (Wang & Ba 2019 SLBO (Luo et al. 2018).
ME-TRPO (Kurutach al. 2018)).
In Sec. 5.2 we focus on prior methods that use ensembles. MBPO (Janner et al. 2019) is model-based method that uses an ensemble of dynamics models for both actor updates and critic updates. REDQ (Chen et al. 2021) is model-free method which achieves sample efficiency on- par with model-baset methods through the use ensembles Q-functions. We also compare to TD3 (Fujimoto et al. 2018) and SAC (Haarnoja al. 2018); while typically not sample efficient, these methods can achieve good performance asymptotially.
5.1 Is THE ALM OBJECTIVE USEFUL?
We start by comparing ALM with the baselines on the locomotion benchmark proposed by Wang et al. 2019). In this benchmark, methods evaluated based the policy return after training for
7

---

Published as a conference paper at ICLR 2023
2e5 environment steps. The Truncated/Ant-V" v2 and TruneatedHumannidd-- -2 tasks included in this benchmark are easier than the standard Ant- and Humanoid- v2 task, which prior model-based methods struggle to solve (Janner al.] 2019: Chua et al.] 2018 Amos et al. 2020: Shen et al.) et al. (018)/Buckman et al. 2018). The results, shown in
2020; Rajeswaran et al. Table 1 show that ALM achieves better performance than the prior methods 4/s tasks. The results for SAC-SVG are taken from Amos et al. (2020) and the rest are from Wang et al. (2019). Because SAC-SVG is structurally similar ALM, the better results from ALM highlight the importance of training the representations and the model using the same objective as the policy.
5.2 CAN ALM ACHIEVE GOOD PERFORMANCE WITHOUT ENSEMBLES?
Prior methods such as MBPO and REDQ use ensembles to achieve SOTA sample efficiency at the cost of long training times. We hypothesize that the self-consistency property of ALM will make the latent-dynamics simpler, allowing it to achieve good sample efficiency without the use of ensembles. Our next experiment studies whether ALM can achieve the benefits of ensembles without the computational costs. As shown in Figure 3 ALM matches the sample complexity of REDQ and MBPO. but requires 6x less wall-clock time train. Note that MBPO fails to solve the highest- -ebbees-dimensionall tasks, Humanoid-v2 and Ant-v2 (R: and I. We optimized the ensemble training in the official REDQ code to be parallelized, leading to an increase in training speeds by 3x, which is still 2x slower than our method (which does not employ parallelization optimization).
5.3 WHY DOES ALM WORK?
To better understand why ALM achieves high sample complexity without ensembles, we analyzed the Q-values, ablation experiments, and visualized the learned representations.
Analyzing the Q-values. One way of inter- preting ALM is that uses model and an augmented reward function to obtain better es- timates of the Q-values, which are used to train the policy. In contrast, REDQ uses a minimum of a random subset of the Q-function ensem- ble and SAC-AVG (baseline used in REDQ pa- per (Chen et al. 2021)) uses an average value of the Q-function ensemble obtain a low vari- ance estimate of these Q-values. While ensem- bles can be an effective way to improve the estimates of neural networks (Garipov et al. 2018 Abdar
ALM(3)
REDQ
SAC-AVG
Humanoid-v2
Humanoid-v2
1.0
average bias
0.5
5td
steps
105
env steps
1e5
Figure 4: Analyzing Q-values. See text for details.
iet al. (02T), we hypothesize that our latent-space model might be more eifective approach in the RL setting because it incorporates the dynamics, while also coming at much lower computational cost.
To test this hypothesis, we measure the bias of the Q-values, as well as the standard deviation of that bias, following the protocol Chen et al. (202T); Fujimooo et al. (2018). See Appendix andd Figure for details. The positive bias will tell us whether the Q-values overestimates the true returns, while the standard deviation of this bias is more relevant for the purpose of selecting actions. We see from Fig. 4 that the standard deviation of the bias is lower for ALM than for REDQ and SAC-AVG, suggesting that the actions maximizing our objective are similar to actions that maximize true returns.
Ablation experiments. In our first ablation experiment, we compare ALM to ablations that sepa- rately remove the KL term and the value term from the encoder objective (Eq. 7 and remove the classifier term from the policy objective (Eq. [8). As shown in Fig. [aa the KL term, which is a purely self supervised objective Grill et al. (2020). is crucial for achieving good performance. The classifier term stabilizes learning (especially on Ant and Walker), while the value term has little effect. We hypothesize that the value term may not be necessary because its effect, driving exploration, may already be incorporated by Q-value overestimation (a common problem for RL algorithms (Sutton & Barto (2018); Fujimoto et al. (2018)) A second ablation experiment (Fig. 5b) shows the performance of ALM for different numbers of unrolling steps (K). We perform a third ablation experiment of ALM(3), which uses the TD3 actor loss for training the policy. This ablation investigates whether the representations learned by ALM(3) are beneficial for model-free free RL. Prior work (Gupta et al.)
2017 Eysenbach al. 2021b Zhang et al. 2020) has shown that representations learning can
8

---

Published as a conference paper at ICLR 2023
ALM131 Classiffer
KL value
ALM(3
ALMIS]
ALM model
TD3
ALM(3)
(ours)
HalfCheetah-v2
HalfCheetah-42
HallCheetah-w2
average returns
12
1.0
180
0.3
0.1
0.2
0.5
25
Ant-22
162
Ant-22
162
Anti2
average returs
365
steps
345
165
(a) Terms in the objective.
(b) Sequence length.
(c) Model free RL.
Figure 5: Ablation experiments (Left) Comparison of ALM (3) with no value term for encoder, no KI. term for the encoder and no classifier hased rewards for the policy. Results reflect the importance of temporal consistency terms, especially for training the encoder. (Center) Comparison of ALM(K) for different values of K and baselines SAC. Using architectures that support larger values of K could promise further improvements performance. (Right) Representation learning objective of ALM(3) leads to higher sample efficiency for model-free RL. To ensure the validity of these results, we implemented TD3 (TD3 (ours)), which uses the same architecture, exploration and learning parameters as our method. Ablation results for other environments can be found in Fig. T0a
10b
10c
facilitate properties like exploration, generalization and transfer. fig. 5e we find that the end to end representation learning of ALM(3) achieves high returns faster than standard model-free RI..
CONCLUSION
This paper introduced ALM, an objective for model-based RL that jointly optimizes representations, latent-space models, and policies all using the same objective. This objective mends the objective mismatch problem and results a method where the representations, model, and policy all cooperate to maximize the expected returns. Our experiments demonstrate the benefits of such joint optimization: it achieves better performance than baselines that use separate objectives, and achieves the benefits of ensembles without their computational costs.
At a high level, our end-to-end method is reminiscent of the success deep supervised learning. Deep learning methods promise to learn representations in an end-to-end fashion, allowing researchers to avoid manual feature design. Similarly, our algorithm suggests that the algorithmic components themselves, like representations and models, can be learned in an end-to-end fashion to optimize the desired objective.
Limitations and future work. The main limitation of our practical method is complexity: while simpler than prior model-based methods, it has more moving parts than model-free algorithms. One potential future work that can directly stem from ALM an on-policy version of ALM, where Equation
8 can be calculated and simultaneously optimized with the encoder, model and policy, without using a classifier. Another limitation is that in our experiments, the value term in the equation 7 does not yield any benefit. An interesting direction is to investigate the reason behind this, or find tasks where an optimistic latent space is beneficial. While our theoretical analysis takes an important step in constructing lower bounds for model-based RL, it leaves many questions open, such as accounting for function approximation and exploration. Nonctheless, we believe that our proposed objective and method are not only practically usefull but may provide template for designing even better model- based methods with learned representations.

---

Published as a conference paper at ICLR 2023
Acknowledgments. The authors thank Shubham Tulsiani for helpful discussions throughout the project and feedback on the paper draft. We thank Melissa Ding, Jamie Gregory, Midhun Sreekumar, Srikanth Vidapanakal, Ola Electric and CMU SCS for helping to set up the compute necessary for running the experiments. We thank Xinyue Chen and Brandon Amos for answering questions about the basclines used in the paper, and Nicklas Hansen for helpful discussions on model-based RL.
REFERENCES
Romina Ahachi, Mohammad Ghavamzadeh, and Amir-massoud Farahmand. Policy- model learning for policy gradient methods, 2020.
00030
Moloud Abdar, Farhad Pourpanah, Sadiq Hussain, Dana Rezazadegan, Li Liu, Mohammad Ghavamzadeh, Paul Fieguth, Xiaochun Cao, Abbas Khosravi, U. Rajendra Acharya, Vladimir Makarenkov, and Saeid Nahavandi. A review of uncertainty quantification in deep learning: Techniques, applicatioon and challenges. Information Fusion, 76:243-297, 2021. ISSN 1566-2535. doi: http:://bb..rrr/1..001..... 2021 005.008. httLitttpss//ww..cieeeedddd htt.://weeesattcccoooooooccaaaaa httpp///.............aaaaa/////.
com
Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Martin Riedmiller. Maximum a posteriori policy optimisation, 2018. /tt//httss////rx/.oorr///
1806 06920
Brandon Amos, Ivan Dario Jimenez Rodriguez, Jacob Sacks, Byron Boots, and J. Zico Kolter. Differentiable mpe for end-to-end planning and control, 2018. Bttp//ttsss//aax...crr///a//////
Brandon Amos, Samuel Stanton, Denis Yarats, and Andrew Gordon Wilson. On the model stochasiic value gradient for continuous reinforcement learning, 2020. URL BtRLLtttss///xxii/i
abs 200.
Dilip Arumugam and Benjamin Van Roy. Deciding what to model: Value-equivalent sampling for reinforce- ment learning. In Alice II Oh, Alekh Agarwal, Danielle Belgrave. and Kyunghyun Cho (eds.). Advances in Neural Information Processing Systems, 2022. URL btttp:://penzeviiey wttpl/wwwwee/ttoutteeee id=
[ORXbI 1 TELP
Hagai Attias. Planning by probabilistic inference. Christopher M. Bishop and Brendan J. Frey (eds.), Proceedings of the Ninth International Workshop on Artificial Intelligence and Statistics, volume R4 of Proceedings of Machine Learning Research, pp. 9-16. PMLR, 03-06 Jan 2003. URL Intpps
//proceecings.o mlr. www.l..ress.e.mmme prwss/ r4 24/attias03a htm.
Reissued by PMLR 01 April 2021.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization, 2016 URL https
06450
Homanga Bharadhwaj. Mohammad Babaeizadeh, Dumitru Erhan, and Sergey Levine. Information prioritization through empowerment in visual model-based rl. In International Conference on Learning Representations, 2021.
Matthew Botvinick and Marc Toussaint. Planning inference. Trends in Cognitive Sciences, 16(10):485 488, 2012. ISSN 1364-6613, doi: https://bo.oop//00001.......... 2012..........2....... URL https://www
scicncedirect. wtt. /o htt.oo/c...........ooo/o.aaaaaa.
Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sample-efficient rein- forcement learning with stochastic ensemble value expansion. In S. Bengio, H. Wallach, H. Larochelle. K. Grauman, N. Cesa- BasacBianchi, and R. Garnett Advances Neural Information Processing Systems, vol- ume 31. Curran Associates, Inc., 2018. 4dd.2225443ec33bb Pap.-br---ppp..... Paper.pd. pdd
//tp:c/e/ro/ceeee//e ing........ neurrps.... pctpp/pp///1.////////
fi 111e/F0222880578664
Lars Buesing, Theophane Weber, Sebastien Racanicre, S. M. Ali Eslami, Danilo Rezende, David P. Reichert, Fabio Viola, Frederic Besse, Karol Gregor, Demis Hassabis, and Daan Wierstra. Learning and querying fast generative models for reinforcement learning, 2018. WRRLbtttss//arxiv wttp:/xwwwwvwrrr//e//eeeeeee/eee 00000 63006
Xinyue Chen, Che Wang, Zijian Zhou, and Keith Ross. Randomized ensembled double q-learning: Learning fast without a model, 2021. URL httt:tt:::
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcemmett learning handful of trials using probabilistic dynamics models, 2018. htttp:://aaxiv..or//////////////
leamnnn in a
12114
Ignasi Clavera, Violet Fu, and Pieter Abbeel. Model-angmented actor-critic: Backpropagating through paths, 2020. wtt.httts://aaxiv.orrr.
10

---

Published as a conference paper at ICLR 2023
Mare Peter Deisenroth and Carl Edward Rasmussen. Pilco: model-based and dataefficient approach policy search. In In Proceedings of the Eventy-Eightn International Conference on Machine Learning (ICML, 2011.
Fei Deng, Ingook Jang, and Sungjin Ahn. Dreamerpro: Reconstruucioon-fess model based reinforcement learning with prototypical representations, 2021. Utl.attps///xxxv.orr////o// http:://aaaaaa.....///////111..0 14565
Pierluca D'Oro, Alberto Maria Metelli, Andrea Tirinzoni, Matteo Papini, and Marcello Restelli. Gradient-aware model-based policy search. Proceedings the AAAI Conference Artificial Intelligence, 34(04))3801-3800 apr 2020, doi: 10. 1699/669/iaii... v34904.5791.
URL https://doi..c.. .rg. 100.1.66099aaa. www.raaaai.....00..0. v34104]
5791
Yonathan Efroni, Dipendra Misra, Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Provable rl with exogenous distractors via multistep inverse dynamics. arXiv preprint v:Xiv:2110.0884 2021.
Benjamin Eysenbach, Swapnil Asawa, Shreyas Chaudhari, Sergey Levine, Ruslan Salakhutdinov. Off- Dynamics Reinforcement Learning: Training for Transfer with Domain Classifiers. arXiv e-prints, art. MeXiv:2006113114 June 2020.
Benjamin Eysenbach, Alexander Khazatsky, Sergey Levine, and Ruslan Salakhutdinov. Mismatched no more: Joint model-policy optimization for model-based rl, 2021a. /t//ppttss///rrrxxrrt///////1/1 02758
Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. Robust predictable control, 2021b. URL
https:
03214
Amir-massoud Farahmand. Iterative value-aware model learning. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Dessa-Binnchi, and R. Garnett (eds.), Advances Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL /ttpps://pxoceeenngss www.inuuri.... www.reue.....ppppaee 2018/E3 1ellee//a /a2 7d96 (52555220335550 2e98 2ew.22221111111eeeeeeee
Amir-Massoud Farahmand, Andre Barreto, and Daniel Nikovski. Value-Aware Loss Function for Model-based Reinforcement Learning. In Aarti Singh and Jerry Zhu (eds.). Proceedings of the 20th International Conference
on Artificial Intelligence and Statistics, volume 54 of Proceedings of Machine Learning Research, pp. 1486 1494. PMLR, 20-22 22 Apr 2017. URL / /tttps::/procee http://poooccd.n..ooo mir. www.sp.rcssvv5/. Earaana ahmandl 7 html
Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael I. Jordan, Joseph Gonzalez, and Sergey Levine. Model- based value estimation for efficient model-free reinforcement learning, 2018. URL htttp:////a
org/
abs/ org/abs/ /1803.110.... 00101
Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1587 1596. PMLR, 10 Jul 2018. URL hssps://p/ http::///rcccdnneeeee http:///......pp./eeeeeeem
mlr. wwwppppsssssssp (tt:00/ffujmmmooeoo0
html
Carlos E. Garcia, David M. Prett, and Manfred Morari. Model predictive control: Theory and practice survey. Autom., 25(3):335-34) 338, 1989. URI. hhtp: /ttp://dbbp http//pppppp.i........ cemii-rries wwt.:eeeeeeeeeeeoeeeeeeeeeeee httprr/naaaaattttaaaaaaaa cmat ca/
aautomat ca25. cww.0..tt.aa.....mmmmmmmmmm
htm FGa PPM8
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry Vetrov, and Andrew Gordon Wilson. Loss surfaces, mode connectivity, and fast ensembling of dnns, 2018. URL.attpss:/aaxxv..ro org abs / 11////80.00 10026
Anirudh Goyal, Shagun Sodhani, Jonathan Binas, Xue Bin Peng, Sergey Levine, and Yoshua Bengio. Re- inforcement learning with competitive ensembles of information-coostrainedd primitives, 2019. URL
hittps://arxivs http://aawwwvvoo/////////////
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H. Richemond, Elena Buchatskaya, Carl Doersch, Bemardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Koray Kavukcuoglu, Rémi Munos, and Michal Valko. Bootstrap your own latent: A new approach to self-supervised learning, 2020. URL/ittpss///rrrv.oor////oa
07733
abs 2211 03506 Christopher Grimm, André Barreto, Satinder Singh, and David Silver. The value equivalence principle for model-based reinforcement learning, 2020. URL/Attps://axxly.
Abhishek Gupta, Coline Devin, YuXuan Liu, Pieter Abheel, and Sergey Levine Learning invariant feature spaces transfer skills with reinforcement learning, 2017. URI
11

---

Published as a conference paper at ICLR 2023
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Desa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. / .nw.urrriiin.g.....iigggg neurips
ce/paper/
Tuomas Haarooja Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor, 2018. URL
abs/
18011012200
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels, 2018. URL
1811
04551
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination, 2019. URL
1912. 01 603
Danijar Hafner, Timothy Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models, 2020. URL
201/ 02193
Jessica B. Hamrick, Abram Friesen, Feryal Behbahani, Arthur Guez, Fahio Viola, Sims Witherspoon, Thomas Anthony, Lars Buesing, Petar Veličkovié, and Théophame Weber. On the role of planning in model-based deep reinforcement lcarning, 2020. URL wttp:htpp////m httpa::/aaxxvvvv http:///wwwwvv.oo/aaaa////// /201 04021
Nicklas A Hansen, Hao Su, and Xiaolong Wang. Temporal difference learning for model predictive control. In Kamalika Chaudhuri. Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 8387-8406. PMIR, 17-23 Jul 2022. URL hhttps://pprceeeteeeeeee ngs
m1 r. press/vsllll 62/haassn222a html
Nicolas Heess, Greg Wayne, David Silver, Timothy Lillicrap. Yuval Tassa, and Tom Erez Learnnng continuous control policies by stochastic value gradients, 2015. ///tp:////////xii http///aawvvvrrgg////a///////111
Taher Jafferjee. Ehsan Imani, Erin Talvitie, Martha White, and Micheal Bowling. Hallucinating value: A pitfall of dyna-style planning with imperfect environment models, 2020. Uttppsssss////xx..oorg abs/ 2006.04363
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When trust your model: Model-based policy optimization, 2019. URL htttps://arxi...o wwtpr/iwwwwo//ab////aa/aaaa 1900108253
Michael Janner, Igor Mordatch, and Sergey Levine. Gamma-models: models: Generative temporal difference learning for infinite- horizon prediction. In H. Larochelle, M. Ranzato, Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1724 1735. Cur- ran Associates, Inc., 2020. URL /ttpss//procced ngs. wwt.inss....pppppeeeeeeee.peeeee pttpp/ppppp//2///1111111 .pdf
22ffb096822
a 6bbbb362289--
Michael I Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul, An introduction to variational methods for graphical models. Machine learning, 37(2): 1999.
Joshua Joseph, Alborz Geramifard, John W. Roberts, Jonathan P. How, and Nicholas Roy. Reinforcement learning with misspecified model classes. In 2013 IEEE International Conference on Robotics and Automation, pp. 939 946, 2013. doi: 10.1 1091199//...2 20636666.....666666666
Lukasz Kaiser, Mohammad Bahacizadeh, Piotr Milos, Blazej Osinski, Roy H Camphell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model-based reinforcement learning for atari, 2019. URL jttpss
l/arxiv.org/ wwwp/////////////////.///eeeeee/ abs/ 900.0000379
Hilbert J Kappen, Vicenç Gómez, and Manfred Opper. Optimal control as graphical model inference problem. Machine learning, 1592),1 2012.
Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model-ensemble trust-region policy optimization, 2018. http:tts:/// arxiv www.///gg//b.......o.../e..eeset 1atbs.///..... 10592
Nathan Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in model- reinforcement learning. ArXiv, abs/2002. 04523, 2020.
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, 2020.
12

---

Published as a conference paper at ICLR 2023
Sergey Levine. Aviral Kumar, George Tucker, and Justin Fu, Offline reinforcement learning: Tutorial, review, and perspectives on open problems, 2020. URL
01643
Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez. Yuval Tassa. David
Silver. and Daan Wierstra. Continuous control with deep reinforcement learning, 2015. URL attps
//arxiv.org/ abs/ 1ee.//5599.02200 02971
Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian. Trevor Darrell, and Tengyu Ma. Algorithmic framework for based deep reinforcement learning with theoretical guarantees. 2018. URLatts:://arxiy
org 8877 0007000557 03858
Luke Metz, C. Daniel Freeman, Samuel S. Schoenholz, and Tal Kachman. Gradients are not all you need, 2021.
URL htt.h:ttt::/aatt http://aaxvvvv... www.aawwwww....aaaaaaaaaao
05803
Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford. Kinematic state abstraction and provably efficient rich-observation reinforcement learning. In International conference on machine learning, pp. 696669666999 1971. PMLR, 2020.
Andrew S. Morgan, Daljeet Nandha, Georgia Chalvatzaki, Carlo D'Eramo, Aaron M. Dollar, and Jan Peters. Model predictive actor-critic: Accelcrating robot skill acquisition with deep reinforcement learning. In 2021 IEEE International Conference on Robotics and Automation (ICRA). IEEE, may 2021. doi: 10.1 sra48506.2021..33 99561298. /ttphttts////i..ooo http://dtt...g....////oo// http:///...12/0rra444888000 2021 2021. 9561298
Tung D Nguyen, Rui Shu, Tuan Pham, Hung Bui, and Stefano Ermon. Temporal predictive coding for model- based planning in latent space. In International Conference Machine Learning, pp. 8130-8139, PMLR, 2021.
Eygenii Nikishin, Romina Abachi, Rishabh Agarwal, and Pierre-Luuc Bacon, Control-oriented model-assee reinforcement learning with implicit differentiation, 2021. URL http:t/e//aaaii...rrr////////1//
03273
Junhyuk Oh, Xiaoxiao Guo, Honglak Lee, Richard Lewis, and Satinder Singh. Action-conditional video prediction using deep networks in atari games. In Proceedings of the 28th International Conference Neural Information Processing Systems Volume NIPS'15, 15, pp. 2863-2871, Cambridge, MA, USA, 2015. MIT Press.
Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network, 2017. URL https://axxxv org/abs/1707.. 03497
Masashi Okada and Tadahiro Taniguchi. Dreaming: bodel-baaed learning by latent imagination without reconstruction, 2020.
14535
Masashi Okada, Luca Rigazio, and Takenobu Aoshima. Path integral networks: End-to-end diffcrentiable 1706.776600 09597
optimal control, 2017. URL http:://wwwx......a.oo./
Paavo Parmas, Carl Edward Rasmussen, Jan Peters, and Kenji Doya. Pipps: Flexible model-based policy search robust to the curse of chaos, 2019. URL htt.thttt///axiiii http://axxvvv..//aaa/// www.obrgg////.o.ooooo.oooo...sra /000 01240
Jan Peters, Katharina Muelling, and Yasemin Altun. Relative entropy policy search. In AAAI, 2010.
Rafael Figueiredo Prudencio, Marcos R. O. A. Maximo, and Esther Luna Colombini. A survey on offline reinforeement lcarning: Taxonomy, review, and open problems, 2022. httpbttss///axiii org/ abs/
2203.001 387
Sébastien Racanière, Théophane Weber, David P. Reichert, Lars Buesing, Arthur Guez, Danilo Jimenez Rezende, Adrià Puigdomênech Badia, Oriol Vinyals, Nicolas Manfred Otto Heess, Yujia Li, Razvan Pascanu, Peter W. Battaglia, Demis Hassabis, David Silver, and Daan Wierstra. reinforcement learning. ArXiv, abs/1707. 00000,0603, 2017.
augmented agents for deep
Aravind Rajeswaran, Igor Mordatch, and Vikash Kumar. A game thcoretic framework for model based reinforeement learning, 2020.
Kate Rakelly, Abhishek Cupta, Carlos Florensa, and Sergey Levine. Which mutual- information representation learning objectives are sufficient for control? Advances in Neural Information Processing Systems, 34: 26345-26357, 2021.
Julian Schrittwieser, loannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis. Thore Graepel, Timothy Lillicrap, and David Silver. Master- ing atari, go, chess and shogi by planning with learned model. Nature, 588 7888)78)66004-666 609, dec 2020. doi: 10. 1038/s41 1064-1866-000 586-020-03000 020-03051-4 URL / .rrg 10. 1038 2 586 020 03051-4
1036
13

---

Published as a conference paper at ICLR 2023
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High--dieensiinnalll continuous control using generalized advantage estimation, 2015. URL http://UUUtttt////eeeete /ttpps///aa.....og///////////// 02438
Jian Shen, Han Zhao, Weinan Zhang, and Yong Yu. Model-based policy optimization with unsupervised model adaptation, 2020. Bttps:/ts://ax
Harshit Sikchi, Wenxuan Zhou, and David Held. Learning off-policy with online planning, 2020. URL
10066
Jost Tobias Springenberg. Nicolas Heess. Daniel Mankowitz. Josh Merel, Arunkumar Byravan, Abbas Abdolmaleki. Jackie Kay. Jonas Degrave, Julian Schrittwieser, Yuval Tassa, Jonas Buchli, Dan Belov. and Martin Riedmiller. Local search for policy iteration in continuous control, 2020. URL https
//arxiv www../org/aoooo abs/ 2016 2010.05545
Richard S. Sutton. Dyna, an integrated architecture for learning, planning, and reacting. SIGART Bulletin, 2(4): 160-163 1991.
Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. The MIT Press, second edition, 2018. URLphtt://Lncomppeee. wwt.dddww......ooooooeoooooooooo not
httpa///e////ooooooooooooooooooo http///oo/oohbbooo--ooooooooe html
Aviv Tamar, YI WU, Garrett Thomas, Sergey Levine, and Pieter Abbeel. Value iteration networks. In D. Lee,
M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016. URL
www.ggs.n..r...... cc/
Russ Tedrake. Underactuated Robotics. Course Notes for MIT 6. 2022. URL BRLLattt:///uaet http:://underactuteee
t. .ddu
Scbastian Thrun and Anton Schwartz. Issues in using function approximation for reinforeement learning. In Proceedings of the 1993 Connectionist Models Summer School Hillsdale, NJ. Lawrence Erlbaum, volume pp. 1 9, 1993.
Mare Toussaint. Rohot trajectory optimization using approximate inference. In Proceedings of the 26th Annual International Conference on Machine Learning, ICML '09, pp. 1049-1056, New York, NY, USA, 2009, Association for Computing Machinery. ISBN 978 605585161.. doi: 10.1145/1533374115533000 URL
fttps://dci..roge 1115:/111111111313311133 11053334..555508 1553508
Claas Vnelcker, Victor I Liao, Animesh Garg, and Amir-massoud Farahmand. Value gradient weighted model- based reinforcement learning, 2022. h htt.tttss///xxt......o
ata///b//22000 0 464
Tingwu Wang and Jimmy Ba. Exploring model-basee planning with policy networks, 2019. URL ht
//arxiv. abs/ 1906. 08649
Tingwu Wang, Xuchan Bao, Ignasi Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shunshi Zhang, Guodong Zhang, Pieter Abbeel, and Jimmy Ba. Benchmarking modcl-based reinforeement learning, 2019.
URL //arivv...vvv/////
www.oorgg//.....eeoooo www.bb//.......... 02057
Florian Wenzel, Kevin Roth, Bastiaan S. Veeling, Jakuh Swigtkoski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sehastian Nowozin. How good the bayes posterior in deep neural networks really?, 2020. htRl:/tttp.:/aariiorr///////
02405
Kevin Xie, Homanga Bharadhwaj, Danijar Hafner, Animesh Garg, and Florian Shkurti. Latent skill planning for exploration and transfer. In International Conference on Learning Representations, 2020.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-angmented reinforcement Icarning, 2021. URL /Rtp/t/teee///rrrv...ro//bbbb
Tom Zahavy, Brendan O' O'Donoghue, Guillaume Desjardins, and Satinder Singh. Reward is enough for convex mdps, 2021. htttl::/tss:/attt wttp:/wwwww..rr//////.//.oo Utt.hrrcpe//aarrcc..rrraaa/////p 00661
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representa- tions for reinforcement learning without reconstruction, 2020. 10742
14

---

Published as a conference paper at ICLR 2023
Outline of Appendices. In Appendix A include all the proofs. In Appendix various components used by ALM and the baselines. Eopendix/Cincludeee additional learning details. In Appendix D we have mentioned implementation details. In Appendix we have details about additional experiments. In Appendix [ we have a summary of experiments that were tried and did not help. Lastly, Appendix compaares objective to the MnM objective from prior work.
compare
A PROOFS
A.I HELPER LEMMAS
Lemma A.1. Let Pk (K( be a truncated geometric distribution.
(1
(1-7))) -
H E (0, K - 1] K H> K
(k(
7
Given discount factor 7 E (0, 1) and a random variable 24, we have the following identity
EP)(H)
Σ
Zt
(K(
Σ
zt
H=0
t=0
K-1
H
K
= =(1 -1-11 7)
+ K
#=0
1=0
1=0
=3o((1 - (1 + 7
K-1) + + ) xi ((1 (7 + *
+ xx (xx(**)
=o((1 77) - 7k
01((1 (1- 71 - --)-1)
+ K)
1 1-9 + { + ..... x)(
?
xxx
Lemma A.2. Let (K(H) be truncated geometric distribution and PolT H) be the distribution over H + 1 length trajectories:
(T(T] H) = po(seee (00 =0)
II p(s: 18--1, (1-1, ze)e 81)
Then the RL objective can be re-writlen in the following way.
(1 - 7) Σ 7/((s, ax)
Epa(T)
K
E. (r) (r)
(1-7) Σ
x x-
r at) + + Q(sK ok)
= (m) [po(H]H- [n[ < K - 13r(sH + 1(H K}Q(sH ax)I]
The Q function is defined Q(HH, .g) (p)(() [[1 7) E 7'r(s14 a1+H).. This lemma helps us to interpret the discounting in RL as sampling from a truncated geometric distribution over future time steps.
Lemma A.3. Let p(x) be a distribution over R. The following optimization problem can be solved using the methods of Lagrange multipliers to obtain optimal solution analytically.
max Ep(x)) [f(x) log p(z)]
p(x)
such that J ((z))x
15

---

Published as a conference paper at ICLR 2023
Starting by writing the Lagrangian for this problem, then differentiating it and then equating it to
[(p(s), A) Bp(a)) [f logp(x)) - 21 p(z)do - 1)
V v)x) = f(x) - p())(x) - -1 X
0 f(z) - log ((r) - -1
p p²()) - ((x)-1-A
We find the value of À by substituting the value of p(x) from (b) in the equality constraint.
I [ =1
elta _ f e
Substituting (c) to remove the dual variable from (b), obtain
d el(
v²(x) Jetldds
A.2 A LOWER BOUND FOR K-STEP LATENT-SPACE
In this section we present the proof of Theorem
3.1 We restate the theorem for clarity.
Theorem 3.1. For any representation e se). latent-space model m 24 az), policy (a zz) and E N, the ALM objective C corresponds to a lower bound on the expected return objective:
1 1-y exp(Cx) Epo(T)
Σ Ex(r(s,
Note that scaling the rewards by a constant factor (1 7) does not change the RI, problem and finding policy maximize the log of the expected returns is the same as finding a policy to maximize expected returns, because log is monotonic.
We want estimate the RL objective with trajectories sampled from different distribution go(T) which leads to an algorithm that avoids sampling high-dimensional observations.
g (T) = Po(s))cs 8o) II p(s-1 81, a.)mm Zt, ( 2e)
1=0
When used as trajectory generative models, (T( predicts representations 2t using current state 8t. Whereas (T(T) predicts by unrolling the learned latent-model recurrently 2;-1 (and a,-1). Similar to variational inference, (T( can be interpreted as the posterior and (T( as the recurrent prior. Since longer recurrent predictions of a learned model significantly diverge from the true ones, we parameterize g to support an arbitrary length of model rollouts (K) during planning:
K
4o (T) = (= 80) T (ao I 20) [ p(st st-1, 1)-1) Zt-1, ae- 1 ze)
Proof. We derive lower bound the RL objective for K-step latent rollouts.
16

---

Published as a conference paper at ICLR 2023
log Epp(r)
(1 (1-7) Σ ((((0))
R
log (rg(r)
(1-?)
Σ
Ar(s, a,) + +
1=0
log (PX((M)
B
K 1)((s,(()) + K)Q(sH an)
=log| log /ʃ PR(H))POT| H H) (V) drdH
PI((H))OB / asll
| H H) PO(T
H1)
( H = =H) (V) (V) drdH
/ʃ PK((H)qU(TT)) log
PT(T H = H)
(o(T HI 11) (1) drdH
/ʃ PK Pk( |
HH | oo)
Σ log cs(2141 88+1) log fo(m1 (2))) log (V) drdH
Iʃ ( [
(
H
log e (24+1
(s+1)) log (2441 24, an) + log (1) drdH
K-1
*fastt
Y' (log c0(2141 $441) log mo 21, ar)) + 7)> log (oogr(s,,a) ))) + + log Qlsk ax)dT
Bgg(()
logr(s ar) + log Q(sk ox) 7' (log (2441 | 8441) - log mp(2111 22, 0z) + log
In (a), we start with the K-step version of the RL objective. For (b), we use Lemma A.2 For (d), we use Jensen's inequality and multiply by {v(r[==H)
- For (e), we apply Jensen's inequality. For (f), since all the terms inside the summation only depends on the first H steps of the trajectory, we change the integration from over H length to over infinite length
trajectories. For (g), we change the order of integration. For (h), we use Lemma A.1 and
the fact that (P) [log(I(H < K - 13r(ss, an) + 1(H K)Q(sH, an))) Σ >)m² log ((te, (t) + log ((zk,ax).
A.3 A LOWER BOUND FOR LATENT- LATENT-SPACE LAMBDA WEIGHTED SVG(K)
Using on-policy the data, the policy returns can be estimated using the k-step RL estimator (RI.*) which uses the sum of rewards for the first k steps and truncates it with the Q-function.
k-1
RL k (1
rr ar) + +~*((sk,00)
Like previous works (Schulman et al., 2015 Hafner et al. 2019). we write the RL objective as A-weighted average k-step returns for different horizons (different values of k), to substantially reduce the variance of the estimated returns at the cost of some bias.
Epe(r) (1 (1-7) 7) Σ 7rr at) = Eps(r)
(1 - A) Σ A RL
17

---

Published as a conference paper at ICLR 2023
Our K-step lower bound C involves rolling out the model on-policy for K time-steps. A larger value of K reduces the dependency on a learned Q- Q-function and hence reduces the estimation bias at the cost of higher variance.
Lemma A.A. We show that a x- X-weighted average of our lower bounds for different horizons (different values of K) is lower bound the RL objective.
log Bp)(r)
(1 (1-?) Σ 7()(s))
po(r)
(1 1) Σ 1 RL#
k=1
log(1 - 1)
B Bpo(r) [RLM
k=1
log 18 Po (k)
[pp(()) [
E P (A) log [pp(T)) [RL []]]
2 (1 A)
-ー log [ppp(T) [RL*]]
00 (1 A)
In (a) we use X-weighted average estimation of the RL objective. In (6) we use linearity of expectation. In (c). we note that for every k, the coefficient of RL is actually the probability of k: under the geometric distribution. Hence, we rewrite it as an expectation over the geometric distribution. In (d). we use Jensen's 'ss inequality. (e) we write out the probability values of the geometric distribution. In (J), we use Theorem 3.1 for every value of k.
A.4 DERIVING THE OPTIMAL LATENT DYNAMICS AND THE OPTIMAL DISCOUNT DISTRIBUTION
We define K (H) be a learned discount distribution over the first K 1 timesteps (0, 1.... K}. We use this learned discount factor when using data from imaginary rollouts. We start by lower bounding the RI. objective and derive the optimal discount distribution (H) and the optimal latent dynamics distribution q" (T H) that maximizes this lower bound.
a log Er (1))
Epo(+|H=H)
1111 K K - 13r(sH,an) + 11H = KK)IIKK an)
b log /f
(H) (H.
H)
Pk (H) Hybb drdHl
H)
? Taxx(H)) / golt H) log Pk (x(H)) H)) drdH
log
K (H) H)
Given a horizon E f1. ... K}. the optimal dynamics qd (F H) can be calculated analytically:
PolT Hy
(r H) d
J Hyddr'
18

---

Published as a conference paper at ICLR 2023
We substitute this value of q" in equation (c):
PolT
H)g
log Pk (K(H) Hyp J Po(T') Hyghr'
drdH
p Hybdr
(o. H)
1 *f (6x log
Yo.k (H) ʃ po(t' H)
(k(H)
dH
We now calculate the optimal discount distribution analytically 12 (1I):
(k(H) Po(T') H)
2o.k (H) _ - po(T' H H)ddT
In (a), we rewrite the RL objective using Lemma A.2 In (b), we multiply by xa.x(H))e H)
Ja.k (7) Hy- In
(c), we use Jensen's 's inequality. In (d) and (g) we use the method of Lagrange multipliers to derive optimal distributions. We use Lemma A.3 for this result. Below we write the optimal latent-space in terms of rewards and Q-function.
Writing out the optimal latent-space distribution: Markovian and do not match MDP dynamics, but are optimistic towards high-return trajectories.
The optimal latent-dynamics are non-
Po(TlH)r(smaa)) Epg (( Pa (Q(en an) Eps (Q(6))
H E [1, - -
(r H)
H= K
Writing out the optimal discount distribution:
(1-7)> Ep (r(enaa))
H K0,K -
Bpg ((6).a))
(H)
Epg. (Qssnaa))
== H > K
Bpp [(16..a6]
A.5 TIGHTENING THE K-STEP LOWER BOUND USING THE OPTIMAL DISCOUNT DISTRIBUTION AND THE OPTIMAL LATENT DYNAMICS
We start from equation (f) from Appendix A.4 the previous proof which is lower bound the RL objective. To derive (f), we have already substituted the optimal latent-dynamics. We now substitute the value of 70.x (H) and verify that the lower bound (f) leads to the original objective, suggesting that the bound is tight.
/ No.k (H) log
Jo.K (H) [ 'oll' H))dd') dh
(H)(H)
/
Px(H) J Po(T' Hyddr'
PutH EK Pk(H) Po(T') Hyddd' PxtHTJ] H-H
[ polt' Hyjjdr") He
log
ΣØ Px (K(H)) polTT' HybdT
log log(Et (E Pk(H) S po(T' HYddr')
PK(H)) Pi(tt f H Hi dH
K Σ PK(H)) HyddT'
log Epp(t)
(1 (1->))_ Ertrot,))))
In (h) we substitute the value of Y (H) and cancel out common terms. Since the integration in (i) is over H. we bring out all the terms that do not depend on H. For (k), we use the result from Lemma A.2
Hence we have proved that the bound becomes tight when using the optimal discount and trajectory distributions.
19

---

Published as a conference paper at ICLR 2023
A.6 TIGHTNESS OF LOWER BOUND WITH LENGTH ROLLOUTS K
Proof. Proving L 2 CK+1 for all K E will do.
C
=Bag (r)
Y' eloges(z4+1 |8441) log mo(21+1 2t, 0z) (1 - 7) log logr(st, at)) +7 log ((sk,ak)
Bog (r)
/f(8, at 8141) + + log (x0x (11 ((1-)) -))
())
ff(s, ae, 8141) +7 E [Flsk ak+,11)) +7 x log Q(SK+),
Bggti % (T)
Σ Tfs,, ae 8141) t k+2 log ((5K+11 CK+1
A.7 A LOWER BOUND FOR LATENT-SPACE SVG(K) IN OFFLINE RI.
In the offline RL setting Levinee et al. 2020 Prudencio et al. 2022). we have access to static dataset of trajectories from the environment. These trajectories are collected from one or more unknown policies. We derive a lower bound similar to Theorem 3.1 The main difference is that in Theorem 3.1 once a policy was updated (ø, - 1++1), we were able to collect new data using it, while in offline RL we have re-use the same static data. Similar to Theorem 3.1 define the distribution over trajectories P(T):
Polb(T) A Po(so)
II (=0
p(s (1, st)es(2 ss)
such that the offline dataset consists of trajectories sampled from this true distribution. We do not assume access to the data collection policies. Rather, To (b, s,) is the behavior cloning policy obtained from the offline dataset. We want to estimate the RL objective with trajectories sampled from different distribution qa(T):
(T) = po(so)e so) II p( | | ( 2,, a, (a 2x), t=
which leads algorithm that avoids sampling high-dimensional observations.
20

---

Published as a conference paper at ICLR 2023
Similar to Theorem 3.1
want to find encoder c0(2) st), policy Ta(a, 8t), and model
(2+1111 2e, ae) to maximize the RL objective:
max log E (p))(T))
(1-?) Σ 7 ai)
K-1
= log E (r)(())
(1-7)
Σ >'r(8), 0x) + + ((tt
1l0g/ // PR(H))OS(T | H 1H= -H) (V) drdHl
PM((H)log/ / (0(T H- H) Po.b(T (T H = H)
(w) drdH
( H = [H)
/ʃ | | H) log
PT. H = H)
(V) drdH
(T
H = H)
/ʃ PR(H)) | 110==0) -
(
Σ log( e cs(2t+1 mo+1)mo(a441 1841)) ( + log (V) drdll
Tes as (To) (u (2411)
Mg (=(+1
& [
H
log( ca(2
S1+1)Tb (ac-1 88422
+ log (V) drdH
ma( m
2r, (1) TT ze+1)
1 ffast (o(T))
K-1
Σ
log( co(z041
m=( Zes ag az) (0c+1 2(41) + (1 >in (sgr(s +n log ((sk,ax)d
$1+1)=b(as)
s641) )(tt))))
))
1=0
K
log( ca(z++1 se41), +n log( To To(ae+1 8t+1), +(1 - 7)n (ogr(ss, ar) + 7 log Q(sk ax)
1 o
Σ
mz(z1+1 Ze, ax)
Th (a(+1
z0+1)
consistency
Behaviour cloning
The only main difference between this proof and the proof A.2 of Theorem 3.1 in step (f), where canceling out the common terms of p(T) and g(T) leaves an additional behavior cloning term. This derivation theoretically backs the additional behavior cloning term used in prior representation learning methods (Abdolmaleki et al. 2018 Peters et al.) (010) for the offline RL setting.
OMISSION OF THE LOG OF REWARDS AND FUNCTION IS EQUIVALENT TO SCALING KL A.8
IN EQUATION g
While our main results omit the logarithmic transformation of the rewards and Q function, in this section we describe that this omission is approximately equivalent to scaling the KL coefficient in Eq.
Using these insights, we applied ALM, with the log of rewards, to a transformed MDP with a shifted reward function. large enough constant, a, was added to all original rewards(which doesn't 't change the optimal policy assuming that there are no terminal states)
"new r'
Taking log of this new reward is equal to the reward of the original objective, scaled by the constant,
allog(r + a log(a)) x
The additive term can be ignored because it won't it contribute to the optimization. We plot both y and y = (log(r + a) log(a)), to show that they are very similar for commonly used values of rewards in Figure 6 The scaling constant can be interpreted as the weight of the log of rewards,
21

---

Published as a conference paper at ICLR 2023
relative to the KL term in the ALM objective. Hence changing this value is approximately equivalent to scaling the KL coefficient. The results, shown in Fig 12 show that this version of ALM which includes the logarithm transformation performs at par with the version of ALM without the logarithm. This results shows that we can add back the logarithm to ALM without hurting performance. For both Figure and Figure T2 we used the value of a 10000.
ALM objective using the transformed reward
K
2 (log 60(411 86-1) - log ms (2t+1 24, ()) + (1 - log""" (st, (+), + + log Quew (SK: ax)
4 (r)
K
Y' (log (a 1841) log m =t, + (1 7)7* (r(s), + log(a)
E B (r)
Σ tm
t ((sx,ax) loss(a))
? (log c0(2141 | (841)) log mo( | 21, a)) + (1 7)/(( (()),)tt) + 7 (lxx,ax)
2 (T)
Bg (r)
y'a (log ea (=(+11 | 84+1) - log m 2t, (z) + (1 - 7) at) + (
In (a), we write the ALM objective using the new reward function. In (b), we use the fact that log(Tnew($t,, at)) ~ r(s..))) + log(a)). and log Qnew (SK, ax) = (Q(a + log(a) for a large enough constant a. (See explanation above). In (c), we remove the extra constants which add up to 0. In (d), note that scaling the rewards and Q function by 1/a, is equivalent to scaling the KL term by a, which a large enough constant.
Original reward vs ALM reward on transformed MDP
atlog(r a) log(a))
17.5
reward output
10.0
7.5
2.5
0.0
0.0
2.5
5.0
7.5
10.0
12.5
15.0
17.5
reward input
Figure 6: Analyzing logarithmic reward transformation: While theoretical derivation motivates returns, we show that appropriately scaling and shifting rewards
logarithmic transformation makes that logarithmic transformation have no effect. This throws light on the fact that omitting the logarithmic transformation can he thought of as scaling the two components in our augmented reward function Eq 6
B COMPARISON TO BASELINES
Here provide quick comparison baselines. conceptually in terms of number of gradient updates per env step, usage of ensembles are used for dynamics model or Q-function, and objectives for representations, model, and policy.
22

---

Published as a conference paper at ICLR 2023
Table Table showing conceptual comparisons ALM baselines.
Ensembles Representations
Model
Policy
UTD
ALM
Lower bound
Lower bound
Lower bound
3
SAC-SVG
X
MLE
MLE
Actor-Critic (stochastic value gradients)
4
MBPO
MLE
Soft Actor-Critte (dyna-style)
20-40
REDQ
X
Soft Actor-Critic
20
C ADDITIONAL LEARNING DETAILS
Estimating the Q-function and reward function. Unlike most MBRL algorithms, the Q function Qo(2t, 0x) is learned using transitions 8t, at, Tt, 8t+1 from the real environment only, using the standard TD loss:
LQ (ze ~ eal: 8c), at, Ft; Fe41 es(: 8t+1)) (Qu(ze, 0r) (re + 7Quum (z++1, (zx-1))))
(10)
The TD target is computed using target Q-function (Fujimoto al. 2018 We learn the reward function ro r= az) using data 8t, de. T from the real environment only. For ro(=,, at), by minimizing the mean squared error between true and predicted rewards:
L (Zt ~ epl: 8c), oe, (ra (2t, ax) rz)??
(11)
Unlike prior representation learning methods (Zhang et al. 2020). do not use the reward or Q function training signals to train the encoder. The encoder, model, and policy are optimized using the principled joint objective only.
D IMPLEMENTATION DETAILS
We implement ALM using DDPG (Lillicrap al. 2015) the base algorithm. Following prior svg methods (Amos et al. 2020). we parameterize the encoder, model, policy, reward, classifier and Q-function 2-layer neural networks, all with 512 hidden units except the model which has 1024 hidden units. The model and the encoder output multivariate Gaussian distribution over the latent-space with diagonal covariance. Like prior work (Hansen al., 2002)] Yarats al. (021)) we apply layer normalization (Ba et al., 2016) to the value function and rewards. Similar to prior work (Schulman al., 2015) Hafner al. 2019), we reduce variance of the policy objective in Equation 8] by computing an exponentially
average of the objective for rollouts of length to K. This average is also a lower bound (Appendix A.3). train the policy, reward, classifier and Q-function we use the representation sampled from the target encoder. For exploration, we use added normal noise with a linear schedule for the std (Yarats et al. (021). All hyperparameters are listed in Table The brief summary of all the neural networks, their loss functions and the inputs to their loss functions are listed in Table
E ADDITIONAL EXPERIMENTS
Analyzing the learned representations. Because the ALM objective optimizes the encoder and model to be self-consistent, expect the ALM dynamics model remain accurate for longer rollouts than alternative methods. We test this hypothesis using an optimal trajectory from the HalfCheetah-v2 task. Starting with the representation of the initial state, autoregressively unroll the learned model, comparing each prediction 2 to the true representation (obtained by applying the encoder observation sy). Fig. (left) visualizes the first coordinate of the representation, and shows that the model learned via ALM remains accurate for ~ 20 steps. The ablation of ALM that removes the KI. term diverges after just two steps.
Bias and Variance of the Lower Bound. We recreate the experimental setup of (Chen al. 2021) to evaluate the average and the std value of the bias between the Monte-Carlo returns and the estimated returns (lower bound LK (s, a) for our experiments). Similar to (Chen et al. 2021) since
23

---

Published as a conference paper at ICLR 2023
Table 3: default set of hyper-parameters used in our experiments.
Hyperparameters
Value
Discount (n)
0.99
Warmup steps
5000
Soft update rate (7)
0.005
Weighted target parameter (A)
0.95
Replay Buffer
10 for humanoid
10 otherwise
Batch size
512
Learning rate
le-4
Max grad norm Latent dimension Coefficient classifier rewards
100.0
50
0.1
Exploration stddev. clip
0.3
Exploration stddev. schedule
linear(1.0 0.1, 100000)
Table 4: Neural networks used by ALM
Neural Network
Loss Function
Inputs to Loss
Encoder: (s) +
joint lower bound (Eq. 5
{st, ass 8141 14+--1
Model: mz(2,a) a) + x'
Policy: (=) +
joint lower bound (Eq. 5)
Zt
Q F-uncttion Q0(2, a) R
TD-loss (Eq.
2t, Oc, Ptss 2141
2t, at 2+11 ep (8t+1) 2641 mo dc)
Classifier: Co(2, a, x²) E (0,
cross entropy loss (Eq. 9)
Reward Function: re(2,a) a)
IR
MSE (Eq. 11
02,
the true returns change training progresses analyze the mean and std value of normalized bias, which is defined as: (C (s, a) (s, a))/|E> a as), This helps us to evaluate bias relative the scale of average Monte-Carlo returns of the current policy. Results for Ant-v2 and Walker2d-v2 are presented in Fig.
Additional Ablations.
In Table 5 we train ALM using the soft actor critic entropy bonus and compare with SAC-SVG.
In Figure we show that ALM is robust to range of coefficients for the KL term in Equation
In Figure 12
we incorporate the logarithmic transformation of the reward function based on Appendix A.S. and show that does not hurt performance.
In Figure we compare ALM version MnM [Esenbachh al. 202la) to show that learning the encoder is necessary for good performance on high dimensional tasks.
In Figure we show that ALM achieves higher asymptotic returns when compared to SAC.
In Figure
15
compare ALM to version of ALM which uses reconstruction loss to learn
the encoder.
Figure 16
we show that ALM works well even when using a linear classifier.
In Figure we add noise in the MDP dynamics to show the performance of ALM with varying alcatoric uncertainty.
In Figure 18 and Figure 19 we compare ALM to a version of ALM which additionally optimizes the encoder and the latent space model predict the value function.
24

---

Published as a conference paper at ICLR 2023
In Figure 20 show the average final episodic returns achieved by ALM for different coefficients for the KL term in Equation g
ALMM3
No
... Ground Iruin
1e1
2th sepreertaiioo condiaate 0101
0.15
10
20
ALM whor
0.05
trjectoryy
steps
Figure 7: Analyzing the learned representations: (Left-top) The ground truth representations obtained from the respective trained encoders on the same optimal trajectory. (Left-bottom) Without the KL term, representations learnt are degenerate, i.e. they correspond to same value for different states. (Right) The KL term in the ALM objective, trains the model to reduce the future K step prediction errors. The latent-space model is accurately able to approximate the true representations upto 20 rollout steps.
ALME3I
REDQ
SAC-AVG
overoge blss
210
3.0
6.0
CTV STOpS
Steps
CRV stops
steps
Figure 8: Bias and Variance of the Lower Bound: In accordance with what theory suggests, the joint objective biased estimate the true returns. The std values are uniform throughout training and consistently lower than REDQ, which could be the reason behind the sample efficiency of ALM(3)
ALM(3)
MBPO
SAC
REDQ
103
TruncaeedHumanndd-222
163
TruncaledAnt-v2
103
Walker2d-v2
average returns
eni steps
env steps
env steps
Figure Additional sample efficiency experiments.: ALM generally matches the sample efficiency of MBPO and REDQ at fraction of the computation complexity.
25

---

Published as a conference paper at ICLR 2023
AUM(3]
no KL
ALMEL:
ALMOSI SAC
classifier
value
Humaneidav2
Hummaniio-22
arenge relurs
ALM no model
TD3 SAC
TD3 lours)
Welker2d-v2
Walker2d-u
Waiker2dw2
2v0r retums
Steps
stops
any stops
(a) Additional results for compar- ing different components of the joint objective on Humanoid-v2 and Walker-v2. The KL consis- tency term is the most important contributing factor.
(b) Additional results for compar- ing different horizon lengths. All horizon lengths generally perform better than sac, but horizon length 3 performs best.
(c) Additional results for com- paring ALM(3)'s representations with pure model-free RL meth- ods.
Figure 10: Additional ablation experiments.
0.5
0.2
0.1
0.15
0.05
0.0
Humanoid-v2
Truncaeddfuumnnidee
Ant-22
Truncated/ot--22
averago retumss
CTW steps
CIVV steps
env steps
166
cnV steps
Figure 11: ALM is robust across different values of the coefficient for the KL term in Equation In our implementation, we deviate from the value 1, hecause leads to relatively gradual increase in returns on some environments. This is expected because higher coefficient the KI. term leads higher compression Eysenbach) at. 202Th). We add that prior work variational inference Wenzel et al. (020) also finds that scaling the KI. term can improve results.
Table 5:
Where possible, have tried to the same hyperparameters and architectures SAC- SVG. We use the same training hyperparameters like update to data collection ratio, batch size and rollout length when compared to sac-svg. Both methods the same policy improvement technique: stochastic value gradients. The two exceptions are decisions that simplify our method: unlike SAC-SVG, we use feedforward dynamics model instead of RNN; unlike SAC-SVG, we use simple random noise instead of a more complex entropy schedule for exploration. To test whether the soft actor critic's entropy, used in SAC-SVG can be a confounding factor causing SAC-SVG to perform worse than ALM, we compare a version of ALM which uses soft actor critic entropy bonus like the SAC-SVG.
Method
T-Hmannoid-22
T-Ant-v2
ALM-ENT(3)
4747 + 900
4921 + 128
ALM(3)
5306
4887 + 5306 437
SAC-SVG(3)
501 = 307
SAC-SVG(2)
472
3833 +
26

---

Published as a conference paper at ICLR 2023
ALM(3)
Transformed rewards
Semi log rewards
Log rewards
149
Humanold-v2
1s3
TruncatedFumanoide22
104
HaltCheetan--2
average retums
0.8
0.6
0.4
env Steps
Tw5
env steps
1e5
env steps
Figure 12:
Using the logarithm does not actually hurt performance. Our theory suggests that we should take the logarithm of the reward function and Q-function. Naively implemented, this logarithmic transformation (pink) performs much worse than omitting the transformation (green). We also see that using a log of rewards for only training the encoder and model does affect the performance (blue). We hypothesize that the non linearity of log(x) reward values makes the Q-values similar for different actions. However, by transforming the reward function (which does not change the optimization problem), we are able to include the seooetically-ssugeeee logarithm while retaining high performance (red). See Appendix A.8 for more details.
ALM(3)
Identity encoder
TruncatedHHmmanoid--
TruncatedAnt-v2
average returns
6
env steps
env steps
Comparison with MnM ablation. The main difference between ALM and prior joint optimization method (MnM), is that ALM learns the encoder function. Replacing learned encoder with an identify function yields method that resembles MnM, and performs much worse. This result supports our claim that RL methods that use latent-space models can significantly outperform state-space models.
Figure 13:
ALM(3)
SAC
Humanoid-v2
1e3 TruncatedHumannidd-y
average returns
steps
1e6
steps
1e6
Figure 14: Asymptotic performance. Even after I million environment steps, ALM still outporforms the SAC baselinc.
27

---

Published as a conference paper at ICLR 2023
ALM(3)
Reconstruction MBRL
163
104
HalfCheetah-v2
163
TruncaeddHumanoiévv2
12
average returns
1.
06-
02
1e5
165
105
Ant: v2
Humanaid x2
average retums
105
env steps
Figure 15: Importance of an aligned objective. Comparison with version of ALM that learns policies using the objective, but learns representations and latent-space models to maximize likelihood. This further verifies claims made Section 5.1 that using an aligned objective for joint optimization leads to superior performance
ALM(3)
Linear classifier
1+1
2
average returns
0.8
Iw5
165
TruncatediHumannidd v2
TruncatedAnt-v2
average returns
steps
165
185
Figure 16: Robustness classifier ALM retains high degree of performance when the classifier restricted to be linear function, showing how ALM can be robust to errors the classifier.
28

---

Published as a conference paper at ICLR 2023
0.01
0.001
0.0
HaltCheetah-v2
1.2
average retures
1.0
0.6
0.4
0.2
0.0
steps
Figure 17: Robustness towards various levels aleatoric uncertainty. We created version HalfCheetah-v2 task with varying levels of aleatoric noise.
ALM(3)
ALM(3) Critic
Gritic
1+1
Humanold-v2
HalfChhetah-vz
1n3
averege reaurs
0.
0.0
165
165
165
TruncatodHumanodd. v2
TruncatedAnt-v22 .x2
average reaums
M MMAAN
WwM
env
145
steps
145
Figure 18: Effect of learning to predict the value function.
29

---

Published as a conference paper at ICLR 2023
ALM(3)
Latent space alm critic
space critic
Humanoidivv
HaifCheetah-v2
Ant-v2
average returns
1.0
02
1+3
TruncadedHumanoid-22
163
TruncatedAnt-42
average returs
env steps
env steps
Figure 19: Effect of learning the encoder as well as latent-space model to predict the value function.
0.5
0.2
0.1
0.15
0.05
0.0
Humanoid-v2
Truncatedllumaniddiv2
Ant-v2
Truncateddnt-v22
Cunulative training training returts
2000
0000
3520
2000
1000
3430
2550
2420
3220
3030
2700
2830
2500
3700
0.3
25
02
03
0.2
Figure 20: Returns at different values for classifier coefficient.
30

---

Published as a conference paper at ICLR 2023
F COMPARISON TO PRIOR WORK ON LOWER BOUND FOR RL (MNM)
Our approach of deriving an evidence lower bound on the RL objective is similar to prior work EE senbach al. 2021a). la). In this section we briefly go over the connection between our method and Eysenbach et al. 202 ap. The lower bound presented in (Eysenbach et al. 2021a) is a special case of the lower bound in 3.1 By taking the limit K - 3o and assuming identity function for the encoder, we exactly reach the lower bound presented in Eysenbach et al. (2021a). By using bottleneck policy (policy with an encoder), ALM(K) learns to represent the observations according to their importance in the control problems rather than trying to reconstruct every observation feature with high fidelity. This is supported by the fact that AL. solves Humanoid-v2 and Ant-v2, which were not solvable by prior methods like MBPO and MnM. By using the model for K steps, we have a parameter to explicitly control the planning Q-learnig bias. Hence, the policy faces a penalty (intrinsic reward term) only for the first K steps rather than the entire horizon (Eysenbach al. 2021a), which could lead to lower returns on the training tasks (Eysenbach et al. 2021b).
G FAILED EXPERIMENTS
Experiments that we tried and that did not help substantially:
Training the critic using data generated from the model rollouts. The added complexity did not add much benefit.
Value expansion:
Warm up steps: training.
Training the policy using real data for a fixed time-steps at the start of
Horizon scheduling:
Scheduling the sequence length from to K at the start of training.
Exponential discounting: Down-scaling the learning rate of future time-steps using tempo- ral discount factor, to avoid exploding or vanishing gradients.
Experiments that were tried and found to help:
Target encoder: Using target encoder for the KL term in Eq. 7 helpped reduce variance in episode returns.
Elu activation: Switching from Relu to Elu activations for all networks for ALM resulted in more stable and sample efficient performance across all tasks.
31


Reviews:
[{"review_id": "ojZDTcrdh0", "reviewer": "Reviewer_JaMR", "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 builds on the latent variable view of reinforcement learning (RL) to provide a model-based RL method that jointly learns representations of observations, their transition dynamics (the representations), and the optimal policy.\n\nIt introduces the Aligned Latent Model objective, which is derived from the typical latent-variable view of RL, and a particularly defined variational distribution. This is then used in a DDPG style RL algorithm.\n\nExperiments compare against several baselines, including SAC-SVG on several complex tasks. Key questions regarding the usefulness of the contribution is analysed using appropriate experiments, including ablations.", "summary_of_the_review": "In general, the paper is written well, and introduces its concepts cleanly. \n\nSome details are used without defining terms appropriately. \n\nIn order for me to improve my rating, I would like the questions / concerns I had raised addressed. In particular, I am worried that some conclusions are drawn without appropriate control over the experiments. ", "strength_and_weaknesses": "### Strengths\n- The theory is built nicely, by introducing the notation first, followed by a high level outline before diving into the maths.\n- The technique appears to have some good benefits.\n- Discussions are nice, with good explanations of points made. \n- Good literature review.\n- It was good that network hyperparameters were chosen to be consistent with baselines.\n\n### Weaknesses\n- See questions below.", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "kf3UjjdDsc", "reviewer": "Reviewer_11kg", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "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": "A theoretical sound model based rl algorithm derivation that can process high dimensional inputs (i.e. images) that derive one objetive function to learn the latent space of the observations, the model and the policy, which usually is done using separate objectives. This work also provides a practical algorithm, that was shown empirically to have. greater sample efficiency that SOTA. The work also provide a lower bound on the RL loss.", "summary_of_the_review": "This work derive a lower bound on the RL loss with a strong theoretical , well supported, derivation, constructing a practical model based RL algorithm, that have higher sample efficiency than SOTA, its end-to-end learning solution, that can handle high dimensional state spaces learning the a latent representation for this states,  ", "strength_and_weaknesses": "Very strong derivation on the objective and losses. End-to-end learning also is another strength of the paper. High sample efficacy for high dimensional states also is a strength. Outperforming model free RL on high dimensional states also seems to be novel, without using pre-training or other self-supervised technique,\n\nWeakness: Classifier training time seems not to be taken in account, also if there are situation where learning the classifier is difficult and how this could impact the overall algorithm. Automatic K selection would be a good addition into the algorithm.\n\nIt would interesting to see the performance of algorithm with different domains that capture different levels of aleatoric uncertainty.\n\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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "WQVtP8uW6sC", "reviewer": "Reviewer_PJmH", "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 authors propose to train a policy, a state-encoder, and a dynamic model in model-based RL with the novel evidence lower bound on expected returns loss. In this way policy, encoder, and dynamic model have the same objective to maximize which differs this work from the previous ones. \n", "summary_of_the_review": "Overall the paper is well-written and motivated. The experimental results are very interesting. However, in my opinion, some key experiments and baselines (e.g. [1] Learning and Planning in Complex Action Spaces) are missing.", "strength_and_weaknesses": "In model-based RL algorithms like SAC-SVG (Amos et al., 2020 [2]) one needs to train an actor as well as a critic in a way that actor-model output should be the input to the critic-model to enable backpropagation to the actor through the critic. This makes the boundaries between an actor and a critic vague. For example, the representation model and dynamic model can be parts of the actor and optimized with actor loss, or be parts of the critic and optimized with critic loss. In more complex cases these models can be jointly optimized with critic and actor losses, or even with some auxiliary losses (such as state reconstructions), or with some combinations of the before-mentioned losses.  \n\nOrthogonal to these choices is the decision whether to use evidence lower bound on these objectives or not. As far as I know, this paper is the first work that successfully explores the possibility of training a representation model and a dynamic model only with actor loss and it did so using evidence lower bound.\n\nThe authors discuss that the proposed objective leads to self-consistency for the actor. However, I want to point out that the paper's objective leads to inconsistency for a critic as the authors train the critic and the reward subnetwork on top of representations that are optimized for the actor and not for the critic. Plus, I believe it is an open question whether the actor's consistency or the critic's consistency is more important, or whether one needs to jointly optimize representation with the actor's and the critic's objectives. It would be great if the authors can comment on this. Plus, I recommend adding experiments if it is possible: \n1) try to optimize the representation model and the dynamic model with critic and reward losses but optimize the policy model with objective (8) derived in the paper. \n2) try to optimize the representation model and the dynamic model with the actor (objective (8)), critic, and reward objective jointly.\n3) the same as (1) and (2) but change only the representation model objective as it seems to be easier with the current coding. \n\nIn the paper, the main baseline (SAC-SVG) uses an auxiliary state-reconstruction objective and performs worse than the proposed method. However, in my opinion, the actor objective is much more aligned with the critic objective than with auxiliary objectives such as state-reconstruction objectives or SSL-objective. So I miss a comparison with the strong baseline that uses only actor and/or critic objectives to train model parts (e.g. with [1]). \n\nStrength\n1) novel evidence low bound on RL objective is derived;\n2) novel loss is well-explained, motivated, discussed, and analyzed. (E.g. that the loss encouraged mode-seeking behavior)\n3) The method shows that axillary objectives (such as state reconstruction) may be hurtful to the performance. \n4) The paper is well-written;\n5) The authors show that the proposed algorithm beats similar SAC-like algorithms that use auxiliary objectives as state reconstruction to learn representations. Plus, the proposed method is less computationally demanding than methods that use ensembles (MBPO, REDQ), but archives similar or better performance. \n6) An ablation study is present.\n\nMinor issues and questions\n1) Authors claim that the proposed method is robust to the number of rollouts, but in Fig.5b (bottom) ALM(5) performance is quite low. \n2) Why not use the same parametrization, objective (8) with the classifier, to train the encoder and dynamics model?\n3) Maybe a small table that shows the main differences between the proposed algorithm and the baselines would be helpful;\n4) In the introduction, the authors mention that MuZero uses value prediction to train a representation model and a dynamic model. I would like to point out that MuZero uses not only value prediction loss but also use the loss for the policy model (actor objective) to train these models. So representation model and a dynamic model are jointly optimized with actor and critic losses in MuZero and I argue this makes representation consistent with both actor and critic objectives.\n\n[1] Hubert et al, Learning and Planning in Complex Action Spaces, 2021.\n\n[2] Amos et al, On the model-based stochastic value gradient for continuous reinforcement learning, 2020.\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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "9YsB7K_ASBP", "reviewer": "Reviewer_km6q", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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 a new objective for model-based RL that lower bounds the true objective and can be used to optimize the encoder, model, and the policy at once. The proposed model shows higher sample efficiency compared to the baselines.", "summary_of_the_review": "The paper proposes a nice approach to unify (and simplify) model-based RL objective. With some additional few tweaks, the proposed methods shows high empirical performance compared to the baselines. ", "strength_and_weaknesses": "Strengths\n- Recent works (especially on offline RL) have shown that we can find simpler alternatives to current RL objectives without compromising performance. This paper adds up nicely to those works by showing that we can also simplify model-based RL objectives.\n- The work proposes a principled way to train state encoders for RL.\n- The proposed method has higher sample efficiency compared to the baselines.\n\nWeaknesses\n- As the authors note, it seems the method still requires a few tweaks to work well empirically. For example, we need to omit the log of the true rewards and scale the KL term in the policy objective to 0.1. While the authors provide a brief intuition on why those modifications are needed, I think the authors should provide more concrete analysis (e.g., empirical results) on what problems the original formula have and how the modifications fix them. Also, it would be better if the authors provide ablation results on those modifications. For example, does the performance drop drastically if the scale of the KL term changes (to 0.05, 0.2, 0.5, ...) ?\n- The compute comparison vs. REDQ  on Figure 3 seems to be misleading. First, less runtime does not necessarily mean less computational cost. Second, if the authors had used the official implementation of REDQ [1], it should be emphasized that this implementation is very inefficient in terms of runtime. In detail, the implementation feeds forward to each Q-network one at a time while this feed-forward process is embarrassingly parallelizable. The runtime of REDQ will drop significantly if we parallelize this process.\n- The ablation experiments seem to show that the value term for the encoder is not necessary. It would be better to provide an explanation on this result.\n- Some of the equations seem to have typos. 1) On equation (4), the first product and the second product have exactly the same form. 2) On algorithm 1 Line 8, shouldn't we use s_n instead of s_t?\n\nQuestions\n- I am curious of the asymptotic performance of the proposed method. If possible, can the authors provide average return results with more env steps?\n\n[1] https://github.com/watchernyu/REDQ", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "5c3lKMHuyY", "reviewer": "Reviewer_RXsd", "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 proposes a proxy model-based RL objective which is proved to be the lower bound of the usual RL objective. The proxy objective can be optimized with respect to the policy, the latent-space model, and the representation. Experiments on a range of continuous control tasks show that a practical algorithm loosely based on the proxy objective achieves better sample complexity than several competitive existing model-based and model-free algorithms.", "summary_of_the_review": "Overall, this is a nice and well-written paper. It develops an extension of an existing objective (Eysenbach et al., 2021). It has strong empirical results for the practical algorithm it proposes. However, in my opinion, there is a major weakness - the practical algorithm deviates from the theoretical analysis without providing enough empirical/theoretical analysis.\n\nSome questions to be addressed that impact the score:\n1. What empirical evidence motivates using a latent-space model instead of a state-space model? How does it compare to MnM (Eysenbach et al., 2021)?\n2. What empirical/theoretical evidence motivates omitting the log of rewards and Q function? Could you provide further justification for it?\n3. How does the coefficient of the KL term affect performance?\n\nSome suggestions for improvement that do not impact the score:\n1. The definition of the Q function is in the appendix. It should be introduced in the main text before using it in Eq (5).\n2. In Eq (7), I believe the KL term is reversed. It should be written as $KL(m_\\phi(\\cdot | z_i, a_i) || e_{\\phi_\\text{targ}}(\\cdot | s_{i+1}))$. Otherwise, it is inconsistent with Eq (6) and the discussion in Sec 3.4.\n3. Inaccurate/misleading statement after Eq (8): “we reduce the variance of this objective by computing this objective for multiple horizons and then taking an average.” However, in the experiment section, only fixed horizons are used.\n4. Typo in the paragraph before the conclusion section: “we that”.\n5. Typo in Lemma A.2: $Q(s_t, a_t)$ should be $Q(s_K, a_K)$.\n6. In the proof of Theorem 3.1, I suggest changing ${d\\tau dH}$ to ${dH d\\tau}$ in equation g for better clarity.\n7. In the proof of Theorem 3.1, it should be $H-1$ instead of $\\max(0, H-1)$ above the summation. The latter will incorrectly calculate $\\log e_\\phi(z_1 | s_1) - \\log m_\\phi(z_1 | z_0, a_0)$ twice.\n8. Appendix B, the comparison to prior work on lower bounds for RL (MNM), should be moved to the main text since it is the most relevant related work.\n9. It would be nice if there were an analysis of the gap of the lower bound with respect to the model error.\n\nBenjamin Eysenbach, Alexander Khazatsky, Sergey Levine, and Ruslan Salakhutdinov. Mismatched no more: Joint model-policy optimization for model-based rl, 2021. URL https://arxiv.org/abs/2110.02758.\n", "strength_and_weaknesses": "The strength of the paper is that the proposed practical algorithm achieves superior performance on some continuous control tasks without using the common ensemble technique for MBRL. In addition, an abundance of ablation experiments is performed to validate parameter K's sensitivity, the components' effectiveness, and the learned values and representations.\n\nHowever, there is a major weakness in this paper. The algorithm is inconsistent with the theory. Notably, the log of rewards and Q function in Eq (5) is omitted in the actual implementation. In addition, an extra hyperparameter is added to tune the effect of the KL term. More importantly, there is no ablation study on these two important design changes.", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster