Loading post…
Hugging FaceJohnny Nuñez Cano, Asier Arranz, Rishabh Chadha, Ben OliveriTutorial
How to Use NVIDIA Warp and MjWarp to Accelerate Robotics Simulation and Learning Workflows
Summary
NVIDIA Warp is a Python framework for writing GPU-accelerated kernels that compiles statically typed Python to CUDA. MuJoCo Warp (MJWarp) implements MuJoCo physics on Warp, enabling batched simulation of hundreds to thousands of independent robot environments on NVIDIA GPUs. The core API transition is small: mjw.put_model and mjw.put_data replace mujoco.MjModel and mujoco.MjData, and mjw.step advances the entire batch. Proper throughput measurement requires GPU synchronization and warm-up to avoid timing queue latency instead of actual GPU work. The post demonstrates moving an SO-101 follower arm from a single CPU world to 2,048 parallel MJWarp environments, validating parity and measuring aggregate world-steps per second.
Context
This is the second article in the State of Simulation for Physical AI series. The first article mapped the robot-simulation landscape. This post prepares and scales the simulation environment but does not train a policy. Later posts cover Newton and Isaac Lab integration layers.
Approach / What changed
The post describes moving an SO-101 follower arm from a single MuJoCo CPU world to as many as 2,048 parallel MJWarp environments on GPU. It covers the core API transition (mjw.put_model, mjw.put_data, mjw.step), batching with np.tile for identical starting states, and proper throughput measurement using CUDA Graphs, warm-up, and GPU synchronization. The workflow validates one MuJoCo world, migrates it to MJWarp, forms a batch, verifies parity, and measures aggregate world-steps per second.
Takeaways
MJWarp enables parallel simulation of hundreds to thousands of environments on GPU, improving aggregate throughput for reinforcement learning and large-scale sampling.
The core API transition from MuJoCo to MJWarp is small: use mjw.put_model, mjw.put_data, and mjw.step instead of mujoco.mj_step.
Proper throughput measurement requires GPU synchronization and warm-up; a naive timer measures queue time, not GPU work.
FAQ
What is the difference between MuJoCo and MJWarp?
MuJoCo runs one or a few CPU worlds with low latency per step. MJWarp implements MuJoCo physics on NVIDIA GPUs via Warp, allowing hundreds or thousands of independent worlds to be advanced in a single batched step, improving aggregate throughput for reinforcement learning and large-scale sampling.
How do I measure MJWarp throughput correctly?
GPU launches are asynchronous, so a naive timer measures queue time, not GPU work. Warm up first to pay compilation costs, then synchronize immediately before and after the timed region. Use wp.synchronize() and time.perf_counter() to measure actual GPU execution.
What batch sizes does MJWarp support?
The post demonstrates scaling from 1 to 2,048 parallel environments for an SO-101 arm. The scaling_study.py script sweeps batch sizes like 1, 64, 1024, 2048, and 8192. Results depend on scene, simulation settings, and hardware; the measured curve shows where additional worlds improve throughput and where memory or compute limits reduce benefit.
Source and freshness
- Source
- huggingface.co
- Published
- Added to Yomu
Summaries are generated from the original by a language model. Check the original before relying on it.
Keep reading
Related reading
The Good Stuff
Paid picks we use ourselves. Links carry a referral code — a subscription through one earns us a commission at no extra cost to you.
The best engineering reads, once a week
One useful email. Unsubscribe anytime.