Original: M ISMATCHED N O M ORE :
Revised: M ISMATCHED N O M ORE :

Original: J OINT M ODEL -P OLICY O PTIMIZATION FOR M ODEL -B ASED RL
Revised: J OINT M ODEL -P OLICY O PTIMIZATION FOR M ODEL -B ASED RL

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

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 RL or planning.
Revised: Many model-based reinforcement learning (RL) methods follow a similar template: ﬁt a model to previously observed data, and then use data from that model for RL or planning.

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

Original: MSE) are not necessarily better for control: an RL agent may seek out the small fraction of states where an accurate model makes mistakes, or it might act in ways that do not expose the errors of an inaccurate model.
Revised: MSE) are not necessarily better for control: an RL agent may seek out the small fraction of states where an accurate model makes mistakes, or it might act in ways that do not expose the errors of an inaccurate model.

Original: As noted in prior work, there is an objective mismatch: models are useful if they yield good policies, but they are trained to maximize their accuracy, rather than the performance of the policies that result from them.
Revised: As noted in prior work, there is an objective mismatch: models are useful if they yield good policies, but they are trained to maximize their accuracy, rather than the performance of the policies that result from them.

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

Original: This joint optimization mends the objective mismatch in prior work.
Revised: This joint optimization mends the objective mismatch in prior work.

Original: Our objective is a global lower bound on expected return, and this bound becomes tight under certain assumptions.
Revised: Our objective is a global lower bound on expected return, and this bound becomes tight under certain assumptions.

Original: The resulting algorithm (MnM) is conceptually similar to a GAN: a classiﬁer distinguishes between real and fake transitions, the model is updated to produce transitions that look realistic, and the policy is updated to avoid states where the model predictions are unrealistic.
Revised: The resulting algorithm (MnM) is conceptually similar to a GAN: a classiﬁer distinguishes between real and fake transitions, the model is updated to produce transitions that look realistic, and the policy is updated to avoid states where the model predictions are unrealistic.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

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

Original: Unfortunately, the accuracy of the learned model does not directly correlate with whether the model-based RL algorithm will receive high reward (Farahmand et al., 2017; Lambert et al., 2020).
Revised: Unfortunately, the accuracy of the learned model does not directly correlate with whether the model-based RL algorithm will receive high reward (Farahmand et al., 2017; Lambert et al., 2020).

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

Original: Alternatively, a model with large errors may yield a policy that attains high return if the model errors occur in states that the policy never visits.
Revised: Alternatively, a model with large errors may yield a policy that attains high return if the model errors occur in states that the policy never visits.

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

Original: Typical model-based methods train a model using data sampled from the real dynamics (e.g., using maximum likelihood), but apply these models by using data sampled from the learned dynamics (Deisenroth & Rasmussen, 2011; Hafner et al., 2019; Janner et al., 2019; Williams et al., 2017).
Revised: Typical model-based methods train a model using data sampled from the real dynamics (e.g., using maximum likelihood), but apply these models by using data sampled from the learned dynamics (Deisenroth & Rasmussen, 2011; Williams et al., 2017; Janner et al., 2019; Hafner et al., 2019).

Original: Prior work has identiﬁed this objective mismatch issue (Farahmand et al., 2017; Lambert et al., 2020;
Revised: Prior work has identiﬁed this objective mismatch issue (Farahmand et al., 2017; Luo et al., 2019;

Original: Luo et al., 2019): the model is trained using one objective, but the policy is trained using a different objective.
Revised: Lambert et al., 2020): the model is trained using one objective, but the policy is trained using a different objective.

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

Original: This paper aims to answer the following question: How should we train a dynamics model so that it produces high-return policies when used for model-based RL?
Revised: This paper aims to answer the following question: How should we train a dynamics model so that it produces high-return policies when used for model-based RL?

Original: In this paper, we propose a model-based RL algorithm where the model and policy are jointly optimized with respect to the same objective.
Revised: In this paper, we propose a model-based RL algorithm where the model and policy are jointly optimized with respect to the same objective.

Original: This objective is a lower bound on the expected return under the true environment dynamics, and the bound becomes tight under certain assumptions.
Revised: Our objective is a lower bound on the expected return under the true environment dynamics; a slightly more complicated version of this lower bound becomes tight under certain assumptions.

Original: Structurally, our algorithm resembles a generative adversarial network (a GAN), in that the model is trained using a discriminator that distinguishes between real and fake transitions.
Revised: Structurally, our algorithm resembles a generative adversarial network (a GAN), in that the model is trained using a discriminator that distinguishes between real and fake transitions.

Original: Crucially, this same discriminator isalso included in the objective for the policy, and both the model and policy are jointly trained to maximize reward and minimize discriminator accuracy.
Revised: This same discriminator is included in the objective for the policy, and both the model and policy are jointly trained to maximize reward and minimize discriminator accuracy.

Original: Thus, the model and policy cooperate to produce realistic and high-reward trajectories against an adversary that attempts to classify whether the trajectory is realistic.
Revised: Thus, the model and policy cooperate to produce realistic and high-reward trajectories.

Original: Our method stands in contrast to standard model-based RL methods, where it is more common to pit the model against the policy (Bagnell et al., 2001; Nilim & El Ghaoui, 2003; Ross & Bagnell, 2012).
Revised: Our method stands in contrast to standard model-based RL methods, where it is more common to pit the model against the policy (Bagnell et al., 2001; Nilim & El Ghaoui, 2003; Ross & Bagnell, 2012).

Original: Aninteresting consequence of maximizing the lower bound is that the dynamics model does not learn the true dynamics, but rather learns “optimistic” dynamics that facilitate exploration.
Revised: An consequence of maximizing the lower bound is that the dynamics model does not learn the true dynamics, but rather learns optimistic dynamics that facilitate exploration.

Original: We illustrate the method in Figure 1.

Original: The main contribution of this work is an algorithm, Mismatched no More (MnM) for model-based RL that provably maximizes a lower bound on expected reward.
Revised: The main contribution of this work is an algorithm, Mismatched no More (MnM), for model-based RL that provably maximizes a lower bound on expected reward.

Original: Importantly, this bound becomes tight at optimality under certain assumptions.
Revised: Importantly, this bound becomes tight at optimality under certain assumptions.

Revised: To the best of our knowledge, this is the ﬁrst model-based RL objective that is a global lower bound on expected return, and that involves optimizing the model and policy using the same objective.

Original: Our algorithm has the unique property of jointly optimizing the policy and model using the same objective.
Revised: Our algorithm has the unique property of jointly optimizing the policy and model using the same objective.

Original: Across a range of tasks, we demonstrate that our method is competitive with prior state-of-the-art methods on benchmark tasks; on certain hard exploration tasks, our method outperforms prior methods based on maximum likelihood estimation.
Revised: Across a range of tasks, we demonstrate that our method is competitive with prior state-of-the-art methods on benchmark tasks; on certain hard exploration tasks, our method outperforms prior methods based on maximum likelihood estimation.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Model-based RL methods typically ﬁt a dynamics model to observed transitions, and then apply an
Revised: Model-based RL methods typically ﬁt a dynamics model to observed transitions, and then apply an

Original: RL method to that learned model.
Revised: RL method to that learned model.

Original: Most of these methods use maximum likelihood to ﬁt the dynamics model, and then use RL to maximize the expected return under samples from that model (Chua et al., 2018; Deisenroth & Rasmussen, 2011; Hafner et al., 2019; Janner et al., 2019; Williams et al., 2017).
Revised: Most of these methods use maximum likelihood to ﬁt the dynamics model, and then use RL to maximize the expected return under samples from that model (Deisenroth & Rasmussen, 2011; Williams et al., 2017; Chua et al., 2018; Hafner et al., 2019; Janner et al., 2019).

Original: The observation that this maximum likelihood objective is not aligned with the RL objective has been noted in prior work (Farahmand et al., 2017;Lambert et al., 2020; Luo et al., 2019; Talvitie, 2014;
Revised: The observation that this maximum likelihood objective is not aligned with the RL objective has been noted in prior work (Ziebart, 2010; Talvitie, 2014; Farahmand et al., 2017; Luo et al., 2019; Lambert et al., 2020).

Original: Ziebart, 2010) where it is referred to as the objective mismatch problem.
Revised: This issue is referred to as the objective mismatch problem: the model and policy (or planner) are optimized using different objectives. This problem arises in almost all model-based RL approaches, including those that train the model to predict the value function (Oh et al., 2017; Schrittwieser et al., 2020) or that perform planning (Chua et al., 2018; Schrittwieser et al., 2020).

Original: One strategy for mitigating this problem is to modify the model training to improve model accuracy under multi-step rollouts (Asadi et al., 2018; 2019; Farahmand et al., 2017; Joseph et al., 2013; Talvitie, 2014; Venkatraman et al., 2016).
Revised: One strategy for mitigating this problem is to modify the model training to improve model accuracy under multi-step rollouts (Joseph et al., 2013; Talvitie, 2014; Venkatraman et al., 2016; Farahmand et al., 2017; Asadi et al., 2018; 2019).

Original: A second strategy is penalize the policy for taking transitions where the model is inaccurate (Kidambi et al., 2020; Luo et al., 2019; Sorg et al., 2010; Yu et al., 2020b; 2021).
Revised: A second strategy is penalize the policy for taking transitions where the model is inaccurate (Sorg et al., 2010; Kidambi et al., 2020; Yu et al., 2020b; 2021; Luo et al., 2019).

Original: Similar to all these prior methods, our approach will also use a modiﬁed reward function to train the policy, but it will also modify how the model is trained such that the model and policy optimize the same objective.
Revised: Similar to all these prior methods, our approach will also use a modiﬁed reward function to train the policy, but it will also modify how the model is trained such that the model and policy optimize the same objective.

Original: Like Luo et al. 
Revised: A third strategy is to directly optimize the model such that it produces good policies (Okada et al., 2017; Amos et al., 2018; Srinivas et al., 2018; D’Oro et al., 2020; Nikishin et al., 2021), as theoretically analyzed in Grimm et al. While our aim is the same as these prior methods, our approach will not require differentiating through unrolled model updates or optimization procedures.

Original: (2019), we show that our objective is a lower bound on  expected return, but our bound will hold for all dynamics models and policies (i.e., globally), rather than a narrow region around previous models and policies (i.e., locally).
Revised: Our work builds on prior work that proposes model-based RL objectives that are lower bounds on the true, expected returns. Kearns & Singh (2002) provide a lower bound that holds globally, but is only computable in tabular settings.

Original: A third strategy is to directly optimize the model such that it produces good policies (Amos et al., 2018; Okada et al., 2017; Srinivas et al., 2018). While our aim is the same as these prior methods, our approach will not require differentiating through unrolled model updates or optimization procedures.
Revised: Luo et al. (2019) provide a lower bound that can be efﬁciently estimated, but which only holds for nearby policies and models. Our lower bound combines the strengths of these prior works, providing a lower bound that holds globally and can be efﬁciently estimated in MDPs with continuous states and actions.

Original:  Our theoretical derivation builds on prior work that casts model-based RL as a two-player game tween a model-player and a policy-player (Bagnell et al., 2001; Nilim & El Ghaoui, 2003; Rajeswaran et al., 2020; Ross et al., 2011).
Revised: Unlike any lower bounds from prior work, ours mends the objective mismatch problem. Our theoretical derivation builds on prior work that casts model-based RL as a two-player game tween a model-player and a policy-player (Bagnell et al., 2001; Nilim & El Ghaoui, 2003; Ross et al., 2011; Rajeswaran et al., 2020).

Original: However, whereas prior work pits model and policy compete against one another, our formulation will result in a cooperative game, wherein the model and policy players cooperate in optimizing the same objective (a lower bound on the expected return).
Revised: However, whereas prior work pits model and policy compete against one another, our formulation will result in a cooperative game, wherein the model and policy players cooperate in optimizing the same objective (a lower bound on the expected return).

Original: Our approach, though structurally resembling a GAN, is different from prior work that simply replaces a maximum likelihood model with a GAN model (Bai et al., 2019;
Revised: Our approach, though structurally resembling a GAN, is different from prior work that simply replaces a maximum likelihood model with a GAN model (Bai et al., 2019;

Original: Chen et al., 2020; Kurutach et al., 2018).
Revised: Chen et al., 2020; Kurutach et al., 2018).

Original: The most similar prior work is VMBPO (Chow et al., 2020).
Revised: The most similar prior work is VMBPO (Chow et al., 2020).

Original: The mechanics of our method are similar, also learning a dynamics model using a classiﬁer that distinguishes real versus generated rollouts.
Revised: The mechanics of our method are similar, also learning a dynamics model using a classiﬁer that distinguishes real versus generated rollouts.

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

Original: This different objective can be expressed as the expected return plus the variance of the return, so VMBPO has the undesirable property of preferring policies that receive slight lower return if the variance of the return is much larger (see Appendix A.1).
Revised: This different objective can be expressed as the expected return plus the variance of the return, so VMBPO has the undesirable property of preferring policies that receive slight lower return if the variance of the return is much larger (see Appendix A.1).

Original: Indeed, most of the components of our method, including classiﬁers and
Revised: Indeed, most of the components of our method, including classiﬁers and

Original: GAN-like models, have been used in prior work, main contribution of our paper is a precise recipe for combining these components in a way that provably maximizes expected return.
Revised: GAN-like models, have been used in prior work, main contribution of our paper is a precise recipe for combining these components in a way that provably maximizes expected return.

Original:  A U NIFIED O BJECTIVE FOR M ODEL -B ASED RL
Revised:  A U NIFIED O BJECTIVE FOR M ODEL -B ASED 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 to learn a control policy π θ ( a t | s t ) with parameters θ that maximizes the expected discounted return:
Revised: Our aim is to learn a control policy π θ ( a t | s t ) with parameters θ that maximizes the expected discounted return:

Original: [Equation]
Revised: [Equation]

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

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

Original: We then deﬁne R ( τ ) (cid:44) (cid:80) ∞
Revised: We then 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: Finally, we deﬁne two distributions over trajectories.
Revised: Finally, we deﬁne two distributions over trajectories.

Original: First, p π ( τ ) is the distribution over trajectories when policy π θ interacts with dynamics p ( s (cid:48) | s, a ) ; q π ( τ ) is the distribution over trajectories when policy π θ interacts with the learned dynamics q θ ( s t+1 | s t , a t ) : p π ( τ ) =
Revised: First, p π ( τ ) is the distribution over trajectories when policy π θ interacts with dynamics p ( s (cid:48) | s, a ) ; q π ( τ ) is the distribution over trajectories when policy π θ interacts with the learned dynamics q θ ( s t

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

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

Revised: [Equation]

Original: π θ ( a t | s t ) .
Revised: = p [Equation] [Equation]

Original: Desiderata.
Revised: Desiderata.

Original: Our aim is to design an objective L ( θ ) that can be used to jointly optimize both the policy  π θ ( a t | s t )  and the dynamics model, q ( s t +1 | s t , a t ) . Such joint optimization would avoid the objective mismatch problem from prior work. This objective should be a lower bound on the expected return.
Revised: Our aim is to design an objective L ( θ ) that can be used to jointly optimize both the policy ( π θ ( a t | s t ) ) and the dynamics model ( q ( s t +1 | s t , a t ) ), and which is a lower bound on the expected return in the true environment.

Original: Moreover, the lower bound should become tight at optimality.

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: Our objective will be the policy’s reward when interacting with the learned model, but using a different reward function.
Revised: Our objective will be the policy’s reward when interacting with the learned model, but using a different reward function.

Original: The new reward function augments the task reward with an additional term that measures the difference between the learned model and the real environment.
Revised: The new reward function augments the task reward with an additional term that measures the difference between the learned model and the real environment.

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

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: This objective should be understood as maximizing an augmented reward under the learned dynamics.
Revised: This objective maximizes an augmented reward under the learned dynamics.

Original: The augmented reward function ˜ r penalizes the policy for taking transitions that are unlikely under the true dynamics model, similar to prior work (Eysenbach et al., 2020a; Yu et al., 2021).
Revised: The augmented reward function ˜ r penalizes the policy for taking transitions that are unlikely under the true dynamics model, similar to prior work (Eysenbach et al., 2020a; Yu et al., 2021).

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

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

Original: For the policy, this optimization entails performing RL to maximize the modiﬁed reward using samples from the learned model; the only difference from prior work is the modiﬁcation to the reward function.
Revised: For the policy, this optimization entails performing RL to maximize the modiﬁed reward using samples from the learned model; the only difference from prior work is the modiﬁcation to the reward function.

Original: Training the dynamics model using this objective is very different from standard maximum likelihood training, and instead resembles a GAN.
Revised: Training the dynamics model using this objective is very different from standard maximum likelihood training, and instead resembles a GAN.

Original: The model is optimized to sample trajectories that both have high reward (i.e., log r is large) and are similar to real dynamics (i.e., log pq is large).
Revised: The model is optimized to sample trajectories that both have high reward (i.e., log r is large) and are similar to real dynamics (i.e., log pq is large).

Original: This objective differs from VMBPO (Chow et al., 2020) by taking the log( · ) of the original reward functions; our experiments demonstrate that excluding this component invalidates our lower bound and results in learning suboptimal policies.
Revised: This objective differs from VMBPO (Chow et al., 2020) by taking the log( · ) of the original reward functions; our experiments demonstrate that excluding this component invalidates our lower bound and results in learning suboptimal policies.

Original: Our objective has two properties that make it particularly useful.
Revised: Our objective has two properties that make it particularly useful.

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

Original: Note that this is very different from prior work, where training the model to be more accurate (increase likelihood) can decrease the policy’s expected return under that model.
Revised: Note that this is very different from prior work, where training the model to be more accurate (increase likelihood) can decrease the policy’s expected return under that model.

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

Original: This property gives us a guarantee on how well the learned policy will perform when deployed on the real environment.
Revised: This property gives us a guarantee on how well the learned policy will perform when deployed on the real environment.

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 A.3.
Revised: The proof is presented in Appendix A.3.

Original: Note that the expected return under the learned model, which most prior model-based RL methods use to train the policy, is not a lower bound on the expected return.
Revised: Note that the expected return under the learned model, which most prior model-based RL methods use to train the policy, is not a lower bound on the expected return.

Original: Unlike the objective proposed in (Luo et al., 2019), which is also a lower bound, our bound bound holds for all dynamics functions and policies, not just those in a small region around earlier dynamics functions and policies.

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

Original: How tight is this lower bound?
Revised: and efﬁciently-computable (unlike Kearns & Singh (2002))

Original: Showing that a boundis tight is crucial for making the claim that better dynamics models will produce better policies. Our experiments show that optimizing this  lower bound yields policies that achieve high reward across a wide range of tasks.
Revised: lower bound for model-based RL. Our experiments show that optimizing this ﬁrst lower bound yields policies that achieve high reward across a wide range of tasks, even though this lower bound may not be tight.

Original: Tightening the lower bound.
Revised: Tightening the lower bound.

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

Original: This new lower bound will be more complex than the one introduced above and we have not yet successfully designed an algorithm for maximizing it.
Revised: This new lower bound will be more complex than the one introduced above and we have not yet successfully designed an algorithm for maximizing it.

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

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

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

Original: Similar learned discount factors have been studied in previous work on model-free RL (Rudner et al., 2021).
Revised: Similar learned discount factors have been studied in previous work on model-free RL (Rudner et al., 2021).

Original: We deﬁne this objective as follows:
Revised: We deﬁne this objective as follows:

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: θ ( t ) (cid:19) , and Γ θ ( t ) =
Revised: and Γ θ ( t ) =

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

Original: =
Revised: =

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

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

Original: Lemma 3.2.
Revised: Lemma 3.2.

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

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

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

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

Original: [Equation]
Revised: [Equation]

Original: The proof is presented in Appendix 3.2.
Revised: The proof is presented in Appendix A.4.

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

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

Original: We leave an empirical investigation of the more complex lower bound L γ  and non-Markovian models to future work.
Revised: We leave an empirical investigation of the more complex lower bound, L γ , and non-Markovian models to future work.

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 not equal to the true environment dynamics.
Revised: In stochastic environments, the dynamics model that optimizes this lower bound is not equal to the true environment dynamics.

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

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

Original: [Equation]
Revised: p ( τ ) R ( τ ) (cid:82) p ( τ (cid:48) ) R ( τ (cid:48) ) dτ (cid:48)

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

Original: We hypothesize that the optimism in the dynamics model will accelerate policy optimization, a hypothesis we test in Sec. 5.1.
Revised: We hypothesize that the optimism in the dynamics model will accelerate policy optimization, a hypothesis we 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 Lemma 3.1?

Original: This is not quite what our method does.
Revised: This is not quite what our method does.

Original: Rather, our method estimates the augmented reward using the optimistic dynamics model, and the reward augmentation compensates for the optimism in the dynamics model.
Revised: Rather, our method estimates the augmented reward using the optimistic dynamics model, and the reward augmentation compensates for the optimism in the dynamics model.

Original:  M ISMATCHED N O M ORE
Revised:  M ISMATCHED N O M ORE

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

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

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

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

Original: We address this challenge by learning a classiﬁer (Sec. 4.1), and then describe the precise update rules for the policy, dynamics model, and classiﬁer (Sec. 4.2). 
Revised: We address this challenge by learning a classiﬁer (Sec. 4.1), and then describe the precise update rules for the policy, dynamics model, and classiﬁer (Sec. 4.2). [Figure] [Figure]

Original: 4. E STIMATING THE A UGMENTED R EWARD F UNCTION
Revised: E STIMATING THE A UGMENTED R EWARD F UNCTION

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

Original: This approach is similar to GANs (Goodfellow et al., 2014) and similar to prior work in RL (Eysenbach et al., 2020a; Yu et al., 2021).
Revised: This approach is similar to GANs (Goodfellow et al., 2014) and similar to prior work in RL (Eysenbach et al., 2020a; Yu et al., 2021).

Original: We use C φ ( s t , a t , s t +1 ) ∈ [0 , 1] to denote the classiﬁer.
Revised: We use C φ ( s t , a t , s t +1 ) ∈ [0 , 1] to denote the classiﬁer.

Original: For an optimal classiﬁer, we can use the classiﬁer’s predictions to estimate the augmented reward function:
Revised: For an optimal classiﬁer, we can use the classiﬁer’s predictions to estimate the augmented reward function:

Original: [Equation]
Revised: [Equation]

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

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

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

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

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

Original: 4. U PDATING THE M ODEL , P OLICY , AND C LASSIFIER
Revised: 4. U PDATING THE M ODEL , P OLICY , AND C LASSIFIER

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

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

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

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

Original: [Equation]
Revised: [Equation]

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

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

Original: [Equation]
Revised: [Equation]

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

Original: ˜ r ( s t , a t , s model t +1 ) 
Revised: ˜ r ( s t , a t , s model t +1 ) +

Original: + γ E π ( a t +1 | s model t +1 ) (cid:2) Q ψ ( s model t +1 , a t +1 ) (cid:3) .
Revised: γ E π [Equation] a t +1 | s model t +1 ) (cid:2) Q ψ ( s model t +1 , a t +1 ) (cid:3) .

Original: The augmented reward, ˜ r , is estimated using the learned classiﬁer.
Revised: The augmented reward, ˜ r , is estimated using the learned classiﬁer.

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

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

Original: The policy is trained to maximize the expected future (augmented) return, as estimated by the Q function:
Revised: The policy is trained to maximize the expected future (augmented) return, as estimated by 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: Following prior work (Fujimoto et al., 2018), we maintain two Q functions and two target Q functions, use the minimum of the two target Q functions to compute the TD target.
Revised: Following prior work (Fujimoto et al., 2018), we maintain two Q functions and two target Q functions, use the minimum of the two target Q functions to compute the TD target.

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

Original: Updating the classiﬁer.
Revised: Updating the classiﬁer.

Original: We train the classiﬁer to distinguish real versus model transitions using the standard cross entropy loss:
Revised: We train the classiﬁer to distinguish real versus model transitions using the standard cross entropy loss:

Original: [Equation]
Revised: [Equation]

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

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

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

Original: [Equation]
Revised: [Equation]

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

Original: This approximation is standard in prior work on GANs (Goodfellow et al., 2014) and adversarial inference (Donahue et al., 2016; Dumoulin et al., 2016).
Revised: This approximation is standard in prior work on GANs (Goodfellow et al., 2014) and adversarial inference (Donahue et al., 2016; Dumoulin et al., 2016).

Original: The procedure for optimizing the dynamics model and the classiﬁer resembles a GAN (Goodfellow et al., 2014): the classiﬁer is optimized to distinguish real transitions from model transitions, and the model is updated to fool the classiﬁer (and increase rewards).
Revised: The procedure for optimizing the dynamics model and the classiﬁer resembles a GAN (Goodfellow et al., 2014): the classiﬁer is optimized to distinguish real transitions from model transitions, and the model is updated to fool the classiﬁer (and increase rewards).

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

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

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

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

Original: Algorithm summary.
Revised: Algorithm summary.

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

Original: 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, create an additional classiﬁer network.
Revised: First, create an additional classiﬁer network.

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

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

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

Original: Following prior work (Janner et al., 2019), we learn an additional model to predict the true environment.
Revised: Following prior work (Janner et al., 2019), we learn a neural network to predict the true environment rewards.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We present two sets of experiments.
Revised: We present two sets of experiments.

Original: Our ﬁrst set of experiments studies the importance of different components of MnM .
Revised: Our ﬁrst set of experiments studies the importance of different components of MnM .

Original: Second, we study how MnM compares with prior model-based RL algorithms on challenging, robotic control tasks.
Revised: Second, we study how MnM compares with prior model-based RL algorithms on challenging, robotic control tasks.

Original: To compare policies learned by different algorithms, we will evaluate the policies using the true environment dynamics, not the learned dynamics model.
Revised: To compare policies learned by different algorithms, we will evaluate the policies using the true environment dynamics, not the learned dynamics model.

Original: 5. U NDERSTANDING THE L OWER B OUND AND THE L EARNED D YNAMICS
Revised: 5. U NDERSTANDING THE L OWER B OUND AND THE L EARNED D YNAMICS

Original: In this section, we begin by studying the seemingly contradictory attributes of our method: optimistic dynamics and pessimistic policies, and end by conﬁrming that together these components optimize an increasingly tight lower bound on the expected return.
Revised: In this section, we begin by studying the seemingly contradictory attributes of our method: optimistic dynamics and pessimistic policies, and end by conﬁrming that together these components optimize an increasingly tight lower bound on the expected return.

Original: Our theory suggests that MnM should work best in settings with stochastic dynamics and challenging exploration requirements, as the dynamics model should tilt the true dynamics to make the stochasticity more favorable for solving the task.
Revised: Our theory suggests that MnM should work best in settings with stochastic dynamics and challenging exploration requirements, as the dynamics model should tilt the true dynamics to make the stochasticity more favorable for solving the task.

Original: We use a 10x10 gridworld with highly stochastic dynamics and a sparse reward function.
Revised: We use a 10x10 gridworld with highly stochastic dynamics and a sparse reward function.

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

Original: 2a (Left) , show that MnM outperforms both
Revised: 2a (Left) , show that MnM outperforms both

Original: Q-learning and VMBPO.
Revised: Q-learning and VMBPO.

Original: In line with our theory, the dynamics learned by MnM (Fig.
Revised: In line with our theory, the dynamics learned by MnM (Fig.

Original: 2a (Right) ) alter
Revised: 2a (Right) ) alter

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: the environment stochasticity to lead the agent towards the goal, increasing the probability of collecting high-reward experience.
Revised: the environment stochasticity to lead the agent towards the goal, increasing the probability of collecting high-reward experience.

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

Original: While VMBPO also learns optimistic dynamics, it omits log-transformation of the reward function (which encourages pessimistic behavior), a difference that has a large effect on this task.
Revised: While VMBPO also learns optimistic dynamics, it omits log-transformation of the reward function (which encourages pessimistic behavior), a difference that has a large effect on this task.

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

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

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

Original: We will use the gridworld shown in Fig. 2b, which contains obstacles that occur at a ﬁner resolution than the model can detect.
Revised: We will use the gridworld shown in Fig. 2b, which contains obstacles that occur at a ﬁner resolution than the model can detect.

Original: When the “low resolution” dynamics model makes predictions for states near the obstacle, it will average together some states with obstacles and some states without obstacles.
Revised: When the “low resolution” dynamics model makes predictions for states near the obstacle, it will average together some states with obstacles and some states without obstacles.

Original: Thus, the model will (incorrectly) predict that the agent always has some probability of moving in each direction, even if that direction is actually blocked by an obstacle.
Revised: Thus, the model will (incorrectly) predict that the agent always has some probability of moving in each direction, even if that direction is actually blocked by an obstacle.

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

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

Original: We then test the importance of the logarithmic transformation by comparing MnM to VMBPO, which includes the classiﬁer term but omits the logarithmic transformation.
Revised: We then test the importance of the logarithmic transformation by comparing MnM to VMBPO, which includes the classiﬁer term but omits the logarithmic transformation.

Original: We hypothesize that VMBPO’s deviation from our lower bound will cause it to exhibit risk seeking behavior.
Revised: We hypothesize that VMBPO’s deviation from our lower bound will cause it to exhibit risk seeking behavior.

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 (top) where numbers indicate the reward at each state.
Revised: 3a (top) where numbers indicate the reward at each state.

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

Original: While MnM
Revised: While MnM

Original: MnM learns the reward-maximizing strategy, VMBPO learns a policy that goes to the right state and receives lower returns (with much higher variance).
Revised: MnM learns the reward-maximizing strategy, VMBPO learns a policy that goes to the right state and receives lower returns (with much higher variance).

Original: Finally, we study how the MnM objective compares to alternative objectives.
Revised: Finally, we study how the MnM objective compares to alternative objectives.

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

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

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

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

Original: Ablations of MnM to omit the reward augmentation or even just the log transformation (i.e., VMBPO) overestimate the expected return.
Revised: Ablations of MnM to omit the reward augmentation or even just the log transformation (i.e., VMBPO) overestimate the expected return.

Original: 5. C OMPARISONS O N R OBOTICS T ASKS
Revised: 5. C OMPARISONS O N R OBOTICS T ASKS

Original: Our second set of experiments use a range of continuouscontrol robotic tasks to answer two questions
Revised: Our next experiments use continuous-control robotic tasks to answer two questions.

Original: We  investigate whetherMnM performs at least comparably with prior work.
Revised: We ﬁrst investigate whether performs at least comparably with prior work.

Original: We then study tasks with sparse rewards and more challenging exploration, where we suspect the optimistic dynamics learned by MnM may be beneﬁcial.
Revised: We then study tasks with sparse rewards and more challenging exploration, where we suspect the optimistic dynamics learned by MnM may be beneﬁcial.

Original: We illustrate a subset of the environments in Fig.
Revised: We illustrate a subset of the environments in Fig.

Original: 4 and include implementation details in Appendix C.2.
Revised: 4 and include implementation details in Appendix C.2.

Original: One detail of note is that we omit the reward augmentation  for MnM during these experiments, as it hinders exploration leading to lower returns.
Revised: One detail of note is that we omit the reward augmentation (Eq. 3) for MnM during these experiments, as it hinders exploration leading to lower returns.

Original: We use MBPO (Janner et al., 2019) as a baseline for model-based RL because it achieves state-ofthe-art results and is a prototypical example of model-based RL algorithms that use maximum likelihood models.
Revised: We use MBPO (Janner et al., 2019) as a baseline for modelbased RL because it achieves state-of-the-art results and is a prototypical example of model-based RL algorithms that use maximum likelihood models.

Original: Our ﬁrst comparison uses three locomotion tasks from the OpenAI Gym benchmark (Brockman et al., 2016), which has become the standard benchmark for model-based RL algorithms.
Revised: Our ﬁrst comparison uses three locomotion tasks from the OpenAI Gym benchmark (Brockman et al., 2016), which has become the standard benchmark for model-based RL algorithms.

Original: Thes tasks have dense rewards and pose no signiﬁcant exploration challenge, so we do not expect MnM to outperform prior methods.
Revised: Thes tasks have dense rewards and pose no signiﬁcant exploration challenge, so we do not expect MnM to outperform prior methods.

Original: The results (Fig. 5 (top left) ) show that MnM performs roughly on par with MBPO.
Revised: The results (Fig. 5 (top left) ) show that MnM performs roughly on par with MBPO.

Original: Tasks with sparse rewards, complex contact dynamics, and those requiring hard exploration often present a challenge for model-based RL algorithms, which are liable to exploit inaccuracies in the learned dynamics model.
Revised: Tasks with sparse rewards, complex contact dynamics, and those requiring hard exploration often present a challenge for model-based RL algorithms, which are liable to exploit inaccuracies in the learned dynamics model.

Original: Our next experiment evaluates MnM on control tasks used in prior work that demonstrate these properties.
Revised: Our next experiment evaluates MnM on control tasks used in prior work that demonstrate these properties.

Original: These tasks the sparse-reward cartpole task from the DM Control benchmark (Tassa et al., 2018), four manipulation tasks from the Metaworld benchmark (Yu et al., 2020a), and four dexterous manipulation tasks from the ROBEL benchmark (Ahn et al., 2020).
Revised: These tasks the sparse-reward cartpole task from the DM Control benchmark (Tassa et al., 2018), four manipulation tasks from the Metaworld benchmark (Yu et al., 2020a), and four dexterous manipulation tasks from the ROBEL benchmark (Ahn et al., 2020).

Original: See
Revised: See

Original: [Figure]
Revised: [Figure]

Original: Appendix C for experiment details.
Revised: Appendix C for experiment details.

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

Original: 5 indicate that the complex environment dynamics of these tasks can cause prior model-based algorithms (MBPO) to perform worse than model-free algorithms (SAC), both in terms of asymptotic performance and sample complexity.
Revised: 5 indicate that the complex environment dynamics of these tasks can cause prior model-based algorithms (MBPO) to perform worse than model-free algorithms (SAC), both in terms of asymptotic performance and sample complexity.

Original: Nonetheless, we observe that MnM frequently outperforms all prior methods and, more importantly, it consistently does well across all tasks.
Revised: Nonetheless, we observe that MnM frequently outperforms all prior methods and, more importantly, it consistently does well across all tasks.

Original: We include ablation experiments in Appendix B.1.
Revised: We include ablation experiments, including a comparison to

Revised: VAML (Farahmand et al., 2017), in Appendix B.

Original: To investigate the beneﬁts of MnM over prior model-based methods, we logged the Q values throughout training and visualized them for the metaworld-drawer-open-v2 task in
Revised: To investigate the beneﬁts of MnM over prior model-based methods, we logged the Q values throughout training and visualized them for the metaworld-drawer-open-v2 task in

Original: Fig.
Revised: Fig.

Original: For fair comparison, we use Q values corresponding to just the task reward, omitting the augmented reward used by MnM. Fig.
Revised: For fair comparison, we use Q values corresponding to just the task reward, omitting the logarithmic transformation and classiﬁer term typically used by MnM. Fig.

Original: 6 shows that MnM yields Q values that are more accurate and more stable than MBPO.
Revised: 6 shows that MnM yields Q values that are more accurate and more stable than MBPO.

Original: This ﬁgure suggests that MBPO may be exploiting inaccuracies in the learned model.
Revised: This ﬁgure suggests that MBPO may be exploiting inaccuracies in the learned model.

Original: Finally, we visualize the dynamics learned by MnM on two robotic control tasks.
Revised: Finally, we visualize the dynamics learned by MnM on two robotic control tasks.

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

Original: However, our implementation relies on function approximation (neural networks) to learn the dynamics, and the limited capacity of function approximators makes otherwisedeterministic dynamics appear stochastic.
Revised: However, our implementation relies on function approximation (neural networks) to learn the dynamics, and the limited capacity of function approximators makes otherwise-deterministic dynamics appear stochastic.

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

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

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

Original: We expect that the implicit stochasticity caused by function approximation to be especially important for real-world tasks, where the complexity of the real dynamics often dwarfs the capacity of the learned dynamics model.
Revised: We expect that the implicit stochasticity caused by function approximation to be especially important for real-world tasks, where the complexity of the real dynamics often dwarfs the capacity of the learned dynamics model.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: The main contribution of this paper is an approach to model-based RL where the policy and dynamics model are jointly optimized using the same objective.
Revised: The main contribution of this paper is an approach to model-based RL where the policy and dynamics model are jointly optimized using the same objective.

Original: Unlike prior work, our objective is a global lower bound on the standard expected return objective.
Revised: Unlike prior work, our objective is a global lower bound on the standard expected return objective.

Original: Our approach not only tells users how to train their dynamics model, but also guarantees to them that updating their model (using our objective) will result in a better policy.
Revised: Our approach not only tells users how to train their dynamics model, but also guarantees to them that updating their model (using our objective) will result in a better policy.

Original: We therefore believe that this joint optimization will ease and accelerate the design of future model-based RL algorithms.
Revised: We therefore believe that this joint optimization will ease and accelerate the design of future model-based RL algorithms.

Original: We suspect that the tools presented in this paper may prove useful for solving tasks that require extensive exploration or long-horizon planning.
Revised: We suspect that the tools presented in this paper may prove useful for solving tasks that require extensive exploration or long-horizon planning.

Original: For example, models learned by MnM may be more capable of sampling realistic long trajectories, opening a door to future model-based algorithms that apply on-policy algorithms to sampled experience.

Original: Limitations.
Revised: Limitations.

Original: The main limitations of our method are complexity and stability.

Original: Compared with methods based on maximum likelihood, our method requires learning an additional classiﬁer and balancing the capacity of that classiﬁer against the capacity of the dynamics model.
Revised: Compared with methods based on maximum likelihood, our method requires learning an additional classiﬁer and balancing the capacity of that classiﬁer against the capacity of the dynamics model.

Original: In Appendix C, we describe implementation details we found useful for stabilizing training of our method.
Revised: Unlike VMBPO, our method does not have a single hyperparameter to control the gap between the proposed objective and the expected return objective.

Original: Ethics Statement.
Revised: Ethics Statement.

Original: This paper works to advance model based reinforcement learning, which can enable data driven policies to learn from smaller datasets than their model-free counterparts.
Revised: This paper works to advance model based reinforcement learning, which can enable data driven policies to learn from smaller datasets than their model-free counterparts.

Original: In doing so, this expands the range of applications for which reinforcement learning can be used.
Revised: In doing so, this expands the range of applications for which reinforcement learning can be used.

Original: Many of these low data regimes, which can include things such as health care, require not just data efﬁciency but also safe online ﬁne-tuning.
Revised: Many of these low data regimes, which can include things such as health care, require not just data efﬁciency but also safe online ﬁne-tuning.

Original: It is important to note that although our method optimizes a lower-bound on the expected returns in theory, this lower-bound utilizes function approximators, which do not come with guarantees in practice.
Revised: It is important to note that although our method optimizes a lower-bound on the expected returns in theory, this lower-bound utilizes function approximators, which do not come with guarantees in practice.

Original: Reproducibility.
Revised: Reproducibility.

Original: Our appendix contains proofs for all theoretical claims, the hyperparameters and implementation details of each algorithm used, ablation experiments illustrating the expected behavior of our method under varying circumstances of interest, implementation details that we did and did not ﬁnd helpful, and ﬁnally the necessary details to recreate each environment used in the paper.
Revised: Our appendix contains proofs for all theoretical claims, the hyperparameters and implementation details of each algorithm used, ablation experiments illustrating the expected behavior of our method under varying circumstances of interest, implementation details that we did and did not ﬁnd helpful, and ﬁnally the necessary details to recreate each environment used in the paper.

Original: Furthermore, we will be publicly releasing the code for our didactic experiments
Revised: Furthermore, we will be publicly releasing the code for our didactic experiments

Original: A P ROOFS AND A DDITIONAL A NALYSIS A. VMBPO M AXIMIZES AN U PPER B OUND ON R ETURN
Revised: A P ROOFS AND A DDITIONAL A NALYSIS A. VMBPO M AXIMIZES AN U PPER B OUND ON R ETURN

Original: While MnM aims to maximize the (log) of the expected return [Equation] VMBPO aims to maximize the expected exponentiated return:
Revised: While MnM aims to maximize the (log) of the expected return, VMBPO aims to maximize the expected exponentiated return:

Original: [Equation]
Revised: [Equation]

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

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

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

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

Original: Rather, it corresponds to maximizing a sum of the expected return an the variance of the return (Mihatsch & Neuneier, 2002, Page 272):
Revised: Rather, it corresponds to maximizing a sum of the expected return and the variance of the return (Mihatsch & Neuneier, 2002, Page 272):

Original: [Equation]
Revised: [Equation]

Original: Thus, in environments with stochastic dynamics or rewards (e
Revised: Thus, in environments with stochastic dynamics or rewards (e

Original: [Equation]
Revised: [Equation]

Original: g., the didactic example in Fig. 3)
Revised: g., the didactic example in Fig. 3)

Original: [Equation]
Revised: [Equation]

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

Original: We note that the expected exponentiated return is an upper bound on the expected return:
Revised: We note that the expected exponentiated return is an upper bound on the expected return:

Original: [Equation]
Revised: [Equation]

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

Original: The bound holds with a strict inequality in almost all MDPs.
Revised: The bound holds with a strict inequality in almost all MDPs.

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

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

Original: A. H ELPER L EMMAS
Revised: A. H ELPER L EMMAS

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

Original: Lemma A.1.
Revised: Lemma A.1.

Original: Deﬁne p
Revised: Deﬁne p

Original: ( H ) =
Revised: ( H ) =

Original: G EOM (1 − γ ) as the geometric distribution.
Revised: G EOM (1 − γ ) as the geometric distribution.

Original: Let discount factor γ ∈ (0 , 1) and random variable x t be given.
Revised: Let discount factor γ ∈ (0 , 1) and random variable x t be given.

Original: Then the following identity holds:
Revised: Then the following identity holds:

Original: [Equation]
Revised: [Equation]

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

Original: Proof.
Revised: Proof.

Original: [Equation]
Revised: [Equation]

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

Original: Lemma A.2.
Revised: Lemma A.2.

Original: Deﬁne p
Revised: Deﬁne p

Original: ( H ) =
Revised: ( H ) =

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

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

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

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

Original: The second identity writes the expectations as integrals
Revised: The second identity writes the expectations as integrals

Original: [Equation]
Revised: [Equation]

Original: which will make future analysis clearer.
Revised: which will make future analysis clearer.

Original: A. P ROOF OF L EMMA 3. 
Revised: A. P ROOF OF L EMMA 3.

Original: Proof.
Revised: Proof.

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

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

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

Original: For (c) , we applied Jensen’s inequality again.
Revised: For (c) , we applied Jensen’s inequality again.

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

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

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

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

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

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

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

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

Original: We have thus obtained the desired result.
Revised: We have thus obtained the desired result.

Original: A. P ROOF OF L EMMA 3. 
Revised: A. P ROOF OF L EMMA 3.

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

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

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

Original: We then proceed, following many of the same steps as for the proof of
Revised: We then proceed, following many of the same steps as for the proof of

Original: Lemma
Revised: Lemma

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

Original: 1 | s H ,a H ))
Revised: 1 | s H ,a H ))

Original: + γ θ ( H )( H log γ − log γ θ ( H ) + log r ( s H ,a H )) (cid:35) ( j ) =
Revised: + γ θ ( H )( H log γ − log γ θ ( H ) + log r ( s H ,a H )) (cid:35) ( j ) =

Original: E qθ ( τ ) (cid:34) ∞ (cid:88)
Revised: E qθ ( τ ) (cid:34) ∞ (cid:88)

Original: H =0 γ θ ( H ) (cid:18) 1 − Γ θ ( H − 1) γ θ ( H ) (log p ( s H +

Revised: H =0 γ θ ( H ) (cid:18) 1 − Γ θ ( H − 1) γ θ ( H ) (log p ( s H +

Original: |
Revised: |

Original: s H ,a H ) − log q θ ( s H +

Revised: s H ,a H ) − log q θ ( s H +

Original: | s H ,a H ))
Revised: | s H ,a H ))

Original: + H log γ − log γ θ ( H )
Revised: + H log γ − log γ θ ( H )

Original: + log r ( s H ,a H ) (cid:19)(cid:35) .
Revised: + log r ( s H ,a H ) (cid:19)(cid:35) .

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

Original: = 1 .
Revised: = 1 .

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

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

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

Original: q ( τ |
Revised: q ( τ |

Original: H ) γ θ ( H ) .
Revised: H ) γ θ ( H ) .

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

Original: =
Revised: =

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

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

Original: For (e) , we recall the deﬁnition q θ ( τ ) =
Revised: For (e) , we recall the deﬁnition q θ ( τ ) =

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

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

Original: For (g) , we regrouped the difference of dynamics terms.
Revised: For (g) , we regrouped the difference of dynamics terms.

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

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

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

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

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

Original: Thus, we have obtained the desired result.
Revised: Thus, we have obtained the desired result.

Original: We now prove Lemma 3.2, showing that
Revised: We now prove Lemma 3.2, showing that

Original: Eq. 4 becomes tight at optimality.
Revised: Eq. 4 becomes tight at optimality.

Original: Proof.
Revised: Proof.

Original: [Equation]
Revised: [Equation]

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

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

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

Original: p ( τ ) p ( H ) .
Revised: p ( τ ) p ( H ) .

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

Original: [Equation]
Revised: [Equation]

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

Original: We then substitute Eq.
Revised: We then substitute Eq.

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

Original: and simplify the resulting expression.
Revised: and simplify the resulting expression.

Original: [Table]
Revised: [Table]

Original: log r ( s H ,a H ) + (cid:40)(cid:40)(cid:40)(cid:40) log(1 − γ ) + log R ( τ ) + (cid:40)(cid:40)(cid:40)(cid:40)(cid:40) log r ( s H ,a H ) dτ (cid:1) dH − (cid:40)(cid:40)(cid:40)(cid:40) log(1 − γ )
Revised: log r ( s H ,a H ) + (cid:40)(cid:40)(cid:40)(cid:40) log(1 − γ ) + log R ( τ ) + (cid:40)(cid:40)(cid:40)(cid:40)(cid:40) log r ( s H ,a H ) dτ (cid:1) dH − (cid:40)(cid:40)(cid:40)(cid:40) log(1 − γ )

Original: = (cid:90)(cid:90)
Revised: = (cid:90)(cid:90)

Original: q θ ( τ ) γ θ ( H | τ )(log p ( τ ) − log q θ ( τ ) + log R ( τ ) dτ ) dH
Revised: q θ ( τ ) γ θ ( H | τ )(log p ( τ ) − log q θ ( τ ) + log R ( τ ) dτ ) dH

Original: = (cid:90) q θ ( τ )(log p ( τ ) − log q θ ( τ ) + log R ( τ )) dτ.
Revised: = (cid:90) q θ ( τ )(log p ( τ ) − log q θ ( τ ) + log R ( τ )) dτ.

Original: In the ﬁnal line we have removed the integral over H because none of the integrands depend on H .
Revised: 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 θ ( τ ) :
Revised: At this point, we can solve analytically for the optimal trajectory distribution, q θ ( τ ) :

Original: [Equation]
Revised: [Equation]

Original: We then substitute Eq.
Revised: We then substitute Eq.

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

Original: [Equation]
Revised: [Equation]

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

Original: A. A LOWER BOUND FOR GOAL - REACHING TASKS .
Revised: A. A LOWER BOUND FOR GOAL - REACHING TASKS .

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

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

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

Original: [Equation]
Revised: [Equation]

Original: We refer the reader to Eysenbach et al.
Revised: We refer the reader to Eysenbach et al.

Original: (2020b)
Revised: (2020b)

Original: for a more detailed discussion of this objective.
Revised: for a more detailed discussion of this objective.

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

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

Original: Lemma A.3.
Revised: Lemma A.3.

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

Original: Then the following bound holds for any dynamics q ( s t +1 | s t , a t ) and policy π ( a t | s t ) :
Revised: Then 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: where ˜ r g ( s t , a t , s t +1 )
Revised: where ˜ r g ( s t , a t , s t +1 )

Original: (cid:44) (1 − γ ) log p ( s t +

Revised: (cid:44) (1 − γ ) log p ( s t +

Original: =
Revised: =

Original: s g | s t , a t )+log p ( s t +1 | s t , a t ) − log q ( s t +1 | s t , a t ) .
Revised: s g | s t , a t )+log p ( s t +1 | s t , a t ) − log q ( s t +1 | s t , a t ) .

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

Original: The ﬁrst term in the reward function, the log probability of reaching the commanded goal one time step in the future, is similar to prior work (Rudner et al., 2021).
Revised: The ﬁrst term in the reward function, the log probability of reaching the commanded goal one time step in the future, is similar to prior work (Rudner et al., 2021).

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

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

Original: Proof.
Revised: Proof.

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

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

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

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

Original: A. D ERIVATION OF M ODEL O BJECTIVE (E Q . 10)
Revised: A. D ERIVATION OF M ODEL O BJECTIVE (E Q . 10)

Original: Our lower bound depends on entirely trajectories sampled from the learned dynamics.
Revised: 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.
Revised: In this section, we show how the same objective can be expressed as an expectation of transitions.

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

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

Original: [Equation]
Revised: [Equation]

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

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

Original: Removing these terms, we arrive at our model training objective (Eq. 10)
Revised: Removing these terms, we arrive at our model training objective (Eq. 10)

Original: B A DDITIONAL E XPERIMENTSB. A BLATION E XPERIMENTS
Revised: B A DDITIONAL E XPERIMENTS

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

Revised: MBPO (Janner et al., 2019) uses a standard maximum likelihood model.

Revised: We implement a version of VAML (Farahmand et al., 2017), which augments the maximum likelihood loss with an additional temporal difference loss; the model should predict next states that have low Bellman error.

Original: We next run an ablation experiment to study the importance of  a few key design decisions.
Revised: Finally, we compare to a variant of the MBPO maximum likelihood model that weights transitions based on the Q values, an idea discussed (but not actually implemented) in Lambert et al. We implement this value weighting method by computing the Q values for the current states and computing a softmax over the batch dimension to obtain per-example weights. We use the DClawScrewFixed-v0 task for this experiment.

Revised: The results, shown in Fig. 8, show that MnM outperforms these alternative approaches.

Revised: We observe that the value-weighting performs slightly better than the standard maximum likelihood model, while the VAML method performs noticeable worse.

Revised: We next run an ablation experiment to study the importance of a few key design decisions.

Original: We compare MnM with ablations that omit the reward augmentation and the model optimism term.
Revised: We compare MnM with ablations that omit the reward augmentation and the model optimism term.

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

Original: 8 indicate that most of the beneﬁt of MnM comes from using a GAN-like model.
Revised: 9 indicate that most of the beneﬁt of MnM comes from using a GAN-like model.

Original: Because the dynamics of these tasks are nearly deterministic, it is not surprising that the optimistic dynamics and the reward augmentation have only a small effect.
Revised: Because the dynamics of these tasks are nearly deterministic, it is not surprising that the optimistic dynamics and the reward augmentation have only a small effect.

Revised: With the implementation details described in Appendix C, we found that the MnM dynamics model trained stably, despite resembling a GAN.

Revised: In Fig.

Revised: 10, we plot the validation MSE of the MnM model throughout training, observing that it decreases monotonically.

Revised: Different lines correspond to different random seeds, and the dashed line corresponds to the minimum MSE of a maximum likelihood model (a MBPO model).

Original:  C I MPLEMENTATION D ETAILSC. A LGORITHMS
Revised: Note that the MnM model is not trained with this MSE objective, but with the GAN-like objective in Eq. It is therefore not surprising that MnM does not perform as well on this objective as the maximum likelihood model. C I MPLEMENTATION D ETAILS

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

Revised: C. A LGORITHMS

Original: Value Iteration (Fig.
Revised: Value Iteration (Fig.

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

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

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

Revised: Experiments were stopped when the MnM dynamics model (which depends on the value function) changed by less

Original: SAC
Revised: SAC

Original: MBPO
Revised: MBPO

Original: MnM
Revised: MnM

Original: HalfCheetah-vHopper-vWalker2d-vdm-cartpole-swingup_sparse   metaworld-window-open-v    metaworld-door-open-v    metaworld-reach-v    metaworld-drawer-open-v
Revised: HalfCheetah-vHopper-vWalker2d-vdm-cartpole-swingup_sparse   metaworld-window-open-v    metaworld-door-open-v    metaworld-reach-v    metaworld-drawer-open-v

Original: DClawPoseRandom-v
Revised: DClawPoseRandom-v

Original: DClawTurnRandom-v
Revised: DClawTurnRandom-v

Original: DClawScrewFixed-v
Revised: DClawScrewFixed-v

Original: DClawScrewRandom-vτ = 0 . We found that the value iteration version of MnM diverged without this Polyak averaging step. Experiments were stopped when the MnM dynamics model (which depends on the value function) changed by less than 1e-6 across iterations, as measured using an L 0 norm.
Revised: DClawScrewRandom-v

Original: Q-learning (Fig. 2)
Revised: Q-learning (Fig. 2)

Original: For the experiments with Q-learning (both with and without the MnM components), we performed (cid:15) -greedy exploration with (cid:15) = 0 .
Revised: 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 .
Revised: We used a learning rate of 1 e − 2 .

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

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

Revised: For VMBPO we used η = 1 .

Original: SAC for continuous control tasks.
Revised: SAC for continuous control tasks.

Original: We used the SAC implementation from TF-Agents (Guadarrama et al., 2018) with the default hyperparameters.
Revised: We used the SAC implementation from TF-Agents (Guadarrama et al., 2018) with the default hyperparameters.

Original: MBPO for continuous control tasks.
Revised: MBPO for continuous control tasks.

Original: We implement MBPO on top of the SAC implementation from TF-Agents (Guadarrama et al., 2018).
Revised: We implement MBPO on top of the SAC implementation from TF-Agents (Guadarrama et al., 2018).

Original: Unless otherwise mentioned, we take the default parameters from this implementation.
Revised: 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.
Revised: We use an ensemble of 5 dynamics models, each with 4 hidden layers of size 256.

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

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

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

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

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

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

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

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

Original: We use a batch size of 256.
Revised: We use a batch size of 256.

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

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

Original: We use a learning rate of 3e-4 for all components.
Revised: 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.
Revised: To stabilize learning, we maintain a target dynamics model using an exponential moving average ( τ = 0 . 001 ), and use this target dynamics model to sample transitions for training.

Original: We update the model, policy, and value functions at the same rate we sample experience from the learned model, which is more frequently than we collect experience from the real environment (see
Revised: We update the model, policy, and value functions at the same rate we sample experience from the learned model, which is more frequently than we collect experience from the real environment (see

Original: Table 1).
Revised: Table 1).

Original: MnM for Continuous Control Tasks We implement MnM on top of the SAC implementation from
Revised: MnM for Continuous Control Tasks We implement MnM on top of the SAC implementation from

Original: TF-Agents (Guadarrama et al., 2018).
Revised: TF-Agents (Guadarrama et al., 2018).

Original: Unless otherwise mentioned, we take the default parameters from this implementation.
Revised: 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 follow the same training protocol.
Revised: Our model architecture is exactly the same as our MBPO implementation, and we follow the same training protocol.

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

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

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

Original: We add input noise with σ = 0 .
Revised: We add input noise with σ = 0 .

Original: 1 while training the classiﬁer.
Revised: 1 while training the classiﬁer.

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

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

Original: Following the suggestion of prior work (Salimans et al., 2016), we use one-sided label smoothing with value 0.1, only smoothing the negative predictions and not the positive predictions.
Revised: Following the suggestion of prior work (Salimans et al., 2016), we use one-sided label smoothing 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.
Revised: We found that gradient penalties and spectral normalization decreased performance.

Original: We found that automatically tuning the classiﬁer input noise also decreased performance.
Revised: We found that automatically tuning the classiﬁer input noise also decreased performance.

Original: We found that mixup had little effect.
Revised: We found that mixup had little effect.

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

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

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

Original: To further improve stability, we compute the value term in the model loss by taking the minimum over two target value functions (like TD3 (Fujimoto et al., 2018)).
Revised: To further improve stability, we compute the value term in the model loss by taking the minimum over two target value functions (like TD3 (Fujimoto et al., 2018)).

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

Original:  C. E NVIRONMENTS
Revised: . C. E NVIRONMENTS

Original: Gridworld for Fig.
Revised: Gridworld for Fig.

Original: 2a.
Revised: 2a.

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

Original: 2a.
Revised: 2a.

Original: There are four discrete actions, corresponding to moving to the four adjacent cells.
Revised: There are 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.
Revised: 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.
Revised: 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.
Revised: 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 .
Revised: Episodes have 200 steps and we use a discount γ = 0 .

Original: Gridworld for Fig.
Revised: Gridworld for Fig.

Original: 2b.
Revised: 2b.

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

Original: 2b.
Revised: 2b.

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

Original: The dynamics are deterministic.
Revised: The dynamics are deterministic.

Original: The agent starts in the top-left cell.
Revised: 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.
Revised: 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 .
Revised: We use γ = 0 .

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

Original: We implementing the aliasing by averaging together the dynamics for each block of 3 × 3 states.
Revised: We implementing the aliasing by averaging together the dynamics for 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 moving to state (3, 4)).
Revised: Importantly, the averaging was done to the relative dynamics (e.g., action 1 corresponds to move right) not the absolute dynamics (e.g., action 1 corresponds to moving to state (3, 4)).

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

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

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

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

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

Original: Gridworld for Fig. 3a.
Revised: Gridworld for Fig. 3a.

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

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

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

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

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

Original: We use γ = 0 .
Revised: We use γ = 0 .

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

Original: Gridworld for Fig.
Revised: Gridworld for Fig.

Original: 3b.
Revised: 3b.

Original: This task used the same dynamics as Fig.
Revised: This task used the same dynamics as Fig.

Original: 2a.
Revised: 2a.

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

Original: We estimate the more complex lower bound (Eq. 4) by also learning the discount factor.
Revised: We estimate the more complex lower bound (Eq. 4) by also learning the discount factor.

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

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

Original: These tasks are taken directly from the
Revised: These tasks are taken directly from the

Original: OpenAI benchmark (Brockman et al., 2016) without modiﬁcation.
Revised: OpenAI benchmark (Brockman et al., 2016) without modiﬁcation.

Original: dm-cartpole-swingup_sparse .
Revised: dm-cartpole-swingup_sparse .

Original: This task is taken directly from the DM Control benchmark (Tassa et al., 2018) without modiﬁcation.
Revised: This task is taken directly from the DM Control benchmark (Tassa et al., 2018) without modiﬁcation.

Original: metaworld-window-open-v2 .
Revised: metaworld-window-open-v2 .

Original: This task is based on the window-open-v2 task from the Metaworld benchmark (Yu et al., 2020a).
Revised: This task is based on the window-open-v2 task from the Metaworld benchmark (Yu et al., 2020a).

Original: To increase the difﬁculty of this task, we set random_init=
Revised: To increase the difﬁculty of this task, we set random_init=

Original: True and use a sparse reward function.
Revised: True and use a sparse reward function.

Original: The sparse reward function is +0 when the window is within 3 units of the (default) open position and -1 otherwise.
Revised: The sparse reward function is +0 when the window is within 3 units of the (default) open position and -1 otherwise.

Original: metaworld-door-open-v2 .
Revised: metaworld-door-open-v2 .

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

Original: Metaworld benchmark (Yu et al., 2020a).
Revised: Metaworld benchmark (Yu et al., 2020a).

Original: To increase the difﬁculty of this task, we set random_init=
Revised: To increase the difﬁculty of this task, we set random_init=

Original: True and use a sparse reward function.
Revised: True and use a sparse reward function.

Original: The sparse reward function is +1 when the window is within 3 units of the (default) open position and +0 otherwise.
Revised: The sparse reward function is +1 when the window is within 3 units of the (default) open position and +0 otherwise.

Original: metaworld-reach-v2 .
Revised: metaworld-reach-v2 .

Original: This task is based on the reach-v2 task from the Metaworld benchmark (Yu et al., 2020a).
Revised: This task is based on the reach-v2 task from the Metaworld benchmark (Yu et al., 2020a).

Original: To increase the difﬁculty of this task, we set random_init=
Revised: To increase the difﬁculty of this task, we set random_init=

Original: True and use a sparse reward function.
Revised: True and use a sparse reward function.

Original: The sparse reward function is +0 when the window is within 1 units of the goal and -1 otherwise.
Revised: The sparse reward function is +0 when the window is within 1 units of the goal and -1 otherwise.

Original: metaworld-drawer-open-v2 .
Revised: metaworld-drawer-open-v2 .

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

Original: Metaworld benchmark (Yu et al., 2020a).
Revised: Metaworld benchmark (Yu et al., 2020a).

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

Original: DClawPoseRandom-v0 ,
Revised: DClawPoseRandom-v0 ,

Original: DClawTurnRandom-v0 ,
Revised: DClawTurnRandom-v0 ,

Original: DClawScrewFixed-v0 ,
Revised: DClawScrewFixed-v0 ,

Original: DClawScrewRandom-v0 .
Revised: DClawScrewRandom-v0 .

Original: These tasks are taken directly from the ROBEL benchmark (Ahn et al., 2020) without modiﬁcation.
Revised: These tasks are taken directly from the ROBEL benchmark (Ahn et al., 2020) without modiﬁcation.