NVIDIA Warp and MJWarp Run 2,048 Parallel Robot Simulations on One GPU

NVIDIA has published the second article in its State of Simulation for Physical AI series. It walks through MuJoCo Warp (MJWarp), a GPU implementation of MuJoCo physics built on the NVIDIA Warp kernel framework, and moves an SO-101 follower arm from a standard CPU workflow to as many as 2,048 parallel simulation worlds on one GPU. The question, the authors argue, is no longer how fast one world advances but how many worlds can advance at once.

Warp is a Python framework for writing high-performance, GPU-accelerated kernels. Developers author statically typed kernels in Python; Warp compiles them for CPU or CUDA and caches the result. The pitch combines CUDA-level speed through just-in-time compilation and CUDA Graphs, built-in vectors, matrices, quaternions and hash grids, and differentiable kernels with DLPack-style interop, so simulation can sit inside a training loop.

How MJWarp batches physics

MJWarp keeps the model format: an MJCF file loads in MuJoCo, and mjw.put_model uploads it to the device, where one mjw.step call advances a whole batch of independent worlds. Four parameters set capacity: nworld, the number of parallel environments; nconmax, expected contacts per world; naconmax, a global contact limit that takes precedence when defined; and njmax, the per-world constraint ceiling. The gain is aggregate throughput rather than single-world latency, which suits reinforcement learning and large-scale sampling.

Validate, then measure

The walkthrough proceeds through checks: build a CPU baseline, move one world to the GPU and confirm parity, size the contact and constraint buffers against the most contact-heavy moment of the task, then reallocate to 2,048 worlds. The example is pick-and-place: an SO-101 arm stacks a red 44 mm cube on a blue cube, judged by a horizontal centre error of 15 mm or less and a vertical separation of 35 to 55 mm. An overflow is reported rather than raised, so a rollout that runs past one is unreliable for benchmarking.

Measurement is the last gate. GPU launches are asynchronous, so a naive timer measures how fast Python queued work, not how fast the GPU finished; the article warms up first and synchronizes around the timed region. Results are quoted in world-steps per second and milliseconds per batched step. Warp 1.15 added opt-in deterministic execution, and a later instalment will port the same environment to Newton.