Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation

TL;DR

Brax is an open-source, high-performance, differentiable physics engine built in JAX that enables scalable reinforcement learning by integrating environment simulation and learning algorithms on accelerators.

Contribution

Brax introduces a novel, efficient physics simulation library in JAX with integrated reinforcement learning algorithms, facilitating large-scale, accelerated policy training.

Findings

Achieves high performance and parallelism on accelerators.

Supports seamless integration of environment simulation and learning algorithms.

Enables training of policies on MuJoCo-like tasks in minutes.

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.