NVIDIA Warp 与 MJWarp：加速机器人模拟与学习工作流
经典 MuJoCo 提供了基于 CPU 的快速机器人模拟，适用于开发、测试和控制机器人，并能在多核 CPU 间并行采样。然而，随着机器学习负载的增长，核心问题从“单个世界运行得有多快”转变为“能同时运行多少个世界”。GPU 加速使得大批量推进这些世界成为可能，同时将模拟和训练数据保持在设备端。
基于 NVIDIA Warp 构建的 MuJoCo Warp (MJWarp) 将兼容的 MuJoCo 模型带入 GPU 规模模式。本文以 SO-101 跟随臂为例，展示如何从熟悉的 MuJoCo 工作流迁移至多达 2,048 个并行 MJWarp 环境，并探讨实现这一过渡的技术细节与验证步骤。
Asier Arranz asiernvidia Rishabh Chadha rchadha-nv Ben Oliveri BenOliveri
这是关于物理 AI 模拟状态系列的第二篇文章。第一篇描绘了机器人模拟的整体景观，而本篇重点在于准备和扩展模拟环境，不涉及策略制定。后续文章将涵盖 Newton 和 Isaac Lab 的集成计划。
技术栈概览
| 层级 | 在堆栈中的角色 |
|---|---|
| NVIDIA Warp | Python 内核语言：单指令多线程 (SIMT)，自动微分，PyTorch/JAX 交互 |
| MJWarp | 在 Warp 上实现的 MuJoCo 物理引擎：支持相同的 MJCF 格式，提供批量 GPU 吞吐量 |
| 场景示例 (SO-101) | 使用熟悉的 Zoo/Robot Studio 资产及任务几何结构 |
| 下一步 (Newton/Isaac Lab) | 多求解器 API，USD 支持，传感器，管理器，训练循环 |
选择指南：
| 需求场景 | 推荐方案 |
|---|---|
| 单机器人 MPC / 遥操作 | MuJoCo CPU |
| 原始 MuJoCo 物理的最大吞吐量 | MJWarp (或 mjlab) |
| JAX 训练配方 | MuJoCo Playground / MJX (impl='warp') |
| 多求解器 + Isaac Lab 集成 | Newton (本系列下一篇文章) |
基础组件：NVIDIA Warp
NVIDIA Warp 是一个用于编写高性能、GPU 加速内核的 Python 框架。它允许开发者在 Python 中创建静态类型内核，并将其编译为 CPU 或 CUDA 执行代码。首次运行时构建并缓存本地模块，后续运行则复用该模块。其核心是一种面向性能的 Python 子集，而常规 Python 负责配置、内存分配和启动编排。
以下是一个简单的机器人导向内核示例，演示了在重力作用下推进点位置的过程。每个逻辑线程处理一个点，因此同一代码可从两个点扩展到数百万个点。
Warp 的三大核心价值主张：
| 支柱 | 优势 |
|---|---|
| 性能 | 通过 JIT 编译、内核融合和 CUDA Graphs 实现接近原生 CUDA的速度 |
| 易用性 | 纯 Python 创作体验，内置向量、矩阵、四元数、BVH、哈希网格、稀疏矩阵等原语 |
| 能力 | 可微分内核及 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())
机器人应用中的关键特性
- 明确的并行工作划分：识别当前逻辑线程所处理的点、接触、刚体或世界实例。
- 显式设备数组：数组驻留在选定设备上。对 CUDA 数组调用 .numpy() 会触发同步并复制到 CPU 内存，并非零拷贝路径。对于固定设备的 PyTorch 或 JAX 管道，建议使用 Warps 框架适配器或 DLPack 兼容共享。
- 可组合的内核启动：程序可按顺序启动多个内核，并将底层 CUDA 工作捕获到 Graph 中以减少重复提交开销。Graph 捕获重放现有缓冲区，但不会合并任意内核。
可微分性与确定性
尽管本文未涉及，但值得了解 Warp 的两个高级功能。曲线内核是可微分的：wp.Tape 记录上下文内的前向内核启动，并在调用反向传播时重放其辅助操作。这使得团队能够构建可微分几何、CFD 和自定义物理（包括 CAE 工作流）。此外，Warp 1.15 引入了确定性执行支持：默认情况下 GPU 原子操作依赖调度器，导致重复启动结果略有差异；启用确定性模式可在模拟、验证和回归测试中换取可重复排序，但会牺牲部分性能。请注意，这些是 Warp 的功能，不保证整个 MJWarp 的分化或决定性。
安装尝试：pip install warp-lang (≥1.15 版本支持 GPU 确定性)，随后使用 python -m warp.examples.browse 或教程笔记本。
MJWarp 详解
机器人模拟器反复计算未来状态：基于当前的关节位置、速度、控制和接触力。本文中，“一个世界”指独立复制的场景及其状态。例如，一个世界包含 SO-101 手臂伸向立方体，另一个世界包含相同手臂从略微不同的姿态开始。
MuJoCo 和 MJWarp 可运行相同的兼容机器人和任务。MuJoCo 适合开发和检查一两个 CPU 世界。MJWarp 则是 NVIDIA Warp 实现的 MuJoCo 物理管线，它将模型和一批独立状态放置在 NVIDIA GPU 上；调用 mjw.step 即可推进整个批次。
对于单个世界，MJWarp 的价值不一定体现在更快的单步时间，而在于能同时推进数百或数千步，从而让 GPU 拥有足够的并行工作量以提高总吞吐量（即每秒完成的世界步数）。这对强化学习和大规模采样尤为有利，因为收集经验比最小化单个环境的延迟更重要。
本文涵盖以下内容：
- 验证一个 MuJoCo 世界。
- 将其迁移至 MJWarp，形成批量环境。
- 检查并正确测量性能。
求解器调整、雅可比表示及专业的多 GPU 或确定性主题不在本次迁移讨论范围内。
关键区别定义：
- 延迟：单个模拟步骤的墙钟时间。
- 总吞吐量：每秒完成的世界步总数。
基本用法：结构、批量大小与最小步骤
核心 API 的迁移改动很小：
| MuJoCo 主机工作流 | MJWarp 工作流 |
|---|---|
| 无对应操作 | mjw.put_model(mjm) 创建设备模型 |
| 无对应操作 | mjw.put_data(mjm, mjd, ...) 保留并批量化现有状态 |
mujoco.mj_step(mjm, mjd) |
mjw.step(m, d) 推进 d 中的所有世界 |
主机数组如 mjd.ctrl |
批量化设备数组如 d.ctrl，形状为 (nworld, nu) |
若意图使用默认/全新状态，请使用 mjw.make_data()。若需精确迁移已初始化的 MuJoCo 状态，请使用 mjw.put_data()。
分配批量资源时需定义以下参数：
| 参数 | 含义 |
|---|---|
| nworld | 并行环境总数 |
| nconmax | 预期单个世界内的接触数（总容量 ≈ nconmax * nworld） |
| naconmax | 替代设置：所有环境合计的全局最大接触数（若两者均定义，此项优先） |
| njmax | 单个世界内约束数的硬性上限 |
性能调优
- CUDA Graph 捕获：
mjw.step涉及多次内核启动；建议捕获一次，频繁重放：
with wp.ScopedCapture() as capture:
mjw.step(m, d)
wp.capture_launch(capture.graph)
- 紧密设置 nconmax / naconmax / njmax：内存和工作量随这些值缩放。可使用
mjwarp-testspeed --measure_alloc进行调优，并在mjwarp-viewer中监控溢出情况。
其他调优考量：在确定接触和约束缓冲区大小后，测试求解器迭代限制而不改变任务行为。网格和 CCD 设置可能会增加内存使用；当测量的接触计数允许时，可通过 nccdmax / naccdmax 减少 CCD 缓冲区分配。MJWarp 的紧凑求解器使用 MuJoCo 的 Newton 约束求解器和休眠机制，而非独立的 Newton 物理引擎框架。紧凑求解器和多 GPU 配置超出本指南范围，请参考 MJWarp 性能调优文档。
在 MJWarp 物理基础上训练策略的途径：
- 通过 Newton 集成 Isaac Lab
- mjlab (直接在 MJWarp + PyTorch 上使用 Manager API)
- 通过 MJX (impl='warp') 使用 MuJoCo Playground
安装/试用：pip install mujoco-warp · mjwarp-viewer path/to/scene.xml · Colab 教程
迁移工作流：从 MuJoCo 场景到 MjWarp
1. 建立 MuJoCo CPU 基线
场景描述：此处尚未涉及 MJWarp 特定内容：一个 SO-101 机械臂、一张桌子和两个待堆叠的立方体，均以标准 MJCF 格式编写。
图 2. SO-101 抓取放置场景，由 MuJoCo CPU 模拟渲染。任务是抓住红色 44mm 立方体并将其堆叠在蓝色立方体上；相同的机器人和场景用于 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">
<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 name="blue_cube" pos="0.33 0.06 0.046">
<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"/>
对于 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
