Accelerated Policy Learning with Parallel Differentiable Simulation

Deep reinforcement learning can generate complex control policies , but requires large amounts of training data to work effectively . Recent work has attempted to address this issue by leveraging differentiable simulators . However , inherent problems such as local minima and exploding/vanishing numerical gradients prevent these methods from being generally applied to control tasks with complex contact-rich dynamics , such as humanoid locomotion in classical RL benchmarks . In this work we present a high-performance differentiable simulator and a new policy learning algorithm ( SHAC ) that can effectively leverage simulation gradients , even in the presence of non-smoothness . Our learning algorithm alleviates problems with local minima through a smooth critic function , avoids vanishing/exploding gradients through a truncated learning window , and allows many physical environments to be run in parallel . We evaluate our method on classical RL control tasks , and show substantial improvements in sample efficiency and wall-clock time over state-of-the-art RL and differentiable simulation-based algorithms . In addition , we demonstrate the scalability of our method by applying it to the challenging high-dimensional problem of muscle-actuated locomotion with a large action space , achieving a greater than 17× reduction in training time over the best-performing established RL algorithm . More visual results are provided at : https : //sites.google.com/view/shac . 1 INTRODUCTION . Learning control policies is an important task in robotics and computer animation . Among various policy learning techniques , reinforcement learning ( RL ) has been a particularly successful tool to learn policies for systems ranging from robots ( e.g. , Cheetah , Shadow Hand ) ( Hwangbo et al. , 2019 ; Andrychowicz et al. , 2020 ) to complex animation characters ( e.g. , muscle-actuated humanoids ) ( Lee et al. , 2019 ) using only high-level reward definitions . Despite this success , RL requires large amounts of training data to approximate the policy gradient , making learning expensive and time-consuming , especially for high-dimensional problems ( Figure 1 , Right ) . The recent development of differentiable simulators opens up new possibilities for accelerating the learning and optimization of control policies . A differentiable simulator may provide accurate first-order gradients of the task performance reward with respect to the control inputs . Such additional information potentially allows the use of efficient gradient-based methods to optimize policies . However , despite the availability of differentiable simulators , it has not yet been convincingly demonstrated that they can effectively accelerate policy learning in complex high-dimensional and contact-rich tasks , such as some traditional RL benchmarks ( Freeman et al. , 2021 ) . There are several reasons for this : 1 . Local minima may cause gradient-based optimization methods to stall . 2 . Numerical gradients may vanish/explode along the backward path for long trajectories . 3 . Discontinuous optimization landscapes can occur during policy failures/early termination . Because of these challenges , previous work has been limited to the optimization of open-loop control policies with short task horizons ( Hu et al. , 2018 ; Huang et al. , 2021 ) , or the optimization of policies for relatively simple tasks ( e.g. , contact-free environments ) ( Mora et al. , 2021 ; Du et al. , 2021 ) . In this work , we explore the question : Can differentiable simulation accelerate policy learning in tasks with continuous closed-loop control and contact-rich dynamics ? Inspired by actor-critic RL algorithms ( Konda & Tsitsiklis , 2000 ) , we propose an approach to effectively leverage differentiable simulators for policy learning . We alleviate the problem of local minima by using a critic network that acts as a smooth surrogate to approximate the underlying noisy reward landscape ( Figure 2 ) . In addition , we use a truncated learning window to shorten the backpropagation path to address problems with vanishing/exploding gradients and reduce memory requirements . Finally , the adopted terminal critic function allows us to support early termination , improving the learning efficiency when a policy failure occurs . A further challenge with differentiable simulators is that the backward pass typically introduces some computational overhead compared to optimized forward-dynamics physics engines . To ensure meaningful comparisons , we must ensure that our learning method not only improves sampleefficiency , but also wall-clock time . GPU-based physics simulation has shown remarkable effectiveness for accelerating model-free RL algorithms ( Liang et al. , 2018 ; Allshire et al. , 2021 ) , given this , we develop a GPU-based differentiable simulator that can compute gradients of standard robotics models over many environments in parallel . Our PyTorch-based simulator allows us to connect high-quality simulation with existing algorithms and tools . To the best of our knowledge , this work is the first to provide a fair and comprehensive comparison between gradient-based and RL-based methods , where fairness is defined as ( a ) benchmarking on both RL-favored tasks and differentiable-simulation-favored tasks , ( b ) testing complex tasks ( i.e. , contact-rich tasks with long task horizons ) , ( c ) comparing to the state-of-the-art implementation of RL algorithms , and ( d ) comparing both sample efficiency and wall-clock time . We evaluate our method on standard RL benchmark tasks , as well as a high-dimensional character control task with over 150 actuated degrees of freedom ( Figure 1 ) . We refer to our method as Short-Horizon ActorCritic ( SHAC ) , and our experiments show that SHAC outperforms state-of-the-art policy learning methods in both sample-efficiency and wall-clock time . 2 RELATED WORK . Differentiable Simulation Physics-based simulation has been widely used in the robotics field ( Todorov et al. , 2012 ; Coumans & Bai , 2016 ) . More recently , there has been interest in the construction of differentiable simulators , which directly compute the gradients of simulation outputs with respect to actions and initial conditions . These simulators may be based on auto-differentiation frameworks ( Griewank & Walther , 2003 ; Heiden et al. , 2021 ; Freeman et al. , 2021 ) or analytic gradient calculation ( Carpentier & Mansard , 2018 ; Geilinger et al. , 2020 ; Werling et al. , 2021 ) . One challenge for differentiable simulation is the non-smoothness of contact dynamics , leading many works to focus on how to efficiently differentiate through linear complementarity ( LCP ) models of contact ( Degrave et al. , 2016 ; de Avila Belbute-Peres et al. , 2018 ; Werling et al. , 2021 ) or leverage a smooth penalty-based contact formulation ( Geilinger et al. , 2020 ; Xu et al. , 2021 ) . Deep Reinforcement Learning Deep reinforcement learning has become a prevalent tool for learning control policies for systems ranging from robots ( Hwangbo et al. , 2019 ; OpenAI et al. , 2019 ; Lee et al. , 2020 ; Andrychowicz et al. , 2020 ) , to complex animation characters ( Peng et al. , 2018 ; 2021 ; Liu & Hodgins , 2018 ; Lee et al. , 2019 ) . Model-free RL algorithms treat the underlying dynamics as a black box in the policy learning process . Among them , on-policy RL approaches ( Schulman et al. , 2015 ; 2017 ) improve the policy from the experience generated by the current policy , while off-policy methods ( Lillicrap et al. , 2016 ; Mnih et al. , 2016 ; Fujimoto et al. , 2018 ; Haarnoja et al. , 2018 ) leverage all the past experience as a learning resource to improve sample efficiency . On the other side , model-based RL methods ( Kurutach et al. , 2018 ; Janner et al. , 2019 ) have been proposed to learn an approximated dynamics model from little experience and then fully exploit the learned dynamics model during policy learning . The underlying idea of our method shares similarity with some prior model-based methods ( Hafner et al. , 2019 ; Clavera et al. , 2020 ) that both learn in an actor-critic mode and leverage the gradient of the simulator “ model ” . However , we show that our method can achieve considerably better wall-clock time efficiency and policy performance using parallel differentiable simulation in place of a learned model . A detailed note on comparison with model-based RL is presented in Appendix A.4.6 . Differentiable Simulation based Policy Learning The recent development of differentiable simulators enables the optimization of control policies via the provided gradient information . Backpropagation Through Time ( BPTT ) ( Mozer , 1995 ) has been widely used in previous work to showcase differentiable systems ( Hu et al. , 2018 ; 2020 ; Liang et al. , 2019 ; Huang et al. , 2021 ; Du et al. , 2021 ) . However , the noisy optimization landscape and exploding/vanishing gradients in long-horizon tasks make such straightforward first-order methods ineffective . A few methods have been proposed to resolve this issue . Qiao et al . ( 2021 ) present a sample enhancement method to increase RL sampleefficiency for the simple MuJoCo Ant environment . However , as the method follows a model-based learning framework , it is significantly slower than state-of-the-art on-policy methods such as PPO ( Schulman et al. , 2017 ) . Mora et al . ( 2021 ) propose interleaving a trajectory optimization stage and an imitation learning stage to detach the policy from the computation graph in order to alleviate the exploding gradient problem . They demonstrate their methods on simple control tasks ( e.g. , stopping a pendulum ) . However , gradients flowing back through long trajectories of states can still create challenging optimization landscapes for more complex tasks . Furthermore , both methods require the full simulation Jacobian , which is not commonly or efficiently available in reverse-mode differentiable simulators . In contrast , our method relies only on first-order gradients . Therefore , it can naturally leverage simulators and frameworks that can provide this information . 3 METHOD . 3.1 DIFFERENTIABLE DYNAMICS SIMULATION . Conceptually , we treat the simulator as an abstract function st+1 = F ( st , at ) that takes a state s from a time t→ t+ 1 , where a is a vector of actuation controls applied during that time-step ( may represent joint torques , or muscle contraction signals depending on the problem ) . Given a differentiable scalar loss function L , and its adjoint L∗ = ∂L∂st+1 , the simulator backward pass computes : ∂L ∂st = ( ∂L ∂st+1 ) ( ∂F ∂st ) , ∂L ∂at = ( ∂L ∂st+1 ) ( ∂F ∂at ) ( 1 ) Concatenating these steps allows us to propagate gradients through an entire trajectory . To model the dynamics function F , our physics simulator solves the forward dynamics equations Mq̈ = JT f ( q , q̇ ) + c ( q , q̇ ) + τ ( q , q̇ , a ) , ( 2 ) where q , q̇ , q̈ are joint coordinates , velocities and accelerations respectively , f represents external forces , c includes Coriolis forces , and τ represents joint-space actuation . The mass matrix M , and Jacobian J , are computed in parallel using one thread per-environment . We use the composite rigid body algorithm ( CRBA ) to compute articulation dynamics which allows us to cache the resulting matrix factorization at each step ( obtained using parallel Cholesky decomposition ) for re-use in the backwards pass . After determining joint accelerations q̈ we perform a semi-implicit Euler integration step to obtain the updated system state s = ( q , q̇ ) . For simple environments we actuate our agents using torque-based control , in which the policy outputs τ at each time-step . For the more complex case of muscle-actuation , each muscle consists of a list of attachment sites that may pass through multiple rigid bodies , and the policy outputs activation values for each muscle . A muscle activation signal generates purely contractive forces ( mimicking biological muscle ) with maximum force prescribed in advance ( Lee et al. , 2019 ) . Analytic articulated dynamics simulation can be non-smooth and even discontinuous when contact and joint limits are introduced , and special care must be taken to ensure smooth dynamics . To model contact , we use the frictional contact model from Geilinger et al . ( 2020 ) , which approximates Coulomb friction with a linear step function , and incorporate the contact damping force formulation from Xu et al . ( 2021 ) to provide better smoothness of the non-interpenetration contact dynamics : fc = ( −kn + kdḋ ) min ( d , 0 ) n , ft = −tmin ( kt‖vt‖ , µ‖fc‖ ) ( 3 ) where fc , ft are the contact normal force and contact friction force respectively , d and ḋ are the contact penetration depth and its time derivative ( negative for penetration ) , n is the contact normal direction , t is the contact tangential direction , vt is the relative speed at the contact point along the contact tangential direction , and kn , kd , kt , µ are contact stiffness , contact damping coefficient , friction stiffness and coefficient of friction respectively . To model joint limits , a continuous penalty-based force is applied : flimit = { klimit ( qlower − q ) , q < qlower klimit ( qupper − q ) , q > qupper ( 4 ) where klimit is the joint limit stiffness , qlower and qupper are the lower and upper bounds for the joint angles , and q represents the joint angle . We build our differentiable simulator on PyTorch ( Paszke et al. , 2019 ) and use a source-code transformation approach to generate forward and backward versions of our simulation kernels ( Griewank & Walther , 2003 ; Hu et al. , 2020 ) . We parallelize the simulator over environments using distributed GPU kernels for the dense matrix routines and evaluation of contact and joint forces .

Summary: This paper focuses on the setting of reinforcement learning with differentiable simulation, where the authors assume access to a simulator which can compute the derivatives of the next state and reward with respect to the current state and action exactly. Under this setup, the authors propose an actor-critic style reinforcement learning algorithm that leverages the simulator derivative information to achieve better training sample efficiency and wall-clock time. Specifically, the authors expand the value function of the policy in h steps, and combine the sum of the simulator gradient of these h steps and the gradient of the learned value function at the h + 1 step to obtain the overall gradient for training the policy. The value function is then trained in the same way as in other actor-critic style methods.  The authors evaluate the proposed algorithm empirically on a few domains and the experiment results suggested that the proposed method achieves superior sample efficiency and wall-clock time compared to prior methods.