S IMPLIFYING M ODEL - BASED RL: L EARNING R EPRESENTATIONS , L ATENT - SPACE M ODELS , AND P OLICIES WITH O NE O BJECTIVE
Anonymous authors
Paper under double-blind review
A BSTRACT
While reinforcement learning (RL) methods that learn an internal model of the environment have the potential to be more sample efﬁcient 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 lowdimensional 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-efﬁciency of the best prior model-based and model-free RL methods.
While such sample efﬁcient methods typically are computationally demanding, our method attains the performance of
SAC in about 50% less wall-clock time.
I NTRODUCTION
While RL algorithms that learn an internal model of the world can learn more quickly than their model-free counterparts (Hafner et al., 2018; Janner et al., 2019), ﬁguring out exactly what these models should predict has remained an open problem: the real world and even realistic simulators are too complex to model accurately.
Although model errors may be rare under the training distribution, a learned RL agent will often seek out the states where an otherwise accurate model makes mistakes (Jafferjee et al., 2020).
Simply training the model with maximum likelihood will not, in general, produce a model that is 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 is especially important in settings with high-dimensional observations, which are challenging to predict with high ﬁdelity.
[Figure]
Prior model-based methods have coped with the difﬁculty to model high-dimensional observations 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 on the states visited by the policy.
Hence, another way of reducing prediction 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 a model-based RL algorithm that automatically learns compact yet sufﬁcient representations for model-based reasoning?
In this paper, we present a simple yet principled answer to this question by devising a single objective that jointly optimizes the three components of the model-based algorithm: the representation, the model, and the policy.
As shown in Fig.
1, this is in contrast to prior methods, which use three separate objectives.
We build upon prior work that views model-based RL 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 et al., 2021a).
This is 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, a model that predicts those representations, and a policy that acts based on those representations.
To the best of our knowledge, this objective is the ﬁrst lower bound for a model-based RL method with a latent-space model.
Across a range of continuous control tasks, we demonstrate that ALM achieves higher sample efﬁciency 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., 2018; Janner et al., 2019) or decision-time planning (Deisenroth & Rasmussen, 2011; Sikchi et al., 2020; Morgan et al., 2021), our open-source implementation performs updates 10 × and 6 × faster than MBPO (Janner et al., 2019) and REDQ (Chen et al., 2021) respectively, and achieves near-optimal returns in about 50% less time than SAC.
R ELATED W ORK
Prior model-based RL methods use models in many ways, using it to search for optimal action sequences (Garcia et al., 1989; Springenberg et al., 2020; Hafner et al., 2018; Chua et al., 2018;
Hafner et al., 2019; Xie et al., 2020), to generate synthetic data (Sutton, 1991; Luo et al., 2018; Hafner et al., 2019; Janner et al., 2019;
Shen et al., 2020), to better estimate the value function (Deisenroth &
Rasmussen, 2011; Chua et al., 2018; Buckman et al., 2018; 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 al., 2019; Clavera et al., 2020;
Amos et al., 2020), our approach uses model rollouts to estimate the value function for a policy gradient.
Because learning a model of high-dimensional observations is challenging, many prior model-based methods ﬁrst learn a 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 et al., 2018; 2019; 2020), value and action prediction (Oh et al., 2017; Schrittwieser et al., 2020; Grimm et al., 2020), planning performance (Tamar et al., 2016; Racanière et al., 2017;
Okada et al., 2017), or self-supervised learning (Deng et al., 2021; Nguyen et 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: 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 a 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; Voelcker et al., 2022) or directly optimize the model to produce high-return policies (Eysenbach et al., 2021a; Amos et al., 2018; Nikishin et al., 2021).
However, effectively addressing the objective mismatch problem for latent-space models remains an open problem.
Our method makes progress on this problem by proposing a single objective to be used for jointly optimizing the model, policy, and representation.
Because all components are 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 in a model-based setting.
A U NIFIED O BJECTIVE FOR L ATENT -S PACE M ODEL -B ASED RL
We ﬁrst introduce notation, then provide a high-level outline of the objective, and then derive our objective.
Sec. 4 will discuss a practical algorithm based on this objective.
3. P RELIMINARIES
The agent interacts with a Markov decision process (MDP) deﬁned by states s t , actions a t , an initial state distribution p 0 ( s ) , a dynamics function p ( s t +1 | s t , a t ) , a positive reward function r ( s t , a t ) ≥and a discount factor γ ∈ [0 , 1) .
The RL objective is to learn a policy π ( a t | s t ) that maximizes the discounted sum of expected rewards within an inﬁnite-horizon episode:
[Equation]
The factor of (1 − γ ) does not change the optimal policy, but simpliﬁes the analysis (Janner et al., 2020; Zahavy et al., 2021).
We consider policies that are factored into two parts: an observation encoder e φ ( z t | s t ) and a representation-conditioned policy π φ ( a t | z t ) .
Our analysis considers inﬁnite-length trajectories τ , which include the actions a t , observations s t , and the corresponding observation representations z t : τ (cid:44) ( s 0 , a 0 , z 0 , s 1 , a 1 , z 1 , · · · ) .
To simplify notation, we write the discounted sum of rewards as R ( τ ) (cid:44) (1 − γ ) (cid:80) ∞
t =0 γ t r ( s t , a t ) .
Lastly, we deﬁne the Q-function of a policy parameterized by φ , as Q ( s t , a t ) =
E s t
+1 ,a t ,z t ∼ p ( ·| s t ,a t ) ,π φ ( ·| z t ) ,e φ ( ·| s t )
[ R ( τ )] .
3. M ETHOD O VERVIEW
Our method consists of three components, shown in Fig.
The ﬁrst component is an encoder e φ ( z t | s t ) , which takes as input a high-dimensional observation s t and produces a compact representation z t .
This representation should be as compact as possible, while retaining the bits for selecting good actions and for predicting the Q-function.
The second component is a dynamics model of representations, m φ ( z t +1 | z t , a t ) , which takes as input the representation of the current observation and the action and predicts the representation of the next observation.
The third component is a policy π φ ( a t | z t ) , which takes representations as inputs 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.
3. D ERIVING THE O BJECTIVE
To derive our objective, we build on prior work (Toussaint, 2009; Kappen et al., 2012) and view the
RL objective as a latent-variable problem, where the return R ( τ ) is the likelihood and the trajectory τ is the latent variable.
Different from prior work, we 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. 1) in terms of trajectories as E p ( τ )
[ R ( τ )] by deﬁning the distribution over trajectories p ( τ ) as
[Equation]
Estimating and optimizing this objective directly is challenging because drawing samples from p ( τ ) requires interacting with the environment, an expensive operation.
What we would like to do instead is estimate this same expectation via trajectories sampled from a different distribution, q ( τ ) .
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 et al., 1999):
[Equation]
This lower bound resolves a ﬁrst problem, allowing us to 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 difﬁcult 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 a function of these representations, we are able to estimate this lower bound without sampling high-dimensional observations.
Further, we carefully parameterize the learned distribution q ( τ ) to support an arbitrary length of model rollouts ( K ), which allows us to estimate the lower bound accurately:
[Equation]
While it may seem strange that the future representations sampled from q Kφ ( τ ) 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 to high dimensional states from the environment’s dynamics function.
Combining the lower bound (Eq. 3) with this choice of parameterization, we obtain the following objective for model-based RL:
[Equation]
where ˜ r ( s t , a t , s t +
[Equation]
=
[Equation]
This objective is an evidence lower bound on the RL objective (see Proof in Appendix A.2).
Theorem 3.1.
The following bound holds for any representation e φ ( z t | s t ) , latent-space model m φ ( z t +1 | z t , a t ) , policy π φ ( a t | z t ) and K ∈ N
[Equation]
3. C ONNECTIONS WITH P RIOR W ORK
In this section, we provide intuition for our objective and relate it to objectives in prior work.
We start by looking at the augmented reward.
The ﬁrst 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 a 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., 2021b; Bharadhwaj et al., 2021;
Rakelly et al., 2021) that regularize the encoder against a prior, to limit the number of bits used from the observations.
Taken together, all the components are trained to make the model self-consistent with the policy and representation.
The last part of our objective is 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 a ﬁxed
Algorithm
The ALM objective can be optimized with any RL algorithm.
We present an implementation based on DDPG (Lillicrap et al., 2015).
1: Initialize the encoder e φ z | s
) , model m φ ( z t +
| z t , a t ) , policy π φ ( a
| z t ) , classiﬁer
C θ ( z t +1 , a t , z t ) , reward r θ ( z t , a t ) , Q-function Q θ ( z t , a t ) , replay buffer B 2: for n = 1 , · · · , N do do 3:
Select action a n = π φ ( a n | e φ ( s n ))
+ N using the current policy and exploration noise N .
Execute action a n and observe reward r n and next state s n +1 .
Store transition ( s n , a n , r n , s n +1 ) in B ; sample length- K sequence ( s i , a i , s i +1 } t + K − 1 i =
t ∼ B 6:
Compute lower bound using off-policy actions: L Ke φ ,m φ (( s i , a i , s i +1 } t + K − 1 i = t ) (cid:46)7:
Update encoder and model by gradient ascent on off-policy lower bound: L Ke φ ,m φ 8:
Compute lower bound using on-policy model-based trajectories:
L Kπ φ (( s i = t ))
(cid:46)9:
Update policy by gradient ascent on on-policy lower bound: L Kπ φ 10:
Update classiﬁer, Q-function and reward by gradient descent on: L C θ , L Q θ , L r θ (cid:46) 10, 11,number of rollouts rather than the entire horizon, like SVG style updates (Amos et al., 2020; 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 difﬁcult to learn.
A larger value of K decreases this bias by reducing the dependency on a learned Q-function.
While the lower bound in Theorem 3.1 is not tight, we can include a learnable discount factor such that it becomes tight (see Appendix A.5).
In our experiments 4, we ﬁnd that the objective in Theorem 3.1 is still a good estimate of true expected returns.
In Appendix A.4, we derive a closed loop form for the optimal latent-dynamics and show that they are biased towards high-return trajectories: they reweight the true probabilities of trajectories with their rewards.
We also derive a lower bound for the model-based ofﬂine RL setting, obtaining a similar objective with an additional behavior cloning term (Appendix A.7).
A P RACTICAL A LGORITHM
We now describe a practical method to jointly train the policy, model, and encoder using the lower bound ( L Kφ ).
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 are aligned.
For training the encoder and model (latent-space learning phase), q Kφ ( τ ) is unrolled using actions from a replay buffer, whereas for training the policy (planning phase), q Kφ ( τ ) is unrolled using actions imagined from the latest policy.
To estimate the lower bound using just representations, our method also learns to predict the reward and Q-function from the learned representations z t using real data only (see Appendix C for details).
Algorithm 1 provides pseudocode.
Maximizing the lower bound with respect to the encoder and latent-space model.
To train the encoder and the latent-space model, we estimate the lower bound L Kφ using K-length sequences of transitions { s i , a i , s i +1 } t + K − 1 i =
t sampled from the replay buffer:
L
[Equation]
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. 2) 1 .
We ﬁnd that using a target encoder e φ targ ( z t | s t ) to calculate the KL consistency term leads to stable learning.
Maximizing the lower bound with respect to the policy.
The latent-space model allows us to evaluate the lower bound for the current policy by generating on-policy trajectories.
Starting from a sampled state s t the latent-space model is recurrently unrolled using actions from the current policy to generate a K -length trajectory of representations and actions ( z t : t + K , a t : t + K ) .
Calculating the intrinsic term in the augmented reward (log e φ ( z t +1 | s t +1 ) − log m φ ( z t +1 | z t , a t )) for on-policy actions is challenging, as we do not have access to the next high dimensional state s t +1 .
Following prior work (Eysenbach et al., 2020; Eysenbach et al., 2021a), we learn a classiﬁer to differentiate between representations sampled from the encoder e φ ( z t +1 | s t +1 ) and the latent-space model m φ ( z t +1 | z t , a 1 ) to estimate this term (see Appendix C for details).
We train the latent policy by recurrently backpropagating stochastic gradients (Heess et al., 2015; Amos et al., 2020; Hafner et al., 2019) of the lower bound evaluated on this trajectory:
L
[Equation]
Estimating the augmented reward function.
During policy training, estimating the intrinsic reward, log e φ ( z t +1 | s t +1 ) − log m φ ( z t +1 | z t , a t ) , is challenging because we do not have access to the next high dimensional state ( s t +1 ) .
Following prior work (Eysenbach et al., 2020; Eysenbach et al., 2021a), we note that a learned classiﬁer between representations sampled from the encoder e φ ( z t +1 | s t +1 ) versus the latent-space model m φ ( z t +1 | z t , a 1 ) can also be used to estimate the difference between log-likelihoods under them, which is exactly equal to the augmented reward function:
[Equation]
Here, C θ ( z t +1 , a t , z t ) ∈ [0 , 1] is a learned classiﬁer’s prediction of the probability that z t +1 is sampled from the encoder conditioned on the next state after starting at ( z t , a t ) pair.
The classiﬁer is trained via the standard cross entropy loss:
[Equation]
where z t +1 ∼ e φ ( · | s t +1 ) is the real next representation and ˆ z t +1 ∼ m φ ( · | z t , a t ) is the imagined next representation, both from the same start ( z t , a t ) pair.
Differences between theory and experiments.
While our theory suggests a coefﬁcient c = 1 , we 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 c ?? .
We also omit the log of the true rewards in both Eq. 9 and 7.
We show that this change is equivalent to the ﬁrst one for all practical purposes (see Appendix D.1.)
E XPERIMENTS
Our experiments focus on whether jointly optimizing the model, representation, and policy yields beneﬁts relative to prior methods that use different objectives for different components.
We use
SAC-SVG (Amos et al., 2020) as the main baseline, as it structurally resembles our method but uses different objectives and architectures.
While design choices like ensembling (MBPO, REDQ) are orthogonal to our paper’s contribution, we nonetheless show that ALM achieves similar sample efﬁciency MBPO and REDQ without requiring ensembling; as a consequence, it achieves good performance in ∼ 6 × 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 ﬁve random seeds.
Where possible, we used hyperparameters from prior works; for example, our network dimensions were taken from Amos et al.
Additional implementation details and hyperparameters are in Appendix D and a summary of successful and failed experiments are in Appendix E.
Baselines.
We provide a quick conceptual comparison to baselines in Table 5.
In Sec. 5.1, we will compare with the most similar prior methods, SAC-SVG (Amos et al., 2020) and SVG (Heess
[Table]
[Figure]
et al., 2015).
Like
, these use to
SAC-SVG also maintains a hidden representation, using a GRU, to make recurrent predictions in the observation space.
While SAC-SVG learns the model and representations using a reconstruction objective, ALM trains these components with the same objective as the policy.
The dynamics model from SAC-SVG bears a resemblance to Dreamer-v2 (Hafner et al., 2020) 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 et al., 2018)).
In Sec.
5.2, we focus on prior methods that use ensembles.
MBPO (Janner et al., 2019) is a model-based method that uses an ensemble of dynamics models for both actor updates and critic updates.
REDQ (Chen et al., 2021) is a model-free method which achieves sample efﬁciency onpar with model-based methods through the use of ensembles of Q-functions.
We also compare to
TD3 (Fujimoto et al., 2018) and SAC (Haarnoja et al., 2018); while typically not sample efﬁcient, these methods can achieve good performance asymptotically.
5. I S THE ALM OBJECTIVE USEFUL ?
We start by comparing ALM with the baselines on the locomotion benchmark proposed by Wang et al.
In this benchmark, methods are evaluated based on the policy return after training for 2e5 environment steps.
The TruncatedAnt-v2 and TruncatedHumanoid-v2 tasks included in this benchmark are easier than the standard Ant-v2 and Humanoid-v2 task, which prior model-based methods struggle to solve (Janner et al., 2019;
Chua et al., 2018; Amos et al., 2020; Shen et al., 2020; Rajeswaran et al., 2020; Feinberg et al., 2018; Buckman et al., 2018).
The results, shown in
Table 1, show that ALM achieves better performance than the prior methods on 4 / 5 tasks.
The results for SAC-SVG are taken from Amos et al.
and the rest are from Wang et al.
Because
SAC-SVG is structurally similar to ALM, the better results from ALM highlight the importance of training the representations and the model using the same objective as the policy.
5. C AN ALM ACHIEVE GOOD PERFORMANCE WITHOUT ENSEMBLES ?
Prior methods such as MBPO and REDQ use ensembles to achieve SOTA sample efﬁciency 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 efﬁciency without the use of ensembles.
Our next experiment studies whether ALM can achieve the beneﬁts of ensembles without the computational costs.
As shown in Figure 3, ALM matches the sample complexity of
REDQ and MBPO, but requires ∼ 6 × less wall-clock time to train.
Note that MBPO fails to solve the highest-dimensional tasks, Humanoid-v2 and Ant-v2 ( R 376 and R 111 ).
We optimized the ensemble training in the ofﬁcial REDQ code to be parallelized, leading to an increase in training speeds by 3 × , which is still 2 × slower than our method (which does not employ parallelization optimization).
5. W HY DOES ALM WORK ?
To better understand why ALM achieves high sample complexity without ensembles, we analyzed the Q-values, ran ablation experiments, and visualized the learned representations.
Analyzing the Q-values.
One way of interpreting ALM is that it uses a model and an augmented reward function to obtain better estimates 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 ensemble and SAC-AVG (baseline used in REDQ paper (Chen et al., 2021)) uses an average value of the Q-function ensemble to obtain a low variance estimate of these Q-values.
While ensembles can be an effective way to improve the estimates of neural networks (Garipov et al., 2018; Abdar et al., 2021), we hypothesize that our latent-space model might be a more effective approach in the RL setting because it incorporates the dynamics, while also coming at a 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 of Chen et al.
; Fujimoto et al.
See Appendix D.2 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 S.D. of the bias is lower for ALM than for REDQ and SAC-AVG, suggesting that the actions maximizing the lower bound are similar to actions that maximize true returns.
Ablation experiments.
In our ﬁrst ablation experiment, we compare ALM to ablations that separately remove the KL term and the value term from the encoder objective (Eq. 7), and remove the classiﬁer term from the policy objective (Eq. 9).
As shown in Fig.
5a, the KL term, which is a purely self supervised objective Grill et al. (2020), is crucial for achieving good performance. The classiﬁer 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.
A second ablation experiment (Fig. 5b) shows the performance of ALM for different number 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 beneﬁcial for model-free RL.
Prior work (Gupta et al., 2017; Eysenbach et al., 2021b;
Zhang et al., 2020) has shown that representations learning can facilitate properties like exploration, generalization and transfer.
In ﬁg. 5c, we ﬁnd that the end to end representation learning of ALM(3) achieves high returns faster than standard model-free RL.
[Figure]
C ONCLUSION
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 in a method where the representations, model, and policy all “cooperate” to maximize a lower bound on the expected returns.
Our experiments demonstrate the beneﬁts of such joint optimization: it achieves better performance than baselines that use separate objectives, and it achieves the beneﬁts of ensembles without their computational costs.
At a high level, our end-to-end method is reminiscent of the success deep supervised learning.
Supervised learning methods are dependent on their input features, and end-to-end approaches have allowed researchers to avoid manual feature design.
Similarly, the success of model-based methods depends on what model is used and what state representations it predicts.
Our approach shows how those components can likewise be learned in an end-to-end fashion.
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 is an on-policy version of ALM, where
Equation 9 can be calculated and simultaneously optimized with the encoder, model and policy, without using a classiﬁer is possible.
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.
Nonetheless, we believe that our proposed objective and method are not only practically useful, but may provide a template for designing even better model-based methods with learned representations
A P ROOFS A. H ELPER L EMMAS
Lemma A.1.
Let P K ( H ) be a truncated geometric distribution.
[Equation]
Given discount factor γ ∈ (0 , 1) and a random variable x t , we have the following identity
E P
[Equation]
Lemma A.2.
Let P K ( H ) be a truncated geometric distribution and p φ ( τ | H ) be the distribution over H + 1 length trajectories:
[Equation]
Then the RL objective can be re-written in the following way.
[Equation]
The Q function is deﬁned as Q ( s H , a H ) =
E p φ ( τ ) [(1 − γ ) (cid:80) ∞ t =0 γ t r ( s t + H , a t + 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 n .
The following optimization problem can be solved using the methods of Lagrange multipliers to obtain an optimal solution analytically.
[Equation]
such that
[Equation]
Starting out by writing the Lagrangian for this problem, then differentiating it and then equating it to
[Equation]
[Equation]
We ﬁnd the value of λ by substituting the value of p ( x ) from ( b ) in the equality constraint.
[Equation]
Substituting ( c ) to remove the dual variable from ( b ) , we obtain
[Equation]
A. 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.
The following bound holds for any representation e φ ( z t | s t ) , latent-space model m φ ( z t +1 | z t , a t ) , policy π φ ( a t | z t ) and K ∈ N
[Equation]
Note that scaling the rewards by a constant factor (1 − γ ) does not change the RL problem and ﬁnding a policy to maximize the log of the expected returns is the same as ﬁnding a policy to maximize expected returns, because log is monotonic.
We want to estimate the RL objective with trajectories sampled from a different distribution q φ ( τ ) which leads to an algorithm that avoids sampling high-dimensional observations.
[Equation]
When used as trajectory generative models, p φ ( τ ) predicts representations z t using current state s t .
Whereas q φ ( τ ) predicts z t by unrolling the learned latent-model recurrently on z t − 1 (and a t − 1 ).
Similar to variational inference, p φ ( τ ) can be interpreted as the posterior and q φ ( τ ) as the recurrent prior.
Since longer recurrent predictions of a learned model signiﬁcantly diverge from the true ones, we parameterize q to support an arbitrary length of model rollouts ( K ) during planning:
[Equation]
Proof.
We derive a lower bound on the RL objective for K-step latent rollouts.
[Table]
[Equation]
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 q φ ( τ | H = H ) q φ ( τ | H = H ) .
For ( e ) , we apply Jensen’s inequality.
For ( f ) , since all the terms inside the summation only depends on the ﬁrst H steps of the trajectory, we change the integration from over H length to over inﬁnite length trajectories.
For ( g ) , we change the order of integration.
For ( h ) , we use Lemma A.1 and the fact that E P K ( H )
[log( 1 { H ≤ K − 1 } r ( s H , a H ) + 1 { H = K } Q ( s H , a H ))]
= (cid:80) K − 1 t =
(1 − γ ) γ t log r ( s t , a t ) + γ K log Q
( z K , a K ) .
A. A LOWER BOUND FOR LATENT - SPACE LAMBDA WEIGHTED SVG(K)
Using on-policy the data, the policy returns can be estimated using the k-step RL estimator (RL k ) which uses the sum of rewards for the ﬁrst k steps and truncates it with the Q-function.
[Equation]
Like previous works (Schulman et al., 2015; Hafner et al., 2019), we write the RL objective as λ -weighted average of 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.
[Equation]
Our K-step lower bound L Kφ involves rolling out the model on-policy for K time-steps.
A larger value of K reduces the dependency on a learned Q-function and hence reduces the estimation bias at the cost of higher variance.
Lemma A.4.
We show that a λ -weighted average of our lower bounds for different horizons (different values of K) is a lower bound on the RL objective.
[Equation]
In ( a ) we use λ -weighted average estimation of the RL objective.
In ( b ) we use linearity of expectation.
In ( c ) , we note that for every k , the coefﬁcient of RL k is actually the probability of k − 1 under the geometric distribution.
Hence, we rewrite it as an expectation over the geometric distribution.
In ( d ) , we use Jensen’s inequality.
In ( e ) we write out the probability values of the geometric distribution.
In ( f ) , we use Theorem 3.1 for every value of k .
A. D ERIVING THE OPTIMAL LATENT DYNAMICS AND THE OPTIMAL DISCOUNT DISTRIBUTION
We deﬁne γ φ,K ( H ) to be a learned discount distribution over the ﬁrst K + 1 timesteps { 0 , 1 , . . .
, K } .
We use this learned discount factor when using data from imaginary rollouts.
We start by lower bounding the RL objective and derive the optimal discount distribution γ ∗ φ,K ( H ) and the optimal latent dynamics distribution q ∗ ( τ | H ) that maximizes this lower bound.
[Equation]
Given a horizon H ∈ {
[Equation] [Equation]
K }
[Equation]
the optimal dynamics q ∗ ( τ | H ) can be calculated analytically:
[Equation]
We substitute this value of q ∗ in equation ( c ) :
[Equation]
We now calculate the optimal discount distribution analytically γ ∗ φ,K ( H
[Equation] [Equation]
In ( a ) , we rewrite the RL objective using Lemma A.2.
In ( b ) , we multiply by γ
[Equation]
K
[Equation]
H
[Equation]
H )
[Equation]
In ( c ) , we use Jensen’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 : The optimal latent-dynamics are nonMarkovian and do not match MDP dynamics, but are optimistic towards high-return trajectories.
[Equation]
Writing out the optimal discount distribution :
[Equation]
A. T IGHTENING 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 a lower bound on the RL objective.
To derive ( f )
[Equation]
we have already substituted the optimal latent-dynamics.
We now substitute the value of γ ∗ φ,K ( H ) and verify that the lower bound ( f ) leads to the original objective, suggesting that the bound is tight.
[Equation]
[Table]
( h ) we substitute the value of γ ∗ φ,K ( 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.
A. T IGHTNESS OF LOWER BOUND WITH LENGTH OF ROLLOUTS K
Proof.
Proving L K ≥ L K +1 for all K ∈ N will do.
[Equation]
A. A LOWER BOUND FOR LATENT - SPACE SVG(K) IN OFFLINE RL
In the ofﬂine RL setting (Levine et al., 2020; Prudencio et al., 2022), we have access to a 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 ( φ t → φ t +1 ) , we were able to collect new data using it, while in ofﬂine
RL we have to re-use the same static data.
Similar to Theorem 3.1, we deﬁne the distribution over trajectories p
( τ ) :
[Equation]
such that the ofﬂine dataset consists of trajectories sampled from this true distribution.
We do not assume access to the data collection policies.
Rather, π b ( a t | s t ) is the behavior cloning policy obtained from the ofﬂine dataset.
We want to estimate the RL objective with trajectories sampled from a different distribution q φ ( τ ) :
[Equation]
which leads to an algorithm that avoids sampling high-dimensional observations.
Similar to Theorem 3.1, we want to ﬁnd an encoder e φ ( z t | s t ) , a policy π φ ( a t | s t ) , and a model m φ ( z t +1 | z t , a t ) to maximize the RL objective:
[Equation]
[Table]
[Equation] [Table] [Equation] [Table]
[Equation] [Table] [Equation] [Table]
[Equation]
[Table]
[Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]

The only main difference is between this proof and the proof A.2 of Theorem3.1 is in step ( f ) , where canceling out the common terms of p ( τ ) and q ( τ ) 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., 2010) for the ofﬂine RL setting.
B C OMPARISON TO PRIOR WORK ON LOWER BOUND FOR RL (M N M)
Our approach of deriving an evidence lower bound on the RL objective is similar to prior work (Eysenbach et al., 2021a).
In this section we brieﬂy go over the connection between our method and Eysenbach et al.
(2021a).
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 → ∞ and assuming an identity function for the encoder, we exactly reach the lower bound presented in Eysenbach et al.
(2021a).
By using a 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 ﬁdelity.
This is supported by the fact that ALM 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-learning bias.
Hence, the policy faces a penalty (intrinsic reward term) only for the ﬁrst K steps rather than the entire horizon (Eysenbach et al., 2021a), which could lead to lower returns on the training tasks (Eysenbach et al., 2021b).
C A DDITIONAL LEARNING DETAILS
Estimating the Q-function and reward function.
Unlike most MBRL algorithms, the Q function
Q θ ( z t , a t ) is learned using transitions s t , a t , r t , s t +1 from the real environment only, using the
[Table]
standard TD loss:
[Equation]
The TD target is computed using a target Q-function (Fujimoto et al., 2018).
We learn the reward function r θ ( z t , a t ) using data s t , a t , r t from the real environment only.
For r θ ( z t , a t ) , by minimizing the mean squared error between true and predicted rewards :
[Equation]
[Equation] Unlike prior representation learning methods (Zhang et al., 2020), we 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, as described in the next paragraph.
D I MPLEMENTATION DETAILS
We implement ALM using DDPG (Lillicrap et al., 2015) as the base algorithm.
Following prior svg methods (Amos et al., 2020), we parameterize the encoder, model, policy, reward, classiﬁer and Q-function as 2-layer neural networks, all with 512 hidden units except the model which has 1024 hidden units.
The model and the encoder output a multivariate Gaussian distribution over the latent-space with diagonal covariance.
Like prior work (Hansen et al., 2022; Yarats et al., 2021), we apply layer normalization (Ba et al., 2016) to the value function and rewards.
Similar to prior work (Schulman et al., 2015; Hafner et al., 2019), we reduce variance of the policy objective in
Equation 9, by computing an exponentially-weighted average of the objective for rollouts of lengthto
K. This average is also a lower bound (Appendix A.3).
To train the policy, reward, classiﬁer 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., 2021).
All hyperparameters are listed in
Table 2.
The brief summary of all the neural networks, their loss functions and the inputs to their loss functions are listed in Table 3.
Analyzing the learned representations.
Because the ALM objective optimizes the encoder and model to be self-consistent (see Sec. 3.4), we expect the ALM dynamics model to 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, we autoregressively unroll the learned model, comparing each prediction z t to the true representation (obtained by applying the encoder to observation s t ).
Fig.
6 (left) visualizes the ﬁrst coordinate of the representation, and shows that the model learned via ALM remains accurate for ∼ 20 steps.
The ablation of ALM that removes the KL term diverges after just two steps.
[Table]
[Figure]
D. O MISSION OF THE LOG OF REWARDS AND Q FUNCTION IS EQUIVALENT TO SCALING KL IN E QUATION
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 coefﬁcient in
Eq. 9.
This result means that our theoretical results still hold approximately when the logarithmic transformation is omitted.
We applied ALM, with the log of rewards, to a transformed MDP with a shifted reward function.
A large enough constant, a, was added to all original rewards(which doesn’t change the optimal policy assuming that there are no terminal states)
[Equation]
Taking a log of this new reward is equal to the reward of the original objective, scaled by the constant,
[Equation]
The additive term can be ignored because it won’t contribute to the optimization.
We plot both y
= r and y = a (log( r + a ) − log( a )) , to show that they are very similar for commonly used values of rewards in Figure 7.
The scaling constant a can be interpreted as the weight of the log of rewards, relative to the KL term in the ALM objective.
Hence changing this value is approximately equivalent to scaling the KL coefﬁcient.
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 7 and Figure 12, we used the value of a = 10000 .
17.5 reward input 0.
r e w a r d o u t p u t
Original reward vs ALM reward on transformed MDP r a(log(r + a) - log(a))
D. B IAS AND V ARIANCE OF THE L OWER B OUND
We recreate the experimental setup of (Chen et al., 2021) to evaluate the average and the std value of the bias between the Monte-Carlo returns and the estimated returns (lower bound L Kφ ( s, a ) for our experiments).
Similar to (Chen et al., 2021), since the true returns change as training progresses we analyze the mean and std value of normalized bias, which is deﬁned as: ( L Kφ ( s, a ) −
Q π ( s, a )) / | E s (cid:48) ,a (cid:48) ∼ π ( s (cid:48) , a (cid:48) )
| .
This helps us to evaluate bias relative to the scale of average MonteCarlo returns of the current policy.
Results for Ant-v2 and Walker2d-v2 are presented in Fig. 8.
E F AILED E XPERIMENTS
Experiments that we tried and that did not help substantially:
• Value expansion:
Training the critic using data generated from the model rollouts.
The added complexity did not add much beneﬁt.
• Warm up steps: Training the policy using real data for a ﬁxed time-steps at the start of training.
• Horizon scheduling: Scheduling the sequence length from 1 to K at the start of training.
• Exponential discounting: Down-scaling the learning rate of future time-steps using
a temporal discount factor, to avoid exploding or vanishing gradients.
Experiments that were tried and found to help:
• Target encoder:
Using a target encoder for the KL term in Eq.
7 helped reduce variance in episode returns.
• Elu activation: Switching from Relu to Elu activations for all networks for ALM resulted in more stable and sample efﬁcient performance across all tasks.
[Figure]
for compar- ing different components of the joint objective on Humanoid-v2 and Walker-v2.
The KL consistency term is the most important contributing factor.
[Figure]
performs
[Figure]
[Figure] [Figure]
[Table]
[Figure] [Figure]
[Figure]
r
[Figure]
[Figure] [Figure] [Figure]
[Figure]
F C OMPARISON TO BASELINES
Here we provide a quick comparison to baselines, conceptually in terms of number of gradient updates per env step, usage of ensembles are used for dynamics model or Q-function
[Table]
and objectives for representations, model, and policy.
[Table]