mujorax 0.1.6

Mujorax is a lightweight open-source JAX-native MuJoCo environment suite for single-agent Reinforcement Learning (RL), built on top of Envrax. It wraps MuJoCo Playground environments with Envrax's JaxEnv so you can use them with envrax.make, envrax.make_vec, and the rest of Envrax's tooling.

It comes with 25 environments from the DM Control Suite. All environment logic follows a stateless functional design that builds on top of the MJX, JAX, and Chex packages to benefit from JAX accelerator efficiency.

Why Mujorax?

Envrax provides a JAX-native Gymnasium-style API standard for RL environments, but it doesn't ship with any environments of its own. One of the biggest spaces in RL is robotics, and the gold-standard physics engine for this is MuJoCo. This makes it the perfect fit for one of the first Envrax environment suites!

MuJoCo Playground is Google DeepMind's open-source library of MuJoCo environments, built on top of MJX (MuJoCo's JAX port that preserves the simulator's full physics fidelity). It already solves the hard parts: research-validated reward and termination logic for DM Control, locomotion, and manipulation environments. The only catch is that its environments expose a Brax-style MjxEnv API, which doesn't quite fit Envrax's API standard.

Rather than reinventing the wheel, Mujorax acts as a thin, type-safe wrapper around the MuJoCo Playground environments to maximise their benefits while maintaining Envrax's API standard, making it completely plug-and-play with Envrax's toolkit.

Requirements

• Python 3.13+
• JAX 0.9+ (CPU, CUDA, or TPU backend)

Installation

pip install mujorax

Or with uv:

uv add mujorax

Quick Start

import jax
import mujorax  # registers the suite at import
import envrax

env = envrax.make("mjx/cartpole_balance-v0")
obs, state = env.reset(jax.random.PRNGKey(0))

action = env.action_space.sample(jax.random.PRNGKey(1))
obs, state, reward, done, info = env.step(state, action)

Vectorised rollouts work the same way:

env = envrax.make_vec("mjx/cartpole_balance-v0", n_envs=128)
obs, state = env.reset(jax.random.PRNGKey(0))  # obs.shape == (128, 5)

You can also use make_multi to utilise several heterogeneous environments at once:

env = envrax.make_multi([
    "mjx/cartpole_balance-v0",
    "mjx/cheetah_run-v0",
])
obs_list, state_list = env.reset(jax.random.PRNGKey(0))  # one entry per env

Or, the make_multi_vec method for vectorised parallel copies of each environment:

env = envrax.make_multi_vec(
    ["mjx/cartpole_balance-v0", "mjx/cheetah_run-v0"],
    n_envs=64,
)
obs_list, state_list = env.reset(jax.random.PRNGKey(0))
# each entries obs.shape == (64, *single_obs.shape)

Mujorax auto-detects whether a CUDA backend is available; on CPU-only systems it transparently falls back to MJX's pure-JAX physics implementation.

You can override this choice through this MjxPlaygroundConfig(config_overrides={"impl": ...}) if needed.

Environments

All environments share canonical IDs in the form mjx/<name>-v0. Here's the full list of supported environments:

Acknowledgements

Mujorax wouldn't be possible without these incredible projects:

• MuJoCo Playground — the underlying environment implementations.
• MuJoCo and MJX — the physics engine and JAX bindings.
• Envrax — the registry and base environment API standard.

❤️ Thank you to all the developers involved - you guys are awesome! ❤️