NVIDIA Warp and MuJoCo Warp (MJWarp) Enable 2,048‑World GPU Simulation for Robotics

TL;DR

NVIDIA Warp and its MuJoCo‑compatible layer MJWarp let you batch thousands of robot simulation worlds on a single GPU, achieving orders‑of‑magnitude higher aggregate throughput than classic CPU‑based MuJoCo.

Stack Overview

Layer Role
NVIDIA Warp Python kernel language that compiles to CUDA, provides autodiff, and interoperates with PyTorch/JAX
MJWarp GPU implementation of MuJoCo physics that preserves the MJCF format and supports batched stepping
Your scene (SO‑101) Standard Menagerie / Robot Studio assets and task geometry
Next (Newton / Isaac Lab) Multi‑solver API, USD assets, sensors, managers, and training loops

When to use each component

Need Use
Single‑robot MPC or teleoperation Classic MuJoCo CPU
Maximum raw MuJoCo physics throughput MJWarp (or mjlab)
JAX‑based training recipes MuJoCo Playground / MJX (implementation = warp)
Multi‑solver + Isaac Lab integration Newton (future post)

NVIDIA Warp: A Minimal GPU Kernel

Warp lets you write statically‑typed Python kernels that JIT‑compile to native CUDA. The example below integrates point positions under gravity, scaling from two points to millions without changing the kernel logic.

import numpy as np, warp as wp

@wp.kernel
def integrate(positions: wp.array[wp.vec3],
                  velocities: wp.array[wp.vec3],
                  dt: float):
    i = wp.tid()
    velocities[i] += wp.vec3(0.0, 0.0, -9.81) * dt
    positions[i] += velocities[i] * dt

wp.init()
device = "cuda:0" if wp.is_cuda_available() else "cpu"
start = np.array([[0.0, 0.0, 0.5], [0.2, 0.0, 0.5]], dtype=np.float32)
positions = wp.array(start, dtype=wp.vec3, device=device)
velocities = wp.zeros_like(positions)

wp.launch(integrate, dim=len(start), inputs=[positions, velocities, 0.01], device=device)
wp.synchronize_device(device)
print(positions.numpy())

Why this matters for robotics

  • Explicit parallel work – wp.tid() identifies the logical thread (point, contact, body, or world).
  • Explicit device arrays – Arrays live on the selected device; .numpy() triggers a host copy, while DLPack adapters enable zero‑copy sharing with PyTorch/JAX.
  • Composable launches – Multiple kernels can be captured into a CUDA graph, reducing launch overhead.

Differentiability & Determinism (Warp capabilities)

  • Differentiable kernels – A wp.Tape records forward launches and automatically computes adjoints on backward(). This enables differentiable geometry, CFD, and custom physics pipelines.
  • Deterministic execution – Introduced in Warp 1.15; deterministic mode trades a small performance hit for reproducible atomic ordering, useful for regression testing. These features are optional and not required for the SO‑101 walkthrough.

What Is MuJoCo Warp (MJWarp)?

MJWarp is a GPU‑backed implementation of MuJoCo’s physics pipeline built on Warp. It accepts the same MJCF models, places the model and a batch of independent states on the GPU, and advances the entire batch with a single mjw.step call.

  • Latency vs. Throughput – MJWarp does not guarantee a faster single‑world step; its advantage is aggregate throughput: the total number of world‑steps completed per second when many worlds run in parallel. This is the key metric for reinforcement‑learning and large‑scale sampling.

Migrating a MuJoCo Scene to MJWarp

The guide walks through three stages: establishing a CPU baseline, validating a one‑world GPU rollout, and scaling to 2,048 worlds.

1. Establish a MuJoCo CPU Baseline

  • Load the SO‑101 pick‑and‑place MJCF (standard Menagerie assets) with mujoco.MjModel.from_xml_path.
  • Use a 50 Hz control loop with 10 physics sub‑steps (timestep = 0.002 s).
  • Verify task success by checking horizontal error ≤ 0.015 m and vertical separation between 0.035 m and 0.055 m after the cubes settle.

2. Validate One‑World MJWarp Parity

import warp as wp, mujoco_warp as mjw
wp.init()
device = wp.get_device()

# Upload model and allocate a single‑world batch
m = mjw.put_model(mjm)
d = mjw.make_data(mjm, nworld=1, nconmax=spec.nconmax, njmax=spec.njmax)

# Seed state from the CPU rollout
wp.copy(d.qpos, wp.array(mjd.qpos[None, :], dtype=wp.float32, device=device))
wp.copy(d.qvel, wp.array(mjd.qvel[None, :], dtype=wp.float32, device=device))
wp.copy(d.ctrl, wp.array(mjd.ctrl[None, :], dtype=wp.float32, device=device))

mjw.forward(m, d)
  • The host loop mirrors the CPU workflow, copying controls to the device each sub‑step and pulling back qpos/qvel for downstream calculations.
  • This path is not a performance benchmark; it validates that the GPU simulation reproduces the CPU result.

3. Size Contact & Constraint Buffers

  • MJWarp pre‑allocates contact (nconmax) and constraint (njmax) buffers per world. Insufficient capacity triggers overflow warnings and may corrupt trajectories.
  • Use mjwarp-testspeed --measure_alloc to measure actual usage and raise the limits accordingly (e.g., nconmax=128, njmax=300 for the SO‑101 task).

4. Scale to 2,048 Worlds

nworld = 2_048
d = mjw.make_data(m, nworld=nworld, nconmax=spec.nconmax, njmax=spec.njmax)

# Replicate the initial state across the batch
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, device=device))
wp.copy(d.ctrl, wp.array(np.tile(mjd.ctrl, (nworld, 1)), dtype=wp.float32, device=device))

mjw.forward(m, d)

# Capture the step into a CUDA graph for low‑overhead replay
with wp.ScopedCapture() as cap:
    mjw.step(m, d)
step_graph = cap.graph
  • The captured graph reuses the same buffers; only d.ctrl needs updating between replays.

5. Measure Aggregate Throughput

# Warm‑up (compilation, allocation)
for _ in range(10):
    wp.capture_launch(step_graph)
wp.synchronize()

t0 = time.perf_counter()
for _ in range(200):
    wp.capture_launch(step_graph)
wp.synchronize()
elapsed = time.perf_counter() - t0
world_steps = 200 * nworld
print(f"{world_steps/elapsed:,.0f} world‑steps/second")
  • Report both world‑steps/second and ms per batched step alongside the batch size.
  • Use the provided scaling_study.py script to sweep batch sizes (1, 64, 1 024, 2 048, 8 192) and observe where throughput saturates.

Getting Started

  • Warp – pip install warp-lang (≥ 1.15 for deterministic mode). Explore examples with python -m warp.examples.browse.
  • MJWarp – pip install mujoco-warp; launch a viewer with mjwarp-viewer path/to/scene.xml.
  • SO‑101 resources – NVIDIA’s Sim‑to‑Real course and Menagerie assets.
  • Training on top of MJWarp – mjlab, MuJoCo Playground (MJX impl='warp'), or Isaac Lab via the upcoming Newton integration.

What’s Next

The next post will embed MJWarp inside the Newton physics engine, adding multi‑format assets, swappable solvers, sensors, and an Isaac Lab training pipeline. It will reuse the same SO‑101 task (and optional reBot variant) to illustrate the additional integration steps.

References

  • State of Simulation for Physical AI – first blog in the series.
  • NVIDIA Warp GitHub and documentation (v1.15.0 adds GPU determinism).
  • MuJoCo Warp repository and official docs.
  • mjlab (arXiv:2601.22074) and MuJoCo Playground.
  • NVIDIA SO‑101 sim‑to‑real learning path.
  • Upcoming Newton repository.

Sources