brax

library0.14.2pypypi✓ verified 84d ago

A differentiable physics engine written in JAX, used for robot locomotion, reinforcement learning, and simulation-based research. Current version: 0.14.2. Released under Apache 2.0, with updates approximately every 2-3 months.

ai-ml

pip install brax

github ↗homepage ↗

INSTALL

IMPORT

SIG · BRAX

brax

ai-mlpythonv0.14.2

harness data pending

Install & Compatibility

Where this runs

No compatibility data collected yet for this library.

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

brax

✓ import brax

Pipelines (env wrapper)

✓ from brax.io import json

Environments

✓ from brax.envs import create_environment

Creates an Ant environment in Brax with batching.

import brax from brax import envs import jax env = envs.create_environment('ant', batch_size=8) state = env.reset(rng=jax.random.PRNGKey(0)) print(state.obs.shape)

Debug

Known issues

breakingBrax 0.10+ removed the `env` function from `brax.envs`. Use `create_environment` instead.

fix

Replace `env = brax.envs.env('ant')` with `env = brax.envs.create_environment('ant')`.

affects: >=0.10

breakingBrax 0.14 dropped support for Python <3.11. Ensure Python 3.11 or later.

fix

Upgrade Python to 3.11+.

affects: >=0.14

gotchaBrax environments return batches by default. For single-environment simulation, set `batch_size=None` (or omit).

fix

Use `create_environment('ant', batch_size=None)` for unbatched.

affects: all

deprecatedThe `jax.config.update('jax_platform_name', 'cpu')` pattern is discouraged. Use `JAX_PLATFORM_NAME=cpu` environment variable.

fix

Set environment variable: `JAX_PLATFORM_NAME=cpu`

affects: all

Errors

Common errors & fixes

ModuleNotFoundError: No module named 'brax.envs'

Brax version <0.10 had different module structure, or install is outdated.

fix

Upgrade brax: `pip install --upgrade brax`

AssertionError: Pipeline is not a valid Brax environment

Trying to use `env = brax.envs.env('humanoid')` which is removed in 0.10+.

fix

Use `env = brax.envs.create_environment('humanoid')`

ValueError: The batch shape is not compatible with the environment

Passing a batch of states to a non-batched environment or vice versa.

fix

Ensure `batch_size` matches between environment creation and state objects.

Upgrade

Version history

0.14.2latest on PyPI · released Mar 15, 2026

Audit

Dependencies

jaxrequiredCore dependency for differentiable physics

jaxlibrequiredJAX backend

flaxoptionalNeural network library for RL policies

optaxoptionalOptimizer for training

Agent activity

33 hits · last 30 days

node

OpenAI (training)

Resources

githubhttp://github.com/google/brax ↗

homepagehttp://github.com/google/brax ↗