NVIDIA Warp 与 MjWarp 加速机器人仿真工作流

MuJoCo Warp（MJWarp）基于 NVIDIA Warp 构建，可将兼容的 MuJoCo 模型带入 GPU 规模化运行阶段。在本文中，我们将把一个 SO-101 机械臂从熟悉的 MuJoCo 工作流迁移到多达 2048 个并行 MJWarp 环境，并探讨使这一迁移成为可能的技术与验证步骤。

图 1. MJWarp 如何连接 Python 与 GPU 仿真。MuJoCo 加载并编译 MJCF 模型；MJWarp 使用 NVIDIA Warp 实现物理引擎，后者编译 CUDA 内核以在 NVIDIA GPU 上推进仿真状态。

这是《Physical AI 仿真现状》系列文章的第二篇。第一篇文章梳理了机器人仿真领域的全景图。本文我们将准备并扩展仿真环境；不涉及策略训练。后续的 Newton 和 Isaac Lab 部分将介绍下一层级的集成。

从一个实用的 Warp 内核开始

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

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

Warp 的三个核心价值主张如下：

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() 时按逆序重放其伴随（adjoint），这也是多个团队在 Warp 中构建可微几何、计算流体力学（CFD）和自定义物理引擎的原因，包括用于仿真与设计优化的 CAE 工作流。Warp 1.15 还引入了确定性执行支持：GPU 原子操作默认依赖调度器，因此同一内核的重复启动可能略有差异，可选的确定性模式会牺牲部分性能以换取仿真、验证和回归测试中的可重现顺序。这些是 Warp 的特性，而非整个 MJWarp 推出的可微性或确定性保证。详情请参阅 Warp 文档中关于可微性和确定性执行的内容。

体验 Warp：pip install warp-lang（≥ 1.15 支持 GPU 确定性），然后运行 python -m warp.examples.browse，或参阅教程笔记本。

什么是 MuJoCo Warp（MJWarp）？

机器人仿真器反复计算下一步会发生什么：给定当前的关节位置、速度、控制量和接触状态，它将场景推进一个小时间步。在本文中，一个 world 意味着该场景及其状态的一个独立副本。一个 world 可能包含 SO-101 机械臂伸手去抓取一个立方体；另一个 world 可以包含从略微不同姿态开始的同一机械臂。

MuJoCo 和 MJWarp 可以运行相同的兼容机器人和任务，但它们的工作组织方式不同。MuJoCo 自然适用于开发和检查一到数个 CPU world。MJWarp 是 MuJoCo 物理引擎管道的 NVIDIA Warp 实现，它将模型和一批独立状态放置在 NVIDIA GPU 上；一次调用 mjw.step 即可推进整个批次。

MJWarp 的价值不一定体现在单个 world 的更快步骤上。它真正价值在于能够同时推进数百或数千个 world，使 GPU 获得足够的并行工作来提升聚合吞吐量（aggregate throughput），即每秒完成的世界步（world-step）总数。这有利于强化学习和大尺度采样，因为在这些场景中收集经验比最小化单个环境的延迟更为重要。

本文涵盖以下内容：

• 验证一个 MuJoCo world
• 将其迁移到 MJWarp，形成批次
• 验证结果并正确测量

求解器调优、雅可比（Jacobian）表示以及专门的多 GPU 或确定性话题不属于本次迁移的必要内容，可以单独讨论。

然后，明确区分以下概念：

延迟（Latency）：一次仿真步的挂钟时间。

聚合吞吐量（Aggregate throughput）：每秒测量到的挂钟时间内完成的 world-step 总数。

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

核心 API 迁移很简单：

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

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

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 物理引擎上训练策略：

• mjlab：直接在 MJWarp + PyTorch 上的管理器 API
• MuJoCo Playground：通过 MJX（impl='warp'）

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

将 MuJoCo 场景迁移到 MJWarp 的工作流

建立 MuJoCo CPU 基线

建立 MuJoCo CPU 基线

场景。这里还没有任何 MJWarp 特定的内容：一个 SO-101 机械臂、一张桌子和两个待堆叠的立方体，以普通的 MJCF 格式编写。

图 2. SO-101 取放场景，从 MuJoCo CPU 仿真渲染。任务是抓住红色 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="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 # 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)

牢记这个结构：每帧计算一次控制量，物理模拟步进 sim_substeps 次。Gate 2 只改变内层循环，这就是迁移易于审查的原因。

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

mjm.opt.timestep = frame_dt / sim_substeps # 50 Hz × 10 substeps -> 0.002 s

没有这行代码，后续所有测量都会继承这个不匹配：奇偶性比较、以"模拟秒"计的吞吐量数字，以及任何动作速率不再匹配部署的学习策略。

检查立方体是否堆叠成功。对于 44 mm 的立方体，成功变成两个可测量条件：水平中心误差 xy_err ≤ 0.015 m（在立方体中心之间测量）和立方体中心之间的垂直间距 0.035 m ≤ dz ≤ 0.055 m（一个立方体边长，留有沉降的余量）。在立方体沉降后评估这两个条件；成功的进程退出本身并不代表任务成功。

从配套仓库检出运行 CPU 任务。发布阻断项：在发布这些说明之前确认可访问的仓库 URL 和固定的依赖及资产版本；下面的仓库占位符不是可执行 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

运行结束时打印上述两个数字（堆叠检查：xy_err=… dz=…），这是文章其余部分进行比较的断言。旁边的 so101_pick_place.py 是同一程序，物理步进部分留作练习。

机械臂直接来自固定在已知良好提交上的 MuJoCo Menagerie，因为 Menagerie 资产会发生变化，所以将场景视为模板。可选的 reBot 变体。配套代码还暴露了 --robot rebot 选项，带有单独的场景布局、夹爪和容量限制配置（nconmax=256，njmax=500）。本教程使用 SO-101。在报告 reBot 资产和任务结果之前单独验证它们。

验证单世界 MJWarp 奇偶性

验证单世界 MJWarp 奇偶性

首先在 GPU 上运行一个世界，同时主机仍在循环中，这样你可以观察同一任务在同一个查看器中的情况，并比较相同的两个数字。上传模型，分配批处理状态，从初始化的主机状态播种，然后在下之前运行一次前向传递：

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)

每个设备数组都带有一个领先的世界维度，这就是主机状态被索引为 mjd.qpos[None, :]（形状 (1, nq) 而不是 (nq,) 的原因。扩展到数千个世界后只改变那个领先维度，而不改变调用。mjw.put_model() 也兼作兼容性检查：如果模型使用了不支持的功能，它会抛出异常而不是静默丢弃。

显式播种这三个字段是透明的做法，它清楚地表明了什么数据跨设备传输；mjw.put_data(mjm, mjd, nworld=…) 在一次调用中携带整个初始化结构传输过去。

帧循环然后是 Gate 1 循环，将其内部步进重定向到 GPU 并镜像回来：

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)

.numpy() 读取会同步并在每个子步将数据复制回主机，所以这是一个任务验证路径，而非吞吐量基准。它将逆运动学、查看和任务检查保留在主机上。复制 qpos 和 qvel 后，调用 mujoco.mj_forward(mjm, mjd) 来刷新派生的主机量（如 mjd.xpos），然后再将它们用于控制、查看或堆叠检查。在循环之后读取这些字段不会自动刷新它们。Gate 4 从吞吐量路径中移除这些每步主机复制操作。

调整接触和约束容量

调整接触和约束容量

MJWarp 在步进之前分配接触和约束缓冲区。超出这些容量会使受影响的 rollout 无效，无法用于验证或基准测试，即使执行继续（带有溢出警告而非异常）。增加相关限制并重新运行任务。更大的缓冲区会使用更多 GPU 内存，因此在收紧分配之前验证整个任务的容量。

为正在模拟的机器人和任务设置接触和约束限制。SO-101 配置使用 nconmax=128 和 njmax=300 作为起始容量。在任务接触最密集的部分检查这些限制是否足够：

d = mjw.make_data(mjm, nworld=nworld, nconmax=spec.nconmax, njmax=spec.njmax)

根据任务接触最密集的时刻来调整它们，对于拾放任务，是两个夹爪和桌面同时接触立方体的瞬间，而不是机械臂在自由空间悬停的时候。溢出是报告而非抛出：默认的 Option.warn_overflow 下，MJWarp 打印增加预算的提示（"narrowphase overflow - please increase nconmax to …"）到运行脚本的终端或查看器，并在 Data.overflow 中标记受影响的 world 以便你在步进后读取回来。只有 mjw.put_data 会直接抛出错误，因为它可以将预算与已有的 MuJoCo 状态进行比较。mjwarp-testspeed --measure_alloc 报告场景实际消耗的接触和约束数量，并在任何 world 溢出时立即中止 rollout，并在输出中包含违规的 world ID。将这些报告视为失败：提高限制并重新运行，然后再信任轨迹或基准测试；每当模型、碰撞几何或任务发生变化时，再次收紧。

扩展到 2,048 个世界

扩展到 2,048 个世界

一旦单世界奇偶性通过，就在目标规模上重新分配，并将初始化状态复制到整个批次。与 Gate 2 相比，有两件事发生变化：nworld，以及每步没有任何数据跨越 PCIe 总线。

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 给每个世界相同的起始状态，这是吞吐量测量的正确基线；每个世界独立随机化则会在设备上写入 d.qpos 的不同行。

CUDA Graph 重用此处捕获的模型和数据缓冲区。在重放之间就地更新 d.ctrl，并在更换缓冲区、更改 nworld 或重建模型后捕获新的 graph。Graph 捕获需要 CUDA。

图 3. 使用相同的兼容模型，将 SO-101 任务从一个 CPU 世界扩展到 2,048 个独立 GPU 状态。一次 MJWarp 步骤即可推进整个批次。此概念图突出展示聚合吞吐量，以每墙上时钟秒的世界步数来衡量。

GPU 启动是异步的，因此简单的计时器测量的是 Python 排队工作的速度，而不是 GPU 完成的速度。首先需要预热——首次启动需要支付内核编译和分配的开销——然后在计时代码区域的前后立即同步：

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:,.0f} world-steps/second")

同时报告聚合的每秒世界步数和每个批处理步数的毫秒数，以及批量大小。使用测量曲线来识别额外的世界在哪些地方提高吞吐量，在哪些地方内存或计算限制会降低收益。结果取决于场景、仿真设置和硬件；单世界延迟比较无法确立批处理吞吐量。

若要在自己的硬件上查看该曲线，scaling_study.py 会扫描批量大小并打印 ms/step 以及吞吐量和加速比：

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

Warp（内核层）pip install warp-lang → python -m warp.examples.browse → 文档 · GitHub

MJWarp（GPU MuJoCo）pip install mujoco-warp → mjwarp-viewer benchmarks/humanoid/humanoid.xml → 文档 · GitHub · Colab 教程

SO-101 背景 SO-101 sim-to-real 课程 · Physical AI 学习路径

基于 MJWarp 训练 mjlab · MuJoCo Playground · Isaac Lab + Newton（后续文章）

本文涵盖了原始 Warp → MJWarp：GPU 内核、批处理步进以及使用 mjw.step 的 SO-101 场景。

接下来，我们将把同一个 MJCF 环境迁移到 Newton 中，使用 MuJoCo Warp 作为其刚体求解器（newton.solvers.SolverMuJoCo）。Newton 将管理模型、状态、控制和接触，而 MJWarp 在其下运行。

你还将看到 Newton 新增的功能：多格式资产、可交换求解器、传感器/IK 辅助工具以及 Isaac Lab 路径。

迁移指南继续使用相同的 SO-101 任务及其可选的 reBot 配置文件，说明 Newton 和单独的 Isaac Lab 集成所需的变化。

如果你使用 Warp 或 MJWarp 构建了任何东西，请在链接的仓库上提 issue，或在 Discord NVIDIA Omniverse 上联系我们。

博客 1：Physical AI 仿真现状概述 — 本系列第 1 篇。

NVIDIA Warp — GitHub · 文档 · v1.15.0 发布（GPU 确定性）· 确定性执行指南

MuJoCo Warp — GitHub · 官方 MJWarp 文档

使用 NVIDIA Warp 为 AI 构建加速、可微分的计算物理代码

在 Warp 1.5.0 中引入基于瓦片的编程

mjlab · arXiv:2601.22074

NVIDIA SO-101 sim-to-real 课程

Newton 下一篇：MJWarp 作为 SolverMuJoCo 以及迁移此环境

更多作者文章

了解谁在何时发言：使用 NVIDIA Nemotron 3 Diari ization 构建实时多说话人 AI

使用 NVIDIA Magpie TTS 构建低延迟多语言语音代理：开放权重和完整部署控制

el manejo de los agentes de ia para el entrenamiendo de rbajadores en simulacion con ia a sido bastante buena la ia ya que aprte de eso nos ahyuda hacerlo mas rapido y me jorar las cosas que nosotros como personas nos ayuda

· 注册或登录以发表评论