Year: 2021

Abstract:
Model-based reinforcement learning provides an efficient mechanism to find the optimal policy by interacting with the learned environment. In addition to treating the learned environment like a black-box simulator, a more effective way to use the model is to exploit its differentiability. Such methods require the gradient information of the learned environment model when calculating the policy gradient. However, since the error of gradient is not considered in the model learning phase, there is no guarantee for the model’s accuracy. To address this problem, we first analyze the convergence rate for the policy optimization methods when the policy gradient is calculated using the learned environment model. The theoretical results show that the model gradient error matters in the policy optimization phrase. Then we propose a two-model-based learning method to control the prediction error and the gradient error. We separate the different roles of these two models at the model learning phase and coordinate them at the policy optimization phase. After proposed the method, we introduce the directional derivative projection policy optimization (DDPPO) algorithm as a practical implementation to find the optimal policy. Finally, we empirically demonstrate the proposed algorithm has better sample efficiency when achieving a comparable or better performance on benchmark continuous control tasks.

heading: 1 INTRODUCTION
text: Reinforcement learning (RL) is a powerful technique for solving the sequential decision making problems (Li, 2018; Sutton & Barto, 1998). Recent work on model-based RL (Nagabandi et al., 2018; Luo et al., 2018; Kurutach et al., 2018; Wang et al., 2019; Janner et al., 2019; Pan et al., 2020), has shown the power of first learning the environment model and then use it to do the policy optimization. Several methods are proposed to achieve the goal of getting similar performance by using fewer data, such as ensembles (Kurutach et al., 2018), probabilistic models (Chua et al., 2018), and meta-learning (Clavera et al., 2018).
In addition to treating the learned environment like a black-box simulator, a more effective way of using the model is to exploit its differentiability (Heess et al., 2015; Clavera et al., 2019; D’ Oro & Jaśkowski, 2020; Amos et al., 2021), which is mainly focused in this paper. To get the policy updating direction, these methods compute the gradient of the objective with respect to the policy directly by back-propagating through the model. Therefore, the model gradient is used in the calculation and its error will influence the accuracy of the policy gradient. However, since the traditional model learning only aims to get the accurate prediction for the next state and the reward, there is no guarantee for the accuracy of the model gradient. In other words, there is an objective mismatch (Lambert et al., 2020) that the algorithm requires the accurate model gradient, but we only learn to decrease the prediction error.
In this paper, to address these problems, we first theoretically analyze the problem and then propose our solution based on the theoretical results. First of all, we present the convergence rate analysis for the policy optimization algorithms in what the policy gradient is calculated using the learned environment model. By taking the model gradient error into account, we can see that the gradients of the transition and reward models matter in the policy optimization. Specifically, the bias of the estimated policy gradient used to update the policy is not only introduced by the prediction error of the learned model but also introduced by the gradient error of the learned model. Furthermore, the
policy gradient bias due to the different types of model error will finally influence the convergence rate of the policy optimization process.
Then, inspired by the theoretical results, we propose the two-model-based learning method. According to the policy gradient bias and convergence rate analysis, in order to optimize the policy efficiently and accurately, we need the learned environment model both with small prediction error and small gradient error. Therefore, we propose to set separate models for different purposes of usage. In the model learning phase, the prediction model aims to reduce the prediction error, and the gradient model focuses on minimizing gradient error. In the policy optimization phase, we will use the prediction model to rollout the data and use the gradient model to calculate the policy gradient.
To make the proposed method applicable, we introduce the directional derivative projection policy optimization (DDPPO) algorithm. Our first goal is to use data to estimate the gradient or Jacobian matrix for the environment model and use the estimator to learn the model’s gradient explicitly. The challenge is that the state and action are usually with high dimensions and directly estimating the gradient or Jacobian matrix using data is intractable. Thus, we first estimate the directional derivative using data and then project the model’s gradient or Jacobian matrix into these directions. Minimizing the error between the estimated directional derivative and the projection value, we can learn the model’s gradient. Secondly, after we learn the environment model with a more accurate gradient, we can leverage two-model-based learning method to do the policy optimization.
Finally, we conduct experiments on the simple environments and the benchmark MuJoCo continuous control environments. The experimental results verify our theoretical findings and demonstrate the effectiveness of the proposed method.
Our main contribution can be summarized as follows:
1. We theoretically depict how the different model errors influence the convergence rate of the model-based policy optimization algorithm. The result shows that the gradient error of the model indeed matters in the convergence of the policy optimization.
2. We propose the two-model-based learning method and the practical DDPPO algorithm which learn and use two environment models (prediction model used for rollout and gradient model used to provide the gradient information) for the model-based policy optimization.
3. Empirically, we can achieve better sample efficiency in the MuJoCo continuous control tasks than state-of-the-art model-based and model-free methods .

heading: 2 PRELIMINARIES
text: Reinforcement Learning: We consider a discrete-time Markov decision process (MDP) M, defined by tuple (S,A, f, r, γ, p0), where S and A are the state and action spaces, respectively. Here, f : st+1 = f(st, at, ϵt) is the transition distribution, r : st = r(st, at) is a reward function, p0 is the initial state distribution and γ represents the discount factor. We define the return (J) as the expected sum of discounted rewards along a trajectory (R = ∑∞ t=0 γ
tr(st, at)). The goal of reinforcement learning is to find a policy πϕ that maximizes the expected sum of discounted rewards, i.e.,
max ϕ J(ϕ) = max ϕ EπϕR = max ϕ Eπϕ [ ∞∑ t=0 γtr(st, at) ] . (1)
Model-based RL is characterized by learning the transition model using data collected from interaction with the environment. Typically, we use a parametric function f̃θ to denote the transition model trained by a neural network. The state predicted by the learned model is defined as s̃t+1 and s̃t+1 ∼ f̃θ. Similarly, we use a parametric function r̃θ to denote the reward model. Generally we use M to denote the true model and M̃ the learned model. Then s̃t+1, r̃t ∼ M̃θ(st, at). The model is trained by supervised learning, typically, via maximum likelihood:
JM̃p(θp) = E [ log M̃p(st+1, rt|st, at) ] . (2)
where (st, at, rt, st+1) are data collected from interaction with the environment.
Policy Optimization According to the Policy Gradient Calculated Using Model: After we learned the model, we can use the model to rollout trajectories and use the data to calculate the policy gradient for the purpose of policy optimization. Specifically, denote the policy parameters as ϕ, we can calculate the return J using generated data as Jπ(ϕ) = E [ ∑∞ t=0 γ
tr̃t], where s̃t+1, r̃t ∼ M̃(s̃t, ãt) and ãt ∼ πϕ(s̃t). To optimize the objective, One can use the traditional policy optimization methods such as A3C (Mnih et al., 2016), PPO (Schulman et al., 2017), SAC (Haarnoja et al., 2018). Beyond that, many previous works argue that we can use the model to calculate the policy gradient by back-propagating through the model (Heess et al., 2015; Clavera et al., 2019; Amos et al., 2021) leveraging the reparameterization trick (Kingma & Welling, 2014).
Specifically, since we have learned the environment model, we can calculate the gradient of the objective w.r.t. ϕ by chain rule as shown in Equation 3 in which the gradients drtdat , drt dst , dskdst+1 , dst+1 dat need to be calculated using the learned model.
dJ dϕ = E ∞∑ t=0 dJ dat dat dϕ = E ∞∑ t=0 ( γt drt dat + ∞∑ k=t+1 γk drk dsk dsk dst+1 dst+1 dat ) dat dϕ . (3)
The expectation is taken over randomness from random policy and transition. To optimize the policy, we can update the policy parameters towards the gradient directions. Since we don’t have the true gradient, we use learned model to estimate the policy gradient denoted as dR M̂
dϕ . αt is learning rate.
ϕt+1 = ϕt + αt dRM̂
dϕ . (4)
Notations : For simplicity and clarity, we introduce our notations in the following. We use fraction notation d·d· to represent the gradient calculated by taking gradient on the true model and use the ”hat” version d̂·d· to represent the gradient calculated on the learned model. We use the ”tilde” notation ·̃ to represent the state, action or reward generate using the learned model and the without-tilde notation to represent those variables generate using the true model. For example the notation dst+1dst to denote the gradient function value calculated by taking gradient for the output of the transition model s′ (s′ = f(s, a, ϵ)) with respect to its input s given input value equals s = st. We use the notation ϵf to represent the transition model prediction error which means, for ∀ s, a, we have∥∥∥f̃(s, a)− f(s, a)∥∥∥ ≤ ϵf . We use the notation ϵgf to represent the gradient error of the transition model. For ∀ s, a, we have∥∥∥∥∥ d̂s′ds − ds′ds ∥∥∥∥∥+ ∥∥∥∥∥ d̂s′da − ds′da ∥∥∥∥∥ = ∥∥∥∥∥df̂(s, a)ds − df(s, a)ds ∥∥∥∥∥+ ∥∥∥∥∥df̂(s, a)da − df(s, a)da
∥∥∥∥∥ ≤ ϵgf . (5) Similarly, we use the notation ϵgr to represent the gradient error of the reward model.

heading: 3 RELATED WORK
text: Due to the space limitation, we clarify several highly related works in this section and put the detailed related work into the Appendix A. Overall speaking, most of the model-based policy optimization algorithms aim to calculate the policy gradient accurately and efficiently. These works can be divided into two categories, i.e., how to learn the model and how to use the model.
Learn the better model: Chua et al. (2018) design an ensemble-based neural network to learn the true model. Zhang et al. (2021a) design the dropout mechanisms to improve the robustness of the learned model. Zhang et al. (2020b) design the auto-regressive network structure to improve the accuracy of the model prediction.
Consider the decision-making phase when learning the model: These works are always called Decision-aware model learning (DAML). D’Oro et al. (2020) use policy gradient to get the reweighted model prediction error for each sample. Farahmand et al. (2017) use the value function to get the re-weighted model prediction error for each sample.
Use the learned model to argument data: Janner et al. (2019) use the learned model to rollout the augmented trajectories and use both the generated trajectories and the true trajectories to train the SAC algorithm. Yu et al. (2020); Lee et al. (2020) extend the similar idea to the offline model-based RL setting.
To calculate the policy gradient by back-propagating through the model: SVG (Heess et al., 2015) presents a framework for learning continuous control policies using back-propagating through the model. More recent works (Clavera et al., 2019; Amos et al., 2021) extend the idea by combining the learned Q-function and the SOTA policy optimization algorithms and thus achieve better performance. MAGE (D’ Oro & Jaśkowski, 2020) computes the policy gradient using value gradient and proposes to reduce the value gradient through TD learning.
Different from the above method, we further consider the influence of the model gradient error with respect to the convergence of the policy optimization. We will use multiple function approximators to separate different requirements for the model at the learning and using phase. Moreover, we can leverage the insight inside of these related works to better design our model-based algorithms.

heading: 4 MODEL-BASED POLICY OPTIMIZATION BY CONSIDERING THE MODEL
text: GRADIENT ERROR
In this section, we will analyze the influence of the model error in the policy optimization process. First of all, for the policy optimization algorithms that using policy gradient calculated according to the learned model, we theoretically show that both the prediction error and the gradient error of the learned model matter for the policy gradient bias and the convergence rate. Secondly, based on the theoretical findings, we propose our method to learn the accurate model for policy optimization. Thirdly, we show the implementation details of our method and propose the DDPPO algorithm.

heading: 4.1 CONVERGENCE RATE FOR MODEL-BASED POLICY OPTIMIZATION
text: In this section, we will prove the finite sample bound for the model-based policy optimization algorithms. Due to the space limitation, we put the full proof into the Appendix C.
Due to the space limitation, we simply put the assumptions that is necessary for understanding the theoretical results here and put the detailed assumptions into Appendix C. Assumption 1. The transition model, the reward model and the policy function are all Lipschitz continuous. The state space and the reward are both bounded. For a fixed policy, the stationary distribution exists and has the uniform periodicity property:
dTV (P (st+τ ∈ ·|s0 = s), µθ(·)) ≤ mρτ ∀τ ≥ 0, ∀s ∈ S. (6)
First of all, we analyze the bias of the policy gradient introduced by the model error. Recall the policy gradient formulation (Equation 3) in the section 2. In the model-based setting, we need to use the learned model to calculate the gradient value ( drtdat , drt dst , dskdst+1 , dst+1 dat
) in Equation 3. Since the model error is unavoidable, it is necessary for us to analyze the bias of the policy gradient introduced by the model error. Theorem 1. Suppose assumptions hold. The policy gradient bias ϵM can be bounded as
ϵM = ∥∥∥∥∥E ( dRM̂ dϕ − dR M dϕ )∥∥∥∥∥ ≤ C6ϵgr + C7ϵgf + C8ϵf , (7) where C6, C7, C8 are three constants related to the Lipschitz constants of the transition model and the reward model and γ.
Remark: The above theorem shows that the difference between the policy gradient calculated using two different models can be upper bounded by the difference of the gradient of the reward modelϵgr , the gradient of the transition modelϵgf , and the prediction of the transition modelϵf . The result suggests that gradient error matters in estimating the policy gradient. Now, one natural conjecture is that the gradient error will influence the convergence of the policy optimization process. The next theorem depicts the influence rigorously.
Secondly, we analyze the convergence rate for the policy gradient algorithms. Since it is always a non-convex optimization problem, the commonly used measure is the speed of the gradient norm converges to zero. Recall the update rule (Equation 4). We update the policy parameters by following the estimated policy gradient direction. Intuitively, the model error will influence the convergence. Theorem 2. Suppose assumptions 3, 4, 5, 6,7, and 8 holds. LJ is the Lipschitz constant of the objective function for the parameter ϕ and Lgs is the smoothness constant of the objective function for the parameter ϕ. Superpose the policy parameters are updated as shown in Equation 4 and the learning rate is set as αt. We have
min 0<t<T E ∥∥∥∥ dJdϕt ∥∥∥∥2 ≤ 2∑T
t=1 αt
( 2UJ + 2τα0L 2 J +
T∑ t=1 α2t ( L2JL g J 2τ + LgJα 2 tL 2 R )) + 2LJmρ τ + ϵ2M .
(8)
Remark: The above theorem shows the rate of the model-based policy gradient method converges to its stationary point. This setting is commonly used in the previous papers that analyze the convergence rate for the actor critic method with function approximation and non-convex optimization (Wu et al., 2020; Kumar et al., 2019; Qiu et al., 2021; Ghadimi & Lan, 2013). The results show that if we set the learning rate αt that satisfies the condition limT→∞ ∑T t=1 α 2 t∑T
t=1 αt = 0, the policy optimization process will converge to near its stationary point with rate O (∑T t=1 α 2 t∑T
t=1 αt
) . As an example, if we
set the learning rate as αt = 1√t , we can get the convergence with rate O( 1√ T ).
Remark: The bias is introduced by the policy gradient bias and the Markov noise. The bias introduced by the Markov noise can be arbitrarily small if we set τ sufficiently large which doesn’t change the main order of the convergence rate. The bias introduced by the policy gradient bias is analyzed in Theorem 1. According to the theoretical results, to reduce policy gradient bias, we need to control the model prediction error and the model gradient error in the model learning phase.

heading: 4.2 TWO-MODEL-BASED LEARNING POLICY OPTIMIZATION
text: The theoretical result in the previous subsection indicates that both the model prediction error and the model gradient error are crucial for policy optimization. In this subsection, we will discuss how to learn and how to use the model taking both errors into account. Specifically, our approach repeatedly alternates between collecting samples from the environment, training the model, and optimizing the policy.
Model Learning: Different from the classical model-based methods, we design a two-model-based learning method to achieve the goal. We call these two models prediction models and the gradient models respectively. Our goal is to use the prediction models to predict the next state and the reward and to use the gradient models to calculate the policy gradient. Training of the prediction model is commonly through supervised learning, e.g., maximum likelihood with early stopping on a validation set (Janner et al., 2019; Clavera et al., 2019). Training of the gradient model is more complex. For the purpose of reducing the gradient error, we need to use the data collected from the environment to estimate the gradient of the environment model and construct the mean square gradient error for the parameter optimization.
Policy Optimization: To update the policy, we first generate trajectories by unrolling the latest policy on the prediction model. Then, we calculate the gradient of the objective by back-propagating through the gradient model. Finally, we update the parameters in the policy by following the policy gradient direction.

heading: 4.3 DIRECTIONAL DERIVATIVE PROJECTION POLICY OPTIMIZATION: A PRACTICAL
text: IMPLEMENTATION
In this subsection, we propose Directional Derivative Projection Policy Optimization (DDPPO) as an implementation of our proposed two-model-based learning method. The key question is how to use the data collected from the environment to estimate the gradient or Jacobian of the environment model. Note that the dimension of the Jacobian of the learned environment model is the product of the state’s dimension and the action’s dimension, which is usually large. Directly estimate the
Algorithm 1 Directional Derivative Projection Policy Optimization 1: Initialize the policy πϕ, predictive model M̃pθp , gradient model M̃ g θg
, value funcrion Qψ 2: Initialize environment replay buffer Denv , model replay buffer Dmodel 3: repeat 4: Take an action in real environment with policy πϕ, and add transition (s, a, s′) to Denv 5: Train predictive model M̃pθp and update θp by maximizing JM̃p(θp) 6: Train gradient model M̃gθg and update θg by maximizing JM̃g (θg) 7: for Nrollouts steps do 8: Sample st uniformly from Denv 9: Perform k-step model rollout on M̃pθp starting from st with policy πϕ
10: Add samples to Dmodel 11: end for 12: D ← Denv ∪ Dmodel 13: for NupdateQ steps do 14: Update ψ using data form D: ψ ← ψ − λQ∇ψJQ(ψ) 15: end for 16: for Nupdate Policy steps do 17: Sample trajectories of length H with policy π and predictive model M̃pθp 18: Update ϕ using sampled trajectories and gradient model M̃gθg : ϕ← ϕ+ λπ∇πJπ(ϕ) 19: end for 20: until The policy performs well in real environment M 21: return Optimal policy πϕ∗
gradient or Jacobian matrix is intractable. For model-based RL, we cannot estimate gradients by perturbing one dimension of a state or action and estimate the partial derivative, because we can only sample data from the environment, and this is contrary to the purpose of efficiency. Besides these, we will also discuss how to design two models to rollout and calculate gradients respectively when optimizing the policy. The overall algorithm pseudo-code is shown in Algorithm 1. More details for explaining the pseudo-code can be found in Appendix B.1.
Model Learning: Prediction loss and the gradient loss.
The prediction model M̃p is trained following Janner et al. (2019) via maximum likelihood (Equation 2) To improve the ability of models to portray complex environment, we use a bootstrap ensemble of models {M̃θ1 , . . . , M̃θB} which is consistent with Janner et al. (2019); Clavera et al. (2019).
Different from the prediction model, the gradient modelM̃g aims to better characterize the gradient of the environment instead of an accurate prediction. Hence, we hope to trade-off the original maximum likelihood objective by a gradient loss term that constraints the Jacobian of the learned model. Although the model’s gradient or Jacobian matrix is difficult to estimate directly on account of the high dimensions, the directional derivative can be estimated by finite difference using data sampled. Meanwhile, the directional derivative of the learned model can be calculated by projecting the model’s gradient or Jacobian matrix into a certain direction.
Hence in this paper, we use the error between the estimated directional derivative and the projection value to constrain the learned model’s gradient. Note that for an input x = (s, a), and the corresponding output M(x) with dimension d, the projection of the Jacobian along a vector v, i.e., ∇xM(x) · v||v|| , can be estimated by finite difference, M(x+v)−M(x) ||v|| . Therefore, we use the finite difference as a target to constraint the Jacobian of the gradient model. Since the finite difference converges to the directional derivative when the norm of v goes to zero, the data point x+ v should be close to x. So in this paper, we sample the nearest n data points of x to calculate n directional derivatives in n different directions and use the finite difference to constraint the Jacobian of the learned model. Sampling the nearest points is for a more accurate estimation of the directional derivatives since the finite difference is closer to the true directional derivative if the two inputs are closer. Then, the gradient loss item can be constructed as follows:
R = 1
n n∑ i=1 1 d ||∇xM̃(x) · vi − (M̃(xi)− M̃(x))||2, (9)
where vi denotes directions xi − x, and xi are the nearest n points of x in the buffer. Intuitively, the gradient at x can be constrained if there are enough points close enough to x. And a relatively small gradient loss item means the learned model should have a gradient such that its first-order approximation is consistent with the given data. By leveraging the gradient loss term, the gradient mode M̃g , is trained via a multi-objective loss:
JM̃g (θg) = E [ log M̃g(st+1, rt|st, at) ] + w ·R. (10)
Note that here we don’t prefer to learn gradient model using gradient loss term only. We can just use the data in the buffer to estimate a few of the directional derivatives since the points sampled should be close to the x to make the estimation correct. So here we also use the prediction loss to constrain the hypothesis space of the learned model to prevent large bias local optimal. Finally, the trained gradient model and the prediction model will be used for policy optimization.
Model Usage: Policy optimization by back-propagating through model. In our DDPPO algorithm, to train the policy, the objective function we use is constructed by following SAC (Haarnoja et al., 2018), since SAC can achieve outstanding performance. The policy is learned by rolling out the prediction model for H steps and maximizing the following objective function:
Jπ(ϕ) = E [ H−1∑ t=0 γt(r̃t − α log πϕ(at|s̃t)) + γH(Q̃(s̃H , aH)− α log πϕ(aH |s̃H)) ] , (11)
where s̃t+1 ∼ M̃pψp , s̃0 random sampled from the environment buffer Denv , a ∼ πϕ, α is the temperature hyperparameter proposed in SAC. Note that the policy can be learned by just computing the derivative of the objective since the learned prediction model M̃p is differentiable. However, in our approach, we replace the gradient ∇s̃t,atM̃p(s̃t, at) by ∇s̃t,atM̃g(s̃t, at) to calculate the full gradient of the objective, believing in that the gradient model can capture the true gradient better than the prediction model.
As for the Q-value function, it can be trained by minimizing the Bellman residual. In practice, we follow SAC (Haarnoja et al., 2018) that trains two Q-functions to mitigate positive bias which is known to degrade the performance of value-based methods (Hasselt, 2010; Fujimoto et al., 2018), and make use of target Q-function which has been shown to stabilize training (Mnih et al., 2015).

heading: 5 EXPERIMENT RESULTS
text: The goal of our experimental evaluation is to answer the following questions: (1) How does our algorithm perform on benchmark reinforcement learning tasks, compared to state-of-the-art modelbased and model-free algorithms? (2) Does it matter to explicitly learn the gradient of the model by using the gradient loss? (3) How does the sensitivity of our algorithm when changing the hyperparameters n and w that are related to the gradient model learning? (4) Why build two models? To answer these questions, we evaluate our approach on model-based continuous control benchmarks in the MuJoCo (Todorov et al., 2012). Experimental details can be found in Appendix B.2.

heading: 5.1 COMPARISON WITH STATE-OF-THE-ARTS
text: To answer the first question, we compare our method against state-of-the-art model-based and model-free methods in terms of sample complexity and asymptotic performance. Model-free baseline includes SAC (Haarnoja et al., 2018), which is a widely accepted baseline. For model-based methods, we select MBPO (Janner et al., 2019) and STEVE (Buckman et al., 2019) as our baseline which both use short-horizon model-based rollouts. We also compare to model-based methods SVG (Heess et al., 2015), which firstly uses the derivative of models to learn the policy, and MAAC, (Clavera et al., 2019) similar to SVG, however entirely relies on the predictions of the model, removing the need for likelihood ratio terms. MAGE (D’ Oro & Jaśkowski, 2020) is added to comparison since this method focuses on the gradient of the Q-function.
The learning curves for all methods, shown in Figure 1, highlight the strength of our method in terms of performance and sample complexity. In all the MuJoCo simulator environments, especially in two higher-dimensional tasks, Ant and Humanoid, our method learns faster and has better efficiency than previous model-based and model-free methods. Note that MAGE learns fastest on Humanoid, and the trick of this algorithm can also be applied to our method. However, even without the help of tricks in MAGE, we can still achieve better asymptotic properties. We get a return around 5400 in the Humanoid environment using just over 100 thousand interactions with the environment, while around three times as many interactions are needed for MBPO to achieve the same performance.

heading: 5.2 GRADIENT ERROR ANALYSIS
text: In this subsection, we answer the second question, i.e., Does it matter to explicitly learn the gradient of the model by using the gradient loss? Firstly, we aim to analyze the difference between the gradient of the true model and the learned model. We collect experience using a random policy and train the model with and without gradient loss respectively. Then we calculate the MSE error of the gradient of the model. Figure 2 shows how the gradient error and the prediction error changes during training of the model on two simple environments that the true gradient is analytical available. From the results in Figure 2 we can see that the prediction error for both models decreases to zero while the gradient error of the model without gradient loss keeps larger than that of the model with gradient loss. The learning curves highlight the strength of adding a gradient loss to the model in terms of high prediction accuracy and accurate gradient estimation under the simple case.

heading: 5.3 ABLATION STUDY
text: To answer the third question, i.e., how does the sensitivity of our algorithm when changing the hyperparameters, and to show the results in a comprehensive way, we use a low-dimensional environment, Hopper, and a high-dimensional environment, Ant, as our representative and then conduct two sets of experiments in them. Figure 3 shows the results. The left column shows how the efficiency of our algorithm changes by tuning n. The results demonstrate that, regardless of the value n, our method is better than our variant without gradient loss. The middle column shows the results for tuning w. The results indicate that the weight is recommended to be smaller for a higher-dimensional environment and the performance is not sensitive to the number of directional derivatives used to construct the gradient loss, even only one directional derivative used can also help to improve the performance.
Now we answer the last question, i.e., why build two models. To investigate the importance of the gradient loss and the important role of two-model-based, we construct the ablation study by removing the gradient loss in our method and retrain the policy (without gradient loss) to see the importance of considering the gradient loss. To see the impact of training two models for prediction and gradient calculation respectively, we design a variant of our method that trains only one model with gradient loss used for both purposes (single model). The experimental results are shown in the right column of Figure 3. The results underpin the importance of considering the gradient error when training the model if back-propagating through it. On the other hand, the results highlight the strength of the two-model-based structure. We also draw the prediction error of both models in our algorithm and the results are shown in Appendix B.3. The prediction error of the prediction model is lower than that of the gradient model, which further validates the idea of two-model-based learning. To address the concern that the improvement of the performance may be caused by more parameters, we design a variant of our model that trains a model with double parameters, however without gradient loss. The experimental results are shown in Figure 4 that yields the importance of considering gradient information.

heading: 6 CONCLUSIONS AND FUTURE WORK
text: In this paper, we theoretically analyze how the model gradient error and the model prediction error influence the policy gradient bias and the convergence rate of model-based policy optimization. The theoretical results motivate us to design a two-model-based learning method for policy optimization. We design the learning and using methods for the two models respectively by taking the model gradient information into account and propose DDPPO algorithms as a practical implementation. Our algorithm can achieve superior sample efficiency than state-of-the-art model-based reinforcement learning algorithms on challenging high-dimensional continuous control tasks. For future work, it is enticing to design a better gradient estimator in order to optimize the gradient model more efficiently. The idea about learning the model according to its usage instead of only considering its prediction error is also general enough to be applied to other problems beyond reinforcement learning.

heading: A RELATED WORK
text: Reinforcement learning (RL) is a powerful technique for solving the sequential decision making problems (Li, 2018; Sutton & Barto, 1998). The sequential decision-making problem is always formulated as the Markov decision process (MDP) framework. The agent uses policy to determine its action according to the current state of the environment, interacts with the environment, and gets the immediate reward from the environment as feedback. Then, the environment changes to the next state. The goal for reinforcement learning is to find the optimal policy that can achieve the highest expectation of the summation of the future reward. Recent work on model-based RL (Nagabandi et al., 2018; Luo et al., 2018; Kurutach et al., 2018; Wang et al., 2019; Janner et al., 2019; Pan et al., 2020), has shown the power of first learning the environment model and then use it to do the policy optimization. These methods all achieve similar performance by using fewer data compare to their model-free counterparts.
Model-free and Model-based Method: Model-free approaches, as a general-purpose tool for learning complex policies (Mnih et al., 2015; Lillicrap et al., 2016; Haarnoja et al., 2018), has the problem of low-efficiency (Janner et al., 2019), which limits the application in real-world physical systems where data collection can be an arduous process. Model-based Reinforcement Learning attempts to reduce sample complexity while maintaining the asymptotic performance. The learned model can be viewed as a black-box simulator and then used for training a model-free policy (Nagabandi et al., 2018; Luo et al., 2018; Kurutach et al., 2018; Wang et al., 2019; Janner et al., 2019; Pan et al., 2020). Tools like ensembles (Kurutach et al., 2018), probabilistic models (Chua et al., 2018), and meta-learning (Clavera et al., 2018) are widely used in model-based RL. A major problem of model-based RL is model-bias (Deisenroth & Rasmussen, 2011) caused by a compounding error effect of long-term predictions. To address this problem, a popular approach is to use interpolation between different horizon predictions (Buckman et al., 2019; Janner et al., 2019) and interpolating between model and real data (Kalweit & Boedecker, 2017). Different from the traditional model based method, we argue to learn the model beyond the accurate prediction and to use the model more efficiently by calculating the policy gradient analytical through the learned model.
Improved Model-based Methods: Several issues have been studied to achieve better performance in model-based RL. Malik et al. (2019) explores which uncertainties are needed for model-based RL. Further research implements a masking mechanism based on the model’s uncertainty estimation to decide whether the model should be used or not (Pan et al., 2020). Dropout mechanisms are also applied to model-based RL algorithms to improve the robustness while maintaining high sampling efficiency (Zhang et al., 2021b). Efficient explorations is also considered in model-based RL (Curi et al., 2020; Song & Sun, 2021; Yao et al., 2021). For model learning, Ke et al. (2018) use latent variables to carry future information that improves long-term predictions, Tomar et al. (2021) use model-invariance state abstraction to improve generalization. Kégl et al. (2021) firstly design metrics to evaluate the various popular generative models when using them on the control problem. Zhang et al. (2021a) shows that tuning of hyperparameters dynamically improves the performance compared to using static hyperparameters which are kept fix for the whole training. This operation can be seen as a simple treatment of model uncertainty. Considering that we also need the prediction model at the model learning and using phase, we can leverage the useful ideas inside of these methods to help us to build a robust prediction model and use it efficiently.
Decision-aware Model Learning:Decision-aware model learning (DAML) (Lovatto et al., 2020; Nair et al., 2020; Abachi et al., 2021; D’Oro et al., 2020; Farahmand et al., 2017; Modhe et al., 2021) considers the problem that how we can learn the model that can be better leveraged by the policy optimization or Q learning. D’Oro et al. (2020) using policy gradient to re-weighted the model prediction error for each sample. Farahmand et al. (2017) use the value function to re-weighted the model prediction error for each sample. Please note that, although the main goal about learning the model by considering the follow-up decision making problems seems very similar, our ideas and methods are totally different from DAML. DAML is still the traditional rollout-based model-based method with a re-weighting prediction loss.
Differentiable Model-based Method: A more effective way of using the model is to exploit its differentiability. SVG (Heess et al., 2015) presents a framework for learning continuous control policies using back-propagating through the model. However, SVG just uses real trajectories leveraging likelihood ratio term which in turn increases the variance of the gradient estimate. Amos
et al. (2021) extend SVG further leveraging the soft update (Haarnoja et al., 2018) and model-based value expansion (Feinberg et al., 2018). Different from SVG, MAAC (Clavera et al., 2019) entirely relies on the predictions of the model, removing the need for likelihood ratio terms, and achieves better performance than MBPO leveraging model-based value expansion for updating policy and Qfunction. The plan horizon in these methods allows to trade-off between the accuracy of the learned model and the accuracy of the learned Q-function. In addition to back-propagating through the path, MAGE (D’ Oro & Jaśkowski, 2020) uses the differentiability of the learned model to compute gradient targets in temporal difference learning. Our method is different from these methods. Instead of Different from these methods, our main contribution is that we propose to use two models to learn the model prediction and the model gradient respectively. By explicitly construct the gradient estimator, our gradient model can produce a more accurate policy gradient and thus benefits for the policy optimization.
Outside of RL: Several algorithms outside of RL adopt gradient penalties into the loss function when training a neural network, known as double backpropagation (Drucker & Le Cun, 1992). Sobolev training (Czarnecki et al., 2017) is most related to this topic, which tries to learn both the value and gradient of a target function during supervised training. However, Sobolev training requires the ground-truth gradient available. In the model-based reinforcement learning setting, we do not know the ground-truth gradient.

heading: B ALGORITHM DETAILS
text: B.1 IMPLEMENTATION OF DDPPO

heading: B.1.1 DIRECTIONAL DERIVATIVE
text: The directional derivative of a function f with respect to a vector v at a point (e.g., position) x may be denoted by ∇vf(x) that is a limit: ∇vf(x) = limh→0 f(x+hv)−f(x)h . If the function f is differentiable at x then the directional derivative exists along any vector v and we have: ∇vf(x) = ∇f(x) · v|v| . Note that if f is a scalar function and the dimension of x is high, then ∇f(x) is a vector and difficult to estimate, however ∇vf(x) is a scalar, and easy to estimate by finite-difference. This intuition helps us to design the constraint on the gradient of the learned model using data samples.

heading: B.1.2 BACK-PROPAGATING THROUGH GRADIENT MODEL
text: In our method, we use the prediction model to rollout, however use the gradient to calculate the gradient when performing policy optimization. This procedure is like SVG (Heess et al., 2015) that gets predictions using the true environment however gets the gradient using a learned model. Note that reparameterization and important sampling weights are needed for this method. In our paper, we use the learned model as a deterministic model that only uses the mean outputted as the prediction. Now we show an example of calculating the gradient. When predicting, we first receive a state s and use the policy to get action a, then using the prediction model to get the next state and reward s′, r. Note that we need to memory the computation graph for getting the action. Then we use the gradient model to calculate the gradient by feeding the (s, a) to it and back-propagation. Combined with the memorized computing graph we can get the final policy gradient.
B.1.3 IMPLEMENTATION
Our algorithm can be divided into three main parts: model learning, policy optimization, and Qfunction learning. And our approach alternates between collecting samples from the environment, training the model, updating Q-function, and policy optimization. The overall algorithm pseudocode is demonstrated in Algorithm 1.
Firstly, we collect data using the current policy and add it into the environment replay buffer Denv . Then the prediction model is trained using data from Denv . The gradient model is also trained with the multi-objective loss. The third step is to perform k step model rollout on the prediction model M̃pθp with the current policy, where k is increased over time which proposed by Janner et al. (2019) to achieve better performance. Those samples are appended to the replay buffer Dmodel, together with samples from Denv , used to update the Q-function. To update the policy, firstly, we obtain
imaginary trajectories by unrolling the latest policy on the prediction model M̃pθp for H steps from a randomly sampled state in Denv . Then we update policy by calculating the gradient of the objective in equation 11, while using the gradient model M̃gθg instead of M̃ p θP .
In our approach, we apply the soft Q-update (Haarnoja et al., 2018) which minimizes the following objective:
JQ = E [( Q̃(st, at)− ( r(st, at) + γQ̃(st+1, at+1)− α log(πϕ(at+1|st+1)) ))2] . (12)
The temperature hyperparameter is automatically tuned by adjusting the expected entropy over the visited states to match a target value (Haarnoja et al., 2018) with the following objective:
Jα = E [ −α log πϕ(at, st)− αH̄ ] (13)
where H̄ is the target entropy.

heading: B.2 EXPERIMENT SETTINGS

heading: B.2.1 ENVIRONMENT SETTINGS
text: We evaluate our approach on six continuous control benchmark tasks in the MuJoCo (Todorov et al., 2012) simulator in our experiments: InvertedPendulum-v2, Hopper-v2, Walker2d-v2, HalfCheetahv2, Ant-v2, Humanoid-v2. Standard full-length versions of these tasks are used. Note that the Ant and Humanoid environments are truncated observations which is consistent with MBPO (Janner et al., 2019). The code for building Ant and Humanoid environment is provided by Janner et al. (2019).1
B.2.2 HYPERPARAMETER SETTINGS
0 25k 50k 75k 100k Steps
0
500
1000
1500
2000
2500
3000
3500
Av er
ag e
Re tu
rn
Hopper
ours(H=1) ours(H=2) ours(H=3) MBPO Table 1 shows the hyperparameters used for DDPPO to achieve results shown in Figure 1. G1 is the number of updates for policy per environment step, G2 is the number of updates for critic per environment step. H is the rollout length for updating policy. n is the number of directional derivatives used to construct the gradient loss. w is the coefficient in front of the gradient loss when updating the gradient model. Note that H is used to trade-off between the accuracy of the learned model and the accuracy of the learned Q-function. As shown in the right figure, the results by setting a small H is close to the results of MBPO, while a larger H may cause worse results. We try different steps and propose 2 for Hopper and 3 for other environments. We employ 5000 (200 for the InvertedPendulum environments) warmup steps of interaction with the environment before starting to update the actor and the critic. The network structure of both the prediction model and the gradient model is the same
1https://github.com/JannerM/mbpo/tree/master/mbpo/env
as MBPO. Both models are trained after every 250 steps of interaction with the environment. Every time after model training, we perform rollout for k steps and add the collected data to the buffer used for updating the critic. k is increased with the number of interactions. For InvertedPendulum, Ant, HalfCheetah, and Walker2D k keeps fixed. For Hopper k increases from 1 to 15 between epoch 25 to epoch 155. For Humanoid k increases from 1 to 25 between epoch 25 to epoch 155. The policy and Q-value function are updated 10 times per environment step which is less than the number of MBPO. For updating the policy we rollout on the prediction model for H steps with batch size 256.

heading: B.2.3 EXPERIMENT DETAILS FOR BASELINES
text: For MBPO, we directly use the reported number given by Janner et al. (2019)2; For SAC we use the codes and hyperparameters available3; For MAGE, we use the codes available4 and hyperparameters from the author. For MAAC, the model structure and hyperparameters such as learning rate are the same as MBPO, and the extra hyperparameter, H , we try different steps and set 2 for hopper and 3 for others which is consistent with ours. For other baselines, we use the codes or results from Wang et al. (2019)5. Note that the number of layers and the number of nodes of Q value functions are all the same for SAC, MBPO, MAAC, and ours. The number of layers and the number of nodes of the learned model are all the same for MBPO, MAAC, and ours. Finally, we keep unimportant hyperparameters of our method as consistent as possible to SAC or MBPO to keep it as fair as possible.

heading: B.3 MORE RESULTS FOR ABLATIONS
text: The results in Figure 3 show the importance of considering the gradient error when training the model if back-propagating through it and highlight the designed two-model-based learning procedure. We extend the experiments in Section 5.3 to more environments. The results in Figure 5 show the performance of our method compared to the variants of our method. The results in Figure 6 show that the prediction models in our algorithm have lower prediction error than the gradient models, which confirm our conjecture that the prediction and gradient calculation should be done by the corresponding model.
2https://github.com/JannerM/mbpo 3https://github.com/pranz24/pytorch-soft-actor-critic 4https://github.com/nnaisense/MAGE 5https://github.com/WilsonWangTHU/mbbl
25000 50000 75000 100000 Env Steps
0.000
0.005
0.010
0.015
0.020
0.025
0.030
Pr ed
ict iv
e er
ro r o
n va
lid at
io n Hopper
10000 15000 20000 25000 Env Steps
0.050
0.075
0.100
0.125
0.150
0.175
0.200
Pr ed
ict iv
e er
ro r o
n va
lid at
io n Ant
0 25000 50000 75000 100000 Env Steps
0.2
0.4
0.6
0.8
1.0
Pr ed
ict iv
e er
ro r o
n va
lid at
io n HalfCheetah
5000 10000 15000 20000 Env Steps
0.0000 0.0001 0.0002 0.0003 0.0004 0.0005 0.0006
Pr ed
ict iv
e er
ro r o
n va
lid at
io n InvertedPendulum
0 25000 50000 75000 Env Steps
0.2
0.4
0.6
0.8
1.0
Pr ed
ict iv
e er
ro r o
n va
lid at
io n Walker2d
50000 100000 150000 Env Steps
0.1
0.2
0.3
0.4
0.5
Pr ed
ict iv
e er
ro r o
n va
lid at
io n Humanoid
gradient model prediction model

heading: C PROOFS OF THE THEOREM
text: In this section, we present the proof of the main theorems in section 3.
First of all, we summarize the necessary assumptions here. These assumptions are commonly used for analyzing reinforcement learning algorithms. Assumption 2 (Stationary). For ∀t1 < t2 and s0t1 , s 1 t1 , denote τ = t2 − t1 there exist L1(τ) < ∞,
τ = t2 − t1, such that ∥∥∥∥E( s1t2−s0t2s1t1−s0t1 )∥∥∥∥ ≤ L1(τ). And L1(τ) can be uniformly upper bounded by L1 < ∞. Assumption 3 (Bounded region). There exist a finite D < ∞ such that: (Wu et al., 2020; Wang et al., 2017)
∥s1 − s2∥ ≤ D for ∀s1, s2 ∈ S. (14) Assumption 4 (Lipschitz). The transition model and the reward model are both Lipschitz continuous with Lipschitz constant Lf and Lr respectively (Clavera et al., 2019). Assumption 5 (Uniform periodicity). For a fixed policy πθ(·|s0), denote µθ as the stationary distribution induced by the policy and the transition probability measure P (·|s, a). Then there exists m < ∞, C0 < ∞, and ρ ∈ (0, 1) such that:(Wu et al., 2020; Zhang et al., 2020a; Qiu et al., 2021)
dTV (P (st+τ ∈ ·|s0 = s), µθ(·)) ≤ mρτ ∀τ ≥ 0, ∀s ∈ S. (15) Assumption 6 (Lipschitz policy). Let πϕ(a|s) be a policy parameterized by θ. there exist constants such that for all given state s and action a we have:(Wu et al., 2020; Clavera et al., 2019; D’ Oro & Jaśkowski, 2020) ∥∥∥∥dπϕ(a|s)dϕ
∥∥∥∥ ≤ Lπ. (16) Assumption 7 (Step-size). The step-size αt is non-increasing and non-negative (Wang et al., 2017; Wu et al., 2020). Assumption 8 (Bounded reward). The reward is bounded in the region [0, rmax] (Clavera et al., 2019; Wu et al., 2020).
Here, we show some lemmas that is important for the proof.
Lemma 1. Suppose assumptions hold, denote the C0 = min(Dϵf , 1−Lt+1f 1−Lf ) for ∀t we have
∥s̃t − st∥ ≤ C0ϵf (17) Lemma 2. Denote C1 = tL21, C2 = tC0LfL21. We have
δ ( E ( dst ds0 )) = ∥∥∥∥∥E ( d̂st ds0 − dst ds0 )∥∥∥∥∥ ≤ C1ϵgf + C2ϵf (18) Lemma 3. If constent 0 < γ < 1, for arbitrary constant C, we have:
∞∑ t=0 tγtC = γ (1− γ)2 C (19)
Lemma 4. Denote constants C3 = L11−γ and C4 = LrL
2 1γ
(1−γ)2 ,C5 = C0L1Lf (1−γ) + C0LfLrγ (1−γ)2 , then we have:
δ
( dR
ds
) = ∥∥∥∥∥E ( dRM̂ ds − dR M ds )∥∥∥∥∥ ≤ C3ϵgr + C4ϵgf + C5ϵf (20) Lemma 5. Considering a Markov process M with stationary distribution µ. Data ξt are sampled from the M . There are two functions of parameter ϕ: G(ϕt, ξt) and g(ϕt) where
g(ϕt) = Eξ∼µG(ϕt, ξ). (21) Denote the upper bound for the function G as L1. Denote the Lipschitz constant for the function G as L2.
ϕt+1 = ϕt − αtG(ϕt, ξt) (22) For all 0 < T < ∞, we can proof that:∥∥∥∥∥ T∑ t=1 αt ⟨g(ϕt), G(ϕt, ξt)− g(ϕt)⟩ ∥∥∥∥∥ ≤ 2τα0L21 + T∑ t=τ+1 α2tL 2 1L2τ + T∑ t=1 αtL1mρ τ (23)
Lemma 6. Denote UJ as the upper bound of the expected sum of discounted rewards J . Then UJ ≤ rmax1−γ Lemma 7. There exist two constants LJ and LgJ such that for all ϕ1, ϕ2, we have
∥J(ϕ1)− J(ϕ2)∥ ≤ LJ ∥ϕ1 − ϕ2∥ (24)∥∥∥∥ dJdϕ1 − dJdϕ2 ∥∥∥∥ ≤ LgJ ∥ϕ1 − ϕ2∥ (25)
Lemma 8. There exist constants LR such that for all ϕ1, ϕ2, we have
E ∥R(ϕ1)−R(ϕ2)∥ ≤ LR ∥ϕ1 − ϕ2∥ (26) (27)
Lemma 9. Suppose the function J is smooth with smoothness constant LgJ
J(ϕ2) ≥ J(ϕ1) + ⟨ dJ
dϕ1 , ϕ2 − ϕ1⟩ − LgJ 2 ∥ϕ2 − ϕ1∥ (28)
Lemma 10. If x, y, z > 0 and √ x ≤ √y + √ z, then we have
x ≤ 2y + 2z (29)

heading: C.1 PROOF OF THEOREM 1
text: Overall speaking, the main idea behind the proof of Theorem 1 can be summarized as follows.
Considering that to calculate a gradient value of a function output w.r.t its input, we need to specify the function and the input value. Therefore, as in the Equation 3, both the function(f̃ and r̃) to be taking gradient and the input value s̃t are different from the ground-truth. So we need to first bound the error of the prediction state(Lemma 1) and then bound the error of the gradient of the learned model in true state (Lemma 4).
Denote C6 = LfLπ+LfC3Lπγ 1−γ , C7 = γLfLπC4+L1Lπ 1−γ , C8 = C0Lf+γC5LfLπ+L1LπLfC0
1−γ , we can prove the theorem by using previous lemmas:∥∥∥∥∥E ( dRM̂ dϕ − dR M dϕ )∥∥∥∥∥ (30) =
∥∥∥∥∥E ( ∞∑ t=0 γt ( d̂r̃t dϕ − drt dϕ ) + γ ∞∑ t=0 γt ( d̂R ds̃t+1 d̂s̃t+1 dãt dãt dϕ − dR dst+1 dst+1 dat dat dϕ ))∥∥∥∥∥ (31) ≤ 1 1− γ (LfC0ϵf + ϵ g r)Lπ + ∥∥∥∥∥E ( γ ∞∑ t=0 γt ( δ ( dR ds ) d̂s̃t+1 dãt + dR dst+1 δ ( ds da )) Lπ
)∥∥∥∥∥ (32) ≤ 1 1− γ (LfC0ϵf + ϵ g r)Lπ + γ 1− γ LfLπ(C3ϵ g r + C4ϵ g f + C5ϵf ) + 1 1− γ L1Lπ(LfC0ϵf + ϵ g f )
(33)
=C6ϵ g r + C7ϵ g f + C8ϵf . (34)
The first equation is the definition. The second inequality follows the assumption 2, the right term . The third term follows the lemma 4

heading: C.2 PROOF OF THE THEOREM 2
text: First of all, we briefly summarize the proof idea here. We first decompose the difference of objective function brought by parameter updating into two terms according to its smooth assumption(Equation 35). The first term is an inner product term. The inter product term can be further decompose into two terms(Equation 36). One is about the model error and the other one is about the non-i.i.d. data. We can carefully bound each terms and then combine them all together to proof the theorem. Readers can refer to the related papers Wang et al. (2017) Wu et al. (2020) Ghadimi & Lan (2013) to get more insight of the proof.
According to the lemma 9, we can have:
J(ϕt+1) ≥ J(ϕt) + αt
〈 dJ
dϕt , dRM̂ dϕ
〉 − LJ
2 ∥ϕt+1 − ϕt∥2, (35)
where dR M̂
dϕ is the estimated policy gradient calculated using learned model and used to update the policy parameter from ϕt to ϕt+1.
Denote h1t = dRM̂ dϕ − dRM dϕ , h 2 t = dRM dϕ .
We can further decompose the value improvement as J(ϕt+1) ≥ J(ϕt) + αt (〈 dJ
dϕt , h1t
〉 + 〈 dJ
dϕt , h2t
〉) −
LgJ 2 ∥ϕt+1 − ϕt∥2. (36)
According to the Cauchy–Schwarz inequality we have
E 〈 dJ
dϕt , h1t
〉 ≥ − √ E ∥∥∥∥ dJdϕt ∥∥∥∥2 E ∥h1t∥2 ≥ − √ E ∥∥∥∥ dJdϕt ∥∥∥∥2(C6ϵgr + C7ϵgf + C8ϵf ). (37)
Rearrange the term we have E 〈 dJ
dϕt , h2t
〉 = E 〈 dJ
dϕt , dRM dϕt − dJ dϕt
〉 + E 〈 dJ
dϕt , dJ dϕt
〉 . (38)
Note that the first term in the right hand side is 0 if the data is I.I.D. sampled from the stationary distribution. If the data is sampled from the markov process, the data is non-i.i.d.. we can further bound the first term using the Lemma 5
T∑ t=1 αtE 〈 dJ dϕt , dRM dϕt − dJ dϕt 〉 ≥ − ( 2τLJα0 + T∑ t=1 α2tL 2 JL g Jτ + T∑ t=1 αtLJmρ τ ) . (39)
The second term is
E 〈 dJ
dϕt , dJ dϕt
〉 = E ∥∥∥∥ dJdϕt ∥∥∥∥2 . (40)
Noticing that
ϕt+1 − ϕt = αt dRM̂
dϕt . (41)
Combine Equations (37), (39), (40), (41), we can get
αtB1 − αtB3 √ B1 ≤ J(ϕt+1)− J(ϕt) + αtB4 +B5, (42)
where B1 = ∥∥∥ dJdϕt ∥∥∥2, B2 = 2τLJα0 +∑Tt=1 α2t (L2JLgJτ + LgJLR22 ) +∑Tt=1 αtLJmρτ , B3 =
C6ϵ g r +C7ϵ g f +C8ϵf , B4 = E 〈 dJ dϕt , dR M dϕt − dJdϕt 〉 B5 = LgJα 2 tLR 2
2 Notice B1, B4, B5 are functions of t.
Taking summarizing on both side, we have T∑ t=1 αtB1 − T∑ t=1 αtB3 √ B1 ≤ J(ϕT+1)− J(ϕ0) + T∑ t=1 αtB4 + T∑ t=1
B5, (43)∑T t=1 αtB1∑T t=1 αt − ∑T t=1 αtB3 √ B1∑T t=1 αt ≤ 2UJ +B2∑T t=1 αt . (44)
We have ∑T t=1 αt( √ B1 − B32 )
2∑T t=1 αt ≤ 2UJ +B2∑T t=1 αt + B23 4 . (45)
Therefore,
min 1≤k≤T
( √ B1 −
B3 2 )2 ≤ 2UJ +B2∑T
t=1 αt +
B23 4 . (46)
According to the lemma 10 We can prove that
min 1≤k≤T
B1 ≤ 4UJ + 2B2∑T
t=1 αt +B23 . (47)

heading: C.3 PROOF OF THE USEFUL LEMMAS
text: Proof of Lemma 1. According to the transition model definition and its Lipschitz continuous assumption, On the one hand, since the state are bound in the range with diameter D, we can see that ∥s̃t − st∥ ≤ D On the other hands,
∥s̃t − st∥ ≤ ϵf + Lf ∥s̃t−1 − st−1∥ ≤ ( 1− Lt+1f 1− Lf ) ϵf . (48)
Proof of the Lemma2.∥∥∥∥∥E ( d̂s̃t ds̃t−1 · · · · · d̂s̃1 ds̃0 ) E ( dst dst−1 · · · · · ds1 ds0 )∥∥∥∥∥ (49) =
∥∥∥∥∥E t∑ i=1 ( Πi−1j=0 dsj dsj−1 )( d̂s̃i ds̃i−1 − dsi dsi−1 )( Πtj=i+1 d̂s̃j ds̃j−1 )∥∥∥∥∥ (50) ≤E
t∑ i=1 (∥∥∥∥Πi−1j=0 dsjdsj−1 ∥∥∥∥) (∥∥∥∥∥ d̂s̃ids̃i−1 − d̂sidsi−1 ∥∥∥∥∥+ ∥∥∥∥∥ d̂sidsi−1 − dsidsi−1 ∥∥∥∥∥ )(∥∥∥∥∥Πtj=i+1 d̂sjdsj−1 ∥∥∥∥∥ ) (51)
≤ t∑ i=1 L21(ϵ g f + C0ϵfLf ) (52) ≤tL21ϵ g f + tC0LfL 2 1ϵf . (53)
Proof of the Lemma 3. Let s = ∑∞ t=0 tγ tC, then we can get that γs = ∑∞ t=0 tγ
t+1C =∑∞ t=1 (t− 1)γtC. So, s − γs = (1 − γ)s = ∑∞ t=1 γ
tC = γ1−γC. At last we can get s = γ(1−γ)2C.
Proof of the Lemma 4.∥∥∥∥∥E ( dRM̂ ds − dR M ds )∥∥∥∥∥ (54) =
∥∥∥∥∥E ∞∑ t=0 γt d̂r̃t ds̃t d̂s̃t s̃0 − E ∞∑ t=0 γt drt dst dst s0 ∥∥∥∥∥ (55) ≤E
∞∑ t=0 γt (∥∥∥∥∥ d̂s̃tds̃0 δ ( drt dst )∥∥∥∥∥+ ∥∥∥∥drtdst δ ( dst ds0 )∥∥∥∥ )
(56)
≤ ∞∑ t=0 γt ( L1(ϵ g r + LrC0ϵf ) + Lr ( C1ϵ g f + C2ϵf )) (57)
= 1
1− γ (L1(ϵ
g r + LrC0ϵf )) + Lr ∞∑ t=0 γt ( C1ϵ g f + C2ϵf ) (58)
= 1
1− γ (L1(ϵ
g r + LrC0ϵf )) +
γ
(1− γ)2 Lrϵ
g fL 2 1 +
γ
(1− γ)2 C0LfLrϵf (59)
=C3ϵ g r + C4ϵ g f + C5ϵf . (60)

heading: Proof of the Lemma 5.
text: T∑ t=1 αt ⟨g(ϕt), G(ϕt, ξt)− g(ϕt)⟩ (61)
= T∑ t=1 αt (⟨g(ϕt), G(ϕt, ξt)−G(ϕt, ξt+τ )⟩+ ⟨g(ϕt), G(ϕt, ξt+τ )− g(ϕt)⟩) . (62)
We first bound the first term in the right hand side.
∥ T∑ t=1 αt (⟨g(ϕt), G(ϕt, ξt)−G(ϕt, ξt+τ )⟩)∥ (63)
≤∥ τ∑ t=1 αt ⟨g(ϕt), G(ϕt, ξt)⟩∥+ ∥ T∑ t=T−τ+1 αt ⟨g(ϕt), G(ϕt, ξt+τ )⟩∥ (64)
+ ∥ T∑
t=τ+1
αt ⟨g(ϕt), G(ϕt, ξt)−G(ϕt+τ , ξt)⟩∥ (65)
≤2τα0L21 + ∥ T∑
t=τ+1
αt ⟨g(ϕt), G(ϕt, ξt)−G(ϕt+τ , ξt)⟩∥ (66)
≤ 2τα0L21 + T∑
t=τ+1 αt∥ t+τ−1∑ k=t ⟨g(ϕt), G(ϕk, ξt)−G(ϕk+1, ξt)⟩∥ (67)
≤ 2τα0L21 + T∑
t=τ+1 αtL2∥ t+τ−1∑ k=t ⟨g(ϕt), ϕk − ϕk+1⟩∥ (68)
≤ 2τα0L21 + T∑
t=τ+1
α2tL 2 1L2τ. (69)
(70)
Now, we will bound the second term in equation in expectation.
∥E ( T∑ t=1 αt (⟨g(ϕt), G(ϕt, ξt+τ )− g(ϕt)⟩) ) ∥ (71)
= T∑ t=1 αt (⟨g(ϕt),E (G(ϕt, ξt+τ )− g(ϕt)|Ft)⟩) (72)
≤ T∑ t=1 αtL1∥ ∫ G(ϕt, ξ) (P (ξt+τ ∈ ξ|Ft)− µ(ξ)) dξ∥ (73)
≤ T∑ t=1 αtL1dTV (P (ξt+τ ∈ ·|Ft), µ(·)) (74)
≤ T∑ t=1 αtL1mρ τ . (75)

heading: Proof of the Lemma 6.
text: J = E ∞∑ t=1 γtrt ≤ rmax 1− γ = UJ . (76)
Proof of the Lemma 7. The first inequality comes from Lemma B.4 in Wu et al. (2020). The second inequality comes from Lemma 3.2 in Zhang et al. (2020a).
Proof of the Lemma 8. ∥∥∥∥E(dRMds )∥∥∥∥ (77)
= ∥∥∥∥∥E ∞∑ t=0 γt drt dst dst s0 ∥∥∥∥∥ (78) ≤E
∞∑ t=0 γt ∥∥∥∥drtdst dsts0 ∥∥∥∥ (79) ≤
∞∑ t=0 γtL1Lr (80)
= 1
1− γ L1Lr. (81)
∥∥∥∥EdRMdϕ ∥∥∥∥ (82)
= ∥∥∥∥∥E ( ∞∑ t=0 γt ( drt dϕ ) + γ ∞∑ t=0 γt ( dR dst+1 dst+1 dat dat dϕ ))∥∥∥∥∥ (83) ≤ 1 1− γ LrLπ + ∥∥∥∥∥Eγ ∞∑ t=0 γt ( dR dst+1 dst+1 dat dat dϕ
)∥∥∥∥∥ (84) ≤ 1 1− γ LrLπ + γ (1− γ)2 L1LπLfLr (85)
=LR. (86)
Proof of the Lemma 9. The conclusion comes from the Equation 4.4.9 in Nesterov et al. (2018).
Proof of the Lemma 10. Since √ x and √ y + √ z are both positive, then we have x ≤ (√y + √ z)2, i.e., x ≤ y + z + 2√yz, note that 2√yz ≤ y + z, then we have x ≤ 2y + 2z.

Review summaries:
Many recent model-based reinforcement learning algorithms exploit the differentiability of the learned model by computing gradients through it. Nonetheless, the standard approach for model learning is to minimize a loss on the prediction of the next environment state which, when the underlying dynamics is complex, does not guarantee that also the resulting gradient information will be accurate. The paper proposes a method for model learning which explicitly tries to estimate the Jacobians of the underlying dynamics, by leveraging a loss function on the directional derivative and a nearest neighbour approximation. Additionally, the paper proposes to use a model learned with a combination of this new loss and a prediction loss for backpropagation only, while training a traditional model for the generation of trajectories.

This paper seeks to improve performance of model-based policy optimization methods by taking advantage of the world-model's differentiability. It starts by showing theoretically that error in estimating the model's gradient contributes to bias in the learned policy. The paper then proposes directly optimizing the jacobian of the world model using a sample-based strategy in order to directly control this bias term.

This paper considers learning a model for reinforcement learning and proposes to also learn the gradients of the model in addition to just the predictions. They approximate the model gradient using the nearest samples from the replay buffer in (12) to learn the model gradient (13), which they use to replace the model gradients in the value expansion in (14). They evaluate on MuJoCo environments (Fig 1) and show some key ablations (Fig 3) on the number of samples used for the gradient estimates, and the weight term in (12), and combined/separate models.