Brax - A Differentiable Physics Engine for Large Scale Rigid Body SimulationDownload PDF

Published: 29 Jul 2021, Last Modified: 04 May 2025NeurIPS 2021 Datasets and Benchmarks Track (Round 1)Readers: Everyone
Keywords: reinforcement learning, physics, differentiable, jax, robotics
TL;DR: We present a new differentiable rigid body physics engine in JAX and several reimplementations of common RL algorithms that can compile to and run on the same accelerator.
Abstract: We present Brax, an open source library for \textbf{r}igid \textbf{b}ody 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 MuJoCo-like tasks in minutes.
URL: https://github.com/google/brax
Supplementary Material: zip
Community Implementations: [![CatalyzeX](/images/catalyzex_icon.svg) 1 code implementation](https://www.catalyzex.com/paper/brax-a-differentiable-physics-engine-for/code)
5 Replies

Loading

OpenReview is a long-term project to advance science through improved peer review with legal nonprofit status. We gratefully acknowledge the support of the OpenReview Sponsors. © 2025 OpenReview