brax


Namebrax JSON
Version 0.11.0 PyPI version JSON
download
home_pagehttp://github.com/google/brax
SummaryA differentiable physics engine written in JAX.
upload_time2024-10-07 21:18:55
maintainerNone
docs_urlNone
authorBrax Authors
requires_pythonNone
licenseApache 2.0
keywords jax reinforcement learning rigidbody physics
VCS
bugtrack_url
requirements No requirements were recorded.
Travis-CI No Travis.
coveralls test coverage No coveralls.
            <img src="https://github.com/google/brax/raw/main/docs/img/brax_logo.gif" width="336" height="80" alt="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 in [JAX](https://github.com/google/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.

<img src="https://github.com/google/brax/raw/main/docs/img/humanoid_v2.gif" width="160" height="160"/><img src="https://github.com/google/brax/raw/main/docs/img/a1.gif" width="160" height="160"/><img src="https://github.com/google/brax/raw/main/docs/img/ant_v2.gif" width="160" height="160"/><img src="https://github.com/google/brax/raw/main/docs/img/ur5e.gif" width="160" height="160"/>

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 as
    [PPO](https://github.com/google/brax/blob/main/brax/training/agents/ppo),
    [SAC](https://github.com/google/brax/blob/main/brax/training/agents/sac),
    [ARS](https://github.com/google/brax/blob/main/brax/training/agents/ars), and
    [evolutionary strategies](https://github.com/google/brax/blob/main/brax/training/agents/es).
*   Learning algorithms that leverage the differentiability of the simulator, such as [analytic policy gradients](https://github.com/google/brax/blob/main/brax/training/agents/apg).

## One API, Four Pipelines

Brax offers four distinct physics pipelines that are easy to swap:

* [MuJoCo XLA - MJX](https://mujoco.readthedocs.io/en/stable/mjx.html) - a JAX
reimplementation of the MuJoCo physics engine.
* [Generalized](https://github.com/google/brax/blob/main/brax/generalized/)
calculates motion in [generalized coordinates](https://en.wikipedia.org/wiki/Generalized_coordinates)
using dynamics algorithms similar to [MuJoCo](https://mujoco.org/) and [TDS](https://github.com/erwincoumans/tiny-differentiable-simulator).
* [Positional](https://github.com/google/brax/blob/main/brax/positional/)
uses [Position Based Dynamics](https://matthias-research.github.io/pages/publications/posBasedDyn.pdf),
a fast but stable method of resolving joint and collision constraints.
* [Spring](https://github.com/google/brax/blob/main/brax/spring/) provides
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 Cloud

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

* [Brax Basics](https://colab.research.google.com/github/google/brax/blob/main/notebooks/basics.ipynb) introduces the Brax API, and shows how to simulate basic physics primitives.
* [Brax Training](https://colab.research.google.com/github/google/brax/blob/main/notebooks/training.ipynb) 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 MuJoCo XLA - MJX](https://colab.sandbox.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb) demonstrates training in Brax using the `MJX` physics simulator.
* [Brax Training with PyTorch on GPU](https://colab.research.google.com/github/google/brax/blob/main/notebooks/training_torch.ipynb) demonstrates how Brax can be used in other ML frameworks for fast training, in this case PyTorch.

## 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
```

You may also install from [Conda](https://docs.conda.io/en/latest/) or [Mamba](https://github.com/mamba-org/mamba):

```
conda install -c conda-forge brax  # s/conda/mamba for mamba
```

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:

```
learn
```

Training on NVidia GPU is supported, but you must first install
[CUDA, CuDNN, and JAX with GPU support](https://github.com/google/jax#installation).

## 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
](https://arxiv.org/abs/2106.13281), which appeared in the [Datasets and Benchmarks Track](https://neurips.cc/Conferences/2021/CallForDatasetsBenchmarks) at [NeurIPS 2021](https://nips.cc/Conferences/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.11.0},
  year = {2021},
}
```

## Acknowledgements

Brax 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 from [Tiny Differentiable Simulator](https://github.com/erwincoumans/tiny-differentiable-simulator).
* Baruch Tabanpour, a colleague who helped launch brax v2 and overhauled the contact library.
* [Shixiang Shane Gu](https://sites.google.com/corp/view/gugurus) and [Hiroki Furuta](https://frt03.github.io/), who contributed BIG-Gym and Braxlines, and a scene composer to Brax.
* Our awesome [open source collaborators and contributors](https://github.com/google/brax/graphs/contributors).  Thank you!

            

Raw data

            {
    "_id": null,
    "home_page": "http://github.com/google/brax",
    "name": "brax",
    "maintainer": null,
    "docs_url": null,
    "requires_python": null,
    "maintainer_email": null,
    "keywords": "JAX reinforcement learning rigidbody physics",
    "author": "Brax Authors",
    "author_email": "no-reply@google.com",
    "download_url": "https://files.pythonhosted.org/packages/43/47/a3d0d90fd92369183eeba258c78385067dd1df872a0ffab26adc8911dfa9/brax-0.11.0.tar.gz",
    "platform": null,
    "description": "<img src=\"https://github.com/google/brax/raw/main/docs/img/brax_logo.gif\" width=\"336\" height=\"80\" alt=\"BRAX\"/>\n\nBrax is a fast and fully differentiable physics engine used for research and\ndevelopment of robotics, human perception, materials science, reinforcement\nlearning, and other simulation-heavy applications.\n\nBrax is written in [JAX](https://github.com/google/jax) and is designed for use\non acceleration hardware. It is both efficient for single-device simulation, and\nscalable to massively parallel simulation on multiple devices, without the need\nfor pesky datacenters.\n\n<img src=\"https://github.com/google/brax/raw/main/docs/img/humanoid_v2.gif\" width=\"160\" height=\"160\"/><img src=\"https://github.com/google/brax/raw/main/docs/img/a1.gif\" width=\"160\" height=\"160\"/><img src=\"https://github.com/google/brax/raw/main/docs/img/ant_v2.gif\" width=\"160\" height=\"160\"/><img src=\"https://github.com/google/brax/raw/main/docs/img/ur5e.gif\" width=\"160\" height=\"160\"/>\n\nBrax simulates environments at millions of physics steps per second on TPU, and includes a suite of learning algorithms that train agents in seconds\nto minutes:\n\n*   Baseline learning algorithms such as\n    [PPO](https://github.com/google/brax/blob/main/brax/training/agents/ppo),\n    [SAC](https://github.com/google/brax/blob/main/brax/training/agents/sac),\n    [ARS](https://github.com/google/brax/blob/main/brax/training/agents/ars), and\n    [evolutionary strategies](https://github.com/google/brax/blob/main/brax/training/agents/es).\n*   Learning algorithms that leverage the differentiability of the simulator, such as [analytic policy gradients](https://github.com/google/brax/blob/main/brax/training/agents/apg).\n\n## One API, Four Pipelines\n\nBrax offers four distinct physics pipelines that are easy to swap:\n\n* [MuJoCo XLA - MJX](https://mujoco.readthedocs.io/en/stable/mjx.html) - a JAX\nreimplementation of the MuJoCo physics engine.\n* [Generalized](https://github.com/google/brax/blob/main/brax/generalized/)\ncalculates motion in [generalized coordinates](https://en.wikipedia.org/wiki/Generalized_coordinates)\nusing dynamics algorithms similar to [MuJoCo](https://mujoco.org/) and [TDS](https://github.com/erwincoumans/tiny-differentiable-simulator).\n* [Positional](https://github.com/google/brax/blob/main/brax/positional/)\nuses [Position Based Dynamics](https://matthias-research.github.io/pages/publications/posBasedDyn.pdf),\na fast but stable method of resolving joint and collision constraints.\n* [Spring](https://github.com/google/brax/blob/main/brax/spring/) provides\nfast and cheap simulation for rapid experimentation, using simple impulse-based\nmethods often found in video games.\n\nThese pipelines share the same API and can run side-by-side within the same\nsimulation.  This makes Brax well suited for experiments in transfer learning\nand closing the gap between simulation and the real world.\n\n## Quickstart: Colab in the Cloud\n\nExplore Brax easily and quickly through a series of colab notebooks:\n\n* [Brax Basics](https://colab.research.google.com/github/google/brax/blob/main/notebooks/basics.ipynb) introduces the Brax API, and shows how to simulate basic physics primitives.\n* [Brax Training](https://colab.research.google.com/github/google/brax/blob/main/notebooks/training.ipynb) introduces Brax's training algorithms, and lets you train your own policies directly within the colab. It also demonstrates loading and saving policies.\n* [Brax Training with MuJoCo XLA - MJX](https://colab.sandbox.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb) demonstrates training in Brax using the `MJX` physics simulator.\n* [Brax Training with PyTorch on GPU](https://colab.research.google.com/github/google/brax/blob/main/notebooks/training_torch.ipynb) demonstrates how Brax can be used in other ML frameworks for fast training, in this case PyTorch.\n\n## Using Brax Locally\n\nTo install Brax from pypi, install it with:\n\n```\npython3 -m venv env\nsource env/bin/activate\npip install --upgrade pip\npip install brax\n```\n\nYou may also install from [Conda](https://docs.conda.io/en/latest/) or [Mamba](https://github.com/mamba-org/mamba):\n\n```\nconda install -c conda-forge brax  # s/conda/mamba for mamba\n```\n\nAlternatively, to install Brax from source, clone this repo, `cd` to it, and then:\n\n```\npython3 -m venv env\nsource env/bin/activate\npip install --upgrade pip\npip install -e .\n```\n\nTo train a model:\n\n```\nlearn\n```\n\nTraining on NVidia GPU is supported, but you must first install\n[CUDA, CuDNN, and JAX with GPU support](https://github.com/google/jax#installation).\n\n## Learn More\n\nFor a deep dive into Brax's design and performance characteristics, please see\nour paper, [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation\n](https://arxiv.org/abs/2106.13281), which appeared in the [Datasets and Benchmarks Track](https://neurips.cc/Conferences/2021/CallForDatasetsBenchmarks) at [NeurIPS 2021](https://nips.cc/Conferences/2021).\n\n## Citing Brax\n\nIf you would like to reference Brax in a publication, please use:\n\n```\n@software{brax2021github,\n  author = {C. Daniel Freeman and Erik Frey and Anton Raichuk and Sertan Girgin and Igor Mordatch and Olivier Bachem},\n  title = {Brax - A Differentiable Physics Engine for Large Scale Rigid Body Simulation},\n  url = {http://github.com/google/brax},\n  version = {0.11.0},\n  year = {2021},\n}\n```\n\n## Acknowledgements\n\nBrax has come a long way since its original publication.  We offer gratitude and\neffusive praise to the following people:\n\n* Manu Orsini and Nikola Momchev who provided a major refactor of Brax's\ntraining algorithms to make them more accessible and reusable.\n* Erwin Coumans who has graciously offered advice and mentorship, and many\nuseful references from [Tiny Differentiable Simulator](https://github.com/erwincoumans/tiny-differentiable-simulator).\n* Baruch Tabanpour, a colleague who helped launch brax v2 and overhauled the contact library.\n* [Shixiang Shane Gu](https://sites.google.com/corp/view/gugurus) and [Hiroki Furuta](https://frt03.github.io/), who contributed BIG-Gym and Braxlines, and a scene composer to Brax.\n* Our awesome [open source collaborators and contributors](https://github.com/google/brax/graphs/contributors).  Thank you!\n",
    "bugtrack_url": null,
    "license": "Apache 2.0",
    "summary": "A differentiable physics engine written in JAX.",
    "version": "0.11.0",
    "project_urls": {
        "Homepage": "http://github.com/google/brax"
    },
    "split_keywords": [
        "jax",
        "reinforcement",
        "learning",
        "rigidbody",
        "physics"
    ],
    "urls": [
        {
            "comment_text": "",
            "digests": {
                "blake2b_256": "11de134be9686bdc31943f9fd30de10ff32a0a2c2c2d995fd89f97c954525a1f",
                "md5": "86f31f703a2009b65f08c32064a7388e",
                "sha256": "46e10afb9185d77c0f5180cb098beeb67339b9bb8a35bdf5556f76ba3875abcf"
            },
            "downloads": -1,
            "filename": "brax-0.11.0-py3-none-any.whl",
            "has_sig": false,
            "md5_digest": "86f31f703a2009b65f08c32064a7388e",
            "packagetype": "bdist_wheel",
            "python_version": "py3",
            "requires_python": null,
            "size": 998555,
            "upload_time": "2024-10-07T21:18:53",
            "upload_time_iso_8601": "2024-10-07T21:18:53.805333Z",
            "url": "https://files.pythonhosted.org/packages/11/de/134be9686bdc31943f9fd30de10ff32a0a2c2c2d995fd89f97c954525a1f/brax-0.11.0-py3-none-any.whl",
            "yanked": false,
            "yanked_reason": null
        },
        {
            "comment_text": "",
            "digests": {
                "blake2b_256": "4347a3d0d90fd92369183eeba258c78385067dd1df872a0ffab26adc8911dfa9",
                "md5": "77177bf8e4d916c77e339ab09b9f62c0",
                "sha256": "7efae65345d8a4700bf24af607bf021df488ef1bc008f0579c9118c7e9e54dc1"
            },
            "downloads": -1,
            "filename": "brax-0.11.0.tar.gz",
            "has_sig": false,
            "md5_digest": "77177bf8e4d916c77e339ab09b9f62c0",
            "packagetype": "sdist",
            "python_version": "source",
            "requires_python": null,
            "size": 767216,
            "upload_time": "2024-10-07T21:18:55",
            "upload_time_iso_8601": "2024-10-07T21:18:55.838453Z",
            "url": "https://files.pythonhosted.org/packages/43/47/a3d0d90fd92369183eeba258c78385067dd1df872a0ffab26adc8911dfa9/brax-0.11.0.tar.gz",
            "yanked": false,
            "yanked_reason": null
        }
    ],
    "upload_time": "2024-10-07 21:18:55",
    "github": true,
    "gitlab": false,
    "bitbucket": false,
    "codeberg": false,
    "github_user": "google",
    "github_project": "brax",
    "travis_ci": false,
    "coveralls": false,
    "github_actions": true,
    "lcname": "brax"
}
        
Elapsed time: 0.35970s