Build1 publisher3 min readPublished
MJWarp scales one SO-101 arm to 2,048 worlds by changing what the thread index means
NVIDIA's walkthrough moves a MuJoCo follower arm up to 2,048 parallel GPU environments, where the kernel body barely changes and most of the work is model compatibility and an opt-in determinism mode.
The Engineer · Build desk
What happened
- NVIDIA's walkthrough takes an SO-101 follower arm out of a familiar MuJoCo workflow and up to as many as 2,048 parallel MJWarp environments, covering the validation steps for the move.
- MuJoCo still loads and compiles the MJCF model, while MJWarp implements the physics in NVIDIA Warp, which compiles the CUDA kernels that advance simulation states on the GPU.
- Warp's GPU atomics are scheduler-dependent by default, so repeated launches of one kernel can differ slightly, and Warp 1.15 added opt-in deterministic modes that trade some performance for reproducible ordering.
- Calling .numpy() on a CUDA array synchronizes and copies it to CPU memory, and device-resident PyTorch or JAX pipelines are pointed at Warp's framework adapters or DLPack sharing instead.
- The post is the second in NVIDIA's State of Simulation for Physical AI series; it prepares and scales the environment without training a policy, and leaves Newton and Isaac Lab to later installments.
Compiled by The EngineerSomething wrong?How this is made
Why it matters
- constraint A model MuJoCo compiles today is not automatically one MJWarp will run, so the schedule risk on a port sits in model preparation.
- decision Someone on the team has to pick the tolerance within which a GPU trajectory must match the CPU one, because NVIDIA declines to promise determinism across an entire MJWarp rollout.
- cost Compilation is paid on the first launch and cached afterwards, so every kernel edit means a rebuild before the next run. That cost lands on iteration time.
- precedent With the integration layers deferred to the Newton and Isaac Lab installments, a team adopting MJWarp now writes its own connection from these environments to a learning loop.
Warp's kernel language is a performance-oriented subset of Python, and everything around it stays ordinary Python: configuration, allocation, launch orchestration [5][7].
The worked example advances point positions under gravity. It adds a vector with z of -9.81 times dt to velocity, then adds velocity times dt to position. dt is 0.01, and a start array holds two points at z = 0.5 [17]. Read the order of the two lines. Velocity updates first, so the first step moves the point by the velocity it has just acquired: -9.81 * 0.01 gives -0.0981, and 0.5 - 0.0981 * 0.01 leaves z at 0.499019 [18]. The launch dimension is the length of that array, so two logical threads run [19].
wp.tid() returns the index of the point, contact, body or world that the current logical thread owns [8]. In the article a world is one independent copy of the scene and its state. The SO-101 arm reaches for a cube in one; in another, the same arm starts from a slightly different pose [15]. NVIDIA's claim for the small kernel is that one thread per point scales from two points to millions without introducing GPU terminology into the control flow [20].
I would treat 2,048 as a property of that arm and that scene. Two conditions have to hold before the number transfers to another robot. The model has to sit in MJWarp's compatible set. MJWarp is described as the path for compatible MuJoCo models, and the article's framing is that MuJoCo and MJWarp run the same compatible robot and task while organizing the work differently [1][16]. Per-world step cost then has to be close to the SO-101's. NVIDIA gives no GPU and no steps-per-second figure for the 2,048-world configuration [21].
Validation is where I would put the time on a port. NVIDIA writes that differentiability and deterministic execution are Warp capabilities, "not guarantees of differentiability or determinism for an entire MJWarp rollout" [13]. Read that before writing a regression test that compares GPU trajectories bit for bit against the CPU model. The tolerance is yours to choose, and the version pin is already in the install line: warp-lang 1.15 or later if you want GPU determinism at all [14].
Graph capture only helps with dispatch overhead. It replays launches against existing buffers to cut repeated dispatch cost, and it does not fuse arbitrary kernels [10].
The two Warp features most likely to interest a robotics team go unused in this workflow. A wp.Tape records the forward kernel launches made inside its context and replays their adjoints in reverse when backward() is called. That is how teams build differentiable geometry, CFD and custom physics in Warp, including CAE workflows for simulation and design optimization [11].
What to watch
- The Newton and Isaac Lab installments, where the training integration layers are supposed to land.
- Whether deterministic execution ever extends from individual Warp kernels to a full MJWarp rollout.
- Which MuJoCo model features stay outside MJWarp's compatible set as the port matures.