A differentiable physics engine written in JAX.
Project description
BRAX
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:
- Baseline learning algorithms such as PPO, SAC, and evolutionary strategies.
- Experimental algorithms such as Variational GCRL, Adversarial Inverse RL, and State Marginal Matching
- Learning algorithms that leverage the differentiability of the simulator, such as analytic policy gradients.
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 Training introduces Brax environments and training algorithms, and lets you train your own policies directly within the colab.
- Brax + PyTorch demonstrates how to operate Brax environments performantly from PyTorch.
Using Brax locally
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:
learn
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:
@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.1.0},
year = {2021},
}
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
File details
Details for the file brax-0.0.4.tar.gz.
File metadata
- Download URL: brax-0.0.4.tar.gz
- Upload date:
- Size: 76.7 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/3.4.2 importlib_metadata/4.6.4 pkginfo/1.7.1 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.1 CPython/3.9.6
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 | 76dd527569ff0c8e076a104864092f6d13006f3f0132f3d25912d62d7fd6b0e9 |
|
| MD5 | 77af5dec39de49dec5b37827dbe45d6b |
|
| BLAKE2b-256 | e060be51db03cec82031e5cb0e28730a2894db70c79476a234d2a3a463604237 |
Provenance
File details
Details for the file brax-0.0.4-py3-none-any.whl.
File metadata
- Download URL: brax-0.0.4-py3-none-any.whl
- Upload date:
- Size: 117.8 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/3.4.2 importlib_metadata/4.6.4 pkginfo/1.7.1 requests/2.26.0 requests-toolbelt/0.9.1 tqdm/4.62.1 CPython/3.9.6
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 | a587f3ac81d804094bc3e5758448e404b422ab3a0796300aab7bf54d11ca4ada |
|
| MD5 | c64d651b1a105a532221c22d466b26e3 |
|
| BLAKE2b-256 | 74d436a0070abd8d73a9d0c61ceefe43744e48e3736fb331079ae6e16c3d8403 |
