Accelerating Robotics Simulation and Learning with NVIDIA Warp and MuJoCo Warp

Accelerating Robotics Simulation and Learning with NVIDIA Warp and MuJoCo Warp MuJoCo Warp (MJWarp) leverages the NVIDIA Warp framework to port MuJoCo physics pipelines to NVIDIA GPUs, enabling massive parallel world-stepping for reinforcement learning and robotics workflows. NVIDIA Warp acts as a high-performance Python kernel language with CUDA compilation, automatic differentiation, and framework interoperability, while MJWarp executes batched physics simulations directly on device buffers. NVIDIA Warp compiles statically typed Python kernels to native CUDA modules, supporting JIT compilation, kernel fusion, CUDA Graphs, and PyTorch/JAX interop. MJWarp implements MuJoCo's physics pipeline on GPUs, stepping large batches of environments concurrently via `mjw.step()`. Buffer capacities (`nconmax`, `naconmax`, `njmax`) must be configured explicitly, as buffer overflows invalidate simulation rollouts. Accurate throughput and latency measurement requires explicit warm-up runs, synchronization barriers (`wp.synchronize()`), and scaling sweeps across batch sizes. MJWarp integrates with learning platforms such as mjlab, MuJoCo Playground via MJX, and Isaac Lab via the Newton rigid-body solver. NVIDIA Warp Architecture and Capabilities NVIDIA Warp provides a Python-based interface for writing high-performance, GPU-accelerated SIMT kernels compiled via JIT into CUDA modules. It includes built-in primitives for spatial math, bounding volume hierarchies, hash grids, and sparse matrices. Kernels support automatic differentiation via tape-based forward recording and reverse adjoint execution. Deterministic execution modes were introduced in Warp 1.15 to address scheduler-dependent GPU atomic operations, trading execution speed for reproducible ordering. Warp also provides DLPack-style zero-copy interop with PyTorch and JAX, whereas operations like `.numpy()` require explicit synchronization and memory copies to the host. MJWarp Physics Execution and API MJWarp ports compatible MuJoCo models directly to NVIDIA GPUs. The workflow relies on loading standard MJCF files, where `mjw.put_model()` compiles supported features onto the device or raises an error if an incompatible feature is encountered. State data is managed through `mjw.make_data()` for fresh states or `mjw.put_data()` to preserve initialized host states. A single call to `mjw.step(m, d)` advances all parallel environments simultaneously using batched control arrays with shape `(nworld, nu)`. MJWarp's compact solver employs MuJoCo's native Newton constraint solver and sleeping mechanisms rather than the standalone Newton physics framework. Buffer Management and Constraint Tuning GPU memory allocation in MJWarp requires pre-allocating contact and constraint capacities across the environment batch using parameters such as `nworld`, per-world contact limits (`nconmax`), global contact caps (`naconmax`), and constraint limits (`njmax`). Because contact and constraint arrays are pre-allocated, exceeding these limits causes overflows that invalidate rollouts. By default, MJWarp logs buffer overflow warnings and flags affected worlds in simulation data. Developers can inspect capacities and overflows using `mjwarp-testspeed --measure_alloc` and the interactive `mjwarp-viewer`. Benchmarking and CUDA Graph Optimization Asynchronous GPU kernel execution requires warm-up steps to account for initial kernel compilation and allocation overheads. Timed benchmarks must call `wp.synchronize()` immediately before and after measured execution loops, including after `wp.capture_launch` graph replays. Simulation performance is characterized by reporting batched latency (milliseconds per step), batch size, and aggregate throughput in world-steps per second. While CUDA Graph capture allows in-place updates to control arrays between replays, rebuilding the model, resizing environments, or modifying allocated buffers requires capturing a new graph. Ecosystem and Framework Integrations MJWarp serves as an execution backend across diverse simulation and policy-training stacks. For pure MuJoCo setups, single-robot control remains suited for CPU execution, while high-throughput sampling shifts to MJWarp or mjlab. In machine learning workflows, JAX-based training is supported via MuJoCo Playground with MJX using the Warp implementation. Additionally, the Newton framework can manage robot models, assets, and inverse kinematics helpers while using MJWarp as its rigid-body solver (`newton.solvers.SolverMuJoCo`), offering a bridge to Isaac Lab workflows.

基于 NVIDIA Warp 与 MJWarp 加速机器人仿真与强化学习流程指南 NVIDIA Warp 与基于其构建的 MuJoCo Warp (MJWarp) 将传统 CPU 驱动的 MuJoCo 仿真迁移至 GPU，实现大规模并行物理模拟，并无缝对接主流机器学习训练流水线。 NVIDIA Warp 采用静态类型 Python 编写原生 CUDA 核函数，具备 JIT 编译、自动微分与 DLPack 互操作能力 MJWarp 负责将 MuJoCo 的 MJCF 模型与物理计算移至 GPU，通过单次 step 调用推进批量环境 仿真资源需严格分配接触与约束预算上限（如 nconmax、njmax），超限会导致当前 rollout 失效 上层训练支持接入 Isaac Lab（通过 Newton 框架）、mjlab 以及 MuJoCo Playground / MJX 基准测试需依赖预热与显式同步（wp.synchronize），并结合批大小评估聚合吞吐量 NVIDIA Warp 的架构与核心特性 NVIDIA Warp 是一个用于编写 GPU 加速核函数的高性能 Python 框架，开发者使用纯 Python 编写静态类型代码，系统会即时编译（JIT）为 CPU 或 CUDA 模块并进行缓存。Warp 支持向量、四元数、BVH 与哈希网格等原生结构，并提供可微计算（wp.Tape）与 DLPack 级别的跨框架数据共享。从 1.15 版本开始，Warp 引入了可选择的确定性执行模式，以微小性能折损换取确定的 GPU 调度顺序；同时需要注意，CUDA 数组调用 .numpy() 会触发内存同步与复制，并非零拷贝路径。 MJWarp 物理引擎执行机制 MJWarp 是 MuJoCo 物理流水线的 Warp 实现，由 MuJoCo 解析并编译 MJCF 模型，MJWarp 负责在 GPU 上推进状态。调用 mjw.step 即可同时更新整批环境。若模型使用了不支持的特性，mjw.put_model 会直接报错；若需将模拟状态迁移至设备端，fresh 状态使用 mjw.make_data，继承现有状态则需调用 mjw.put_data。当 qpos 与 qvel 拷回主机后，需执行 mujoco.mj_forward 重新计算由其导出的主机端变量（如 mjd.xpos）。 接触与约束缓冲区的容量管理 MJWarp 要求在运行前预先配置批次环境数（nworld）、单环境接触上限（nconmax）或全局接触上限（naconmax）以及单环境约束上限（njmax）。由于内存开销随预算缩放，需要借助 mjwarp-testspeed 的 --measure_alloc 选项及可视化工具调优参数。默认配置下超出容量会在终端提示并标记 Data.overflow，超限会导致测试验证失效。此外，CUDA Graphs 会固定复用捕获时的模型与数据缓冲区，原地更新控制变量 d.ctrl 有效，但更改批大小或重构模型后必须重新捕获。 上层学习框架与基准评估 在训练策略层面，单机器人控制和遥操作适合经典的 MuJoCo CPU 方案，纯物理的高吞吐场景推荐 MJWarp（或 mjlab），JAX 体系可采用 MuJoCo Playground / MJX，而结合 Isaac Lab 的多解算器流程则能通过 Newton 框架将 MJWarp 作为刚体解算器。进行性能评估时，由于 GPU 调度异步，必须先通过预热消除内核编译与缓存开销，并在时间记录窗口前后调用 wp.synchronize()，综合汇报批大小、单步时延（ms/step）以及每秒完成的总世界步数（throughput）。

Classic MuJoCo provides fast CPU-based robot simulation for developing, testing, and controlling robots and it can parallelize sampling across CPU cores. But as learning workloads grow, the question shifts from how quickly one world can run to how many worlds can run at once. GPU acceleration makes it possible to advance those worlds in large batches while keeping simulation and learning data close to the device. MuJoCo Warp (MJWarp) , built on NVIDIA Warp , takes compatible MuJoCo models into that GPU-scale regime. In this article, we will move an SO-101 follower arm from a familiar MuJoCo workflow to as many as 2,048 parallel MJWarp environments and examine the technology and validation steps that make the transition possible. Figure 1. How MJWarp connects Python to GPU simulation. MuJoCo loads and compiles the MJCF model; MJWarp implements the physics in NVIDIA Warp, which compiles CUDA kernels to advance simulation states on NVIDIA GPUs. This is the second article in our State of Simulation for Physical AI series. The first article mapped the robot-simulation landscape. Here, we prepare and scale the simulation environment; we do not train a policy. The later Newton and Isaac Lab installments cover the next integration layers. Putting it together Layer Role in the stack NVIDIA Warp Python kernel language: single instruction, multiple threads (SIMT), autodiff, PyTorch/JAX interop MJWarp MuJoCo physics on Warp: same MJCF, batched GPU throughput Your scene (SO-101) Familiar Menagerie / Robot Studio assets + task geometry Next (Newton / Isaac Lab) Multi-solver API, USD, sensors, managers, training loops Decision shortcut: If you need… Reach for… Single-robot MPC / teleop MuJoCo CPU Max throughput on raw MuJoCo physics MJWarp (or mjlab ) JAX training recipes MuJoCo Playground / MJX (impl='warp') Multi-solver + Isaac Lab integration Newton — next post in this series Start with one useful Warp Kernel NVIDIA Warp is a Python framework for writing high-performance, GPU-accelerated kernels. Warp lets developers author statically typed kernels in Python and compiles them for CPU or CUDA execution. The first launch builds and caches a native module; later launches reuse it. The kernel language is a performance-oriented subset of Python, while ordinary Python remains responsible for configuration, allocation, and launch orchestration. This small robotics-oriented kernel advances point positions under gravity. One logical thread handles one point, so the same code scales from two points to millions without introducing GPU terminology into the control flow. The three value propositions of Warp are: Pillar What you get Performance Native-CUDA speed via JIT compilation, kernel fusion, and CUDA Graphs Ease of use Pure Python authoring with built-in vectors, matrices, quaternions, BVHs, hash grids, sparse matrices, and tile primitives Capability Differentiable kernels and DLPack-style interop so simulation can sit inside an ML training loop 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()) Three properties make this useful in robotics: Explicit parallel work. wp.tid() identifies the point, contact, body, or world owned by the current logical thread. Explicit device arrays. An array lives on the selected device. Calling .numpy() on a CUDA array synchronizes and copies it to CPU memory; it is not a zero-copy path. For a device-resident PyTorch or JAX pipeline, use Warp’s framework adapters or DLPack-compatible sharing instead. Composable kernel launches. A program can launch a sequence of focused kernels and capture supported CUDA work into a graph to reduce repeated dispatch overhead. Graph capture replays launches against existing buffers; it does not fuse arbitrary kernels. Differentiability and Determinism. Two further Warp capabilities are worth knowing, even though neither is used in the SO-101 workflow in this article. Warp kernels are differentiable : a wp.Tape records the forward kernel launches made inside its context and replays their adjoints in reverse when backward() is called, which is why teams build differentiable geometry, CFD, and custom physics in Warp, including CAE workflows for simulation and design optimization. Warp also supports deterministic execution , introduced in Warp 1.15 : GPU atomics are scheduler-dependent by default, so repeated launches of the same kernel can differ slightly, and the opt-in deterministic modes trade some performance for reproducible ordering in simulation, validation, and regression tests. These are Warp capabilities, not guarantees of differentiability or determinism for an entire MJWarp rollout. See the Warp documentation on differentiability and deterministic execution for the details. Try Warp: pip install warp-lang (≥ 1.15 for GPU determinism), then python -m warp.examples.browse, or the tutorial notebooks . What is MuJoCo Warp (MJWarp)? A robot simulator repeatedly computes what happens next: given the current joint positions, velocities, controls, and contacts, it advances the scene by one small timestep. In this article, a world means one independent copy of that scene and its state. One world might contain the SO-101 arm reaching for a cube; another can contain the same arm starting from a slightly different pose. MuJoCo and MJWarp can run the same compatible robot and task, but they organize the work differently. MuJoCo naturally suits developing and inspecting one or a few CPU worlds. MJWarp is a NVIDIA Warp implementation of MuJoCo ’s physics pipeline that places the model and a batch of independent states on NVIDIA GPUs; one call to mjw.step advances the entire batch. MJWarp’s value is not necessarily a faster step for one world. It is the ability to advance hundreds or thousands together, giving the GPU enough parallel work to improve aggregate throughput , the total world-steps completed per second. That favors reinforcement learning and large-scale sampling, where collecting experience matters more than minimizing one environment’s latency. This blog covers the following: validate one MuJoCo world, move it to MJWarp, form a batch, verify it, and measure it correctly. Solver tuning, Jacobian representation, and specialized multi-GPU or determinism topics are not required for this migration and can be covered separately. Then, the distinction is precise: Latency is wall-clock time for one simulation step. Aggregate throughput is the total number of world-steps completed per measured wall-clock second. Basic usage: structs, batch sizes, and a minimal step The core API transition is small: MuJoCo host workflow MJWarp workflow mujoco.MjModel mjw.put_model(mjm) creates a device model mujoco.MjData mjw.put_data(mjm, mjd, ...) preserves and batches an existing state mujoco.mj_step(mjm, mjd) mjw.step(m, d) advances every world in d Host arrays such as mjd.ctrl Batched device arrays such as d.ctrl with shape (nworld, nu) Use mjw.make_data() when default/fresh state is intended. Use mjw.put_data() when the exact initialized MuJoCo state must cross the migration boundary. Allocating batched resources requires defining the following parameters (refer to Batch sizes ): Parameter Meaning nworld Total number of parallel environments nconmax Expected contacts per individual world (overall capacity ≈ nconmax * nworld) naconmax Alternative setting: global maximum contacts across all environments combined (takes precedence if both are defined) njmax Hard upper limit on constraints per world Performance tuning 1. CUDA graph capture :mjw.step is many kernel launches; capture once, replay often: with wp.ScopedCapture() as capture: mjw.step(m, d) wp.capture_launch(capture.graph) 2. Size nconmax / naconmax / njmax tightly : memory and work scale with them. Tune with mjwarp-testspeed: --measure_alloc and watch overflows in mjwarp-viewer. Additional tuning considerations. After sizing contact and constraint buffers, test solver iteration limits without changing task behavior. Meshes and CCD settings can increase memory use; nccdmax / naccdmax can reduce CCD buffer allocation when the measured contact counts allow it. MJWarp’s compact solver uses MuJoCo’s Newton constraint solver and sleeping, not the separate Newton physics-engine framework. Compact-solver and multi-GPU configuration are beyond this walkthrough; consult the MJWarp performance-tuning documentation. To train policies on MJWarp physics: Isaac Lab via Newton mjlab (manager API directly on MJWarp + PyTorch) MuJoCo Playground via MJX (impl='warp') Install / try: pip install mujoco-warp · mjwarp-viewer path/to/scene.xml · Colab tutorial Learn by doing. Explore our hands-on tutorial notebooks , from MuJoCo fundamentals and robot pick-and-place to GPU simulation with MJWarp and CPU/GPU benchmarking. Workflow to migrate a MuJoCo scene to MjWarp Establish a MuJoCo CPU baseline The scene. Nothing here is MJWarp-specific yet: an SO-101 arm, a table, and two cubes to stack, written as ordinary MJCF. Figure 2. SO-101 pick-and-place scene, rendered from the MuJoCo CPU simulation. The task is to grasp the red 44 mm cube and stack it on the blue cube; the same robot and scene are used for MJWarp validation. < 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" /> < body name = "red_cube" pos = "0.33 -0.13 0.046" > < 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" /> </ body > < body name = "blue_cube" pos = "0.33 0.06 0.046" > < 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" /> </ body > </ worldbody > </ mujoco > For an MJCF box, the size values are half-extents: size=”0.022 …” defines a cube with 44 mm edges. The task uses this size for its success thresholds. The arm base is at the origin, its reach is along +X, and the cubes are arranged along Y. In the companion repository this file is generated rather than hand-written: resolve_pick_place_scene() copies the Menagerie arm into .generated/, fills the table and cube coordinates from a robot profile, and writes scene_pick_place.xml. The walkthrough uses the SO-101 profile; the optional reBot variant is described below. Loading it. Compilation and stepping are ordinary MuJoCo: import mujoco mjm = mujoco.MjModel.from_xml_path( "scene_pick_place.xml" ) mjd = mujoco.MjData(mjm) fps = 50 # controller rate sim_substeps = 10 # physics steps per control frame frame_dt = 1.0 / fps mjm.opt.timestep = frame_dt / sim_substeps controller = PickPlaceController(spec=spec) # waypoints + damped-least-squares IK for _ in range ( 600 ): # 600 control frames ctrl = controller.step(mjm, mjd, frame_dt) for _ in range (sim_substeps): mjd.ctrl[: mjm.nu] = ctrl mujoco.mj_step(mjm, mjd) Keep that shape in mind: compute controls once per frame, step physics sim_substeps times. Gate 2 changes only the inner loop, which is what makes the migration easy to review. Match the simulation and control rates. At 50 control frames per second and 10 physics substeps per frame, use a physics timestep of 0.002 seconds. Set it before the CPU rollout and before uploading the model with mjw.put_model so both backends advance the same simulated time: mjm.opt.timestep = frame_dt / sim_substeps # 50 Hz × 10 substeps -> 0.002 s Without that line, every later measurement inherits the mismatch: parity comparisons, throughput numbers quoted as “simulated seconds,” and any learned policy whose action rate no longer matches deployment. Check whether the cubes are stacked successfully. With 44 mm cubes, success becomes two measurable conditions: a horizontal center error of xy_err ≤ 0.015 m (measured between the cube centers) and a vertical separation of 0.035 m ≤ dz ≤ 0.055 m between cube centers (one cube edge, with slack for settling). Evaluate both conditions after the cubes have settled; a successful process exit alone does not establish task success. Run the CPU task from the companion checkout. Publication blocker: confirm the accessible repository URL and pinned dependency and asset versions before publishing these instructions; the repository placeholder below is not an executable URL. git clone https://github.com/NVIDIA/accelerated-computing-hub.git blogs cd blogs/tutorials/sim2real-blogs/notebooks/mujoco uv venv --python 3.12 && source .venv/bin/activate uv pip install -r requirements.txt python part1/solutions/so101_pick_place_solution.py The run ends by printing the two numbers above (stack check: xy_err=… dz=…), which is the assertion the rest of the article compares against. so101_pick_place.py next to it is the same program with the physics steps left as exercises. The arm comes straight from MuJoCo Menagerie pinned to a known-good commit, since Menagerie assets change, so treat the scene as a template. Optional reBot variant. The companion code also exposes --robot rebot with a separate profile for the scene layout, gripper, and capacity limits (nconmax=256, njmax=500). This walkthrough uses SO-101. Validate the reBot asset and task separately before reporting its results. Validate one-world MJWarp parity Run one world on the GPU first, with the host still in the loop, so you can watch the same task in the same viewer and compare the same two numbers. Upload the model, allocate batched state, seed it from the initialized host state, and run one forward pass before stepping: wp.init() import mujoco_warp as mjw device = wp.get_device() m = mjw.put_model(mjm) d = mjw.make_data(mjm, nworld= 1 , nconmax=spec.nconmax, njmax=spec.njmax) 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 carries a leading world dimension, which is why the host state is indexed as mjd.qpos[None, :], shape (1, nq) instead of (nq,). Scaling to thousands of worlds later changes only that leading dimension, not the calls. mjw.put_model() also doubles as a compatibility check: it raises if the model uses unsupported features rather than silently dropping them. Seeding the three fields explicitly is the transparent option, and it makes clear exactly what crosses to the device; mjw.put_data(mjm, mjd, nworld=…) carries the whole initialized struct over in one call instead. The frame loop is then the Gate 1 loop with its inner step redirected to the GPU and mirrored back: 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) The .numpy() reads synchronize and copy data to the host on every substep, so this is a task-validation path, not a throughput benchmark. It keeps inverse kinematics, viewing, and task checks on the host. After copying qpos and qvel, call mujoco.mj_forward(mjm, mjd) to refresh derived host quantities such as mjd.xpos before using them for control, viewing, or the stack check. Reading those fields after the loop does not refresh them automatically. Gate 4 removes these per-step host copies from the throughput path. Size contact and constraint capacity MJWarp allocates contact and constraint buffers before stepping. Exceeding those capacities invalidates the affected rollout for verification or benchmarking, even when execution continues with an overflow warning rather than an exception. Increase the relevant limit and rerun the task. Larger buffers use more GPU memory, so verify capacity over the full task before tightening the allocation. Set contact and constraint limits for the robot and task being simulated. The SO-101 profile uses nconmax=128 and njmax=300 as starting capacities. Check that these limits are sufficient during the most contact-heavy part of the task: d = mjw.make_data(mjm, nworld=nworld, nconmax=spec.nconmax, njmax=spec.njmax) Size them against the most contact-heavy moment of the task, for pick-and-place, the instant both jaws and the table touch a cube, not the arm hovering in free space. An overflow is reported rather than raised: with Option.warn_overflow at its default, MJWarp prints the budget to increase (“narrowphase overflow - please increase nconmax to …”) to the terminal running your script or the viewer, and flags the affected worlds in Data.overflow for you to read back after a step. Only mjw.put_data raises an error outright, because it can compare the budgets against a MuJoCo state it already holds. mjwarp-testspeed --measure_alloc reports the contacts and constraints a scene actually consumed, and it aborts the rollout with the offending world IDs as soon as any world overflows. Treat those reports as failures: raise the limit and re-run before trusting either the trajectory or the benchmark, then tighten again whenever the model, collision geometry, or task changes. Scale to 2,048 worlds Once one-world parity passes, reallocate at the target size and replicate the initialized state across the batch. Two things change relative to Gate 2: nworld, and the fact that nothing crosses the PCIe bus per step. 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 starting state, which is the right baseline for a 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, and capture a new graph after replacing buffers, changing nworld, or rebuilding the model. Graph capture requires CUDA. Figure 3. Scaling the SO-101 task from one CPU world to 2,048 independent GPU states using the same compatible model. A single MJWarp step advances the full batch. This conceptual illustration highlights aggregate throughput, measured as world-steps per wall-clock second. Verify, then measure GPU launches are asynchronous, so a naive timer measures how fast Python queued work, not how fast the GPU finished it. Warm up first — the first launches pay kernel compilation and allocation — then synchronize immediately before and after the timed region: import time for _ in range ( 10 ): # warm-up: compilation, allocation, caches 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 total = 200 * nworld print ( f" {total / elapsed:, .0 f} world-steps/second" ) Report both aggregate world-steps per second and milliseconds per batched step, together with the batch size. Use the measured curve to identify where additional worlds improve throughput and where memory or compute limits reduce the benefit. Results depend on the scene, simulation settings, and hardware; a one-world latency comparison does not establish batched throughput. To see that curve on your own hardware, scaling_study.py sweeps the batch size and prints ms/step alongside throughput and speedup: cd part2 python solutions/so101_mjwarp_solution.py # needs CUDA python scaling_study.py --worlds 1 64 1024 2048 8192 --steps 100 Get started Warp (kernel layer) pip install warp-lang → python -m warp.examples.browse → docs · GitHub MJWarp (GPU MuJoCo) pip install mujoco-warp → mjwarp-viewer benchmarks/humanoid/humanoid.xml → docs · GitHub · Colab tutorial SO-101 context SO-101 sim-to-real course · Physical AI learning paths Train on top of MJWarp mjlab · MuJoCo Playground · Isaac Lab + Newton (upcoming posts) What’s next This post covered raw Warp → MJWarp : GPU kernels, batched stepping, and an SO-101 scene using mjw.step. Next, we will port the same MJCF environment into Newton , using MuJoCo Warp as its rigid-body solver (newton.solvers.SolverMuJoCo). Newton will manage the model, state, controls, and contacts, while MJWarp runs underneath. You will also see what Newton adds: multi-format assets, swappable solvers, sensors/IK helpers, and an Isaac Lab path. The migration guide continues with the same SO-101 task and its optional reBot profile, explaining the changes required by Newton and the separate Isaac Lab integration. If you build something with Warp or MJWarp, open an issue on the linked repositories or find us on Discord NVIDIA Omniverse . References Blog 1: * The State of Simulation for Physical AI: An Overview* — Post 1 of this series. NVIDIA Warp — GitHub · Documentation · v1.15.0 release (GPU determinism) · Deterministic execution guide MuJoCo Warp — GitHub · Official MJWarp docs Build Accelerated, Differentiable Computational Physics Code for AI with NVIDIA Warp Introducing Tile-Based Programming in Warp 1.5.0 mjlab · arXiv:2601.22074 MuJoCo Playground NVIDIA SO-101 sim-to-real course Newton next post: MJWarp as SolverMuJoCo and porting this environment