Original: Joint Model-Policy Optimization of a Lower Bound for Model-Based RL
Revised: Joint Model-Policy Optimization of a Lower Bound for Model-Based RL

Original: Abstract
Revised: Abstract

Original: Many model-based reinforcement learning (RL) methods follow a similar template:ﬁt a model to previously observed data, and then use data from that model for RLor planning.
Revised: Many model-based reinforcement learning (RL) methods follow a similar template: ﬁt a model to previously observed data, and then use data from that model for RL or planning.

Original: However, models that achieve better training performance (e.g., lower
Revised: However, models that achieve better training performance (e.g., lower

Original: MSE) are not necessarily better for control: an RL agent may seek out the smallfraction of states where an accurate model makes mistakes, or it might act in waysthat do not expose the errors of an inaccurate model.
Revised: 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.

Original: As noted in prior work, thereis an objective mismatch: models are useful if they yield good policies, but they aretrained to maximize their accuracy, rather than the performance of the policies thatresult from them.
Revised: 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.

Original: In this work, we propose a single objective for jointly training themodel and the policy, such that updates to either component increase a lower boundon expected return.
Revised: In this work, we propose a single objective for jointly training the model and the policy, such that updates to either component increase a lower bound on expected return.

Original: To the best of our knowledge, this is the ﬁrst lower bound formodel-based RL that holds globally and can be efﬁciently estimated in continuoussettings; it is the only lower bound that mends the objective mismatch problem.
Revised: To the best of our knowledge, this is the ﬁrst lower bound for model-based RL that holds globally and can be efﬁciently estimated in continuous settings; it is the only lower bound that mends the objective mismatch problem.

Original: Aversion of this bound becomes tight under certain assumptions.
Revised: A version of this bound becomes tight under certain assumptions.

Original: Optimizing thisbound resembles a GAN: a classiﬁer distinguishes between real and fake transitions,the model is updated to produce transitions that look realistic, and the policy isupdated to avoid states where the model predictions are unrealistic.
Revised: Optimizing this bound resembles a GAN: a classiﬁer distinguishes between real and fake transitions, the model is updated to produce transitions that look realistic, and the policy is updated to avoid states where the model predictions are unrealistic.

Original: Numericalsimulations demonstrate that optimizing this bound yields reward maximizingpolicies and yields dynamics that (perhaps surprisingly) can aid in exploration.
Revised: Numerical simulations demonstrate that optimizing this bound yields reward maximizing policies and yields dynamics that (perhaps surprisingly) can aid in exploration.

Original: Wealso show that a deep RL algorithm loosely based on our lower bound can achieveperformance competitive with prior model-based methods, and better performanceon certain hard exploration tasks.
Revised: We also show that a deep RL algorithm loosely based on our lower bound can achieve performance competitive with prior model-based methods, and better performance on certain hard exploration tasks.

Original:  Introduction
Revised:  Introduction

Original: Much of the appeal of model-based RL is that model learning is a simple and scalable supervisedlearning problem.
Revised: Much of the appeal of model-based RL is that model learning is a simple and scalable supervised learning problem.

Original: However, even after learning a very accurate model, it is hard to say whether thatmodel will actually be useful for model-based RL [17, 29].
Revised: However, even after learning a very accurate model, it is hard to say whether that model will actually be useful for model-based RL [17, 29].

Original: For example, a model might make smallmistakes in critical states that cause a policy to take suboptimal actions.
Revised: For example, a model might make small mistakes in critical states that cause a policy to take suboptimal actions.

Original: Alternatively, a model withlarge errors may yield a policy that attains high return if the model errors occur in states that thepolicy never visits.
Revised: 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.

Original: The underlying problem is that dynamics models are trained differently from how they are used.
Revised: The underlying problem is that dynamics models are trained differently from how they are used.

Original: Typical model-based methods train a model using data sampled from the real dynamics (e.g., usingmaximum likelihood), but apply these models by using data sampled from the learned dynamics [11,23, 24, 48].
Revised: Typical model-based methods train a model using data sampled from the real dynamics (e.g., using maximum likelihood), but apply these models by using data sampled from the learned dynamics [11, 23, 24, 48].

Original: Prior work has identiﬁed this objective mismatch issue [17, 29, 30]: the model is trainedusing one objective, but the policy is trained using a different objective.
Revised: Prior work has identiﬁed this objective mismatch issue [17, 29, 30]: the model is trained using one objective, but the policy is trained using a different objective.

Original: Designing an objectivefor model training that is guaranteed to improve the expected reward remains an open problem.
Revised: Designing an objective for model training that is guaranteed to improve the expected reward remains an open problem.

Original: So, how should we train a dynamics model so that it produces high-return policies when used formodel-based RL?
Revised: So, how should we train a dynamics model so that it produces high-return policies when used for model-based RL?

Original: The key idea in this paper is to view model-based RL as a latent-variable problem: the latent variableis the trajectory and the cumulative reward is interpreted as the probability that the trajectory solvesthe task.
Revised: The key idea in this paper is to view model-based RL as a latent-variable problem: the latent variable is the trajectory and the cumulative reward is interpreted as the probability that the trajectory solves the task.

Original: Inferring the latent variable corresponds to learning a dynamics model.
Revised: Inferring the latent variable corresponds to learning a dynamics model.

Original: Latent variablemodels are typically learned via an evidence lower bound, and we show how a similar evidencelower bound provides yields a new objective for model-based RL.
Revised: Latent variable models are typically learned via an evidence lower bound, and we show how a similar evidence lower bound provides yields a new objective for model-based RL.

Original: In the same way that the evidencelower bound is a joint optimization problem over two variables, our objective will jointly optimizethe model and the policy using the same objective: to produce realistic and high-return trajectories.
Revised: In the same way that the evidence lower bound is a joint optimization problem over two variables, our objective will jointly optimize the model and the policy using the same objective: to produce realistic and high-return trajectories.

Original: Our objective differs from standard model-based RL objectives, where it is more common to pit themodel against the policy [5, 33, 37]. 
Revised: Our objective differs from standard model-based RL objectives, where it is more common to pit the model against the policy [5, 33, 37]. A 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.

Original: The main contribution of this work is an objective for model-based RL.
Revised: The main contribution of this work is an objective for model-based RL.

Original: To the best of our knowledge,this is the ﬁrst lower bound for model-based RL that holds globally (unlike Luo et al.
Revised: To the best of our knowledge, this is the ﬁrst lower bound for model-based RL that holds globally (unlike Luo et al.

Original: [30]) and canbe efﬁciently estimated in continuous settings (unlike Kearns and Singh [26]).
Revised: [30]) and can be efﬁciently estimated in continuous settings (unlike Kearns and Singh [26]).

Original: It is the ﬁrst lowerbound that jointly optimizes the model and policy using the same objective.
Revised: It is the ﬁrst lower bound that jointly optimizes the model and policy using the same objective.

Original: We also present a morecomplex version of this bound that becomes tight under some assumptions.
Revised: We also present a more complex version of this bound that becomes tight under some assumptions.

Original: Through numericalsimulations in simple tasks, we demonstrate that optimizing the bound yields reward-maximizingpolicies and yields an optimistic dynamics model that can aid exploration.
Revised: Through numerical simulations in simple tasks, we demonstrate that optimizing the bound yields reward-maximizing policies and yields an optimistic dynamics model that can aid exploration.

Original: We also demonstrate thatour bound gracefully accounts for function approximation error in the model – the policy is penalizedfor taking transitions that the model cannot represent.
Revised: We also demonstrate that our bound gracefully accounts for function approximation error in the model – the policy is penalized for taking transitions that the model cannot represent.

Original: Finally, we show that we can use parts of ourtheoretically-motivated objective to design a practically-applicable deep RL method that, despitedeviating from the theory, can match the performance of prior model-based methods.
Revised: Finally, we show that we can use parts of our theoretically-motivated objective to design a practically-applicable deep RL method that, despite deviating from the theory, can match the performance of prior model-based methods.

Original:  Related Work
Revised:  Related Work

Original: Most model-based RL methods use maximum likelihood to ﬁt the dynamics model and then use RLto maximize the expected return under samples from that model [10, 11, 23, 24, 48].
Revised: Most model-based RL methods use maximum likelihood to ﬁt the dynamics model and then use RL to maximize the expected return under samples from that model [10, 11, 23, 24, 48].

Original: As noted inprior work, this maximum likelihood objective is not aligned with the RL objective: models thatachieve higher likelihood do not necessarily produce better policies [17, 29, 30, 45, 52].
Revised: As noted in prior work, this maximum likelihood objective is not aligned with the RL objective: models that achieve higher likelihood do not necessarily produce better policies [17, 29, 30, 45, 52].

Original: This issue isreferred to as the objective mismatch problem: the model and policy (or planner) are optimized usingdifferent objectives.
Revised: This issue is referred to as the objective mismatch problem: the model and policy (or planner) are optimized using different objectives.

Original: This problem arises in almost all model-based RL approaches, including thosethat train the model to predict the value function [34, 41] or that perform planning [10, 41].
Revised: This problem arises in almost all model-based RL approaches, including those that train the model to predict the value function [34, 41] or that perform planning [10, 41].

Original: Some prior work addresses this problem by decreasing the rewards at states where the model isinaccurate [27, 30, 42, 49, 51], a strategy that our method will also employ.
Revised: Some prior work addresses this problem by decreasing the rewards at states where the model is inaccurate [27, 30, 42, 49, 51], a strategy that our method will also employ through a discriminator. While some of these methods also use discriminators in this manner, ours is the ﬁrst to optimize a lower bound on returns.

Original: Other work modiﬁesthe model objective to include multi-step rollouts [3, 4, 17, 25, 45, 47].
Revised: Other work modiﬁes the model objective to include multi-step rollouts [3, 4, 17, 25, 45, 47].

Original: Our method will modifythe model objective in a different way, so that the model objective is exactly the same as the policyobjective.
Revised: Our method will modify the model objective in a different way, so that the model objective is exactly the same as the policy objective.

Original: Some prior work directly optimizes the model to produce good policies [2, 13, 32, 35, 43],as theoretically analyzed in Grimm et al.
Revised: Some prior work directly optimizes the model to produce good policies [2, 13, 32, 35, 43], as theoretically analyzed in Grimm et al.

Original: While our aim is the same as these prior methods, ourapproach will not require differentiating through unrolled model updates or optimization procedures. 
Revised: While our aim is the same as these prior methods, our approach will not require differentiating through unrolled model updates or optimization procedures. [Table]

Original: Our work builds on prior lower bounds for model-based RL.
Revised: Our work builds on prior lower bounds for model-based RL.

Original: Kearns and Singh [26] provide a lowerbound that holds globally, but that is only computable in tabular settings.
Revised: Kearns and Singh [26] provide a lower bound that holds globally, but that is only computable in tabular settings.

Original: Luo et al.
Revised: Luo et al.

Original: [30] provide alower bound that can be efﬁciently estimated, but which only holds for nearby policies and models.
Revised: [30] provide a lower bound that can be efﬁciently estimated, but which only holds for nearby policies and models.

Original: As shown in Table 1, our bound combines the strengths of these prior works, providing a lower boundthat holds globally and can be efﬁciently estimated in MDPs with continuous states and actions.
Revised: As shown in Table 1, our bound combines the strengths of these prior works, providing a lower bound that holds globally and can be efﬁciently estimated in MDPs with continuous states and actions.

Original: Unlike prior work, our bound also mends the objective mismatch problem.
Revised: Unlike prior work, our bound also mends the objective mismatch problem.

Original: [Table] Our theoretical derivationbuilds on prior work thatcasts model-based RL as atwo-player game between amodel-player and a policyplayer [5, 33, 36, 38].
Revised: Our theoretical derivation builds on prior work that casts model-based RL as a two-player game between a model-player and a policy-player [5, 33, 36, 38].

Original: Whereasprior work pits model andpolicy against one another, ourformulation will result in acooperative game: the modeland policy cooperate to optimize the same objective (a lower bound on the expected return).
Revised: Whereas prior work pits model and policy against one another, our formulation will result in a cooperative game: the model and policy cooperate to optimize the same objective (a lower bound on the expected return).

Original: Ourapproach, though structurally resembling a GAN, is different from prior work that replaces amaximum likelihood model with a GAN model [6, 8, 28].
Revised: Our approach, though structurally resembling a GAN, is different from prior work that replaces a maximum likelihood model with a GAN model [6, 8, 28].

Original: The most similar prior work is VMBPO [9], which also jointly optimizes the model and the policyusing the same objective.
Revised: The most similar prior work is VMBPO [9], which also jointly optimizes the model and the policy using the same objective.

Original: However, while our objective is a lower bound on expected return, VMBPOmaximizes a different, risk-seeking objective, which is an upper bound on expected return (see
Revised: However, while our objective is a lower bound on expected return, VMBPO maximizes a different, risk-seeking objective, which is an upper bound on expected return (see

Original: Appendix B.1).
Revised: Appendix B.1).

Revised: 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 B.1).

Original: Indeed, while most of the components of our method (e.g., classiﬁers, GAN-likemodels) have been used in prior work, our paper is the ﬁrst to provide a  recipe for combining thesecomponents into an objective that is a provable lower bound.
Revised: Indeed, while most of the components of our method (e.g., classiﬁers, GAN-like models) have been used in prior work, our paper is the ﬁrst to provide a precise recipe for combining these components into an objective that is a provable lower bound.

Original:  A Uniﬁed Objective for Model-Based RL
Revised:  A Uniﬁed Objective for Model-Based RL

Original: Notation.
Revised: Notation.

Original: We focus on the Markov decision process with states s t , actions a t , initial state distribution p 0 ( s 0 ) , positive reward function r ( s t , a t ) > 0 , and dynamics p ( s t +1 | s t , a t ) .
Revised: We focus on the Markov decision process with states s t , actions a t , initial state distribution p 0 ( s 0 ) , positive reward function r ( s t , a t ) > 0 , and dynamics p ( s t +1 | s t , a t ) .

Original: Our aim is tolearn a control policy π θ ( a t | s t ) with parameters θ that maximizes the expected discounted return:
Revised: Our aim is to learn a control policy π θ ( a t | s t ) with parameters θ that maximizes the expected discounted return:

Original: [Equation]
Revised: [Equation]

Original: We use transitions ( s t , a t , r t , s t +1 ) collected from the (real) environment to train the dynamics modelq θ ( s t +1 | s t , a t ) , and use transitions sampled from this learned model to train the policy.
Revised: We use transitions ( s t , a t , r t , s t +1 ) collected from the (real) environment to train the dynamics model q θ ( s t +1 | s t , a t ) , and use transitions sampled from this learned model to train the policy.

Original: To simplifynotation, we will deﬁne a trajectory τ (cid:44) ( s 0 , a 0 , s 1 , a 1 , · · · ) as a sequence of states and actionsvisited in an episode.
Revised: To simplify notation, we will deﬁne a trajectory τ (cid:44) ( s 0 , a 0 , s 1 , a 1 , · · · ) as a sequence of states and actions visited in an episode.

Original: We deﬁne R ( τ ) (cid:44) (cid:80) ∞
Revised: We deﬁne R ( τ ) (cid:44) (cid:80) ∞

Original: t
Revised: t

Original: =0 γ t r ( s t , a t ) as the discounted return of a trajectory.
Revised: =0 γ t r ( s t , a t ) as the discounted return of a trajectory.

Original: We deﬁne two distributions over trajectories: p π ( τ ) and q π ( τ ) for when policy π θ interacts withdynamics p ( s t +1 | s t , a t ) and q θ ( s t +1 | s t , a t ) , respectively:p π ( τ ) = p 0 ( s 0 ) ∞ (cid:89) t = p ( s t +1 | s t , a t )
Revised: We deﬁne two distributions over trajectories: p π ( τ ) and q π ( τ ) for when policy π θ interacts with dynamics p ( s t +1 | s t , a t ) and q θ ( s t +1 | s t , a t ) , respectively:

Original: π θ ( a t | s t ) , q π ( τ ) = p 0 ( s 0 ) ∞ (cid:89) t = q θ ( s t +1 | s t , a t ) π θ ( a t | s t ) .
Revised: [Equation]

Original: Desiderata.
Revised: Desiderata.

Original: Our aim is to design an objective L ( θ ) with two properties.
Revised: Our aim is to design an objective L ( θ ) with two properties.

Original: First , this objective shouldbe a lower bound on the expected return in the true environment: if the policy does well in the learnedmodel, we are guaranteed that the policy will also do well in the true environment.
Revised: First , this objective should be a lower bound on the expected return in the true environment: if the policy does well in the learned model, we are guaranteed that the policy will also do well in the true environment.

Original: The expectedreturn under the learned model, which most prior model-based RL methods use to train the policy, isnot a lower bound on the expected return [26, 30].
Revised: 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 [26, 30].

Original: While prior work has made strides in developinglower bounds for model-based RL, even the best lower bounds do not hold for all models [30] or arelimited to tabular settings [26].
Revised: While prior work has made strides in developing lower bounds for model-based RL, even the best lower bounds do not hold for all models [30] or are limited to tabular settings [26].

Original: Second , this objective should be the same for the policy and the dynamics model, such that updatesto the model would improve the policy, and vice versa.
Revised: Second , this objective should be the same for the policy and the dynamics model, such that updates to the model would improve the policy, and vice versa.

Original: This desiderata is important because priorwork has found that training the model to be more accurate (increase likelihood) can decrease thepolicy’s expected return under that model [17, 29, 30, 45, 52].
Revised: This desiderata is important because prior work has found that training the model to be more accurate (increase likelihood) can decrease the policy’s expected return under that model [17, 29, 30, 45, 52].

Original: One prior method (VMBPO [9]) doestrain the model and policy with the same objective, but this objective corresponds to a risk-seeking,upper-bound on the expected returns.
Revised: One prior method (VMBPO [9]) does train the model and policy with the same objective, but this objective corresponds to a risk-seeking, upper-bound on the expected returns.

Original: An objective for model-based RL.
Revised: An objective for model-based RL.

Original: We now introduce an objective that achieves these aims.
Revised: We now introduce an objective that achieves these aims.

Original: Thekey idea to deriving this result is to take a probabilistic perspective on decision making: we viewthe trajectory as an unobserved random variable, and the reward function as the probability that atrajectory solves a task.
Revised: The key idea to deriving this result is to take a probabilistic perspective on decision making: we view the trajectory as an unobserved random variable, and the reward function as the probability that a trajectory solves a task.

Original: Then, the problem of inferring this random variable is equivalent to learningthe dynamics model.
Revised: Then, the problem of inferring this random variable is equivalent to learning the dynamics model.

Original: Using an evidence lower bound, we obtain an objective for jointly training themodel and policy.
Revised: Using an evidence lower bound, we obtain an objective for jointly training the model and policy.

Original: We provide the full derivation in Appendix B.3.
Revised: We provide the full derivation in Appendix B.3.

Original: Our resulting objective is the policy’s reward when interacting with the learned model, but using anew reward function.
Revised: Our resulting objective is the policy’s reward when interacting with the learned model, but using a new reward function.

Original: The new reward function combines the task reward with a term that measuresthe difference between the learned model and the real environment.
Revised: The new reward function combines the task reward with a term that measures the difference between the learned model and the real environment.

Original: We deﬁne our objective
Revised: We deﬁne our objective

Original: [Equation]
Revised: [Equation]

Original: where the modiﬁed reward function is deﬁned as
Revised: where the modiﬁed reward function is deﬁned as

Original: [Equation]
Revised: [Equation]

Original: Intuitively, the new reward function ˜ r penalizes the policy for taking transitions that are unlikelyunder the true dynamics model, similar to prior work [15, 46, 49].
Revised: Intuitively, the new reward function ˜ r penalizes the policy for taking transitions that are unlikely under the true dynamics model, similar to prior work [15, 46, 49].

Original: Later, we will show that we canestimate this augmented reward without knowing the true environment dynamics by using a classiﬁer.
Revised: Later, we will show that we can estimate this augmented reward without knowing the true environment dynamics by using a classiﬁer.

Original: We will optimize this lower bound with respect to both the policy π θ ( a t | s t ) and the dynamics modelq θ ( s t +1 | s t , a t ) .
Revised: We will optimize this lower bound with respect to both the policy π θ ( a t | s t ) and the dynamics model q θ ( s t +1 | s t , a t ) .

Original: For the policy, we maximize the modiﬁed reward using samples from the learnedmodel; the only difference from prior work is the modiﬁcation to the reward function.
Revised: For the policy, we maximize the modiﬁed reward using samples from the learned model; the only difference from prior work is the modiﬁcation to the reward function.

Original: Training thedynamics model using this objective is very different from standard maximum likelihood training.
Revised: Training the dynamics model using this objective is very different from standard maximum likelihood training.

Original: The model is optimized to sample trajectories that are similar to real dynamics (like a GAN) and thathave high reward (unlike a GAN).
Revised: The model is optimized to sample trajectories that are similar to real dynamics (like a GAN) and that have high reward (unlike a GAN).

Original: This objective differs from VMBPO [9] by taking the log( · ) of theoriginal reward functions; our experiments demonstrate that excluding this component invalidates ourlower bound and results in learning suboptimal policies (Fig. 3a).
Revised: This objective differs from VMBPO [9] by taking the log( · ) of the original reward functions; our experiments demonstrate that excluding this component invalidates our lower bound and results in learning suboptimal policies (Fig. 3a).

Original: Our objective achieves the two desiderata.
Revised: Our objective achieves the two desiderata.

Original: First , our objective is a lower bound on the expectedreturn.
Revised: First , our objective is a lower bound on the expected return.

Original: To state this result formally, we will take the logarithm of the expected return.
Revised: To state this result formally, we will take the logarithm of the expected return.

Original: Of course,maximizing the log( · ) of the expected return is equivalent to maximizing the expected return.
Revised: Of course, maximizing the log( · ) of the expected return is equivalent to maximizing the expected return.

Original: Theorem 3.1.
Revised: Theorem 3.1.

Original: The following bound holds for any dynamics q ( s t +1 | s t , a t ) and policy π ( a t | s t ) : 
Revised: The following bound holds for any dynamics q ( s t +1 | s t , a t ) and policy π ( a t | s t ) : log E [Equation] (cid:21)

Original: [Equation]
Revised: [Equation]

Original: The proof is presented in Appendix B.3.
Revised: The proof is presented in Appendix B.3.

Original: To the best of our knowledge, this is the ﬁrst global(unlike Luo et al.
Revised: To the best of our knowledge, this is the ﬁrst global (unlike Luo et al.

Original: [30]) and efﬁciently-computable (unlike Kearns and Singh [26]) lower bound formodel-based RL.
Revised: [30]) and efﬁciently-computable (unlike Kearns and Singh [26]) lower bound for model-based RL.

Original: The model and the policy are trained using the same objective: updating the modelnot only increases the objective for the model, but also increases the objective for the policy.
Revised: The model and the policy are trained using the same objective: updating the model not only increases the objective for the model, but also increases the objective for the policy.

Original: Sec. 4 will introduce an algorithm to maximize this lower bound.
Revised: Sec. 4 will introduce an algorithm to maximize this lower bound.

Original: While this lower bound may not betight, experiments in Sec.
Revised: While this lower bound may not be tight, experiments in Sec.

Original: 5 demonstrate that optimizing this lower bound yields policies that achievehigh reward across a wide range of tasks.
Revised: 5 demonstrate that optimizing this lower bound yields policies that achieve high reward across a wide range of tasks.

Original: In Appendix A, we propose a variant of this bound thatdoes become tight at convergence.
Revised: In Appendix A, we propose a variant of this bound that does become tight at convergence.

Original: Because this tight bound is more complex, we focus on the simplebound in this paper.
Revised: Because this tight bound is more complex, we focus on the simple bound in this paper.

Original: The optimal dynamics are optimistic.
Revised: The optimal dynamics are optimistic.

Original: We now return to analyzing the simpler lower bound ( L ( θ )in Eq. 2).
Revised: We now return to analyzing the simpler lower bound ( L ( θ ) in Eq. 2).

Original: In stochastic environments, the dynamics model that optimizes this lower bound is notequal to the true environment dynamics.
Revised: In stochastic environments, the dynamics model that optimizes this lower bound is not equal to the true environment dynamics.

Original: Rather, it is biased towards sampling trajectories withhigh return.
Revised: Rather, it is biased towards sampling trajectories with high return.

Original: Ignoring parametrization constraints, the dynamics model that optimizes our lowerbound is q ∗ ( τ ) =
Revised: Ignoring parametrization constraints, the dynamics model that optimizes our lower bound is q ∗ ( τ ) =

Original: [Equation]
Revised: [Equation]

Original: (proof in Appendix B.4.).
Revised: (proof in Appendix B.4.).

Original: While it may seem surprising that theobjective-optimizing dynamics would differ from the true dynamics, this result is analogous to a
Revised: While it may seem surprising that the objective-optimizing dynamics would differ from the true dynamics, this result is analogous to a

Original: VAE, where the ELBO-optimizing encoder differs from the prior.
Revised: VAE, where the ELBO-optimizing encoder differs from the prior.

Original: The optimism in the dynamicsmodel may accelerate policy optimization, a hypothesis we will test in Sec. 5.1.
Revised: The optimism in the dynamics model may accelerate policy optimization, a hypothesis we will test in Sec. 5.1.

Original: Would the optimistic dynamics overestimate the policy’s return, violating Theorem 3.1?
Revised: Would the optimistic dynamics overestimate the policy’s return, violating Theorem 3.1?

Original: While usingthe optimistic dynamics with the original reward function will overestimate the true return, using theoptimistic dynamics with the augmented reward function yields a valid lower bound.
Revised: While using the optimistic dynamics with the original reward function will overestimate the true return, using the optimistic dynamics with the augmented reward function yields a valid lower bound.

Original: We demonstratethis effect in Fig.
Revised: We demonstrate this effect in Fig.

Original: 3b.
Revised: 3b.

Original:  Practical Optimization of the Lower Bound
Revised:  Practical Optimization of the Lower Bound

Original: The previous section presented a single (global) lower bound ( L from Eq. 2) for jointly optimizingthe policy and the dynamics model.
Revised: The previous section presented a single (global) lower bound ( L from Eq. 2) for jointly optimizing the policy and the dynamics model.

Original: In this section, we develop a practical algorithm for optimizingthis lower bound.
Revised: In this section, we develop a practical algorithm for optimizing this lower bound.

Original: The main challenge in optimizing this bound is that the augmented reward functiondepends on the transition probabilities of the real environment, p ( s t +1 | s t , a t ) , which are unknown.
Revised: The main challenge in optimizing this bound is that the augmented reward function depends on the transition probabilities of the real environment, p ( s t +1 | s t , a t ) , which are unknown.

Original: We address this challenge by learning a classiﬁer (Sec. 4.1).
Revised: We address this challenge by learning a classiﬁer (Sec. 4.1).

Original: Of course, Theorem 3.1 only holds if theclassiﬁer is Bayes-optimal.
Revised: Of course, Theorem 3.1 only holds if the classiﬁer is Bayes-optimal.

Original: We then describe the precise update rules for the policy, dynamics model,and classiﬁer (Sec. 4.2).
Revised: We then describe the precise update rules for the policy, dynamics model, and classiﬁer (Sec. 4.2).

Original: 4.1 Estimating the Augmented Reward Function
Revised: 4.1 Estimating the Augmented Reward Function

Original: To estimate the augmented reward function, which depends on the transition probabilities ofthe real environment, we learn a classiﬁer that distinguishes real transitions from fake transi- tions.
Revised: To estimate the augmented reward function, which depends on the transition probabilities of the real environment, we learn a classiﬁer that distinguishes real transitions from fake transitions.

Original: This approach is similar to GANs [19] and similar to prior work in RL [15, 49].
Revised: This approach is similar to GANs [19] and similar to prior work in RL [15, 49].

Original: We use
Revised: We use

Original: C φ ( s t , a t , s t +1 ) ∈ [0 , 1] to denote the classiﬁer, which we train to distinguish real versus modeltransitions using the standard cross entropy loss:
Revised: C φ ( s t , a t , s t +1 ) ∈ [0 , 1] to denote the classiﬁer, which we train to distinguish real versus model transitions using the standard cross entropy loss:

Original: [Equation]
Revised: [Equation]

Original: Note that the real transition ( s real t , a real t , s real t +1 ) and model transition ( s real t , a real t , s model t +1 ) have the sameinitial state and initial action.
Revised: Note that the real transition ( s real t , a real t , s real t +1 ) and model transition ( s real t , a real t , s model t +1 ) have the same initial state and initial action.

Original: Once trained, we can use the classiﬁer’s predictions to estimate theaugmented reward function:
Revised: Once trained, we can use the classiﬁer’s predictions to estimate the augmented reward function:

Original: [Equation]
Revised: [Equation]

Original: The approximation above reﬂects function approximation error in learning the classiﬁer.
Revised: The approximation above reﬂects function approximation error in learning the classiﬁer.

Original: Now thatwe can estimate the augmented reward function, we can apply any RL method to maximize theaugmented reward under transitions sampled from the dynamics model.
Revised: Now that we can estimate the augmented reward function, we can apply any RL method to maximize the augmented reward under transitions sampled from the dynamics model.

Original: The following sectiondescribes a particular instantiation using an off-policy RL algorithm.
Revised: The following section describes a particular instantiation using an off-policy RL algorithm.

Original: 4.2 Updating the Model and Policy
Revised: 4.2 Updating the Model and Policy

Original: [Figure]
Revised: [Figure]

Original: We now present our complete method, whichtrains three components: a classiﬁer, a policy,and a dynamics model.
Revised: We now present our complete method, which trains three components: a classiﬁer, a policy, and a dynamics model.

Original: Our method alternatesbetween (1) updating the policy (by performing RL using model experience with augmentedrewards) and (2) updating the dynamics modeland classiﬁer (using a GAN-like objective).
Revised: Our method alternates between (1) updating the policy (by performing RL using model experience with augmented rewards) and (2) updating the dynamics model and classiﬁer (using a GAN-like objective).

Original: Indescribing the loss functions below, we use thesuperscripts ( · ) real and ( · ) model to denote transitions that have been sampled from the trueenvironment dynamics or the learned dynamicsfunction.
Revised: In describing the loss functions below, we use the superscripts ( · ) real and ( · ) model to denote transitions that have been sampled from the true environment dynamics or the learned dynamics function.

Original: To reduce clutter, we omit the superscripts when unambiguous.
Revised: To reduce clutter, we omit the superscripts when unambiguous.

Original: Updating the policy.
Revised: Updating the policy.

Original: The policy is optimized to maximize the augmented reward on transitionssampled from the learned dynamics model.
Revised: The policy is optimized to maximize the augmented reward on transitions sampled from the learned dynamics model.

Original: While this optimization can be done using any RLalgorithm, including on-policy methods, we will focus on an off-policy actor-critic method.
Revised: While this optimization can be done using any RL algorithm, including on-policy methods, we will focus on an off-policy actor-critic method.

Original: We deﬁne the Q function as sum of augmented rewards under the learned dynamics model:
Revised: We deﬁne the Q function as sum of augmented rewards under the learned dynamics model:

Original: [Equation]
Revised: [Equation]

Original: We approximate the Q function using a neural network Q ψ ( s t , a t ) with parameters φ .
Revised: We approximate the Q function using a neural network Q ψ ( s t , a t ) with parameters φ .

Original: We train the Q206 function using the TD loss on transitions sampled from the learned dynamics model:
Revised: We train the Q function using the TD loss on transitions sampled from the learned dynamics model:

Original: [Equation]
Revised: [Equation]

Original: 7) where (cid:98)·(cid:99) sg is the stop-gradient operator and y t =
Revised: where (cid:98)·(cid:99) sg is the stop-gradient operator and y [Equation]

Original: ˜ r ( s t , a t , s model
Revised: =

Original: [Equation] 
Revised: [Equation] [Equation]

Original: ) (cid:3)is the TD target.
Revised: is the TD target.

Original: The augmented reward ˜ r is estimated using the learned classiﬁer (Eq. 5).
Revised: The augmented reward ˜ r is estimated using the learned classiﬁer (Eq. 5).

Original: To estimate209 the corresponding value function, we use a 1-sample approximation: V ψ ( s t ) =
Revised: To estimate the corresponding value function, we use a 1-sample approximation: V ψ ( s t ) =

Original: Q ψ ( s t , a t ) wherea t ∼ π θ ( a t | s t )) .
Revised: Q ψ ( s t , a t ) where a t ∼ π θ ( a t | s t )) .

Original: The policy is trained to maximize the Q function:
Revised: The policy is trained to maximize the Q function:

Original: [Equation]
Revised: [Equation]

Original: In our implementation, we regularize the policy by adding an additional entropy regularizer.
Revised: In our implementation, we regularize the policy by adding an additional entropy regularizer.

Original: Followingprior work [18], we maintain two Q functions and two target Q functions, using the minimum of thetwo to compute the TD target.
Revised: Following prior work [18], we maintain two Q functions and two target Q functions, using the minimum of the two to compute the TD target.

Original: See Appendix D for details.
Revised: See Appendix D for details.

Original: Updating the dynamics model.
Revised: Updating the dynamics model.

Original: To optimize the dynamics model, we rewrite the lower bound in215 terms of a single transition (derivation in Appendix B.6):
Revised: To optimize the dynamics model, we rewrite the lower bound in terms of a single transition (derivation in Appendix B.6):

Original: [Equation]
Revised: [Equation]

Original: This loss is an approximation of our original lower bound (Eq. 2) because we estimate the differencein dynamics using a classiﬁer.
Revised: This loss is an approximation of our original lower bound (Eq. 2) because we estimate the difference in dynamics using a classiﬁer.

Original: This approximation is standard in prior work on GANs [19] andadversarial inference [12, 14].
Revised: This approximation is standard in prior work on GANs [19] and adversarial inference [12, 14].

Original: Intuitively, the procedure for optimizing the dynamics model and the classiﬁer resembles a GAN [19]:the classiﬁer is optimized to distinguish real transitions from model transitions, and the model isupdated to fool the classiﬁer (and increase rewards).
Revised: Intuitively, the procedure for optimizing the dynamics model and the classiﬁer resembles a GAN [19]: the classiﬁer is optimized to distinguish real transitions from model transitions, and the model is updated to fool the classiﬁer (and increase rewards).

Original: However, our method is not equivalent to simplyreplacing a maximum likelihood model with a GAN model .
Revised: However, our method is not equivalent to simply replacing a maximum likelihood model with a GAN model .

Original: Indeed, such an approach would notoptimize a lower bound on expected return.
Revised: Indeed, such an approach would not optimize a lower bound on expected return.

Original: Rather, our model objective includes an additional valueterm and our policy objective includes an additional classiﬁer term.
Revised: Rather, our model objective includes an additional value term and our policy objective includes an additional classiﬁer term.

Original: These changes enable the modeland policy to optimize the same objective, which is a lower bound on expected return.
Revised: These changes enable the model and policy to optimize the same objective, which is a lower bound on expected return.

Original: Algorithm summary.
Revised: Algorithm summary.

Original: We summarize the method in Alg. 1 and provide an illustration in Fig.
Revised: We summarize the method in Alg. 1 and provide an illustration in Fig.

Original: We call the method M ISMATCHED NO M ORE (MnM) because the policy and model optimize thesame objective, thereby resolving the objective mismatch problem noted in prior work.
Revised: We call the method M ISMATCHED NO M ORE (MnM) because the policy and model optimize the same objective, thereby resolving the objective mismatch problem noted in prior work.

Original: While themodel and policy are optimized using the same objective, that objective can stop being a lower boundif the learned classiﬁer is not Bayes-optimal.
Revised: While the model and policy are optimized using the same objective, that objective can stop being a lower bound if the learned classiﬁer is not Bayes-optimal.

Original: Implementing MnM on top of a standard model-based RL algorithm is straightforward.
Revised: Implementing MnM on top of a standard model-based RL algorithm is straightforward.

Original: First, createan additional classiﬁer network.
Revised: First, create an additional classiﬁer network.

Original: Second, instead of using the maximum likelihood objective to trainthe model, use the GAN-like objective in Eq.
Revised: Second, instead of using the maximum likelihood objective to train the model, use the GAN-like objective in Eq.

Original: 9 to update both the model and the classiﬁer.
Revised: 9 to update both the model and the classiﬁer.

Original: Third, addthe classiﬁer’s logits to the predicted rewards (Eq. 5).
Revised: Third, add the classiﬁer’s logits to the predicted rewards (Eq. 5).

Original: Following prior work [24], we learn a neuralnetwork to predict the true environment rewards.
Revised: Following prior work [24], we learn a neural network to predict the true environment rewards.

Original: 4.3 A Note about Exploration
Revised: 4.3 A Note about Exploration

Original: The classiﬁer term in the augmented reward (Eq. 3) is a double edged sword.
Revised: The classiﬁer term in the augmented reward (Eq. 3) is a double edged sword.

Original: Our theoretical derivationsuggests that this term should appear, and our didactic experiments demonstrate that removing thisterm results in suboptimal behavior.
Revised: Our theoretical derivation suggests that this term should appear, and our didactic experiments demonstrate that removing this term results in suboptimal behavior.

Original: Experiments also show that this term can effectively combaterrors in the learned model.
Revised: Experiments also show that this term can effectively combat errors in the learned model.

Original: Prior work in ofﬂine RL has found similar model-error reward termscritical for achieving good performance in the ofﬂine setting [27, 42, 49, 51].
Revised: Prior work in ofﬂine RL has found similar model-error reward terms critical for achieving good performance in the ofﬂine setting [27, 42, 49, 51].

Original: However, when scaling MnM to continuous control tasks in the online setting, we found thatincluding this term hurts performance.
Revised: However, when scaling MnM to continuous control tasks in the online setting, we found that including this term hurts performance (Fig. 10).

Original: This makes sense: this classiﬁer term is exactly the oppositeof exploration objectives based on model error (e.g., [44]), and penalizes the policy for performingexploration.
Revised: This makes sense: this classiﬁer term is exactly the opposite [Figure] of exploration objectives based on model error (e.g., [44]) , and penalizes the policy for performing exploration.

Original: Our continuous control experiments in Sec.
Revised: Our continuous control experiments in Sec.

Original: 5.2 will deviate from our theoreticalderivation: they will use the model objective suggested by our theory but the same policy objectiveas prior.
Revised: 5.2 will deviate from our theoretical derivation: they will use the model objective suggested by our theory, but the same policy objective as prior.

Original: That is, the policy maximizes r ( s t , a t ) instead of ˜ r ( s t , a t ) (Eq. 3)
Revised: That is, the policy maximizes r ( s t , a t ) instead of ˜ r ( s t , a t ) (Eq. 3).

Original: We call this method“MnM-approx many theoretical model-based RL papers also ﬁnd it necessary toimplement practical algorithms that differ from the theory [30, 46, [Figure] [Figure]
Revised: We call this method “MnM-approx.” We note that many theoretical model-based RL papers also ﬁnd it necessary to implement practical algorithms that differ from the theory [30, 46, 49].

Original: [Figure]

Original:  Numerical Simulations
Revised:  Numerical Simulations

Original: The primary aim of our experiments is to verify that our objective is a valid lower bound and thatmaximizing the bound yields reward-maximizing policies.
Revised: The primary aim of our experiments is to verify that our objective is a valid lower bound, and that maximizing the bound yields reward-maximizing policies.

Original: We study these questions in Sec. 5.1,where we will use tabular problems so that we can study the objective in the absence of functionapproximation error.
Revised: We study these questions in Sec. 5.1, where we will use tabular problems so that we can study the objective in the absence of function approximation error.

Original: Then, in Sec. 5.2, we show adapting part of our objective into a scalable modelbased RL algorithm yields a method that, while deviating from the theory, can achieve competitiveresults on benchmark tasks.
Revised: Then, in Sec. 5.2, we show adapting part of our objective into a scalable modelbased RL algorithm yields a method that, while deviating from the theory, can achieve competitive results on benchmark tasks.

Original: 5.1 Understanding the Lower Bound and the Learned Dynamics
Revised: 5.1 Understanding the Lower Bound and the Learned Dynamics

Original: Our ﬁrst experiments test whether optimizing the lower bound produces optimal policies and studyhow the components of the lower bound.
Revised: Our didactic experiments test whether optimizing the lower bound produces optimal policies and study how the components of the lower bound.

Original: We use tabular domains in this section as they allow us toanalytically compute the optimal policy for comparison and allow us to evaluate the lower bound inthe absence of function approximation error.
Revised: We use tabular domains in this section, as they allow us to analytically compute the optimal policy for comparison, and allow us to evaluate the lower bound in the absence of function approximation error.

Original: Our ﬁrst experiment compares the lower bound to an oracle method that applies Q-learning to aperfect dynamics model.
Revised: Our ﬁrst experiment compares the lower bound to an oracle method that applies Q-learning to a perfect dynamics model.

Original: In comparison to this baseline, MnM learns a dynamics model using the
Revised: In comparison to this baseline, MnM learns a dynamics model using the

Original: GAN-like objective (Eq. 9) and maximizes a modiﬁed reward function (Eq. 5).
Revised: GAN-like objective (Eq. 9) and maximizes a modiﬁed reward function (Eq. 5).

Original: We also compare to
Revised: We also compare to

Original: VMBPO, which learns a dynamics model similar to MnM but omits the log-transformation of thereward function; this transformation ordinarily encourages pessimistic behavior.
Revised: VMBPO, which learns a dynamics model similar to MnM but omits the log-transformation of the reward function; this transformation ordinarily encourages pessimistic behavior.

Original: For this experiment,we use a 10 × 10 gridworld with stochastic dynamics and sparse reward, shown in Fig.
Revised: For this experiment, we use a 10 × 10 gridworld with stochastic dynamics and sparse reward, shown in Fig.

Original: 2a (Left) .
Revised: 2a (Left) .

Original: The results, shown in Fig.
Revised: The results, shown in Fig.

Original: 2a, show that MnM outperforms both Q-learning with the correct modeland VMBPO on this task.
Revised: 2a, show that MnM outperforms both Q-learning with the correct model and VMBPO on this task.

Original: We hypothesize that VMBPO performs poorly on this task because itmaximizes an upper bound on performance; and conﬁrm this in the following experiments.
Revised: We hypothesize that VMBPO performs poorly on this task because it maximizes an upper bound on performance; and conﬁrm this in the following experiments.

Original: We hypothesize that MnM outperforms Q-learning with the correct model because it learns “optimistic”dynamics.
Revised: We hypothesize that MnM outperforms Q-learning with the correct model because it learns “optimistic” dynamics.

Original: We test this hypothesis by visualizing the dynamics model learned by MnM (Fig.
Revised: We test this hypothesis by visualizing the dynamics model learned by MnM (Fig.

Original: 2a (Left) ).
Revised: 2a (Left) ).

Original: While the true environment dynamics have stochasticity that moves the agent in a random directionwith equal probability, the MnM dynamics model biases this stochasticity to lead the agent towardsthe goal (blue arrows point towards the goal).
Revised: While the true environment dynamics have stochasticity that moves the agent in a random direction with equal probability, the MnM dynamics model biases this stochasticity to lead the agent towards the goal (blue arrows point towards the goal).

Original: Of course, we use the true environment dynamics, notthe optimistic dynamics model, for evaluating the policies.
Revised: Of course, we use the true environment dynamics, not the optimistic dynamics model, for evaluating the policies.

Original: Our augmented reward function contains two crucial components, (1) the classiﬁer term and (2) thelogarithmic transformation of the reward function.
Revised: Our augmented reward function contains two crucial components, (1) the classiﬁer term and (2) the logarithmic transformation of the reward function.

Original: We test the importance of the classiﬁer term incorrecting for inaccurate models.
Revised: We test the importance of the classiﬁer term in correcting for inaccurate models.

Original: To do this, we limit the capacity of the MnM dynamics model[Figure] so that it makes “low-resolution” predictions, forcing all states in 3 × 3 blocks to have the samedynamics.
Revised: To do this, we limit the capacity of the MnM dynamics model so that it makes “low-resolution” predictions, forcing all states in 3 × 3 blocks to have the same dynamics.

Original: We will use the gridworld shown in Fig.
Revised: We will use the gridworld shown in Fig.

Original: 2b (Left) , which contains obstacles that occurat a ﬁner resolution than the model can detect.
Revised: 2b (Left) , which contains obstacles that occur [Figure] at a ﬁner resolution than the model can detect.

Original: When the “low resolution” dynamics model makespredictions for states near the obstacle, it will average together some states with obstacles and somestates without obstacles.
Revised: 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.

Original: Thus, the model will (incorrectly) predict that the agent always has someprobability of moving in each direction, even if that direction is actually blocked by an obstacle.
Revised: 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.

Original: However, the classiﬁer (whose capacity we have also limited) detects that the dynamics model isinaccurate in these states, so the augmented reward is much lower at these states.
Revised: However, the classiﬁer (whose capacity we have also limited) detects that the dynamics model is inaccurate in these states, so the augmented reward is much lower at these states.

Original: Thus, MnM is ableto solve this task despite the inaccurate model; an ablation of MnM that removes the classiﬁer termattempts to navigate through the wall and fails to reach the goal.
Revised: Thus, MnM is able to solve this task despite the inaccurate model; an ablation of MnM that removes the classiﬁer term attempts to navigate through the wall and fails to reach the goal.

Original: Like MnM, VMBPO includes a classiﬁer term in the reward function but omits the logarithmictransformation, a difference we expect to cause VMBPO to prefer suboptimal, risk-seeking policies.
Revised: Like MnM, VMBPO includes a classiﬁer term in the reward function but omits the logarithmic transformation, a difference we expect to cause VMBPO to prefer suboptimal, risk-seeking policies.

Original: To test this hypothesis, we use the 3-state MDP in Fig.
Revised: To test this hypothesis, we use the 3-state MDP in Fig.

Original: 3a (Right) , where numbers indicate the rewardat each state.
Revised: 3a (Right) , where numbers indicate the reward at each state.

Original: While moving to the right state yields slightly higher rewards, “wind” knocks the agentout of this state with probability 50% so the reward-maximizing strategy is to move to the left state.
Revised: While moving to the right state yields slightly higher rewards, “wind” knocks the agent out of this state with probability 50% so the reward-maximizing strategy is to move to the left state.

Original: While MnM learns the reward-maximizing strategy, VMBPO learns a policy that goes to the rightstate and receives lower returns.
Revised: While MnM learns the reward-maximizing strategy, VMBPO learns a policy that goes to the right state and receives lower returns.

Original: Finally, we verify Theorem 3.1 by comparing the MnM objective to the true expected return.
Revised: Finally, we verify Theorem 3.1 by comparing the MnM objective to the true expected return.

Original: Wealso compare to the objective from VMBPO, which looks similar to the MnM objective but omitsthe logarithmic transformation; our theory predicts that the VMBPO objective will therefore be anupper bound on the expected return (see Appendix B.1).
Revised: We also compare to the objective from VMBPO, which looks similar to the MnM objective but omits the logarithmic transformation; our theory predicts that the VMBPO objective will therefore be an upper bound on the expected return (see Appendix B.1).

Original: We use the gridworld from Fig.
Revised: We use the gridworld from Fig.

Original: 2a and use aversion of MnM based on value iteration to avoid approximation error.
Revised: 2a and use a version of MnM based on value iteration to avoid approximation error.

Original: Plotting the MnM objective in
Revised: Plotting the MnM objective in

Original: Fig.
Revised: Fig.

Original: 3b (Right) , we observe that it is always a lower bound on the (log) expected return, as predictedby our theory.
Revised: 3b (Right) , we observe that it is always a lower bound on the (log) expected return, as predicted by our theory.

Original: Also as predicted by the theory, the VMBPO objective overestimates the expectedreturn, illustrating the importance of the logarithmic transformation. 
Revised: Also as predicted by the theory, the VMBPO objective overestimates the expected return, illustrating the importance of the logarithmic transformation. Of the oft-cited beneﬁts of model-based RL is that the learned dynamics model can be re-used to solve new tasks.

Original: 5.2 Comparisons On Higher-Dimensional Tasks
Revised: MnM presents a twist on that story, because MnM does not learn the true environment dynamics but rather learns optimistic dynamics (see Sec. 3). In Appendix C (Fig. 13), we examine MnM’s effectiveness at transferring dynamics to different tasks in the stochastic gridworld from Fig.

Original:  Our next experiments use continuous-control robotic tasks to study whether the new model-learningobjective suggested by our lower bound can be stably applied to higher-dimensional control tasks.
Revised: 2a. We ﬁnd that the (optimistic) dynamics learned by MnM do not slow learning on dissimilar tasks, but can accelerate learning of challenging, similar tasks. 5.2 Comparisons On Higher-Dimensional Tasks Our next experiments use continuous-control robotic tasks to study whether the new model-learning objective suggested by our lower bound can be stably applied to higher-dimensional control tasks.

Original: While the MnM-approx method tested in this section deviates from our theoretical derivation (see
Revised: While the MnM-approx method tested in this section deviates from our theoretical derivation (see

Original: Sec. 4.3), these experiments are nonetheless useful for providing preliminary evidence about whetherthe proposed model objective can be efﬁciently estimated and stably optimized.
Revised: Sec. 4.3), these experiments are nonetheless useful for providing preliminary evidence about whether the proposed model objective can be efﬁciently estimated and stably optimized.

Original: We use MBPO [24] as a baseline for model-based RL because it achieves state-of-the-art resultsand is a prototypical example of model-based RL algorithms that use maximum likelihood models.
Revised: We use MBPO [24] 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.

Original: Because our method differs from the baseline (MBPO) along only one dimension now (new modelupdate, same policy update), these experiments will directly test the utility of our proposed modelupdate.
Revised: Because our method differs from the baseline (MBPO) along only one dimension now (new model update, same policy update), these experiments will directly test the utility of our proposed model update.

Original: Experimental details are in Appendix D. 
Revised: Experimental details are in Appendix D. [Figure] [Figure]

Original: We ﬁrst use three locomotion tasks from the OpenAI Gym benchmark [7] to compare MnM-approxto MBPO and VMBPO.
Revised: We ﬁrst use three locomotion tasks from the OpenAI Gym benchmark [7] to compare MnM-approx to MBPO and VMBPO.

Original: The VMBPO curves are taken directly from that paper.
Revised: The VMBPO curves are taken directly from that paper.

Original: As shown in Fig.
Revised: As shown in Fig.

Original: 4a,
Revised: 4a,

Original: MnM-approx performs roughly on par with MBPO and outperforms [9], a more recent model-based[Figure] method also addresses the objective mismatch problem by maximizing an upper bound on expectedreturns (1-sided p-values: p = 0 . 04 , 0 . 00 , 0 . 03 ).
Revised: MnM-approx performs roughly on par with MBPO and outperforms [9], a more recent model-based method also addresses the objective mismatch problem by maximizing an upper bound on expected returns (1-sided p-values: p = 0 . 04 , 0 . 00 , 0 . 03 ).

Original: We next use the ROBEL manipulation benchmark [1] to compare how MnM-approx and MBPOhandle tasks with more complicated dynamics.
Revised: We next use the ROBEL manipulation benchmark [1] to compare how MnM-approx and MBPO handle tasks with more complicated dynamics.

Original: As shown in Fig.
Revised: As shown in Fig.

Original: 4b, MBPO struggles to learn threeof the four tasks, likely because the dynamics are hard.
Revised: 4b, MBPO struggles to learn three of the four tasks, likely because the dynamics are hard.

Original: In contrast, MnM-approx solves these tasks,likely because the GAN-like model is more accurate.
Revised: In contrast, MnM-approx solves these tasks, likely because the GAN-like model is more accurate.

Original: MnM-approx outperforms MBPO on all tasks( p ≤ 0 . 03 ) and outperforms the model-free baseline on 2/4 tasks ( p ≤ 0 . 02 ).
Revised: MnM-approx outperforms MBPO on all tasks ( p ≤ 0 . 03 ) and outperforms the model-free baseline on 2/4 tasks ( p ≤ 0 . 02 )

Original: We next compare to prior methods for addressing the objective mismatch problem using the
Revised: We next compare to prior methods for addressing the objective mismatch problem using the

Original: DClawScrewFixed-v0 task.
Revised: DClawScrewFixed-v0 task.

Original: We compare to VAML [17] and the value-weighted maximumlikelihood approach proposed in Lambert et al.[29].
Revised: We compare to VAML [17] and the value-weighted maximum likelihood approach proposed in Lambert et al.

Original: As shown in Fig.
Revised: As shown in Fig.

Original: 5, these alternative modellearning objectives perform poorly on this task.
Revised: 6, these alternative model learning objectives perform poorly on this task.

Original: To better understand why MnM-approxsometimes outperforms the maximum likelihood baseline (MBPO), we visualizedthe Q-values throughout training.
Revised: To better understand why MnM-approx sometimes outperforms the maximum likelihood baseline (MBPO), we visualized the Q-values throughout training.

Original: We usedmetaworld-drawer-open-v2 , a task wherewe found a noticeable difference in the performance between MBPO and MnM-approx.
Revised: We used metaworld-drawer-open-v2 , a task where we found a noticeable difference in the performance between MBPO and MnM-approx.

Original: Fig.
Revised: Fig.

Original: 6a shows that MnM-approx yields Q values that are more accurate and more stable than MBPO,perhaps because MBPO learns a policy that exploits inaccuracies in the learned model.
Revised: 7a shows that MnM-approx yields Q values that are more accurate and more stable than MBPO, perhaps because MBPO learns a policy that exploits inaccuracies in the learned model.

Original: To study the stability of MnM-approx, we plot the validation MSE of the MnM-approx modelwhen training on the DClawScrewRandom-v0 task.
Revised: To study the stability of MnM-approx, we plot the validation MSE of the MnM-approx model when training on the DClawScrewRandom-v task.

Original: As shown in Fig.
Revised: As shown in Fig.

Original: 6b, the MSE decreases stably,indicating that the adversarial nature of the MnM-approx model training does not create instabilities. 
Revised: 7b, the MSE decreases stably, indicating that the adversarial nature of the MnM-approx model training does not create instabilities. In Appendix C (Fig. 8, 9) we visualize Q

Original:  In Appendix C (Fig. 7), we  visualize the dynamicsmodel learned by the deep RL implementationof MnM-approx.
Revised: [Figure] value stability and MSE model error on additional environments. In Appendix C (Fig. 12), we study the effect of model horizon on MnM in a continuous control task, and ﬁnd that performance degrades linearly as we increase this value. We visualize the dynamics learned by MnM on two robotic control tasks in Fig 5.

Original: As in Fig.
Revised: 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-deterministic dynamics appear stochastic. On the Pusher-v2 task, the MnM dynamics cause the puck to move towards the robot arm even before the arm has come in contact with the puck. While this movement is not physically realistic, it may make the exploration problem easier.

Original: 2a, we ﬁnd that the  dynamics are optimistic, biasing the agent towardsoutcomes with high rewards.
Revised: On the HalfCheetah-v2 task, the MnM dynamics increase the probability that the agent remains upright after tripping, likely making it easier for the agent to learn how to run.

Revised: 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.

Original:  Conclusion
Revised:  Conclusion

Original: The main contribution of this paper is an objective for model-based RL that is both efﬁcient tocompute and globally valid.
Revised: The main contribution of this paper is an objective for model-based RL that is both efﬁcient to compute and globally valid.

Original: Moreover, the model and policy are jointly optimized using this sameobjective, thereby addressing the objective mismatch problem.
Revised: Moreover, the model and policy are jointly optimized using this same objective, thereby addressing the objective mismatch problem.

Original: This joint optimization will ease andaccelerate the design of future model-based RL algorithms. 
Revised: This joint optimization will ease and accelerate the design of future model-based RL algorithms. The main limitation of this paper is the classiﬁer term.

Revised: Because this term is estimated, the objective used in practice can fail to be a lower bound when the classiﬁer is not Bayes-optimal.

Revised: Additionally, this classiﬁer term can hinder exploration and degrade performance on continuous control tasks in the online setting.

Revised: We encourage future investigations into joint optimization based approaches that do explicitly address the role of exploration, perhaps including the dataset as an additional optimization variable.

Original: The main limitation of this paper is  the classiﬁer term.
Revised: Acknowledgements. We thank Laura Smith, Marvin Zhang and Dibya Ghosh for helpful discussions, and thank Anusha Nagabandi and Michael Janner for help with experiments. We thank Yinlam Chow for help sharing the VMBPO baselines and thank Danijar Hafner for reviewing a draft of the paper. This material is supported by the Fannie and John Hertz Foundation and the NSF GRFP (DGE1745016).

Original: Because this term is estimated, the objectiveused in practice can fail to be a  lower bound when the classiﬁer is not Bayes-optimal.
Revised: We now introduce a modiﬁcation to our lower bound that does make the bound tight. This new lower bound will be more complex than the one introduced above and we have not yet successfully designed an algorithm for maximizing it. Nonetheless, we believe that presenting the bound may prove useful for the design of future model-based RL algorithms.

Original: Additionally,this classiﬁer term can hinder exploration and degrade performance on continuous control tasks in theonline setting.
Revised: We will use L γ ( θ ) to denote this new lower bound. In addition to the policy and dynamics, this bound will also depend on a time-varying discount, γ θ ( t ) , in place of the typical γ t term.

Original:  We encourage future investigations into joint optimization based approaches that doexplicitly address the role of exploration, perhaps including the dataset asan additional optimizationvariable.
Revised: Similar learned discount factors have been studied in previous work on model-free RL [39]. We deﬁne this objective as follows: [Equation] where the augmented reward is now deﬁned as

Original: Checklist
Revised: [Equation]

Original: For all authors...
Revised: and Γ θ ( t ) = (cid:80) tt (cid:48) =

Original: (a) Do the  main claims made in the abstract and introduction accurately reﬂect the paper’s
Revised: γ θ ( t (cid:48) ) is the CDF of the learned discount function (i.e., γ θ ( t ) is a probability distribution over t .). 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 A.1.

Original: contributions and scope?
Revised: Let an arbitrary policy π ( a t | s t ) be given. The objective L γ ( θ ) is also a lower bound on the expected return objective, log E π [ (cid:80) ∞ t = [Equation] γ t r ( s t , a t )] ≥ L γ ( θ ) , and this bound becomes tight at optimality:

Original: [Yes] Our main claim is that our proposed objective is a
Revised: [Equation] The proof is presented in Appendix B.4. One important limitation of this result is that the learned dynamics that maximize this lower bound to make the bound tight may be non-Markovian. Intriguingly, this analysis suggests that using non-Markovian models, such as RNNs and transformers, may accelerate learning on Markovian tasks.

Original: global lower bound on  the expected return (proof in Appendix B).
Revised: This paper does not propose an algorithm for optimizing this more complex lower bound. B Proofs and Additional Analysis B.1 VMBPO Maximizes an Upper Bound on Return While MnM aims to maximize the (log) of the expected return, VMBPO aims to maximize the expected exponentiated return:

Original: (b) Did you describe the limitations of your work?
Revised: [Equation] where η > 0 is a temperature term used by VMBPO. Note that maximizing the log of the expected return, as done by MnM, is equivalent to maximizing the expected return, as the function log( · ) is monotone increasing. However, maximizing the log of the expected exponentiated return, as done by

Revised: VMBPO, is not equivalent to maximizing the expected return.

Revised: Rather, it corresponds to maximizing a sum of the expected return and the variance of the return [31, Page 272]:

Revised: [Equation]

Revised: Thus, in environments with stochastic dynamics or rewards (e.g., the didactic example in Fig. 3a),

Revised: VMBPO will prefer to receive lower returns if the variance of the returns is much higher.

Original: [Yes] See Sec. 6.
Revised: We note that the expected exponentiated return is an upper bound on the expected return: [Equation] This statement is a direct application of Jensen’s inequality. The bound holds with a strict inequality in almost all MDPs.

Revised: The one exception is trivial MDPs where all trajectories have exactly the same return.

Revised: Of course, even a random policy is optimal for these trivial MDPs.

Revised: B.2 Helper Lemmas

Revised: We start by introducing a simple identity that will help handle discount factors in our analysis.

Revised: Lemma B.1.

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: Deﬁne p ( H ) = G EOM (1 − γ ) as the geometric distribution. Let discount factor γ ∈ (0 , 1) and random variable x t be given.

Revised: Then the following identity holds:

Revised: [Equation]

Revised: The proof involves substituting the deﬁnition of the Geometric distribution and then rearranging terms.

Revised: Proof.

Revised: [Equation]

Revised: The second helper lemma describes how the discounted expected return objective can be written as the expected terminal reward of a mixture of ﬁnite-length episodes.

Revised: Lemma B.2.

Revised: Deﬁne p

Revised: ( H ) =

Revised: G EOM (1 − γ ) as the geometric distribution, and p ( τ | H ) as a distribution over length- H episodes.

Revised: We can then write the expected discounted return objective as follows:

Original: [Yes] While
Revised: [Equation] Proof.

Original: RL algorithms might be applied to applications with positive and negative societal implications, our theoretical analysis does not dictate these applications.
Revised: The ﬁrst identity follows from the deﬁnition of the geometric distribution. The second identity writes the expectations as integrals, which will make future analysis clearer.

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to
Revised: Proof.

Original: them?
Revised: [Table]

Original: [Yes]
Revised: [Equation]

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
Revised: [Table]

Original: [Yes]
Revised: [Equation]

Original: (b) Did you include complete proofs of all theoretical results?
Revised: [Table]

Original: [Yes] See Appendix B.
Revised: [Equation] [Table]

Original: If you ran experiments...
Revised: [Equation]

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?
Revised: [Table]

Original: [No] Experimental
Revised: [Equation]

Original: details are included in Appendix D. Code will be released upon publications.(b) Did you specify all the training details (e.g., data splits, hyperparameters,
Revised: [Table]

Original: how they
Revised: [Equation]

Original: were chosen)?
Revised: [Table]

Original: [Yes]
Revised: [Equation]

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
Revised: [Table]

Original: [Yes] Plots show the mean and standard deviation across the random seeds.
Revised: [Equation]

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
Revised: [Table]

Original: [No] Proprietary.
Revised: [Equation]

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
Revised: [Table] [Equation]

Original: [N/A] [Table] 
Revised: [Table] [Equation]

Original: (b) Did you mention the license of the assets?
Revised: [Table] [Equation] [Table] [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Original: [N/A]
Revised: [Equation] [Table] [Equation] [Table]

Original: (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: [Equation] [Table] [Equation] For (a) , we applied Lemma B.2.

Original: [N/A]
Revised: For (b) , we applied Jensen’s inequality.

Original: [Table]
Revised: For (c) , we applied Jensen’s inequality again. For (d) , we substituted the deﬁnitions of p θ ( τ | H ) and q θ ( τ | H ) . For (e) , we noted that the term inside the summation only depends on the ﬁrst H steps of the trajectory, so collecting longer trajectories will not change the result. This allows us to rewrite the integral as an expectation using a single inﬁnite-length trajectory.

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re
Revised: For (f) , we recalled the deﬁnition q θ ( τ ) = q θ ( τ = H = ∞ ) and swap the order of integration. For (g) , we express the inner integral over p ( H ) as an expectation.

Original: [Table] 
Revised: For (h) , we applied the identity from Lemma B.1. For (i) , we moved the constant log(1 − γ ) back inside the integral and rewrote the integral as an expectation. We have thus obtained the desired result. B.4 Proof of Lemma A.

Original: using/curating?
Revised: Before presenting the proof of Theorem 3.1 itself, we show how we derived the lower bound in this more general case. While this step is not required for the proof, we include it because it sheds light on how similar lower bounds might be derived for other problems. We deﬁne γ θ ( H ) to be a learned distribution over horizons H . We then proceed, following many of the same steps as for the proof of Theorem 3.1.

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Original: [N/A]
Revised: [Table] [Equation] [Table] For (a) , we applied Lemma B.2 and multiplied the integrand by q θ ( τ | H = H ) γ θ ( H ) q θ ( τ | H = H ) γ θ ( H )

Original: [Table]
Revised: = 1 .

Original: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable
Revised: For (b) , we applied Jensen’s inequality. For (c) , we factored p ( τ, H ) =

Original: information or offensive content?
Revised: p ( τ, H ) p ( H ) and q θ ( τ, H ) = q ( τ | H ) γ θ ( H ) .

Original: [N/A]
Revised: Note that under the joint distribution p ( τ, H ) , the horizon H ∼ p ( H ) = G EOM (1 − γ ) is independent of the trajectory, τ . For (d) , we rewrote the expectation as an expectation over a single inﬁnite-length trajectory and simpliﬁed the summand.

Original: [Table]
Revised: For (e) , we recall the deﬁnition q θ ( τ ) = q θ ( τ = H = ∞ ) .

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: For (f) , we rewrote the integral as an expectation and wrote out the deﬁnition of the geometric distribution, p ( H ) . For (g) , we regrouped the difference of dynamics terms. For (h) , we noted used the fact that (cid:80) H − 1 t =0 γ θ ( t ) + γ ∞ t =

Original: [Table]
Revised: H γ θ ( t ) = 1 .

Original: (a) Did you include the full text of instructions given to participants and screenshots, if
Revised: For (i) , we substituted the deﬁnition of the CDF function. For (j) , we rearranged terms so that all were multiplied by the discount γ θ ( H ) . Thus, we have obtained the desired result.

Original: [Table]
Revised: We now prove Lemma A.1, showing that Eq. 10 becomes tight at optimality. Proof.

Revised: [Equation]

Revised: For (a) , we undo some of the simpliﬁcations above, going back to Eq.

Revised: For (b) , we factor q θ ( τ, H ) =

Revised: q θ ( τ ) γ θ ( H | τ ) and p ( τ, H ) =

Revised: p ( τ ) p ( H ) .

Revised: At this point, we can solve analytically for the optimal discount distribution, γ θ ( H | τ ) :

Revised: [Equation]

Revised: In the second equality, we substitute the deﬁnition of R ( τ ) .

Revised: We then substitute Eq.

Revised: 15 into our expression for L γ ( θ )

Revised: and simplify the resulting expression.

Revised: [Equation]

Revised: In the ﬁnal line we have removed the integral over H because none of the integrands depend on H .

Revised: At this point

Revised: [Equation]

Revised: we can solve analytically for the optimal trajectory distribution, q θ ( τ ) :

Revised: [Equation]

Revised: [Equation]

Revised: We then substitute Eq.

Revised: 16 into our expression for L γ ( θ ) , and simplify the resulting expression:

Revised: [Equation]

Revised: [Table]

Revised: We have thus shown that the lower bound L γ becomes tight when we use the optimal distribution over trajectories q θ ( τ ) and optimal learned discount γ θ ( H | τ ) .

Revised: B.5 A lower bound for goal-reaching tasks.

Revised: Many RL problems can be better formulated as goal-reaching problems, a formulation that does not require deﬁning a reward function.

Revised: We now introduce a variant of our method for goal-reaching tasks.

Revised: Using ρ π ( s t + ) to denote the discounted state occupancy measure of policy π , we deﬁne the goal-reaching objective as maximizing the probability density of reaching a desired goal s g :

Revised: [Equation]

Revised: We refer the reader to Eysenbach et al.

Revised: [16] for a more detailed discussion of this objective.

Revised: For simplicity, we assume that the goal is ﬁxed, noting that the multi-task setting can be handled by conditioning the policy on the commanded goal.

Revised: Similar to Theorem 3.1, we can construct a lower bound on the goal-conditioned RL problem:

Revised: Lemma B.3.

Revised: Let initial state distribution p 1 ( s 1 ) , real dynamics p ( s t +1 | s t , a t ) , reward function r ( s t , a t ) > 0 , discount factor γ ∈ (0 , 1) , and goal g be given.

Original: applicable?
Revised: Then the following bound holds for any dynamics q ( s t +1 | s t , a t ) and policy π ( a t | s t ) : [Equation] where [Equation]

Revised: The proof, presented below, is similar to the proof of Theorem 3.1.

Revised: The ﬁrst term in the reward function, the log ratio of reaching the commanded goal one time step in the future, is similar to prior work [39].

Revised: The correction term log p − log q incentivizes the policy to avoid transitions where the model is inaccurate, and can be estimated using a separate classiﬁer.

Revised: One important aspect of this goal-reaching problem is that it is entirely data-driven, avoiding the need for any manually-designed reward functions.

Original: [N/A]
Revised: Proof. [Equation] [Figure] [Equation]

Revised: Similar to the more complex lower bound presented in Eq.

Revised: 10, this lower bound on goal-reaching can be modiﬁed (by learning a discount factor) to become a tight lower bound.

Revised: The resulting objective would resemble a model-based version of the algorithm from Rudner et al.

Original: [Table]
Revised: [Equation]

Revised: B.6 Derivation of Model Objective (Eq. 9)

Revised: Our lower bound depends on entirely trajectories sampled from the learned dynamics.

Revised: In this section, we show how the same objective can be expressed as an expectation of transitions.

Revised: This expression is easier to optimize

Revised: [Equation]

Revised: as it does not require backpropagating gradients through time.

Revised: We start by writing our lower bound, conditioned on a current state s t .

Revised: [Equation]

Revised: , s t +1 ) − (1 − γ ) log(1 − γ ) + γV ( s t +1 ) | s t (cid:21)

Revised: In (a) , we substituted the deﬁnition of the augmented return.

Revised: For the purpose of optimizing the dynamics model, we can ignore all terms that do not depend on s t +1 .

Original: (b) Did you describe any potential participant risks, with links to Institutional Review
Revised: Removing these terms, we arrive at our model training objective (Eq. 9) C Additional Experiments Fig. We compare MnM to a number of alternative model learning methods.

Revised: MBPO [24] uses a standard maximum likelihood model.

Revised: We implement a version of VAML [17], which augments the maximum likelihood loss with an additional temporal difference loss; the model should predict next states that have low Bellman error.

Revised: Finally, we compare to 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 et al.

Revised: 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.

Revised: Fig.

Revised: In this experiment, we show the model MSE throughout the training of MnM-approx, similar to Fig.

Revised: 7b.

Revised: While MnM-approx does not optimize for MSE, the MSE is nonetheless a rough barometer for whether the model is training stably.

Revised: On the DClawScrewFixed-v0 task, we observe that 33 seeds

Revised: [Figure]

Revised: [Figure]

Revised: all train stably, converges towards (though not quite reaching) the MSE of the maximum likelihood model.

Revised: On the metaworld-drawer-open-v2 task, the results are a bit more complicated, withseeds achieving a model MSE much lower than the maximum likelihood model, while two seeds seem to have failed to converge.

Revised: Overall, these plots indicate that MnM-approx often trains stably, but some seeds can fail to converge on some environments.

Revised: Fig.

Revised: In this experiment, we show the Q-values throughout the training of MnM-approx, similar to Fig.

Revised: 7a.

Revised: In both environments, we observe that the Q-values from MnMapprox are more stable than the Q-values from MBPO.

Revised: Note how the Q-values from MBPO tend to peak early during training, suggesting that they have overestimated the true returns and then correct towards a less inﬂated estimate of the agent’s expected return

Revised: Fig.

Revised: This is a plot from a preliminary version of

Revised: MnM, when we were testing the effect of the classiﬁer term.

Revised: In experiments like these, we found that the classiﬁer term signiﬁcantly hurt performance, motivating us to not include the classiﬁer term in the “MnM-Approx” method used in the continuous control experiments.

Revised: Fig.

Revised: In this experiment, we ablate the two factors of the MnM augmented reward (Eq. 3, using the stochastic gridworld from Fig.

Revised: 2a.

Revised: The ﬁrst ablation removes the logarithm, while the second ablation removes the classiﬁer.

Revised: The results, shown in Fig. 11, indicate that the logarithm term is crucial for getting good performance, but that removing classiﬁer term has a relative small effect, and may even boost performance by a small margin.

Revised: Not to read into these results too much, the aliasing experiment in Fig.

Revised: 2b has already demonstrated that the classiﬁer term is critical for ensuring good performance in the presence of function approximation.

Original: [Table]
Revised: [Figure] Fig.

Revised: In our experiments, we used very short model rollouts, only collecting a single transition from the model.

Revised: Fig.

Revised: 12 shows an ablation experiment on the metaworld-drawer-open-v2 task where we increased the model horizon, ﬁnding that it uniformly degrades performance.

Revised: Fig.

Revised: In this experiment, we examine the transferability of MnM’s learned dynamics model to new tasks.

Revised: When learning a similar task, we might expect that the optimistic dynamics produced by MnM would be more useful than the true environment dynamics.

Revised: However, when learning dissimilar tasks, the optimism for one task might hinder exploration towards solving the new task, and might be worse than using the true environment dynamics.

Revised: We tested these hypotheses using the stochastic gridworld from Fig.

Revised: 2a.

Revised: We took the dynamics learned by MnM (Fig.

Revised: 2a (right))

Revised: and used it to solve new tasks, deﬁned by placing the goal in different locations.

Revised: As shown in Fig.

Revised: 13 (right), tasks B and C are similar to the original task, in that they involve navigating around the wall; task A is not as similar to the original task, but is much easier.

Revised: To study the effectiveness of transferring the dynamics, we applied Q-learning to each of these tasks, either using the (incorrect) optimistic dynamics learned from the original task, or by using the true environment dynamics.

Revised: We show the results in Fig. 13.

Revised: On Tasks A and B, we observe little difference between Q-learning with the transferred dynamics model versus the true dynamics model.

Revised: However, on task C, only by using the transferred dynamics is Q-learning able to solve this challenging task.

Revised: In summary, these results suggest that the bias of the learned dynamics does not seem to hurt when learning dissimilar tasks, but can accelerate learning of challenging, similar tasks.

Revised: D Implementation Details

Revised: All experiments were run on at least three random seeds.

Revised: Although we were limited by computational constraints, we ﬁnd that most of our conclusions hold with p < 0 .

Revised: 05 , as noted in the main text.

Original: Board (IRB) approvals, if applicable?
Revised: Value Iteration (Fig. 2a (right) , Fig. 2b, Fig. 3b)

Revised: For the tabular experiments that perform value iteration, we perform Polyak averaging of the policy and learned dynamics model with parameter τ = 0 .

Revised: We found that the value iteration version of MnM diverged without this Polyak averaging step.

Revised: Experiments were stopped when the MnM dynamics model (which depends on the value function) changed by less than 1e-6 across iterations, as measured using an L 0 norm.

Revised: For VMBPO we used η = 1 .

Revised: Q-learning (Fig. 2a)

Revised: For the experiments with Q-learning (both with and without the MnM components), we performed (cid:15) -greedy exploration with (cid:15) = 0 .

Revised: We used a learning rate of 1 e − 2 .

Revised: For this task alone, we compute the MnM dynamics analytically by combining the true environment dynamics with the learned value function, allowing for clearer theoretical analysis.

Revised: For fair comparison, all

Revised: [Table]

Revised: methods receive the same amount of data, perform the same number of updates, and are evaluated using the real environment dynamics.

Revised: For VMBPO we used η = 1 .

Revised: SAC for continuous control tasks.

Revised: We used the SAC implementation from TF-Agents [21] with the default hyperparameters.

Revised: MBPO for continuous control tasks.

Revised: We implement MBPO on top of the SAC implementation from TF-Agents [21].

Revised: Unless otherwise mentioned, we take the default parameters from this implementation.

Revised: We use an ensemble of 5 dynamics models, each with 4 hidden layers of size 256.

Revised: The dynamics model predicts the whitened difference between the next state and the current state.

Revised: That is, to obtain the prediction for the next state, the predictions are scaled by a per-coordinate variance, shifted by a per-coordinate mean, and then added to the current state.

Revised: These whitened predictions are clipped to have a minimum standard deviation of 1e-5; without this, we found that the MBPO model resulted in numerical instability.

Revised: The model is trained using the standard maximum likelihood objective, with all members of the ensemble being trained on the same data.

Revised: To sample data from this model we perform 1-step rollouts, starting at states visited in the true dynamics.

Revised: We perform one batch of rollouts in parallel using a batch size of 256.

Revised: To sample the corresponding action, with probability 50% we take the action that was executed in the true dynamics; with probability 50% we sample an action from the current policy.

Revised: We found that this modiﬁcation slightly improves the results of MBPO.

Revised: We use a batch size of 256.

Revised: We have two replay buffers: the model replay buffer has size 256e3 and the replay buffer of real experience has size 1e6.

Revised: At the start of training, we collect 1e4 transitions from the real environment, train the dynamics model on this experience for 1e5 batches, and only then start training the policy.

Revised: We use a learning rate of 3e-4 for all components.

Revised: To stabilize learning, we maintain a target dynamics model using an exponential moving average ( τ = 0 . 001 ), and use this target dynamics model to sample transitions for training.

Revised: 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 2).

Revised: MnM for Continuous Control Tasks We implement MnM on top of the SAC implementation from TF-Agents [21].

Revised: Unless otherwise mentioned, we take the default parameters from this implementation.

Revised: Our model architecture is exactly the same as our MBPO implementation, and we follow the same training protocol.

Revised: Unlike MBPO, MnM also learns a classiﬁer for distinguishing real versus model transitions.

Revised: The classiﬁer architecture is a 2 layers neural network with 1024 hidden units in each layer.

Revised: We found that this large capacity was important for stable learning.

Revised: We add input noise with σ = 0 .

Revised: 1 while training the classiﬁer.

Revised: We whiten the inputs to the classiﬁer by subtracting a coordinate-wise mean and dividing by a coordinate-wise standard deviation.

Revised: When training the classiﬁer, we take samples from both the dynamics model and the target dynamics model as negative examples, ﬁnding that this stabilizes learning somewhat.

Original: [N/A]
Revised: Following the suggestion of prior work [40], we use one-sided label smoothing with value 0.1, only smoothing the negative predictions and not the positive predictions. We found that gradient penalties and spectral normalization decreased performance. We found that automatically tuning the classiﬁer input noise also decreased performance. We found that mixup had little effect.

Revised: We found that the loss would often plateau around 1e4 batches, but would eventually start decreasing again after 2e4 - 2e5 batches.

Revised: Like the MBPO model, we ﬁrst collect 1e4 transitions of experience from the real environment using a random policy, then train the dynamics model and classiﬁer for 1e5 batches, and only then start updating the policy.

Revised: Because the Q values are poor at the start of training, we only add the value term to the model loss (resulting in the optimistic dynamics model) after 2e5 batches (1e5 batches of model training, then 1e5 batches of model+policy training).

Revised: To further improve stability, we compute the value term in the model loss by taking the minimum over two target value functions (like

Revised: TD3 [18]).

Revised: We update the model, classiﬁer policy, and value functions at the same rate we sample experience from the learned model, which is more frequently than we collect experience from the real environment (see Table 2).

Revised: D.2 Environments

Revised: This section provides details for the environments used in our experiments.

Revised: We visualize the environments in Fig.

Revised: Gridworld for Fig.

Revised: 2a.

Revised: This task is a 10 × 10 gridworld with obstacles shown in Fig.

Revised: 2a.

Revised: There are four discrete actions, corresponding to moving to the four adjacent cells.

Revised: With probability 50%, the agent’s action is ignored and a random action is taken instead.

Revised: The agent starts in the top-left cell.

Revised: The agent receives a reward of +0.001 at each time step and a reward of +10 when at the goal state.

Revised: Episodes have 200 steps and we use a discount γ = 0 .

Original: [Table]
Revised: Gridworld for Fig. 2b. This task is a 15 × 15 gridworld with obstacles shown in Fig. 2b.

Revised: There are four discrete actions, corresponding to moving to the four adjacent cells.

Revised: The dynamics are deterministic.

Revised: The agent starts in the top-left cell.

Revised: The reward is +1.0 at every state except the goal state, where the agent receives a reward of +100.0.

Revised: We use γ = 0 .

Revised: 9 and compute optimal policies analytically using value iteration.

Revised: We implement the aliasing by averaging together the dynamics for each block of 3 × 3 states.

Revised: Importantly, the averaging was done to the relative dynamics (e.g., action 1 corresponds to move right) not the absolute dynamics (e.g., action 1 corresponds to moving to state (3, 4)).

Revised: To handle edge effects, we modiﬁed the averaged dynamics so that the agent could not exit the gridworld.

Revised: We computed the classiﬁer analytically using the true and learned dynamics models.

Revised: However, the augmented rewards become inﬁnite because the learned model assigns non-zero probability to transitions that cannot occur under the true dynamics.

Revised: This is not a failure of our theory, as the optimal policy would choose to never visit these states, but it presents a challenge for optimization.

Revised: We therefore added label smoothing with parameter 0.7 to the classiﬁer.

Revised: Gridworld for Fig.

Revised: 2a This task is a 10 × 10 gridworld with obstacles shown in Fig.

Revised: 2a (Right) .

Revised: There are four discrete actions, corresponding to moving to the four adjacent cells.

Revised: With probability 90%, the agent’s action is ignored and a random action is taken instead.

Revised: The agent starts in the top-left cell.

Revised: The rewards depend on the Manhattan distance to the goal: transitions that lead away from the goal have a reward of +0.001, transitions that do not change the distance to the goal have a reward of +1.001, and transitions that decrease the distance to the goal have a reward of +2.001.

Revised: We use γ = 0 .

Revised: 5 and compute values and returns analytically using value iteration.

Revised: Gridworld for Fig.

Original: (c) Did you include the estimated hourly wage paid to participants and  the total amount
Revised: 3b. This task used the same dynamics as Fig. 2a. The one change is that the agent receives a reward of +1.0 at each time step and a reward of +10.0 at the goal state.

Original: [Table]
Revised: We estimate the more complex lower bound (Eq. 10) by also learning the discount factor. However, since we currently do not have a method for learning non-Markovian dynamics to fully optimize this lower bound, we do not expect the lower bound to become tight. HalfCheetah-v2, Hopper-v2, Walker2d-v2 . These tasks are taken directly from the OpenAI benchmark [7] without modiﬁcation.

Original: spent on participant compensation?
Revised: metaworld-drawer-open-v2 . This task is based on the drawer-open-v2 task from the Metaworld benchmark [50]. To increase the difﬁculty of this task, we remove the reward shaping term ( reward_for_caging ) and just optimize the reward for opening the drawer ( reward_for_opening ).

Revised: DClawPoseRandom-v0 , DClawTurnRandom-v0 , DClawScrewFixed-v0 , DClawScrewRandom-v0 .

Original: [N/A]
Revised: These tasks are taken directly from the ROBEL benchmark [1] without modiﬁcation. D.3 Differences between MnM and MnM-Approx Both MnM and MnM-Approx learn the same dynamics classiﬁer, and both use the same GAN-like objective (Eq. 9) to update the model. Both perform the same policy updates.

Revised: The difference is the reward function used to update the Q-function:

Original: [Table]
Revised: [Equation]