google/brax

What is it?

Key features

• 01Runs rigid-body simulation in JAX on GPU and TPU hardware, scaling across multiple devices.
• 02Provides four swappable physics pipelines (MJX, generalized, positional, spring) under one API.
• 03Includes PPO, SAC, ARS, evolutionary strategies, and analytic policy gradients for training.
• 04Supports differentiable simulation so gradients can flow through physics steps.
• 05Offers Colab notebooks for basics, training, MJX, and PyTorch integration.

Use cases

• Training locomotion or manipulation policies for simulated robots with reinforcement learning.
• Running large batches of parallel physics simulations on accelerator hardware.
• Experimenting with differentiable physics and analytic policy gradient methods.
• Using a simulator as a training environment inside PyTorch or other ML frameworks.

Quick start

1. 1

   Create and activate a virtual environment

   python3 -m venv env && source env/bin/activate
2. 2

   Upgrade pip

   pip install --upgrade pip
3. 3

   Install Brax from PyPI

   pip install brax
4. 4

   Train a model

   learn

Things to know

• !Only brax/training is actively maintained; brax/envs is superseded by MuJoCo Playground.
• !Training on NVIDIA GPU requires separately installing CUDA, cuDNN and GPU-enabled JAX.

In collections

Topics

FAQ

What is brax?

A JAX-based rigid-body physics engine that runs massively parallel, differentiable simulation on accelerators for robotics and RL research. Brax is a physics engine written in JAX and built for accelerator hardware such as GPUs and TPUs. It targets research in robotics, reinforcement learning, human perception, and materials science, and can scale simulation across multiple devices.

Is brax free to use? What license does it have?

brax is open source under the Apache-2.0 license. Check the license terms before using it in commercial products.

Is brax still maintained?

Actively maintained. The last commit to brax was on Oct 20, 2026.

What can I use brax for?

Training locomotion or manipulation policies for simulated robots with reinforcement learning. Running large batches of parallel physics simulations on accelerator hardware. Experimenting with differentiable physics and analytic policy gradient methods. Using a simulator as a training environment inside PyTorch or other ML frameworks.

How do I install brax?

Create and activate a virtual environment: python3 -m venv env && source env/bin/activate Upgrade pip: pip install --upgrade pip Install Brax from PyPI: pip install brax Train a model: learn

Where can I get brax?

The source code and installation instructions are on GitHub: https://github.com/google/brax