A CCELERATED P OLICY L EARNING WITH P ARALLEL D IFFERENTIABLE S IMULATION
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 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 efﬁciency 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.
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 timeconsuming, 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.
As recently Freeman et al. (2021) show, 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 task horizons (Hu et al., 2019; 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 complex 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 resulted by complex dynamics and occurrences of policy failures (Figure 2).
In addition, we use a truncated learning window to shorten the backpropagation path to address problems with vanishing/exploding gradients, reduce memory requirements, and improve learning efﬁciency.
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), 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 ﬁrst to provide a fair and comprehensive comparison between gradient-based and RL-based policy learning 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 task 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 (some of tasks are shown in 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., 2019; 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;
Xu et al., 2019; Lee et al., 2020; Andrychowicz et al., 2020; Chen et al., 2021a), 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.
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., 2019; 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 ﬁrst-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 sampleefﬁ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 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 ﬂowing 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 efﬁciently available in reverse-mode differentiable simulators.
In contrast, our method relies only on ﬁrst-order gradients.
Therefore, it can naturally leverage 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.
To model the dynamics function F , our physics simulator solves the forward dynamics equations
[Equation]
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
[Figure]
from Xu et al.
(2021) to provide better smoothness of the non-interpenetration contact dynamics:
[Equation]
where f c , f t 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, 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:
[Equation]
where k limit is the joint limit stiffness, and [ q lower , q upper ] is the bound for the joint angle q .
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.
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 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, inspired by previous work (Parmas et al., 2018), we plot the optimization landscape in Figure 2 (Left) for a humanoid locomotion problem with a 1000 -step task horizon.
Speciﬁcally, we take a trained policy, perturb the value of a single parameter θ k in the neural network, and 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 .
We provide more landscape analysis in Appendix A.5.
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 work successfully.
3. S HORT -H ORIZON A CTOR -C RITIC (SHAC)
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
Learning episode of horizon length 𝐻 𝜋 !
ℱ 𝑎 !
𝑆 " 𝑆 # 𝑎 " 𝑎 # 𝑆 $ 𝑎 $ 𝑆 % 𝑎 % …… 𝑎 "#$ 𝑆 !"
# 𝑎 "#% 𝑆 !"
$ 𝑎 "#& 𝑆 ’ 𝑆 !"
% 𝑆 !
𝑆 !"& 𝑎 "#’ ……
𝑉(𝑠 ! )
Learning episode of short horizon length ℎ 𝜋 !
ℱ 𝑎 !
𝑆 " 𝑎 " 𝑆 !
𝑎 (#’ 𝑆 (#’
Envs …… ……
𝑆 $’"& 𝑎 $’"& 𝑉(𝑠 "! )
Learning episode of short horizon length ℎ
Reset 𝑆 ) 𝑎 ) 𝑆 )* 𝑆 )* 𝑎 )* several sub-windows of smaller horizons across learning episodes (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 task horizon H .
At each step, an action vector a t is computed by a feedback policy π θ ( a t | s t ) .
While our method does not constrain the policy to be deterministic or stochastic, we use the stochastic policy in our experiments to facilitate extra exploration.
Speciﬁcally, the action is sampled by a t ∼ N ( µ θ ( s t ) , σ θ ( s t )) .
The transition function F is modeled by our differentiable simulation (Section 3.1).
A single-step reward r t =
R ( s t , a t ) is received at each step.
The goal of the problem is to ﬁnd the policy parameters θ that maximize the expected ﬁnite-horizon reward.
Our method works in an on-policy mode as follows.
In each learning episode, the algorithm samples
N trajectories { τ i } of short horizon h (cid:28) H in parallel from the simulation, which continue from the end states of the trajectories in the previous learning episode.
The following policy loss is then computed:
[Equation]
where s it and a it are the state and action 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 trajectory 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 for the stochastic policy.
For details of gradient computation, see Appendix A.1.
Our algorithm then updates the policy using one step of Adam (Kingma & Ba, 2014).
The differentiable simulator plays a critical role here, as it allows us to fully utilize the underlying dynamics linking states and actions, as well as optimize the policy, producing better short-horizon reward inside the trajectory and a more promising terminal state for the sake of long-term performance.
We note that the gradients are cut off between learning episodes to prevent unstable gradients during long-horizon backpropagation.
After we update the policy π θ , we use the trajectories collected in the current learning episode to train the value function V φ .
The value function network is trained by the following MSE loss:
[Equation]
Initialize policy π θ , value function V φ , and target value function V φ (cid:48)
← V φ .
for learning episode ← 1 , 2 , ..., M do
Sample N short-horizon trajectories of length h by the parallel differentiable simulation from the end states of the previous trajectories.
Compute the policy loss L θ deﬁned in Eq.
5 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. 7.
Fit the value function V φ using the critic loss deﬁned in Eq. 6.
Update target end for 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]
where G kt = (cid:16) (cid:80) k − 1 l =
γ l r t + l (cid:17) + γ k V φ ( s t + k ) is the k -step return from time t .
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.
6 does not ﬂow through the states and actions in Eq.
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. 5) and to estimate state values (Eq. 7).
In addition, we apply observation 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 learning 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 episodes avoid numerical problems when backpropagating the gradient through deeply nested update chains.
Finally, the use of short-horizon episodes 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.2.
We also experiment our method with a ﬁxed hyperparameter setting and with a deterministic policy choice.
Due to the space limit, the details of those experiments are provided in Appendix A.4.
4. B ENCHMARK C ONTROL P ROBLEMS
For comprehensive evaluations, we select six broad control tasks, including ﬁve classical RL tasks across different complexity levels, as well as one high-dimensional control task with a large action
[Figure]
space.
All tasks have stochastic initial states to further improve the robustness of the learned policy.
We introduce four representative tasks in the main paper, and leave the others in the Appendix.
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 .
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.3.
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 use high-performance implementations from RL games (Makoviichuk & Makoviychuk, 2021).
To achieve state-of-the-art performance, we follow Makoviychuk et al.
(2021): all simulation, reward and observation data remain on the GPU and are shared as PyTorch tensors between the RL algorithm and simulator.
The 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-batches 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 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.4.
We observe that our method 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 accurately estimate the policy gradient near the optimum, especially in complex problems.
We also compare our method to another model-free method REDQ (Chen et al., 2021b) in Appendix A.4.
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., 2019; 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 not efﬁciently available 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 lead to a higher wall-clock 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 SEMBPO 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.
5 to account for the long-term performance of the policy after the short episode horizon.
In this experiment, we evaluate the importance of this term.
By removing the terminal critic from Eq. 5, 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]
B e s t
R e w a r d
T r a i n i n g T i m e ( m i n )
Study of short horizon length h .
The choice of horizon length h is important for the performance of our method.
h cannot be too small, as it will result in worse value estimation by td- λ (Eq. 7) 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 lessfrequent policy updates.
Empirically, we ﬁnd that a short horizon length h = 32 with N = 64 parallel trajectories works well for all tasks in our experiments.
We conduct a study of short horizon length on the Ant task to show the inﬂuence of this hyperparameter.
We run our algorithm with six short horizon lengths h = 4 , 8 , 16 , 32 , 64 , 128 .
We set the corresponding number of parallel trajectories N = 512 , 256 , 128 , 64 , 32 ,for the variant, such that each one generates the same amount of samples in single learning 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 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 the developed 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.
A CKNOWLEDGEMENTS
We thank the anonymous reviewers for their helpful comments in revising the paper.
The majority of this work was done during the internship of Jie Xu at NVIDIA, and we thank Tae-Yong Kim for his mentorship. This work is also partially supported by Defense Advanced Research Projects
Agency (FA8750-20-C-0075)
A A PPENDIX A. P OLICY L OSS G RADIENT C OMPUTATION
We minimize the following policy loss to improve the policy network π θ :
[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]
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 ∂π θ ( s ) ∂ s t .
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 learning episodes for CartPole Swing Up problem and for 2000 learning episodes for other ﬁve 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.2. 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 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 not efﬁciently available in a reverse-mode differentiable simulator (including ours).
For SE-MBPO, we adapt the implementation released by Qiao et al.
(2021) into our benchmark problem.
A.2. 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 for Figure 4 in Table 1 for better reproducibility.
For other problems ( i.e. ,
HalfCheetah and Hopper), refer to our released code for the adopted hyperparameters.
[Table]
[Table]
. H YPERPARAMETERS S ETTINGS
Since our benchmark problems span across a wide range of complexity and problem dimentionality, the optimal hyperparameters in expected will be different across problems.
We conduct an extensive hyperparameter search for all algorithms, especially for PPO and SAC, and report their best performing hyperparameter settings for fair comparison.
For PPO, 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 3.
[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 [Table] best hyperparameter settings in Figure 4, and report the hyperparameters for each problem in Table 4.
[Table]
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
[Table]
in general, we found the set of hyperparameters in Table 5 works reasonably well on all problems (leading to slightly slower convergence speed in CartPole Swing Up and Ant tasks due to large α for target value network).
More analysis is provided in Appendix A.4.3.
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 six benchmark tasks, including ﬁve 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.3. 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]
[Table]
[Table]
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.
A.3. H ALF C HEETAH
In this problem [Table] a two-legged cheetah robot is controlled to run forward as fast as possible.
We have 17 -dimensional observation space as shown in Table 7 and 6 -dimensional action to control the torque applied to each joint.
[Table]
The single-step reward is deﬁned as:
[Equation]
(14) where v x is the forward velocity.
The initial state is randomly sampled.
The task horizon is 1000 steps and there is no early termination in this environment.
A.3. H OPPER
In this problem, a jumping hopper is controlled to run forward as fast as possible.
We have 11dimensional observation space as shown in Table 8 and 3 -dimensional action to control the torque applied to each joint.
[Table]
The single-step reward is deﬁned as:
[Equation]
15) where R v = v x is the forward velocity, R height is designed to penalize the low height state, deﬁned by:
[Equation] [Equation]
and the R angle is designed to encourage the upper body of the hopper to be as upward as possible, deﬁned by
[Equation]
The initial state is randomly sampled.
The task horizon is 1000 steps and early termination is triggered when the height of the hopper is lower than − 0 .
45 m.
A.3. A NT
In this problem, a four-legged ant is controlled to run forward as fast as possible.
We have 37dimensional observation space as shown in Table 9 and 8 -dimensional action to control the torque applied to each joint.
The single-step reward is deﬁned as:
[Equation]
(19) 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.3. 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 10 and a 21 -dimensional action vector to control the torque applied to each joint. [Table] [Table]
The single-step reward is deﬁned as:
[Equation]
20) where R v = v x is the forward velocity, R up is the projection of torso on the upward direction encouraging the agent to be vertically stable, R heading is the projection of torso on the forward direction encouraging the agent to run straight forward, and R height is deﬁned by:
[Equation]
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.3. 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 11 and a 152 -dimensional action vector.
The single-step reward is deﬁned as:
[Equation]
23) 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
[Table]
forward, and R height is deﬁned by:
[Equation]
The initial state is randomly sampled.
The task horizon is 1000 steps and early termination is triggered when the height of pelvis is lower than 0 .
46 m.
A. M ORE R ESULTS A.4. F ULL L EARNING C URVE C OMPARISON
In Section 4.3, we provide the truncated version of the learning curves in Figure 4 for better visualization.
We provide the full version of the learning curves in Figure 7, where the curve for each algorithm continues to the end of the training of that algorithm.
In addition, we include the experiment results for the other two classic RL environments, HalfCheetah and Hopper .
From the ﬁgure, we can see that our method shows extreme improvements in sample efﬁciency over other methods and constantly achieves better policies given the same amount of training time.
A.4. W ALL -C LOCK T IME B REAKDOWN OF T RAINING
We provide the detailed wall-clock time performance breakdown of a single learning episode of our method in Table 12.
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.
A.4. 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]
[Table]
A.4. 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
[Figure]
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 works reasonably well with deterministic policy, and sometimes the deterministic policy even outperforms the stochastic policy ( e.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.4. C OMPARISON TO REDQ
In this section, we compare our method to another advanced model-free reinforcement learning algorithm, namely REDQ (Chen et al., 2021b).
REDQ is an off-policy method and is featured by its high sample efﬁciency.
We experiment the ofﬁcially released code of REDQ on four of our 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 offpolicy 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 highdimensional problems, Humanoid and Humanoid MTU problems.
A.4. C OMPARISON TO M ODEL -B ASED RL
In this section, we compare our method to a model-based reinforcement learning algorithm.
In a typical model-based reinforcement learning algorithm, the learned model can be used in two ways: (1) data augmentation, (2) policy gradient estimation with backpropagation through a differentiable model.
In the data augmentation category, the learned model is used to provide rollouts around the observed environment transitions, as introduced in the DYNA algorithm (Sutton, 1990).
The core idea being that the model can be used to provide additional training samples in the neighborhood of observed environment data.
The core idea in DYNA was extended by STEVE (Buckman et al., 2018) and
MBPO (Janner et al., 2019), which use interpolation between different horizon predictions compared to a single step roll-out in DYNA.
However, the performance of such approaches rapidly degrades with increasing model error.
Notably, as noted by Janner et al.
(2019), empirically the one-step rollout is a very strong baseline to beat in part because error in model can undermine the advantage from model-based data-augmentation.
In contrast the model can be used for policy-value gradient estimation, which originates from
Nguyen & Widrow (1990) and Jordan & Rumelhart (1992).
This idea has been extended multiple times such as in PILCO (Deisenroth & Rasmussen, 2011), SVG (Heess et al., 2015), and MAAC (Clavera et al., 2020).
Despite the core method being very similar to the original idea, the difference between recent methods such as SVG and MAAC emerges in the details of what data is used in estimation.
SVG uses only real samples, while MAAC uses both real and model roll-outs.
But most of these methods still rely on model error being low for policy and value computation.
While SHAC resembles these gradient-based methods, our model usage fundamentally differs from previous approaches.
SHAC is the ﬁrst method that combines actor-critic methods with analytical gradient computation from a differentiable simulator model.
The analytical gradient solves the model learning issues which plague methods both in model-based data-augmentation and direct gradient estimators.
At the same time, SHAC prevents exploding/vanishing gradients in BPTT through truncated roll-outs and a terminal critic function.
SHAC offers improvement over prior model-based methods in both improved efﬁciency as well as scale-up to higher dimensional problems, beyond the ones reported in these methods.
To further illustrate the advantage of our method over model-based RL, we conduct the experiments to compare to Model-Based Policy Optimization (MBPO)
(Janner et al., 2019).
We run the ofﬁcially released code of MBPO on four of our benchmark problems: CartPole Swing Up, Ant, Humanoid, and Humanoid MTU , and provide the comparison between our method and MBPO in Figure 10.
As shown in the ﬁgure, the MBPO with the default parameters does not work well on our benchmark problems except CartPole Swing Up .
We hypothesize that it is because of the hyperparameter sensitivity of model-based RL method.
Although a fair comparision can be further conducted by tuning the hyperparameters of MBPO, an implicit comparison can be made through the comparison with REDQ.
In REDQ paper (Chen et al., 2021b), it is reported that REDQ is up to 75% faster in wall-clock time than MBPO.
This indirectly indicates that our method has signiﬁcant advantages in the wall-clock time efﬁciency over MBPO.
A. M ORE O PTIMIZATION L ANDSCAPE A NALYSIS
In Figure 2, we show the original landscape of the problem and the surrogate landscape deﬁned by
SHAC.
However the smoothness of the surrogate landscape sometimes 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 landscape, we plot the comparison of the computed single-weight analytical gradient and the single-
[Figure]
weight 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.
To better visualize the landscape, we also compare the landscape surfaces in Figure 13.
Speciﬁcally, we experiment on the landscapes for the Humanoid and Ant problems.
For Humanoid , we select a policy during training, and for Ant we use the optimal policy from training for study.
We randomly select two directions (cid:126) θ 1 and (cid:126) θ 2 in the policy network parameter space, and evaluate the policy variations along these two directions with weights θ (cid:48)
=
θ + k 1 (cid:126)θ 1 + k 2 (cid:126)θ 2 , k 1 , k 2 ∈ [ − 1 , 1] .
The ﬁgure shows that the short-horizon episodes and the terminal value critic of SHAC produce a surrogate landscape which is an approximation of the original landscape but is much smoother.
[Figure] [Figure]