PaperlayerView exact revision on arXiv

Source-linked AI summary

Accelerated Policy Learning with Parallel Differentiable Simulation

Jie Xu, Viktor Makoviychuk, Yashraj Narang, Fabio Ramos, Wojciech Matusik, Animesh Garg, Miles Macklin

arXiv:2204.07137v1cs.LGcs.AIcs.GRcs.RO

TL;DR

Control-policy learning with reinforcement learning is expensive in data and time, while differentiable-simulation methods face local minima, unstable gradients, and contact-related non-smoothness. The paper introduces SHAC with a parallel differentiable simulator, a smooth critic, and truncated learning windows, and reports higher sample and wall-clock efficiency across tasks, including high-dimensional locomotion.

  • Problem

    Reinforcement learning requires large training datasets, while differentiable-simulation policy optimization struggles with local minima, vanishing or exploding gradients, and non-smooth contact-rich dynamics.

  • Method

    SHAC combines a smooth critic surrogate, truncated learning windows, and parallel GPU-based differentiable simulation to learn control policies from simulation gradients.

  • Results

    SHAC outperforms state-of-the-art RL and gradient-based methods in sample efficiency and wall-clock time across benchmark and high-dimensional control tasks.

  • Takeaways & Limitations

    The method scales to high-dimensional control, where PPO and SAC struggle and SHAC benefits from more accurate differentiable-simulation gradients.

  • Takeaways & Limitations

    Backward computation remains consistently around 2× the forward-pass time, leaving wall-clock efficiency dependent on faster differentiable-simulator gradient computation.

Abstract

from arXiv · show

Deep reinforcement learning can generate complex control policies, but requires large amounts of training data to work effectively. Recent work has attempted to address this issue by leveraging differentiable simulators. However, inherent problems such as local minima and exploding/vanishing numerical gradients prevent these methods from being generally applied to control tasks with complex contact-rich dynamics, such as humanoid locomotion in classical RL benchmarks. In this work we present a high-performance differentiable simulator and a new policy learning algorithm (SHAC) that can effectively leverage simulation gradients, even in the presence of non-smoothness. Our learning algorithm alleviates problems with local minima through a smooth critic function, avoids vanishing/exploding gradients through a truncated learning window, and allows many physical environments to be run in parallel. We evaluate our method on classical RL control tasks, and show substantial improvements in sample efficiency and wall-clock time over state-of-the-art RL and differentiable simulation-based algorithms. In addition, we demonstrate the scalability of our method by applying it to the challenging high-dimensional problem of muscle-actuated locomotion with a large action space, achieving a greater than 17x reduction in training time over the best-performing established RL algorithm.

1 INTRODUCTION

Differentiable simulation could accelerate control-policy learning, but local minima, unstable gradients, and discontinuities limit its use on long-horizon contact-rich tasks. SHAC addresses these challenges with a smooth critic, truncated backpropagation, and parallel differentiable simulation, outperforming established methods in sample efficiency and wall-clock time.

  • RL learns complex robot and animation control policies from rewards but requires substantial training data, making high-dimensional problems expensive and time-consuming.
  • Local minima, vanishing or exploding gradients, and discontinuous landscapes impede gradient-based policy optimization.
  • Prior differentiable-simulation work largely focused on short-horizon open-loop policies or relatively simple contact-free tasks.
  • SHAC uses a smooth critic surrogate for noisy rewards and a truncated learning window to reduce long backward paths, memory requirements, and gradient instability.
  • Parallel GPU-based differentiable simulation enables comparisons across standard RL benchmarks and a character-control task with over 150 actuated degrees of freedom.
  • SHAC outperforms state-of-the-art policy-learning methods in both sample efficiency and wall-clock time.

2 RELATED WORK

Differentiable simulation provides gradients for policy optimization, while deep RL treats dynamics as a black box and model-based methods learn approximate dynamics. Existing differentiable-policy methods remain challenged by non-smooth contact dynamics, noisy landscapes, and unstable gradients on long horizons.

  • Differentiable Simulation: Differentiable simulators compute gradients of simulation outputs with respect to actions and initial conditions, using autodifferentiation or analytic calculations.
  • Differentiable Simulation: Contact dynamics are non-smooth, motivating methods based on differentiating complementarity models or smoothing contact with penalty formulations.
  • Deep Reinforcement Learning: Model-free RL treats system dynamics as a black box, whereas model-based RL learns an approximate dynamics model and exploits it during policy learning.
  • Differentiable Simulation based Policy Learning: Backpropagation through time and related first-order methods become ineffective on long-horizon tasks because optimization landscapes are noisy and gradients can vanish or explode.

3 METHOD

The method combines differentiable simulation with short-horizon actor-critic learning to address non-convex, discontinuous landscapes and unstable long-horizon gradients. A GPU-parallel simulator supports gradient-based policy optimization across complex contact-rich dynamics.

  • GPU-based differentiable dynamics simulation: The simulator models forward dynamics, contact, joint limits, and torque- or muscle-based actuation while supporting differentiable state transitions.It treats the simulator as st+1 = F(st, at), with parallel computation of articulated dynamics and forces.
  • GPU-based differentiable dynamics simulation: Smoothed contact and penalty-based joint-limit models improve the differentiability of otherwise non-smooth physical dynamics.Contact friction is approximated with a linear step function, while contact damping and continuous joint-limit penalties provide smoother dynamics.
  • GPU-based differentiable dynamics simulation: Distributed GPU kernels parallelize simulation across environments, including dense matrix routines and contact and joint-force evaluation.This parallelization supports collecting many trajectories while retaining differentiable simulation for policy updates.
  • Challenges in differentiable policy learning: Long-horizon feedback control creates highly non-convex, discontinuous reward landscapes where gradient-based optimization can stall and gradients can become numerically unstable.For humanoid locomotion with a 1000-step horizon, the backpropagated gradient norm exceeds 10^6.
  • Short-Horizon Actor-Critic: SHAC jointly learns an actor and critic, divides trajectories into short windows, and uses differentiable simulation within each window plus terminal value estimates across windows.The method samples short-horizon trajectories in parallel and uses multi-step rewards with a learned terminal value to train the policy.

4 EXPERIMENTS

The experiments compare SHAC with RL and differentiable-simulation baselines across classical and high-dimensional control tasks, measuring sample efficiency, wall-clock time, scalability, and ablations. SHAC improves sample efficiency and scales to muscle-actuated humanoid control, while results also expose trade-offs involving simulator overhead and design choices.

  • Experimental setup: The evaluation covers six control tasks, including five classical RL tasks and one high-dimensional task with a large action space.Representative tasks include CartPole Swing Up, Ant, Humanoid, and 152-muscle-tendon-unit Humanoid MTU control.
  • Benchmark results: SHAC shows significant sample-efficiency improvements over PPO and SAC on three classical RL problems, especially as problem dimension increases.The differentiable simulator supplies analytical gradients, whereas PPO and SAC estimate policy gradients from Monte Carlo samples.
  • Benchmark results: Wall-clock performance is closer across methods than sample efficiency, with model-free baselines competitive on simpler environments and SHAC initially slower than PPO.The comparison accounts for the higher per-iteration cost and backward-pass overhead of differentiable simulation.
  • Comparisons with differentiable-simulation methods: BPTT succeeds on contact-free CartPole Swing Up but quickly falls into local minima on contact-rich tasks, while PODS converges to a local optimum on CartPole Swing Up.PODS is less competitive in sample and wall-clock efficiency because of extra line-search simulations and slow imitation learning.
  • Comparisons with differentiable-simulation methods: SE-MBPO requires 8 hours to reach a reasonable Ant policy, whereas SHAC takes less than 15 minutes to reach the same gait level in the reported Ant problem.The adapted SE-MBPO implementation did not successfully optimize the policy after considerable hyperparameter tuning.
  • High-dimensional scalability: On Humanoid MTU, SHAC reaches PPO's reward in around 35 minutes and 1.7M simulation steps, yielding over 17× and 30× wall-clock improvements over PPO and SAC.The task uses a 152-dimensional action space; after 1.5 hours, SHAC finds a policy with twice the reward of the best RL policy.
  • Ablations: Removing the terminal critic prevents the method from learning a reasonable policy because optimization then focuses on short-horizon reward without long-term behavior.The ablation is algorithmically equivalent to short-window BPTT with discounted reward calculation.
  • Ablations: A horizon length h = 32 with N = 64 parallel trajectories works well across the reported tasks, while horizon length affects value estimation, gradient use, and update frequency.The study evaluates h = 4, 8, 16, 32, 64, 128 with corresponding parallel trajectory counts chosen to equalize samples per learning episode.

5 CONCLUSION AND FUTURE WORK

The method improves policy-learning efficiency by combining smooth critics, truncated learning windows, and parallel differentiable simulation. Its fixed short horizon works well empirically, while adaptive horizon scheduling remains future work.

  • Too-small horizons worsen value estimation, whereas too-large horizons create an ill-posed optimization landscape and increase training time.
  • SHAC combines a smooth critic, truncated learning window, and parallel differentiable simulation to improve sample and wall-clock efficiency.The method is especially effective as problem complexity increases.
  • Future work could replace the fixed short horizon with an adaptive schedule that varies during learning.

A.1 POLICY LOSS GRADIENT COMPUTATION

The policy-loss gradient is computed by differentiating through the simulator across short trajectory windows, using reverse-mode computation and reparameterized stochastic actions.

  • Each trajectory is indexed by τ_i, with s_t and a_t denoting its state and action at time step t.
  • The differentiable simulator is treated as a layer in the computation graph for the policy loss, enabling reverse-mode gradient computation with PyTorch.
  • Adjoints are propagated backward from the terminal time t0 + h through the preceding steps t0 ≤ t < t0 + h.
  • The policy-loss gradient is obtained from the computed adjoints across the trajectory set.
  • For stochastic policies sampled from N(µθ(s_t), σθ(s_t)), reparameterization sampling enables gradients through the policy actions.

A.2 EXPERIMENT SETUP

Experiments use matched hardware and extensive hyperparameter searches to compare algorithms under consistent wall-clock conditions, with different episode budgets for CartPole Swing Up and the other tasks.

  • All algorithms run on the same TITAN X GPU and Intel Xeon(R) E5-2620 CPU for fair wall-clock comparisons.
  • The study conducts extensive hyperparameter searches and reports each algorithm's best settings on every problem.
  • SHAC runs for 500 learning episodes on CartPole Swing Up and 2000 episodes on each of the other five problems.

A.2.1 BASELINE ALGORITHM IMPLEMENTATIONS

The baseline implementations use GPU-oriented neural networks and parallel simulation, with task-specific modifications for BPTT, PODS, and SE-MBPO.

  • Policy and Value Network Structures: All algorithms use multilayer perceptrons with ELU hidden-layer activations and layer normalization, while network sizes vary by task.
  • Codebase and Implementations: PPO and SAC use high-performance RL Games implementations with simulation, rewards, and observations kept on the GPU as shared PyTorch tensors.
  • Codebase and Implementations: BPTT is modified to use 64-step windows for CartPole and 128-step windows for other tasks, with parallel simulation for concurrent trajectory sampling.
  • Codebase and Implementations: PODS uses a first-order implementation because its second-order version requires a full state-action Jacobian that reverse-mode simulation does not efficiently provide.
  • Codebase and Implementations: SE-MBPO is adapted from the implementation released by Qiao et al. (2021).

A.2.2 IMPLEMENTATION DETAILS OF OUR METHOD

The method uses multilayer perceptrons for policy and value networks, with ELU hidden-layer activations and layer normalization.

  • Policy and value functions use multilayer perceptrons with ELU activations in hidden layers and layer normalization.

A.2.3 HYPERPARAMETERS SETTINGS

The study searches hyperparameters separately across problems and reports the best settings for fair comparisons, while also identifying a broadly effective general configuration.

  • Extensive hyperparameter searches cover PPO, SAC, and the proposed method across problems to support fair comparisons.The proposed method searches short horizon length, parallel environment count, and policy and critic learning rates.
  • The reported results use each algorithm's best-performing hyperparameter settings for each problem.
  • A general hyperparameter setting works reasonably well across problems, with slightly slower convergence in CartPole Swing Up and Ant.The slower convergence is attributed to a large α for the target value network.

A.3 BENCHMARK CONTROL PROBLEMS

The evaluation spans classical control benchmarks and a high-dimensional muscle-actuated humanoid task, testing sample efficiency, gradient behavior, and wall-clock performance. The method uses truncated horizons and a terminal critic to obtain smoother, more stable optimization while scaling to 152-dimensional actions.

  • Benchmark control problems: The benchmark suite contains six tasks: five classical RL problems and one high-dimensional problem with differentiable reward functions.
  • Benchmark control problems: The humanoid muscle-actuation task has 53 observations and 152 muscle-tendon-unit actions, making it the most complex scalability test.
  • Results: SHAC shows extreme sample-efficiency improvements over other methods across the evaluated learning curves.
  • Wall-clock performance: Backward simulation takes roughly 2× the forward time, creating remaining wall-clock overhead in the differentiable simulator.
  • Wall-clock performance: 27 hours versus 13 minutes: REDQ takes longer than SHAC to reach 5000 rewards on Ant.
  • Optimization design: SHAC combines actor-critic learning with analytical simulator gradients and uses truncated rollouts plus a terminal critic to prevent exploding or vanishing BPTT gradients.
  • Optimization design: SHAC gradients are more stable than BPTT gradients during Humanoid training, where gradient explosion appears in BPTT.
  • Optimization design: Short-horizon episodes and a terminal value critic produce a smoother surrogate landscape than the original landscape.
Loading 2204.07137v1…