Section 17.4: MJX/Brax-training and JAX RL

A Careful Control Loop

For MJX/Brax-training and JAX RL, GPU RL depends on simulator fidelity, PPO rollout semantics, reward terms, and reset distribution being versioned in the same training artifact.

This section develops the JAX-native training contract for MJX and Brax. We focus on vectorized environments, jit boundaries, vmap or batched stepping, lax.scan rollout collection, and explicit PRNG splitting.

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?

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 a compiled function that maps policy parameters, simulator state, and random keys to new state and experience tensors. The builder must know which dimensions are static, which arrays are donated or reused, and how many bytes the rollout buffer consumes before training begins.

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.

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.

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.