How to Run 2,048 MuJoCo Robot Simulations in Parallel on a Single GPU with NVIDIA Warp

Why GPU-Scale Simulation Matters

Classic MuJoCo runs fast on a CPU and parallelizes sampling across cores. But modern robot learning workloads need thousands of worlds stepping simultaneously. NVIDIA Warp and its MuJoCo implementation, MJWarp, let you keep the same MJCF model while advancing hundreds or thousands of independent states per GPU call.

This tutorial walks through migrating an SO-101 pick-and-place scene from CPU MuJoCo to 2,048 parallel MJWarp environments. You will validate one world first, then scale and benchmark correctly.

Key distinction: latency is one step's wall-clock time; aggregate throughput is total world-steps per second. MJWarp optimizes the latter.

For a broader view of GPU utilization strategies, see our AI ROI and cost management blueprint.

Step 1: Validate One World on CPU

Start with the familiar MuJoCo loop. The task: grasp a 44 mm cube and stack it on a blue cube.

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  # 0.002 s

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)

Success criteria: xy_err <= 0.015 and 0.035 <= dz <= 0.055 after settling.

Step 2: Move to MJWarp with One World

import warp as wp
import mujoco_warp as mjw

wp.init()
device = wp.get_device()
m = mjw.put_model(mjm)  # raises on unsupported features
d = mjw.make_data(mjm, nworld=1, nconmax=128, njmax=300)

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. Keep the host in the loop for validation:

def simulate_frame():
    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)  # refresh derived quantities

Watch for overflow warnings: increase nconmax/njmax if MJWarp reports capacity issues.

Step 3: Scale to 2,048 Worlds and Capture a CUDA Graph

import numpy as np

nworld = 2048
d = mjw.make_data(mjm, nworld=nworld, nconmax=128, njmax=300)
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

Step 4: Benchmark Correctly

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()  # without this you time the queue, not the work
elapsed = time.perf_counter() - t0
print(f"{200 * nworld / elapsed:,.0f} world-steps/second")

Report both world-steps per second and milliseconds per batched step, alongside batch size. A single-world latency comparison does not establish throughput.

Limitations and Caveats

• Single-world latency can regress. Kernel launch overhead and host-device copies may make one GPU world slower than CPU MuJoCo. MJWarp's value is batching.
• Buffer sizing is critical. Exceeding nconmax/njmax invalidates the rollout even if execution continues with a warning. Size against the most contact-heavy moment (both jaws and table touching a cube).
• Determinism is opt-in. Warp 1.15+ supports deterministic execution, but it trades performance for reproducible ordering; it is not guaranteed across an entire MJWarp rollout by default.
• First launch compiles kernels. Expect several minutes of startup delay on first run.

Next Steps

• Port the same MJCF into Newton using newton.solvers.SolverMuJoCo for multi-solver support and Isaac Lab integration.
• Try mjlab or MuJoCo Playground for policy training on top of MJWarp physics.
• Explore Warp's differentiable kernels for gradient-based simulation and design optimization.

For scaling GPU workloads beyond simulation, see our deep dive on maximizing GPU utilization for LLM inference with NVIDIA Runai and NIM.

Takeaway

MJWarp lets you keep your MJCF model and scale from one CPU world to thousands of GPU worlds. Validate parity first, size your contact buffers against the worst-case moment of the task, and benchmark aggregate throughput with proper synchronization. The migration is small in API surface but requires discipline in measurement and capacity planning.

Start with pip install mujoco-warp and the official Colab tutorial, then scale your own scene.