How to Batch-Simulate Robot Physics on a GPU With NVIDIA Warp and MuJoCo Warp

If you have ever trained a robot to walk in simulation before trying it on real hardware, you have run into a wall: a physics engine like MuJoCo can simulate one robot extremely well, but reinforcement learning needs thousands or millions of attempts to learn anything. Stepping one simulation at a time on a CPU, even a fast one, is nowhere near enough throughput. The fix that robotics labs actually use today is to run many independent copies of the same physics scene at once on a GPU, so a single graphics card advances thousands of “worlds” in the time it would take a CPU to advance a handful.

In this tutorial you will build that pipeline from scratch: a small physics scene defined in MuJoCo’s native XML format, ported to the GPU with MuJoCo Warp (MJWarp), and scaled up to tens of thousands of parallel copies. Every command in this tutorial was run on a real NVIDIA GPU, and every number you will see (timings, memory use, even a floating-point rounding error) is copied from that real output, not invented. Along the way you will also learn to write your own GPU kernel with NVIDIA Warp, the Python-to-CUDA compiler that MJWarp itself is built on.

What you will accomplish and why it matters

By the end of this tutorial you will have:

• Confirmed your GPU is visible to Warp and written a small GPU-accelerated kernel by hand.
• Built a two-link pendulum physics scene, simulated it on the CPU, then ported the identical scene to the GPU and proven the two produce matching results.
• Measured, on real hardware, exactly how many parallel physics “worlds” you need before a GPU actually starts winning a race against a plain CPU loop.
• Learned two real, reproducible gotchas: a silent timing bug from GPU kernels running asynchronously, and a one-time compilation delay that catches almost everyone the first time they run code like this on a new machine.

A few terms, defined before you need them:

• Kernel: a small function that runs once per “thread.” On a GPU, thousands of threads run the same kernel code at the same time, each operating on its own slice of data.
• SIMT (single instruction, multiple threads): the execution model behind that idea. Every thread runs the identical instructions, but each thread has its own inputs, so the effect is like running a Python for loop instantly in parallel instead of one iteration at a time.
• MJCF: MuJoCo’s native XML file format for describing a physics scene: bodies, joints, geometry, and how they connect.
• World (in this context): one independent copy of a physics scene. “Batching 4,096 worlds” means simulating 4,096 separate pendulums (or robots, or whatever your scene contains) at the same time, each with its own state.

Prerequisites

• A Windows or Linux machine with an NVIDIA GPU and a recent driver. This tutorial used a fairly modest workstation card, an NVIDIA RTX A400 with 4 GB of VRAM, CUDA 13.2, driver 595.79. You do not need a data-center GPU for any of this; if nvidia-smi shows any modern NVIDIA card, you are set.
• Python 3.10 or newer and pip. This tutorial used Python 3.13.14.
• Comfort with basic Python (functions, loops, NumPy arrays) and running commands from a terminal.
• No prior experience with CUDA, GPU programming, or physics simulation is assumed. Both are explained as you go.

Step 1: Confirm your GPU and install the packages

First, check that your system actually sees an NVIDIA GPU and driver:

nvidia-smi

Expected output includes a table naming your GPU, its driver version, and its CUDA version, something like:

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 595.79                 Driver Version: 595.79         CUDA Version: 13.2     |
+-----------------------------------------+------------------------+----------------------+
|   0  NVIDIA RTX A400              WDDM  |   00000000:02:00.0 Off |                  N/A |
| 30%   31C    P8            N/A  /   50W |     344MiB /   4094MiB |      0%      Default |
+-----------------------------------------+------------------------+----------------------+

If this command fails or shows no GPU, stop here and install or update your NVIDIA driver first; nothing else in this tutorial will work without it. Once confirmed, install the three packages you will use, ideally inside a fresh virtual environment:

python -m venv venv
venv\Scripts\activate        # on Linux/macOS: source venv/bin/activate
pip install warp-lang mujoco mujoco-warp

Three different, easy-to-confuse names, so it is worth being precise about what each one does:

• warp-lang installs the warp Python package: a framework that JIT-compiles ordinary-looking Python functions into CUDA kernels. It is the compiler underneath everything else here.
• mujoco is Google DeepMind’s physics engine itself: a CPU library for simulating rigid bodies, joints, and contacts, defined via MJCF files.
• mujoco-warp is a GPU reimplementation of MuJoCo’s physics, built on top of Warp. It reads the exact same MJCF models as regular MuJoCo, but steps many copies of a scene in parallel on the GPU.

Step 2: Confirm Warp can see your GPU

Create a file named check_gpu.py:

import warp as wp

wp.init()
print("CUDA available:", wp.is_cuda_available())
print("Devices:", wp.get_devices())

Run it:

python check_gpu.py

Real output from this tutorial’s machine:

Warp 1.17.0 initialized:
   CUDA Toolkit 12.9, Driver 13.2
   Devices:
     "cpu"      : "Intel64 Family 6 Model 198 Stepping 2, GenuineIntel"
     "cuda:0"   : "NVIDIA RTX A400" (4 GiB, sm_86, mempool enabled)
CUDA available: True
Devices: ['cpu', 'cuda:0']

Two devices show up: "cpu" and "cuda:0". This matters because every Warp kernel can run on either device unchanged. That flexibility is genuinely useful: you can develop and debug logic on the CPU, where errors are easier to read, then switch a single string to run the identical code at GPU scale. If CUDA available prints False, Warp has fallen back to CPU-only mode; double-check your driver installation before continuing.

Step 3: Write and run your first Warp kernel

Before touching a full physics engine, it helps to see the raw mechanics of a GPU kernel on something small. The example below simulates several projectiles launched at different speeds, each affected by gravity and simple linear air drag. Save it as first_kernel.py:

import numpy as np
import warp as wp


@wp.kernel
def simulate_projectile(
    positions: wp.array(dtype=wp.vec3),
    velocities: wp.array(dtype=wp.vec3),
    drag: float,
    dt: float,
):
    i = wp.tid()
    gravity = wp.vec3(0.0, 0.0, -9.81)
    v = velocities[i]
    drag_accel = -drag * v
    velocities[i] = v + (gravity + drag_accel) * dt
    positions[i] = positions[i] + velocities[i] * dt


wp.init()
device = "cuda:0" if wp.is_cuda_available() else "cpu"
print("Running on:", device)

n = 5
start_pos = np.zeros((n, 3), dtype=np.float32)
start_vel = np.array([[2.0 + 0.5 * i, 0.0, 5.0] for i in range(n)], dtype=np.float32)

positions = wp.array(start_pos, dtype=wp.vec3, device=device)
velocities = wp.array(start_vel, dtype=wp.vec3, device=device)

dt = 1.0 / 60.0
for _ in range(60):
    wp.launch(simulate_projectile, dim=n, inputs=[positions, velocities, 0.05, dt], device=device)

wp.synchronize_device(device)
print("Final positions after 60 steps (1 second):")
print(positions.numpy())

Run it with python first_kernel.py. Real captured output:

Running on: cuda:0
Final positions after 60 steps (1 second):
[[ 1.9499899   0.         -0.03102732]
 [ 2.4374874   0.         -0.03102732]
 [ 2.9249852   0.         -0.03102732]
 [ 3.4124823   0.         -0.03102732]
 [ 3.8999798   0.         -0.03102732]]

Walking through what happened, and why:

• @wp.kernel marks a function to be compiled to native CPU or CUDA code the first time it runs. Inside a kernel you write a restricted, statically typed subset of Python: no arbitrary Python objects, only Warp’s own array, vector, and scalar types.
• wp.tid() returns the “thread index,” meaning which particle this particular invocation of the kernel is responsible for. There is no explicit loop anywhere in this code; wp.launch(..., dim=n, ...) tells Warp to run the kernel n times, and on the GPU those n invocations can execute simultaneously instead of one after another.
• Device arrays (wp.array(..., device=device)) live entirely on whichever device you chose. Calling .numpy() on a CUDA array triggers a real copy from GPU memory back to your CPU’s RAM; that copy is not free, and doing it inside a hot loop (rather than once, at the end, as done here) is a common way to accidentally make GPU code slower than CPU code.
• wp.synchronize_device(device) blocks until every kernel launched on that device has actually finished. This is easy to skip and looks harmless when you do; Step 6 shows exactly what goes wrong if you forget it.

Verify this step worked: the five projectiles were launched with slightly different horizontal speeds (2.0, 2.5, 3.0, 3.5, 4.0 m/s) and identical vertical motion, so after exactly one second you should see five different x-positions with an identical z-position, which is exactly what the output above shows.

Step 4: Build a physics scene and simulate it on the CPU

Now build an actual physics scene: a two-link pendulum. This is deliberately simple (two hinges, two capsule-shaped links, a floor) so you can reason about every number that comes out of it. Save this as pendulum.xml:

<mujoco model="double_pendulum">
  <option timestep="0.005" gravity="0 0 -9.81"/>
  <worldbody>
    <light diffuse=".8 .8 .8" pos="0 0 3" dir="0 0 -1"/>
    <geom type="plane" size="2 2 0.1" rgba=".9 .9 .9 1"/>
    <body name="link1" pos="0 0 2">
      <joint name="hinge1" type="hinge" axis="0 1 0" damping="0.02"/>
      <geom type="capsule" fromto="0 0 0 0 0 -0.5" size="0.03" rgba=".8 .2 .2 1" mass="1.0"/>
      <body name="link2" pos="0 0 -0.5">
        <joint name="hinge2" type="hinge" axis="0 1 0" damping="0.02"/>
        <geom type="capsule" fromto="0 0 0 0 0 -0.5" size="0.03" rgba=".2 .2 .8 1" mass="1.0"/>
      </body>
    </body>
  </worldbody>
</mujoco>

A quick read of the structure: <worldbody> is the top-level container. link1 is a body attached to the world by a hinge joint that rotates around the y-axis; link2 is nested inside link1, attached by its own hinge, which is how MJCF expresses a kinematic chain: each body’s position is relative to its parent, so link2 swings from wherever link1’s end happens to be.

Now simulate it on the CPU. Save as cpu_baseline.py:

import mujoco

mjm = mujoco.MjModel.from_xml_path("pendulum.xml")
mjd = mujoco.MjData(mjm)
mjd.qpos[:] = [0.3, 0.1]

for _ in range(200):
    mujoco.mj_step(mjm, mjd)

print("CPU final qpos after 200 steps:", mjd.qpos)

Run with python cpu_baseline.py. Real captured output:

CPU final qpos after 200 steps: [-0.2367546 -0.1022122]

MjModel holds the static, unchanging description of the scene (loaded once from your XML). MjData holds the mutable state: joint positions (qpos), velocities, and everything else that changes as the simulation runs. qpos here has two entries, one angle per hinge, because this model has two joints. mj_step is MuJoCo’s own name for what most physics engines call a “tick”: compute forces, solve any contacts, integrate the equations of motion forward by one timestep (0.005 seconds here, so 200 steps covers one simulated second).

Verify this step worked: starting both hinges at a small nonzero angle (0.3 and 0.1 radians) and letting gravity pull the pendulum down and around should leave it somewhere clearly different from where it started, which is exactly what the negative final angles above show. This CPU run is also your ground truth: the next step exists specifically to prove the GPU port reproduces it.

Step 5: Port the scene to the GPU and prove it matches

MuJoCo Warp reads the same MJCF file, so porting a scene means uploading the loaded CPU model and state into GPU-resident arrays, then stepping those arrays instead. Save as gpu_parity.py:

import numpy as np
import mujoco
import mujoco_warp as mjwarp
import warp as wp

wp.init()
mjm = mujoco.MjModel.from_xml_path("pendulum.xml")

# CPU reference (identical to Step 4)
mjd = mujoco.MjData(mjm)
mjd.qpos[:] = [0.3, 0.1]
for _ in range(200):
    mujoco.mj_step(mjm, mjd)
cpu_final = mjd.qpos.copy()

# GPU port: one world, same initial condition
m = mjwarp.put_model(mjm)
d = mjwarp.put_data(mjm, mujoco.MjData(mjm), nworld=1)
d.qpos.assign(np.array([[0.3, 0.1]], dtype=np.float32))
for _ in range(200):
    mjwarp.step(m, d)
wp.synchronize()
gpu_final = d.qpos.numpy()[0]

print("CPU final qpos:", cpu_final)
print("GPU final qpos:", gpu_final)
print("Max abs difference:", np.max(np.abs(cpu_final - gpu_final)))

Run with python gpu_parity.py. Real captured output:

CPU final qpos: [-0.2367546 -0.1022122]
GPU final qpos: [-0.23675461 -0.10221219]
Max abs difference: 1.4060637626434058e-08

mjwarp.put_model(mjm) converts the static MuJoCo model into a GPU-resident structure once. mjwarp.put_data(mjm, mjd, nworld=N) allocates GPU arrays sized for N independent copies of that model’s state; here nworld=1 means exactly one world, deliberately, so this run is a fair apples-to-apples comparison against Step 4. mjwarp.step(m, d) is the GPU equivalent of mj_step.

The two final answers are not bit-for-bit identical (CPU MuJoCo computes in 64-bit floating point by default, while MJWarp’s GPU arrays are 32-bit), but they agree to about one part in a hundred million, which is exactly the kind of rounding-level difference you should expect and hope for. This step is the single most important habit in this whole tutorial: never trust that a “GPU port” of a physics or numerical pipeline is correct just because it runs without errors. Always run one identical case through both the trusted CPU path and the new GPU path first, and only move on once the two agree to a sensible tolerance. If they had disagreed by, say, 0.1 instead of 0.00000001, that would mean something in the port is wrong, not that the GPU is simply “a bit different.”

Step 6: Scale up to thousands of parallel worlds

The entire point of MJWarp is nworld being large. Save as batch_scaling.py:

import time
import numpy as np
import mujoco
import mujoco_warp as mjwarp
import warp as wp

wp.init()
mjm = mujoco.MjModel.from_xml_path("pendulum.xml")
m = mjwarp.put_model(mjm)

# warm up so kernel compilation doesn't pollute the timing (see Step 8)
warm = mjwarp.put_data(mjm, mujoco.MjData(mjm), nworld=8)
for _ in range(5):
    mjwarp.step(m, warm)
wp.synchronize()

rng = np.random.default_rng(42)
for nworld in [64, 256, 1024, 4096, 16384, 32768]:
    d = mjwarp.put_data(mjm, mujoco.MjData(mjm), nworld=nworld)
    init = rng.uniform(-0.5, 0.5, size=(nworld, 2)).astype(np.float32)
    d.qpos.assign(init)
    wp.synchronize()

    t0 = time.perf_counter()
    for _ in range(500):
        mjwarp.step(m, d)
    wp.synchronize()  # without this, the timer stops before the GPU actually finishes
    t1 = time.perf_counter()

    rate = 500 * nworld / (t1 - t0)
    print(f"nworld={nworld:>6}: {t1 - t0:.3f}s for 500 steps -> {rate:,.0f} world-steps/sec")

Run with python batch_scaling.py. Real captured output from this tutorial’s 4 GB card:

nworld=    64: 1.355s for 500 steps -> 23,624 world-steps/sec
nworld=   256: 1.646s for 500 steps -> 77,771 world-steps/sec
nworld=  1024: 1.765s for 500 steps -> 290,012 world-steps/sec
nworld=  4096: 1.825s for 500 steps -> 1,121,904 world-steps/sec
nworld= 16384: 2.710s for 500 steps -> 3,023,396 world-steps/sec
nworld= 32768: 4.477s for 500 steps -> 3,659,838 world-steps/sec

Two things worth noticing before Step 7 explains why. First, throughput keeps climbing as nworld grows, but not proportionally: going from 4,096 to 32,768 worlds (8 times more work) only took roughly 2.5 times longer, meaning the GPU handled most of that extra work almost for free. Second, checking GPU memory with nvidia-smi in a separate terminal, during an equivalent run against these same batch sizes on this same machine, showed usage climbing from an idle 344 MiB to about 932 MiB at 32,768 worlds, out of the 4,094 MiB available on this modest card. There was room to push further; 32,768 tiny two-joint worlds is nowhere near this GPU’s actual ceiling.

The comment above wp.synchronize() flags a genuine gotcha. Warp launches GPU kernels asynchronously: the Python call to mjwarp.step(...) returns immediately, while the actual computation continues on the GPU in the background. If you stop your timer before calling wp.synchronize(), you measure how long it took Python to submit the work, not how long the GPU took to finish it, and your benchmark will report a number that is too fast. On this same machine, timing 500 steps at nworld=16384 without a trailing synchronize measured 2.55 seconds, versus the honest 2.82 seconds with one; the gap grows the more work is still queued up when your clock stops. Always synchronize immediately before you read a wall-clock timer around GPU work.

Step 7: Compare against a naive CPU loop and find the real crossover point

Numbers on their own do not tell you whether the GPU is actually a good idea for this model; you need something to compare against. The obvious CPU alternative is simply looping over independent MjData objects one at a time. Save as cpu_naive_loop.py:

import time
import numpy as np
import mujoco

mjm = mujoco.MjModel.from_xml_path("pendulum.xml")
rng = np.random.default_rng(42)

for nworld in [64, 256]:
    inits = rng.uniform(-0.5, 0.5, size=(nworld, 2))
    datas = [mujoco.MjData(mjm) for _ in range(nworld)]
    for i, d in enumerate(datas):
        d.qpos[:] = inits[i]

    t0 = time.perf_counter()
    for _ in range(500):
        for d in datas:
            mujoco.mj_step(mjm, d)
    t1 = time.perf_counter()

    rate = 500 * nworld / (t1 - t0)
    print(f"CPU loop, nworld={nworld:>4}: {t1 - t0:.3f}s for 500 steps -> {rate:,.0f} world-steps/sec")

Run with python cpu_naive_loop.py. Real captured output:

CPU loop, nworld=  64: 0.047s for 500 steps -> 676,000 world-steps/sec
CPU loop, nworld= 256: 0.195s for 500 steps -> 656,095 world-steps/sec

This is the finding most beginners do not expect: at small batch sizes, the plain single-threaded CPU loop is dramatically faster than the GPU, roughly 28 times faster at nworld=64 (676,000 versus 23,624 world-steps per second) and still about 8 times faster at nworld=256. The GPU only pulls ahead once there is enough parallel work to justify its overhead. Re-running Step 6’s batch sizes alongside this CPU number (which stays essentially flat, since the loop is sequential and does not benefit from more worlds) gives a clear crossover:

Testing the gap more finely between 1,024 and 4,096 worlds on this same machine, the two curves cross somewhere around 2,000 to 2,600 worlds for this specific two-joint model. Your exact crossover point will depend on your GPU and your model’s size and complexity, but the shape of the curve is the real lesson: a GPU has fixed per-launch overhead (dispatching work, synchronizing, driver bookkeeping) that a single-threaded CPU loop simply does not pay. With too little batched work, that overhead dominates and the CPU wins outright. Only once you have enough independent worlds to keep the GPU genuinely busy does its parallelism start paying for itself, after which it keeps pulling further and further ahead. This is precisely why reinforcement learning pipelines that need thousands of parallel rollouts reach for tools like MJWarp, while a single robot being teleoperated in real time is usually still best served by plain CPU MuJoCo.

Step 8: Two more gotchas worth knowing before you rely on this

You already saw the asynchronous-timing gotcha in Step 6. Here is the other one that will surprise you the first time it happens.

Clear Warp’s on-disk kernel cache and NVIDIA’s driver-level compute cache (on Windows, these live under %LOCALAPPDATA%\NVIDIA\warp\Cache and %APPDATA%\NVIDIA\ComputeCache), then time the very first call into MJWarp on this now genuinely clean machine:

COLD (caches cleared): imports + put_model + put_data + first step() = 21.14s total

Running the identical script again immediately afterward, with both caches now populated:

WARM (caches populated): imports + put_model + put_data + first step() = 1.13s total

A roughly twenty-second difference for running exactly the same code twice in a row. The first time MJWarp needs a particular combination of kernels for your specific model, Warp has to generate CUDA source, invoke NVIDIA’s compiler, and cache the result; this tutorial’s own terminal showed over thirty separate "Module ... load on device 'cuda:0' took N ms (compiled)" lines scroll by during that first cold run. Two practical consequences follow. First, never include your first call to put_model, put_data, or step inside a timed benchmark; that is exactly what the “warm up” loop in Step 6’s script exists to avoid. Second, do not panic the first time you run GPU physics code on a new machine (or after a package upgrade) and it appears to hang for twenty seconds; it is compiling, not broken, and every run after that first one will be fast again because the compiled code is cached on disk.

Confirm everything works end to end

Before moving on, you should be able to check off all of the following, each backed by output you personally saw:

• nvidia-smi and wp.is_cuda_available() both confirm your GPU is visible.
• A hand-written Warp kernel (Step 3) produces physically sensible output: five projectiles with different horizontal speeds land at five different positions after one second.
• The same pendulum model, stepped 200 times, gives matching final joint angles on the CPU and on the GPU, agreeing to roughly 1e-8 (Step 5).
• You can reproduce a batch-scaling table like the one in Step 7 on your own hardware, and locate roughly where your own GPU starts outperforming a naive CPU loop for your own model.
• You know to synchronize before timing GPU work, and to warm up before benchmarking it.

From here, a natural next step is training an actual reinforcement learning policy against a batched MJWarp environment instead of just measuring raw physics throughput; the Warp documentation covers its automatic differentiation support (wp.Tape), which lets you backpropagate gradients through a simulation, something not used anywhere in this tutorial but directly relevant if you move on to gradient-based control or trajectory optimization. If you want to go deeper on the physics engine itself rather than its GPU port, the official MuJoCo documentation covers contacts, actuators, and sensors in far more depth than this tutorial’s bare hinges. And if you have not yet installed a local model runner for other AI experiments on the same machine, this site’s own guide to installing Ollama is a good companion piece, since GPU driver troubleshooting tends to come up in both contexts.