Section 17.4: MJX/Brax-training and JAX RL

"When physics and gradients compile together, the boundary between simulation and learning disappears into a single kernel."

A JIT-Compiled Robot Physicist

When ETH Zurich's RSL group trained an ANYmal quadruped to walk in minutes (Rudin et al., 2022), the result was reproducible only because the MJCF contact model, the PPO clip ratio, the foot-clearance and joint-torque reward terms, and the terrain-and-pose reset distribution were all pinned in one training artifact. Change the friction coefficient on a Brax flat-ground reset, or let the velocity-tracking reward drift, and the same compiled rollout produces a policy that walks in sim but face-plants on the real ANYmal-C. MJX and Brax make these four pieces (simulator fidelity, PPO rollout semantics, reward terms, reset distribution) part of a single versioned, device-resident kernel rather than four separately-tuned scripts.

The JAX-native training contract for MJX and Brax rests on one mindset shift. JAX is a Python library for high-performance numerical computing that compiles array operations to GPU or TPU kernels and provides automatic differentiation; here it lets physics and learning share one compiled program. As Figure 17.4A illustrates, the shift is to think in batched arrays first: fixed shapes, explicit random keys, and rollout buffers sized before the first compile. The moving parts are vectorized environments, jit boundaries, vmap or batched stepping, lax.scan rollout collection, and explicit pseudo-random number generator (PRNG) key splitting. Concretely: jit is the JAX decorator that traces a Python function once and compiles it to a fused GPU or TPU kernel; vmap takes a function written for one environment and automatically maps it over a batch dimension without a Python loop; and lax.scan is the compiled equivalent of a for-loop, carrying state across steps while writing each step's output into a pre-allocated buffer. Figure 17.4B lays out how these pieces connect inside a single compiled kernel, from PRNG splitting through the scan-based rollout to the PPO update that feeds back into the next iteration.

The key question is practical: can the rollout loop be compiled once, run many times with fixed tensor shapes, and still produce independent training and evaluation episodes?

This approach buys throughput and reproducibility. A conventional Python RL loop that calls a C++ simulator wastes a large fraction of its wall-clock time on Python overhead, host-device data copies, and sequential environment stepping. With 8,192 environments on a GPU, each Python call becomes a bottleneck. MJX and Brax remove that bottleneck: they express physics and RL as JAX array transformations that compile into a single device kernel. No Python loop runs in the hot path, no host-to-device copy happens per step, and the rollout stays entirely on the accelerator until the update finishes. To make that concrete: in practice, teams report a locomotion curriculum that takes on the order of 14 hours with a sequential Python-MuJoCo loop finishing in under 45 minutes with a compiled MJX rollout on the same GPU, roughly a 20x wall-clock reduction from one architectural decision; the exact ratio depends on horizon length, environment count, and how much of the original loop's time was Python overhead versus physics computation itself. Put another way, a hyperparameter search that would have required 50,000 sequential episodes to cover the same wall-clock budget can fit in roughly 300 compiled rollouts under favorable conditions, because each rollout steps all 8,192 environments in parallel inside a single device kernel. Gains of this magnitude are typical for throughput-bound locomotion workloads rather than guaranteed for every task; this generally changes what research is feasible to iterate on.

Theory

MJX provides MuJoCo-style physics through JAX arrays, while Brax provides a JAX-native physics and RL stack designed for massively parallel simulation. In both cases, the important unit is not an object-oriented environment instance; it is a batched simulator state transformed by pure functions.

A typical rollout uses the single-kernel training loop: a compiled function mapping policy parameters, simulator state, and random keys to new state and experience tensors. Before training begins, the builder must pin down three things: which dimensions are static, which arrays are donated (marked so JAX can overwrite their device memory in place instead of allocating a fresh buffer) or reused, and how many bytes the rollout buffer consumes.

Choose MJX when your task requires MuJoCo's contact model, articulated body semantics, or existing MJCF assets. MJX preserves those modeling concepts while making the step function JAX-compatible. Choose Brax when you prototype from scratch and want the full RL environment API (reset, step, reward, done) in a single JAX-native package with less setup friction. MJX costs you coverage: the JAX backend does not support every MuJoCo feature. Brax costs you fidelity: its simplified contact model may not match hardware behavior when precise contact forces matter. Teams building locomotion policies often start with Brax for fast iteration, then switch to MJX when they need higher-fidelity contact dynamics for sim-to-real transfer. A policy that runs perfectly in simulation but collapses on hardware is not a trained policy; it is a well-compiled guess.

Worked Example

Code Fragment 17.4.1 estimates the rollout buffer before writing a JAX training loop. This is the first sanity check for whether a planned MJX or Brax experiment fits on the accelerator.

# Estimate device memory for a JAX-native rollout buffer.
# Static shapes should be chosen before jit compilation and training launch.
num_envs = 8192
horizon = 32
obs_dim = 64
act_dim = 12
float_bytes = 4

buffers = {
    "observations": horizon * num_envs * obs_dim * float_bytes,
    "actions": horizon * num_envs * act_dim * float_bytes,
    "rewards": horizon * num_envs * float_bytes,
    "dones": horizon * num_envs * float_bytes,
    "values": horizon * num_envs * float_bytes,
}

total_mb = sum(buffers.values()) / 1_000_000
for name, size in buffers.items():
    print(f"{name}: {size / 1_000_000:.1f} MB")
print(f"rollout buffer total: {total_mb:.1f} MB")

Expected output: the trace should expose which tensor dominates memory. A JAX run that recompiles or spills buffers because shapes were guessed late will lose the throughput advantage this section is trying to teach.

Before reading on, guess: if a sequential Python-MuJoCo loop requires 14 hours to train a quadruped locomotion policy, how much of that time is JAX compilation overhead versus actual physics steps on the GPU?

The memory estimate above sizes the buffer; the following skeleton is the compiled rollout itself, the piece this section's title promises and that the Algorithm box below formalizes. It shows the actual wiring of vmap, lax.scan, and jit around a policy and an MJX or Brax environment, so the shapes and calls named earlier have one concrete home before the abstract algorithm.

# Minimal compiled rollout skeleton (illustrative, not a runnable full trainer).
# vmap parallelizes reset/step across environments; lax.scan collects the horizon;
# jit fuses both into one device-resident kernel.
import jax
from jax import lax

def rollout(params, env, key, num_envs, horizon):
    reset_keys = jax.random.split(key, num_envs)
    state = jax.vmap(env.reset)(reset_keys)

    def scan_step(carry, _):
        state, step_key = carry
        step_key, action_key = jax.random.split(step_key)
        actions = policy_apply(params, state.obs, action_key)  # vmapped internally
        next_state = jax.vmap(env.step)(state, actions)
        done_mask = 1.0 - next_state.done
        reward_masked = next_state.reward * done_mask
        return (next_state, step_key), (state.obs, actions, reward_masked, next_state.done)

    (final_state, _), trajectory = lax.scan(
        scan_step, (state, key), xs=None, length=horizon
    )
    return trajectory  # (obs, actions, rewards, dones), each shape (horizon, num_envs, ...)

compiled_rollout = jax.jit(rollout, static_argnums=(3, 4))

Practical Recipe

1. Choose static shapes for environment count, horizon, observation groups, action dimension, and privileged critic state.
2. Split PRNG keys per environment and per rollout step so randomization does not synchronize.
3. Put reset logic into compiled array operations where possible.
4. Profile compile time separately from steady-state training throughput.
5. Evaluate on a separate key stream and a separate task panel, then save both key roots.

The idea in this section becomes useful when accelerator constraints are treated as part of the algorithm. A JAX RL loop is fast because it exposes the whole rollout as a regular computation graph. That same regularity means dynamic episode bookkeeping, variable-size observations, and host-side logging must be designed carefully.

The graduate-level habit is to report compile time, steady-state steps per second, memory footprint, shape choices, and key schedule separately. Otherwise a result may confuse a better algorithm with a better compilation path.

Knowing which quantities to report separately is only half the discipline; the table below names the specific tools that produce each of them, so you can map every reporting field to the library call responsible for it.

Once those tools are chosen, the next step is to pin down how they will be configured for a given run, and that record is what a manifest captures. A robust implementation starts with a JAX run manifest. It records the shapes that must stay static, the keys that drive randomness, and the metrics that separate compilation overhead from steady-state throughput.

1. Freeze batch, horizon, observation, action, and privileged-state shapes before the first benchmark.
2. Record root PRNG keys for training, curriculum randomization, and evaluation separately.
3. Measure first-call compile time and post-compile steady-state throughput as different fields.
4. Save memory estimates for rollout buffers and optimizer state.
5. Compare algorithms only when they share the same static shapes and held-out key stream.

# Record the JAX-specific fields that make an MJX or Brax run auditable.
# Static shapes and separate PRNG roots prevent hidden recompilation and leakage.
from dataclasses import dataclass, asdict

@dataclass
class JaxRlManifest:
    simulator: str
    batch_shape: tuple[int, int]
    obs_dim: int
    action_dim: int
    train_key_root: int
    eval_key_root: int

    def as_row(self) -> dict[str, object]:
        return asdict(self)

manifest = JaxRlManifest(
    simulator="MJX",
    batch_shape=(32, 8192),
    obs_dim=64,
    action_dim=12,
    train_key_root=1704001,
    eval_key_root=1704999,
)
print(manifest.as_row())

When a JAX-native run fails, separate numerical, compilation, and RL causes. Check NaNs, dtype choices, action clipping, reset masks, key reuse, recompilation counts, and host-device transfers before blaming the policy architecture.