Mismatched No More: Joint Model-Policy Optimization for Model-Based RL
Authors: Benjamin Eysenbach
Alexander Khazatsky
Sergey Levine
Ruslan Salakhutdinov
Venue: ICLR 2022 Submitted (2022)
Keywords: reinforcement learning
model-based RL
joint optimization

Abstract: Many model-based reinforcement learning (RL) methods follow a similar template: fit 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 classifier 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.

Under review as a conference paper at ICLR 2022
MISMATCHED No MORE: JOINT MODEL-POLIC' OPTIMIZATION FOR MODEL-BASED RL
Anonymous authors Paper under double-blind review
ABSTRACT
Many model-based reinforcement learning (RL) methods follow similar template: fit 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 single objective for jointly training the model and the policy, such that updates to either component increase a lower bound 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 classifier 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.
INTRODUCTION
Much of the appeal of model-based RL is that model learning 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 (Farahmand al.) 2017; Lambert et al., 2020). For example, model might make small mistakes in critical states that cause 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 model using data sampled from the real dynamics (e.g.. using maximum likelihood), but apply these models by using data sampled from the learned dynam-
ics (Deisenroth & Rasmussen 2011 Williams al. 2017] Janner et al., 2019 Hafner et al. 2019).
Prior work has identified this objective mismatch issue (Farahmand et al. 2017 Luo et al. 2019
Lambert et al.; 2020): the model is trained using objective, but the policy is trained using differ- ent objective. Designing an objective for model training that is guaranteed to improve the expected reward remains an open problem. This paper aims 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, propose a model- RL algorithm where the model and policy are jointly optimized with respect to the same objective. Our objective is a lower bound on the expected return under the true environment dynamics; slightly more complicated version of this lower bound becomes tight under certain assumptions. Structurally, our algorithm resembles generative adversarial network GAN), in that the model is trained using discriminator that distinguishes between real and fake transitions. This same discriminator is 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. Our method stands contrast to standard model-based RL methods, where it is more common to pit the model against the policy (Bagnell et al.] 2001 Nilim & El Ghaoui 2003 Ross & Bagnell 2012). An 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.

---

Under review as a conference paper at ICLR 2022
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. To the best of our knowledge, this is the first model-based RL objective that is a global lower bound on expected return, and that involves optimizing the model and policy using the same objective. 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 competitive with prior state-of-the-an methods on benchmark tasks; on certain hard exploration tasks, our method outperforms prior methods based on maximum likelihood estimation.
2 RELATED WORK
Model-based RL methods typically fit a dynamics model to observed transitions, and then apply an RL method to that learned model. Most of these methods use maximum likelihood to fit the dynamics model, and then use RL to maximize the expected return under samples from that model (Deisenroth
& Rasmussen 2011 Williams et al. 2017] Chua et al., [0118] Hafner et al. 2019 Janner et al. 2019). The observation that this maximum likelihood objective is not aligned with the RL objective has been
noted in prior work (Ziebart 2010; Talviitie, 2014 Farahmand et al.) 2017 Luo et al. 2019 Lambert
ct af. 2020). This issue is referred to as the objective mismatch problem: the model and policy (or planner) optimized using different objectives. This problem arises in almost all model-based
RL approaches, including those that train the model to predict the value function (Oh et al. 2017
Schrittwieser et al.) (020)) or that perform planning (Chua et al. 2018 Schrittwieser et al. 2020).
One strategy for mitigating this problem is to modify the model training to improve model accuracy under multi-step rollouts (Joseph et al., 2013 "alliite (2014) Venkatraman et al., 2016; Farahmand cet al. 2017] Asadi et al. 2018|2019) A second strategy is penalize the policy for taking transitions where the model is inaccurate (Sorg al., 2010: Kidambi al., 2020 Yu et al. 2020b T Luo et al. 2019). Similar to all these prior methods, our approach will also use a modified 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. A third strategy is to directly optimize the model such that it produces good policies (Okda al. 2017, Amos et al. 2018 Srinivas et al. 2018 D'Oro et al., 2020 Nikishin et al.) (021)) as theoretically analyzed in Grimm et al. (2020). While our aim is the same as these prior methods, our approach will not require differentiating through unrolled model updates optimization procedures.
Our work builds on prior work that proposes model-based RI. objectives that are lower bounds on the true, expected returns. Kearns & Singh] (2002) provide lower bound that holds globally, but is only computable in tabular settings. Luo et al. (2019) provide a lower bound that can be efficiently estimated, but which only holds for nearby policies and models. Our lower bound combines the strengths of these prior works, providing a lower bound that holds globally and can be efficiently estimated in MDPs with continuous states and actions. Unlike any lower bounds from prior work, ours mends the objective mismatch problem.
Our theoretical derivation builds on prior work that casts model-based RL two-player game tween a model-player and a policy-player (Bagnell et al.) [001] Nilim & El Ghaoui 2003 Ross et al., 2011 Rajeswaran et al. 2020). 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 GAN model (Bai et al. 2019 Chen et al. 2020 Kurutach et al. 2018)
The most similar prior work is VMBPO (Chow et al. 2020). The mechanics of our method are similar, also learning a dynamics model using a classifier 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..I). Indeed, most of the components of our method, including classifiers and GAN-like models, have been used in prior work, main contribution of our paper is a precise recipe for combining these components in way that provably maximizes expected return.
2

---

Under review as a conference paper at ICLR 2022
3 A UNIFIED OBJECTIVE FOR MODEL-BASED RL
Notation. We focus on the Markov decision process with states st, actions ae, initial state distribution po(so), positive reward function r(st, at) > 0, and dynamics p(s++1 st, dz). Our aim is to learn control policy (a(a; | st) with parameters 9 that maximizes the expected discounted return:
Σ Exr((se.a))) an)
max E
We use transitions (8t, de; I', St,,1) collected from the (real) environment train the dynamics model go(st+1 st, at), and use transitions sampled from this learned model to train the policy. To simplify notation, we will define a trajectory T to be the sequence of states and actions visited in an episode:
т _ (so: ag, $1, ..). We then define R(T) _ Ei= y'r(st, at) as the discounted return of a trajectory. Finally, we define two distributions over trajectories. First, p=((r) is the distribution over trajectories when policy To interacts with dynamics p(s' s, a); q" (T) is the distribution over trajectories when policy Th interacts with the learned dynamics (0(+(+1 St, dz):
0c
p" (T) po( II p( sc, (e:Tr)) 1=0
st), (T) po (o II
(0(+)) | |s,, a (T)(( | st).
==
Desiderata. Our aim is to design an objective (()) that can be used to jointly optimize both the policy (No(ae) se)) and the dynamics model (g(S1+1 8t, at)), and which is lower bound on the expected return in the true environment.
An objective for model-based RL. We now introduce an objective that achieves these aims. Our objective will be the policy's '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 define our objective
(() A (r
Σ 7'F(se) ae (61))
where the modified reward function is defined as
(((sct] st, st,
F(5z, ae, $141) (1 7) log r(st, ar) + log
(($c+1 8t, st, - (1 7) log(1 - 7)).
(3)
This objective maximizes an augmented reward under the learned dynamics. The augmented reward function F penalizes the policy for taking transitions that are unlikely under the true dynamics model, similar to prior work (Eysenbach estimate this augmented reward without knowing the true environment dynamics by using GAN- like classifier. We will optimize this lower bound with respect to both the policy ae(a, st) and the dynamics model (0 8t, az). For the policy, this optimization entails performing RI. to maximize the modified reward using samples from the learned model; the only difference from prior work is the modification 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 is large) and similar is large). This objective differs from VMBPO (Chow et al. (020)) by
al. 2020al Yu et al. 202T). Later, we will show that we can
to real dynamics (i.e., log q 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 ss expected return under that model. Second, our objective is 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, course, maximizing the log(.) of the expected return is equivalent to maximizing the expected return.
will take the logarithm of the expected return. Of
Theorem 3.1. The following bound holds for any dynamics q(
st, at,az) and policy x(a st):
Σ > 7rr(st, an) 2 L()).
log E,
0=0
3

---

Under review as a conference paper at ICLR 2022
The proof is presented in Appendix A.33 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. To the best of our knowledge. this is the first global (unlike Luo et al. (2019)) and efficiently-comuuaable (unlike Kearns & Singh 2002 lower bound for model-based RI..
Sec. 4 will introduce an algorithm maximize this lower bound. While this lower bound may not be tight, experiments in Sec. demonstrate that optimizing this first lower bound yields policies that achieve high reward across wide range of tasks.
Tightening the lower bound. We now introduce a modification 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 (0) to denote this new lower bound. In addition to the policy and dynamics, this bound will also depend on time-varying discount, (o()), in place of the typical 7 term. Similar learned
discount factors have been studied in previous work on model-free RI. (Rudner et al. (021). We
define this objective as follows:
~ 2 Eg
vo (o(t)r ae, ac,)))) Σ t=
(4)
where the augmented reward is now defined
(Se, de, 1t+1) & log r(se, ae) + 1 I'olt- 1) (L)
p( St, oe St at a1. + log (t)
log
90((5+1 8t, ass 8t-1, at-1,
and (o())= Ei (tt))) ( is the CDF of the learned discount function (i.e., (t is probability distribution over 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 arbitrary policy π T(a | ) be given. The objective L,( is also lower bound on the expected return objective, log EE, [Er=o 7'r(se, ax) 2 (0), and this bound becomes tight at optimality:
rrr(s,, an) mi)
log Em
max (-(0)-
(F)-no(t)
The proof is presented in Appendix A.4 One important limitation of this result is that the learned dynamics that maximize this lower bound make the bound tight may be non-Markovian. Intrigu- ingly, this analysis suggests that using non-Markovian models, such as RNNs and transformers, may accelerate learning on Markovian tasks. This paper does not propose an algorithm for optimizing this more complex lower bound.
The optimal dynamics are optimistic. We now return to analyzing the simpler lower bound (C(O) in Eq. (). 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
(T) = p(r') (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
R(T)
Would the optimistic dynamics overestimate the policy's '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 the dynamics model.
4 MISMATCHED NO MORE
The previous section presented single (global) lower bound (C 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 method MISMATCHED NO MORE (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
4

---

Under review as a conference paper at ICLR 2022
Classifier
Learned dynamics model
policy
log ai)),a)) log
Learned dynamics model
policy
log Lggu((s, stg) log
Learned
policy
logg r(es, Rts) - log
dynamics model
a model-based RL algorithm that learns a policy, dynamics model, and classifier. The classifier distinguishes real transitions from model transitions. The policy and dynamics model are jointly optimized to sample transitions that yield high return and realistic, as estimated by the classifier.
Figure Mismatched
the transition probabilities of the real environment, p($++1 8t, az). which are unknown. We address this challenge by learning classifier (Sec. 4.1 and then describe the precise update rules for the policy, dynamics model, and classifier (Sec.
4.1 ESTIMATING THE AUGMENTED REWARD FUNCTION
To estimate the augmented reward function, which depends the transition probabilities of the real environment, we learn a classifier that distinguishes real transitions from fake transitions. This approach is similar to GANs (Goodfellow al., 2014) and similar prior work in RL (Eysenbach) it al. 2020a Yu et al. 2021). We use Cg(st, St+1) E [0, 1] to denote the classifier, For an optimal classifier, we can use the classifier's predictions to estimate the augmented reward function:
Co(St, Oe: St+1)
F(st; ar, 8t41) ((sr + log Co( Oes 8141)
(5)
sslog p(*c (8evar) -og do als,,0,)
The approximation above reflects function approximation error in learning the classifier.
We present complete method, which trains three components: classifier, policy, and 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 classifier (using a GAN-like objective). In describing the loss functions below, we use the superscripts (-) 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.2 UPDATING THE MODEL, POLICY, AND CLASSIFIER
Updating the policy. The policy is optimized 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 off-policy actor-critic method.
We define the Q function as sum of augmented rewards under the learned dynamics model:
(ssi 0z)
E x(aglos).
Σ
"F(ss
St,
(6)
go(*s (1lg
Life
We approximate the Q function using neural network ( ax) with parameters ø. We train the Q function using the TD loss on transitions sampled from the learned dynamics model:
Lo (st, ae, Tzs 9t+1
model,
vb) (Qu(se, an) [vr]s))
where 1. is the stop- gradientt operator and yt F(s,, a,, model + 1nfferrlssttt uls (s (s model
a.+1)))
The augmented reward, 7. estimated using the learned classifier. To estimate the corresponding value function, we use I-sample approximation: (s,ss)) = (st, a₂ (a(a st)). The policy trained to maximize the expected future (augmented) return, as estimated by the Q function:
max C, (s:; Ery(a,(s)) ( ot)
(
In our implementation, regularize the policy by adding additional entropy regularizer. Following prior work (Fujimoto et al. 2018), 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 for details.
5

---

Under review as a conference paper at ICLR 2022
Algorithm Mismatched More (MnM) an algorithm for model-based RI.. The method alternates between training the policy on experience from the learned dynamics model with augmented rewards and updating the modeltclassifier using a GAN-liG -ike loss. Updates are gradient steps with the Adam optimizer.
1: while not converged do
2: Sample experience from learned model and modify rewards using the classifier (Eq. 5
Update policy and Q function using the model experience and augmented rewards (Eq.s and
4: Update model and classifier using GAN-like losses (Eq.s |9aad (T].
5: (Infrequently) Sample experience from real model.
6: return policy nT st)..
Updating the classifier. We train the classifier distinguish real versus model transitions using the standard cross entropy loss:
max La(st geal ot 1841, 6) & log Co(st
model
( or 5841) + log (1 Colo 81.11
gcal real model
Note that the real transition (staal atta o8/11 and model transition (sfral, agaa o841 have the same initial state and initial action.
Updating the dynamics model. To optimize the dynamics model, we rewrite the lower bound in terms of single transition (derivation in Appendix A.6):
Ca(stal real medaly S141
real 0) = E mertl (#e+4 soll (27)) (s roal
V (5++1)
model
+ log 1 Ca(s d S11
(10)
The approximation above reflects approximation error in learning the optimal classifier. This approxi- mation is standard in prior work on GANs (Goodfellow et al., 2014) and adversarial inference (Don- ahue et al. 20166] Dumoulin et al. (016). The procedure for optimizing the dynamics model and the classifier resembles a GAN (Goodfellw et al. 2014): the classifier is optimized to distinguish real transitions from model transitions, and the model is updated to fool the classifier (and increase rewards). However, our method is not equivalent to simply replacing a maximum likelihood model with GAN model. Indeed, such 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 classifier term. These changes enable the model and policy optimize the same objective, which is a lower bound on expected return.
Algorithm summary. We summarize the method in Alg. T and provide an illustration in Fig. Implementing MnM on top of standard model-based RL algorithm is straightforward. First, create an additional classifier network. Second, instead of using the maximum likelihood objective to train the model, use the GAN-like objective in Eq. TO to update both the model and the classifier. Third, add the classifier's 's logits to the predicted rewards (Eq. 5). Following prior work (Janner et al. 2019). we learn a neural network to predict the true environment rewards.
5 EXPERIMENTS
We present two sets of experiments. Our first 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.1 UNDERSTANDING THE LOWER BOUND AND THE LEARNED DYNAMICS
In this section, we begin by studying the seemingly contradictory attributes of our method: optimistic dynamics and pessimistic policies, and end by confirming 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 stochastic- ity more favorable for solving the task. We use 10x 10 gridworld with highly stochastic dynamics and a sparse reward function. The results, shown in Fig. (Left),, show that MnM outperforms both Q-learning and VMBPO. In line with our theory, the dynamics learned by MnM (Fig. 2a (Right)) alter

---

Under review as a conference paper at ICLR 2022
start
goal
1.25
Correct Model
MnM (ours)
1.00
VMBPO
Return
0.50
0.25
0.00
Nrst
Nrch with an traccurate Model
200
600
800
1000
1200
1400
Nedh with Irsccunati Mudell
Iterations
clweofier
(b) Inaccurate Models Figure 2: Two Didactic Experiments. (Leeft) We apply MnM navigation task with transition noise that moves the agent to neighboring states with equal probability. MnM solves this task more quickly than Q Q-learning and VMBPO. The dynamics learned by MrM are different from the real dynamics, changing the transition noise (blue arrows) lo point towards the goal. (Right) We simulate function approximation by a learning model that forced to make the same predictions for groups of 3 states, resulting model that inaccurate around obstacles. The classifier term compensates for this function approximation error by penalizing the policy for navigating near obstacles.
(a) Stochastic Gridworld
+50
+1 +51
(log) VMBPO objective
2.45
Overestimate
(log) Expected Return
Underestimate
2.40
Return
400
MnM objective
2.35
200
2.30
Q-learning
MnM
VMBPO
5
10
15
Iterations
(a) Testing for risk seeking behavior.
Comparing objectives.
Figure 3: Lower Bounds and Risk Seeking. (Left) On a simple 3-state MDP with stochastic transition in one state (red arrows), MnM converges the reward-maken mwaadd-aximizzing policy while VMBPO learns strategy with lower rewards and higher variance (as predicted by theory). (Right) We apply value iteration to the gridworld from Fig. 2a analytically compute various objectives. predicted by our theory, MrM objective is lower bound on the expected return, whereas the VMBPO objective overestimates the expected return.
the environment stochasticity to lead the agent towards the goal, increasing the probability of collect- ing 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 the reward function (which encourages pessimistic behavior), difference that has large effect this task.
Our augmented reward function contains two crucial components, (/) the classifier term and (2) the logarithmic transformation of the reward function. We test the importance of the classifier term in correcting for inaccurate models. To do this, we limit the capacity of the MnM dynamics model that it makes "low-resolution" predictions, forcing all states in 3 3 blocks have the same dynamics. We will use the gridworld shown Fig. 2bb which contains obstacles that occur at a finer 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 classifier (whose capacity have not limited) detects that the dynamies 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 classifier term attempts to navigate through the wall and fails to reach the goal.
We then test the importance the logarithmic transformation by comparing MnM to VMBPO, which includes the classifier 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, the 3-state MDP in Fig. Ba (top) where numbers indicate the reward at each state. While moving to the right state yields slightly higher rewards, "wind" knocks the agent out this state with probability 50% the reward-maximizan strategy is move the left state. While MnM
7

---

Under review as a conference paper at ICLR 2022
MnM (ours)
MBPO
SAC
VMBPO
Hoppero2
1000
everage retm
12500
220:1-
0202
00:00
1008
1093
28
0.2
600
mettower d wiitoo 1po +2
345
metsowrd d doer octn 27
-202
mettawnld reach.u2
2V2rge raturn
930
-1232
1200
-1.00
-1433
031-
1500
1200
3300
ovenge
505
2000
1003
3300
230
environment steps
envtronment steps
onuironment steps
enviromment steps
We compare MnM MBPO and SAC simulated control tasks. On the henchmark locomation tasks (top left), MnM performs comparably with MBPO. On many of the other tasks with sparse rewards that pose exploration challenge, MnM outperforms both MBPO and the model-free baseline. These experiments suggest that maximizing a well defined bound on expected return, as done by our method, can lead to improved performance on difficult tasks.
Figure 5: Comparison Robotics Tasks:
Mn. learns the reward-maxiiing receives lower returns (with much higher variance).
strategy, VMBPO learns a policy that goes to the right state and
Finally, we study how the MnM objective compares to alternative objectives. We use the gridworld from Fig. [2 and use version of MnM based on value iteration to avoid approximation error. Plotting the MnM objective in Fig. 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.c.. VMBPO) overestimate the expected return.
5.2 COMPARISONS ON ROBOTICS TASKS
Our next experiments use continuous-control robotic tasks to answer two questions. We first investigate whether MnM performs at least comparably with prior work. We then study tasks with sparse rewards and more challenging explo- ration, where we suspect the optimistic dynamics learned by MnM may be beneficial. We illustrate a subset of the environments in Fig. 4 include implementation details in Appendix C.2 One detail of note is that we omit the reward augmentation (Bq. 3 for MnM during these exper- iments, as it hinders exploration leading to lower returns. We use MBPO (Janner et al.) (019)) 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.
Figure 4: Environments: Our experi- ments included tasks from four benchmarks: (clockwise from top-left) OpenAl Gym, DM Control, Metaworld, and ROBEL..
Our first comparison uses three locomotion tasks from the OpenAI Gym benchmark (Brockman et al.) 2016), which has become the standard benchmark for model- RL algorithms. Thes tasks have dense rewards and pose no significant 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 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 (Tassa et al. 2018). four manipulation tasks from the Metaworld benchmark YYu et al. 2020ab, and four dexterous manipulation tasks from the ROBEL benchmark (Ahn et al. 2020). See
8

---

Under review as a conference paper at ICLR 2022
MnM Dynamics:
Real Dynamics:
Figure 7: Optimistic Dynamies: (Leftt) On the Pusher-va task, the MnM dynamies model makes the puck move towards the puck move towards the gripper before being grasped. (Right) On the #alECheetat-v2 task, the MnM dynamies model helps the agent stay upright after tripping.
Appendix for experiment details. The results shown in Fig. 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, observe that MnM frequently outperforms all prior methods and, more importantly, it consistently does well across all tasks. We include ablation experiments, including a comparison to VAML (Farahmand al. 2017). in Appendix B
To investigate the benefits of MnM over prior model-based methods, we logged the Q val- ues throughout training and visualized them for the deetawoolddddrrawe---eeee
MBPO
values
task in Fig. 6 For fair comparison, we use Q values cor- responding to just the task reward, omitting the logarithmic transformation and classifier term typically used by MnM. Fig. shows that MnM yields Q values that are more accurate and more stable than MBPO. This figure suggests that MBPO may be exploiting inaccuracies in the learned model.
30*
101
gradient updates
Figure 6: Model exploitation: The very large Q values of MBPO suggest model exploitation, which our method appears to avoid.
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 otherwise-detenn deterministic dynamics appear stochastic. On the Fusher-vz 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 not physically realistic, it may make the exploration problem easier. On the Hal fChethhhhhhh 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.
CONCLUSION
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 train their dynamics model, but also guarantees to them that updating their model (using our objective) will result in 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.
Limitations. Compared with methods based on maximum likelihood, our method requires learning an additional classifier and balancing the capacity of that classifier against the capacity of the dynamics model. Unlike VMBPO, our method does not have a single hyperparameter to control the gap between the proposed objective and the expected return objective.

---

Under review as a conference paper at ICLR 2022
Ethics Statement. This paper works to advance model based reinforcement learning, which can enable data driven policies to learn from smaller datasets than their model-free counterparts. In doing so, this expands the range of applications for which reinforcement learning can be used. Many of these low data regimes, which can include things such as health care, require not just data efficiency but also safe online fine-tuning. It is important note that although our method optimizes lower-bound on the expected returns in theory, this lower-bound utilizes function approximators, which do not come with guarantees practice.
Reproducibility. Our appendix contains proofs for all theoretical claims, the hyperparameters and implementation details of each algorithm used, ablation experiments illustrating the expected behavior of our method under varying circumstances of interest, implementation details that we did and did not find helpful, and finally the necessary details to recreate each environment used in the paper. Furthermore, we will be publicly releasing the code for our didactic experiments.
REFERENCES
Michael Ahn, Henry Zhu, Kristian Hartikainen, Hugo Ponte, Abhishek Gupta, Sergey Levine, and Vikash Kumar. Robel: Robotics benchmarks for learning with low-cost robots. In Conference on Robot Learning, pp. 1300-1313, PMLR, 2020.
B Amos, I Rodriguez, J Sacks, B Boots, and Z Kolter. Differentiable mpe for end-to-end planning and control. Advonces neural information processing systems, 2018.
Kavosh Asadi, Dipendra Misra, and Michael Littman. Lipschitz continuity in model-based reinforcement learning. In International Conference on Machine Learning, pp. 264 273. PMLR, 2018.
Kavosh Asadi, Dipendra Misra, Seungehan Kim, and Michel L Littman. Combating the compounding-crvs problem with multi-step model. arXiv preprint arXiv:1905.10 2019.
J. Andrew Bagnell, Andrew Y Ng. and Jeff G Schneider. Solving uncertain markov decision processes. 2001.
Xueying Bai, Jian Guan, and Hongning Wang. A model-based reinforcement learning with adversarial training for online recommendation. Advances in Neural Information Processing Systems, 32: 10:15575-,,746, 2019.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint Xiv: 540, 2016.
Song Chen, Junpeng Jiang, Xiaofang Zhang, Jinjin Wu, and Gongzheng Lu. Gan-based planning model in deep reinforcement learning. International Conference on Artificial Neural Networks, pp. 323-334. Springer, 2020.
Yinlam Chow, Brandon Cui, MoonKyung Ryu, and Mohammad Ghavamzadeh. Variational model-based policy optimization. arXiv preprint arXiv:2006.05443, 2020.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful of trials using probabilistic dynamics models. In Advances Neural Information Processing Systems, pp. 4754 4765, 2018.
Mare Peter Deisenroth and Carl Edward Rasmussen. PILCO: model-based and data-cfficient approach to policy search. In International Conference on Machine Learning (ICML), pp. 465-472, 2011.
Jeff Donahue, Philipp Krähenbühh, and Trevor Darrell. Adversarial feature learning. arXiv:1605.09782. 2016.
arXiv preprint
Pierluca D' Oro, Alberto Maria Metelli, Andrea Tirinzoni, Matteo Papini, and Marcello Restelli. Gradient-aware model-hased policy search. In Proceedings of the AAAI Conference Artificial Intelligence, volume 34, pp. 3801-3808, 2020.
Vincent Dumoulin, Ishmael Beighazi, Ben Poole, Olivier Mastropictro, Alex Lamb, Martin Arjovsky, and Aaron Courville. Adversarially Iearned inference. arXiv preprint 160iv:1606.00 00704, 2016.
Benjamin Eysenbach, Shreyas Chaudhari, Swapnil Asawa, Sergey Levine. and Ruslan Salakhuidinov. Off- dynamics reinforcement learning: Training for transfer with domain classifiers. In International Conference on Learning Representations, 2020a.
Benjamin Eysenhach, Ruslan Salakhutdinov, and Sergey Levine. C-learning: I Learning to achieve goals via recursive classification. In International Conference on Learning Representations, 2020b.
10

---

Under review as a conference paper at ICLR 2022
Amir-massoud Farahmand, Andre Barreto, and Daniel Nikovski. Value-aware loss function for model-based reinforcement learning. In Artificial Intelligence and Statistics, pp. 1486-1494, 2017.
Scott Fujimoto, Herke van Hoof, and David Meger. Addressing Function Approximation Error in Actor-Critic Methods, International Conference on Machine Learning (ICMI., 2018.
Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville,, and Yoshua Bengio. Generative Adversarial Nets. In Advances in Neural Information Processing Systems (NIPS), 2014, URL at ptttts://aaaat wwtp:////ppppppttt /1400 2661. pdf
Christopher Grimm, André Barreto, Satinder Singh, and David Silver. The value equivalence principle for model-based reinforcement learning. arXiv preprint arXiv:201 2020.
Sergio Guadarrama, Anoop Korattikara, Oscar Ramirez, Pablo Castro, Ethan Holly, Sam Fishman, Ke Wang. Ekaterina Gonina, Neal Wu, Efi Kokiopoulou. Luciano Sbaiz, Jamie Smith, Gábor Bartók, Jesse Berent, Chris Harris. Vincent Vanhoucke, and Eugene Brevdo. TF-Agents: A library for reinforcement learning in tensorflow, 2018. [Online; accessed 35-Jume-2019)
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream control: leamming hehaviors by latent imagination. In International Conference on Learning Representations, 2019.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When trust your model: Model-based policy optimization. In Advances in Neural Information Processing Systems, pp. 12519 12530, 2019.
Joshua Joseph, Alborz Geramifard, John W Roberts, Jonathan P How, and Nicholas Roy. Reinforcement learning with misspecified model classes. In 2013 IEEE International Conference on Robotics and Automation, pp. 939-946, IEEE, 2013.
Michael Kearns and Satinder Singh. Near-optimal reinforcement learning in polynomial time. Machine learning. 49(2-3):209-233 2002.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and T. Joachims. Morel Mooel-bbsed offline reinforcement learning. ArXiv, abs/2005 05951, 2020.
Thanard Kurutach, Aviv Tamar, Ge Yang, Stuart J Russell, and Pieter Abbeel. Learning plannable representations with causal infogan. In Advances in Neural Information Processing Systems, pp. 8733-8744, 2018
Nathan Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in model- reinforcement learning. In Learning for Dynamics and Control, pp. 761-770. PMLR, 2020.
Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorithmic framework for model-hased deep reinforcement learning with theoretical guarantees. In ICLR (Poster), 2019.
Oliver Mihatsch and Ralph Neuneier. Risk-sensitive reinforcement learning. Machine learning, 49(2):667-290, 2002.
Evgenii Nikishin, Romina Abachi, Rishabh Agarwal, and Pierre-Luc Bacon. Control-oriented model-based reinforcement learning with implicit diffcrentiation. arXiv preprint arXiv:2006.00 03273, 2021.
Arnab Nilim and Laurent El Ghaoui. Robustness markov decision problems with uncertain transition matrices. In NIPS, pp. 839-846. Citeseer, 2003.
Junhyuk Oh, Satinder Singh, and Honglak Lec. Value prediction network. arXiv preprint 170.:00470999 2017.
Masashi Okada, Luca Rigazio, and Takenobu Aoshima. Path integral networks: End-to-end differentiable optimal control. arXiv preprint (xXXx::700009597 2017.
Aravind Rajeswaran, Igor Mordatch, and Vikash Kumar. A game theoretic framework for model based reinforcement learning. In International Conference on Machine Learning, pp. 7953-7963. PMLR, 2020.
Stéphane Ross, Geoffrey J Gordon, and J Andrew Bagnell. A Reduction of Imitation Leamingg and Structured Prediction to No-Regret Online Learning. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2011.
Stéphane Ross and Drew Bagnell. Agnostic system identification for model-based reinforcement learning. In ICML, 2012.
Tim GJ Rudner, Vitchyr H Pong, Rowan McAllister, Yarin Gal, and Sergey Levine. Outcome-driven reinforce- ment learning via variational inference. arXiv preprint arXiv:2104. 100000004 2021.

---

Under review as a conference paper at ICLR 2022
Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans, Advances neural information processing systems, 29:2234 2242, 2016.
Julian Schrittwieser, Ioannis Antonoglou. Thomas Hubert, Karen Simonyan, Laurent Sifre. Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis. Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839))600--69)) 2020.
Jonathan Sorg, Satinder Singh, and Richard Lewis. Internal rewards mitigate agent boundedness, In ICML, 2010.
Aravind Srinivas, Allan Jabri, Pieter Abbeel, Sergey Levine. and Chelsea Finn. Universal planning networks: Learning generalizable representations for visuomotor control. In International Conference on Machine Learning, pp. 4732-4741. PMLR, 2018.
Erik Talvitie. Model regularization for stable sample rollouts. In UAL pp. 780-789, 2014.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmalcki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint 1801:080000069) 2018.
Arun Venkatraman, Roberto Capobianco. Lerrel Pinto, Martial Hebert, Daniele Nardi, and Andrew Bagnell. Improved learning of dynamics models for control. In International Symposium on Experimental Robotics. pp. 703-713. Springer, 2016.
Grady Williams, Nolan Wagener. Brian Goldfain, Paul Drews, James M Rehg. Byron Boots. and Evange- los A Theodorou. Information Theoretic MPC for Conference on Robotics and Automation (ICRA), 2017.
Reinforcement Learning. In International
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsca Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. Conference on Robot Learning, pp. 1094- 1100. PMLR, 2020a.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-base offline policy optimization. arXiv preprint MxXi::2005.1325555 2020b.
Tianhe Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, Chelsea Finn. Combo: Conservative offline model-based policy optimization. arXiv preprint arXiv:2 2021.
Brian D Ziebart, Modeling purposeful adaptive behavior with the principle maximum causal entropy. 2010.
12

---

Under review as a conference paper at ICLR 2022
A PROOFS AND ADDITIONAL ANALYSIS
A.1 VMBPO MAXIMIZES AN UPPER BOUND ON RETURN
While MnM aims maximize the (log) of the expected return, VMBPO aims to maximize the expected exponentiated return:
MnM:
log Eg
Σ Err(s) 1=0
VMBPO:
log Bx T
Y' v'r(snai)
where 1/ > 0 is 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 the expected exponentiated return, as done by VMBPO, is not equivalent to maximizing the expected return. Rather, it corresponds to maximizing sum of the expected return and the variance of the return (Mihatsch Neuneier
2002
Page 272):
log E ( = E Σ r an) +
" Var,
Σ (rr(s,,))) L1=0
+ +
Thus, in environments with stochastic dynamics or rewards (e.g., the didactic example in Fig. 3 VMBPO will prefer receive lower returns if the variance of the returns much higher. We note that the expected exponentiated return is an upper bound on the expected return:
log E,r [e (((())))) ( 2 E
Σ 1=0
rrrss, on)
This statement is 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 random policy optimal for these trivial MDPs.
A.2 HELPER LEMMAS
We start by introducing a simple identity that will help handle discount factors in our analysis.
Lemma A.1. Define p(H) = GEOM(1 7) as the geometric distribution. Let discount factor
y (0, 1) and random variable x4 be given. Then the following identity holds:
1
E
Ep(H)
Σ
x
xxx xt.
The proof involves substituting the definition of the Geometric distribution and then rearranging terms.
Proof.
h
00
H
Ep(H)
Σ
It
= =1 - 7)
Σ Et
1==
==
= (1 - 7) (xo + (20 + 21 + + x2) +1 +...)
= (1 - 7) (xo(1 + + + 11 72² + ...) +...)
1
72
(1 7) zD y
x1 1 1-9
+ x2
1-7 +...
Y Tz.
==
The second helper lemma describes how the discounted expected return objective can be written as the expected terminal reward of a mixture of finite-length episodes.
13

---

Under review as a conference paper at ICLR 2022
Lemma A.2. Define H)) GEOM(1 ?) the geometric distribution, and p(T H) as distribution over length-H episodes. We can then write the expected discounted return objective as follows:
R p H=00)
Σ 7 r(s a+)
=
By(H)) (BtriHH (r(ss, an)]]
(11)
lf p(H))(T HH)r(s,,aa))))) addH.
(12)
Proof. The first identity follows from the definition of the geometric distribution. The second identity writes the expectations as integrals, which will make future analysis clearer.
A.3
PROOF OF LEMMA 3.1
Proof.
Σ ((1())))))
(8) log /f p(H)p(r H an)drr
1-0
log /ʃ p(H)
plr H H)
H) vo(r H H)r((n log(1 -n)
gu(t
2 ((H) lon
p(r H go(r H Hrj(s, ag)dd) dH log(1 -7)
go(r HH= H)
- /ʃ ((H))o(T)
H (logp(7 H) go H) + log r(sn as)) drdH log(1 -7)
@ /ʃ P(H)RO(T H. H)
logu(se:1 (.) log log g0(4141 (())))) - loassforr + log logr(en. drdH - log(1 - 7)
2 [ʃ M = oo)
log posp(1 (1,a)) log ( ae)))) log logr(ss,
drdH loR(1 (7))
1 f (o(r)) S (IM))
Σ
logp(ee11 (8))))) log ( secras) + logr(sH; aHdr - loR(1 -n)
log (8,a))) 1o Reraa)) + logn(ss, dr log(1
(0)
20 (T)
(log ac) 1o ssraa)) + 7) logr(sH; ax)) log(1
(gg(r)
p(se4 (s))0)) 1o (ee)))) + (1 - -n) log r(s.. -1 7) log(1 - 130
For (b), we applied Jensen's inequality. For (c), we applied Jensen's inequality again. For (d), substituted the definitions of po(T H) and go(T H). For (e), noted that the term inside the summation only depends on the first steps of the trajectory, so collecting longer trajectories will change the result. This allows using a single infinite-length trajectory. For (f) we recalled the definition go(7) = go(r = H = x)) and swap the order of integration. For (g), we express the inner integral over p(H) an expectation. For (h), we applied the identity from Lemma A.I For (i), we moved the constant log - 7) back inside the integral and rewrote the integral as an expectation. We have thus obtained the desired result.
For (a), we applied Lemma A.2
rewrite the integral
expectation
A.4 PROOF OF LEMMA 3.2
Before presenting the proof of Lemma 3.1 itself, show how derived the lower bound in this more general case. While this step is not required for the proof, we include because sheds light on how similar lower bounds might be derived for other problems. We define ye (H) to be learned distribution over horizons H. We then proceed, following many of the same steps for the proof of
14

---

Under review as a conference paper at ICLR 2022
Lemma 3.1
log Bm
Σ 2/r(s) ar) 1-0
( log /I
Mir, H) go(r, H)
go(r, H)r(sw, (arrss,aa)dddde an)drdH log(1
qu(t, H) (log ((),H) H) log go( H) log rsn, asr)dr) dHl log(1 7)
((2
(0())007 H)
log p(e441 (e, log go(se-1 Be, ax) + logp(H) - log (o()) + log nsw, an an)dr) log(l - 7)
( / (o ( | = =0) Σ (H(H))
((2
Σ log p(#c+1 [t, log go eeraa)) + log (HH) log rogr, an)dr) log(1 -7)
H=o
(0(T))
Σ 70(11)
Σ
se,ae) log ( s() logp(H) (og + log logr(sn.ag)dr Ing(1 -7)
1 (xp(t)
Σ T0(H)
pogp(s [()) log go ec.
+oog(t log log logr(en.
loggt
-=
Eng(r)
(())
(log p(s (g) log (o( [88. (H) (H log log (o( logr(sn
(hi
(ap(t)
g(t)
(logp(sH-1 (x) log (x.aa)) (H) log log (H rss ax))
R-o
2-
(7)
log log log r(s To(H 1)) (logp(sH+1 (s,aa) log (s6241 (x)) + Tie (H) (#f log -
(1- (1- Te(H - (log |sg ax) log (M5441 (g)) ) + Hf log (H) + log (s (e (II)
(xg(r)
Σ (o(H))
For (b), we For (a), we applied Lemma A.2 and multiplied the integrand by 2a(r H (HH
4e(7 H= H=
applied Jensen's 's inequality. For (c), we factored p(r, H) = p(T, H HH) H) and qo H) = g(T (w()H)) Note that under the joint distribution p(T, H). the horizon H p(H) GEOM - -) is independent of the trajectory, a single lnfinite-length trajectory and simplified the summand. For (e), we recall the definition expectation and wrote out the
For (d), we rewrote the expectation as expectation over
(T) g = H oo). For (f). we rewrote the integral definition of the geometric distribution, p(H). For (g). we regrouped the difference of dynamics terms. For (h), we noted used the fact that E (t) + T (o))
1. For (i), we substituted definition of the CDF function. For (j), we rearranged terms so that all were multiplied by the discount 7o (H). Thus, we have obtained the desired result. We now prove Lemma 3.2 showing that
Eq. becomes tight at optimality.
Proof.
2,10) (2) Iʃ oo(t, H) (logp(r, H) log (e (oggrsss.aaa) dH log(1 - 7)
/ʃ 96 (H (r) (logp(r) + (H) - log (o()) - 7H( (r) + (ogg(ss.aagaee dH log(( - 7)
(14)
For (a), we undo some of the simplifications above, going back to Eq. 13 For (b), we factor (T, H)- = (H T) and p(T, H) p(T)p(H). At this point, we can solve analytically for the optimal discount distribution, HH T):
((H)r(HHaa)) No (H (T) (H" (s) (1-7)R(T))
pl (H)r(HH.
(15)
In the second equality, substitute the definition of R(7). We then substitute Eq. 15 into expression for (0) and simplify the resulting expression.
(0) /] (H( (r) (loRw(r) + lountHry log (e( loan(tT) - log.zeen:a77 Jow[1==75 3J + log R(r) + logr(enr aH Joull
/] (o( (log((r) (r) log R(r)dr) d.H
(r) (lo) log (e) log R(7)) dr.
15

---

Under review as a conference paper at ICLR 2022
In the final line 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, qo (T):
R(T)
o (T)
(16)
We then substitute Eq. 16 into our expression for (0), and simplify the resulting expression:
2. (o())) [oartr] Losp(F] log log
p(r') jar' + Log 1o8-1177) dr
= log p(r))( = l,
Σ Ex'r(ven 2r)()))
We have thus shown that the lower bound C., becomes tight when we use the optimal distribution over trajectories qo (T) and optimal learned discount 7o (H T).
A.5 A LOWER BOUND FOR GOAL-REACHING TASKS.
Many RL problems can be better formulated as goal-reaching problems, formulation that does not require defining a reward function. We now introduce a variant of our method for goal-reaching tasks. Using p² (st+) to denote the discounted state occupancy measure of policy we define the goal-reaching objective as maximizing the probability density of reaching a desired goal 89:
max log p* ($t+ = =8)).
(17)
We refer the reader to Eysenbach et al. (0200b for a more detailed discussion of this objective. For simplicity, we assume that the goal fixed, 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 lower bound on the
RI. problem:
Lemma A.3. Let initial state distribution (i(s)). real dynamics p(5 St, az), reward function r(8,, at) > O, discount factor 7 (0, 1), and goal be given. Then the following bound holds for any dynamics g(8t+1 81, a4) and policy x(az 8:):
En'r(se (n)eaa))
logp" (8L+ 8g) 2 (yyo(r)
(18)
where (8(8 at, 81+1) A (1-7) log 8gg 8t, (z) log pl (8e41
88, (z) - log ((8441 8t, oz).
The proof, presented below, is similar to the proof of Lemma 3.1 The first term in the reward function, the log probability of reaching the commanded goal one time step in the future, is similar to prior work (Rudner et al,, (202T). The correction term log log incentivizes the policy to avoid transitions where the model is inaccurate, and can be estimated using separate classifier. One important aspect of this goal-reaching problem is that it is entirely data-driven, avoiding the need for any manually-designe reward functions.
Proof.
log p (ec+ =,,)= (2) log p(H)p(T H = H)p(s, | sx addd
log /f p(H) p(r н1 H) go(r ( | HH== H) play I ax)
go(s, en,an) qo(*g | sn, an)drd - log(1 - 7)
qa(7
1 // p(H)oo( H H) (loge(r | H = H) log go H log p(as ss asr) log (a(ey (er, drdH log(1 7))
(9) 1ʃ M(Hyoot H = oo)
log plses (.)0)) log ( se) + log ( lo Ris R2))df log(1 - 7)
(T( n(()) Σ Jonn(se11 (..0)) - log ne ())) logp(s, I (H.og) log (e (3. - log(( -n)
00(1)) (T) 7" (logp(set1 (e) - 1og 19,, (lox (2))0)) log (e =ever)) dr log(1
(4) (xg(+)
(logp(0e41 I (e))ee) log ge(eess revor) + (1 (oog (e) go(o) ae) (1 7) log(1 7))
16

---

Under review as a conference paper at ICLR 2022
Similar to the more complex lower bound presented in Eq. 4 this lower bound on goal- reaching can be modified (by learning a discount factor) to become a tight lower bound. The resulting objective would resemble model-based version of the algorithm from Rudner al. (02TT.
A.6 DERIVATION OF MODEL OBJECTIVE (EQ. 10
Our lower bound depends 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 St.
Σ f' "F(ss,ar)
T(ae se), de (=(
[[oe; (1,411) + (8(+1)) st]
x(oele)). go(ss Maya)
(1 log + log Calse, a2, S421) (1 - -7) log(1 -7) + (V(++11) (8t, 01, 8141)
(#) (1p))a))
In (a), we substituted the definition of the augmented return. For the purpose of optimizing the dynamics model, we can ignore all terms that do not depend on 8+41- Removing these terms, we arrive model training objective (Eq. 10
B ADDITIONAL EXPERIMENTS
AverageReturn
MnM (ours)
4000
MBPO
SAC
2000
VAML
Velue-Weighted
0.00
0.25
0.50
0.75
1.00
1.25
1.50
Env steps
1e5
Figure 8: Alternative Model Learning Objectives: Using the DClawScrewFixed task, we compare MnM and MBPO (Janner et al. 2019) to two additional model learning objectives suggested in the literature. VAML (Farahmand al., 2017) and value-weighted maximum likelihood (Lambert et al. 2020). MnM (our method) outperforms these altemative approaches.
We compare MnM number of alternative model learning methods. MBPO (Janerr al. 2019 uses standard maximum likelihood model. We implement version of VAMI (Farahmand et al. 2017). which augments the maximum likelihood loss with an additional temporal difference loss; the model should predict next states that have low Bellman error. Finally,
compare a variant of the MBPO maximum likelihood model that weights transitions based on the Q values, an idea discussed (but not actually implemented) in Lambert al. (220). We implement this value weighting method by computing the Q values for the current states and computing a softmax over the batch dimension to obtain per-example weights.
We use the DClawScrewt xxed-v0 task for this experiment. The results, shown in Fig. 8 show that MnM outperforms these alternative approaches. We observe that the value-weighte performs
17

---

Under review as a conference paper at ICLR 2022
slightly better than the standard maximum likelihood model, while the VAML method performs noticeable worse.
metaworts
712
average retur
MnM
600
MBPO SAC
600
without
cptimism sugmentation
200
400
without model
optimism
300
environmeent steps
environmeett steps
Figure 9: Ablations Experiments:
Compared with MBPO (orange line), MnM uses a GAN-like model (red line) with a model optimism term and modifies the reward function.
We next run an ablation experiment to study the importance of few key design decisions. We compare MnM with ablations that omit the reward augmentation and the model optimism term. The results shown Fig. indicate that most of the benefit MnM comes from using GAN-like model. Because the dynamics of these tasks nearly deterministic, it is not surprising that the optimistic dynamics and the reward augmentation have only a small effect.
0.030
Validation Model MSE
0.015
0.005
0.000
0.5
1.0
2.5
3.0
Training steps
166
Figure 10: MnM trains stably. Despite resembling GAN, the MnM dynamics model trains stably, with validation MSE decreasing steadily throughout training. Different colors correspond to different random seeds of MnM. dashed line corresponds to minimum validation MSE maximum likelihond dynamics model.
With the implementation details described in Appendix C we found that the MnM dynamics model trained stably, despite resembling GAN. In Fig. TO we plot the validation MSE of the MnM model throughout training, observing that it decreases monotonically. Different lines correspond different random seeds, and the dashed line corresponds to the minimum MSE of a maximum likelihood model (a MBPO model). Note that the MnM model is not trained with this MSE objective, but with the GAN-like objective in Eq. I0 It is therefore not surprising that MnM does not perform well on this objective as the maximum likelihood model.
C IMPLEMENTATION DETAILS
All experiments were run on at least three random seeds.
C.I ALGORITIMS
For the tabular experiments that perform value iteration, we perform Polyak averaging of the policy and learned dynamics model with parameter T = 0.5. We found that the value iteration version MnM diverged without this Polyak averaging step. Experiments function) changed by less than le-6 across iterations, as measured using an Lo norm. For VMBPO used
Value Iteration (Fig. 2a (right), Fig. 2b Fig. 3
stopped when the MnM dynamics model (which depends the value
18

---

Under review as a conference paper at ICLR 2022
SAC
MBPO
MnM
HalfCheetah-v2
40
20
Hopper-v2
40
20
Walker2d-v2
40
20
dm dth.caatpolle-ssiinguppsppp
1
metssworlddwinnddw---errrrv
20
20
20
metaworldddoor-openn
40
40
40
metaworrld-reach-v?
20
20
20
metaworld-drawer-peee
me
40
40
40
DCLawPoserRandom-10
20
20
20
DClawTurnKandom-vv
40
40
40
DclawScrewFixed-vo
40
40
40
DClawSorewRandom--
40
40
40
Table I: Gradient updates per real environment step: This parameter was separately tuned for each method and each environment.
Q-learning (Fig. 2
For the experiments with Q-learning (both with and without the MnM compo-
nents), we performed e-greedy exploration with 0.5. We used learning rate of le - 2. For this task alone, we compute the MnM dynamics analytically by combining the 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. For VMBPO we used 1.
SAC for continuous control tasks. We used the SAC implementation from TF-Agents (Guadar rama et al. 2018) with the default hyperparameters.
MBPO for continuous control tasks. We implement MBPO on top of the SAC implementation from TF-Agents (Guadarrama al. 2018). Unless otherwise mentioned, take the default parameters from this implementation. We use an ensemble 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 per-coordinate mean, and then added the current state. These whitened predictions are clipped to have 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 the true dynamics. We perform batch of rollouts in parallel using batch size of 256. To sample the corresponding action, with probability 50% we take the action that was executed the true dynamics; with probability 50% sample action from the current policy. We found that this modification 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 le6. At the start of training, we collect le4 transitions from the real environment, train the dynamics model on this experience for le5 batches, and only then start training the policy. We use learning all components. To stabilize learning, we maintain a target dynamics model using an exponential moving average (r 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 D).
of 3e-4 for
MnM for Continuous Control Tasks We implement MnM top of the SAC implementation from TF-Agents (Guadarrama et al. 2018). Unless otherwise mentioned, we take the default parameters from this implementation. Our model architecture is exactly the same MBPO implementation, and we follow the same training protocol.
Unlike MBPO, MnM also learns classifier for distinguishing real versus model transitions. The classifier architecture is a 2 layers neural network with 1024 hidden units in each layer. We found that this large capacity important for stable learning. We add input noise with 0.1 while
19

---

Under review as a conference paper at ICLR 2022
Figure 11: Environments: Our experiments look at three locomotion tasks from OpenAl Gym Brockman al. 2016), the inverted pendulum task from DM Control TTssa ct al. 2018). four manipulation tasks from Metaworld (Yu et al. 2020a and four dextrous manipulation tasks from Robel (Ahn al. 2020). The Robel tasks use the same dynamics but different reward functions, so we only include an image of one task.
training the classifier. We whiten the inputs to the classifier by subtracting coordinate-wise mean and dividing by a coordinate-wise standard deviation. When training the classifier, we take samples from both the dynamics model and the target dynamics model as negative examples, finding that this stabilizes learning somewhat. Following the suggestion of prior work (Salimans et al.) (016)) 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 classifier input noise also decreased performance. We found that mixup had little effect. We found that the loss would often plateau around Ie4 batches, but would eventually start decreasing again after 2c4 2e5 batches.
Like the MBPO model, we first collect le4 transitions of experience from the real environment using a random policy, then train the dynamics model and classifier for le5 batches, and only then start updating the policy. Because the Q values poor at the start of training, term to the model loss (resulting in the optimistic dynamics model) after 2e5 batches (1e5 batches of model training. then le5 batches of podel+policy training). To further improve stability, compute the value term in the model loss by taking the minimum over two target value functions (like TD3 (Fujimoto al. 2018 ). We update the model, classifier policy, and value functions the same rate we sample experience from the learned model, which is more frequently than we collect experience from the real environment (see Table
only add the value
C.2 ENVIRONMENTS
Gridworld for Fig. Za This task is 10 10 gridworld with obstacles shown Fig. 2a There are four discrete actions, corresponding to moving to the four adjacent cells. With probability 50% the agent's action ignored and a random action taken instead. The agent starts in the top-left cell. The agent receives a reward of +0.001 at each time step and reward of + 10 when at the goal state. Episodes have 200 steps and we use discount ? 0.9.
Gridworld for Fig. 2b This task is a 15 x 15 gridworld with obstacles shown in Fig, 2b There four discrete actions, corresponding 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 7 0.9 and compute optimal policies analytically using value iteration. We implementing the aliasing by averaging together the dynamics for each block 3 x states. Importantly, the averaging was done the relative dynamics (e.g., action corresponds to move right) not the absolute dynamics (e.g., action corresponds to moving to state (3, 4)). To handle edge effects, we modified the averaged dynamics so that the agent could not exit the gridworld. We computed the classifier analytically using the true and learned dynamics models. However, the augmented rewards become infinite because the learned model assigns non-zero probability transitions that cannot occur under the true dynamics. This is not 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 classifier.
moving to the four adjacent cells. The dynamics
20

---

Under review as a conference paper at ICLR 2022
This task is a 10 x 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 s action is ignored and 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 reward of +0.001, transitions that do not change the distance to the goal have reward of +1.001, and transitions that decrease the distance to the goal have reward of +2.001. We use 7 0.5 and compute values and returns analytically using value iteration.
Gridworld for Fig. 3a
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 reward of + +1000 the goal state. We estimate the more complex lower bound (Eq. by also learning the discount factor. However, since 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.
Hal F22 Hopper-v2, Walker2d-v2. These tasks are taken directly from the OpenAl benchmark (Brockman et al. 2016) without modification.
dm-cartpole-swingue Saapppssarreeeee This task is taken directly from the DM Control bench- mark TTassa et al. 2018) without modification.
metaworrddwiiddww-pppn----
This task is based the window-open-v2 task from the Metaworld benchmark (Yu et al.) 2020a). To increase the difficulty of this task, we set random_init=Te and sparse reward function. The sparse reward function is +0 when the window is within 3 units of the (default) open position and otherwise.
2etaaoorldddorr--pen--2 This task based on the door-open-v2 task from the Metaworld benchmark (Yu et al., (2020a). To increase the difficulty of this task, we set random_init-Tru and use a sparse reward function. The sparse reward function is when the window is within 3 units of the (default) open position and +0 otherwise.
metahorlddreeac--v..
This task is based on the reach-v2 task from the Metaworld bench- mark (Yu et al., al., (0200)) To increase the difficulty of this task, set random init==rue and use a sparse reward function. The sparse reward function is +0 when the window is within I units of the goal and -1 otherwise.
This task is based on the drawer-open-ve task from the Metaworld benchmark (Yu et al. 2020a). To increase the difficulty of this task, we remove the reward shaping term (reward_for lareaad__oor_rrro_ee forccdiiigraiinnnglll and just optimize the reward for opening the drawer (reward__fo__oeening])
DClaaWoseRandom-vv, OClawscrevRandoo-vvo
DClawScreeFixed-Wg These tasks are taken directly from the ROBEL benchmark (Ahn)
DClawTurnRandom-y
al. 2020) without modification.
21


Reviews:
[{"review_id": "SVNmfI1ADor", "reviewer": "Reviewer_uKVv", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Objective mismatch problem is important in MBRL. And the proposed MnM is motivated by the idea of GAN to solve this problem. The whole paper is easy to follow and I also appreciate the theorems in the paper. However, I would like to see more comparisons, both theoretical and experimental, that if MnM performs better than the VAML framework [1,2] and the simple weighting strategy in [3]. Both these works are addressing the same problem, so it would be important to directly compare them, especially considering that MnM is a more complicated algorithm.\n\n[1] Value-Aware Loss Function for Model-based Reinforcement Learning. Amir-massoud Farahmand et al. AISTATS 2017\\\n[2] Iterative Value-Aware Model Learning. Amir-massoud Farahmand et al. NeurIPS 2018 \\\n[3] Objective Mismatch in Model-based Reinforcement Learning. Nathan Lambert et al. L4DC 2020", "summary_of_the_paper": "This paper concerns the objective mismatch problem in model-based RL that the model is optimized for prediction accuracy, not a good performance. The authors propose Mismatched no More (MnM) that maximizes a lower bound on expected reward in an adversarial training manner. They also give a global lower bound that holds for any dynamics. Experiments also validate the effectiveness of MnM.", "summary_of_the_review": "The overall paper is motivated and clear, The concern I have is the advantage over previous works that address exactly the same problem with relatively simple mechanisms.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Ar24yIo1Di7", "reviewer": "Reviewer_yCzC", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The method is well motivated, and the proposed augmented reward function is novel. The experiments are also well conducted, and the visualization of the learned dynamics provides intuitive illustrations that it tends to help the policy to succeed. The overall performance of MnM is also demonstrated to perform at least as well compared with SAC and MBPO.\n\nDespite these advantages mentioned above, I have some comments below that I’d like to hear from the authors in the responses.\n\n1. The augmented reward in Eq. (3) is informative, while in Theorem 3.1, the equality of the lower bound is hard to hold. In the proof of Theorem 3.1, Jensen’s inequality is applied, twice, to move the logarithm inside the expectation. For Jensen’s inequality, the equality holds when the function (logarithm here) is affine or the variable inside the function (the reward r here) is constant. Obviously, neither is the case here. Therefore, at least for the reward function in Eq. (3), maximizing over its cumulative expectation does not guarantee the optimality in the original expected return.\n\n2. Although, immediately, the following contexts about “Tightening the lower bound” gives a new reward function by employing a learnable $\\gamma_{\\theta}$, I cannot agree that this is really significant. From the proof of Lemma 3.2, the equality holds when the two analytical solutions are satisfied for $\\gamma_{\\theta}$ and the dynamics model, and the two analytical solutions are independent with each other. Then, for the learnable $\\gamma_{\\theta}$, it indeed plays a role of reward shaping that shapes the reward in Eq. (3) probably nonlinearly to reach the logarithm of the original expected return (probably a way to approximate the affine function to remove the Jensen’s gap?). Also, in the experiments, this new lower bound and reward function are not applied. So, I agree Lemma 3.2 demonstrates that there exists some way to modify the reward in Eq. (3) to have a tighter lower bound, while I think this is a benefit from reward shaping, instead of a thoroughly informative new reward function.\n\nMinor:\nIn the proof of Lemma 3.1, line (d) in the equation, the reduced $\\pi_{\\theta}(a_t|s_t)$ missed the logarithm.\n\nOverall, I think the proposed method is novel and the experiments give nice explanations for the learned dynamics in tending to assist the policy learning, and MnM indeed has its superiority compared to prior model-based RL methods.\n", "summary_of_the_paper": "This paper proposes a new model-based RL method by devising a new reward function, which incorporates a term measuring the difference between the learned dynamics model and the true dynamics probability, in addition to the logarithm of the original environment reward. Maximizing this augmented reward function is proved to maximizing a lower bound of the logarithm of the original expected return. The reward function is not practically usable since the true dynamics probability is not known. To practically apply the proposed reward function, a GAN-like classifier is introduced to differentiate the transition generated from the dynamics model and the true dynamics. The experiments show that under an unified objective, the learned dynamics tends to assign higher transition probability for some (s,a,s’) if this transition can assist the policy to achieve a higher return in the true environments.", "summary_of_the_review": "A good paper that introduces a novel model-based RL method, where the policy learning and dynamics modeling share a unified objective.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "AYzNlGj_7zQ", "reviewer": "Reviewer_7QkN", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Writing: The paper is written clearly, although it has a few typos. \n\n1. The link to the proof of Lemma 3.2 is wrong. \n2. Eq (8): What is $\\tilde Q_\\phi$? Is it $Q_\\psi$​?\n3. Appendix A.3, (d), the first cancelled term should be $\\log \\pi_\\theta(a_t | s_t)$​. \n4. Appendix A.1, \"Rather, it corresponds to maximizing a sum of the expected return **and** the *variance* of the return\"\n\nNovelty: The proposed method is similar to a prior work,  VMBPO. Both share similar components, with similar method of learning. However, there are a few difference which separates them apart. One of the most important differences is that MnM takes the logarithm of rewards, and derives a lower bound of the log of expected return, while VMBPO derives a variational lower bound of an upper bound of the log of the expected return. Given VMBPO as a prior work, The theoretical result is not so surprising. If the technique of logarithm of reward can be well justified, I think the empirical novelty is good. \n\nOther comments: \n\n> Following prior work (Janner et al., 2019), we learn an additional model to predict the true environment.\nWhat does it mean by \"additional model\"? \n\nIs the dynamics model learned well? From my personel research experience, learning the model using a GAN-like objective leads to an inaccurate model and sometimes produces unrealistic states. The signals from the discriminator are much weaker than the signals from supervised learning. What if we train $C$​ and $\\pi$​ with MnM, but train maximum likelihood model $q$​​? \n\nIn Figure 6: What does it mean by \"For fair comparison, we use Q values corresponding to just the task reward\"? MBPO uses $r$ while MnM uses $\\log r$, so they indeed differ a lot. \n\nAs the new reward function $\\tilde r$ takes the logarithm of old reward function $r$, it seems that (a) it's tricky to apply logarithm to non-positive numbers; (b) The new reward function $\\hat r$  is not invariant to the scale of the old reward function $r$, i.e., shifting $r$ can lead to a very different $\\tilde r$. Do all these $\\tilde r$​ work in practice? \n\nTheoretically, there is indeed a difference whether we transform the reward. But I'm not very convinced that this technique is essential in practice. The authors constructs a 3-state MDP in Figure 3, but some details are missing, e.g., how is $\\eta$​​​​ in VMBPO chosen? I agree that VMBPO optimizes an upper bound approximately, but it can somehow control overestimation by adjusting $\\eta$​, while I don't see how MnM controls underestimation. Moreover, what happens if we don't apply this technique in practice? I expected to see the ablation in a more complex environment. \n\n", "summary_of_the_paper": "This paper studies the the problem of objective mismatch problem in model-based reinforcement learning, which states that the goal of the model is different from the goal of policy optimization. To solve this issue, this paper proposes a single objective to train both the policy and the dynamics model, by deriving a lower bound of the (log of) expected return and optimizing the lower bound jointly. Based on the objective, the paper proposes the Mismatched no More (MnM) algorithm, which use a GAN-style objective to train the model. Experiments show that the proposed algorithm can match the performance of baseline algorithm (MBPO) when exploration is not a big issue, while outperforms it significantly for sparse-reward tasks. ", "summary_of_the_review": "This paper is written well and the idea sounds very interesting and promising. However, as stated in the review, I'm not fully concerned by the log reward technique. Overall, I think there is more merit in this paper than the flaw so I recommend weak acceptance. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "AA_eC9RLdjW", "reviewer": "Reviewer_D2vZ", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Novelty\n\nTo the best of the reviewer's knowledge, the idea is novel. The authors propose a lower bound on the logarithm of expected returns that the policy and the model can jointly optimize against an adversary that tries to distinguish between real and model rollouts. As a result, the MnM model mitigates the objective mismatch yet produces realistic samples. This contrasts the approach with other non-MLE methods (e.g. see [1]) that address the objective mismatch but produce unrealistic rollouts. The authors might want to emphasize the property more in the paper and study the effect in more detail.\n\nSignificance\n\nThe significance of the paper is limited. The main concern of the reviewer is about the empirical results: \n- First, most experiments appear to have a very limited number of runs. The reviewer did not find the details in the paper explicitly articulating the number of random seeds. The reviewer might infer from the plots that MnM and the baselines were run 2-3 times (e.g. Figure 5, `meta-world-door-open-v2`). Given the high variance of the results, it is hard to make conclusions about the comparison of MnM to the baselines. Reporting the results using 10 random seeds might improve the credibility of the results. See [2] for a discussion about the significance of the experimentation in RL.\n- Second, a few important baselines are missing. In tabular experiments, there is no comparison to an MLE MBRL agent. In continuous control, there is no comparison to any MBRL algorithm that addresses the objective mismatch. Why not use VMBPO in these experiments if you already used it in tabular experiments? Furthermore, Appendix A.4 suggests that the optimal dynamics for the lower bound are achieved by weighting the true dynamics with returns. Why not try return-weighted MLE as a baseline? Including other baselines that address the objective mismatch would further increase the significance of the results (see notes about related work).\n- Lastly, while MnM mitigates the objective mismatch, it is unclear from the results what are the benefits of this. Figure 5 suggests that out of 12 continuous control environments, only on `DClawScrewRandom-v0` we can make a conclusion that MnM significantly outperforms the presented baselines. In 3 MuJoCo environments, the performance of MnM and MBPO does not differ significantly. In the other 8 environments, it is hard to make definite conclusions due to the variance of returns.\n\nThe authors provide theoretical analyses associated with MnM. There are several concerns about this part of the paper as well:\n- The proposed method optimizes a bound $L(\\theta)$ on the log expected returns. The authors claim that “this bound becomes tight under certain assumptions”. To support the claim, the paper introduces another bound $L_\\gamma(\\theta)$ and proves that the supremum of $L_\\gamma(\\theta)$ coincides with the log expected returns. While the reviewer appreciates the similarity of the $L_\\gamma(\\theta)$ with $L(\\theta)$, the new bound relies on (a) non-markovian dynamics (b) **time-varying** and **parameterized** discount factor, limiting the justification of the method by this theoretical result.\n- With the above in mind, it is possible to argue that even a standard MLE-based MBRL algorithm optimizes a lower bound on the expected returns (yet this bound will not be tight) using the Simulation Lemma [3]. See section 2 in [4] for an explanation. Since the bound that the practical algorithm optimizes is not tight, it is unclear to which extent the proposed method mitigates the objective mismatch.\n- Lastly, the analyses make several assumptions that might not be entirely justified. The most noticeable one is the assumption about the positive reward. While it is common to assume that the rewards are *non-negative*, the assumption about positive rewards might be limiting. Moreover, the rewards in certain environments (e.g. `HalfCheetah-v2`) where MnM is tested take negative values.\n\nDetailed notes and questions\n\nThere are several smaller concerns and questions that the authors might address:\n1. The literature review could be improved. In particular, relevant references include [1] that characterize the set of models that are optimal for planning; [5] that optimize the policy and the model using the same objective; [6] that learn the model to directly optimize agent’s performance, not a lower bound on it; [7] that learn the model that is helpful for policy improvement using a weighting scheme.\n2. Since $\\Gamma_\\theta(t)$ is a CDF, does it imply that $\\gamma(t)$ is a distribution?\n3. “One detail of note is that we omit the reward augmentation for MnM during these experiments, as it hinders exploration leading to lower returns.” Is it the entropy augmentation of SAC or the proposed augmentation with the log density ratio?\n4. Is the reward function learned as well?\n\nReferences\n\n[1] Grimm, Christopher, André Barreto, Satinder Singh, and David Silver. \"The value equivalence principle for model-based reinforcement learning.\" arXiv preprint arXiv:2011.03506 (2020).\n\n[2] Henderson, Peter, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. \"Deep reinforcement learning that matters.\" In Proceedings of the AAAI conference on artificial intelligence, vol. 32, no. 1. 2018.\n\n[3] Kearns, Michael, and Satinder Singh. \"Near-optimal reinforcement learning in polynomial time.\" Machine learning 49, no. 2 (2002): 209-232.\n\n[4] Farahmand, Amir-massoud, Andre Barreto, and Daniel Nikovski. \"Value-aware loss function for model-based reinforcement learning.\" In Artificial Intelligence and Statistics, pp. 1486-1494. PMLR, 2017.\n\n[5] Schrittwieser, Julian, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez et al. \"Mastering atari, go, chess and shogi by planning with a learned model.\" Nature 588, no. 7839 (2020): 604-609.\n\n[6] Nikishin, Evgenii, Romina Abachi, Rishabh Agarwal, and Pierre-Luc Bacon. \"Control-Oriented Model-Based Reinforcement Learning with Implicit Differentiation.\" arXiv preprint arXiv:2106.03273 (2021).\n\n[7] D'Oro, Pierluca, Alberto Maria Metelli, Andrea Tirinzoni, Matteo Papini, and Marcello Restelli. \"Gradient-aware model-based policy search.\" In Proceedings of the AAAI Conference on Artificial Intelligence, vol. 34, no. 04, pp. 3801-3808. 2020.\n\n", "summary_of_the_paper": "The paper studies the model learning aspect in model-based reinforcement learning (MBRL). Standard Dyna-like MBRL approaches that train the model using the maximum likelihood estimation (MLE) or its variations. In contrast, the authors propose an algorithm called Mismatched no More (MnM) that optimizes model parameters using a lower bound on the agent’s performance. By doing so, the paper mitigates the known objective mismatch of standard MBRL: MLE-optimized model might not necessarily be useful for optimizing agent’s returns, the overall objective of an RL system. Experimentally, the paper presents two studies. First, in 2 tabular environments, the authors compare MnM to Q-learning and VMBPO, an alternative MBRL algorithm that addresses the objective mismatch. Second, the authors compare MnM to an MLE-like MBRL baseline and a model-free algorithm on 12 continuous control tasks.\n", "summary_of_the_review": "The reviewer recommends rejecting the paper. While the idea is promising, the submission in its current state needs substantial improvements. Addressing the outlined concerns might increase the overall score.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers agree that the proposed method of joint model-policy optimization using a lower bound is novel and interesting and worthwhile pursuing. But all reviewers find a variety of issues in the paper, such that ratings are just above borderline or below. Given all the mixed feedback, it appears that the paper is still a bit premature for publication and could greatly benefit from improvements in a future submission.
Decision: Reject