Original: Plan To Predict: Learning an Uncertainty-Foreseeing Model for Model-Based Reinforcement Learning
Revised: Plan To Predict: Learning an Uncertainty-Foreseeing Model for Model-Based Reinforcement Learning

Original: Abstract
Revised: Abstract

Original: In Model-based Reinforcement Learning (MBRL), model learning is critical since an inaccurate model can bias policy learning via generating misleading samples.
Revised: In Model-based Reinforcement Learning (MBRL), model learning is critical since an inaccurate model can bias policy learning via generating misleading samples.

Original: However, learning an accurate model can be difficult since the policy is continually updated and the induced distribution over visited states used for model learning shifts accordingly.
Revised: However, learning an accurate model can be difficult since the policy is continually updated and the induced distribution over visited states used for model learning shifts accordingly.

Original: Prior methods alleviate this issue by quantifying the uncertainty of model-generated samples.
Revised: Prior methods alleviate this issue by quantifying the uncertainty of model-generated samples.

Original: However, these methods only quantify the uncertainty passively after the samples were generated, rather than foreseeing the uncertainty before model trajectories fall into those highly uncertain regions.
Revised: However, these methods only quantify the uncertainty passively after the samples were generated, rather than foreseeing the uncertainty before model trajectories fall into those highly uncertain regions.

Original: The resulting lowquality samples can induce unstable learning targets and hinder the optimization of the policy.
Revised: The resulting lowquality samples can induce unstable learning targets and hinder the optimization of the policy.

Original: Moreover, while being learned to minimize one-step prediction errors, the model is generally used to predict for multiple steps, leading to a mismatch between the objectives of model learning and model usage.
Revised: Moreover, while being learned to minimize one-step prediction errors, the model is generally used to predict for multiple steps, leading to a mismatch between the objectives of model learning and model usage.

Original: To this end, we propose Plan To Predict (P2P), an MBRL framework that treats the model rollout process as a sequential decision making problem by reversely considering the model as a decision maker and the current policy as the dynamics.
Revised: To this end, we propose Plan To Predict (P2P), an MBRL framework that treats the model rollout process as a sequential decision making problem by reversely considering the model as a decision maker and the current policy as the dynamics.

Original: In this way, the model can quickly adapt to the current policy and foresee the multi-step future uncertainty when generating trajectories.
Revised: In this way, the model can quickly adapt to the current policy and foresee the multi-step future uncertainty when generating trajectories.

Original: Theoretically, we show that the performance of P2P can be guaranteed by approximately optimizing a lower bound of the true environment return.
Revised: Theoretically, we show that the performance of P2P can be guaranteed by approximately optimizing a lower bound of the true environment return.

Original: Empirical results demonstrate that P2P achieves state-of-the-art performance on several challenging benchmark tasks
Revised: Empirical results demonstrate that P2P achieves state-of-the-art performance on several challenging benchmark tasks

Original:  Introduction
Revised:  Introduction

Original: Through learning an approximate dynamics of the environment and then using it to assist policy learning, Model-based Reinforcement Learning (MBRL) methods [Moerland et al., 2020] have been shown to be much more sample-efficient than model-free methods both theoretically and empirically [Nagabandi et al., 2018, Luo et al., 2019, Janner et al., 2019].
Revised: Through learning an approximate dynamics of the environment and then using it to assist policy learning, Model-based Reinforcement Learning (MBRL) methods [Moerland et al., 2020] have been shown to be much more sample-efficient than model-free methods both theoretically and empirically [Nagabandi et al., 2018, Luo et al., 2019, Janner et al., 2019].

Original: The learning of the model is critical for MBRL, since an inaccurate model can bias policy learning through generating misleading pseudo samples, which is also referred to as the model bias issue [Deisenroth and Rasmussen, 2011].
Revised: The learning of the model is critical for MBRL, since an inaccurate model can bias policy learning through generating misleading pseudo samples, which is also referred to as the model bias issue [Deisenroth and Rasmussen, 2011].

Original: [Figure]
Revised: [Figure]

Original: However, learning an accurate model can be difficult, since the update of the policy continually shifts the distribution over visited states, and due to lack of enough training data for the modified distribution, the model can suffer from overfitting and become inaccurate in regions where the updated policy is likely to visit.
Revised: However, learning an accurate model can be difficult, since the update of the policy continually shifts the distribution over visited states, and due to lack of enough training data for the modified distribution, the model can suffer from overfitting and become inaccurate in regions where the updated policy is likely to visit.

Original: This issue is further aggravated when using expressive models like deep neural networks, which are inherently prone to overfitting.
Revised: This issue is further aggravated when using expressive models like deep neural networks, which are inherently prone to overfitting.

Original: Prior methods mitigate this issue by applying uncertainty quantification techniques, typically by using a bootstrapped model ensemble [Kurutach et al., 2018, Chua et al., 2018, Janner et al., 2019].
Revised: Prior methods mitigate this issue by applying uncertainty quantification techniques, typically by using a bootstrapped model ensemble [Kurutach et al., 2018, Chua et al., 2018, Janner et al., 2019].

Original: However, due to the lack of explicit considerations of the aforementioned policy shift issue, these methods can only quantify the uncertainty passively after the samples were generated, and thus fail to prevent the trajectories induced by the current policy from running into regions with high uncertainty.
Revised: However, due to the lack of explicit considerations of the aforementioned policy shift issue, these methods can only quantify the uncertainty passively after the samples were generated, and thus fail to prevent the trajectories induced by the current policy from running into regions with high uncertainty.

Original: This incapability to avoid uncertain regions can result in low quality of the generated samples, which provides unstable learning targets and thus may hinder the learning of the policy.
Revised: This incapability to avoid uncertain regions can result in low quality of the generated samples, which provides unstable learning targets and thus may hinder the learning of the policy.

Original: Moreover, the model in the existing methods is usually used to predict for multiple steps, while the model learning objective is commonly to minimize the one-step prediction errors.
Revised: Moreover, the model in the existing methods is usually used to predict for multiple steps, while the model learning objective is commonly to minimize the one-step prediction errors.

Original: In other words, there exists an objective mismatch [Lambert et al., 2020] between model learning and model usage, which means that the policy optimization generally expects accurate multi-step trajectories (i.e., model usage), yet the model is only estimated to consider the accuracy of immediate predictions (i.e., model learning).
Revised: In other words, there exists an objective mismatch [Lambert et al., 2020] between model learning and model usage, which means that the policy optimization generally expects accurate multi-step trajectories (i.e., model usage), yet the model is only estimated to consider the accuracy of immediate predictions (i.e., model learning).

Original: To gain a conceptual understanding of the above two motivations, consider the scenario in Figure 1(a):
Revised: To gain a conceptual understanding of the above two motivations, consider the scenario in Figure 1(a):

Original: The approximate model
Revised: The approximate model

Original: ˆ P of the true environment dynamics P predicts ˆ P ( s 1 | x 0 ) > 0 , ˆ P ( s ′ 1 | x 0 ) > 0 and ˆ P ( s | x 0 ) = 0 , ∀ s / ∈ { s 1 , s ′ 1 } at the state-action pair x 0 ; and 2)
Revised: ˆ P of the true environment dynamics P predicts ˆ P ( s 1 | x 0 ) > 0 , ˆ P ( s ′ 1 | x 0 ) > 0 and ˆ P ( s | x 0 ) = 0 , ∀ s / ∈ { s 1 , s ′ 1 } at the state-action pair x 0 ; and 2)

Original: The policy is updated from π old to π new to exploit the regions with high values predicted by the value function in the model, and the resulting distribution shift exposes the model’s inaccuracy at s 1 to the current policy, leading to large expected prediction error in the next step, i.e., (cid:80) a ′ π new ( a ′ | s 1 ) D (cid:0) ˆ P ( ·| s 1 , a ′ ) , P ( ·| s 1 , a ′ ) (cid:1) >> (cid:80) a ′ π new ( a ′ | s ′ 1 ) D (cid:0) ˆ P ( ·| s ′ 1 , a ′ ) , P ( ·| s ′ 1 , a ′ ) (cid:1) , where D ( · , · ) is some distance metrics and the symbol >> means significantly greater than.
Revised: The policy is updated from π old to π new to exploit the regions with high values predicted by the value function in the model, and the resulting distribution shift exposes the model’s inaccuracy at s 1 to the current policy, leading to large expected prediction error in the next step, i.e., (cid:80) a ′ π new ( a ′ | s 1 ) D (cid:0) ˆ P ( ·| s 1 , a ′ ) , P ( ·| s 1 , a ′ ) (cid:1) >> (cid:80) a ′ π new ( a ′ | s ′ 1 ) D (cid:0) ˆ P ( ·| s ′ 1 , a ′ ) , P ( ·| s ′ 1 , a ′ ) (cid:1) , where D ( · , · ) is some distance metrics and the symbol >> means significantly greater than.

Original: Now starting from x 0 , under the updated policy π new , predicting s 1 is likely to direct the trajectory into highly uncertain regions with spurious high values (e.g., the grey region in Figure 1(b)), and cause considerably larger error in the subsequent trajectory generation than that of predicting s ′ 1 .
Revised: Now starting from x 0 , under the updated policy π new , predicting s 1 is likely to direct the trajectory into highly uncertain regions with spurious high values (e.g., the grey region in Figure 1(b)), and cause considerably larger error in the subsequent trajectory generation than that of predicting s ′ 1 .

Original: In contrast, if the model can foresee this uncertainty through explicit considerations of the current policy, then it can learn to plan for future predictions and thus jump to s ′ 1 more often to prevent the trajectory from falling into uncertain regions (e.g., the ant falls into the right hand side of the wall in Figure 1(b)).
Revised: In contrast, if the model can foresee this uncertainty through explicit considerations of the current policy, then it can learn to plan for future predictions and thus jump to s ′ 1 more often to prevent the trajectory from falling into uncertain regions (e.g., the ant falls into the right hand side of the wall in Figure 1(b)).

Original: Therefore, to generate trajectories that have small accumulation errors, the model should be aware of the update of the policy and change its predicting “strategy” by taking into account the long-term effect of immediate predictions.
Revised: Therefore, to generate trajectories that have small accumulation errors, the model should be aware of the update of the policy and change its predicting “strategy” by taking into account the long-term effect of immediate predictions.

Original: Inspired by the above intuition, we propose a new framework for MBRL, named Plan To Predict (P2P).
Revised: Inspired by the above intuition, we propose a new framework for MBRL, named Plan To Predict (P2P).

Original: Our core idea is to treat the model rollout process as a sequential decision making problem and learn the model by the corresponding algorithms, such as reinforcement learning and model predictive control.
Revised: Our core idea is to treat the model rollout process as a sequential decision making problem and learn the model by the corresponding algorithms, such as reinforcement learning and model predictive control.

Original: During the learning of the model, the current policy is fixed and serves as the environment dynamics, while at the same time the model becomes a decision maker and is penalized via the prediction errors at each step.
Revised: During the learning of the model, the current policy is fixed and serves as the environment dynamics, while at the same time the model becomes a decision maker and is penalized via the prediction errors at each step.

Original: Through actively interacting with the current policy and minimizing the accumulative prediction errors on the resulting state distribution, the model can quickly adapt to the state distribution shift by considering the long-term effect of the immediate predictions, thus foreseeing the uncertainty of regions that model trajectories are likely to fall into.
Revised: Through actively interacting with the current policy and minimizing the accumulative prediction errors on the resulting state distribution, the model can quickly adapt to the state distribution shift by considering the long-term effect of the immediate predictions, thus foreseeing the uncertainty of regions that model trajectories are likely to fall into.

Original: Such foresight of uncertainty enables the model to plan to predict for multiple steps, which improves the quality of the pseudo samples generated by the model and eventually benefits policy learning.
Revised: Such foresight of uncertainty enables the model to plan to predict for multiple steps, which improves the quality of the pseudo samples generated by the model and eventually benefits policy learning.

Original: Note that the typical treatment of minimizing one-step prediction errors regardless of the current policy is in fact a greedy optimization process, and thus can be viewed as a special case of P2P. Some prior works, such as [Nagabandi et al., 2018, Luo et al., 2019], adopt similar multi-step prediction loss to learn the model, yet these works simply compute the loss on trajectories collected by previous policies and can still suffer from severe distribution shift due to the lack of considerations of the policy update.
Revised: Note that the typical treatment of minimizing one-step prediction errors regardless of the current policy is in fact a greedy optimization process, and thus can be viewed as a special case of P2P. Some prior works, such as [Nagabandi et al., 2018, Luo et al., 2019], adopt similar multi-step prediction loss to learn the model, yet these works simply compute the loss on trajectories collected by previous policies and can still suffer from severe distribution shift due to the lack of considerations of the policy update.

Original: Theoretically, we prove that P2P approximately optimizes a lower bound of the true environment return.
Revised: Theoretically, we prove that P2P approximately optimizes a lower bound of the true environment return.

Original: Different from prior works [Luo et al., 2019, Janner et al., 2019], our theoretical analysis does not scale up the step-wise prediction errors to their maximum over timesteps when considering the impact of model error on the return approximation, thus resulting in a tighter bound of the true return and also justifying the objective mismatch issue.
Revised: Different from prior works [Luo et al., 2019, Janner et al., 2019], our theoretical analysis does not scale up the step-wise prediction errors to their maximum over timesteps when considering the impact of model error on the return approximation, thus resulting in a tighter bound of the true return and also justifying the objective mismatch issue.

Original: Optimizing this bound not only provides a stronger guarantee for policy improvement, but also enables the model to quickly adapt to the current policy.
Revised: Optimizing this bound not only provides a stronger guarantee for policy improvement, but also enables the model to quickly adapt to the current policy.

Original: Empirical results demonstrate that P2P outperforms the existing state-of-the-art methods on several challenging benchmark tasks.
Revised: Empirical results demonstrate that P2P outperforms the existing state-of-the-art methods on several challenging benchmark tasks.

Original: It also deserves to note that, due to the capability to avoid model error, P2P may be particularly suited for offline learning tasks, where incremental data collection is not allowed and thus the model error is harder to be corrected than in online settings, inducing an irreversible bias on policy learning.
Revised: It also deserves to note that, due to the capability to avoid model error, P2P may be particularly suited for offline learning tasks, where incremental data collection is not allowed and thus the model error is harder to be corrected than in online settings, inducing an irreversible bias on policy learning.

Original: Therefore, we also provide an evaluation of P2P in offline learning settings in the experiments, and present an interpretative visualization on how the learning objective of P2P affects the model predictions.
Revised: Therefore, we also provide an evaluation of P2P in offline learning settings in the experiments, and present an interpretative visualization on how the learning objective of P2P affects the model predictions.

Original:  Preliminaries and Notations
Revised:  Preliminaries and Notations

Original: A discrete time Markov desision process (MDP) M is defined by tuple ( S , A , R, P, ρ 0 , γ ) , where S is the state space, A the action space, R ( s, a ) a scalar reward function of s ∈ S and a ∈ A , P ( s ′ | s, a ) ∈ [0 , 1] the transition dynamics which is assumed to be unknown, γ ∈ [0 , 1) the discount factor and ρ 

Revised: A discrete time Markov desision process (MDP) M is defined by tuple ( S , A , R, P, ρ 0 , γ ) , where S is the state space, A the action space, R ( s, a ) a scalar reward function of s ∈ S and a ∈ A , P ( s ′ | s, a ) ∈ [0 , 1] the transition dynamics which is assumed to be unknown, γ ∈ [0 , 1) the discount factor and ρ

Original: the initial state distribution.
Revised: the initial state distribution.

Original: A stochastic policy π ( a | s ) maps states to a probability distribution over the action space.
Revised: A stochastic policy π ( a | s ) maps states to a probability distribution over the action space.

Original: The expected return of π defined by J ( π ) :
Revised: The expected return of π defined by J ( π ) :

Original: = E π,P [ (cid:80) ∞ t =

Revised: = E π,P [ (cid:80) ∞ t =

Original: γ t R ( s t , a t )] is the expectation of the sum of discounted rewards induced by π .
Revised: γ t R ( s t , a t )] is the expectation of the sum of discounted rewards induced by π .

Original: Following the general setting of reinforcement learning, the goal is to search for an optimal policy that maximizes the expected return from the environment, i.e., arg max π J ( π ) .
Revised: Following the general setting of reinforcement learning, the goal is to search for an optimal policy that maximizes the expected return from the environment, i.e., arg max π J ( π ) .

Original: MBRL methods learn a model ˆ P to approximate the unknown dynamics P , and then use this model to aid policy learning.
Revised: MBRL methods learn a model ˆ P to approximate the unknown dynamics P , and then use this model to aid policy learning.

Original: While the model can be utilized in various ways, this paper focuses on the usage of generating pseudo samples to enrich the dataset, which is one of the most common usages of the model.
Revised: While the model can be utilized in various ways, this paper focuses on the usage of generating pseudo samples to enrich the dataset, which is one of the most common usages of the model.

Original: The expected return of π predicted by the model ˆ P is denoted as J ˆ P ( π ) :
Revised: The expected return of π predicted by the model ˆ P is denoted as J ˆ P ( π ) :

Original: = E π, ˆ P [ (cid:80) ∞ t =

Revised: = E π, ˆ P [ (cid:80) ∞ t =

Original: γ t R ( s t , a t )] .
Revised: γ t R ( s t , a t )] .

Original: In our analysis, the reward function is assumed to be available, 2 and the policy obtained in the last iteration is denoted as the data-collecting policy π D ,
Revised: In our analysis, the reward function is assumed to be available, 2 and the policy obtained in the last iteration is denoted as the data-collecting policy π D ,

Original:  The Plan To Predict (P2P) Framework
Revised:  The Plan To Predict (P2P) Framework

Original: In this section, we present a detailed description of our proposed algorithmic framework for MBRL, i.e., Plan To Predict (P2P).
Revised: In this section, we present a detailed description of our proposed algorithmic framework for MBRL, i.e., Plan To Predict (P2P).

Original: In Section 3.1, the model rollout process is reformulated as a sequential decision making problem using the notion of MDP and then a meta-algorithm of the P2P framework is provided as a generic solution to this reformulated problem.
Revised: In Section 3.1, the model rollout process is reformulated as a sequential decision making problem using the notion of MDP and then a meta-algorithm of the P2P framework is provided as a generic solution to this reformulated problem.

Original: In Section 3.2, theoretical results are provided to guarantee the performance of P2P. In Section 3.3, two practical algorithms are developed under the P2P framework by handling some practical issues.
Revised: In Section 3.2, theoretical results are provided to guarantee the performance of P2P. In Section 3.3, two practical algorithms are developed under the P2P framework by handling some practical issues.

Original: 3.1 The Overall Framework
Revised: 3.1 The Overall Framework

Original: We now define the model MDP to reformulate the model rollout problem and outline our overall algorithm design.
Revised: We now define the model MDP to reformulate the model rollout problem and outline our overall algorithm design.

Original: Definition 1.
Revised: Definition 1.

Original: The model MDP M π is defined by tuple ( S m , A m , R m , P m , γ, ρ m 0 ) , where at timestep t , s mt :=
Revised: The model MDP M π is defined by tuple ( S m , A m , R m , P m , γ, ρ m 0 ) , where at timestep t , s mt :=

Original: ( s t , a t ) ∈ S m , a mt :=
Revised: ( s t , a t ) ∈ S m , a mt :=

Original: s t +1 ∈
Revised: s t +1 ∈

Original: A m , R m ( s mt , a mt ) :
Revised: A m , R m ( s mt , a mt ) :

Original: = − (cid:0) ˆ P ( s t +1 | s t , a t ) −
Revised: = − (cid:0) ˆ P ( s t +1 | s t , a t ) −

Original: P ( s t +1 | s t , a t ) (cid:1) 2 , P m ( s mt +

Revised: P ( s t +1 | s t , a t ) (cid:1) 2 , P m ( s mt +

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

Original: = π ( a t +1 | s t +1 )
Revised: = π ( a t +1 | s t +1 )

Original: ˆ P ( s t +1 | s t , a t ) , γ is the discount factor , ρ m 0 :
Revised: ˆ P ( s t +1 | s t , a t ) , γ is the discount factor , ρ m 0 :

Original: =
Revised: =

Original: π ( a 0 | s 0 )
Revised: π ( a 0 | s 0 )

Original: ρ 0 , and ρ 0 is the initial state distribution of the true environment.
Revised: ρ 0 , and ρ 0 is the initial state distribution of the true environment.

Original: In the model learning process of P2P, the roles played by the policy π and the model ˆ P are reversed: π is fixed to serve as the environment dynamics and ˆ P , now as the decision maker, interacts with π and is updated to minimize the accumulative prediction error along the generated trajectories.
Revised: In the model learning process of P2P, the roles played by the policy π and the model ˆ P are reversed: π is fixed to serve as the environment dynamics and ˆ P , now as the decision maker, interacts with π and is updated to minimize the accumulative prediction error along the generated trajectories.

Original: Similar to the optimization of the policy, the objective of this learning process can be written as arg max
Revised: Similar to the optimization of the policy, the objective of this learning process can be written as arg max

Original: ˆ P J π (cid:0) ˆ P (cid:1) where J π (cid:0) ˆ P (cid:1) :
Revised: ˆ P J π (cid:0) ˆ P (cid:1) where J π (cid:0) ˆ P (cid:1) :

Original: = E [ (cid:80) t γ t R m ( s mt , a mt )] .
Revised: = E [ (cid:80) t γ t R m ( s mt , a mt )] .

Original: The overall algorithmic framework of
Revised: The overall algorithmic framework of

Original: P2P is depicted in Algorithm 1, where the main difference compared to most existing MBRL methods lies in Line 5, i.e., the learning of the model.
Revised: P2P is depicted in Algorithm 1, where the main difference compared to most existing MBRL methods lies in Line 5, i.e., the learning of the model.

Original: Specifically, this process can be achieved in a way similar to policy learning, i.e., generating samples by actively interacting with the policy (now viewed as the background environment), and optimizing the expected return J π (cid:0) ˆ P (cid:1) accordingly.
Revised: Specifically, this process can be achieved in a way similar to policy learning, i.e., generating samples by actively interacting with the policy (now viewed as the background environment), and optimizing the expected return J π (cid:0) ˆ P (cid:1) accordingly.

Original: This optimization can in principle be done by any approaches solving sequential decision making problems.
Revised: This optimization can in principle be done by any approaches solving sequential decision making problems.

Original: As two representatives of these approaches, Model Predictive Control (MPC)
Revised: As two representatives of these approaches, Model Predictive Control (MPC)

Original: [Camacho and Alba, 2013] and Reinforcement Learning (RL)
Revised: [Camacho and Alba, 2013] and Reinforcement Learning (RL)

Original: [Sutton and Barto, 2018] are both applied in our practical version of P2P, which are detailed in Section 3.3.
Revised: [Sutton and Barto, 2018] are both applied in our practical version of P2P, which are detailed in Section 3.3.

Original: In addition, one of the biggest gaps between this theoretical construction of the model MDP and the practical implementation may lie in the derivation of R m .
Revised: In addition, one of the biggest gaps between this theoretical construction of the model MDP and the practical implementation may lie in the derivation of R m .

Original: According to Definition 1, the true environment dynamics is required to compute this reward function, while in practice the true dynamics is not accessible in the model learning phase.
Revised: According to Definition 1, the true environment dynamics is required to compute this reward function, while in practice the true dynamics is not accessible in the model learning phase.

Original: Therefore, depending on which underlying approach (i.e., MPC or RL) is chosen for model optimization, we propose two different alternative plans to approximate R m in Section 3.3.
Revised: Therefore, depending on which underlying approach (i.e., MPC or RL) is chosen for model optimization, we propose two different alternative plans to approximate R m in Section 3.3.

Original: 3.2 Theoretical Results
Revised: 3.2 Theoretical Results

Original: In general, MBRL methods learn an approximate model of the real world dynamics and learn the policy by maximizing the expected return predicted by the model.
Revised: In general, MBRL methods learn an approximate model of the real world dynamics and learn the policy by maximizing the expected return predicted by the model.

Original: The better this pseudo return approximates the true environment return, the stronger guarantee the model can provide for policy improvement.
Revised: The better this pseudo return approximates the true environment return, the stronger guarantee the model can provide for policy improvement.

Original: Thus, a critical problem in theory is how to bound the gap between the expected return of the model and the true environment, i.e., (cid:12)(cid:12)
Revised: Thus, a critical problem in theory is how to bound the gap between the expected return of the model and the true environment, i.e., (cid:12)(cid:12)

Original: J ˆ P ( π ) − J ( π ) (cid:12)(cid:12) .
Revised: J ˆ P ( π ) − J ( π ) (cid:12)(cid:12) .

Original: Suppose that this gap can be upper bounded by C , and the policy is updated from π D to π with the expected model return improved by 2 C , then the policy improvement in terms of the true environment can be guaranteed:
Revised: Suppose that this gap can be upper bounded by C , and the policy is updated from π D to π with the expected model return improved by 2 C , then the policy improvement in terms of the true environment can be guaranteed:

Original: [Equation]
Revised: [Equation]

Original: Our main theoretical result is the following theorem:
Revised: Our main theoretical result is the following theorem:

Original: Theorem 1.
Revised: Theorem 1.

Original: The gap between the expected return of the model and the environment is bounded as:
Revised: The gap between the expected return of the model and the environment is bounded as:

Original: [Equation]
Revised: [Equation]

Original: where ϵ π := max s D TV ( π D ( ·| s )
Revised: where ϵ π := [Equation] s [Equation]

Original: ∥ π ( ·| s )) denotes the policy distribution shift, ϵ mt := 
Revised: π ( · [Equation] s )) denotes the policy distribution shift, ϵ mt := E s ∼

Original: E s ∼ ˆ P t − 1 ( s ; π ) (cid:2) max a D TV ( ˆ P ( ·| s, a ) ∥ P ( ·| s, a)) (cid:3) denotes the upper bound of one-step model prediction error at timestep t of the model rollout trajectory,
Revised: [Equation] P [Equation] a

Original:  D TV ( p ∥ q ) refers to the total variation between distribution p and q , R max := max s,a R ( s, a ) , and ˆ P t − 1 ( s ; π ) denotes the state distribution at t under ˆ P and π .
Revised: [Equation] P [Equation] , a )) (cid:3) denotes the upper bound of one-step model prediction error at timestep t of the model rollout trajectory, D TV ( p ∥ q ) refers to the total variation between distribution p and q , R max := max s,a R ( s, a ) , and ˆ P t − 1 ( s, a ; π ) denotes the state-action distribution at t under ˆ P and π .

Original: Proof.
Revised: Proof.

Original: Please refer to Appendix A.
Revised: Please refer to Appendix A.

Original: Transforming the inequality in Eq.
Revised: Transforming the inequality in Eq.

Original: (1) results in J ( π ) ≥ J ˆ P ( π )
Revised: (1) results in J ( π ) ≥ J ˆ P ( π )

Original: − 2 C ( ˆ P, π ) , where C ( ˆ P, π ) denotes the right hand side in Eq.
Revised: − 2 C ( ˆ P, π ) , where C ( ˆ P, π ) denotes the right hand side in Eq.

Original: Thus by iteratively applying the update rule in Eq.
Revised: Thus by iteratively applying the update rule in Eq.

Original: (2), the above lower bound of the policy performance can be monotonically improved:
Revised: (2), the above lower bound of the policy performance can be monotonically improved:

Original: [Equation]
Revised: [Equation]

Original: This update rule is ideal and usually impractical since it involves an exhaustive search in the state and action space to compute C , and requires full-horizon rollouts in the model for estimating the accumulative model errors.
Revised: This update rule is ideal and usually impractical since it involves an exhaustive search in the state and action space to compute C , and requires full-horizon rollouts in the model for estimating the accumulative model errors.

Original: Thus, similar to how algorithms like TRPO [Schulman et al., 2015] approximate their theoretically monotonic version, P2P approximates this update rule by maximizing the expected model return while keeping the accumulative model error small.
Revised: Thus, similar to how algorithms like TRPO [Schulman et al., 2015] approximate their theoretically monotonic version, P2P approximates this update rule by maximizing the expected model return while keeping the accumulative model error small.

Revised: Specifically, via replacing the computationally expensive D TV with an expectation over ˆ P in the definition of ϵ mt , it can be easily derived that maximizing J π ( ˆ P ) defined in Section 3.1 is equivalent to minimizing (cid:80) t γ t ϵ mt .

Original: As for the policy shift term ϵ π , though the bound suggests that this term may also be constrained, we found empirically that it is sufficient to only control the model error.
Revised: As for the policy shift term ϵ π , though the bound suggests that this term may also be constrained, we found empirically that it is sufficient to only control the model error.

Original: This may be explained by the relatively small scale of policy shift with respect to model error, as observed in [Janner et al., 2019].
Revised: This may be explained by the relatively small scale of policy shift with respect to model error, as observed in [Janner et al., 2019].

Original: It is worth noting that the accumulative error here is not simply summing up the multi-step errors on trajectories collected by any other policies as in [Nagabandi et al., 2018, Luo et al., 2019], because the definition of the step-wise model error (i.e., ϵ mt ) in Theorem 1 involves a policy-dependent expectation over the state space, indicating that this error is also affected by the update of the policy.
Revised: It is worth noting that the accumulative error here is not simply summing up the multi-step errors on trajectories collected by any other policies as in [Nagabandi et al., 2018, Luo et al., 2019], because the definition of the step-wise model error (i.e., ϵ mt ) in Theorem 1 involves a policy-dependent expectation over the state space, indicating that this error is also affected by the update of the policy.

Original: To enable quick adaptation to this policy update, P2P minimizes the accumulative error on the state distribution induced by an active interaction between the model and the current policy, and thus better approximates the accumulative error defined in Theorem 1.
Revised: To enable quick adaptation to this policy update, P2P minimizes the accumulative error on the state distribution induced by an active interaction between the model and the current policy, and thus better approximates the accumulative error defined in Theorem 1.

Original: Different from most of prior work [Luo et al., 2019, Janner et al., 2019], in Theorem 1 the stepwise model prediction errors are not scaled up to their maximum over timesteps (i.e., max t ϵ mt ), leading to a tighter error bound.
Revised: Different from most of prior work [Luo et al., 2019, Janner et al., 2019], in Theorem 1 the stepwise model prediction errors are not scaled up to their maximum over timesteps (i.e., max t ϵ mt ), leading to a tighter error bound.

Original: As a result, this bound not only provides stronger guarantee for policy improvement, but also justifies the objective mismatch issue between the model learning and model usage from a theoretical perspective: Scaling up these error terms by such a potentially large magnitude yields a loose bound which leads to algorithms aiming at minimizing the upper bound of individual one-step prediction losses.
Revised: As a result, this bound not only provides stronger guarantee for policy improvement, but also justifies the objective mismatch issue between the model learning and model usage from a theoretical perspective: Scaling up these error terms by such a potentially large magnitude yields a loose bound which leads to algorithms aiming at minimizing the upper bound of individual one-step prediction losses.

Original: However, to obtain better policy improvement, the tighter bound in Eq.
Revised: However, to obtain better policy improvement, the tighter bound in Eq.

Original: (1) implies that a more fine-grained model learning process is desired, which considers consecutive rollout steps and minimizes the accumulative prediction error, i.e., (cid:80) t γ t ϵ mt .
Revised: (1) implies that a more fine-grained model learning process is desired, which considers consecutive rollout steps and minimizes the accumulative prediction error, i.e., (cid:80) t γ t ϵ mt .

Original: Using the notion of model MDP defined in Section 3.1, the model induced by minimizing one-step prediction errors can be intuitively interpreted as a greedy decision maker, which is often considered shortsighted and may easily lead to severe suboptimality.
Revised: Using the notion of model MDP defined in Section 3.1, the model induced by minimizing one-step prediction errors can be intuitively interpreted as a greedy decision maker, which is often considered shortsighted and may easily lead to severe suboptimality.

Original: 3.3 Practical Implementation
Revised: 3.3 Practical Implementation

Original: In this subsection, we instantiate Algorithm 1 by specifying explicit approaches to learn the model and providing solutions to some practical issues.
Revised: In this subsection, we instantiate Algorithm 1 by specifying explicit approaches to learn the model and providing solutions to some practical issues.

Original: Detailed pseudo codes can be found in Appendix E.
Revised: Detailed pseudo codes can be found in Appendix E.

Original: P2P-MPC: Since in our problem formulation the “dynamics” of the model rollout process, i.e., the current policy π , is accessible, one of the simplest yet effective ways to learn the model can be the model predictive control (MPC)
Revised: P2P-MPC: Since in our problem formulation the “dynamics” of the model rollout process, i.e., the current policy π , is accessible, one of the simplest yet effective ways to learn the model can be the model predictive control (MPC)

Original: [Camacho and Alba, 2013], which utilizes the dynamics to plan and optimize for a sequence of actions.
Revised: [Camacho and Alba, 2013], which utilizes the dynamics to plan and optimize for a sequence of actions.

Original: Given the state s mt at step t , the MPC controller first optimizes the sequence of actions a m t : t
Revised: Given the state s mt at step t , the MPC controller first optimizes the sequence of actions a mt : t

Original: + H over a finite horizon H , and then employs the first action of the optimal action sequence, i.e., a mH,t := arg max a mt : t + H E ˆ P (cid:80) t
Revised: + H over a finite horizon H , and then employs the first action of the optimal action sequence, i.e., a mH,t := arg max a mt : t + H E ˆ P (cid:80) t

Original: + H − 1 t ′ =
Revised: + H − 1 t ′ =

Original: t
Revised: t

Original: R m ( s mt ′ , a mt ′ ) .
Revised: R m ( s mt ′ , a mt ′ ) .

Original: Computing theexact arg max is computationally challenging due to the dimension of the state and action spaces, so we adopt the random-sampling shooting method [Rao, 2009] which generates random action sequences, executes them respectively, and chooses the one with highest return predicted by the
Revised: Computing the

Original: [Figure]
Revised: [Figure]

Original: HalfCheetah
Revised: HalfCheetah

Original: [Figure]
Revised: [Figure]

Original: r a g
Revised: r a g

Original: [Figure]
Revised: [Figure]

Original: R e t u rn
Revised: R e t u r

Original: [Figure]
Revised: [Figure]

Original:  dynamics.
Revised: exact arg max is computationally challenging due to the dimension of the state and action spaces, so we adopt the random-sampling shooting method [Rao, 2009] which generates random action sequences, executes them respectively, and chooses the one with highest return predicted by the dynamics.

Original: Besides, the reward R m , which cannot be directly computed by definition, is approximated here using a neural network and trained on the environment dataset.
Revised: Besides, the reward R m , which cannot be directly computed by definition, is approximated here using a neural network and trained on the environment dataset.

Original: Specifically, during each training iteration, P2P-MPC first trains the model via traditional one-step prediction loss, and then trains the ˆ
Revised: Specifically, during each training iteration, P2P-MPC first trains the model via traditional one-step prediction loss, and then trains the ˆ

Original: R m network by taking transitions sampled from the environment dataset as inputs, and the minus prediction errors on these transitions as labels.
Revised: R m network by taking transitions sampled from the environment dataset as inputs, and the minus prediction errors on these transitions as labels.

Original: The prediction error of an environment transition ( s, a, r, s ′ ) is computed via ∥ ˆ s ′
Revised: The prediction error of an environment transition ( s, a, r, s ′ ) is computed via ∥ ˆ s ′

Original: − s ′ ∥
Revised: − s ′ ∥

Original: + ∥ ˆ r − r ∥ , where ˆ s ′ , ˆ r are sampled from ˆ P ( · , ·| s, a ) .
Revised: + ∥ ˆ r − r ∥ , where ˆ s ′ , ˆ r are sampled from ˆ P ( · , ·| s, a ) .

Original: P2P-RL: It is well-known that though effective, RL methods often suffer from high sample complexity.
Revised: P2P-RL: It is well-known that though effective, RL methods often suffer from high sample complexity.

Original: Hence, the model is trained on the environment dataset instead of the samples generated by the interaction of the model and policy.
Revised: Hence, the model is trained on the environment dataset instead of the samples generated by the interaction of the model and policy.

Original: However, there are two issues to be addressed: 1) As the “environment dynamics”, the policy is highly non-stationary due to its continuing update; and 2) The decision maker which generates the dataset is the true environment rather than our model.
Revised: However, there are two issues to be addressed: 1) As the “environment dynamics”, the policy is highly non-stationary due to its continuing update; and 2) The decision maker which generates the dataset is the true environment rather than our model.

Original: The first issue can be addressed by updating the next state in each sampled transitions by applying the current policy, i.e., s mt +1 ← ( s t +1 , π ( s t +1 )) .
Revised: The first issue can be addressed by updating the next state in each sampled transitions by applying the current policy, i.e., s mt +1 ← ( s t +1 , π ( s t +1 )) .

Original: As for the second issue, since the model-generated samples are not used in model learning, it can be seen approximately as an offline learning problem and thus can be addressed by well-developed offline RL methods [Fujimoto et al., 2019, Levine et al., 2020].
Revised: As for the second issue, since the model-generated samples are not used in model learning, it can be seen approximately as an offline learning problem and thus can be addressed by well-developed offline RL methods [Fujimoto et al., 2019, Levine et al., 2020].

Original: Due to the effectiveness and the implementational simplicity, we adopt SAC [Haarnoja et al., 2018] with behavior cloning as the underlying learning algorithm, which bears similarities to the TD3+BC approach proposed by [Fujimoto and Gu, 2021].
Revised: Due to the effectiveness and the implementational simplicity, we adopt SAC [Haarnoja et al., 2018] with behavior cloning as the underlying learning algorithm, which bears similarities to the TD3+BC approach proposed by [Fujimoto and Gu, 2021].

Original: Besides, an approach called DualDICE [Nachum et al., 2019] is also applied to correct the estimation of the state distribution.
Revised: Besides, an approach called DualDICE [Nachum et al., 2019] is also applied to correct the estimation of the state distribution.

Original: Finally, since the samples come from the true environment, the reward R m can be simply approximated by re-predicting the dynamics in the sampled transitions and computing the prediction errors for the current model.
Revised: Finally, since the samples come from the true environment, the reward R m can be simply approximated by re-predicting the dynamics in the sampled transitions and computing the prediction errors for the current model.

Original: Specifically, P2P-RL trains the model on the environment dataset and treats the model learning process as an offline RL problem, where the true dynamics becomes the "decision maker" of the environment dataset in our problem formulation.
Revised: Specifically, P2P-RL trains the model on the environment dataset and treats the model learning process as an offline RL problem, where the true dynamics becomes the "decision maker" of the environment dataset in our problem formulation.

Original: Thus, regarding a transition ( s, a, r, s ′ ) , R m can be directly approximated by computing −∥ ˆ s ′
Revised: Thus, regarding a transition ( s, a, r, s ′ ) , R m can be directly approximated by computing −∥ ˆ s ′

Original: − s ′
Revised: − s ′

Original: ∥ − ∥ ˆ r − r ∥ , where ˆ s ′ , ˆ r ∼ ˆ P ( · , ·| s, a ) .
Revised: ∥ − ∥ ˆ r − r ∥ , where ˆ s ′ , ˆ r ∼ ˆ P ( · , ·| s, a ) .

Original:  Experiments
Revised:  Experiments

Original: In this section, we first evaluate the empirical performance of P2P on a set of MuJoCo continuous control tasks [Todorov et al., 2012] in Section 4.1, and then present an quantitative analysis for the model error in Section 4.2.
Revised: In this section, we first evaluate the empirical performance of P2P on a set of MuJoCo continuous control tasks [Todorov et al., 2012] in Section 4.1, and then present an quantitative analysis for the model error in Section 4.2.

Original: The ablation study provided in Section 4.3 verifies the necessity of an active policy-model interaction when computing the multi-step prediction loss.
Revised: The ablation study provided in Section 4.3 verifies the necessity of an active policy-model interaction when computing the multi-step prediction loss.

Original: Besides, a keyparameter of P2P-MPC is also studied in Section 4.3.
Revised: Besides, a key

Original: Finally, to gain a deeper understanding of how P2P work in offline settings, a visualization of the model’s rollout trajectories is shown in Section 4.4.
Revised: 8 Rollout

Original: [Figure]
Revised: [Figure]

Original: A cc u m u l a t
Revised: A cc u m u l a t

Original: i v e E
Revised: i v e E

Original: rr
Revised: rr

Original: o r
Revised: o r

Original: Hopper       8 Rollout Steps

Revised: Hopper       8 Rollout Steps

Original: A cc u m u l a t
Revised: A cc u m u l a t

Original: i v e E
Revised: i v e E

Original: rr
Revised: rr

Original: o r
Revised: o r

Original: Ant
Revised: Ant

Original: P2P-MPC
Revised: P2P-MPC

Original: P2P-RL
Revised: P2P-RL

Original: MBPO 1 2 3 4 5 6 7 8 9 Rollout Steps

Revised: MBPO 1 2 3 4 5 6 7 8 9 Rollout Steps

Original: A cc u m u l a t
Revised: A cc u m u l a t

Original: i v e E
Revised: i v e E

Original: rr
Revised: rr

Original: o r
Revised: o r

Original: Humanoid 
Revised: Humanoid parameter of P2P-MPC is also studied in Section 4.3. Finally, to gain a deeper understanding of how P2P work in offline settings, a visualization of the model’s rollout trajectories is shown in Section 4.4.

Original: 4.1 Comparisons with Baselines
Revised: 4.1 Comparisons with Baselines

Original: We compare P2P 3 against some state-of-the-art model-free and model-based methods.
Revised: We compare P2P 3 against some state-of-the-art model-free and model-based methods.

Original: The model-free baseline is SAC [Haarnoja et al., 2018], which is the model-free counterpart of P2P.
Revised: The model-free baseline is SAC [Haarnoja et al., 2018], which is the model-free counterpart of P2P.

Original: The model-based baselines include MBPO [Janner et al., 2019], which uses short-horizon rollouts branched from the state distribution of previous policies, SLBO [Luo et al., 2019], which enjoys theoretical performance guarantee and uses model rollouts from the initial state distribution, and STEVE [Buckman et al., 2018], which also generates short rollouts but uses model data for estimating target values instead of policy learning.
Revised: The model-based baselines include MBPO [Janner et al., 2019], which uses short-horizon rollouts branched from the state distribution of previous policies, SLBO [Luo et al., 2019], which enjoys theoretical performance guarantee and uses model rollouts from the initial state distribution, and STEVE [Buckman et al., 2018], which also generates short rollouts but uses model data for estimating target values instead of policy learning.

Original: Note that our framework only focuses on model learning and can be employed as a flexible plug-in component by most MBRL algorithms that generate pseudo samples to assist policy learning.
Revised: Note that our framework only focuses on model learning and can be employed as a flexible plug-in component by most MBRL algorithms that generate pseudo samples to assist policy learning.

Original: In our experiments, we plug the model learning process of P2P into MBPO since it is a widely accepted strong baseline in MBRL.
Revised: In our experiments, we plug the model learning process of P2P into MBPO since it is a widely accepted strong baseline in MBRL.

Original: The implementational details and the hyperparameter settings of P2P, as well as the environment settings, are all described in Appendix B.
Revised: The implementational details and the hyperparameter settings of P2P, as well as the environment settings, are all described in Appendix B.

Original: As shown in Figure 2, P2P-MPC significantly outperforms the baselines on several tasks, especially on Hopper, Humanoid and Ant, which are generally considered to be the most challenging MuJoCo tasks.
Revised: As shown in Figure 2, P2P-MPC significantly outperforms the baselines on several tasks, especially on Hopper, Humanoid and Ant, which are generally considered to be the most challenging MuJoCo tasks.

Original: In Walker2d, the best rollout length is one step in our implementation and MBPO, thus resulting in similar performance between these two methods.
Revised: In Walker2d, the best rollout length is one step in our implementation and MBPO, thus resulting in similar performance between these two methods.

Original: Compared to P2P-MPC, the performance of
Revised: Compared to P2P-MPC, the performance of

Original: P2P-RL is less stable.
Revised: P2P-RL is less stable.

Original: Through an investigation on the model learning process (see Appendix C), we find that P2P-RL sometimes struggles to balance the loss of behavior cloning and RL, leading to the difficulty in hyperparameter tuning and the instability of learning network parameters, while the optimization of P2P-MPC is non-parametric and thus does not suffer from this issue.
Revised: Through an investigation on the model learning process (see Appendix C), we find that P2P-RL sometimes struggles to balance the loss of behavior cloning and RL, leading to the difficulty in hyperparameter tuning and the instability of learning network parameters, while the optimization of P2P-MPC is non-parametric and thus does not suffer from this issue.

Original: 4.2 Model Error Analysis
Revised: 4.2 Model Error Analysis

Original: The theoretical result presented in Section 3.2 implies that to guarantee better policy improvement, the model should learn to minimize the accumulative errors along the trajectories induced by the current policy and the model.
Revised: The theoretical result presented in Section 3.2 implies that to guarantee better policy improvement, the model should learn to minimize the accumulative errors along the trajectories induced by the current policy and the model.

Original: Since P2P is designed based on this principle, the empirical results in Figureverify the theoretical result in terms of the final performance.
Revised: Since P2P is designed based on this principle, the empirical results in Figureverify the theoretical result in terms of the final performance.

Original: To further verify the effectiveness of this algorithmic design, we study the accumulative prediction error on three tasks that require rollouts with relative long horizons, i.e., Hopper, Ant and Humanoid.
Revised: To further verify the effectiveness of this algorithmic design, we study the accumulative prediction error on three tasks that require rollouts with relative long horizons, i.e., Hopper, Ant and Humanoid.

Original: The results shown in Figure 3 together with Figure 2 strongly support our theoretical claims by demonstrating that the algorithm inducing less accumulative error achieves better performance.
Revised: The results shown in Figure 3 together with Figure 2 strongly support our theoretical claims by demonstrating that the algorithm inducing less accumulative error achieves better performance.

Original: Besides, we can also observe that in Hopper the model errors of P2P-MPC in the first few steps are slightly larger than the other two methods.
Revised: Besides, we can also observe that in Hopper the model errors of P2P-MPC in the first few steps are slightly larger than the other two methods.

Original: Nevertheless, as the rollout trajectory goes longer, the accumulative error induced by P2P-MPC grows significantly slower than the other methods.
Revised: Nevertheless, as the rollout trajectory goes longer, the accumulative error induced by P2P-MPC grows significantly slower than the other methods.

Original: This further agrees with the intuition mentioned in Figure 1 by showing that P2P is able to trade the one-step model error for accumulative model error by considering the long-term effect of the immediate prediction.
Revised: This further agrees with the intuition mentioned in Figure 1 by showing that P2P is able to trade the one-step model error for accumulative model error by considering the long-term effect of the immediate prediction.

Original: 4.3 Ablation Study
Revised: 4.3 Ablation Study

Original: Multi-Step Prediction Loss
Revised: Multi-Step Prediction Loss

Original: In Section 3.2, we emphasize that the loss of P2P is essentially different from the multi-step prediction loss used in prior work, since P2P computes the multi-step loss on the state distribution induced from an active interaction between the model and the currentpolicy, and on the contrary, prior work computes this loss on the trajectories collected by previous policies. The empirical results in Figure 4 compare the performances induced by these two kind of losses and highlight the importance of the interactive model learning process in P2P.
Revised: In Section 3.2, we emphasize that the loss of P2P is essentially different from the multi-step prediction loss used in prior work, since P2P computes the multi-step loss on the state distribution induced from an active interaction between the model and the current

Original: [Figure]
Revised: [Figure]

Original: A v e r a g e
Revised: A v e r a g e

Original: R e t u r n
Revised: R e t u r n

Original: Ant
Revised: Ant

Original: [Figure]
Revised: [Figure]

Revised: policy, and on the contrary, prior work computes this loss on the trajectories collected by previous policies.

Revised: The empirical results in Figure 4 compare the performances induced by these two kind of losses and highlight the importance of the interactive model learning process in P2P.

Original: Planning Horizon Different from most RL methods, MPC methods can assign an explicit horizon for planning, hence P2P-MPC is able to set the planning horizon of the model to be exactly the same with the rollout horizon.
Revised: Planning Horizon Different from most RL methods, MPC methods can assign an explicit horizon for planning, hence P2P-MPC is able to set the planning horizon of the model to be exactly the same with the rollout horizon.

Original: This difference may partially explain the superior performance of P2P-MPC against P2P-RL.
Revised: This difference may partially explain the superior performance of P2P-MPC against P2P-RL.

Original: Nevertheless, longer planning horizon incurs higher computational cost, especially in tasks with high-dimensional state spaces.
Revised: Nevertheless, longer planning horizon incurs higher computational cost, especially in tasks with high-dimensional state spaces.

Original: To trade-off the performance and the computational cost, our solution is to simply limit the planning horizon, and the empirical results shown in Figuresuggests that a horizon of 6 is satisfactory for the tasks experimented.
Revised: To trade-off the performance and the computational cost, our solution is to simply limit the planning horizon, and the empirical results shown in Figuresuggests that a horizon of 6 is satisfactory for the tasks experimented.

Original: 4.4 An Interpretative Experiment
Revised: 4.4 An Interpretative Experiment

Original: To gain an intuitive understanding of how the learning objective of P2P affects the model predictions, here we present a visualization of the rollout trajectories in maze2d-medium, an offline learning task in the D4RL benchmark [Fu et al., 2020].
Revised: To gain an intuitive understanding of how the learning objective of P2P affects the model predictions, here we present a visualization of the rollout trajectories in maze2d-medium, an offline learning task in the D4RL benchmark [Fu et al., 2020].

Original: As shown in Figure 6, the environment is slightly modified by lengthening the wall at the middle of the maze and fixing the goal to the top right corner.
Revised: As shown in Figure 6, the environment is slightly modified by lengthening the wall at the middle of the maze and fixing the goal to the top right corner.

Original: Besides, the number of offline samples that involve the grey region are partially discarded to increase this region’s uncertainty.
Revised: Besides, the number of offline samples that involve the grey region are partially discarded to increase this region’s uncertainty.

Original: Figure 6 shows the top four trajectories that are most frequently generated from a fixed starting point using our methods and MOPO
Revised: Figure 6 shows the top four trajectories that are most frequently generated from a fixed starting point using our methods and MOPO

Original: [Yu et al., 2020], a well-known offline model-based RL method that uses one-step error for model learning.
Revised: [Yu et al., 2020], a well-known offline model-based RL method that uses one-step error for model learning.

Original: The visualization clearly demonstrates the different rollout preferences of the two kinds of models: the model of P2P-MPC tends to generate trajectories heading to the regions with low uncertainty, while the model of MOPO fails to prevent the trajectories from being absorbed into the highly uncertain region.
Revised: The visualization clearly demonstrates the different rollout preferences of the two kinds of models: the model of P2P-MPC tends to generate trajectories heading to the regions with low uncertainty, while the model of MOPO fails to prevent the trajectories from being absorbed into the highly uncertain region.

Original: The resulting accumulative model errors are quantified in Figure 7, and as shown in
Revised: The resulting accumulative model errors are quantified in Figure 7, and as shown in

Original: Table 1, the final performance of the induced policy of P2P-MPC is significantly better than MOPO.
Revised: Table 1, the final performance of the induced policy of P2P-MPC is significantly better than MOPO.

Original: Note that in online settings, the model error is possible to be corrected by the newly collected data from the environment, while in offline settings this incremental data-collecting process is not allowed.
Revised: Note that in online settings, the model error is possible to be corrected by the newly collected data from the environment, while in offline settings this incremental data-collecting process is not allowed.

Original: Thus, the model error can have amore detrimental effect on policy learning, which may severely degrade the asymptotic performance as evidenced in Table 1. Additional results on a much harder task (i.e., maze2d-large-v1) are also provided in Table 1 to further verifies the effectiveness of P2P-MPC.
Revised: Thus, the model error can have a

Original: [Figure]
Revised: [Figure]

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

Original: [Figure]
Revised: u

Original: [Figure]
Revised: [Figure]

Original: cc u m u l a t

Original:  iv e
Revised: l a t i

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

Original: P2P
Revised: P2P

Original: [Figure]
Revised: [Figure]

Original: MPC
Revised: MPC

Original: MOPO
Revised: MOPO

Original: [Table]
Revised: [Table]

Original: diverse [Table]
Revised: more detrimental effect on policy learning, which may severely degrade the asymptotic performance as evidenced in Table 1. Additional results on a much harder task (i.e., maze2d-large-v1) are also provided in Table 1 to further verifies the effectiveness of P2P-MPC.

Original:  Related Work
Revised:  Related Work

Original: MBRL aims at improving the sample efficiency of model-free methods while maintaining high asymptotic performance.
Revised: MBRL aims at improving the sample efficiency of model-free methods while maintaining high asymptotic performance.

Original: The research of MBRL can be roughly divided into two lines: the model usage and the model learning.
Revised: The research of MBRL can be roughly divided into two lines: the model usage and the model learning.

Original: The most common usage of the model is generating pseudo samples to enrich the data buffer, so as to reduce the interaction with the environment and thus accelerate policy learning [Sutton, 1990, 1991, Deisenroth et al., 2013, Kalweit and Boedecker, 2017, Luo et al., 2019,
Revised: The most common usage of the model is generating pseudo samples to enrich the data buffer, so as to reduce the interaction with the environment and thus accelerate policy learning [Sutton, 1990, 1991, Deisenroth et al., 2013, Kalweit and Boedecker, 2017, Luo et al., 2019,

Original: Janner et al., 2019, Pan et al., 2020].
Revised: Janner et al., 2019, Pan et al., 2020].

Original: The main challenge of this kind of usage is the compounding model error of long-term predictions, which is also known as the model bias issue [Deisenroth and
Revised: The main challenge of this kind of usage is the compounding model error of long-term predictions, which is also known as the model bias issue [Deisenroth and

Original: Rasmussen, 2011].
Revised: Rasmussen, 2011].

Original: A number of approaches are proposed to address this issue, varying from using
Revised: A number of approaches are proposed to address this issue, varying from using

Original: Dyna-like short-horizon rollouts [Sutton, 1991, Janner et al., 2019], interpolating between real and pseudo samples [Kalweit and Boedecker, 2017], learning an uncertainty-aware probabilistic model ensemble [Chua et al., 2018], to masking the rollouts according to the uncertainty estimation [Pan et al., 2020].
Revised: Dyna-like short-horizon rollouts [Sutton, 1991, Janner et al., 2019], interpolating between real and pseudo samples [Kalweit and Boedecker, 2017], learning an uncertainty-aware probabilistic model ensemble [Chua et al., 2018], to masking the rollouts according to the uncertainty estimation [Pan et al., 2020].

Original: Falling into this category of model usage, the focus of our work is to improve the quality of model-generated trajectories by reducing the accumulative model error along these trajectories.
Revised: Falling into this category of model usage, the focus of our work is to improve the quality of model-generated trajectories by reducing the accumulative model error along these trajectories.

Original: Through an active interaction between the model and the current policy, P2P reinforces a quick model adaptation to the update of policy, so as to foresee and thus avoid the uncertainty in the future trajectory during rollout.
Revised: Through an active interaction between the model and the current policy, P2P reinforces a quick model adaptation to the update of policy, so as to foresee and thus avoid the uncertainty in the future trajectory during rollout.

Original: Hence, one of the major differences between P2P and prior work can be interpreted as: not just to passively estimate the uncertainty after encountering it, but to actively avoid it before this encounter.
Revised: Hence, one of the major differences between P2P and prior work can be interpreted as: not just to passively estimate the uncertainty after encountering it, but to actively avoid it before this encounter.

Original: It is worth noting that there are also methods based on the optimism in face of uncertainty, which try to exploit the uncertainty of the model for sufficient exploration in the real environment, such as [Curi et al., 2021].
Revised: It is worth noting that there are also methods based on the optimism in face of uncertainty, which try to exploit the uncertainty of the model for sufficient exploration in the real environment, such as [Curi et al., 2021].

Original: If the model’s promise of large return turns out to be false, executing the policy trained by such a model will hopefully collect real-world counterfactual data to the previously erroneous model.
Revised: If the model’s promise of large return turns out to be false, executing the policy trained by such a model will hopefully collect real-world counterfactual data to the previously erroneous model.

Original: However, this correcting process itself requires additional samples, which may run counter to the principle of developing model-based methods.
Revised: However, this correcting process itself requires additional samples, which may run counter to the principle of developing model-based methods.

Original: Furthermore, exploring the erroneous part of the model may not only provides the agent with false information about the real environment, but also be prone to the model exploitation issue which severely hurts the asymptotic performance of the policy.
Revised: Furthermore, exploring the erroneous part of the model may not only provides the agent with false information about the real environment, but also be prone to the model exploitation issue which severely hurts the asymptotic performance of the policy.

Original: Therefore, how to carefully exploit the model uncertainty remains an open question.
Revised: Therefore, how to carefully exploit the model uncertainty remains an open question.

Original: In the line of model learning, the most commonly used objective is to simply minimize each one-step prediction error for transitions available in the environment dataset [Kurutach et al., 2018, Chua et al., 2018, Janner et al., 2019].
Revised: In the line of model learning, the most commonly used objective is to simply minimize each one-step prediction error for transitions available in the environment dataset [Kurutach et al., 2018, Chua et al., 2018, Janner et al., 2019].

Original: However, since the data distribution of the environment dataset varies with the update of policy, it is unlikely to learn an accurate model for all the transitions in complex environments.
Revised: However, since the data distribution of the environment dataset varies with the update of policy, it is unlikely to learn an accurate model for all the transitions in complex environments.

Original: Thus, when the model is used to predict for multiple steps, there exists an objective mismatch [Lambert et al., 2020] between model usage and model learning.
Revised: Thus, when the model is used to predict for multiple steps, there exists an objective mismatch [Lambert et al., 2020] between model usage and model learning.

Original: To address this issue, we propose to minimize the accumulative prediction error under the state distribution of the current policy.
Revised: To address this issue, we propose to minimize the accumulative prediction error under the state distribution of the current policy.

Original: The most relevant work to ours includes [Nagabandi et al., 2018, Luo et al., 2019], where a similar multi-step prediction loss is employed for model learning.
Revised: The most relevant work to ours includes [Nagabandi et al., 2018, Luo et al., 2019], where a similar multi-step prediction loss is employed for model learning.

Original: The essential difference between their multi-step loss and ours is clarified in Section 3.1 and 3.2.
Revised: The essential difference between their multi-step loss and ours is clarified in Section 3.1 and 3.2.

Original: Generative models [Kingma and Welling, 2014] have also been utilized to perform multi-step prediction [Mishra et al., 2017, Ha and Schmidhuber, 2018].
Revised: Generative models [Kingma and Welling, 2014] have also been utilized to perform multi-step prediction [Mishra et al., 2017, Ha and Schmidhuber, 2018].

Original: However, the need for a highly exploratory dataset limits their application to relatively simple and low-dimensional tasks.
Revised: However, the need for a highly exploratory dataset limits their application to relatively simple and low-dimensional tasks.

Original: Other learning objectives improve traditional model learning from angles that are less similar to ours, including goal-aware model learning that predicts how far the current state is from the goal state [Nair et al., 2020] and decisionaware model learning that aims to better aids policy learning by re-weighting the model error on different samples [Farahmand et al., 2017, Farahmand, 2018, D’Oro et al., 2020].
Revised: Other learning objectives improve traditional model learning from angles that are less similar to ours, including goal-aware model learning that predicts how far the current state is from the goal state [Nair et al., 2020] and decisionaware model learning that aims to better aids policy learning by re-weighting the model error on different samples [Farahmand et al., 2017, Farahmand, 2018, D’Oro et al., 2020].

Original: In recent years, a series of work has attempted to learn the model by treating the model rollout process as a sequential decision-making problem.
Revised: In recent years, a series of work has attempted to learn the model by treating the model rollout process as a sequential decision-making problem.

Original: While these work shares some similarity with our work, the targeted problems and the proposed solutions are not exactly the same as P2P. Shang et al.
Revised: While these work shares some similarity with our work, the targeted problems and the proposed solutions are not exactly the same as P2P. Shang et al.

Original: [2019] propose an environment reconstruction method which models the influence of the hidden confounder on the environment by treating the platform, the user and the confounder as three agents interacting with each other.
Revised: [2019] propose an environment reconstruction method which models the influence of the hidden confounder on the environment by treating the platform, the user and the confounder as three agents interacting with each other.

Original: They focus on the offline setting (i.e., RL-based recommendation) and simultaneously train the model and the policy using a multi-agent imitation learning method.
Revised: They focus on the offline setting (i.e., RL-based recommendation) and simultaneously train the model and the policy using a multi-agent imitation learning method.

Original: Xu et al.
Revised: Xu et al.

Original: [2020] treat the model as a dual agent and analyze the error bounds of the model.
Revised: [2020] treat the model as a dual agent and analyze the error bounds of the model.

Original: They propose to train the model using imitation learning methods.
Revised: They propose to train the model using imitation learning methods.

Original: Chen et al.
Revised: Chen et al.

Original: [2022] also consider multi-step model error, yet they mainly focus on handling counterfactual data queried by adversarial policies.
Revised: [2022] also consider multi-step model error, yet they mainly focus on handling counterfactual data queried by adversarial policies.

Original: Unlike the above two work [Xu et al., 2020, Chen et al., 2022] that both focus solely on model learning, P2P aims at proposing a general MBRL framework and attempts to adapt the model to the continuously updating policy quickly.
Revised: Unlike the above two work [Xu et al., 2020, Chen et al., 2022] that both focus solely on model learning, P2P aims at proposing a general MBRL framework and attempts to adapt the model to the continuously updating policy quickly.

Original:  Conclusion and Future Work
Revised:  Conclusion and Future Work

Original: In this paper, we propose an MBRL framework named Plan To Predict (P2P), which novelly treats the model rollout process as a sequential decision making problem.
Revised: In this paper, we propose an MBRL framework named Plan To Predict (P2P), which novelly treats the model rollout process as a sequential decision making problem.

Original: In this problem reformulation, the roles played by the model and the policy are reversed, where the model becomes a decision maker and the policy instead serves as the outer environment.
Revised: In this problem reformulation, the roles played by the model and the policy are reversed, where the model becomes a decision maker and the policy instead serves as the outer environment.

Original: Through an active interaction between the model and the current policy, the model can minimize the multi-step accumulative errors on the induced trajectories and thus quickly adapt to the state distribution of the current policy.
Revised: Through an active interaction between the model and the current policy, the model can minimize the multi-step accumulative errors on the induced trajectories and thus quickly adapt to the state distribution of the current policy.

Original: We provide theoretical guarantee for P2P, which also clarifies the difference of our multi-step loss from prior methods and justifies the mismatch of objectives between model learning and model usage.
Revised: We provide theoretical guarantee for P2P, which also clarifies the difference of our multi-step loss from prior methods and justifies the mismatch of objectives between model learning and model usage.

Original: Empirical results on several challenging benchmark tasks verify the effectiveness of P2P. Moreover, observing that the capability to avoid model error makes P2P promising for offline settings, we also provide preliminary results on some offline learning tasks.
Revised: Empirical results on several challenging benchmark tasks verify the effectiveness of P2P. Moreover, observing that the capability to avoid model error makes P2P promising for offline settings, we also provide preliminary results on some offline learning tasks.

Original: While this is not the major focus of this paper, studying how P2P can be applied in offline tasks more properly is left as an important future work.
Revised: While this is not the major focus of this paper, studying how P2P can be applied in offline tasks more properly is left as an important future work.

Original: Acknowledgments and Disclosure of Funding
Revised: Acknowledgments and Disclosure of Funding

Original: The work is supported by the National Natural Science Foundation of China under Grant 62076259.
Revised: The work is supported by the National Natural Science Foundation of China under Grant 62076259.

Original: The authors would like to thank Qian Lin, Zongkai Liu and Xuejing Zheng for the useful discussions, and thank Hanlin Yang for the help in drawing Figure 1 (b).
Revised: The authors would like to thank Qian Lin, Zongkai Liu and Xuejing Zheng for the useful discussions, and thank Hanlin Yang for the help in drawing Figure 1 (b).

Original: Finally, the reviewers and metareviewer are highly appreciated for their constructive feedback on the paper
Revised: Finally, the reviewers and metareviewer are highly appreciated for their constructive feedback on the paper