google/brax: a differentiable physics engine in JAX, now mostly an RL library

What Brax solves, and who the README now says it is for

Brax is a differentiable physics engine written in JAX, aimed at rigid body simulation on accelerator hardware. The pitch is scale: the README states it simulates environments at millions of physics steps per second on TPU, and that it scales to massively parallel simulation across multiple devices without a datacenter. Differentiability is the other half. Because the simulator is written in JAX, gradients flow through the physics, which is what makes analytic policy gradients possible rather than treating the environment as a black box.

The intended audience has narrowed, though. The warning at the top of the README says only brax/training is actively being maintained as of 0.13.0, that users of brax/envs should move to MuJoCo Playground, and that anyone wanting Brax for physics simulation should use MJX (pip install mujoco_mjx) or MuJoCo Warp instead of Brax as a wrapper to MuJoCo. It also floats the possibility that brax may be repurposed purely as an RL library. That is an unusually direct statement of scope for a project this size, and it should shape the decision more than any benchmark figure. If your interest is a fast differentiable simulator, the README is telling you to look elsewhere. If your interest is training agents in JAX, you are in the part that is still maintained.

Four pipelines behind one API: MJX, Generalized, Positional and Spring

The architectural idea is that physics backends are swappable. Brax exposes four pipelines under a shared API, and the README says they can run side by side within the same simulation. That last property is the interesting one: it makes transfer experiments between a cheap simulator and an expensive one a configuration change rather than a rewrite.

The four differ in how they resolve motion. MJX is described as a JAX reimplementation of the MuJoCo physics engine. Generalized computes motion in generalized coordinates using dynamics algorithms similar to MuJoCo and Tiny Differentiable Simulator. Positional uses Position Based Dynamics, which the README characterizes as a fast but stable way to resolve joint and collision constraints. Spring uses simple impulse-based methods, the kind found in video games, for cheap rapid experimentation. The trade-off axis is roughly fidelity against throughput, with Spring at the cheap end and MJX and Generalized at the accurate end. Note that MJX is a separate package pulled in as a dependency, so the MuJoCo path is not self-contained inside Brax.

On top of the simulator sits the training layer, with baseline algorithms PPO, SAC, ARS and evolutionary strategies, plus the differentiable-specific analytic policy gradients. The repository layout reflects this split: the top level holds brax/, notebooks/, datasets/ and docs/, and the packaging config excludes datasets, docs, notebooks and tests from the wheel, so an installed Brax is the library, not the tutorials.

Installing Brax and running your first training job

The README gives a virtual environment install from PyPI. It creates an environment, activates it, upgrades pip and installs the package. The pyproject.toml sets requires-python to >=3.11, so an older interpreter will fail before anything else does.

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

A Conda or Mamba install is also documented. The README writes the Mamba variant as a comment on the same line, so the command is identical apart from the binary name.

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

For a source checkout, clone the repository, change into it, and install in editable mode. This is the path to take if you intend to touch the pipelines themselves rather than just train against them.

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

Once installed, the README's training entry point is a single command. It is terse, and the README does not spell out its flags or which algorithm it defaults to, so treat it as a starting point rather than a configured run.

learn

GPU training is supported but not automatic. The README says you must first install CUDA, CuDNN and JAX with GPU support, and links to the JAX installation page for that. If you skip it, JAX falls back to CPU and the parallel simulation story does not apply. For a guided first run, the README points to Colab notebooks: Brax Basics for the API and physics primitives, Brax Training for the algorithms plus loading and saving policies, a MuJoCo XLA tutorial, and a notebook showing Brax training with PyTorch on GPU.

Where Brax is the wrong tool

The clearest limitation is stated by the maintainers themselves. If you want a physics simulator, the README tells you to use MJX or MuJoCo Warp rather than Brax as a wrapper to MuJoCo. That is not a subtle performance caveat; it is a routing instruction. Building a new simulation-heavy project on the Brax pipelines means building on a surface the project has publicly deprioritized.

The envs situation is similar. The README says users of brax/envs should use MuJoCo Playground instead, and notes that those environments train well with brax/training. So the supported combination is Brax's training stack against someone else's environments, which is a narrower product than the repository's four-pipeline framing suggests.

There is a second class of limitation that follows from the JAX design rather than from maintenance. JAX code is shaped by tracing and functional purity. The README's own framing is that Brax is written in JAX and designed for accelerator hardware, which is a commitment to a particular programming model and a particular dependency stack: the install pulls jax and jaxlib at >=0.4.6 alongside flax, optax, jaxopt, orbax-checkpoint, mujoco, mujoco-mjx and trimesh. If your team is not already comfortable with JAX, the physics is not the hard part.

Finally, versioning. The pyproject.toml classifier reads Development Status :: 4 - Beta, and the release cadence in the repository shows v0.14.0 in December 2025, v0.14.1 in February 2026 and v0.14.2 in March 2026. Nothing in the README describes a deprecation policy or a stability guarantee across minor versions.

Brax against MJX and MuJoCo Warp: the actual difference

The obvious alternative is MJX, and the README makes the comparison for you. MJX is a JAX reimplementation of MuJoCo, distributed as the mujoco_mjx package, and Brax itself depends on it. The difference is ownership and scope: MJX is the supported way to do JAX-based MuJoCo physics, while Brax's own pipelines are the parts the README steers physics users away from. If your goal is accurate rigid body contact under JAX and you do not need Brax's training agents, MJX alone is the shorter path, and it is the one the README names first.

MuJoCo Warp is the other name the README gives, from google-deepmind/mujoco_warp. The README does not describe its internals, so the honest statement is that it is offered as the alternative for physics simulation alongside MJX, and that choosing between them is outside what the README documents.

For the training layer, the comparison is different, because that is the part still maintained. The README's own framing is that MuJoCo Playground environments train well with brax/training, which positions Brax's agents as complementary to the environment libraries rather than competing with them. The practical question is therefore not Brax versus MJX for simulation, but whether you want Brax's PPO, SAC, ARS, ES and APG implementations driving an environment that lives somewhere else.

Maintenance, releases and what the Apache-2.0 licence leaves you to check

The repository is not archived, and the last push was on 2026-09-15. The most recent tagged release is v0.14.2 from 2026-03-15, so the code is moving ahead of the tags. Combined with the README warning that only brax/training is actively maintained as of 0.13.0, the upgrade picture is uneven: expect the training code to track JAX releases, and do not assume the same for the physics pipelines.

Upgrade cost is dominated by the dependency set. Brax pins jax and jaxlib at >=0.4.6 and requires Python 3.11 or newer, and it pulls flax, optax, jaxopt, orbax-checkpoint, mujoco, mujoco-mjx, trimesh and tensorboardX. A JAX or MuJoCo bump can move the whole stack, so pinning Brax means pinning its neighbours too. Checkpoint compatibility is the piece to watch on upgrades, since orbax-checkpoint is a declared dependency and the README's training notebook is where policy saving and loading is demonstrated.

On licensing, the project is Apache-2.0, declared in pyproject.toml via license = { file = "LICENSE" } and listed as an OSI-approved Apache classifier. That is a permissive licence with an explicit patent grant, which is generally what research and commercial users want, but the dependency tree carries its own licences and Brax's Apache-2.0 does not speak for them. This is a description of what the repository declares, not legal advice; if licence compatibility matters for your distribution, review the LICENSE file and the licences of the dependencies yourself.

Editorial conclusion

Adopt Brax if you want PPO, SAC, ARS, evolutionary strategies or analytic policy gradients running against a JAX-native simulator and you are willing to treat brax/training as the supported surface. Do not adopt it as a general MuJoCo wrapper for physics work: the README directs that traffic to MJX (pip install mujoco_mjx) or MuJoCo Warp. Before committing, check which pipeline your model needs, confirm your Python is 3.11 or newer, and read the warning block at the top of the README yourself, since it defines the project's current boundary more clearly than any other page.

Frequently asked questions

Official sources

1. google/brax on GitHub
2. Issues
3. License: Apache-2.0
4. README
5. Releases