Package: brax
Brax is a fast and fully differentiable physics engine used for research and
development of robotics, human perception, materials science, reinforcement
learning, and other simulation-heavy applications.Brax is written inJAXand 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.Brax simulates environments at millions of physics steps per second on TPU, and includes a suite of learning algorithms that train agents in seconds
to minutes:Baseline learning algorithms such asPPO,SAC,ARS, andevolutionary strategies.Learning algorithms that leverage the differentiability of the simulator, such asanalytic policy gradients.One API, Four PipelinesBrax offers four distinct physics pipelines that are easy to swap:MuJoCo XLA - MJX- a JAX
reimplementation of the MuJoCo physics engine.Generalizedcalculates motion ingeneralized coordinatesusing dynamics algorithms similar toMuJoCoandTDS.PositionalusesPosition Based Dynamics,
a fast but stable method of resolving joint and collision constraints.Springprovides
fast and cheap simulation for rapid experimentation, using simple impulse-based
methods often found in video games.These pipelines share the same API and can run side-by-side within the same
simulation.  This makes Brax well suited for experiments in transfer learning
and closing the gap between simulation and the real world.Quickstart: Colab in the CloudExplore Brax easily and quickly through a series of colab notebooks:Brax Basicsintroduces the Brax API, and shows how to simulate basic physics primitives.Brax Trainingintroduces 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 MuJoCo XLA - MJXdemonstrates training in Brax using theMJXphysics simulator.Brax Training with PyTorch on GPUdemonstrates how Brax can be used in other ML frameworks for fast training, in this case PyTorch.Using Brax LocallyTo install Brax from pypi, install it with:python3 -m venv env
source env/bin/activate
pip install --upgrade pip
pip install braxYou may also install fromCondaorMamba:conda install -c conda-forge brax  # s/conda/mamba for mambaAlternatively, to install Brax from source, clone this repo,cdto it, and then:python3 -m venv env
source env/bin/activate
pip install --upgrade pip
pip install -e .To train a model:learnTraining on NVidia GPU is supported, but you must first installCUDA, CuDNN, and JAX with GPU support.Learn MoreFor 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, which appeared in theDatasets and Benchmarks TrackatNeurIPS 2021.Citing BraxIf you would like to reference Brax in a publication, please use:@software{brax2021github,
  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 = {http://github.com/google/brax},
  version = {0.10.0},
  year = {2021},
}AcknowledgementsBrax has come a long way since its original publication.  We offer gratitude and
effusive praise to the following people:Manu Orsini and Nikola Momchev who provided a major refactor of Brax's
training algorithms to make them more accessible and reusable.Erwin Coumans who has graciously offered advice and mentorship, and many
useful references fromTiny Differentiable Simulator.Baruch Tabanpour, a colleague who helped launch brax v2 and overhauled the contact library.Shixiang Shane GuandHiroki Furuta, who contributed BIG-Gym and Braxlines, and a scene composer to Brax.Our awesomeopen source collaborators and contributors.  Thank you!