Model-Augmented Actor-Critic: Backpropagating through Paths

1 INTRODUCTION . Model-based reinforcement learning ( RL ) offers the potential to be a general-purpose tool for learning complex policies while being sample efficient . When learning in real-world physical systems , data collection can be an arduous process . Contrary to model-free methods , model-based approaches are appealing due to their comparatively fast learning . By first learning the dynamics of the system in a supervised learning way , it can exploit off-policy data . Then , model-based methods use the model to derive controllers from it either parametric controllers ( Luo et al. , 2019 ; Buckman et al. , 2018 ; Janner et al. , 2019 ) or non-parametric controllers ( Nagabandi et al. , 2017 ; Chua et al. , 2018 ) . Current model-based methods learn with an order of magnitude less data than their model-free counterparts while achieving the same asymptotic convergence . Tools like ensembles , probabilistic models , planning over shorter horizons , and meta-learning have been used to achieved such performance ( Kurutach et al. , 2018 ; Chua et al. , 2018 ; Clavera et al. , 2018 ) . However , the model usage in all of these methods is the same : simple data augmentation . They use the learned model as a black-box simulator generating samples from it . In high-dimensional environments or environments that require longer planning , substantial sampling is needed to provide meaningful signal for the policy . Can we further exploit our learned models ? In this work , we propose to estimate the policy gradient by backpropagating its gradient through the model using the pathwise derivative estimator . Since the learned model is differentiable , one can link together the model , reward function , and policy to obtain an analytic expression for the gradient of the returns with respect to the policy . By computing the gradient in this manner , we obtain an expressive signal that allows rapid policy learning . We avoid the instabilities that often result from back-propagating through long horizons by using a terminal Q-function . This scheme fully exploits the learned model without harming the learning stability seen in previous approaches ( Kurutach et al. , 2018 ; Heess et al. , 2015 ) . The horizon at which we apply the terminal Q-function acts as a hyperparameter between model-free ( when fully relying on the Q-function ) and model-based ( when using a longer horizon ) of our algorithm . The main contribution of this work is a model-based method that significantly reduces the sample complexity compared to state-of-the-art model-based algorithms ( Janner et al. , 2019 ; Buckman et al. , 2018 ) . For instance , we achieve a 10k return in the half-cheetah environment in just 50 trajectories . We theoretically justify our optimization objective and derive the monotonic improvement of our learned policy in terms of the Q-function and the model error . Furtermore , we experimentally analyze the theoretical derivations . Finally , we pinpoint the importance of our objective by ablating all the components of our algorithm . The results are reported in four model-based benchmarking environments ( Wang et al. , 2019 ; Todorov et al. , 2012 ) . The low sample complexity and high performance of our method carry high promise towards learning directly on real robots . 2 RELATED WORK . Model-Based Reinforcement Learning . Learned dynamics models offer the possibility to reduce sample complexity while maintaining the asymptotic performance . For instance , the models can act as a learned simulator on which a model-free policy is trained on ( Kurutach et al. , 2018 ; Luo et al. , 2019 ; Janner et al. , 2019 ) . The model can also be used to improve the target value estimates ( Feinberg et al. , 2018 ) or to provide additional context to a policy ( Du & Narasimhan , 2019 ) . Contrary to these methods , our approach uses the model in a different way : we exploit the fact that the learned simulator is differentiable and optimize the policy with the analytical gradient . Long term predictions suffer from a compounding error effect in the model , resulting in unrealistic predictions . In such cases , the policy tends to overfit to the deficiencies of the model , which translates to poor performance in the real environment ; this problem is known as model-bias ( Deisenroth & Rasmussen , 2011 ) . The model-bias problem has motivated work that uses meta-learning ( Clavera et al. , 2018 ) , interpolation between different horizon predictions ( Buckman et al. , 2018 ; Janner et al. , 2019 ) , and interpolating between model and real data ( Kalweit & Boedecker , 2017 ) . To prevent model-bias , we exploit the model for a short horizon and use a terminal value function to model the rest of the trajectory . Finally , since our approach returns a stochastic policy , dynamics model , and value function could use model-predictive control ( MPC ) for better performance at test time , similar to ( Lowrey et al. , 2018 ; Hong et al. , 2019 ) . MPC methods ( Nagabandi et al. , 2017 ) have shown to be very effective when the uncertainty of the dynamics is modelled ( Chua et al. , 2018 ; Wang & Ba , 2019 ) . Differentable Planning . Previous work has used backpropagate through learned models to obtain the optimal sequences of actions . For instance , Levine & Abbeel ( 2014 ) learn linear local models and obtain the optimal sequences of actions , which is then distilled into a neural network policy . The planning can be incorporated into the neural network architecture ( Okada et al. , 2017 ; Tamar et al. , 2016 ; Srinivas et al. , 2018 ; Karkus et al. , 2019 ) or formulated as a differentiable function ( Pereira et al. , 2018 ; Amos et al. , 2018 ) . Planning sequences of actions , even when doing model-predictive control ( MPC ) , does not scale well to high-dimensional , complex domains Janner et al . ( 2019 ) . Our method , instead learns a neural network policy in an actor-critic fashion aided with a learned model . In our study , we evaluate the benefit of carrying out MPC on top of our learned policy at test time , Section 5.4 . The results suggest that the policy captures the optimal sequence of action , and re-planning does not result in significant benefits . Policy Gradient Estimation . The reinforcement learning objective involves computing the gradient of an expectation ( Schulman et al. , 2015a ) . By using Gaussian processes ( Deisenroth & Rasmussen , 2011 ) , it is possible to compute the expectation analytically . However , when learning expressive parametric non-linear dynamical models and policies , such closed form solutions do not exist . The gradient is then estimated using Monte-Carlo methods ( Mohamed et al. , 2019 ) . In the context of model-based RL , previous approaches mostly made use of the score-function , or REINFORCE estimator ( Peters & Schaal , 2006 ; Kurutach et al. , 2018 ) . However , this estimator has high variance and extensive sampling is needed , which hampers its applicability in high-dimensional environments . In this work , we make use of the pathwise derivative estimator ( Mohamed et al. , 2019 ) . Similar to our approach , Heess et al . ( 2015 ) uses this estimator in the context of model-based RL . However , they just make use of real-world trajectories that introduces the need of a likelihood ratio term for the model predictions , which in turn increases the variance of the gradient estimate . Instead , we entirely rely on the predictions of the model , removing the need of likelihood ratio terms . Actor-Critic Methods . Actor-critic methods alternate between policy evaluation , computing the value function for the policy ; and policy improvement using such value function ( Sutton & Barto , 1998 ; Barto et al. , 1983 ) . Actor-critic methods can be classified between on-policy and off-policy . On-policy methods tend to be more stable , but at the cost of sample efficiency ( Sutton , 1991 ; Mnih et al. , 2016 ) . On the other hand , off-policy methods offer better sample complexity ( Lillicrap et al. , 2015 ) . Recent work has significantly stabilized and improved the performance of off-policy methods using maximum-entropy objectives ( Haarnoja et al. , 2018a ) and multiple value functions ( Fujimoto et al. , 2018 ) . Our method combines the benefit of both . By using the learned model we can have a learning that resembles an on-policy method while still being off-policy . 3 BACKGROUND . In this section , we present the reinforcement learning problem , two different lines of algorithms that tackle it , and a summary on Monte-Carlo gradient estimators . 3.1 REINFORCEMENT LEARNING . A discrete-time finite Markov decision process ( MDP ) M is defined by the tuple ( S , A , f , r , γ , p0 , T ) . Here , S is the set of states , A the action space , st+1 ∼ f ( st , at ) the transition distribution , r : S ×A → R is a reward function , p0 : S → R+ represents the initial state distribution , γ the discount factor , and T is the horizon of the process . We define the return as the sum of rewards r ( st , at ) along a trajectory τ : = ( s0 , a0 , ... , sT−1 , aT−1 , sT ) . The goal of reinforcement learning is to find a policy πθ : S ×A → R+ that maximizes the expected return , i.e. , maxθ J ( θ ) = maxθ E [ ∑ t γ tr ( st , at ) ] . Actor-Critic . In actor-critic methods , we learn a function Q̂ ( critic ) that approximates the expected return conditioned on a state s and action a , E [ ∑ t γ tr ( st , at ) |s0 = s , a0 = a ] . Then , the learned Q-function is used to optimize a policy π ( actor ) . Usually , the Q-function is learned by iteratively minimizing the Bellman residual : JQ = E [ ( Q̂ ( st , at ) − ( r ( st , at ) + γQ̂ ( st+1 , at+1 ) ) ) 2 ] The above method is referred as one-step Q-learning , and while a naive implementation often results in unstable behaviour , recent methods have succeeded in stabilizing the Q-function training ( Fujimoto et al. , 2018 ) . The actor then can be trained to maximize the learned Q̂ function Jπ = E [ Q̂ ( s , π ( s ) ) ] . The benefit of this form of actor-critic method is that it can be applied in an off-policy fashion , sampling random mini-batches of transitions from an experience replay buffer ( Lin , 1992 ) . Model-Based RL . Model-based methods , contrary to model-free RL , learn the transition distribution from experience . Typically , this is carried out by learning a parametric function approximator f̂φ , known as a dynamics model . We define the state predicted by the dynamics model as ŝt+1 , i.e. , ŝt+1 ∼ f̂φ ( st , at ) . The models are trained via maximum likelihood : maxφ Jf ( φ ) = maxφ E [ log p ( ŝt+1|st , at ) ]

Summary: This paper presents MAAC, a model-based reinforcement learning algorithm which makes use of path-wise derivative to optimize the policy. The learned Q function is used to estimate the terminal rewards and reduce instability and is trained by a STEVE-like style and Clipped Double Q. The dynamics model is learned by the same method in PETS. The policy is trained by directly computing the gradient to improve model-boosted Q. The theory shows that if learned model and learned Q function have similar derivative as the real one, the improvement of policy can be lower bounded. Experiments show that MAAC achieves the state-of-the-art and can be further improved by adding MPC.