NVIDIA Warp + MJWarp：把机器人仿真搬上 GPU，单机可并行 2048 环境

NVIDIA Warp 与 MuJoCo Warp（MJWarp）正在成为物理 AI 仿真栈中的关键加速层。Hugging Face 博客近期发布的《State of Simulation for Physical AI》系列第二篇文章，系统介绍了如何把 CPU 上的经典 MuJoCo 工作流迁移到 GPU 上，单机最多可并行 2,048 个独立环境，从而把「一个世界能跑多快」的问题转向「同时能跑多少个世界」。

背景：从 CPU 并行到 GPU 批处理

经典 MuJoCo 提供基于 CPU 的高速机器人仿真，并能在多核之间并行采样。但当学习负载增大，研究者更关心的是大规模批量推进场景的能力，而 GPU 加速让仿真与学习数据更贴近设备。MJWarp 基于 NVIDIA Warp 构建，把兼容的 MuJoCo 模型带入 GPU 规模。文章以 SO-101 跟随臂为示例，演示从熟悉的 MuJoCo 工作流迁移到 2,048 个并行 MJWarp 环境的过程，并梳理其中涉及的技术与验证步骤。

整个仿真栈的层次划分如下：

NVIDIA Warp：Python 内核语言，基于 SIMT 的并行模型，自带自动微分，与 PyTorch / JAX 互通。
MJWarp：在 Warp 上实现的 MuJoCo 物理引擎，MJCF 模型不变，批量吞吐由 GPU 承担。
场景层（SO-101）：沿用 Menagerie / Robot Studio 的资产和任务几何。
上层集成（Newton / Isaac Lab）：多求解器 API、USD、传感器、管理器与训练循环将在后续文章展开。

NVIDIA Warp 是什么

NVIDIA Warp 是一个用 Python 编写高性能 GPU 加速内核的框架。开发者以静态类型的方式在 Python 中编写内核，Warp 会把它编译为 CPU 或 CUDA 可执行文件；首次启动会构建并缓存原生模块，后续调用直接复用。其语言是面向性能的 Python 子集，而常规 Python 仍负责配置、分配与启动调度。

文章给出一个面向机器人的最小内核：在重力作用下推进点位置。每个逻辑线程负责一个点，因此同一段代码可以从两点扩展到上百万点，且无需在控制流中引入 GPU 术语。

Warp 的三大价值主张：

性能：通过 JIT 编译、内核融合与 CUDA Graphs，达到原生 CUDA 速度。
易用性：纯 Python 编写，自带向量、矩阵、四元数、BVH、哈希网格、稀疏矩阵与 tile 原语。
能力：可微分内核与类 DLPack 的互操作，使仿真可嵌入机器学习训练循环。

两个值得关注的特性是可微分与确定性执行（Warp 1.15 引入）。前者通过 wp.Tape 记录正向内核调用并在 backward() 时反向回放其伴随，常用于可微几何、CFD 与自定义物理；后者通过可选的确定性模式以部分性能换取可复现的调度顺序，便于仿真、验证与回归测试。但需要注意的是，这些是 Warp 自身的能力，并不能保证整个 MJWarp rollout 自动具备可微性或确定性。

MuJoCo Warp（MJWarp）定位

机器人仿真器反复执行同一件事：给定当前的关节位置、速度、控制量与接触，把场景向前推进一步。文章中将一个独立的「世界」定义为该场景及其状态的一份完整副本。一个世界可能包含 SO-101 机械臂去够一个方块，另一个则可能是相同机械臂从一个略微不同的初始位姿出发。

把多个这样的世界并行推进，正是 MJWarp 的核心场景。它复用原有 MJCF 模型描述，把物理计算放到 Warp 内核中执行，再通过 CUDA 在 GPU 上批量推进状态，从而把采样吞吐拉到一个 CPU 难以企及的量级。

选型决策：何时用哪个

文章给出了一张简明的工具选择对照表：

单机 MPC / 遥操：用 MuJoCo CPU，单机延迟敏感场景最简单。
在原始 MuJoCo 物理上追求最大吞吐：选 MJWarp（或 mjlab）。
JAX 训练配方：选 MuJoCo Playground / MJX（impl='warp'）。
需要多求解器 + Isaac Lab 集成：选 Newton，将在该系列下一篇文章展开。

小结

NVIDIA Warp 提供可微、Python 友好、与主流 ML 框架互通的内核层；MJWarp 在此之上把 MuJoCo 物理引擎带入 GPU 批处理模式，使得像 SO-101 这样的机器人学习任务可以在单机上同时跑起 2,048 个并行环境。对于正在搭建物理 AI 仿真栈的团队，理解 Warp 的并行模型、设备数组语义、CUDA Graphs 捕获方式，以及可微分与确定性执行的边界条件，是顺利把仿真搬上 GPU 的前提。系列后续文章将进一步覆盖 Newton 与 Isaac Lab 的集成层。