One Kernel, Two Thousand Worlds

MJWarp runs 2,048 parallel robot simulations on a single GPU call. Here's what that means for how machines learn to move.

By The ai-compute Desk

Somewhere inside a data center, an SO-101 robotic arm is reaching for a cube—as many as 2,048 parallel MJWarp environments running at once, one world containing the arm reaching for a cube and another containing the same arm starting from a slightly different pose, each advancing physics on NVIDIA GPUs with a single call to mjw.step that advances the entire batch. That is the premise of MuJoCo Warp, or MJWarp: not faster simulation of one robot, but the ability to run hundreds or thousands of independent physics worlds in parallel and collect experience at a scale that single-environment simulators cannot approach. [1][2]

MJWarp is a joint project of NVIDIA and Google DeepMind, maintained in the google-deepmind/mujoco_warp GitHub repository. [2] It re-implements MuJoCo's physics pipeline in NVIDIA Warp—a Python framework for writing GPU-accelerated kernels—and places both the model and a batch of independent states directly on NVIDIA GPUs. [1] One call to mjw.step advances the entire batch. The architecture is not a patch on MuJoCo; it is a deliberate redesign of where the work happens.

The distinction matters because robot learning has outgrown the assumptions that classical simulators were built on. Speed for a single environment is no longer the bottleneck—quantity of experience is. When a system needs to see millions of transitions before a policy generalizes, the infrastructure underneath it has to deliver those transitions at scale.

The MuJoCo documentation is precise about what MJWarp trades away to achieve that scale: a simulation step with MJWarp will be less performant than a step with MuJoCo for the same simulation. [2] MuJoCo is optimized for latency—time for one simulation step. MJWarp is optimized for throughput—the total number of simulation steps per unit time. [2] Those are different objectives, and choosing the wrong tool for the wrong task will cost compute either way.

MJWarp scales better than MJX on nearly every workload, in particular complex simulation scenes, according to the official documentation. [2] MJX, MuJoCo's JAX-based batched backend, relies on jax.vmap and jax.pmap to distribute work across CPUs, GPUs, or TPUs. mujoco_warp.step targets NVIDIA GPUs specifically and works well with PyTorch, which matters for teams whose training infrastructure is already PyTorch-native. [2] The CPU rollout option—mujoco.rollout, a Python API for multi-threaded calls to mj_step—can reach high throughput on hardware with fast cores and large thread counts, but applications that require frequent host-to-device transfers, such as reinforcement learning with simulation on CPU and learning on GPU, may be bottlenecked by transfer overhead. [2]

The foundation beneath MJWarp is NVIDIA Warp, a Python framework for writing high-performance, GPU-accelerated kernels. [1] Warp lets developers author statically typed kernels in Python and compiles them for CPU or CUDA execution. [1] The first launch builds and caches a native module; later launches reuse it. [1] In practice this means a robotics engineer writes logic that reads like Python but executes like CUDA, with wp.tid() identifying the point, contact, body, or world owned by the current logical thread. [1] One logical thread handles one point, so the same code scales from two points to millions without introducing GPU terminology into the control flow. [1]

Two further Warp capabilities shape the longer-term roadmap. Warp kernels are differentiable: a wp.Tape records the forward kernel launches made inside its context and replays their adjoints in reverse when backward() is called, which is why teams build differentiable geometry, CFD, and custom physics in Warp, including CAE workflows for simulation and design optimization. [1] Warp also supports deterministic execution, introduced in Warp 1.15: GPU atomics are scheduler-dependent by default, so repeated launches of the same kernel can differ slightly, and the opt-in deterministic modes trade some performance for reproducible ordering in simulation, validation, and regression tests. [1] Neither capability is guaranteed for an entire MJWarp rollout—they are Warp-layer properties, not MJWarp-level guarantees—but they signal the engineering ambitions underneath the project.

Scaling anything across thousands of parallel instances creates its own category of problems. Memory budgets tighten. Buffer overflows that would surface as obvious crashes in a single-world run become silent corruption in a batched one. The engineering discipline required is less about writing kernels and more about knowing exactly how much space each contact and constraint will consume before the first step fires.

MJWarp allocates contact and constraint buffers before stepping, and exceeding those capacities invalidates the affected rollout for verification or benchmarking, even when execution continues with an overflow warning rather than an exception. [1] The parameters nconmax and njmax control how much space is reserved; naconmax is an alternative to nconmax that sets the maximum number of contacts over all worlds—if both nconmax and naconmax are set, nconmax is ignored. [2] The documentation recommends sizing them against the most contact-heavy moment of the task—for pick-and-place, the instant both jaws and the table touch a cube, not the arm hovering in free space. [2]

Memory pressure compounds as batch sizes grow. Sparse Jacobians reduce the footprint for high-degree-of-freedom scenes significantly: in the Aloha clutter benchmark, with nv = 136 and 2048 worlds, the sparse representation of efc.[2] CCD colliders require more memory than primitive colliders, and multiccd requires more memory than CCD, with memory requirements that scale linearly with the number of CCD iterations. [2] MJWarp can scale to scenes with hundreds of degrees of freedom as long as sleeping is enabled and the scene can be partitioned into independent islands with dormant bodies. [2]

The compact solver addresses scenes that are large in total degrees of freedom but small in active ones. It identifies the set of active DOFs for each world, determined from the active islands, compacts them into a contiguous dense workspace of a known maximum size, and executes the Newton constraint solver using GPU-optimized tile operations such as blocked Cholesky factorization. [2] Results scatter back to the global state, freezing the inactive DOFs. [2] By operating on a fixed-size compacted workspace, the solver avoids GPU thread divergence and leverages high-performance tensor operations optimized for fixed tile sizes. [2] Improving performance for single connected mechanisms beyond approximately 60 DoFs remains an active priority. [2]

CUDA graph capture is the performance lever most teams will reach for first. mjw.step comprises many kernel launches; without capture, Warp launches them individually. Capturing the sequence once and replaying it eliminates repeated dispatch overhead. [1][2] GPU launches are asynchronous, so a naive timer measures how fast Python queued work, not how fast the GPU finished it—which means benchmarks that skip the warm-up phase and the synchronization step before and after the timed region are measuring the wrong thing. [1] The NVIDIA tutorial walks the transition from a single CPU world to 2,048 parallel GPU states using the SO-101 arm and a pick-and-place scene, using scaling_study.py to sweep batch sizes and print milliseconds-per-step alongside aggregate throughput. [1]

The migration path from vanilla MuJoCo to MJWarp is smaller than it appears. MuJoCo loads and compiles the MJCF model; MJWarp implements the physics in NVIDIA Warp, which compiles CUDA kernels to advance simulation states on NVIDIA GPUs. [1] The core API transition involves swapping mjw.make_data() for a fresh state or mjw.put_data() to carry an initialized MuJoCo state across the boundary, then calling mjw.put_model() as a compatibility check—it raises if the model uses unsupported features rather than silently dropping them. [1] Every device array carries a leading world dimension; scaling to thousands of worlds later changes only that leading dimension, not the calls. [1]

What the architecture ultimately describes is a new relationship between software abstraction and physical scale. Writing a kernel that advances one point under gravity and a kernel that advances two thousand robot arms in parallel are, in this framework, the same act. The complexity is managed at the infrastructure layer so that researchers can stay in the domain of physics and policy.

MJWarp is well suited for applications where large numbers of samples are required, like reinforcement learning, while MuJoCo is likely more useful for real-time applications like online control or interactive graphical interfaces. [1] The bet is that the gap between simulated experience and physical behavior closes fastest not by refining any single world, but by running more of them.