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

By C. Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, Olivier Bachem. NeurIPS Datasets and Benchmarks, 2021

578 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Science Applications, Computer Vision and Pattern Recognition

View Paper: https://www.semanticscholar.org/paper/13263ec6b83e59442aafe1b45b7554bbba57b680

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Science Applications

In Computer Vision and Pattern Recognition

Citations in the corpus

PDEBENCH: An Extensive Benchmark for Scientific Machine Learning (2022, background): “Differentiable direct PDE solvers are increasingly available, e.g. [37] and frequently built upon neural network technology stacks [13, 3, 18].”
Building a Subspace of Policies for Scalable Continual Learning (2022, methodology): “[12] C Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem.”
Gradients are Not All You Need (2021, methodology): “…physics [Heiden et al., 2021, Hu et al., 2019, Werling et al., 2021, Degrave et al., 2019, de Avila Belbute-Peres et al., 2018, Gradu et al., 2021, Freeman et al., 2021], graphics [Li et al., 2018, Kato et al., 2020], molecular dynamics [Schoenholz and Cubuk, 2020, Hinsen, 2000], differentiating…”
Gymnasium: A Standard Interface for Reinforcement Learning Environments (2025, background): “We expand on this in the following section.”
Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning (2021, result): “We note that others [8] have recently begun attempting an approach similar to Isaac Gym with respect to running end-to-end training on hardware accelerators.”
AlpacaFarm: A Simulation Framework for Methods that Learn from Human Feedback (2023, background): “More broadly, building a simulator environment to enable low-cost experimentation is common in the field of reinforcement learning and robotics [11, 67, 66, 64, 17, 23, 18].”