G RADIENT I NFORMATION M ATTERS IN P OLICY O PTIMIZATION BY B ACK - PROPAGATING THROUGH M ODEL
A BSTRACT
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 proposing 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.
I NTRODUCTION
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 using 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 as 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 & Ja´skowski, 2020; Amos et al., 2021), which is mainly focused in this paper.
To get the policy updating direction, these methods compute the policy gradient 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, the algorithm requires the accurate model gradient, but we only learn to decrease the prediction error which results in an objective mismatch.
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 which 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 learning 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:
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.
We propose the two-model-based learning method and the practical DDPPO algorithm which learns and uses two environment models (prediction model used for rollout and gradient model used to provide the gradient information) for the model-based policy optimization.
Empirically, we can achieve better sample efficiency in the MuJoCo continuous control tasks than state-of-the-art model-based and model-free methods .
P RELIMINARIES
Reinforcement Learning: We consider a discrete-time Markov decision process (MDP) M , de- fined by tuple ( S , A , f, r, γ, p 0 ) , where S and A are the state and action spaces, respectively.
Here, f : s t +
= f
( s t , a t , ϵ t ) is the transition distribution, r : s t
= r ( s t , a t ) is a reward function, p 0 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 = (cid:80) ∞
[Equation] {+=
+}
γ
[Equation]
r ( s t , a t ) ).
The goal of reinforcement learning is to find a policy π ϕ that maximizes the expected sum of discounted rewards, i.e.,
[Equation]
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 θ ( s t , a t ) .
The model is trained by supervised learning, typically, via maximum likelihood:
[Equation]
where ( s t , a t , r t , s t +1 ) are data collected from interaction with the environment.
Policy Optimization
According to the Policy Gradient Calculated Using Model:
After we learn 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 [ (cid:80) ∞ t =0 γ t ˜ r 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 re-parameterization 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 dr t da t , dr k ds k , ds k ds t +1 , ds t +1 da t need to be calculated using the learned model.
[Equation]
The expectation is taken over randomness from random policy and transition.
To optimize the policy, we can update the policy parameters in the gradient direction.
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.
[Equation]
Notations :
We use fraction notation d · d · to represent the gradient calculated by taking gradient on the true model and use the ”hat” version (cid:98) d · d · to represent the gradient calculated on the learned model.
We use the ”tilde” notation ˜ · to represent the state, action or reward generated using the learned model and the without-tilde notation to represent those variables generated using the true model.
For example, the notation ds t +1 ds t denotes 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 = s t .
We use the notation ϵ f to represent the transition model prediction error which means, for ∀ s, a , we have E (cid:13)(cid:13)(cid:13) (
˜ f
( s, a ) − f ( s, a )) (cid:13)(cid:13)(cid:13) ≤ ϵ f .
We use the notation ϵ gf to represent the gradient error of the transition model.
For ∀ s, a , we have
[Equation]
da
[Equation]
Similarly, we use the notation ϵ gr to represent the gradient error of the reward model.
R ELATED W ORK
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 trajectories and use both the generated 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 & Ja´skowski, 2020) computes the policy gradient using value gradient and proposes to reduce the value gradient through TD learning.
Different from the above methods, 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.
M ODEL -B ASED P OLICY O PTIMIZATION BY C ONSIDERING THE M ODEL G RADIENT E RROR
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.
4. C ONVERGENCE R ATE FOR M ODEL -B ASED P OLICY O PTIMIZATION
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 ergodicity property:
[Equation]
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 ( dr t da t , dr k ds k , ds k ds t +1 , ds t +1 da t ) 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
[Equation]
where C 6 , C 7 , C 8 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.
L J is the Lipschitz constant of the objective function for the parameter ϕ
and L gs is the smoothness constant of the objective function for the parameter ϕ .
Suppose the policy parameters are updated as shown in Equation 4 and the learning rate is set as α t .
We have
[Equation]
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 lim T →∞ (cid:80)
Tt =1 α 2 t (cid:80) Tt =
1 α t = 0 , the policy optimization process will converge to near its stationary point with rate O (cid:16) (cid:80)
Tt =1 α
[Equation]
(cid:80) Tt =
1 α t (cid:17) .
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.
4. T WO -M ODEL -B ASED L EARNING P OLICY O PTIMIZATION
The theorem 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 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 model to predict the next state and the reward and to use the gradient model 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.
4. D IRECTIONAL D ERIVATIVE P ROJECTION P OLICY O PTIMIZATION : A P RACTICAL I MPLEMENTATION
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 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.
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.
; Clavera et al.
Different from the prediction model, the gradient model ˜ M 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 projec- tion 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., ∇ x M ( 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:
[Equation]
[Figure]
[Figure]
where v i denotes directions x i − x
[Figure]
and x i are the nearest n points of x in the buffer.
Intuitively
[Figure]
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:
[Equation]
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:
[Equation]
where ˜ s t +1 ∼ ˜ M pψ p , ˜ s 0 random sampled from the environment buffer D env , a ∼ π ϕ , α is the temperature hyperparameter in SAC.
Note that the policy can be learned by 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 ,a t ˜ M p (˜ s t , a t ) by ∇ ˜ s t ,a t ˜ M g (˜ s t , a t ) to calculate the full gradient, 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).
E XPERIMENT R ESULTS
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.
[Figure]
a d i e n t [Figure] d
[Figure]
5. C OMPARISON WITH S TATE - OF - THE -A RTS
To answer the first question
[Figure]
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 & Ja ´ skowski, 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.
5. G RADIENT E RROR A NALYSIS
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.
5. A BLATION S TUDY
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.
[Figure] [Figure]
Now we answer the last question, i.e., why build two models.
To investigate the importance of the gradient loss and role of two-model-based, we remove 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.
C ONCLUSIONS AND F UTURE W ORK
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.
A R ELATED W ORK
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 compared to their model-free counterparts.
Model-free and Model-based Methods: 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 analytically 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) explore 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 exploration 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.
K´egl 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-weight the model prediction error for each sample.
Farahmand et al.
(2017) use the value function to re-weight 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-weighted 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. Recently,
Amos et al.
(2021) extend SVG further leveraging the soft update (Haarnoja et al., 2018) and modelbased 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 Q-function.
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 & Ja´skowski, 2020) uses the differentiability of the learned model to compute gradient targets in temporal difference learning.
Our method is different from these methods.
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.
B A LGORITHM DETAILS B. I MPLEMENTATION OF DDPPO B.1. D IRECTIONAL D ERIVATIVE
The directional derivative of a function f with respect to a vector v at a point (e.g., position) x may be denoted by ∇ v f ( x ) that is a limit: ∇ v f ( x ) =
lim h → 0 f ( x + h v ) − f ( x ) h .
If the function f is differentiable at x then the directional derivative exists along any vector v and we have: ∇ v f ( 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 ∇ v f ( 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.
B.1. B ACK - PROPAGATING THROUGH G RADIENT M ODEL
In our method, we use the prediction model to rollout, however use the gradient model 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. I MPLEMENTATION
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 D env .
Then the prediction model is trained using data from D env .
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.
to achieve better performance.
Those samples are appended to the replay buffer D model , together with samples from D env , used to update the Q -function.
To update the policy, firstly, we obtain
[Table]
H
[Table]
n
[Table]
w
[Table]
imaginary trajectories by unrolling the latest policy on the prediction model
˜ M pθ p for H steps from a randomly sampled state in D env .
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:
[Equation] [Equation]
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:
[Equation]
13) where ¯ H is the target entropy.
B. E XPERIMENT S ETTINGS B.2. E NVIRONMENT S ETTINGS
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 and 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.
B.2. H YPERPARAMETER S ETTINGS
[Figure]
Table 1 shows the hyperparameters used for DDPPO to achieve results shown in Figure 1.
G 1 is the number of updates for policy per environment step, G 2 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 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 epochto 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.
B.2. E XPERIMENT D ETAILS FOR B ASELINES
For MBPO, we directly use the reported number given by Janner et al.
(2019) 2 ; For SAC we use the codes and hyperparameters available 3 ; For MAGE, we use the codes available 4 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 andfor others which is consistent with ours.
For other baselines, we use the codes or results from Wang et al.
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.
B. M ORE RESULTS FOR ABLATIONS
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 Figureshow 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 [Figure] r
[Figure]
e R e t u r
[Figure]
[Figure]
https://github.com/JannerM/mbpohttps://github.com/pranz24/pytorch-soft-actor-critichttps://github.com/nnaisense/MAGE
[Figure]
https://github.com/WilsonWangTHU/mbbl
[Figure]
single model) is not even as good as using the model without consideration of the gradient loss for some environments, such as Hopper, HalfCheetah
[Figure]
and Humanoid.
[Figure]
u
[Figure]
C P ROOFS OF THE THEOREM
In this section
[Figure]
we present the proof of the main theorems in section
[Figure]
First of all, we summarize the necessary assumptions here.
These assumptions are commonly used for analyzing reinforcement learning algorithms.
Assumption 2 (Stationary) .
For ∀ t 1 < t 2 and s 0 t 1 , s 1 t 1 , denote τ = t 2 − t 1 there exist L 1 ( τ ) < ∞ , τ
= t 2 − t 1 , such that (cid:13)(cid:13)(cid:13)(cid:13) E (cid:18)
[Equation]
s [Equation] 1 (cid:19)(cid:13)(cid:13)(cid:13)(cid:13) ≤ L 1 ( τ ) .
And L 1 ( τ ) can be uniformly upper bounded by L 1 < ∞ . Here, || E [Equation]
|| is the 2-norm of the matrix of which the element ij is: E [Equation] (cid:19) , where s t ( i ) is the i-th element in the vector s t .
Assumption 3 (Bounded region) .
There exist a finite D < ∞ such that: (Wu et al., 2020; Wang et al., 2017)
[Equation]
Assumption 4 (Lipschitz) .
The transition model and the reward model are both Lipschitz continuous with Lipschitz constant L f and L r respectively (Clavera et al., 2019).
In addition, we need the derivative of the transition model and the reward model to be Lipschitz continuous with Lipschitz constant L f g and L r g .
Assumption 5 (Uniform ergodicity) .
For a fixed policy π θ ( ·| s 0 ) , denote µ θ as the stationary distribution induced by the policy and the transition probability measure P ( ·| s, a ) .
Then there exists m < ∞ , C 0 < ∞ , and ρ ∈ (0 , 1) such that:(Wu et al., 2020; Zhang et al., 2020a;
Qiu et al., 2021)
[Equation]
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 & Ja´skowski, 2020)
[Equation]
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 , r max ] (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 C
= min( Dϵ f , 1 − L t +1 f 1 − L f ) for ∀ t we have
[Equation]
Lemma 2.
Denote C 1 = tL 21 , C 2 = tC
L f g L 21 .
We have
[Equation]
Lemma 3.
If constent 0 < γ < 1 , for arbitrary constant C , we have:
[Equation]
Lemma 4.
Denote constants C
=
L 1 1 − γ and C
=
L r L 21 γ (1 − γ ) 2 , C 5 = C
L 1 L rg (1 − γ ) + C
L fg L r γ (1 − γ ) 2 , then we have:
[Equation]
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
[Equation]
Denote the upper bound for the function G as L 1 .
Denote the Lipschitz constant for the function G as L 2 .
[Equation]
For all 0 < T < ∞ , we can proof that:
[Equation]
Lemma 6.
Denote U J as the upper bound of the expected sum of discounted rewards J .
Then
U J ≤ r max 1 − γ
Lemma 7.
There exist two constants L J and L gJ such that for all ϕ 1 , ϕ 2 , we have
[Equation]
Lemma 8.
There exist constants L R such that for all ϕ 1 , ϕ 2 , we have
[Equation]
Lemma 9.
Suppose the function J is smooth with smoothness constant
[Equation]
Lemma 10.
If x, y, z > 0 and √ x ≤ √ y + √ z , then we have
[Equation]
C. P ROOF OF THEOREM
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).
Firstly, we define some notations, denote C
[Equation]
, now we can prove the theorem by using previous lemmas:
[Equation]
(cid:32) (cid:92) dR
d ˜ s t + (cid:92) d ˜ s t + d ˜ a t d ˜ a t dϕ − dR ds t + ds t + da t da t dϕ (cid:33)(cid:33)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) (31) ≤ 1 − γ
[Equation]
+ (cid:13)(cid:13) (cid:13) (cid:13) (cid:13) E (cid:32) γ ∞ (cid:88) t
= γ t (cid:32) δ (cid:18) dR ds (cid:19) (cid:92) d ˜ s t + d ˜ a t + dR ds t +1 δ (cid:18) ds da (cid:19)(cid:33)
L π (cid:33)(cid:13)(cid:13) (cid:13) (cid:13) (cid:13) (32) ≤
[Equation]
[Equation]
C 6 ϵ gr + C
{+ϵ g f + C 
+}
ϵ f .
The first equation is the definition.
The second inequality follows the assumption 2, the right term .
The third term follows the lemma
Remark:
This theorem is fit for the case that the Lipschitz constant L f < 1 .
For the dynamics that the gradient is always larger than 1, it is unrealistic and uncontrollable.
For the dynamics that the gradient is always less than 1, theorem 1 and the assumptions are enough.
For the last case, it needs a stronger assumption directly on ∥ ˜ s t − s t ∥ .
C. P ROOF OF THE THEOREM
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.
Wu et al.
Ghadimi & Lan (2013) to get more insight of the proof.
According to the lemma 9, we can have:
[Equation] [Equation]
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
[Equation]
We can further decompose the value improvement as
[Equation]
According to the Cauchy–Schwarz inequality we have
[Equation]
[Equation]
Rearrange the term we have
[Equation]
[Equation]
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
[Equation]
[Equation]
The second term is
[Equation]
dJ dϕ
[Equation]
(cid:13) (cid:13) (cid:13) (cid:13) .
Noticing that
[Equation]
Combine Equations (37), (39), (40), (41), we can get
[Equation]
[Equation]
where B 1 =
[Equation]
Notice B 1 , B 4 , B 5 are functions of t .
Taking summarizing on both side, we have
[Equation]
[Equation]
We have
[Equation]
Therefore,
[Equation]
According to the lemma 10 We can prove that
[Equation]
C. P ROOF OF THE USEFUL LEMMAS
Proof of Lemma 1.
According to the transition model definition and its Lipschitz continuous assumption, on the one hand, since the states are bounded in the range with diameter D [Equation]
we can see that ∥ ˜ s t − s t ∥ ≤ D . On the other hand,
[Equation]
Proof of the Lemma
[Equation]
(50) ≤ E t (cid:88) i =
(cid:18)(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
ds
[Equation]
ds
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:19) (cid:32)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) (cid:92) d ˜ s i d
[Equation]
− (cid:92) ds
[Equation]
[-t (cid:88) i = (cid:18)(cid:13)(cid:13)(cid:13)(cid:13) Π i − 1 j =
 ds j ds j − 1 (cid:13)(cid:13)(cid:13)(cid:13)(cid:19) (cid:32)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) (cid:92) d ˜ s i d ˜ s i − 1 − (cid:92)-] ds [Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) + (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) (cid:92) ds
i ds [Equation] ds
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:33) (cid:32)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) Π
tj =
i + (cid:92) d ˜ s j d ˜ s j − 1 (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:33) (51) ≤ t (cid:88) i =
[Equation]
(52) ≤
[Equation]
gf
[Equation]
21 ϵ
[Equation]
Proof of the Lemma 3.
Let s = (cid:80)
∞
t =0 tγ t C , then we can get that γs = (cid:80) ∞
{+t =
+}
tγ t +1 C
=
(cid:80) ∞ t =1 ( t − 1) γ t C . So, s − γs = (1 − γ ) s = (cid:80) ∞ t =1 γ t C = γ 1 − γ C.
At last we can get s = γ (1 − γ ) 2 C.
Proof of the Lemma 4.
[Equation]
γ
[Equation]
Proof of the Lemma 5.
[Equation]
We first bound the first term in the right hand side.
[Equation]
t ) , G ( ϕ t , ξ t ) − G ( ϕ t + τ , ξ t ) ⟩∥ (64) ≤ 2 τα 0 L 21 + ∥
T (cid:88) t = τ + α t ⟨ g ( ϕ t ) , G ( ϕ t , ξ t ) − G ( ϕ t + τ , ξ t ) ⟩∥ (65) ≤ 2 τα 0 L 21 +
T (cid:88) t = τ + α t ∥ t + τ − 1 (cid:88) k
{+= t ⟨ g ( ϕ t ) , G ( ϕ k , ξ t ) − G ( ϕ k +1 , ξ t ) ⟩∥ (66) ≤ 2 τα 
+}
L 2 1 +
T (cid:88) t = τ + α t L 2 ∥ t + τ − 1 (cid:88) k
{+= t ⟨ g ( ϕ t ) , ϕ k − ϕ k +1 ⟩∥ (67) ≤ 2 τα 
 L 21 + T (cid:88) t = τ + α 2 t L 21 L 2 τ.+}
Now, we will bound the second term in equation in expectation.
[Equation]
Proof of the Lemma 6.
[Equation]
Proof of the Lemma 7.
The first inequality comes from Lemma B.4 in Wu et al.
The second inequality comes from Lemma 3.2 in Zhang et al.
(2020a).
Proof of the Lemma 8.
[Equation]
≤
[-+ (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) E γ ∞ (cid:88) t = γ t (cid:18) dR ds t + ds t + da t da t dϕ (cid:19)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) (84) ≤ 1 − γ L r L π + γ (1 − γ ) 
 L 1 L π L f L r (85) = L R .-][Equation]
Proof of the Lemma 9.
The conclusion comes from the Equation 4.4.9 in Nesterov et al.
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