[Paper Review] Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation

C. Daniel Freeman, Erik Frey+4 more|arXiv (Cornell University)|Jun 24, 2021

Context-Aware Activity Recognition Systems61 citations

TL;DR

Brax is an open-source, differentiable rigid-body physics engine built in JAX that runs on accelerators, enabling large-scale parallel environment simulation and on-device learning with bundled RL algorithms.

ABSTRACT

We present Brax, an open source library for rigid body simulation with a focus on performance and parallelism on accelerators, written in JAX. We present results on a suite of tasks inspired by the existing reinforcement learning literature, but remade in our engine. Additionally, we provide reimplementations of PPO, SAC, ES, and direct policy optimization in JAX that compile alongside our environments, allowing the learning algorithm and the environment processing to occur on the same device, and to scale seamlessly on accelerators. Finally, we include notebooks that facilitate training of performant policies on common OpenAI Gym MuJoCo-like tasks in minutes.

Motivation & Objective

• Motivate faster, more accessible RL research by co-locating physics and learning on accelerators.
• Provide a differentiable, open-source engine to enable gradient-based RL methods.
• Offer a ProtoBuf-based environment specification and Gym-like interfaces for scalable tasks.
• Demonstrate scalability and performance of Brax across accelerators and compare to MuJoCo.
• Survey limitations and future work to guide improvements in differentiable physics engines.

Proposed method

• Implement maximal-coordinate rigid-body dynamics with a QP data primitive and vectorized transforms.
• Use JAX for auto-vectorization, device-parallelism, JIT, and automatic differentiation to run thousands of environments on accelerators.
• Define physics steps as parallel transformations (joints, actuators, colliders) applied to a Brax_system with a second-order symplectic Euler integrator.
• Provide a ProtoBuf-based system specification and gym-like Env for building and solving physics tasks.
• Bundle RL algorithms (PPO, SAC, ES, APG) in JAX to run entirely on accelerators with on-device rollouts and gradient-based updates.
• Demonstrate environment suite including MuJoCo-like Ant, Humanoid, Halfcheetah, Grasp, and Fetch.

Experimental results

Research questions

• RQ1 Can a differentiable physics engine running on accelerators dramatically speed up RL training by co-locating environment simulation and learning?
• RQ2 How does Brax performance and scalability on TPU/GPU compare to traditional CPU-based engines like MuJoCo on common RL benchmarks?
• RQ3 What are the effects of differentiability on learning algorithms (PPO, SAC, ES, APG) when implemented in Brax?
• RQ4 What are the practical limitations and required tuning for large-scale, differentiable rigid-body simulations?
• RQ5 How effective is the ProtoBuf environment specification and gym-like interface for building diverse physically-based tasks?

Key findings

• Brax achieves millions of simulation steps per second on a single accelerator for certain environments.
• On a 4x2 TPUv3, Brax scales to hundreds of millions of environment steps per second by distributing computation across accelerators.
• PPO and SAC implementations compiled for Brax enable on-device rollouts and gradient updates with minimal data transfer overhead, achieving significantly faster training than non-accelerated baselines.
• Brax environments qualitatively match MuJoCo counterparts in learning trajectories and rewards under comparable hyperparameters when evaluated with SAC.
• Compared to MuJoCo A nt, Brax with optimized PPO reaches locomotion in about 10 seconds versus ~30 minutes for a standard PPO implementation.
• Brax demonstrates good momentum conservation and competitive energy/angular momentum behavior relative to other engines.