mujoco-mjx

What it is and what it does

MJX is a physics simulation library that reimplements MuJoCo using JAX, enabling physics computations to run on GPUs and TPUs. It maintains API compatibility with MuJoCo but is not a drop-in replacement—some features are missing. The package is developed and maintained by Google DeepMind and is kept synchronized with MuJoCo releases.

You use MJX when you need differentiable physics for reinforcement learning, want to exploit accelerator hardware for batch simulations, or need to integrate physics into JAX-based machine learning pipelines. It depends on JAX, jaxlib, the base mujoco package, scipy, trimesh, absl-py, and etils.

Use it for:

• Train reinforcement learning policies with differentiable physics gradients flowing through the simulation.
• Run batched physics simulations in parallel on GPU/TPU for robotics research or control optimization.
• Integrate physics simulation into JAX-based neural network training pipelines.
• Prototype physics-based learning algorithms that require automatic differentiation through dynamics.
• Accelerate large-scale multi-environment simulations for population-based training.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

A JAX-based reimplementation of the MuJoCo physics engine that runs physics simulations on accelerators (GPUs/TPUs) with an API compatible with MuJoCo.

Yes, if you are doing machine learning with physics simulation and have access to GPU/TPU hardware. MJX is actively maintained by DeepMind, has no known vulnerabilities, and is permissively licensed. Install only if you need JAX integration and can tolerate that some MuJoCo features are not yet supported; if you need full MuJoCo compatibility, use the base mujoco package instead.

Install

mujoco-mjx on PyPI

pip install mujoco-mjx

uv add mujoco-mjx

poetry add mujoco-mjx

Installing mujoco-mjx

Before you install

Low friction install with a pure-Python wheel. Actively maintained by Google DeepMind with a recent release (17 days old) and high repository activity (14550 stars). Requires modern Python (3.10+) and brings in substantial dependencies including JAX, jaxlib, and the base mujoco package.

License in practice

Licensed under Apache-2.0 (permissive), allowing commercial and private use with attribution. No restrictions on modification or distribution.

Quickstart

pip install mujoco-mjx

from mujoco import mjx
# Load and step a physics model on accelerator
model = mjx.Model.from_xml(xml_string)
data = mjx.Data(model)

Requires Python 3.10 or later. JAX and jaxlib must be installed and configured for your target hardware (CPU, GPU, or TPU).

Verify before relying

• Feature parity with MuJoCo: documentation notes some features are missing; specifics of what is unsupported are not in the fact sheet.
• Performance characteristics: whether acceleration actually improves simulation speed depends on workload and hardware; not quantified in the fact sheet.
• Batch simulation capability: whether the package supports vectorized/batched physics steps is implied but not explicitly confirmed.

Package facts

Evidence: mujoco_mjx-3.11.0-py3-none-any.whl

Tags

More Scientific/Engineering packages