November 3, 2021

441 words 3 mins read



Massively parallel rigidbody physics simulation on accelerator hardware.

repo name google/brax
repo link
language Jupyter Notebook
size (curr.) 135468 kB
stars (curr.) 1029
created 2021-06-02
license Apache License 2.0

Brax is a differentiable physics engine that simulates environments made up of rigid bodies, joints, and actuators. Brax is written in JAX and is designed for use on acceleration hardware. It is both efficient for single-device simulation, and scalable to massively parallel simulation on multiple devices, without the need for pesky datacenters.

Some policies trained via Brax. Brax simulates these environments at millions of physics steps per second on TPU.

Brax also includes a suite of learning algorithms that train agents in seconds to minutes:

Quickstart: Colab in the Cloud

Explore Brax easily and quickly through a series of colab notebooks:

  • Brax Basics introduces the Brax API, and shows how to simulate basic physics primitives.
  • Brax Environments shows how to operate and visualize Brax environments. It also demonstrates converting Brax environments to Gym environments, and how to use Brax via other ML frameworks such as PyTorch.
  • Brax Training with TPU introduces Brax’s training algorithms, and lets you train your own policies directly within the colab. It also demonstrates loading and saving policies.
  • Brax Training with PyTorch on GPU demonstrates how Brax can be used in other ML frameworks for fast training, in this case PyTorch.
  • Brax Multi-Agent measures Brax’s performance on multi-agent simulation, with many bodies in the environment at once.

Using Brax locally

To install Brax from pypi, install it with:

python3 -m venv env
source env/bin/activate
pip install --upgrade pip
pip install brax

Alternatively, to install Brax from source, clone this repo, cd to it, and then:

python3 -m venv env
source env/bin/activate
pip install --upgrade pip
pip install -e .

To train a model:


Training on NVidia GPU is supported, but you must first install CUDA, CuDNN, and JAX with GPU support.

Learn More

For a deep dive into Brax’s design and performance characteristics, please see our paper, Brax – A Differentiable Physics Engine for Large Scale Rigid Body Simulation , to appear in the Datasets and Benchmarks Track at NeurIPS 2021.

Citing Brax

If you would like to reference Brax in a publication, please use:

  author = {C. Daniel Freeman and Erik Frey and Anton Raichuk and Sertan Girgin and Igor Mordatch and Olivier Bachem},
  title = {Brax - A Differentiable Physics Engine for Large Scale Rigid Body Simulation},
  url = {},
  version = {0.0.7},
  year = {2021},
comments powered by Disqus