作者：苏晚把 MuJoCo 搬上 GPU：MJWarp 迁移的四道闸门与三种不报错的错误
NVIDIA 把一台 SO-101 从单个 MuJoCo CPU 世界搬到 2,048 个并行 MJWarp 环境，API 改动只有一张对照表。真正难的是搬完之后怎么确认没搬错：容量溢出只打印警告、.numpy() 静默同步、计时不同步测的是入队速度——三类错误都不抛异常，只让数字失真。
NVIDIA 的仿真团队 9 月 23 日在 Hugging Face 博客上更新了《State of Simulation for Physical AI》系列的第二篇——How to Use NVIDIA Warp and MjWarp to Accelerate Robotics Simulation and Learning Workflows，署名 Johnny Nuñez Cano、Asier Arranz、Rishabh Chadha、Ben Oliveri，四人均来自 NVIDIA。文章拿一台 SO-101 从臂（follower arm，遥操作主从对里跟随主臂动作的那一端）的常规 MuJoCo 流程出发，一路搬到 2,048 个并行的 MJWarp 环境里。
读完整篇会发现，它真正花力气写的并不是"怎么把代码改成 GPU 版"——API 层面的改动小到一张对照表就说完了。它花力气写的是搬完之后你凭什么相信搬对了：哪些错误不会抛异常、哪些测量方式会给出一个漂亮但无意义的数字、哪一行没写会让后面所有指标一起失准。对于国内正在把强化学习采样从 CPU 集群往单机多卡上收的团队，这后半部分才是稀缺的。下面按这条线索拆。
换的不是硬件，是问题
经典 MuJoCo 是一套很快的 CPU 物理引擎，用来开发、调试、控制一台机器人，也能把采样按 CPU 核数铺开。但当学习类工作负载变重，问题就换了：不再是"一个世界能跑多快",而是"能同时跑多少个世界"。
原文把这两件事的口径定得很死，值得照抄进自己的文档：
- 延迟（latency）：单次仿真步的墙上时钟时间。
- 聚合吞吐（aggregate throughput）：每一秒实测墙上时钟内完成的世界步（world-steps）总数。
MJWarp 的价值"不一定是让单个世界的一步变快"，而是能同时推进成百上千个世界，给 GPU 喂足并行度。这句话决定了它适合谁：强化学习和大规模采样看重的是攒经验的速度；单机器人 MPC、遥操作这类看重单环境延迟的活，留在 MuJoCo CPU 上更合适。原文给的选型捷径是一张分诊表——要在原生 MuJoCo 物理上榨吞吐用 MJWarp（或 mjlab），要 JAX 训练配方用 MuJoCo Playground / MJX（impl='warp'），要多求解器加 Isaac Lab 集成则等 Newton。
原文图 1：从 Python 到 GPU 的分层。图片来源：NVIDIA / Hugging Face Blog
Warp：一个 kernel 就能看懂的三条性质
NVIDIA Warp 是用 Python 写高性能 GPU kernel 的框架：静态类型的 kernel 由 Python 写成，编译到 CPU 或 CUDA 执行；首次启动会构建并缓存原生模块，之后复用——所以第一次跑慢不是性能问题，是编译。kernel 语言是 Python 的一个面向性能的子集，而配置、分配、启动编排仍然是普通 Python 的活。
原文给的最小示例是一个在重力下推进质点位置的 kernel：
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())
一个逻辑线程管一个质点，所以同一段代码从两个点扩到上百万个点，控制流里不需要出现任何 GPU 术语。原文提炼出三条对机器人仿真真正有用的性质：
一、并行粒度是显式的。 wp.tid() 标识当前逻辑线程所拥有的那个点、接触、刚体或世界。写代码时先想清楚"一个线程负责什么",后面扩批量就只是换个数字。
二、设备数组是显式的。 数组归属于选定的设备。对一个 CUDA 数组调 .numpy() 会同步并拷回 CPU 内存，这不是零拷贝路径。如果下游是常驻显存的 PyTorch 或 JAX 管线，要走 Warp 的框架适配器或 DLPack 风格的共享，而不是 .numpy() 绕一圈。这一条后面会变成一个具体的坑。
想了解 AI 如何助力您的企业？
免费获取企业 AI 成熟度诊断报告，发现转型机会
