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: They implement a differentiable simulator based on PyTorch (actually it was a bit unclear to me whether they implemented it, or just used Isaac Gym; I have a question about this below, and hope to receive an answer from the authors), as well as a method to utilize the differentiable simulation to optimize the policy more efficiently. The method operates as follows: 1. It computes a batch of episodes each of length H using a stochastic policy. 2. It splits the episodes into subsequences of length h. 3. It adds a terminal value function to each sequence, as well as discount factors on the rewards (starting from no discounting at the beginning of each subsequence), and computes policy gradients via backpropagation for optimization. 4. The value function is learned by computing lambda return targets (note that lambda returns are not used in the policy optimization though).  There is also some discussion about exploding gradients that hinder optimization through long horizons, which is why they truncated their episodes into subsequences.  They perform experiments on cartpole swingup, ant, humanoid and humanoid MTU (a humanoid actuated by 152 muscle tendon units), and compare to PPO, SAC, Backpropagation through time and PODS. The proposed method (SHAC) outperformed the other methods both in terms of number of steps (by a large margin) as well as computation speed (though PPO was faster in terms of wall-clock time in the early stage of learning after which SHAC overtook it). Particularly on the humanoid MTU task, the final performance was much better than the other methods. In terms of wall-clock time, the required computation time was a few minutes for cartpole and ant, and around 1h for the humanoid tasks.  They also performed ablation studies looking into the necessity of the terminal value function (it was necessary) and also the length of the subsequences (intermediate length sequences were optimal).  Recently some similar simulators, e.g. Brax or Isaac Gym have been made, but these works have not yet shown impressive results by taking advantage of differentiating through the model.