Gradient Information Matters in Policy Optimization by Back-propagating through Model

1 INTRODUCTION . 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 . 2 PRELIMINARIES . 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 . 3 RELATED WORK . 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 .

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