Using NVIDIA Warp and MJWarp to Accelerate Robotics Simulation and Learning Workflows

Using NVIDIA Warp and MJWarp to Accelerate Robotics Simulation and Learning Workflows

Classic MuJoCo provides fast CPU-based robot simulation and can parallelize sampling across CPU cores. As learning workloads grow, the important question becomes how many independent worlds can run simultaneously. GPU acceleration advances those worlds in batches while keeping simulation and learning data close to the device.

MuJoCo Warp (MJWarp), built on NVIDIA Warp, brings compatible MuJoCo models to NVIDIA GPUs. This walkthrough moves an SO-101 follower arm from a standard MuJoCo workflow to as many as 2,048 parallel MJWarp environments. It also covers the validation and measurement steps needed for a reliable migration.

This article prepares and scales the simulation environment. It does not train a policy.

Putting the stack together

A practical choice depends on the workload:

Start with a Warp kernel

NVIDIA Warp is a Python framework for writing high-performance, GPU-accelerated kernels. Developers write statically typed kernels in Python, and Warp compiles them for CPU or CUDA execution. The first launch builds and caches a native module, while later launches reuse it.

Warp's kernel language is a performance-oriented subset of Python. Regular Python remains responsible for configuration, memory allocation, and launch orchestration.

The following kernel advances point positions under gravity. Each logical thread handles one point, allowing the same pattern to scale from two points to millions:

import numpy as np
import 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())

Warp's main capabilities can be grouped into three areas:

Properties that matter in robotics

• Explicit parallel work: wp.tid() identifies the point, contact, body, or world handled by the current logical thread.
• Explicit device arrays: Arrays reside on the selected device. Calling .numpy() on a CUDA array synchronizes execution and copies data to CPU memory. It is not a zero-copy operation. Device-resident PyTorch or JAX pipelines should use Warp framework adapters or DLPack-compatible sharing.
• Composable launches: Programs can launch several focused kernels and capture supported CUDA work in a graph to reduce repeated dispatch overhead. Graph capture replays launches against existing buffers, but it does not fuse arbitrary kernels.

Differentiability and determinism

Warp kernels are differentiable. A wp.Tape records forward kernel launches made inside its context and replays their adjoints in reverse when backward() is called. This supports differentiable geometry, computational fluid dynamics, and custom physics workflows, including simulation and design optimization.

Warp also supports deterministic execution, introduced in Warp 1.15. GPU atomic operations are scheduler-dependent by default, so repeated launches can differ slightly. The opt-in deterministic modes trade some performance for reproducible ordering in simulation, validation, and regression tests.

These are Warp capabilities, not guarantees that an entire MJWarp rollout is differentiable or deterministic.

Warp can be installed and explored with:

pip install warp-lang
python -m warp.examples.browse

What is MuJoCo Warp?

A robot simulator repeatedly computes the next state from joint positions, velocities, controls, and contacts. Here, a world means one independent copy of a scene and its state. One world might contain an SO-101 arm reaching for a cube, while another contains the same arm from a different starting pose.

MuJoCo and MJWarp can simulate the same compatible robot and task, but they organize computation differently. MuJoCo is well suited to developing and inspecting one or a few CPU worlds. MJWarp implements MuJoCo's physics pipeline with NVIDIA Warp, placing the model and a batch of independent states on NVIDIA GPUs. A call to mjw.step advances every world in the batch.

MJWarp's primary benefit is not necessarily lower latency for one world. Its benefit is the ability to advance hundreds or thousands of worlds together, providing enough parallel work for higher aggregate throughput, measured as total world-steps completed per second. This is useful for reinforcement learning and large-scale sampling.

The migration has three stages:

1. Validate one MuJoCo world.
2. Move it to MJWarp and create a batch.
3. Verify the results and measure performance correctly.

Two performance terms are important:

• Latency: Wall-clock time for one simulation step.
• Aggregate throughput: Total world-steps completed per measured wall-clock second.

Basic usage, structures, and batch sizes

The central API transition is small:

Use mjw.make_data() for a default or fresh state. Use mjw.put_data() when the exact initialized MuJoCo state must cross the migration boundary.

Batched resources require the following parameters:

Performance tuning

1. Capture CUDA work in a graph. mjw.step launches many kernels, so capture it once and replay it repeatedly:

   with wp.ScopedCapture() as capture:
        mjw.step(m, d)
   wp.capture_launch(capture.graph)
   

2. Size buffers carefully. Memory use and work scale with nconmax, naconmax, and njmax. Use mjwarp-testspeed --measure_alloc and watch for overflows in mjwarp-viewer.

After sizing contact and constraint buffers, test solver iteration limits without changing task behavior. Meshes and continuous collision detection settings can increase memory use. When measured contact counts allow it, nccdmax and naccdmax can reduce CCD buffer allocation.

MJWarp's compact solver uses MuJoCo's Newton constraint solver and sleeping. It is separate from the Newton physics-engine framework. Compact-solver and multi-GPU configuration require separate consideration.

Possible training integrations include:

• Isaac Lab through Newton
• mjlab, which provides a manager API directly on MJWarp with PyTorch
• MuJoCo Playground through MJX with impl='warp'

Migrating a MuJoCo scene to MJWarp

1. Establish a MuJoCo CPU baseline

The initial scene is ordinary MJCF: an SO-101 arm, a table, and two cubes to stack.

<mujoco model="so101_pick_place">
  <include file="so101.xml"/>

  <worldbody>
    <light pos="0.3 0 1.5" dir="0 0 -1" directional="true"/>
    <geom name="floor" type="plane" size="0 0 0.05"/>

    <geom name="table" type="box" pos="0.35 -0.04 0.012"
          size="0.16 0.26 0.012" rgba="0.32 0.32 0.32 1"
          friction="1 0.005 0.0005" condim="3"/>

      <freejoint name="red_cube_joint"/>
      <geom type="box" size="0.022 0.022 0.022" mass="0.08"
            rgba="0.85 0.05 0.04 1" friction="1.2 0.005 0.0005" condim="3"/>

      <freejoint name="blue_cube_joint"/>
      <geom type="box" size="0.022 0.022 0.022" mass="0.08"
            rgba="0.05 0.20 0.90 1" friction="1.2 0.005 0.0005" condim="3"/>
    
  </worldbody>
</mujoco>

For an MJCF box, size values are half-extents. Therefore, size="0.022 ..." defines a cube with 44 mm edges. The task uses this dimension for its success thresholds. The arm base is at the origin, its reach extends along +X, and the cubes are arranged along Y.

In the companion repository, the scene is generated rather than written entirely by hand. resolve_pick_place_scene() copies the Menagerie arm into .generated/, fills table and cube coordinates from a robot profile, and writes scene_pick_place.xml. The walkthrough uses the SO-101 profile. An optional reBot profile is also available.

Loading and stepping remain standard MuJoCo operations:

import mujoco

mjm = mujoco.MjModel.from_xml_path("scene_pick_place.xml")
mjd = mujoco.MjData(mjm)

fps = 50
sim_substeps = 10
frame_dt = 1.0 / fps
mjm.opt.timestep = frame_dt / sim_substeps

controller = PickPlaceController(spec=spec)

for _ in range(600):
    ctrl = controller.step(mjm, mjd, frame_dt)
    for _ in range(sim_substeps):
        mjd.ctrl[: mjm.nu] = ctrl
        mujoco.mj_step(mjm, mjd)

The control structure is important: controls are computed once per frame, and physics advances sim_substeps times. At 50 control frames per second and 10 physics steps per frame, the physics timestep should be 0.002 seconds:

mjm.opt.timestep = frame_dt / sim_substeps

Set this before both the CPU rollout and the model upload with mjw.put_model(). Otherwise, parity comparisons, simulated-time throughput measurements, and action-rate assumptions can become inconsistent.

For 44 mm cubes, stacking success has two measurable conditions:

• Horizontal center error: xy_err <= 0.015 m
• Vertical center separation: 0.035 m <= dz <= 0.055 m

Evaluate both after the cubes have settled. A successful process exit alone does not establish task success.

The SO-101 arm comes from MuJoCo Menagerie, pinned to a known-good commit because Menagerie assets can change. The optional reBot variant uses a separate scene, gripper, and capacity profile with nconmax=256 and njmax=500. It should be validated separately.

2. Validate one-world MJWarp parity

Run one world on the GPU first, while keeping the host loop available for inspection and comparison. Upload the model, allocate batched state, seed it from the initialized host state, and perform a forward pass:

wp.init()
import mujoco_warp as mjw

m = mjw.put_model(mjm)
d = mjw.make_data(mjm, nworld=1, nconmax=spec.nconmax, njmax=spec.njmax)

device = wp.get_device()
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)

Every device array has a leading world dimension. The host state therefore becomes mjd.qpos[None, :], with shape (1, nq) rather than (nq,). Scaling to thousands of worlds changes only that leading dimension.

mjw.put_model() also checks model compatibility and raises an error when unsupported features are used rather than silently dropping them.

The three fields can be seeded explicitly, making the migration boundary clear. Alternatively, mjw.put_data(mjm, mjd, nworld=...) transfers the entire initialized structure in one call.

The frame loop redirects the physics step to the GPU and mirrors state back to the host:

def simulate_frame() -> None:
    ctrl = controller.step(mjm, mjd, frame_dt)
    for _ in range(sim_substeps):
        mjd.ctrl[: mjm.nu] = ctrl
        wp.copy(d.ctrl, wp.array(mjd.ctrl[None, :], dtype=wp.float32, device=device))
        mjw.step(m, d)
        mjd.qpos[:] = d.qpos.numpy()[0]
        mjd.qvel[:] = d.qvel.numpy()[0]
    mujoco.mj_forward(mjm, mjd)

Calling .numpy() synchronizes execution and copies data to the host on every substep. This makes the loop suitable for validation, inverse kinematics, viewing, and task checks, but not for throughput benchmarking.

After copying qpos and qvel, call mujoco.mj_forward(mjm, mjd) to refresh derived host quantities such as mjd.xpos. These values are not refreshed automatically merely by reading them.

3. Size contact and constraint capacity

MJWarp allocates contact and constraint buffers before stepping. If a rollout exceeds those capacities, the affected result is invalid for verification or benchmarking, even when execution continues with an overflow warning. Increase the relevant limit and rerun the task.

The SO-101 profile uses nconmax=128 and njmax=300 as starting capacities:

d = mjw.make_data(
    mjm,
    nworld=nworld,
    nconmax=spec.nconmax,
    njmax=spec.njmax,
)

Capacity should be based on the most contact-heavy part of the task. For pick-and-place, this may be the moment when both jaws and the table touch a cube, rather than a period when the arm is moving through free space.

With Option.warn_overflow at its default, MJWarp reports the needed budget, such as narrowphase overflow - please increase nconmax to ..., and flags affected worlds in Data.overflow. Read this value after stepping. mjw.put_data() can raise an error immediately because it can compare the allocated budgets with a MuJoCo state it already holds.

mjwarp-testspeed --measure_alloc reports contacts and constraints consumed by the scene and aborts when a world overflows. Treat these reports as failures: increase the limit and rerun before trusting a trajectory or benchmark. Recheck capacity whenever the model, collision geometry, or task changes.

4. Scale to 2,048 worlds

After one-world parity passes, allocate the target batch size and replicate the initialized state:

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, device=device),
)
wp.copy(
    d.ctrl,
    wp.array(np.tile(mjd.ctrl, (nworld, 1)), dtype=wp.float32, device=device),
)
mjw.forward(m, d)

with wp.ScopedCapture() as capture:
    mjw.step(m, d)
step_graph = capture.graph

np.tile gives every world the same initial state, which is appropriate for a baseline throughput measurement. Per-world randomization would instead write different rows of d.qpos on the device.

CUDA Graphs reuse the model and data buffers captured here. Update d.ctrl in place between replays. Capture a new graph after replacing buffers, changing nworld, or rebuilding the model. Graph capture requires CUDA.

5. Verify before measuring

GPU launches are asynchronous. A timer that does not synchronize can measure how quickly Python queued work rather than how quickly the GPU completed it. Warm up first, since initial launches include compilation and allocation costs, then synchronize immediately before and after the timed region:

import time

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

total = 200 * nworld
print(f"{total / elapsed:,.0f} world-steps/second")

Report both aggregate world-steps per second and milliseconds per batched step, together with the batch size. Results depend on the scene, simulation settings, and hardware. A one-world latency comparison does not establish batched throughput.

A scaling study can sweep batch sizes and report milliseconds per step, throughput, and speedup:

python solutions/so101_mjwarp_solution.py --headless-steps 600
python scaling_study.py --worlds 1 64 1024 2048 8192 --steps 100

The resulting curve shows where adding worlds improves throughput and where memory or compute limits reduce the benefit.

Getting started

Warp

pip install warp-lang
python -m warp.examples.browse

Documentation: Warp documentation and Warp on GitHub

MJWarp

pip install mujoco-warp
mjwarp-viewer benchmarks/humanoid/humanoid.xml

Documentation: MJWarp documentation, MJWarp on GitHub, and the Colab tutorial

SO-101 and training integrations

• SO-101 sim-to-real course
• Physical AI learning paths
• mjlab
• MuJoCo Playground
• Isaac Lab with Newton

What comes next

This workflow covers raw Warp to MJWarp integration: GPU kernels, batched stepping, and an SO-101 scene using mjw.step.

The next integration layer ports the same MJCF environment into Newton, using MuJoCo Warp as its rigid-body solver through newton.solvers.SolverMuJoCo. Newton manages the model, state, controls, and contacts while MJWarp runs underneath.

That workflow adds multi-format assets, swappable solvers, sensors, inverse-kinematics helpers, and an Isaac Lab path. It continues with the same SO-101 task and optional reBot profile while explaining the changes required by Newton and the separate Isaac Lab integration.

References

• NVIDIA Warp on GitHub
• Warp documentation
• Warp 1.15.0 release
• Warp deterministic execution guide
• MuJoCo Warp on GitHub
• Official MJWarp documentation
• Accelerated and differentiable computational physics with NVIDIA Warp
• Tile-based programming in Warp 1.5.0
• mjlab
• MuJoCo Playground
• NVIDIA SO-101 sim-to-real course
• Newton