Section 11.3: MuJoCo MJX and MuJoCo Warp

"One robot teaches you a behavior. Ten thousand robots teach you the distribution."

A Batch-First Embodied AI Agent

The Problem: One Environment Is Not Enough

Picture a locomotion policy that walks flawlessly in your simulator, then topples the instant the real floor is 10% more slippery than the one virtual world it ever saw, exactly the failure mode Figure 11.3A illustrates: the fix is not a better single rollout but thousands of them, and that is why modern robot learning needs a population of worlds instead of one beautiful rollout. Reinforcement learning needs many episodes to estimate gradients. Domain randomization needs many variations to expose brittle policies. System identification needs parameter sweeps. This changes the simulator question from "Can I step this robot?" to "Can I step a population of worlds while keeping the data on the accelerator?" A single-world baseline might require 48 hours to collect enough experience for a stable locomotion policy; the same policy trained across 4,096 parallel MJX worlds collects equivalent experience in roughly 45 minutes (as of 2024, on representative mid-range GPU hardware), because the batch steps all environments in a single GPU kernel call rather than one after another.

MJX: MuJoCo Through JAX

MJX is useful when the rest of the learning stack is written in JAX or when you want jit, vmap, and accelerator execution to organize simulation at scale. This pattern is called batch-first simulation design, and it means you express the parallelism at the data level rather than through explicit loops. The mental model has two parts: keep the model structure close to MuJoCo, but represent batched simulation state as arrays that JAX can transform. Figure 11.3B diagrams how a single model fans out across the MJX and MuJoCo Warp backends before converging on a shared comparison step.

The sharp edge is that MJX earns its keep on large batches of similar worlds. It is not automatically the best route for a single interactive scene, a debugging session, or a model feature that is not supported by the accelerator backend. Keep a small CPU MuJoCo case nearby, then use MJX when the experiment really needs thousands of parallel states, domain-randomized parameters, or JAX-native policy updates. MJX compiles simulation steps through XLA (Accelerated Linear Algebra, the compiler that JAX uses to turn array operations into fused GPU or TPU kernels), which is why a batched step executes as one kernel launch rather than thousands of small ones.

Code Fragment 1 does not require MJX. It teaches the vectorization pattern MJX makes valuable: one equation runs across many worlds with different parameters.

# Vectorized rollout: one update rule runs across many worlds.
# NumPy stands in for JAX here so the batching idea is easy to inspect.
# Each column represents a different gravity setting.
import numpy as np

heights = np.full(4, 0.25)
velocities = np.zeros(4)
gravities = np.array([-9.60, -9.70, -9.81, -10.00])
dt = 0.02

for _ in range(10):
    velocities = velocities + gravities * dt
    heights = np.maximum(0.0, heights + velocities * dt)

print(np.round(heights, 4))

[0.0388 0.0366 0.0342 0.03  ]

MuJoCo Warp: MuJoCo On NVIDIA Warp

MuJoCo Warp, often written MJWarp, targets NVIDIA hardware through Warp. Where MJX makes simulation feel like JAX, MJWarp brings MuJoCo-style physics into a GPU programming model built for parallel kernels. Reach for it when raw throughput on NVIDIA hardware is the bottleneck, or when the stack connects to Newton (NVIDIA's open-source GPU physics engine built on Warp, aimed at large-scale robot learning rather than the JAX ecosystem MJX targets).

For real-robot deployment, raw throughput is not just a training convenience. Sim-to-real transfer quality depends on how many distinct physical configurations a policy sees before hardware deployment. Higher throughput per GPU-hour lets you sweep more friction coefficients, more link masses, and more terrain profiles before you freeze the policy. A policy underexposed to variation will degrade on the first floor surface that differs from the simulation mean. A policy that has seen only one world is not robust; it merely memorizes a response to a single set of constants. In practice, a Unitree Go2 locomotion policy trained across only 256 parallel worlds typically requires on the order of tens of thousands of episodes before friction generalization stabilizes; the same policy trained across 4,096 worlds typically reaches comparable generalization in a couple hundred episodes on representative training runs, because each batch already contains the full friction distribution rather than a narrow slice of it. These figures are illustrative order-of-magnitude estimates rather than a benchmarked result, and the exact ratio depends on the task, the randomization ranges, and the policy architecture.

How Warp Reaches That Throughput

Warp achieves its throughput by expressing physics kernels directly in a compiled GPU dialect rather than routing through JAX's XLA compiler. Each simulation step launches CUDA kernels that run contact resolution, constraint solving, and state integration entirely on-device, with no Python-level iteration and no CPU-GPU data transfers between steps. MuJoCo Warp maps each environment in the batch to a thread block (a group of GPU threads that execute together and share fast on-chip memory, the standard unit of work scheduling in CUDA-style programming), so contact events in world 1,000 resolve in parallel with those in world 1,001 rather than sequentially.

Differentiability Is Powerful, But Not Magic

Beyond running thousands of worlds in parallel, the JAX-native backend unlocks a second capability that pure throughput engines like Warp do not emphasize: gradients that flow through the physics itself.

Differentiable simulation lets you propagate gradients from a future cost back through physics to motor torques or model parameters. In a concrete locomotion setting, you estimate how a Unitree H1's ankle stiffness affects cumulative foot-strike forces 200 ms ahead, then update that stiffness in one gradient step instead of running a separate parameter sweep. For a Franka Panda grasping task, you differentiate through a 5 mm fingertip contact patch to find the wrist orientation that maximizes grip stability before you attempt the grasp. Contact is where this breaks down. A foot leaving the ground or a finger snapping to a new contact face introduces a step discontinuity, and that discontinuity makes the gradient meaningless or misleading at that instant. Treat simulator gradients as useful local signals for smooth trajectory segments, not as reliable guides through contact events.

A robust workflow tests gradients the same way it tests speed. Compare automatic gradients against finite differences on a tiny scene, perturb contact parameters, and replay closed-loop rollouts after any backend change. If the gradient improves a loss but breaks the rollout under a small friction change, the optimization found a simulator artifact rather than a control insight.

Practical Recipe

1. Choose MJX when the rest of your learner is JAX-first and batching is central.
2. Choose MuJoCo Warp when NVIDIA GPU throughput and Warp or Newton integration matter.
3. Keep a CPU MuJoCo sanity case for small examples and debugging.
4. Run the same seed, model, and metric across backends before comparing speed.
5. Report throughput and behavior, not throughput alone.

Expected output: An accelerator-backend comparison should report model identity, backend versions, batch size, seeds, device, throughput, rollout metric, and a small CPU MuJoCo sanity trace. Speed without behavior matching is not enough evidence.