A CCELERATED P OLICY L EARNING WITH P ARALLEL D IFFERENTIABLE S IMULATION
Anonymous authors
Paper under double-blind review
A BSTRACT
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 SHAC, a short-horizon actor-critic method that successfully leverages parallel differentiable simulation to accelerate policy learning.
Our method 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 efﬁciency and wall-clock time over state-of-the-art RL and differentiable simulationbased 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.
I NTRODUCTION
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 deﬁnitions.
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 ﬁrst-order gradients of the task performance reward with respect to the control inputs.
Such additional information potentially allows the use of efﬁcient 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.
There are several reasons for this:
Local minima may cause gradient-based optimization methods to stall.
Numerical gradients may vanish/explode along the backward path for long trajectories.
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 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 efﬁciency 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 sampleefﬁciency, 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), and we design our method to effectively leverage parallel differentiable simulation using GPUs.
To the best of our knowledge, this work is the ﬁrst to provide a fair and comprehensive comparison between gradient-based and RL-based methods, where fairness is deﬁned as (a) benchmarking on both RL-favored tasks and differentiable-simulation-favored tasks, (b) testing complex tasks ( i.e. , contact-rich tasks with long horizons), (c) comparing to the state-of-the-art implementation of RL algorithms, and (d) comparing both sample efﬁciency 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 Actor Critic (SHAC), and our experiments show that SHAC outperforms state-of-the-art policy learning methods in both sample-efﬁciency and wall-clock time.
R ELATED W ORK
Differentiable Simulation Physics-based simulation has been widely used in the robotics ﬁeld (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 efﬁciently 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 efﬁciency.
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 efﬁciency and policy performance using parallel differentiable simulation in place of a learned model.
Differentiable Simulation based Policy Learning
The recent development of differentiable simulators opens up new possibilities to optimize 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 longhorizon tasks make such straightforward ﬁrst-order methods ineffective.
A few works have been proposed recently to solve this issue.
Qiao et al.
(2021) present a sample enhancement method to increase RL sample-efﬁciency for the simple MuJoCo Ant environment.
However, as the method follows a model-based learning framework, it is signiﬁcantly slower than state-of-the-art on-policy methods such as PPO (Schulman et al., 2017).
Mora et al.
(2021) propose to interleave a trajectory optimization stage and an imitation learning stage to detach the policy from the computation graph so as to alleviate the exploding gradient problem.
They demonstrate their methods on simple control tasks ( e.g., stopping a pendulum).
However, gradients ﬂowing back through long trajectories of states can still create challenging optimization landscapes for more complex tasks.
Furthermore, both methods (Mora et al., 2021; Qiao et al., 2021) require the full simulation Jacobian, which is not commonly or efﬁciently available in reverse-mode differentiable simulators.
In contrast, our method relies only on ﬁrst-order gradients.
Therefore, it can naturally leverage improvements in simulators and frameworks that can provide this information.
M ETHOD 3. D IFFERENTIABLE D YNAMICS S IMULATION
Conceptually, we treat the simulator as an abstract function s t
+1 =
F ( s t , a t ) 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 ∂ s t +1 , the simulator backward pass computes:
[Equation]
Concatenating these steps allows us to propagate gradients through an entire trajectory.
Although our actor-critic method is designed to leverage efﬁcient parallel simulation, it is compatible with any differentiable simulator that can compute these gradients.
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).
To compute articulation dynamics, we use the composite rigid body algorithm (CRBA), which uses dense matrix routines to build system matrices (Featherstone, 2014).
We perform a dense Cholesky decomposition to obtain joint accelerations, before ﬁnal semiimplicit Euler integration steps.
The decomposition result may be cached so that it does not need to be repeated in the backward pass.
We parallelize the simulator over environments using distributed
GPU kernels for the dense matrix routines and evaluation of contact and joint forces.
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.
In addition, we incorporate the contact damping force formulation from Xu et al.
(2021) into our dynamics model to provide better smoothness of the noninterpenetration contact dynamics.
To model joint limits, a continuous penalty-based force is applied instead of enforcing limits as hard constraints.
We provide more simulation details in Appendix A.1.
3. O PTIMIZATION L ANDSCAPE A NALYSIS
Although smoothed physical models improve the local optimization landscape, the combination of forward dynamics and the neural network control policy renders each simulation step non-linear and [Figure] non-
convex. This problem is exacerbated when thousands of simulation steps are concatenated and the actions in each step are coupled by a feedback control policy.
The complexity of the resulting reward landscape leads simple gradient-based methods to easily become trapped in local optima.
Furthermore, to handle agent failure ( e.g., a humanoid falling down) and improve sample efﬁciency, early termination techniques are widely used in policy learning algorithms (Brockman et al., 2016).
Although these have proven effective for model-free algorithms, early termination introduces additional discontinuities to the optimization problem, which makes methods based on analytical gradients less successful.
To analyze this problem, we plot the optimization landscape in Figure 2 (Left) for a humanoid locomotion problem with a 1000 -step horizon.
Speciﬁcally, we take a trained policy, perturb the value of a single parameter θ k in the neural network, and re-evaluate performance for the policy variations.
As shown in the ﬁgure, with long task horizons and early termination, the landscape of the humanoid problem is highly non-convex and discontinuous.
In addition, the norm of the gradient ∂ L ∂θ computed from backpropagation is larger than 10 6 .
Thus, most previous works based on differentiable simulation focus on short-horizon tasks with contact-free dynamics and no early termination, where pure gradient-based optimization ( e.g. , BPTT) can optimize the policy successfully.
[Equation]
To resolve the aforementioned issues of gradient-based policy learning, we propose the ShortHorizon Actor-Critic method (SHAC).
Our method concurrently learns a policy network ( i.e. , actor)
π θ and a value network ( i.e. , critic)
V φ during task execution, and splits the entire task horizon into several sub-windows of smaller horizons (Figure 3).
A multi-step reward in the sub-window plus a terminal value estimation from the learned critic is used to improve the policy network.
The differentiable simulation is used to backpropagate the gradient through the states and actions inside the sub-windows to provide an accurate policy gradient.
The trajectory rollouts are then collected and used to learn the critic network in each learning episode.
Speciﬁcally, we model each of our control problems as a ﬁnite-horizon Markov decision process (MDP) with state space S , action space A , reward function R , transition function F , initial state distribution D s 0 , and horizon H .
Each trajectory starts with an initial state sampled from the distribution s 0 ∼ D s 0 .
At each step, an action vector a t is computed by a feedback policy π θ ( a t | s t ) .
The transition function F is modeled by our differentiable simulation as mentioned in Section 3.
and deﬁnes the next state given the current state and action, s t +1 ← F ( s t , a t ) .
A single-step reward r t =
R ( s t , a t ) is received at each step.
The goal of the problem is then to ﬁnd the policy parameters θ that maximize the expected ﬁnite-horizon reward:
[Equation]
Although our method does not constrain the policy to be deterministic or stochastic, we choose to use the stochastic policy in our experiments for the extra exploration that it provides.
Speciﬁcally, the action at step t is sampled by
[Equation]
Episode of horizon length 𝐻 𝜋 !
ℱ 𝑎 ! 𝑆 " 𝑆 # 𝑎 " 𝑎 # 𝑆 $ 𝑎 $ 𝑆 % 𝑎 % …… 𝑎 "#$ 𝑆 !" # 𝑎 "#% 𝑆 !"
$ 𝑎 "#& 𝑆 ’ 𝑆 !" % 𝑆 ! 𝑆 !"& 𝑎 "#’ …… 𝑉(𝑠 ! )
Episode of short horizon length ℎ 𝜋 !
ℱ 𝑎 ! 𝑆 " 𝑎 " 𝑆 ! 𝑎 (#’ 𝑆 (#’
Envs …… ……
𝑆 $’"& 𝑎 $’"& 𝑉(𝑠 "! )
Episode of short horizon length ℎ
Reset 𝑆 ) 𝑎 ) 𝑆 )* 𝑆 )* 𝑎 )*
In each learning episode, we sample N trajectories { τ i } of short horizon h (cid:28) H in parallel from the simulation.
The following policy loss is then computed:
[Equation]
(4) where s it and a it are the state and actions at step t of the i -th trajectory, and γ < 1 is a discount factor introduced to stabilize the training.
Special handling such as resetting the discount ratio is conducted when task termination happens during sampling.
To compute the gradient of the policy loss ∂ L θ ∂θ , we treat the simulator as a differentiable layer (with backward pass shown in Eq. 1) in the PyTorch computation graph and perform regular backpropagation. We apply reparameterization sampling method to compute the gradient of Eq. 3 for the stochastic policy.
For details of gradient computaion, please see Appendix A.2.
Our algorithm then updates the policy using one step of a gradient-based solver, e.g. , Adam (Kingma & Ba, 2014).
The differentiable simulator plays an important role here, as it allows us to fully utilize the underlying dynamics linking states and actions, and to optimize the policy, producing better short-horizon reward inside the trajectory, and also resulting in a more promising terminal state for the sake of long-term performance.
We note that the trajectories sampled in the next episode will continue from the end state of the previous episode.
However, the gradients are cut off between episodes to prevent unstable gradients during long-horizon backpropagation.
After we update the policy π θ , we use the trajectories collected in the current episode to train the value function V φ .
The value function network is trained by the following MSE loss:
[Equation]
(5) where ˜ V ( s ) is the estimated value of state s , and is computed from the sampled short-horizon trajectories through a td- λ formulation (Sutton et al., 1998), which computes the estimated value by exponentially averaging different k -step returns to balance the variance and bias of the estimation:
[Equation]
Initialize policy π θ , value function V φ , and target value function V φ (cid:48)
← V φ .
for episode ← 1 , 2 , ..., M do
Sample N trajectories of length h by the parallel differentiable simulation.
Compute the policy loss L θ deﬁned in Eq.
4 from the sampled trajectories and V φ (cid:48) .
Compute the analytical gradient ∂ L θ ∂θ and update the policy π θ one step with Adam.
Compute estimated values for all the states in sampled trajectories with Eq. 6.
Fit the value function V φ using the critic loss deﬁned in Eq.
Update target value function: V φ (cid:48) ← αV φ (cid:48)
+ (1 − α )
V φ .
end for
The estimated value ˜ V ( s ) is treated as constant during critic training, as in regular actor-critic RL methods.
In other words, the gradient of Eq.
5 does not ﬂow through the states and actions in Eq. 6.
We further utilize the target value function technique (Mnih et al., 2015) to stabilize the training by smoothly transitioning from the previous value function to the newly ﬁtted one, and use the target value function V φ (cid:48) to compute the policy loss (Eq. 4) and to estimate state values (Eq. 6).
In addition, we apply state normalization as is common in RL algorithms, which normalizes the state observation by a running mean and standard deviation calculated from the state observations in previous episodes.
The pseudo code of our method is provided in Algorithm 1.
Our actor-critic formulation has several advantages that enable it to leverage simulation gradients effectively and efﬁciently.
First, the terminal value function absorbs noisy landscape over long dynamics horizons and discontinuity introduced by early termination into a smooth function, as shown in Figure 2 (Right).
This smooth surrogate formulation helps reduce the number of local spikes and alleviates the problem of easily getting stuck in local optima.
Second, the short horizon trajectories avoid numerical problems with backpropagating the gradient through deeply nested update chains.
Finally, the use of short horizons allows us to update the actor more frequently, which, when combined with parallel differentiable simulation, results in a signiﬁcant speed up of training time.
E XPERIMENTS
We design experiments to investigate ﬁve questions: (1) How does our method compare to the stateof-the-art RL algorithms on classical RL control tasks, in terms of both sample efﬁciency and wallclock time efﬁciency?
(2) How does our method compare to the previous differentiable simulationbased policy learning methods?
(3) Does our method scale to high-dimensional problems?
(4) Is the terminal critic necessary?
(5) How important is the choice of short horizon length h for our method?
4. E XPERIMENT S ETUP
To ensure a fair comparison for wall-clock time performance, we run all algorithms on the same
GPU model (TITAN X) and CPU model (Intel Xeon(R) E5-2620).
Furthermore, we conduct hyperparameter searches for all algorithms and report the performance of the best hyperparameters for each problem.
In addition, we report the performance averaged from ﬁve individual runs for each algorithm on each problem.
The details of the experimental setup are provided in Appendix A.3.
4. B ENCHMARK C ONTROL P ROBLEMS
For comprehensive evaluations, we select four broad control tasks, including three classical RL tasks across different complexity levels, as well as one high-dimensional control task with a large action space.
All tasks have stochastic initial states to further improve the robustness of the learned policy.
Classical Tasks :
We select CartPole Swing Up , Ant and Humanoid as three representative RL tasks, as shown in Figure 1.
Their difﬁculty spans from the simplest contact-free dynamics ( CartPole
Swing Up ), to complex contact-rich dynamics ( Humanoid ).
For CartPole Swing Up , we use H = 240 as the task horizon, whereas the other tasks use horizons of H = 1000 .
Humanoid MTU :
To assess how our method scales to high-dimensional tasks, we examine the challenging problem of muscle-actuated humanoid control (Figure 1, Right).
We use the lower body of the humanoid model from Lee et al.
(2019), which contains 152 muscle-tendon units (MTUs).
Each MTU contributes one actuated degree of freedom that controls the contractile force applied to the attachment sites on the connected bodies.
The task horizon for this problem is H = 1000 .
[Figure]
To be compatible with differentiable simulation, the reward formulations of each problem are deﬁned as differentiable functions.
The details of each task are provided in Appendix A.4.
4. R ESULTS
Comparison to model-free RL.
We compare SHAC with Proximal Policy Optimization (PPO)
(Schulman et al., 2017) (on-policy) & Soft Actor-Critic (SAC) (Haarnoja et al., 2018) (offpolicy).
We select these two baseline algorithms since PPO is the state-of-the-art on-policy method with high wall-clock training time efﬁciency, and SAC is one of the most advanced off-policy methods with high sample efﬁciency.
We use high-performance implementations of both methods available in RL games (Makoviichuk & Makoviychuk, 2021).
To achieve state-of-the-art performance, we follow the approach proposed by Makoviychuk et al.
(2021), where all simulation, reward and observation data remain on the GPU and are shared as PyTorch tensors between the RL algorithm and the parallel simulator.
Both PPO and SAC implementations are parallelized and operate on vectorized states and actions.
With PPO we used short episode lengths, an adaptive learning rate, and large mini-batch sizes during training to achieve the best possible performance.
As shown in the ﬁrst row of Figure 4, our method shows signiﬁcant improvements in sample efﬁciency over PPO and SAC in three classical RL problems, especially when the dimension of the problem increases ( e.g. , Humanoid ).
The analytical gradients provided by the differentiable simulation allow us to efﬁciently acquire the expected policy gradient through a small number of samples.
In contrast, PPO and SAC have to collect many Monte-Carlo samples to estimate the policy gradient.
Model-free algorithms typically have a lower per-iteration cost than methods based on differentiable simulation; thus, it makes sense to also evaluate wall-clock time efﬁciency instead of sampleefﬁciency alone.
As shown in the second row of Figure 4, the wall-clock time performance of PPO,
SAC, and our method are much closer than the sample efﬁciency plot.
Interestingly, the training speed of our method is slower than PPO at the start of training.
We hypothesize that the target value network in our method is far from the real value network initially, requiring sufﬁcient episodes to warm up.
We also note that the backward time for our simulation is consistently around 2 × that of the forward pass.
This indicates that our method still has room to improve its overall wall-clock time efﬁciency through the development of more optimized differentiable simulators with fast backward gradient computation.
We provide a detailed timing breakdown of our method in Appendix A.5.
We also observe that our method consistently achieves better policies than RL methods in all problems.
We hypothesize that, while RL methods are effective at exploration far from the solution, they struggle to estimate the policy gradient with sufﬁcient accuracy near the optimum point, especially in complex problems.
Comparison with previous gradient-based methods.
We compare our approach to three gradientbased learning methods: (1) Backpropagation Through Time (BPTT), which has been widely used in the differentiable simulation literature (Hu et al., 2018; Du et al., 2021), (2) PODS (Mora et al., 2021), and (3) Sample Enhanced Model-based Policy Optimization (SE-MBPO)
(Qiao et al., 2021).
BPTT :
The original BPTT method backpropagates gradients over the entire trajectory, which results in exploding gradients as shown in Section 3.2.
We modify BPTT to work on a shorter window of the tasks ( 64 steps for CartPole and 128 steps for other tasks), and also leverage parallel differentiable simulation to sample multiple trajectories concurrently to improve its time efﬁciency.
As shown in Figure 4, BPTT successfully optimizes the policy for the contact-free CartPole Swing Up task, whereas it falls into local minima quickly in all other tasks involving contact.
For example, the policy that BPTT learns for Ant is a stationary position leaning forward, which is a local minimum.
PODS: We compare to the ﬁrst-order version of PODS, as the second-order version requires the full Jacobian of the state with respect to the whole action trajectory, which is typically unavailable in a reverse-mode differentiable simulator (including ours).
Since PODS relies on a trajectory optimization step to optimize an open-loop action sequence, it is not clear how to accommodate early termination where the trajectory length can vary during optimization.
Therefore, we test PODS performance only on the CartPole Swing Up problem.
As shown in Figure 4, PODS quickly converges to a local optimum and is unable to improve further.
This is because PODS is designed to be a method with high gradient exploitation but little exploration.
Speciﬁcally, the line search applied in the trajectory optimization stage helps it converge quickly, but also prevents it from exploring more surrounding space.
Furthermore, the extra simulation calls introduced by the line search and the slow imitation learning stage make it less competitive in either sample or wall-clock time efﬁciency.
SE-MBPO :
Qiao et al.
(2021) propose to improve a model-based RL method MBPO (Janner et al., 2019) by augmenting the rollout samples using data augmentation that relies on the Jacobian from the differentiable simulator.
Although SE-MBPO shows high sample efﬁciency, the underlying model-based RL algorithm and off-policy training style cause their method to have a higher wallclock time.
As a comparison, the ofﬁcially released code for SE-MBPO takes 8 hours to achieve a reasonable policy in the Ant problem used by Qiao et al.
(2021), whereas our algorithm takes less than 15 minutes to acquire a policy with the same gait level in our Ant problem.
Aiming for a more fair comparison, we adapt their implementation to work on our Ant problem in our simulator.
However, we found that it could not successfully optimize the policy even after considerable hyperparameter tuning.
Regardless, the difference in wall-clock time between two algorithms is obvious, and the training time of SE-MBPO is unlikely to be improved signiﬁcantly by integrating it into our simulation environment.
Furthermore, as suggested by Qiao et al.
(2021), SE-MBPO does not generalize well to other tasks, whereas our method can be successfully applied to various complexity levels of tasks.
Scalability to high-dimensional problems.
We test our algorithm and RL baselines on the Humanoid MTU example to compare their scalability to high-dimensional problems.
With the large 152 -dimensional action space, both PPO and SAC struggle to learn the policy as shown in Figure(Right).
Speciﬁcally, PPO and SAC learn signiﬁcantly worse policies after more than 10 hours of training and with hundreds of millions of samples.
This is because the amount of data required to accurately estimate the policy gradient signiﬁcantly increases as the state and action spaces become large.
In contrast, our method scales well due to direct access to the more accurate gradients provided by the differentiable simulation with the reparameterization techniques.
To achieve the same reward level as PPO, our approach only takes around 35 minutes of training and 1 .
7 M simulation steps.
This results in over 17 × and 30 × wall-clock time improvement over PPO and SAC, respectively, and 382 × and 170 × more sample efﬁciency.
Furthermore, after training for only 1 . 5 hours, our method is able to ﬁnd a policy that has twice the reward of the best-performing policy from the RL methods.
A learned running gait is visualized in Figure 5.
Such scalability to high-dimensional control problems opens up new possibilities for applying differentiable simulation in computer animation, where complex character models are widely used to provide more natural motion.
Ablation study on the terminal critic.
We introduce a terminal critic value in Eq.
4 to account for the long-term performance of the policy after the short horizon.
In this experiment, we evaluate the importance of this term.
By removing the terminal critic from Eq. 4, we get an algorithmic equivalent to BPTT with a short horizon window and discounted reward calculation.
We apply this no-critic variation on all four problems and plot the training curve in Figure 4, denoted by “No
Critic.”
Without a terminal critic function, the algorithm is not able to learn a reasonable policy, as it only optimizes a short-horizon reward of the policy regardless of its long-term behavior.
[Figure]
A small h results in worse value es- timation.
A too large h leads to an ill-posed optimization landscape and longer training time.
Ablation study on short horizon length h .
The choice of horizon length h is important for the performance of our method.
The horizon length cannot be too small, as it will result in worse value estimation by td- λ (Eq. 6) and underutilize the power of the differentiable simulator to predict the sensitivity of future performance to the policy weights.
On the other hand, a horizon length that is too long will lead to a noisy optimization landscape and less-frequent policy updates.
Empirically, we ﬁnd that a short horizon length h =with N = 64 parallel trajectories generally works well for all four tasks in our experiments.
We conduct a hyperparameter study of short horizon length on the Ant task to show the inﬂuence of this hyperparameter.
We run our algorithm with six variations of hyperparameters h
= 4 , 8 , 16 , 32 , 64 , 128 .
We set the corresponding number of parallel trajectories N = 512 , 256 , 128 , 64 , 32 , 16 for the variant, such that each one generates the same amount of samples in an episode.
We run each variant for the same number of episodes M = 2000 with 5 individual random seeds.
In Figure 6, we report the average reward of the best policies from 5 runs for each variant, as well as the total training time.
As expected, the best reward is achieved when h = 16 or 32 , and the training time scales linearly as h increases.
C ONCLUSION AND F UTURE W ORK
In this work, we propose an approach to effectively leverage differentiable simulation for policy learning.
At the core of our method is the use of a critic network that acts as a smooth surrogate to approximate the underlying noisy optimization landscape.
In addition, a truncated learning window is adopted to alleviate the problem of exploding/vanishing gradients during deep backward paths.
Equipped with parallel differentiable simulation, our method shows signiﬁcantly higher sample efﬁciency and wall-clock time efﬁciency over state-of-the-art RL and gradient-based methods, especially when the problem complexity increases.
As shown in our experiments, model-free methods demonstrate efﬁcient learning at the start of training, but SHAC is able to achieve superior performance after a sufﬁcient number of episodes.
A compelling future direction for research is how to combine model-free methods with our gradient-based method in order to leverage the strengths of both.
Furthermore, in our method, we use a ﬁxed and predetermined short horizon length h throughout the learning process; however, future work may focus on implementing an adaptive short horizon schedule that varies with the status of the optimization landscape.
R EPRODUCIBILITY
To aid reproducibility, we provide the details of the simulation implementation in the manuscript and appendix.
The pseudo code along with detailed mathematical equations of our algorithm are presented in Section 3.3.
We also report all the hyperparameters used in each experiment in Appendix
A.3.
Furthermore, we provide the details of the reward function deﬁnition of each problem including the coefﬁcient of each reward term in Appendix A.4.
We plan to open source our code including the parallel differentiable simulation and the policy learning implementation upon the conclusion of double-blind review process.
A A PPENDIX A. D IFFERENTIABLE S IMULATION D ETAILS
Our physics simulator solves the forward dynamics equations
[Equation]
where q , ˙q , ¨q are joint coordinates, velocities and accelerations respectively, f represents external forces, and c includes Coriolis forces.
We discretize Eq.
7 using a semi-implicit time-discretization and solve for the resulting change in joint velocities before updating coordinates.
Many methods exist for solving these dynamics equations, including linear time methods such as the articulated-body algorithm ABA (Featherstone, 2014).
We use the CRBA method which is simpler to implement but also allows us to cache the resulting matrix factorization (obtained using parallel Cholesky decomposition) for re-use in the backwards pass
To model friction contact we use the following formulation:
[Equation]
where f c , f t are the contact normal force and contact friction force respectively, d and ˙
d is the contact penetration depth and its time derivative (negative for penetration), n is the contact normal direction, t is the contact tangential direction, v t is the relative speed at the contact point along the contact tangential direction, and k n , k d , k t , µ are contact stiffness, contact damping coefﬁcient, friction stiffness and coefﬁcient of friction respectively.
To model joint limits, a continuous penalty-based force is applied instead of enforcing limits as hard constraints.
[Equation]
where k limit is the joint limit stiffness, q lower and q upper are the lower and upper bounds for the joint angles, and q is the joint angle.
A. P OLICY L OSS G RADIENT C OMPUTATION
We minimize the following policy loss to improve the policy network π θ : [Equation]
[Equation]
where N is the number of trajectories, h is the short horizon length, t 0 is the starting time of each trajectory, s it and a it are the state and actions at time step t of trajectory τ i .
To compute its gradient, we treat the differentiable simulator as a differentiable layer (with backward pass shown in Eq. 1) in the computation graph for policy loss L θ , and acquiring the gradients ∂
L θ ∂θ by PyTorch with its reverse-mode computation.
Mathematically, we have the following terminal adjoints for each trajectory τ
i
[Equation]
[Equation]
From the last time step t
+ h , we can compute the adjoints in the previous steps t 0 ≤ t < t
+ h in reverse order:
[Equation]
From all the computed adjoints, we can compute the policy loss gradient by
[Equation]
To be noted, for our stochastic policy a t ∼ N ( µ θ ( s t ) , σ θ ( s t )) , we use the reparamterization sampling method, which allows us to compute the gradients ∂π θ ( s ) ∂θ and
[Equation] s [Equation]
A. E XPERIMENT S ETUP
To ensure a fair comparison for wall-clock time performance, we run all algorithms on the same GPU model (TITAN X) and CPU model (Intel Xeon(R) E5-2620).
Furthermore, we conduct extensive hyperparameter searches for all algorithms and report the performances of the best hyperparameter settings on each problem.
For our method, we run for 500 episodes for CartPole Swing Up problem and for 2000 episodes for other three problems.
For baseline algorithms, we run each of them for sufﬁciently long time in order to acquire the policies with the highest rewards.
We run each algorithm with ﬁve individual random seeds, and report the average performance.
A.3. B ASELINE A LGORITHM I MPLEMENTATIONS
Policy and Value Network Structures
We use multilayer perceptron (MLP) as the policy and value network structures for all algorithms.
We use ELU as our activation function for the hidden layers and apply layer normalization to each hidden layer.
Since different problems have different state and action dimensionalities, we use different network sizes for different tasks.
Codebase and Implementations For PPO and SAC, we use high-performance implementations of both methods available in RL games (Makoviichuk & Makoviychuk, 2021).
To achieve state-ofthe-art performance we follow the approach proposed by Makoviychuk et al.
(2021) where all the simulation, reward and observation data stays on GPU and is shared as PyTorch tensors between the
RL algorithm and the parallel simulator.
Both PPO and SAC implementations are parallelized and operate on vectorized states and actions.
With PPO we used short episode lengths, an adaptive learning rate, and large mini-batch sizes during training to achieve the maximum possible performance.
For BPTT, the original BPTT which requires to backprapagate gradients over the entire trajectory and results in exploding gradients problem.
In order to make it work on our long-horizon tasks with complex dynamics, we modify BPTT to work on a shorter window of the tasks ( 64 steps for CartPole and 128 steps for other tasks).
We further improve its wall-clock time efﬁciency by leveraging parallel differentiable simulation to sample multiple trajectories concurrently.
For PODS, due to the lack of released code, we implement our own version of it.
Speciﬁcally, we implement the ﬁrst-order version of PODS proposed by Mora et al.
(2021) since the second-order version requires the full Jacobian of the state with respect to the whole action trajectory, which is typically unavailable ina revser-mode differentiable simulator (including ours).
For SE-MBPO, we adapt the implementation released by Qiao et al.
(2021) into our benchmark problem.
A.3. I MPLEMENTATION D ETAILS OF O UR M ETHOD
Policy and Value Network Structures
We use multilayer perceptron (MLP) as the policy for our method.
We use ELU as our activation function for the hidden layers and apply layer normalization to each hidden layer.
Since different problems have different state and action dimensionalities, we use different network sizes for different tasks.
Basically, we use larger policy and value network structures for problems with higher degrees of freedoms.
Empirically we ﬁnd the performance of our method is insensitive to the size of the network structures.
We report the network sizes that are used in our experiments in Table 1 for better reproducibility.
[Table]
A.3. H YPERPARAMETERS S ETTINGS
Since our four benchmark problems span across a wide range of complexity and problem dimentionality
[Table]
the optimal hyperparameters in expected will be different across problems.
We conduct an extensive hyperparameter search for all algorithms, especially for PPO and SAC
[Table]
and report their best performing hyperparameter settings for fair comparison.
For PPO
[Table]
we adopt the adaptive learning rate strategy implemented in RL games (Makoviichuk & Makoviychuk, 2021).
The hyperparameters of PPO and SAC we used in the experiments are reported in Table 2 and
[Table]
[Table]
[Table]
For our method, we also conduct hyperparameter searches for each problem, including the short horizon length h , the number of parallel environments N , and the policy and critic learning rates.
The learning rates follow a linear decay schedule over episodes.
We report performance from the best hyperparameter settings in Figure 4, and report the hyperparameters for each problem in Table 4.
Although the optimal hyperparameters are different for different problems due to the different complexity of the problem and we report the best setting for fair comparison against baseline algorithms, in general, we found the set of hyperparameters in Table 5 works reasonably well on all problems [Table] (leading to slightly slower convergence speed in CartPole Swing Up and Ant tasks due to large α for target value network) [Table]
[Table]
A. B ENCHMARK C ONTROL P ROBLEMS
We describe more details of our benchmark control problems in this section.
We evaluate our algorithm and baseline algorithms on four benchmark tasks, including three classical RL benchmark problems and one high-dimensional problem.
We deﬁne the reward functions to be differentiable functions in order to incorporate them into our differentiable simulation framework
. A.4. C ART P OLE S WING U P
CartPole Swing Up is one of the simplest classical RL control tasks.
In this problem, the control policy has to swing up the pole to the upward direction and keeps it in that direction as long as possible.
We have 5 -dimensional observation space as shown in Table 6 and 1 -dimensional action to control the torque applied to the prismatic joint of the cart base.
The single-step reward is deﬁned as:
[Equation]
The initial state is randomly sampled.
The task horizon is 240 steps and there is no early termination in this environment so that it can be used to test the algorithms which are not compatible with early termination strategy.
[Table]
In this problem, an four-legged ant is controlled to run forward as fast as possible.
We have 37dimensional observation space as shown in Table 7 and 8 -dimensional action to control the torque applied to each joint.
[Table]
The single-step reward is deﬁned as:
[Equation]
16) where R v = v x is the forward velocity, R up = projection(upward direction) encourages the agent to be vertically stable, R heading =
projection(forward direction) encourages the agent to run straight forward, and R height = h − 0 .
27 is the height reward.
The initial state is randomly sampled.
The task horizon is 1000 steps and early termination is triggered when the height of the ant is lower than 0 .
27 m.
A.4. H UMANOID
In this problem, a humanoid robot is controlled to run forward as fast as possible.
We have a 76dimensional observation space as shown in Table 8 and a 21 -dimensional action vector to control the torque applied to each joint.
The single-step reward is deﬁned as:
[Equation]
(17) where R v = v x is the forward velocity, R up =
projection(upward direction) encourages the agent to be vertically stable, R heading =
projection(forward direction) encourages the agent to run straight forward, and R height is deﬁned by:
[Equation]
[Table]
The initial state is randomly sampled.
The task horizon is 1000 steps and early termination is triggered when the height of torso is lower than 0 .
74 m.
A.4. H UMANOID MTU
This is the most complex problem designed to assess the scalability of the algorithms.
In this problem, a lower body of the humanoid model from Lee et al.
(2019) is actuated by 152 muscle-tendon units (MTUs).
Each MTU contributes one actation degree of freedom that controls the contractile force applied to the attachment sites on the connected bodies.
We have a 53 -dimensional observation space as shown in Table 9 and a 152 -dimensional action vector.
[Table]
The single-step reward is deﬁned as:
[Equation]
(20) where R v = v x is the forward velocity
[Table]
R up = projection(upward direction) encourages the agent to be vertically stable, R heading =
projection(forward direction) encourages the agent to run straight forward
[Table]
and R height is deﬁned by:
[Equation] [Equation]
The initial state is randomly sampled.
The task horizon is 1000 steps and early termination is triggered when the A. M ORE R ESULTS A.5. F ULL L EARNING C URVE C OMPARISON
In Section 4.
[Figure]
we provide the truncated version of the learning curves in Figure 4 for better visualization.
We provide the full version of the learning curves here in Figure
[Figure]
where the curve for each algorithm continues to the end of the training of that algorithm. In addition
[Figure]
we include one more classic RL environment, HalfCheetah .
From the ﬁgure
[Figure]
we can see that our method shows extremely improvements in sample efﬁciency over any other methods [Figure] and constantly achieve the better policies given the same amount of training time [Figure]
A.5. W ALL -C LOCK T IME B REAKDOWN OF T RAINING
We provide the detailed wall-clock time performance breakdown of a single episode of our method in Table 10.
Since we care about the wall-clock time performance of the algorithm, we use this table for a better analysis of which is the time bottleneck in our algorithm.
The performance is measured on a desktop with GPU model TITAN X and CPU model Intel Xeon(R) E5-2620 @ 2.10GHz.
As shown in the table, the forward simulation and backward simulation time scales up when the dimensionality of the problem increases, while the time spent in critic value function training is almost constant across problems.
As expected, the differentiable simulation brings extra overhead for its backward gradient computation.
Speciﬁcally in our differentiable simulation, the backward computation time is roughly 2 × of the forward time.
This indicates that our method still has room to improve its overall wall-clock time efﬁciency through the development of more optimized differentiable simulators with fast backward computation. [Table]
A.5. F IXED N ETWORK A RCHITECTURE AND H YPERPARAMETERS
In the previous results, we ﬁne tune the network architectures of policy and value function for each algorithm on each problem to get their maximal performance for a fair comparison.
In this section, we test the robustness of our method by using a ﬁxed network architecture and ﬁxed set of hyperparameters ( e.g. learning rates) to train on all problems.
Speciﬁcally, we use the same architecture and the hyperparameters used in the Humanoid problem, and plot the training curves in
Figure 8.
[Figure]
[Figure]
A.5. D ETERMINISTIC P OLICY
Our method does not constrain the policy to be stochastic or deterministic.
We choose to use the stochastic policy in the most experiments for the extra exploration that it provides.
In this section, we experiment our method with the deterministic policy choice.
Speciﬁcally, we change the policy of our method in each problem from stochastic policy to deterministic policy while keeping all other hyperparameters such as network dimensions and learning rates the same as shown in Table 4.
The training curves of the deterministic policy is plotted in Figure 9.
The results show that our method with deterministic policy works reasonably well on all problems
[Figure]
and sometimes the deterministic policy even outperforms the stochastic policy ( e
[Figure]
g. Humanoid ).
The small performance drop on the
Ant problem comes from one single random seed (out of ﬁve) which results in a sub-optimal policy
[Figure]
A.5. C OMPARISON TO REDQ
In this section, we compare our method to another advanced model-free reinforcement learning algorithm, namely REDQ (Chen et al., 2021).
REDQ is an off-policy method and is featured by its high sample efﬁciency.
We experiment ofﬁcially released code of REDQ on our four benchmark problems: CartPole Swing Up, Ant, Humanoid, and Humanoid MTU , and provide the comparison between our method and REDQ in Figure 10.
From the plots, we see that REDQ demonstrates its high sample efﬁciency in the CartPole Swing Up and Ant problems due to its off-policy updates and the randomized Q function ensembling strategy.
However, REDQ is extremely slow in wall-clock training time.
For example, in Ant problem, REDQ takes 27 hours to achieve 5000 rewards while
SHAC only takes 13 minutes to achieve the same reward.
Such huge difference in wall-clock time performance renders REDQ a less attractive method than SHAC in the settings where the policy is learned in a fast simulation.
Furthermore, REDQ fails to train on our high-dimensional problems,
Humanoid and Humanoid MTU problems.
A.5. C OMPARISON TO M ODEL -B ASED RL A. L ANDSCAPE A NALYSIS
In Figure 2, we show the original landscape of the problem and the surrogate landscape deﬁned by our SHAC.
However the smoothness of the surrogate landscape does not necessarily results in a smooth gradient when the evaluation trajectory is stochastic and the gradient is computed by averaging over policy stochasticity (Parmas et al., 2018).
To further analyze the gradient of the
[Figure]
landscape, we plot the comparison of the computed single-weight analytical gradient and the singleweight gradient computed by ﬁnite difference of the landscape (shown in Figure 2 Right) in Figure 11.
The plot demonstrates that our computed gradient is well-behaved and close enough to the ﬁnite difference gradient.
[Figure]
Furthermore, we compare the gradient of the BPTT method and the gradient of SHAC in Figure 12, where we picked six different policies during the training for the Humanoid problem, and plot the gradient value distribution of BPTT method and SHAC method for each of them.
From Figure 12, we can clearly see that the gradient explosion problem exists in the BPTT method while the gradient acquired by SHAC is much more stable throughout the training process.
[Figure]