How Does NVIDIA Move a Robot Arm Simulation From One CPU Core to 2,048 GPU Worlds?

NVIDIA Warp is a Python framework that compiles kernels for GPU execution. MuJoCo Warp builds on it to run thousands of independent MuJoCo physics scenes at once on NVIDIA GPUs.

The Stack and the Shift

Classic MuJoCo runs fast on CPU cores. It parallelizes sampling across those cores. But reinforcement learning and large-scale sampling need something different: not lower latency for one world, but higher aggregate throughput across many. Aggregate throughput means total world-steps completed per wall-clock second. MJWarp delivers that by placing a compatible MuJoCo model and a batch of independent states on the GPU. One call to mjw.step advances the entire batch.

The SO-101 follower arm serves as the test case. The scene contains the arm, a table, and two 44 mm cubes to stack. The task: grasp the red cube and place it on the blue one. The same MJCF model runs on both backends. The migration follows four gates: establish a CPU baseline, validate one-world parity on GPU, size contact and constraint buffers, then scale to 2,048 worlds.

Establishing the CPU Baseline

The scene loads from scene_pick_place.xml. For a MuJoCo box, the size values are half-extents. size="0.022 0.022 0.022" defines a cube with 44 mm edges. The arm base sits at the origin. Its reach extends along +X. The cubes sit along Y.

Control runs at 50 Hz. Physics substeps run at 10 per control frame. That yields a physics timestep of 0.002 seconds. The line mjm.opt.timestep = frame_dt / sim_substeps must execute before the CPU rollout and before uploading the model with mjw.put_model. Without it, every later measurement inherits the mismatch.

Success is not a clean process exit. It is two measurable conditions evaluated after the cubes settle. Horizontal center error between cube centers must be xy_err ≤ 0.015 m. Vertical separation between cube centers must be 0.035 m ≤ dz ≤ 0.055 m. The companion repository prints these numbers at the end of the run. That printout is the assertion the rest of the article compares against.

Validating One World on the GPU

Gate 2 runs one world on the GPU with the host still in the loop. The model uploads via m = mjw.put_model(mjm). This call doubles as a compatibility check: it raises if the model uses unsupported features rather than silently dropping them.

Batched state allocates with d = mjw.make_data(mjm, nworld=1, nconmax=spec.nconmax, njmax=spec.njmax). The SO-101 profile uses nconmax=128 and njmax=300 as starting capacities. Host state seeds the device arrays. Every device array carries a leading world dimension. That is why mjd.qpos[None, :] has shape (1, nq) instead of (nq,).

The frame loop mirrors the CPU version. Controls compute on the host. They copy to d.ctrl each substep. mjw.step(m, d) advances the GPU world. Results copy back. mujoco.mj_forward(mjm, mjd) refreshes derived host quantities like mjd.xpos before the stack check. Reading those fields after the loop does not refresh them automatically. This path validates the task. It is not a throughput benchmark. The .numpy() reads synchronize and copy data to the host on every substep.

Sizing Buffers Before Scaling

MJWarp allocates contact and constraint buffers before stepping. Exceeding those capacities invalidates the affected rollout for verification or benchmarking. Execution continues with an overflow warning rather than an exception. The warning prints the budget to increase: "narrowphase overflow - please increase nconmax to …". It also flags affected worlds in Data.overflow.

Larger buffers use more GPU memory. The rule: verify capacity over the full task before tightening the allocation. Size them against the most contact-heavy moment. For pick-and-place, that is the instant both jaws and the table touch a cube. Not the arm hovering in free space.

The tool mjwarp-testspeed --measure_alloc reports actual contact and constraint consumption. It aborts the rollout with offending world IDs as soon as any world overflows. Treat those reports as failures. Raise the limit and re-run before trusting either the trajectory or the benchmark. Then tighten again whenever the model, collision geometry, or task changes.

Scaling to 2,048 Worlds

Scaling changes two things relative to Gate 2: nworld, and the fact that nothing crosses the PCIe bus per step.

nworld = 2_048
d = mjw.make_data(mjm, nworld=nworld, nconmax=spec.nconmax, njmax=spec.njmax)
wp.copy(d.qpos, wp.array(np.tile(mjd.qpos, (nworld, 1)), dtype=wp.float32, device=device))
wp.copy(d.qvel, wp.array(np.tile(mjd.qvel, (nworld, 1)), dtype=wp.float32,