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 RLor 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 smallfraction of states where an accurate model makes mistakes, or it might act in waysthat 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, 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.

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 themodel and the policy, such that updates to either component increase a lower boundon 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 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.

Original: Aversion of this bound becomes tight under certain assumptions.
Revised: Aversion 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 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.

Original: Numericalsimulations demonstrate that optimizing this bound yields reward maximizingpolicies and yields dynamics that (perhaps surprisingly) can aid in exploration.
Revised: Numericalsimulations demonstrate that optimizing this bound yields reward maximizingpolicies 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: 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.

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 supervisedlearning 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 thatmodel 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 smallmistakes 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 withlarge errors may yield a policy that attains high return if the model errors occur in states that thepolicy 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., usingmaximum 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 trainedusing 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 objectivefor 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 formodel-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 variableis the trajectory and the cumulative reward is interpreted as the probability that the trajectory solvesthe 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 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.

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

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 themodel against the policy [5, 33, 37].

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 canbe 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 lowerbound 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 morecomplex 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 numericalsimulations in simple tasks, we demonstrate that optimizing the bound yields reward-maximizingpolicies 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 thatour bound gracefully accounts for function approximation error in the model – the policy is penalizedfor 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 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.

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 RLto 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 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].

Original: This issue isreferred to as the objective mismatch problem: the model and policy (or planner) are optimized usingdifferent objectives.
Revised: This issue isreferred to as the objective mismatch problem: the model and policy (or planner) are optimized usingdifferent 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 thosethat 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 isinaccurate [27, 30, 42, 49, 51], a strategy that our method will also employ.

Original: Other work modiﬁesthe model objective to include multi-step rollouts [3, 4, 17, 25, 45, 47].
Revised: Other work modiﬁesthe 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 modifythe model objective in a different way, so that the model objective is exactly the same as the policyobjective.

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, ourapproach will not require differentiating through unrolled model updates or optimization procedures.

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 lowerbound 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 alower 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 boundthat 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]
Revised: [Table]

Original: 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 derivationbuilds on prior work thatcasts model-based RL as atwo-player game between amodel-player and a policyplayer [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: 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).

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: Ourapproach, though structurally resembling a GAN, is different from prior work that replaces amaximum 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 policyusing 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, VMBPOmaximizes a different, risk-seeking objective, which is an upper bound on expected return (see

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

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 tolearn 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 modelq θ ( 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 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.

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

Original: π ( τ ) = p 0 ( s 0 ) ∞ (cid:89) t = p ( s t +1 | s t , a t )
Revised: π ( τ ) = p 0 ( s 0 ) ∞ (cid:89) t = p ( s t +1 | s t , a t )

Original: π θ ( a t | s t ) , q π ( τ ) =
Revised: π θ ( a t | s t ) , q π ( τ ) =

Original: p 0 ( s 0 ) ∞ (cid:89) t =
Revised: p 0 ( s 0 ) ∞ (cid:89) t =

Original: q θ ( s t +1 | s t , a t )
Revised: q θ ( s t +1 | s t , a t )

Original: π θ ( a t | s t ) .
Revised: π θ ( a t | s t ) .

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

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 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].

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 developinglower bounds for model-based RL, even the best lower bounds do not hold for all models [30] or arelimited 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 updatesto 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 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].

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]) doestrain 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: 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.

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 learningthe 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 themodel 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 anew 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 measuresthe 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 unlikelyunder 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 canestimate 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 modelq θ ( 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 learnedmodel; 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 thedynamics 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 thathave 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 theoriginal reward functions; our experiments demonstrate that excluding this component invalidates ourlower 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 expectedreturn.

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 ) :

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 formodel-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 modelnot 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 betight, 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 achievehigh 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 thatdoes 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 simplebound 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 notequal to the true environment dynamics.

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

Original: Ignoring parametrization constraints, the dynamics model that optimizes our lowerbound is q ∗ ( τ ) =
Revised: Ignoring parametrization constraints, the dynamics model that optimizes our lowerbound 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 theobjective-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 dynamicsmodel may accelerate policy optimization, a hypothesis we will test in Sec. 5.1.

Original: Would the optimistic dynamics overestimate the policy’s return, violating Lemma 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 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.

Original: We demonstratethis effect in Fig.
Revised: We demonstratethis 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 optimizingthe 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 optimizingthis 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 functiondepends 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 theclassiﬁ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 ofthe real environment, we learn a classiﬁer that distinguishes real transitions from fake transi- tions.

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 modeltransitions 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 sameinitial 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 theaugmented 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 thatwe can estimate the augmented reward function, we can apply any RL method to maximize theaugmented reward under transitions sampled from the dynamics model.

Original: The following sectiondescribes a particular instantiation using an off-policy RL algorithm.
Revised: The following sectiondescribes 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, whichtrains 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 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).

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

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 transitionssampled 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 RLalgorithm, 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 Q 206 function using the TD loss on transitions sampled from the learned dynamics model:
Revised: We train the Q 206 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: 7) where (cid:98)·(cid:99) sg is the stop-gradient operator and y t =

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

Original: [Equation]
Revised: [Equation]

Original: ) (cid:3)is the TD target.
Revised: ) (cid:3)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 estimate 209 the corresponding value function, we use a 1-sample approximation: V ψ ( s t ) =
Revised: To estimate 209 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 ) wherea 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: Followingprior work [18], we maintain two Q functions and two target Q functions, using the minimum of thetwo 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 in 215 terms of a single transition (derivation in Appendix B.6):
Revised: To optimize the dynamics model, we rewrite the lower bound in 215 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 differencein 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] andadversarial 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 isupdated 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 simplyreplacing 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 notoptimize 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 valueterm 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 modeland 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 thesame 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 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.

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, createan 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 trainthe 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, addthe 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 neuralnetwork 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 derivationsuggests that this term should appear, and our didactic experiments demonstrate that removing thisterm 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 combaterrors 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 termscritical 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 thatincluding this term hurts performance.

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 oppositeof exploration objectives based on model error (e.g., [44]), and penalizes the policy for performingexploration.

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 theoreticalderivation: they will use the model objective suggested by our theory but the same policy objectiveas 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,
Revised: 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,

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [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 thatmaximizing the bound yields reward-maximizing policies.

Original: We study these questions in Sec.
Revised: We study these questions in Sec.

Original: 5.1,where we will use tabular problems so that we can study the objective in the absence of functionapproximation error.
Revised: 5.1,where we will use tabular problems so that we can study the objective in the absence of functionapproximation 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 competitiveresults 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 ﬁrst experiments test whether optimizing the lower bound produces optimal policies and studyhow 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 toanalytically compute the optimal policy for comparison and allow us to evaluate the lower bound inthe 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 aperfect 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 thereward 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 modeland 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 itmaximizes 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 directionwith equal probability, the MnM dynamics model biases this stochasticity to lead the agent towardsthe 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, notthe 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) thelogarithmic 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 incorrecting for inaccurate models.

Original: To do this, we limit the capacity of the MnM dynamics model
Revised: To do this, we limit the capacity of the MnM dynamics model

Original: [Figure]
Revised: [Figure]

Original: so that it makes “low-resolution” predictions, forcing all states in 3 × 3 blocks to have the samedynamics.
Revised: so that it makes “low-resolution” predictions, forcing all states in 3 × 3 blocks to have the samedynamics.

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 occurat 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 makespredictions for states near the obstacle, it will average together some states with obstacles and somestates 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 someprobability 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 isinaccurate 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 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.

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 logarithmictransformation, 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 rewardat 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 agentout 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 rightstate 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: 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).

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 aversion 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 predictedby 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 expectedreturn, illustrating the importance of the logarithmic transformation.

Original: 5.2 Comparisons On Higher-Dimensional Tasks
Revised: 5.2 Comparisons On Higher-Dimensional Tasks

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

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 whetherthe 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 resultsand 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 modelupdate, same policy update), these experiments will directly test the utility of our proposed modelupdate.

Original: Experimental details are in Appendix D.
Revised: Experimental details are in Appendix D.

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-approxto 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
Revised: MnM-approx performs roughly on par with MBPO and outperforms [9], a more recent model-based

Original: [Figure]
Revised: [Figure]

Original: 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: 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 ).

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 MBPOhandle 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 threeof 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 ScrewFixed-v0 task.
Revised: We next compare to prior methods for addressing the objective mismatch problem using the DClawScrewFixed-v0 task.

Original: We compare to
Revised: We compare to

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

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

Original: 5, these alternative modellearning objectives perform poorly on this task.
Revised: 5, these alternative modellearning 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-approxsometimes outperforms the maximum likelihood baseline (MBPO), we visualizedthe 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 usedmetaworld-drawer-open-v2 , a task wherewe 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: 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.

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 modelwhen training on the DClawScrewRandom-v0 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: 6b, the MSE decreases stably,indicating that the adversarial nature of the MnM-approx model training does not create instabilities.

Original: In Appendix C (Fig. 7), we visualize the dynamics model learned by the deep RL implementationof MnM-approx.
Revised: In Appendix C (Fig. 7), we visualize the dynamics model learned by the deep RL implementationof MnM-approx.

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

Original: 2a, we ﬁnd that the dynamics are optimistic, biasing the agent towardsoutcomes with high rewards.
Revised: 2a, we ﬁnd that the dynamics are optimistic, biasing the agent towardsoutcomes with high rewards.

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 tocompute 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 sameobjective, 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 andaccelerate the design of future model-based RL algorithms.

Original: The main limitation of this paper is the classiﬁer term.
Revised: The main limitation of this paper is the classiﬁer term.

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: Because this term is estimated, the objectiveused in practice can fail to be a lower bound when the classiﬁer is not Bayes-optimal.

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

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

Original: Checklist
Revised: Checklist

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
Revised: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s

Original: contributions and scope?
Revised: contributions and scope?

Original: [Yes] Our main claim is that our proposed objective is a
Revised: [Yes] Our main claim is that our proposed objective is a

Original: global lower bound on the expected return (proof in Appendix B).
Revised: global lower bound on the expected return (proof in Appendix B).

Original: (b) Did you describe the limitations of your work?
Revised: (b) Did you describe the limitations of your work?

Original: [Yes] See Sec. 6.
Revised: [Yes] See Sec. 6.

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: [Yes] While
Revised: [Yes] While

Original: RL algorithms might be applied to applications with positive and negative societal
Revised: RL algorithms might be applied to applications with positive and negative societal

Original: implications, our theoretical analysis does not dictate these applications.
Revised: implications, our theoretical analysis does not dictate these applications.

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

Original: them?
Revised: them?

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes]
Revised: [Yes]

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

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

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?
Revised: mental results (either in the supplemental material or as a URL)?

Original: [No] Experimental
Revised: [No] Experimental

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: 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,

Original: how they
Revised: how they

Original: were chosen)?
Revised: were chosen)?

Original: [Yes]
Revised: [Yes]

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

Original: [Yes] Plots show the mean and standard deviation across the
Revised: [Yes] Plots show the mean and standard deviation across the

Original: random seeds.
Revised: random seeds.

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: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?

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

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

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

Original: [Table]
Revised: [Table]

Original: (b) Did you mention the license of the assets?
Revised: (b) Did you mention the license of the assets?

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

Original: (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: (c) Did you include any new assets either in the supplemental material or as a URL?

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

Original: [Table]
Revised: [Table]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: [Table]
Revised: [Table]

Original: using/curating?
Revised: using/curating?

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

Original: [Table]
Revised: [Table]

Original: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable
Revised: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable

Original: information or offensive content?
Revised: information or offensive content?

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

Original: [Table]
Revised: [Table]

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: [Table]
Revised: [Table]

Original: (a) Did you include the full text of instructions given to participants and screenshots, if
Revised: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: [Table]
Revised: [Table]

Original: applicable?
Revised: applicable?

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

Original: [Table]
Revised: [Table]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review
Revised: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: [Table]
Revised: [Table]

Original: Board (IRB) approvals, if applicable?
Revised: Board (IRB) approvals, if applicable?

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

Original: [Table]
Revised: [Table]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount[Table] spent on participant compensation?
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amount

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

Original: A Tightening the lower bound

Original: We now introduce a modiﬁcation to our lower bound that does make the bound tight.

Original: This new lowerbound will be more complex than the one introduced above and we have not yet successfully designedan algorithm for maximizing it.

Original: Nonetheless, we believe that presenting the bound may prove usefulfor the design of future model-based RL algorithms.

Original: We will use L γ ( θ ) to denote this new lower bound.

Original: In addition to the policy and dynamics, this boundwill also depend on a time-varying discount, γ θ ( t ) , in place of the typical γ t term.

Original: Similar learneddiscount factors have been studied in previous work on model-free RL [39].

Original: We deﬁne this objectiveas follows:

Original: [Equation]

Original: where the augmented reward is now deﬁned as

Original: [Equation]

Original: and Γ θ ( t ) =

Original: (cid:80) tt (cid:48)

Original: =

Original: γ θ ( t (cid:48) ) is the CDF of the learned discount function (i.e., γ θ ( t ) is a probabilitydistribution over t .).

Original: This new lower bound, which differs from our main lower bound by the learnablediscount factor, does provide a tight bound on the expected return objective:

Original: Lemma A.1.

Original: Let an arbitrary policy π ( a t | s t ) be given.

Original: [-The objective L γ ( θ ) is also a lower boundon the expected return objective, log E π [ (cid:80) ∞ t =
-]
Revised: [-The objective L γ ( θ ) is also a lower boundon the expected return objective, log E π [ (cid:80) ∞ t =
-]

Original: γ t r ( s t , a t )] ≥ L γ ( θ ) , and this bound becomes tight 528 at optimality:

Original: [Equation]

Original: The proof is presented in Appendix B.4.

Original: One important limitation of this result is that the learneddynamics that maximize this lower bound to make the bound tight may be non-Markovian.

Original: Intriguingly, this analysis suggests that using non-Markovian models, such as RNNs and transformers, mayaccelerate learning on Markovian tasks.

Original: This paper does not propose an algorithm for optimizing thismore complex lower bound.

Original: B Proofs and Additional AnalysisB.1 VMBPO Maximizes an Upper Bound on Return

Original: While MnM aims to maximize the (log) of the expected return, VMBPO aims to maximize theexpected exponentiated return:

Original: [Equation]

Original: where η > 0 is a temperature term used by VMBPO.

Original: Note that maximizing the log of the expectedreturn, as done by MnM, is equivalent to maximizing the expected return, as the function log( · ) ismonotone increasing.

Original: However, maximizing the log of the expected exponentiated return, as done by

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

Original: Rather, it corresponds to maximizing asum of the expected return and the variance of the return [31, Page 272]:

Original: [Equation]

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

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

Original: We notethat the expected exponentiated return is an upper bound on the expected return:

Original: [Equation]

Original: This statement is a direct application of Jensen’s inequality.

Original: The bound holds with a strict inequalityin almost all MDPs.

Original: The one exception is trivial MDPs where all trajectories have exactly the samereturn.

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

Original: B.2 Helper Lemmas

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

Original: Lemma B.1.

Original: Deﬁne p ( H ) = G EOM (1 − γ ) as the geometric distribution. Let discount factorγ ∈ (0 , 1) and random variable x t be given.
Revised: spent on participant compensation?

Original: Then the following identity holds:

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

Original: The proof involves substituting the deﬁnition of the Geometric distribution and then rearrangingterms.

Original: Proof.

Original: [Equation]

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

Original: Lemma B.2.

Original: Deﬁne p

Original: ( H ) =

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

Original: We can then write the expected discounted return objective asfollows:

Original: [Equation]

Original: Proof.

Original: The ﬁrst identity follows from the deﬁnition of the geometric distribution.

Original: The second identitywrites the expectations as integrals, which will make future analysis clearer.

Original: Proof.

Original: [Table]

Original: [Equation]

Original: For (a) , we applied Lemma B.2.

Original: For (b) , we applied Jensen’s inequality.

Original: For (c) , we applied Jensen’sinequality again.

Original: For (d) , we substituted the deﬁnitions of p θ ( τ | H ) and q θ ( τ | H ) .

Original: For (e) , we notedthat the term inside the summation only depends on the ﬁrst H steps of the trajectory, so collectinglonger trajectories will not change the result.

Original: This allows us to rewrite the integral as an expectationusing a single inﬁnite-length trajectory.

Original: For (f) , we recalled the deﬁnition q θ ( τ ) =

Original: q θ ( τ = H = ∞ )and swap the order of integration.

Original: For (g) , we express the inner integral over p ( H ) as an expectation.

Original: For (h) , we applied the identity from Lemma B.1.

Original: For (i) , we moved the constant log(1 − γ ) backinside the integral and rewrote the integral as an expectation.

Original: We have thus obtained the desiredresult.

Original: B.4 Proof of Lemma A.1

Original: Before presenting the proof of Lemma 3.1 itself, we show how we derived the lower bound in this 575 more general case.

Original: While this step is not required for the proof, we include it because it sheds lighton how similar lower bounds might be derived for other problems.

Original: We deﬁne γ θ ( H ) to be a learneddistribution over horizons H .

Original: We then proceed, following many of the same steps as for the proof of 578 Lemma 3.1.

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: H ,

Original: [Table]+ [Table]
Revised: [Table]

Original: +

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: H

Original: [Table]

Original: For (a) , we applied Lemma B.2 and multiplied the integrand by q θ ( τ | H = H ) γ θ ( H ) q θ ( τ | H = H ) γ θ ( H )

Original: = 1 .

Original: For (b) , weapplied Jensen’s inequality.

Original: For (c) , we factored p ( τ, H ) =

Original: p ( τ, H ) p ( H ) and q θ ( τ, H ) =

Original: q ( τ |

Original: H ) γ θ ( H ) .

Original: Note that under the joint distribution p ( τ, H ) , the horizon H ∼ p ( H )

Original: =

Original: G EOM (1 − γ )is independent of the trajectory, τ .

Original: For (d) , we rewrote the expectation as an expectation overa single inﬁnite-length trajectory and simpliﬁed the summand.

Original: For (e) , we recall the deﬁnitionq θ ( τ ) =

Original: q θ ( τ = H = ∞ ) .

Original: For (f) , we rewrote the integral as an expectation and wrote out thedeﬁnition of the geometric distribution, p ( H ) .

Original: For (g) , we regrouped the difference of dynamicsterms.

Original: For (h) , we noted used the fact that

Original: (cid:80) H − 1 t =0 γ θ ( t ) + γ ∞ t =

Original: H γ θ ( t ) = 1 .

Original: For (i) , we substitutedthe deﬁnition of the CDF function.

Original: For (j) , we rearranged terms so that all were multiplied by thediscount γ θ ( H ) .

Original: Thus, we have obtained the desired result.

Original: We now prove Lemma A.1, showing that

Original: Eq.

Original: 10 becomes tight at optimality.

Original: Proof.

Original: [Equation]

Original: [Equation]

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

Original: For (b) , we factorq θ ( τ, H ) =

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

Original: p ( τ ) p ( H ) .

Original: At this point, we can solve analytically forthe optimal discount distribution, γ θ ( H | τ ) :

Original: [Equation]

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

Original: We then substitute Eq.

Original: 15 into ourexpression for L γ ( θ ) and simplify the resulting expression.

Original: [Equation]

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

Original: At this point, we can solve analytically for the optimal trajectory distribution, q θ ( τ ) :

Original: [Equation]

Original: [Equation]

Original: We then substitute Eq.

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

Original: [Equation]

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

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

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

Original: We now introduce a variant of our method for goal-reachingtasks.

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

Original: [Equation]

Original: We refer the reader to Eysenbach et al.

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

Original: Forsimplicity, we assume that the goal is ﬁxed, noting that the multi-task setting can be handled byconditioning the policy on the commanded goal.

Original: Similar to Lemma 3.1, we can construct a lowerbound on the goal-conditioned RL problem:

Original: Lemma B.3.

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

Original: Then the following bound holds forany dynamics q ( s t +1 | s t , a t ) and policy π ( a t | s t ) :

Original: [Equation]

Original: where

Original: [Equation]

Original: The proof, presented below, is similar to the proof of Lemma 3.1.

Original: The ﬁrst term in the rewardfunction, the log ratio of reaching the commanded goal one time step in the future, is similar to priorwork [39].

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

Original: One important aspect of thisgoal-reaching problem is that it is entirely data-driven, avoiding the need for any manually-designedreward functions.

Original: Proof.

Original: [Equation]

Original: [Equation]

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

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

Original: The resulting objectivewould resemble a model-based version of the algorithm from Rudner et al.

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

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

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

Original: This expression iseasier to optimize, as it does not require backpropagating gradients through time.

Original: We start by writingour lower bound, conditioned on a current state s t .

Original: [Equation]

Original: , a t , s t +1 ) 1 − C φ ( s t , a t , s t +1 ) − (1 − γ ) log(1 − γ ) + γV ( s t +1 ) | s t (cid:21)

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

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

Original: Removing these terms, wearrive at our model training objective (Eq. 9)

Original: C Additional Experiments

Original: [Figure]

Original: Fig.

Original: We visualize the dynamics learned by MnM on two robotic control tasks.

Original: These tasks havedeterministic dynamics, so our theory would predict that an idealized version of MnM would learn adynamics model exactly equal to the deterministic dynamics.

Original: However, our implementation relies onfunction approximation (neural networks) to learn the dynamics, and the limited capacity of functionapproximators makes otherwise-deterministic dynamics appear stochastic.

Original: On the Pusher-v2 task,the MnM dynamics cause the puck to move towards the robot arm even before the arm has come incontact with the puck.

Original: While this movement is not physically realistic, it may make the explorationproblem easier.

Original: On the HalfCheetah-v2 task, the MnM dynamics increase the probability that theagent remains upright after tripping, likely making it easier for the agent to learn how to run.

Original: Weexpect that the implicit stochasticity caused by function approximation to be especially important forreal-world tasks, where the complexity of the real dynamics often dwarfs the capacity of the learneddynamics model.

Original: [Figure]

Original: Fig.

Original: We compare MnM to a number of alternative model learning methods.

Original: MBPO [24] usesa standard maximum likelihood model.

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

Original: Finally, we compare to a variant of the MBPO maximumlikelihood model that weights transitions based on the Q values, an idea discussed (but not actuallyimplemented) in Lambert et al.

Original: We implement this value weighting method by computing the Qvalues for the current states and computing a softmax over the batch dimension to obtain per-exampleweights.

Original: D Implementation Details

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

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

Original: 05 , as noted in the main text.

Original: D.1 Algorithms

Original: Value Iteration (Fig.

Original: 2a (right) , Fig. 2b, Fig.

Original: 3b)

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

Original: We found that the value iteration version of MnM diverged without this Polyak averagingstep.

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

Original: For VMBPOwe used η = 1 .

Original: Q-learning (Fig. 2a)

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

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

Original: For thistask alone, we compute the MnM dynamics analytically by combining the true environment dynamicswith the learned value function, allowing for clearer theoretical analysis.

Original: For fair comparison, allmethods receive the same amount of data, perform the same number of updates, and are evaluatedusing the real environment dynamics.

Original: For VMBPO we used η = 1 .

Original: SAC for continuous control tasks.

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

Original: MBPO for continuous control tasks.

Original: We implement MBPO on top of the SAC implementationfrom TF-Agents [21].

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

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

Original: Thedynamics model predicts the whitened difference between the next state and the current state.

Original: Thatis, 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.

Original: These whitened predictionsare clipped to have a minimum standard deviation of 1e-5; without this, we found that the MBPOmodel resulted in numerical instability.

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

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

Original: We performone batch of rollouts in parallel using a batch size of 256.

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

Original: We found that this modiﬁcation slightly improves the

Original: SAC

Original: MBPO

Original: MnM

Original: HalfCheetah-vHopper-vWalker2d-vmetaworld-drawer-open-v

Original: DClawPoseRandom-v

Original: DClawTurnRandom-v

Original: DClawScrewFixed-v

Original: DClawScrewRandom-v    results of MBPO.

Original: We use a batch size of 256.

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

Original: At the start of training, wecollect 1e4 transitions from the real environment, train the dynamics model on this experience for1e5 batches, and only then start training the policy.

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

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

Original: We update themodel, 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).

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

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

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

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

Original: Theclassiﬁer architecture is a 2 layers neural network with 1024 hidden units in each layer.

Original: We foundthat this large capacity was important for stable learning.

Original: We add input noise with σ = 0 .

Original: 1 whiletraining the classiﬁer.

Original: We whiten the inputs to the classiﬁer by subtracting a coordinate-wise meanand dividing by a coordinate-wise standard deviation.

Original: When training the classiﬁer, we take samplesfrom both the dynamics model and the target dynamics model as negative examples, ﬁnding that thisstabilizes learning somewhat.

Original: Following the suggestion of prior work [40], we use one-sided labelsmoothing with value 0.1, only smoothing the negative predictions and not the positive predictions.

Original: We found that gradient penalties and spectral normalization decreased performance.

Original: We found thatautomatically tuning the classiﬁer input noise also decreased performance.

Original: We found that mixup hadlittle effect.

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

Original: Like the MBPO model, we ﬁrst collect 1e4 transitions of experience from the real environment usinga random policy, then train the dynamics model and classiﬁer for 1e5 batches, and only then startupdating the policy.

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

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

Original: TD3 [18]).

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

Original: Gridworld for Fig.

Original: 2a.

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

Original: 2a.

Original: Thereare four discrete actions, corresponding to moving to the four adjacent cells.

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

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

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

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

Original: Gridworld for Fig.

Original: 2b.

Original: This task is a 15 × 15 gridworld with obstacles shown in Fig.

Original: 2b.

Original: Thereare four discrete actions, corresponding to moving to the four adjacent cells.

Original: The dynamics aredeterministic.

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

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

Original: We use γ = 0 .

Original: 9 and compute optimal policiesanalytically using value iteration.

Original: We implement the aliasing by averaging together the dynamicsfor each block of 3 × 3 states.

Original: 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 movingto state (3, 4)).

Original: To handle edge effects, we modiﬁed the averaged dynamics so that the agent couldnot exit the gridworld.

Original: We computed the classiﬁer analytically using the true and learned dynamicsmodels.

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

Original: This is not a failure of ourtheory, as the optimal policy would choose to never visit these states, but it presents a challenge foroptimization.

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

Original: Gridworld for Fig.

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

Original: 2a (Right) .

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

Original: With probability90%, the agent’s action is ignored and a random action is taken instead.

Original: The agent starts in the top-leftcell.

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

Original: We useγ =

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

Original: Gridworld for Fig.

Original: 3b.

Original: This task used the same dynamics as Fig.

Original: 2a.

Original: The one change is that theagent receives a reward of +1.0 at each time step and a reward of +10.0 at the goal state.

Original: We estimatethe more complex lower bound (Eq. 10) by also learning the discount factor.

Original: However, since wecurrently do not have a method for learning non-Markovian dynamics to fully optimize this lowerbound, we do not expect the lower bound to become tight.

Original: HalfCheetah-v2, Hopper-v2, Walker2d-v2 .

Original: These tasks are taken directly from the OpenAIbenchmark [7] without modiﬁcation.metaworld-drawer-open-v2 .

Original: This task is based on the drawer-open-v2 task from the

Original: Metaworld benchmark [50].

Original: To increase the difﬁculty of this task, we remove the rewardshaping term ( reward_for_caging ) and just optimize the reward for opening the drawer( reward_for_opening ).

Original: DClawPoseRandom-

Original: v

Original: DClawTurnRandom-v0 , DClawScrewFixed-v0 , DClawScrewRandom-v

Original: These tasks are taken directly from the ROBEL benchmark [1] without modiﬁcation.