M ISMATCHED N O M ORE : J OINT M ODEL -P OLICY O PTIMIZATION FOR M ODEL -B ASED RL
A BSTRACT
Many model-based reinforcement learning (RL) methods follow a similar template: ﬁt a model to previously observed data, and then use data from that model for RL or planning.
However, models that achieve better training performance (e.g., lower MSE) are not necessarily better for control: an RL agent may seek out the small fraction of states where an accurate model makes mistakes, or it might act in ways that do not expose the errors of an inaccurate model.
As noted in prior work, there is an objective mismatch: models are useful if they yield good policies, but they are trained to maximize their accuracy, rather than the performance of the policies that result from them.
In this work we propose a single objective for jointly training the model and the policy, such that updates to either component increases a lower bound on expected return.
This joint optimization mends the objective mismatch in prior work. Our objective is a global lower bound on expected return, and this bound becomes tight under certain assumptions.
The resulting algorithm (MnM) is conceptually similar to a GAN: a classiﬁer distinguishes between real and fake transitions, the model is updated to produce transitions that look realistic, and the policy is updated to avoid states where the model predictions are unrealistic.
I NTRODUCTION
Much of the appeal of model-based RL is that model learning is a simple and scalable supervised learning problem.
Unfortunately, the accuracy of the learned model does not directly correlate with whether the model-based RL algorithm will receive high reward [16, 25].
For example, a model might make small mistakes in critical states that cause a policy to take suboptimal actions.
Alternatively, a model with large errors may yield a policy that attains high return if the model errors occur in states that the policy never visits.
The underlying problem is that dynamics models are trained differently from how they are used.
Typical model-based methods train a model using data sampled from the real dynamics (e.g., using maximum likelihood), but apply these models by using data sampled from the learned dynamics [11, 20, 21, 40].
Prior work has identiﬁed this objective mismatch issue [16, 25, 26]: the model is trained using one objective, but the policy is trained using a different objective.
Designing an objective for model training that is guaranteed to improve the expected reward remains an open problem.
This paper aims to answer the following question: How should we train a dynamics model so that it produces high-return policies when used for model-based RL?
In this paper, we propose a model-based RL algorithm where the model and policy are jointly optimized with respect to the same objective.
This objective is a lower bound on the expected return under the true environment dynamics, and the bound becomes tight under certain assumptions.
Structurally, our algorithm resembles a generative adversarial network (a GAN), in that the model is trained using a discriminator that distinguishes between real and fake transitions.
Crucially, this same discriminator is also included in the objective for the policy, and both the model and policy are jointly trained to maximize reward and minimize discriminator accuracy.
Thus, the model and policy cooperate to produce realistic and high-reward trajectories against an adversary that attempts to classify whether the trajectory is realistic.
Our method stands in contrast to standard model-based ∗ contribution. a r X i v : . 02758v [ c s .
L G ]
O c t RL methods, where it is more common to pit the model against the policy [5, 28, 32]. An interesting consequence of maximizing the lower bound is that the dynamics model does not learn the true dynamics, but rather learns “optimistic” dynamics that facilitate exploration.
The main contribution of this work is an algorithm, Mismatched no More (MnM) for model-based RL that provably maximizes a lower bound on expected reward.
Importantly, this bound becomes tight at optimality under certain assumptions.
Our algorithm has the unique property of jointly optimizing the policy and model using the same objective.
Across a range of tasks, we demonstrate that our method is competitive with prior state-of-the-art methods on benchmark tasks; on certain hard exploration tasks, our method outperforms prior methods based on maximum likelihood estimation.
R ELATED W ORK
Model-based RL methods typically ﬁt a dynamics model to observed transitions, and then apply an
RL method to that learned model.
Most of these methods use maximum likelihood to ﬁt the dynamics model, and then use RL to maximize the expected return under samples from that model [10, 11, 20, 21, 40].
The observation that this maximum likelihood objective is not aligned with the RL objective has been noted in prior work [16, 25, 26, 37, 44] where it is referred to as the objective mismatch problem.
One strategy for mitigating this problem is to modify the model training to improve model accuracy under multi-step rollouts [3, 4, 16, 22, 37, 39].
A second strategy is penalize the policy for taking transitions where the model is inaccurate [23, 26, 35, 42, 43]. Similar to all these prior methods, our approach will also use a modiﬁed reward function to train the policy, but it will also modify how the model is trained such that the model and policy optimize the same objective.
Like Luo et al.
[26], we show that our objective is a lower bound on expected return, but our bound will hold for all dynamics models and policies (i.e., globally), rather than a narrow region around previous models and policies (i.e., locally).
A third strategy is to directly optimize the model such that it produces good policies [2, 29, 36].
While our aim is the same as these prior methods, our approach will not require differentiating through unrolled model updates or optimization procedures.
Our theoretical derivation builds on prior work that casts model-based RL as a two-player game tween a model-player and a policy-player [5, 28, 30, 31].
However, whereas prior work pits model and policy compete against one another, our formulation will result in a cooperative game, wherein the model and policy players cooperate in optimizing the same objective (a lower bound on the expected return).
Our approach, though structurally resembling a GAN, is different from prior work that simply replaces a maximum likelihood model with a GAN model [6, 8, 24].
The most similar prior work is VMBPO [9].
The mechanics of our method are similar, also learning a dynamics model using a classiﬁer that distinguishes real versus generated rollouts.
However, while our method maximizes a lower bound on expected return, VMBPO maximizes a different, risk-seeking objective, which is an upper bound on expected return.
This different objective can be expressed as the expected return plus the variance of the return, so VMBPO has the undesirable property of preferring policies that receive slight lower return if the variance of the return is much larger (see Appendix A.1).
Indeed, most of the components of our method, including classiﬁers and
GAN-like models, have been used in prior work, main contribution of our paper is a precise recipe for combining these components in a way that provably maximizes expected return.
A U NIFIED O BJECTIVE FOR M ODEL -B ASED RL
Notation.
We focus on the Markov decision process with states s t , actions a t , initial state distribution p 0 ( s 0 ) , positive reward function r ( s t , a t ) > 0 , and dynamics p ( s t +1 | s t , a t ) .
Our aim is to learn a control policy π θ ( a t | s t ) with parameters θ that maximizes the expected discounted return:
[Equation]
We use transitions ( s t , a t , r t , s t +1 ) collected from the (real) environment to train the dynamics model q θ ( s t +1 | s t , a t ) , and use transitions sampled from this learned model to train the policy.
To simplify notation, we will deﬁne a trajectory τ to be the sequence of states and actions visited in an episode: τ (cid:44) ( s 0 , a 0 , s 1 , a 1 , · · · ) .
We then deﬁne R ( τ ) (cid:44) (cid:80) ∞
t
=0 γ t r ( s t , a t ) as the discounted return of a trajectory.
Finally, we deﬁne two distributions over trajectories.
First, p π ( τ ) is the distribution over trajectories when policy π θ interacts with dynamics p ( s (cid:48) | s, a ) ; q π ( τ ) is the distribution over trajectories when policy π θ interacts with the learned dynamics q θ ( s t +1 | s t , a t ) : p π ( τ ) =
p 0 ( s 0 ) ∞ (cid:89) t = p
( s t +1 | s t , a t )
π θ ( a t | s t ) , q π ( τ ) =
p 0 ( s 0 ) ∞ (cid:89) t =
q θ ( s t +1 | s t , a t )
π θ ( a t | s t ) .
Desiderata.
Our aim is to design an objective L ( θ ) that can be used to jointly optimize both the policy π θ ( a t | s t ) and the dynamics model, q ( s t +1 | s t , a t ) .
Such joint optimization would avoid the objective mismatch problem from prior work.
This objective should be a lower bound on the expected return. Moreover, the lower bound should become tight at optimality.
An objective for model-based RL.
We now introduce an objective that achieves these aims.
Our objective will be the policy’s reward when interacting with the learned model, but using a different reward function.
The new reward function augments the task reward with an additional term that measures the difference between the learned model and the real environment.
We deﬁne our objective
[Equation]
where the modiﬁed reward function is deﬁned as
[Equation]
This objective should be understood as maximizing an augmented reward under the learned dynamics.
The augmented reward function ˜ r penalizes the policy for taking transitions that are unlikely under the true dynamics model, similar to prior work [14, 43].
Later, we will show that we can estimate this augmented reward without knowing the true environment dynamics by using a GAN-like classiﬁer.
We will optimize this lower bound with respect to both the policy π θ ( a t | s t ) and the dynamics model q θ ( s t +1 | s t , a t ) .
For the policy, this optimization entails performing RL to maximize the modiﬁed reward using samples from the learned model; the only difference from prior work is the modiﬁcation to the reward function.
Training the dynamics model using this objective is very different from standard maximum likelihood training, and instead resembles a GAN.
The model is optimized to sample trajectories that both have high reward (i.e., log r is large) and are similar to real dynamics (i.e., log pq is large).
This objective differs from VMBPO [9] by taking the log( · ) of the original reward functions; our experiments demonstrate that excluding this component invalidates our lower bound and results in learning suboptimal policies.
Our objective has two properties that make it particularly useful.
First, the model and the policy are trained using exactly the same objective: updating the model not only increases the objective for the model, but also increases the objective for the policy.
Note that this is very different from prior work, where training the model to be more accurate (increase likelihood) can decrease the policy’s expected return under that model.
Second, our objective is a lower bound on the expected return. This property gives us a guarantee on how well the learned policy will perform when deployed on the real environment. To state this result formally, we will take the logarithm of the expected return. Of course, maximizing the log( · ) of the expected return is equivalent to maximizing the expected return.
Theorem 3.1.
The following bound holds for any dynamics q ( s t +1 | s t , a t ) and policy π ( a t | s t ) :
[Equation]
The proof is presented in Appendix A.3.
Note that the expected return under the learned model, which most prior model-based RL methods use to train the policy, is not a lower bound on the expected return.
Unlike the objective proposed in [26], which is also a lower bound, our bound bound holds for all dynamics functions and policies, not just those in a small region around earlier dynamics functions and policies.
To the best of our knowledge, this is the ﬁrst global lower bound for model-based RL.
How tight is this lower bound?
Showing that a bound is tight is crucial for making the claim that better dynamics models will produce better policies.
Our experiments show that optimizing this lower bound yields policies that achieve high reward across a wide range of tasks. [Figure] [Figure]
Tightening the lower bound.
We now introduce a modiﬁcation to our lower bound that does make the bound tight.
This new lower bound will be more complex than the one introduced above and we have not yet successfully designed an algorithm for maximizing it.
Nonetheless, we believe that presenting the bound may prove useful for the design of future model-based RL algorithms.
We will use L γ ( θ ) to denote this new lower bound.
In addition to the policy and dynamics, this bound will also depend on a time-varying discount, γ θ ( t ) , in place of the typical γ t term.
Similar learned discount factors have been studied in previous work on model-free RL [33].
We deﬁne this objective as follows:
[Equation]
where the augmented reward is now deﬁned as
[Equation]
and Γ θ ( t ) =
(cid:80) tt (cid:48)
=
γ θ ( t (cid:48) ) is the CDF of the learned discount function.
This new lower bound, which differs from our main lower bound by the learnable discount factor, does provide a tight bound on the expected return objective:
Lemma 3.2.
Let an arbitrary policy π ( a t | s t ) be given.
The objective L γ ( θ ) is also a lower bound on the expected return objective, log E π [ (cid:80) ∞ t =
γ t r ( s t , a t )] ≥ L γ ( θ ) , and this bound becomes tight at optimality:
[Equation]
The proof is presented in Appendix 3.2.
One important limitation of this result is that the learned dynamics that maximize this lower bound to make the bound tight may be non-Markovian.
Intriguingly, this analysis suggests that using non-Markovian models, such as RNNs and transformers, may accelerate learning on Markovian tasks.
We leave an empirical investigation of the more complex lower bound L γ and non-Markovian models to future work.
The optimal dynamics are optimistic.
We now return to analyzing the simpler lower bound ( L ( θ ) in Eq. 2). In stochastic environments, the dynamics model that optimizes this lower bound is not equal to the true environment dynamics.
Rather, it is biased towards sampling trajectories with high return.
Ignoring parametrization constraints, the dynamics model that optimizes our lower bound is q ∗ ( τ ) =
[Equation]
(proof in Appendix A.4.).
We hypothesize that the optimism in the dynamics model will accelerate policy optimization, a hypothesis we test in Sec. 5.1.
Would the optimistic dynamics overestimate the policy’s return, violating Lemma 3.1? This is not quite what our method does. Rather, our method estimates the augmented reward using the optimistic dynamics model, and the reward augmentation compensates for the optimism in the dynamics model.
M ISMATCHED N O M ORE
The previous section presented a single (global) lower bound ( L from Eq. 2) for jointly optimizing the policy and the dynamics model.
In this section, we develop a practical algorithm for optimizing this lower bound.
We call our method M ISMATCHED NO M ORE (MnM) because the policy and model optimize the same objective, thereby resolving the objective mismatch problem noted in prior work.
The main challenge in optimizing this bound is that the augmented reward function depends on the transition probabilities of the real environment, p ( s t +1 | s t , a t ) , which are unknown.
We address this challenge by learning a classiﬁer (Sec. 4.1), and then describe the precise update rules for the policy, dynamics model, and classiﬁer (Sec. 4.2).
4. E STIMATING THE A UGMENTED R EWARD F UNCTION
To estimate the augmented reward function, which depends on the transition probabilities of the real environment, we learn a classiﬁer that distinguishes real transitions from fake transitions.
This approach is similar to GANs [18] and similar to prior work in RL [14, 43].
We use C φ ( s t , a t , s t +1 ) ∈ [0 , 1] to denote the classiﬁer.
For an optimal classiﬁer, we can use the classiﬁer’s predictions to estimate the augmented reward function:
[Equation]
The approximation above reﬂects function approximation error in learning the classiﬁer.
We now present our complete method, which trains three components: a classiﬁer, a policy, and a dynamics model.
Our method alternates between (1) updating the policy (by performing RL using model experience with augmented rewards) and (2) updating the dynamics model and classiﬁer (using a GAN-like objective).
In describing the loss functions below, we use the superscripts ( · ) real and ( · ) model to denote transitions that have been sampled from the true environment dynamics or the learned dynamics function.
To reduce clutter, we omit the superscripts when unambiguous.
4. U PDATING THE M ODEL , P OLICY , AND C LASSIFIER
Updating the policy.
The policy is optimized to maximize the augmented reward on transitions sampled from the learned dynamics model.
While this optimization can be done using any RL algorithm, including on-policy methods, we will focus on an off-policy actor-critic method.
We deﬁne the Q function as sum of augmented rewards under the learned dynamics model:
[Equation]
We approximate the Q function using a neural network Q ψ ( s t , a t ) with parameters φ .
We train the Q function using the TD loss on transitions sampled from the learned dynamics model:
[Equation]
(7) where (cid:98)·(cid:99) sg is the stop-gradient operator and y t =
˜ r ( s t , a t , s model t +1 )
+ γ E π [Equation] a t +
[Equation]
s model t + [Equation] (cid:2) Q ψ ( s model t +1 , a t +1 ) (cid:3) .
The augmented reward, ˜ r , is estimated using the learned classiﬁer.
To estimate the corresponding value function, we use a 1-sample approximation: V ψ ( s t ) =
Q ψ ( s t , a t ∼ π θ ( a t | s t )) .
The policy is trained to maximize the expected future (augmented) return, as estimated by the Q function:
[Equation]
In our implementation, we regularize the policy by adding an additional entropy regularizer.
Following prior work [17], we maintain two Q functions and two target Q functions, use the minimum of the two target Q functions to compute the TD target.
See Appendix C for details.
Updating the classiﬁer.
We train the classiﬁer to distinguish real versus model transitions using the standard cross entropy loss:
[Equation]
Note that the real transition ( s real t , a real t , s real t +1 ) and model transition ( s real t , a real t , s model t +1 ) have the same initial state and initial action.
[Table]
Updating the dynamics model.
To optimize the dynamics model, we rewrite the lower bound in terms of a single transition (derivation in Appendix A.6):
[Equation]
The approximation above reﬂects approximation error in learning the optimal classiﬁer.
This approximation is standard in prior work on GANs [18] and adversarial inference [12, 13].
The procedure for optimizing the dynamics model and the classiﬁer resembles a GAN [18]: the classiﬁer is optimized to distinguish real transitions from model transitions, and the model is updated to fool the classiﬁer (and increase rewards).
However, our method is not equivalent to simply replacing a maximum likelihood model with a GAN model .
Indeed, such an approach would not optimize a lower bound on expected return.
Rather, our model objective includes an additional value term and our policy objective includes an additional classiﬁer term.
These changes enable the model and policy to optimize the same objective, which is a lower bound on expected return.
Algorithm summary.
We summarize the method in Alg. 1 and provide an illustration in Fig.
Implementing MnM on top of a standard model-based RL algorithm is straightforward.
First, create an additional classiﬁer network.
Second, instead of using the maximum likelihood objective to train the model, use the GAN-like objective in Eq.
10 to update both the model and the classiﬁer.
Third, add the classiﬁer’s logits to the predicted rewards (Eq. 5).
Following prior work [21], we learn an additional model to predict the true environment.
E XPERIMENTS
We present two sets of experiments.
Our ﬁrst set of experiments studies the importance of different components of MnM .
Second, we study how MnM compares with prior model-based RL algorithms on challenging, robotic control tasks.
To compare policies learned by different algorithms, we will evaluate the policies using the true environment dynamics, not the learned dynamics model.
5. U NDERSTANDING THE L OWER B OUND AND THE L EARNED D YNAMICS
In this section, we begin by studying the seemingly contradictory attributes of our method: optimistic dynamics and pessimistic policies, and end by conﬁrming that together these components optimize an increasingly tight lower bound on the expected return.
Our theory suggests that MnM should work best in settings with stochastic dynamics and challenging exploration requirements, as the dynamics model should tilt the true dynamics to make the stochasticity more favorable for solving the task.
We use a 10x10 gridworld with highly stochastic dynamics and a sparse reward [-upon reaching a goal state (Fig. 
 (Right) ).-]function.
The results, shown in Fig.
2a (Left) , show that MnM outperforms both Q-learning and VMBPO.
In line with our theory, the dynamics learned by MnM (Fig.
2a (Right) ) alter the environment stochasticity to lead the agent towards the goal, increasing the probability of collecting high-reward experience.
Of course, we use the true environment dynamics, not the optimistic dynamics model, for evaluating the policies.
While VMBPO also learns optimistic dynamics, it omits log-transformation of the reward function (which encourages pessimistic behavior), a difference that has a large effect on this task.
Our augmented reward function contains two crucial components, (1) the classiﬁer term and (2) the logarithmic transformation of the reward function. We test the importance of the classiﬁer term in [Figure]
[Figure] correcting for inaccurate models. To do this, we limit the capacity of the MnM dynamics model so that it makes “low-resolution” predictions, forcing all states in 3 × 3 blocks to have the same dynamics.
We will use the gridworld shown in Fig. 2b, which contains obstacles that occur at a ﬁner resolution than the model can detect.
When the “low resolution” dynamics model makes predictions for states near the obstacle, it will average together some states with obstacles and some states without obstacles.
Thus, the model will (incorrectly) predict that the agent always has some probability of moving in each direction, even if that direction is actually blocked by an obstacle. However, the classiﬁer (whose capacity we have not limited) detects that the dynamics model is inaccurate in these states and so the augmented reward is much lower at these states.
Thus, MnM is able to solve this task despite the inaccurate model; an ablation of MnM that removes the classiﬁer term attempts to navigate through the wall and fails to reach the goal. We then test the importance of the logarithmic transformation by comparing MnM to VMBPO, which includes the classiﬁer term but omits the logarithmic transformation.
We hypothesize that VMBPO’s deviation from our lower bound will cause it to exhibit risk seeking behavior.
To test this hypothesis, we use the 3-state MDP in Fig.
3a (top) where numbers indicate the reward at each state. While moving to the right state yields slightly higher rewards, “wind” knocks the agent out of this state with probability 50% so the reward-maximizing strategy is to move to the left state.
While MnM MnM learns the reward-maximizing strategy, VMBPO learns a policy that goes to the right state and receives lower returns (with much higher variance). Finally, we study how the MnM objective compares to alternative objectives.
We use the gridworld from Fig.
2a and use a version of MnM based on value iteration to avoid approximation error. Plotting the MnM objective in Fig.
3b, we observe that it is always a lower bound on the (log) expected return, as predicted by our theory. Ablations of MnM to omit the reward augmentation or even just the log transformation (i.e., VMBPO) overestimate the expected return. 5. C OMPARISONS O N R OBOTICS T ASKS
Our second set of experiments use a range of continuouscontrol robotic tasks to answer two questions. We ﬁrst investigate whether MnM performs at least comparably with prior work. We then study tasks with sparse rewards and more challenging exploration, where we suspect the optimistic dynamics learned by MnM may be beneﬁcial.
We illustrate a subset of the environments in Fig. 4 and include implementation details in Appendix C.2.
One detail of note is that we omit the reward augmentation for MnM during these experiments, as it hinders exploration leading to lower returns.
We use MBPO [21] as a baseline for model-based RL because it achieves state-of-the-art results and is a prototypical example of model-based RL algorithms that use maximum likelihood models.
Our ﬁrst comparison uses three locomotion tasks from the
OpenAI Gym benchmark [7], which has become the standard benchmark for model-based RL algorithms. Thes tasks have dense rewards and pose no signiﬁcant exploration challenge, so we do not expect MnM to outperform prior methods.
The results (Fig. 5 (top left) ) show that MnM performs roughly on par with MBPO.
Tasks with sparse rewards, complex contact dynamics, and those requiring hard exploration often present a challenge for model-based RL algorithms, which are liable to exploit inaccuracies in the learned dynamics model.
Our next experiment evaluates MnM on control tasks used in prior work that demonstrate these properties.
These tasks the sparse-reward cartpole task from the DM Control benchmark [38], four manipulation tasks from the Metaworld benchmark [41], and four dexterous manipulation tasks from the ROBEL benchmark [1].
See Appendix C for experiment details.
The results shown in Fig.
5 indicate that the complex environment dynamics of these tasks can cause prior model-based algorithms (MBPO) to perform worse than model-free algorithms (SAC), both in terms of asymptotic performance and sample complexity.
Nonetheless, we observe that MnM frequently outperforms all prior methods and, more importantly, it consistently does well across all tasks.
We include ablation experiments in Appendix B.1.
[Figure]
To investigate the beneﬁts of MnM over prior model-based methods, we logged the Q values throughout training and visualized them for the metaworld-drawer-open-v2 task in
Fig.
For fair comparison, we use Q values corresponding to just the task reward, omitting the augmented reward used by MnM. Fig.
6 shows that MnM yields Q values that are more accurate and more stable than MBPO.
This ﬁgure suggests that MBPO may be exploiting inaccuracies in the learned model.
Finally, we visualize the dynamics learned by
MnM on two robotic control tasks.
These tasks have deterministic dynamics, so our theory would predict that an idealized version of MnM would learn a dynamics model exactly equal to the deterministic dynamics. However, our implementation relies on function approximation (neural networks) to learn the dynamics, and the limited capacity of function approximators makes otherwisedeterministic dynamics appear stochastic.
On the Pusher-v2 task, the MnM dynamics cause the puck to move towards the robot arm even before the arm has come in contact with the puck.
While this movement is not physically realistic, it may make the exploration problem easier.
On the HalfChetah-v2 task, the MnM dynamics increase the probability that the agent remains upright after tripping, likely making it easier for the agent to learn how to run.
We expect that the implicit stochasticity caused by function approximation to be especially important for real-world tasks, where the complexity of the real dynamics often dwarfs the capacity of the learned dynamics model.
C ONCLUSION
The main contribution of this paper is an approach to model-based RL where the policy and dynamics model are jointly optimized using the same objective.
Unlike prior work, our objective is a global lower bound on the standard expected return objective.
Our approach not only tells users how to train their dynamics model, but also guarantees to them that updating their model (using our objective) will result in a better policy.
We therefore believe that this joint optimization will ease and accelerate the design of future model-based RL algorithms.
We suspect that the tools presented in this paper may prove useful for solving tasks that require extensive exploration or long-horizon planning.
For example, models learned by MnM may be more capable of sampling realistic long trajectories, opening a door to future model-based algorithms that apply on-policy algorithms to sampled experience.
Limitations.
The main limitations of our method are complexity and stability.
Compared with methods based on maximum likelihood, our method requires learning an additional classiﬁer and balancing the capacity of that classiﬁer against the capacity of the dynamics model.
In Appendix C, we describe implementation details we found useful for stabilizing training of our method.
Acknowledgements.
We thank Laura Smith, Marvin Zhang, and Dibya Ghosh for helpful discussions, thank Anusha Nagabandi and Michael Janner for help with experiments, and thank Danijar Hafner for reviewing a draft of the paper.
This material is supported by the Fannie and John Hertz Foundation and the NSF GRFP (DGE1745016).
A P ROOFS AND A DDITIONAL A NALYSIS A. VMBPO M AXIMIZES AN U PPER B OUND ON R ETURN
While MnM aims to maximize the (log) of the expected return, VMBPO aims to maximize the expected exponentiated return:
[Equation]
where η > 0 is a temperature term used by VMBPO.
Note that maximizing the log of the expected return, as done by MnM, is equivalent to maximizing the expected return, as the function log( · ) is monotone increasing.
However, maximizing the log of the expected exponentiated return, as done by
VMBPO, is not equivalent to maximizing the expected return.
Rather, it corresponds to maximizing a sum of the expected return an the variance of the return [27, Page 272]:
[Equation]
Thus, in environments with stochastic dynamics or rewards (e.g., the didactic example in Fig. 3),
VMBPO will prefer to receive lower returns if the variance of the returns is much higher.
We note that the expected exponentiated return is an upper bound on the expected return:
[Equation]
This statement is a direct application of Jensen’s inequality.
The bound holds with a strict inequality in almost all MDPs.
The one exception is trivial MDPs where all trajectories have exactly the same return.
Of course, even a random policy is optimal for these trivial MDPs.
A. H ELPER L EMMAS
We start by introducing a simple identity that will help handle discount factors in our analysis.
Lemma A.1.
Deﬁne p
( H ) =
G EOM (1 − γ ) as the geometric distribution.
Let discount factor γ ∈ (0 , 1) and random variable x t be given.
Then the following identity holds:
[Equation]
The proof involves substituting the deﬁnition of the Geometric distribution and then rearranging terms.
Proof.
[Equation]
The second helper lemma describes how the discounted expected return objective can be written as the expected terminal reward of a mixture of ﬁnite-length episodes.
Lemma A.2.
Deﬁne p
( H ) =
G EOM (1 − γ ) as the geometric distribution, and p ( τ | H ) as a distribution over length- H episodes.
We can then write the expected discounted return objective as follows:
[Equation]
Proof.
The ﬁrst identity follows from the deﬁnition of the geometric distribution.
The second identity writes the expectations as integrals
[Equation]
which will make future analysis clearer.
A. P ROOF OF L EMMA 3.
Proof.
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Table]
For (a) , we applied Lemma A.2.
For (b) , we applied Jensen’s inequality.
For (c) , we applied Jensen’s inequality again.
For (d) , we substituted the deﬁnitions of p θ ( τ | H ) and q θ ( τ | H ) .
For (e) , we noted that the term inside the summation only depends on the ﬁrst H steps of the trajectory, so collecting longer trajectories will not change the result.
This allows us to rewrite the integral as an expectation using a single inﬁnite-length trajectory.
For (f) , we recalled the deﬁnition q θ ( τ ) =
q θ ( τ = H = ∞ ) and swap the order of integration.
For (g) , we express the inner integral over p ( H ) as an expectation.
For (h) , we applied the identity from Lemma A.1.
For (i) , we moved the constant log(1 − γ ) back inside the integral and rewrote the integral as an expectation.
We have thus obtained the desired result.
A. P ROOF OF L EMMA 3.
Before presenting the proof of Lemma 3.1 itself, we show how we derived the lower bound in this more general case.
While this step is not required for the proof, we include it because it sheds light on how similar lower bounds might be derived for other problems.
We deﬁne γ θ ( H ) to be a learned distribution over horizons H .
We then proceed, following many of the same steps as for the proof of Lemma 3.1.
[Equation]
[Figure]
[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]
[Figure]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
( H )( H log γ − log γ θ ( H ) {++ log r ( s H ,a H )) (cid:35) ( j ) = E qθ ( τ ) (cid:34) ∞ (cid:88) H =0 γ θ ( H ) (cid:18) 1 − Γ θ ( H − 1) γ θ ( H ) (log p ( s H +
 |+}
s H ,a H ) − log q θ ( {+s H +
+} | s H ,a H ))
+ H log γ − log γ θ ( H ) + log r ( s H ,a H ) (cid:19)(cid:35) . For (a) , we applied Lemma A.2 and multiplied the integrand by q θ ( τ | H = H ) γ θ ( H ) q θ ( τ | H = H ) γ θ ( H )
= 1 .
For (b) , we applied Jensen’s inequality.
For (c) , we factored p ( τ, H ) =
p ( τ, H ) p ( H ) and q θ ( τ, H ) =
q ( τ |
H ) γ θ ( H ) .
Note that under the joint distribution p ( τ, H ) , the horizon H ∼ p ( H )
=
G EOM (1 − γ ) is independent of the trajectory, τ .
For (d) , we rewrote the expectation as an expectation over a single inﬁnite-length trajectory and simpliﬁed the summand.
For (e) , we recall the deﬁnition q θ ( τ ) =
q θ ( τ = H = ∞ ) .
For (f) , we rewrote the integral as an expectation and wrote out the deﬁnition of the geometric distribution, p ( H ) .
For (g) , we regrouped the difference of dynamics terms.
For (h) , we noted used the fact that (cid:80) H − 1 t =0 γ θ ( t ) + γ ∞ t =
[-Ht =
-]H γ θ ( t ) = 1 .
For (i) , we substituted the deﬁnition of the CDF function. For (j) , we rearranged terms so that all were multiplied by the discount γ θ ( H ) .
Thus, we have obtained the desired result.
We now prove Lemma 3.2, showing that
Eq. 4 becomes tight at optimality.
Proof.
[Equation]
For (a) , we undo some of the simpliﬁcations above, going back to Eq.
For (b) , we factor q θ ( τ, H ) =
q θ ( τ ) γ θ ( H | τ ) and p ( τ, H ) =
p ( τ ) p ( H ) .
At this point, we can solve analytically for the optimal discount distribution, γ θ ( H | τ ) :
[Equation]
In the second equality, we substitute the deﬁnition of R ( τ ) .
We then substitute Eq.
15 into our expression for L γ ( θ )
and simplify the resulting expression.
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
log r ( s H ,a H ) + (cid:40)(cid:40)(cid:40)(cid:40) log(1 − γ ) + log R ( τ ) + (cid:40)(cid:40)(cid:40)(cid:40)(cid:40) log r ( s H ,a H ) dτ (cid:1) dH − (cid:40)(cid:40)(cid:40)(cid:40) log(1 − γ )
= (cid:90)(cid:90)
q θ ( τ ) γ θ ( H | τ )(log p ( τ ) − log q θ ( τ ) + log R ( τ ) dτ ) dH
= (cid:90) q θ ( τ )(log p ( τ ) − log q θ ( τ ) + log R ( τ )) dτ.
In the ﬁnal line we have removed the integral over H because none of the integrands depend on H .
At this point, we can solve analytically for the optimal trajectory distribution, q θ ( τ ) :
[Equation]
We then substitute Eq.
16 into our expression for L γ ( θ ) , and simplify the resulting expression:
[Equation]
We have thus shown that the lower bound L γ becomes tight when we use the optimal distribution over trajectories q θ ( τ ) and optimal learned discount γ θ ( H | τ ) .
A. A LOWER BOUND FOR GOAL - REACHING TASKS .
Many RL problems can be better formulated as goal-reaching problems, a formulation that does not require deﬁning a reward function.
We now introduce a variant of our method for goal-reaching tasks.
Using ρ π ( s t + ) to denote the discounted state occupancy measure of policy π , we deﬁne the goal-reaching objective as maximizing the probability density of reaching a desired goal s g :
[Equation]
We refer the reader to Eysenbach et al.
[15] for a more detailed discussion of this objective.
For simplicity, we assume that the goal is ﬁxed, noting that the multi-task setting can be handled by conditioning the policy on the commanded goal.
Similar to Lemma 3.1, we can construct a lower bound on the goal-conditioned RL problem:
Lemma A.3.
Let initial state distribution p 1 ( s 1 ) , real dynamics p ( s t +1 | s t , a t ) , reward function r ( s t , a t ) > 0 , discount factor γ ∈ (0 , 1) , and goal g be given.
Then the following bound holds for any dynamics q ( s t +1 | s t , a t ) and policy π ( a t | s t ) :
[Equation]
where ˜ r g ( s t , a t , s t +1 ) (cid:44) (1 − γ ) log p ( s t +
=
s g | s t , a t )+log p ( s t +1 | s t , a t ) − log q ( s t +1 | s t , a t ) .
The proof, presented below, is similar to the proof of Lemma 3.1.
The ﬁrst term in the reward function, the log probability of reaching the commanded goal one time step in the future, is similar to prior work [33].
The correction term log p − log q incentivizes the policy to avoid transitions where the model is inaccurate, and can be estimated using a separate classiﬁer.
One important aspect of this goal-reaching problem is that it is entirely data-driven, avoiding the need for any manually-designed reward functions. Proof. [Equation] [Table]
Similar to the more complex lower bound presented in Eq.
4, this lower bound on goal-reaching can be modiﬁed (by learning a discount factor) to become a tight lower bound.
The resulting objective would resemble a model-based version of the algorithm from Rudner et al. A. D ERIVATION OF M ODEL O BJECTIVE (E Q . 10) Our lower bound depends on entirely trajectories sampled from the learned dynamics. In this section, we show how the same objective can be expressed as an expectation of transitions.
This expression is easier to optimize, as it does not require backpropagating gradients through time.
We start by writing our lower bound, conditioned on a current state s t .
[Equation]
In (a) , we substituted the deﬁnition of the augmented return.
For the purpose of optimizing the dynamics model, we can ignore all terms that do not depend on s t +1 . Removing these terms, we arrive at our model training objective (Eq. 10) B A DDITIONAL E XPERIMENTS B. A BLATION E XPERIMENTS We next run an ablation experiment to study the importance of a few key design decisions.
We compare MnM with ablations that omit the reward augmentation and the model optimism term. The results shown in Fig.
8 indicate that most of the beneﬁt of MnM comes from using a GAN-like model. Because the dynamics of these tasks are nearly deterministic, it is not surprising that the optimistic dynamics and the reward augmentation have only a small effect. C I MPLEMENTATION D ETAILS C. A LGORITHMS
Value Iteration (Fig. 2a (right) , Fig. 2b, Fig. 3) For the tabular experiments that perform value iteration, we perform Polyak averaging of the policy and learned dynamics model with parameter
SAC
MBPO
MnM
HalfCheetah-vHopper-vWalker2d-vdm-cartpole-swingup_sparse   metaworld-window-open-v    metaworld-door-open-v    metaworld-reach-v    metaworld-drawer-open-v
DClawPoseRandom-v
DClawTurnRandom-v
DClawScrewFixed-v
DClawScrewRandom-v    τ = 0 . We found that the value iteration version of MnM diverged without this Polyak averaging step. Experiments were stopped when the MnM dynamics model (which depends on the value function) changed by less than 1e-6 across iterations, as measured using an L 0 norm.
Q-learning (Fig. [-
-]2)
For the experiments with Q-learning (both with and without the MnM components), we performed (cid:15) -greedy exploration with (cid:15) = 0 .
We used a learning rate of 1 e − 2 .
For this task alone, we compute the MnM dynamics analytically by combining the true environment dynamics with the learned value function, allowing for clearer theoretical analysis.
For fair comparison, all methods receive the same amount of data, perform the same number of updates, and are evaluated using the real environment dynamics.
SAC for continuous control tasks. We used the SAC implementation from TF-Agents [19] with the default hyperparameters.
MBPO for continuous control tasks.
We implement MBPO on top of the SAC implementation from TF-Agents [19].
Unless otherwise mentioned, we take the default parameters from this implementation.
We use an ensemble of 5 dynamics models, each with 4 hidden layers of size 256.
The dynamics model predicts the whitened difference between the next state and the current state.
That is, to obtain the prediction for the next state, the predictions are scaled by a per-coordinate variance, shifted by a per-coordinate mean, and then added to the current state.
These whitened predictions are clipped to have a minimum standard deviation of 1e-5; without this, we found that the MBPO model resulted in numerical instability.
The model is trained using the standard maximum likelihood objective, with all members of the ensemble being trained on the same data.
To sample data from this model we perform 1-step rollouts, starting at states visited in the true dynamics.
We perform one batch of rollouts in parallel using a batch size of 256.
To sample the corresponding action, with probability 50% we take the action that was executed in the true dynamics; with probability 50% we sample an action from the current policy.
We found that this modiﬁcation slightly improves the results of MBPO.
We use a batch size of 256.
We have two replay buffers: the model replay buffer has size 256e3 and the replay buffer of real experience has size 1e6.
At the start of training, we collect 1e4 transitions from the real environment, train the dynamics model on this experience for 1e5 batches, and only then start training the policy.
We use a learning rate of 3e-4 for all components.
To stabilize learning, we maintain a target dynamics model using an exponential moving average ( τ = 0 . 001 ), and use this target dynamics model to sample transitions for training.
We update the model, policy, and value functions at the same rate we sample experience from the learned model, which is more frequently than we collect experience from the real environment (see Table 1).
MnM for Continuous Control Tasks We implement MnM on top of the SAC implementation from TF-Agents [19].
Unless otherwise mentioned, we take the default parameters from this implementation.
Our model architecture is exactly the same as our MBPO implementation, and we follow the same training protocol.
Unlike MBPO, MnM also learns a classiﬁer for distinguishing real versus model transitions.
The classiﬁer architecture is a 2 layers neural network with 1024 hidden units in each layer.
We found that this large capacity was important for stable learning.
We add input noise with σ = 0 .
1 while training the classiﬁer.
We whiten the inputs to the classiﬁer by subtracting a coordinate-wise mean and dividing by a coordinate-wise standard deviation.
When training the classiﬁer, we take samples from both the dynamics model and the target dynamics model as negative examples, ﬁnding that this stabilizes learning somewhat.
Following the suggestion of prior work [34], we use one-sided label smoothing with value 0.1, only smoothing the negative predictions and not the positive predictions.
We found that gradient penalties and spectral normalization decreased performance.
We found that automatically tuning the classiﬁer input noise also decreased performance.
We found that mixup had little effect.
We found that the loss would often plateau around 1e4 batches, but would eventually start decreasing again after 2e4 - 2e5 batches.
Like the MBPO model, we ﬁrst collect 1e4 transitions of experience from the real environment using a random policy, then train the dynamics model and classiﬁer for 1e5 batches, and only then start updating the policy.
Because the Q values are poor at the start of training, we only add the value term to the model loss (resulting in the optimistic dynamics model) after 2e5 batches (1e5 batches of model training, then 1e5 batches of model+policy training).
To further improve stability, we compute the value term in the model loss by taking the minimum over two target value functions (like TD3 [17]).
We update the model, classiﬁer policy, and value functions at the same rate we sample experience from the learned model, which is more frequently than we collect experience from the real environment (see Table 1).
C. E NVIRONMENTS
Gridworld for Fig.
2a.
This task is a 10 × 10 gridworld with obstacles shown in Fig.
2a.
There are four discrete actions, corresponding to moving to the four adjacent cells.
With probability 50%, the agent’s action is ignored and a random action is taken instead.
The agent starts in the top-left cell.
The agent receives a reward of +0.001 at each time step and a reward of +10 when at the goal state.
Episodes have 200 steps and we use a discount γ = 0 .
Gridworld for Fig.
2b.
This task is a 15 × 15 gridworld with obstacles shown in Fig.
2b.
There are four discrete actions, corresponding to moving to the four adjacent cells.
The dynamics are deterministic.
The agent starts in the top-left cell.
The reward is +1.0 at every state except the goal state, where the agent receives a reward of +100.0.
We use γ = 0 .
9 and compute optimal policies analytically using value iteration.
We implementing the aliasing by averaging together the dynamics for each block of 3 × 3 states. Importantly, the averaging was done to the relative dynamics (e.g., action 1 corresponds to move right) not the absolute dynamics (e.g., action 1 corresponds to moving to state (3, 4)). To handle edge effects, we modiﬁed the averaged dynamics so that the agent could not exit the gridworld. We computed the classiﬁer analytically using the true and learned dynamics models.
However, the augmented rewards become inﬁnite because the learned model assigns non-zero probability to transitions that cannot occur under the true dynamics. This is not a failure of our theory, as the optimal policy would choose to never visit these states, but it presents a challenge for optimization. We therefore added label smoothing with parameter 0.7 to the classiﬁer. Gridworld for Fig. 3a.
This task is a 10 × 10 gridworld with obstacles shown in Fig.
There are four discrete actions, corresponding to moving to the four adjacent cells. With probability 90%, the agent’s action is ignored and a random action is taken instead.
The agent starts in the top-left cell.
The rewards depend on the Manhattan distance to the goal: transitions that lead away from the goal have a reward of +0.001, transitions that do not change the distance to the goal have a reward of +1.001, and transitions that decrease the distance to the goal have a reward of +2.001.
We use γ = 0 .
5 and compute values and returns analytically using value iteration.
Gridworld for Fig. 3b.
This task used the same dynamics as Fig. 2a.
The one change is that the agent receives a reward of +1.0 at each time step and a reward of +10.0 at the goal state.
We estimate the more complex lower bound (Eq. 4) by also learning the discount factor. However, since we currently do not have a method for learning non-Markovian dynamics to fully optimize this lower bound, we do not expect the lower bound to become tight.
HalfCheetah-v2, Hopper-v2, Walker2d-v2 .
These tasks are taken directly from the
OpenAI benchmark without modiﬁcation.
dm-cartpole-swingup_sparse .
This task is taken directly from the DM Control benchmark [38] without modiﬁcation.
metaworld-window-open-v2 .
This task is based on the window-open-v2 task from the Metaworld benchmark [41].
To increase the difﬁculty of this task, we set random_init=
True and use a sparse reward function.
The sparse reward function is +0 when the window is within 3 units of the (default) open position and -1 otherwise.
metaworld-door-open-v2 .
This task is based on the door-open-v2 task from the Meta- world benchmark [41].
To increase the difﬁculty of this task, we set random_init=
True and use a sparse reward function.
The sparse reward function is +1 when the window is within 3 units of the (default) open position and +0 otherwise.
metaworld-reach-v2 .
This task is based on the reach-v2 task from the Metaworld benchmark [41].
To increase the difﬁculty of this task, we set random_init=
True and use a sparse reward function.
The sparse reward function is +0 when the window is within 1 units of the goal and -1 otherwise.
metaworld-drawer-open-v2 .
This task is based on the drawer-open-v2 task from the Metaworld benchmark [41].
To increase the difﬁculty of this task, we remove the reward shaping term ( reward_for_caging ) and just optimize the reward for opening the drawer ( reward_for_opening ).
DClawPoseRandom-v0 ,
DClawTurnRandom-v0 ,
DClawScrewFixed-v0 ,
DClawScrewRandom-v0 .
These tasks are taken directly from the ROBEL benchmark [1] without modiﬁcation.