MuJoCo‘s classic CPU path is fast for one world and parallelizes across cores, but learning workloads shift the question from how quickly one world runs to how many worlds run at once. MuJoCo Warp (MJWarp), built on NVIDIA Warp, takes compatible MuJoCo models into the GPU-batched regime, advancing hundreds or thousands of independent environments in a single mjw.step call. This walkthrough moves an SO-101 follower arm from a standard MuJoCo workflow to 2,048 parallel MJWarp environments, with explicit validation gates along the way. It is the second post in the State of Simulation for Physical AI series and does not train a policy.

NVIDIA Warp is a Python framework for authoring statically typed, GPU-accelerated kernels. Kernels are written in a performance-oriented Python subset, JIT-compiled for CPU or CUDA, and cached after the first launch. Three properties make it useful for robotics: explicit parallel work via wp.tid(), which identifies the point, contact, body, or world owned by a logical thread; explicit device arrays, where .numpy() on a CUDA array synchronizes and copies to host rather than zero-copy sharing, so framework adapters or DLPack-compatible sharing are preferred for device-resident PyTorch/JAX pipelines; and composable kernel launches that can be captured into CUDA Graphs. Warp kernels are also differentiable via wp.Tape, and Warp 1.15 added opt-in deterministic execution, though neither is used in this workflow and neither is guaranteed for a full MJWarp rollout.

The core API transition is small: mujoco.MjModel maps to mjw.put_model(mjm), which also raises on unsupported features; mujoco.MjData maps to mjw.make_data() for fresh state or mjw.put_data() when the exact initialized MuJoCo state must cross the boundary; mujoco.mj_step becomes mjw.step(m, d), advancing every world in the batch. Every device array carries a leading world dimension, so host state is indexed as mjd.qpos[None, :] with shape (1, nq); scaling to thousands of worlds changes only that leading dimension. Allocation requires sizing nworld, nconmax (expected contacts per world), naconmax (global contact maximum, taking precedence if both are set), and njmax (constraint limit per world).

The migration proceeds through gates. First, establish a CPU baseline: an ordinary MJCF scene with the SO-101 arm, a table, and two 44 mm cubes (MJCF box sizes are half-extents), with a 50 Hz controller rate, 10 physics substeps per frame, and a 0.002 s physics timestep set before both the CPU rollout and mjw.put_model so both backends advance the same simulated time. Task success is two measurable conditions after settling: horizontal center error xy_err ≤ 0.015 m and vertical separation 0.035 m ≤ dz ≤ 0.055 m. Second, validate one-world parity on GPU with the host still in the loop, copying qpos, qvel, and ctrl to the device, stepping, and copying results back each substep. This is a task-validation path, not a throughput benchmark; after copying, call mujoco.mj_forward to refresh derived host quantities like mjd.xpos. Third, size contact and constraint buffers against the most contact-heavy moment of the task — for pick-and-place, when both jaws and the table touch a cube. The SO-101 profile starts at nconmax=128 and njmax=300; overflows print a budget to increase rather than raising, flag affected worlds in Data.overflow, and invalidate the rollout, so treat them as failures and re-run.

Fourth, scale to 2,048 worlds: reallocate at target size, replicate the initialized state across the batch with np.tile, and capture the step into a CUDA Graph with wp.ScopedCapture for replay. CUDA Graphs reuse captured model and data buffers; update d.ctrl in place between replays and capture a new graph after replacing buffers, changing nworld, or rebuilding the model. Finally, measure correctly: GPU launches are asynchronous, so warm up for compilation and allocation, synchronize immediately before and after the timed region, and report aggregate world-steps per second alongside milliseconds per batched step and batch size. MJWarp’s value is not necessarily a faster step for one world; it is aggregate throughput across the batch, which favors reinforcement learning and large-scale sampling over single-environment latency. A one-world latency comparison does not establish batched throughput, and results depend on scene, settings, and hardware. The article closes with a decision shortcut: single-robot MPC or teleop on MuJoCo CPU; max raw-physics throughput on MJWarp or mjlab; JAX recipes via MuJoCo Playground or MJX with impl=’warp’; and Newton (covered in the next post) for multi-solver plus Isaac Lab integration, with MJWarp running underneath as newton.solvers.SolverMuJoCo. The companion repository includes scaling_study.py to sweep batch sizes, and a publication note confirms the repo URL and pinned dependency versions must be verified before publishing.