如何使用 NVIDIA Warp 和 MjWarp 加速机器人仿真与学习工作流

经典 MuJoCo 为开发、测试和控制机器人提供了快速的 CPU 仿真，并且能够在多个 CPU 核心之间并行化采样。然而，随着学习工作负载不断增长，问题从单个世界能够运行多快，转变为能够同时运行多少个世界。GPU 加速使得大量世界可以批量推进，同时让仿真和学习数据始终靠近设备端。

MuJoCo Warp（MJWarp）基于 NVIDIA Warp 构建，可将兼容的 MuJoCo 模型带入 GPU 规模化运行的范畴。本文将把一个 SO-101 从臂从熟悉的 MuJoCo 工作流迁移至最多 2,048 个并行 MJWarp 环境，并考察实现这一迁移所需的技术与验证步骤。

图 1. MJWarp 如何将 Python 连接至 GPU 仿真。MuJoCo 加载并编译 MJCF 模型；MJWarp 在 NVIDIA Warp 中实现物理计算，后者会编译 CUDA 内核以在 NVIDIA GPU 上推进仿真状态。

这是“物理 AI 仿真现状”系列的第二篇文章。第一篇文章梳理了机器人仿真的技术版图。本文将准备和扩展仿真环境，而不训练策略。后续关于 Newton 和 Isaac Lab 的文章将介绍下一层集成。

整合各层组件

层级	在技术栈中的角色
NVIDIA Warp	Python 内核语言：单指令多线程（SIMT）、自动微分、PyTorch/JAX 互操作
MJWarp	基于 Warp 的 MuJoCo 物理：相同的 MJCF、批量 GPU 吞吐量
你的场景（SO-101）	熟悉的 Menagerie / Robot Studio 资产与任务几何体
下一步（Newton / Isaac Lab）	多求解器 API、USD、传感器、管理器、训练循环

快速决策：

如果你需要……	建议使用……
单机器人 MPC / 远程操控	MuJoCo CPU
原始 MuJoCo 物理的最高吞吐量	MJWarp（或 mjlab）
JAX 训练方案	MuJoCo Playground / MJX（impl='warp'）
多求解器 + Isaac Lab 集成	Newton —— 本系列下一篇文章

从一个实用的 Warp Kernel 开始

NVIDIA Warp 是一个用于编写高性能 GPU 加速内核的 Python 框架。Warp 让开发者可以使用 Python 编写静态类型内核，并将其编译为 CPU 或 CUDA 执行代码。首次启动会构建并缓存原生模块；后续启动则会复用该模块。内核语言是一个以性能为导向的 Python 子集，而普通 Python 仍负责配置、分配和启动编排。

下面这个面向机器人应用的小型内核会在重力作用下推进点的位置。每个逻辑线程处理一个点，因此相同代码可从两个点扩展到数百万个点，无需在控制流中引入 GPU 术语。

Warp 的三项核心价值主张是：

支柱	你获得的能力
性能	通过 JIT 编译、内核融合和 CUDA Graphs 获得原生 CUDA 级速度
易用性	纯 Python 编写方式，内置向量、矩阵、四元数、BVH、哈希网格、稀疏矩阵和 tile 原语
能力	可微内核和 DLPack 风格互操作，使仿真可置于 ML 训练循环中

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())

三项特性使其适用于机器人领域：

• 显式并行工作。 wp.tid() 用于标识当前逻辑线程负责的点、接触点、刚体或世界。
• 显式设备数组。 数组驻留在选定设备上。对 CUDA 数组调用 .numpy() 会进行同步，并将其复制到 CPU 内存；这并非零拷贝路径。对于驻留在设备端的 PyTorch 或 JAX 流水线，应改用 Warp 的框架适配器或兼容 DLPack 的共享方式。
• 可组合的内核启动。 程序可以启动一系列专注的内核，并将受支持的 CUDA 工作捕获至图中，以降低重复分发开销。图捕获会针对现有缓冲区重放启动操作；它不会融合任意内核。

可微性与确定性。

另有两项 Warp 能力值得了解，尽管本文的 SO-101 工作流并未使用它们。Warp 内核具有可微性：wp.Tape 会记录其上下文中执行的前向内核启动，并在调用 backward() 时以反向顺序重放其伴随计算；因此，团队会使用 Warp 构建可微几何、CFD 和自定义物理系统，其中也包括用于仿真与设计优化的 CAE 工作流。Warp 还支持确定性执行，该功能在 Warp 1.15 中引入：默认情况下，GPU 原子操作依赖调度器，因此对同一内核的重复启动可能略有差异；选择启用确定性模式会以部分性能为代价，为仿真、验证和回归测试提供可复现的执行顺序。这些是 Warp 的能力，而不是对整个 MJWarp rollout 的可微性或确定性的保证。详细信息请参阅 Warp 文档中关于可微性和确定性执行的说明。

试用 Warp：pip install warp-lang（GPU 确定性需要 ≥ 1.15），然后运行 python -m warp.examples.browse，或查看教程笔记本。

什么是 MuJoCo Warp（MJWarp）？

机器人仿真器会重复计算下一时刻会发生什么：给定当前的关节位置、速度、控制输入和接触信息，它会将场景推进一个很小的时间步长。本文中，一个**世界（world）**是指该场景及其状态的一份独立副本。一个世界可能包含正在伸手抓取立方体的 SO-101 机械臂；另一个世界则可以包含以略有不同姿态开始的同一机械臂。

MuJoCo 和 MJWarp 可以运行同一个兼容的机器人与任务，但二者组织计算工作的方式不同。MuJoCo 天然适合开发和检查一个或少数几个 CPU 世界。MJWarp 是在 NVIDIA Warp 中实现的 MuJoCo 物理流水线，它会将模型和一批独立状态放置在 NVIDIA GPU 上；一次调用 mjw.step 就可以推进整个批次。

MJWarp 的价值不一定在于让单个世界的步进更快，而在于能够同时推进数百乃至数千个世界，为 GPU 提供足够多的并行工作，从而提高总吞吐量（aggregate throughput），即每秒完成的世界步进总数。这有利于强化学习和大规模采样场景，在这些场景中，收集经验比最小化单个环境的延迟更重要。

本文涵盖以下内容：

1. 验证一个 MuJoCo 世界；
2. 将其迁移至 MJWarp，并组成批次；
3. 验证结果，并以正确方式测量性能。

求解器调优、雅可比矩阵表示，以及专门的多 GPU 或确定性主题并不是此次迁移的必要内容，可单独讨论。

因此，二者的区别可以精确定义为：

• **延迟（Latency）**是一次仿真步进所需的墙钟时间。
• **总吞吐量（Aggregate throughput）**是在测量到的每墙钟秒内完成的世界步进总数。

基础用法：结构体、批次大小与最小步进

• 核心 API 的迁移幅度很小：

MuJoCo 主机端工作流	MJWarp 工作流
mujoco.MjModel	mjw.put_model(mjm) 创建设备端模型
mujoco.MjData	mjw.put_data(mjm, mjd, ...) 保留并批处理现有状态
mujoco.mj_step(mjm, mjd)	mjw.step(m, d) 推进 d 中的每一个世界
诸如 mjd.ctrl 的主机端数组	形如 (nworld, nu) 的批量设备端数组，如 d.ctrl

当需要默认或全新的状态时，使用 mjw.make_data()。当必须让精确初始化后的 MuJoCo 状态跨越迁移边界时，使用 mjw.put_data()。

分配批量资源需要定义以下参数（参阅批次大小）：

参数	含义
nworld	并行环境总数
nconmax	每个独立世界的预期接触数（总容量约为 nconmax * nworld）
naconmax	替代设置：所有环境合计的全局最大接触数（如两者均被定义，则该值优先）
njmax	每个世界的约束硬上限

性能调优

1. CUDA 图捕获：mjw.step 包含许多内核启动操作；应捕获一次、经常重放：

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

2. **紧凑设置 nconmax / naconmax / njmax：**内存占用和计算工作量会随这些参数增长。使用 mjwarp-testspeed --measure_alloc 调优，并在 mjwarp-viewer 中观察溢出情况。

其他调优考虑事项。 在设定接触和约束缓冲区大小后，应在不改变任务行为的前提下测试求解器迭代上限。网格和 CCD 设置可能增加内存使用；当实测接触数允许时，nccdmax / naccdmax 可以减少 CCD 缓冲区分配。MJWarp 的紧凑求解器使用 MuJoCo 的 Newton 约束求解器和休眠机制，而不是独立的 Newton 物理引擎框架。紧凑求解器和多 GPU 配置超出了本演练范围；请查阅 MJWarp 性能调优文档。

要在 MJWarp 物理上训练策略：

安装 / 试用：pip install mujoco-warp · mjwarp-viewer path/to/scene.xml · Colab 教程

将 MuJoCo 场景迁移至 MjWarp 的工作流

1. 建立 MuJoCo CPU 基线

场景。 此处尚无任何 MJWarp 特有内容：一个 SO-101 机械臂、一张桌子，以及两个用于堆叠的立方体，均以普通 MJCF 编写。

图 2. 从 MuJoCo CPU 仿真渲染的 SO-101 抓取放置场景。任务是抓取红色的 44 mm 立方体，并将其堆叠在蓝色立方体上；同一机器人与场景将用于 MJWarp 验证。

<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>

对于 MJCF 盒体，size 的值是半边长：size="0.022 …" 定义了边长为 44 mm 的立方体。任务使用这一尺寸设定成功阈值。机械臂底座位于原点，其可达范围沿 +X 方向，立方体则沿 Y 方向排列。

在配套代码库中，该文件并非手动编写，而是自动生成：resolve_pick_place_scene() 会将 Menagerie 机械臂复制到 .generated/，从机器人配置文件中填充桌子和立方体坐标，并写入 scene_pick_place.xml。本演练使用 SO-101 配置文件；下文将介绍可选的 reBot 变体。

加载场景。 编译和步进均使用常规 MuJoCo：

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)

请记住这一结构：每帧计算一次控制输入，物理仿真步进 sim_substeps 次。 第 2 个阶段仅改变内层循环，这正是迁移易于审查的原因。

匹配仿真与控制速率。 在每秒 50 个控制帧、每个控制帧 10 个物理子步的条件下，应使用 0.002 秒的物理时间步长。在 CPU rollout 之前，以及通过 mjw.put_model 上传模型之前设置它，以确保两个后端推进相同的仿真时间：

mjm.opt.timestep = frame_dt / sim_substeps 

如果缺少这一行，后续每项测量都会继承这一不匹配：包括一致性对比、以“仿真秒数”表述的吞吐量数据，以及动作频率不再匹配部署环境的任何已训练策略。

**检查是否

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:

1. validate one MuJoCo world,
2. move it to MJWarp, form a batch,
3. 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:

Install / try: pip install mujoco-warp · mjwarp-viewer path/to/scene.xml · Colab tutorial

Workflow to migrate a MuJoCo scene to MjWarp

1. 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 
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)

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 

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

cd /tutorials/sim2real-blogs/notebooks/mujoco
python solutions/so101_pick_place_solution.py --headless-steps 600 --debug

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.

2. 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.

3. 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.

4. 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.

5. 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): 
    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. 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 /tutorials/sim2real-blogs/notebooks/mujoco/part2
python solutions/so101_mjwarp_solution.py --headless-steps 600     # parity, 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