Stabilizing Reinforcement Learning in Differentiable Multiphysics Simulation

TL;DR

Soft Analytic Policy Optimization (SAPO), a maximum entropy first-order model-based actor-critic RL algorithm, which uses first-order analytic gradients from differentiable simulation to train a stochastic actor to maximize expected return and entropy is introduced.

Abstract

Recent advances in GPU-based parallel simulation have enabled practitioners to collect large amounts of data and train complex control policies using deep reinforcement learning (RL), on commodity GPUs. However, such successes for RL in robotics have been limited to tasks sufficiently simulated by fast rigid-body dynamics. Simulation techniques for soft bodies are comparatively several orders of magnitude slower, thereby limiting the use of RL due to sample complexity requirements. To address this challenge, this paper presents both a novel RL algorithm and a simulation platform to enable scaling RL on tasks involving rigid bodies and deformables. We introduce Soft Analytic Policy Optimization (SAPO), a maximum entropy first-order model-based actor-critic RL algorithm, which uses first-order analytic gradients from differentiable simulation to train a stochastic actor to maximize expected return and entropy. Alongside our approach, we develop Rewarped, a parallel differentiable multiphysics simulation platform that supports simulating various materials beyond rigid bodies. We re-implement challenging manipulation and locomotion tasks in Rewarped, and show that SAPO outperforms baselines over a range of tasks that involve interaction between rigid bodies, articulations, and deformables. Additional details at https://rewarped.github.io/.

Figures (13)

• Figure 1: Visualizations of tasks implemented in Rewarped. These are manipulation and locomotion tasks involving rigid and soft bodies. AntRun and HandReorient are tasks with articulated rigid bodies, while RollingFlat, SoftJumper, HandFlip, and FluidMove are tasks with deformables.
• Figure 2: Rewarped tasks training curves. Episode return as a function of environment steps in Rewarped AntRun ($\mathcal{A}\subset\mathbb{R}^{8}$), HandReorient ($\mathcal{A}\subset\mathbb{R}^{16}$), RollingFlat ($\mathcal{A}\subset\mathbb{R}^{3}$), SoftJumper ($\mathcal{A}\subset\mathbb{R}^{222}$), HandFlip ($\mathcal{A}\subset\mathbb{R}^{24}$), and FluidMove ($\mathcal{A}\subset\mathbb{R}^{3}$) tasks. Smoothed using EWMA with $\alpha=0.99$. Mean and 95% CIs over 10 random seeds.
• Figure 3: SAPO ablations -- HandFlip training curves. Episode return as a function of environment steps. Smoothed using EWMA with $\alpha=0.99$. Mean and 95% CIs over 10 random seeds.
• Figure 4: Diagram of computational graph of SAPO. Note that $J(\pi)$ only updates the actor, although gradients are still propagated through the critic.
• Figure 5: Example of SHAC getting stuck in local minima. Episode return as a function of environment steps in DFlex Ant ($\mathcal{A}\subset\mathbb{R}^{8}$). One run (colored in red) quickly plateaus after 1M steps and does not improve. 6 random seeds.