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

来源：https://huggingface.co/blog/nvidia/how-to-use-nvidia-warp-and-mjwarp

经典的 MuJoCo 提供了基于 CPU 的快速机器人仿真 (https://www.nvidia.com/en-us/use-cases/robotics-simulation/)，可用于开发、测试和控制机器人，并且可以在 CPU 核心之间并行采样。但随着学习工作负载的增长，问题从单个世界的运行速度转变为同时能运行多少个世界。GPU 加速使得可以大批量推进这些世界，同时让仿真和学习数据尽可能靠近设备。MuJoCo Warp (MJWarp) (https://mujoco.readthedocs.io/en/latest/mjwarp/) 基于 NVIDIA Warp (https://developer.nvidia.com/warp-python) 构建，将兼容的 MuJoCo 模型带入了 GPU 规模的计算领域。在本文中，我们将把一个 SO-101 机械臂从熟悉的 MuJoCo 工作流迁移到最多 2,048 个并行的 MJWarp 环境中，并探讨使这一转变成为可能的技术和验证步骤。

图片 1 (https://cdn-uploads.huggingface.co/production/uploads/6994dc99f850a10f03fd0b21/8SmNNsqLiBHC0W31zT4cf.png) 图 1. MJWarp 如何将 Python 与 GPU 仿真连接起来。MuJoCo 加载并编译 MJCF 模型；MJWarp 在 NVIDIA Warp 中实现物理引擎，Warp 编译 CUDA 内核以在 NVIDIA GPU 上推进仿真状态。

这是我们的“物理 AI 仿真现状”系列文章的第二篇。第一篇文章概述了机器人仿真领域。本文重点准备和扩展仿真环境；我们并不训练策略。后续的 Newton 和 Isaac Lab 系列将涵盖接下来的集成层次。

https://huggingface.co/blog/nvidia/how-to-use-nvidia-warp-and-mjwarp#putting-it-together整合起来

决策快捷指南：
如果你需要… 请参考…

• 单机器人 MPC / 遥操作 → MuJoCo CPU
• 原始 MuJoCo 物理的最大吞吐量 → MJWarp（或 mjlab (https://github.com/mujocolab/mjlab)）
• JAX 训练方法 → MuJoCo Playground (https://github.com/google-deepmind/mujoco_playground) / MJX (impl=‘warp’)
• 多求解器 + Isaac Lab 集成 → Newton——本系列的下一篇文章

https://huggingface.co/blog/nvidia/how-to-use-nvidia-warp-and-mjwarp#start-with-one-useful-warp-kernel从一个有用的 Warp 内核开始

NVIDIA Warp (https://github.com/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())

https://huggingface.co/blog/nvidia/how-to-use-nvidia-warp-and-mjwarp#three-properties-make-this-useful-in-robotics三个使其在机器人领域有用的特性：

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

https://huggingface.co/blog/nvidia/how-to-use-nvidia-warp-and-mjwarp#differentiability-and-determinism可微分性与确定性。

另外两个 Warp 能力值得一提，尽管它们在本文的 SO-101 工作流中都没有使用。Warp 内核是可微分的：wp.Tape 记录其上下文内进行的前向内核启动，并在调用 backward() 时反向重放其伴随，这就是为什么团队在 Warp 中构建可微分的几何、CFD 和自定义物理，包括用于仿真和设计优化的 CAE 工作流 (https://developer.nvidia.com/topics/cae)。Warp 还支持确定性执行，在 Warp 1.15 中引入：GPU 原子操作默认依赖调度器，因此重复启动同一个内核可能略有不同，而可选的确定性模式以部分性能为代价，在仿真、验证和回归测试中实现可重现的顺序。

这些是 Warp 的能力，并非整个 MJWarp rollout 的可微分性或确定性保证。详情请参阅 Warp 关于可微分性和确定性执行的文档。

尝试 Warp： pip install warp-lang（≥ 1.15 以获得 GPU 确定性），然后 python -m warp.examples.browse，或 教程笔记本 (https://github.com/NVIDIA/accelerated-computing-hub/tree/main/tutorials/warp/notebooks)。

https://huggingface.co/blog/nvidia/how-to-use-nvidia-warp-and-mjwarp#what-is-mujoco-warp-mjwarp什么是 MuJoCo Warp (MJWarp)？

机器人仿真器会反复计算下一步会发生什么：给定当前的关节位置、速度、控制信号和接触，它将场景推进一个小时间步长。在本文中，一个 world 指该场景及其状态的一个独立副本。一个世界可能包含 SO-101 机械臂伸手去抓一个立方体；另一个可以包含同一机械臂从略微不同的姿态开始。MuJoCo 和 MJWarp 可以运行相同的兼容机器人和任务，但它们组织工作的方式不同。

MuJoCo 自然适合在一个或几个 CPU 世界中进行开发和检查。MJWarp 是 MuJoCo 物理流水线的 NVIDIA Warp 实现，它将模型和一批独立状态放在 NVIDIA GPU 上；一次对 mjw.step 的调用就能推进整个批次。

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

本文涵盖以下内容：

1. 验证一个 MuJoCo 世界，
2. 将其迁移到 MJWarp，形成批次，
3. 正确验证和测量它。

求解器调整、雅可比表示以及专门的多 GPU 或确定性主题并非此迁移所必需，可以单独介绍。

然后，区分是精确的：

• 延迟 是一个仿真步的挂钟时间。
• 总吞吐量 是在测量的挂钟秒内完成的总世界步数。

https://huggingface.co/blog/nvidia/how-to-use-nvidia-warp-and-mjwarp#basic-usage-structs-batch-sizes-and-a-minimal-step基本用法：结构体、批大小和最小步

• 核心 API 转换很小：

主机数组，如 mjd.ctrl → 批处理设备数组，如 d.ctrl，形状为 (nworld, nu)

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

分配批处理资源需要定义以下参数（参考 批大小 (https://mujoco.readthedocs.io/en/latest/mjwarp/#batch-sizes)）：

https://huggingface.co/blog/nvidia/how-to-use-nvidia-warp-and-mjwarp#performance-tuning性能调优

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 物理引擎框架。紧凑求解器和多 GPU 配置超出本入门指南范围；请查阅 MJWarp 性能调优文档。

要在 MJWarp 物理上训练策略：

• Isaac Lab (https://github.com/isaac-sim/IsaacLab/tree/feature/newton) 通过 Newton (https://github.com/newton-physics/newton)
• mjlab (https://github.com/mujocolab/mjlab)（直接在 MJWarp + PyTorch 上的管理器 API）
• MuJoCo Playground (https://github.com/google-deepmind/mujoco_playground) 通过 MJX (impl=‘warp’)

安装/尝试： pip install mujoco-warp · mjwarp-viewer path/to/scene.xml · Colab 教程 (https://colab.research.google.com/github/google-deepmind/mujoco_warp/blob/main/notebooks/tutorial.ipynb)

https://huggingface.co/blog/nvidia/how-to-use-nvidia-warp-and-mjwarp#workflow-to-migrate-a-mujoco-scene-to-mjwarp将 MuJoCo 场景迁移到 MjWarp 的工作流

1. https://huggingface.co/blog/nvidia/how-to-use-nvidia-warp-and-mjwarp#establish-a-mujoco-cpu-baseline建立 MuJoCo CPU 基线

场景。 这里还不是 MJWarp 特定的：一个 SO-101 机械臂、一个桌子和两个要堆叠的立方体，作为普通 MJCF 编写。

图片 2 (https://cdn-uploads.huggingface.co/production/uploads/6994dc99f850a10f03fd0b21/Tn07F7pRaVdJXRnB7jSV1.png) 图 2. SO-101 抓放场景，从 MuJoCo CPU 仿真渲染。任务是抓取红色 44 毫米立方体并将其堆叠在蓝色立方体上；相同的机器人和场景用于 MJWarp 验证。

<!-- scene_pick_place.xml -->

对于 MJCF 盒子，size 值是半尺寸：size="0.022 ..." 定义了一个边长为 44 毫米的立方体。任务使用此尺寸作为其成功阈值。

机械臂基座位于原点，其伸展方向沿 +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) # 路径点 + 阻尼最小二乘 IK

for _ in range(600): # 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 # 50 Hz × 10 子步 -> 0.002 s

没有这一行，后续的每个测量都会继承这种不匹配：奇偶性比较、以“仿真秒”为单位的吞吐量数字，以及任何动作速率不再匹配部署的习得策略。

检查立方体是否成功堆叠。 对于 44 毫米的立方体，成功变为两个可测量的条件：水平中心误差 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

运行结束时打印上述两个数字（stack check: xy_err=... dz=...），这是本文其余部分进行比较的断言。旁边的 so101_pick_place.py 是同一个程序，但物理步骤留作练习。

该机械臂直接来自 MuJoCo Menagerie (https://github.com/google-deepmind/mujoco_menagerie/tree/main/robotstudio_so101)，固定在一个已知良好的提交版本，因为 Menagerie 资源会变化，所以将场景视为模板。

可选 reBot 变体。配套代码也提供 --robot rebot，具有用于场景布局、抓手和容量限制（nconmax=256, njmax=500）的单独配置文件。本演练使用 SO-101。在报告其结果之前，请单独验证 reBot 资产和任务。

2. https://huggingface.co/blog/nvidia/how-to-use-nvidia-warp-and-mjwarp#validate-one-world-mjwarp-parity验证单个世界的 MJWarp 一致性

首先在 GPU 上运行 一个 世界，主机仍在循环中，以便您可以在同一个查看器中观察相同的任务并比较相同的两个数字。上传模型，分配…