How to Use NVIDIA Warp and MjWarp to Accelerate Robotics Simulation and Learning Workflows

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

Classic MuJoCo provides fast CPU-based robot simulation for developing, testing, and controlling robots and it can parallelize sampling across CPU cores. But as learning workloads grow, the question shifts from how quickly one world can run to how many worlds can run at once. GPU acceleration makes it possible to advance those worlds in large batches while keeping simulation and learning data close to the device.

经典的 MuJoCo 为机器人开发、测试和控制提供了快速的基于 CPU 的仿真，并且可以在 CPU 核心之间并行化采样。但随着学习工作负载的增加，问题从“一个世界能运行多快”转变为“一次能同时运行多少个世界”。GPU 加速使得在大规模批次中推进这些世界成为可能，同时将仿真和学习数据保持在靠近设备端的位置。

MuJoCo Warp (MJWarp), built on NVIDIA Warp, takes compatible MuJoCo models into that GPU-scale regime. In this article, we will move an SO-101 follower arm from a familiar MuJoCo workflow to as many as 2,048 parallel MJWarp environments and examine the technology and validation steps that make the transition possible.

基于 NVIDIA Warp 构建的 MuJoCo Warp (MJWarp) 将兼容的 MuJoCo 模型带入了 GPU 规模的领域。在本文中，我们将把一个 SO-101 从动机械臂从熟悉的 MuJoCo 工作流迁移到多达 2,048 个并行 MJWarp 环境中，并探讨实现这一转变的技术和验证步骤。

Figure 1. How MJWarp connects Python to GPU simulation.

图 1. MJWarp 如何将 Python 连接到 GPU 仿真。

MuJoCo loads and compiles the MJCF model; MJWarp implements the physics in NVIDIA Warp, which compiles CUDA kernels to advance simulation states on NVIDIA GPUs. This is the second article in our State of Simulation for Physical AI series. The first article mapped the robot-simulation landscape. Here, we prepare and scale the simulation environment; we do not train a policy. The later Newton and Isaac Lab installments cover the next integration layers.

MuJoCo 加载并编译 MJCF 模型；MJWarp 在 NVIDIA Warp 中实现物理计算，后者编译 CUDA 内核以在 NVIDIA GPU 上推进仿真状态。这是我们“物理 AI 仿真现状”系列文章的第二篇。第一篇文章梳理了机器人仿真领域。在这里，我们准备并扩展仿真环境；我们不进行策略训练。后续关于 Newton 和 Isaac Lab 的文章将涵盖接下来的集成层。

Start with one useful Warp Kernel

从一个实用的 Warp 内核开始

NVIDIA Warp is a Python framework for writing high-performance, GPU-accelerated kernels. Warp lets developers author statically typed kernels in Python and compiles them for CPU or CUDA execution. The first launch builds and caches a native module; later launches reuse it. The kernel language is a performance-oriented subset of Python, while ordinary Python remains responsible for configuration, allocation, and launch orchestration.

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

This small robotics-oriented kernel advances point positions under gravity. One logical thread handles one point, so the same code scales from two points to millions without introducing GPU terminology into the control flow.

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

The three value propositions of Warp are: Warp 的三大价值主张是：

Pillar	What you get
Performance	Native-CUDA speed via JIT compilation, kernel fusion, and CUDA Graphs
Ease of use	Pure Python authoring with built-in vectors, matrices, quaternions, BVHs, hash grids, sparse matrices, and tile primitives
Capability	Differentiable kernels and DLPack-style interop so simulation can sit inside an ML training loop

支柱	你将获得
性能	通过 JIT 编译、内核融合和 CUDA Graphs 实现原生 CUDA 速度
易用性	纯 Python 编写，内置向量、矩阵、四元数、BVH、哈希网格、稀疏矩阵和瓦片原语
能力	可微内核和 DLPack 风格的互操作性，使仿真能够嵌入到机器学习训练循环中

Three properties make this useful in robotics:

三个特性使其在机器人领域非常有用：

Explicit parallel work. wp.tid() identifies the point, contact, body, or world owned by the current logical thread.
Explicit device arrays. An array lives on the selected device. Calling .numpy() on a CUDA array synchronizes and copies it to CPU memory; it is not a zero-copy path. For a device-resident PyTorch or JAX pipeline, use Warp’s framework adapters or DLPack-compatible sharing instead.
Composable kernel launches. A program can launch a sequence of focused kernels and capture supported CUDA work into a graph to reduce repeated dispatch overhead. Graph capture replays launches against existing buffers; it does not fuse arbitrary kernels.
显式并行工作。 wp.tid() 标识当前逻辑线程所拥有的点、接触、物体或世界。
显式设备数组。 数组驻留在选定的设备上。在 CUDA 数组上调用 .numpy() 会同步并将其复制到 CPU 内存；这不是零拷贝路径。对于驻留在设备端的 PyTorch 或 JAX 流水线，请改用 Warp 的框架适配器或兼容 DLPack 的共享方式。
可组合的内核启动。 程序可以启动一系列专注的内核，并将支持的 CUDA 工作捕获到图中，以减少重复调度的开销。图捕获（Graph capture）针对现有缓冲区重放启动；它不会融合任意内核。

Differentiability and Determinism

可微性与确定性

Two further Warp capabilities are worth knowing, even though neither is used in the SO-101 workflow in this article. Warp kernels are differentiable: a wp.Tape records the forward kernel launches made inside its context and replays their adjoints in reverse when backward() is called, which is why teams build differentiable geometry, CFD, and custom physics in Warp, including CAE workflows for simulation and design optimization.

还有两个值得了解的 Warp 功能，尽管本文的 SO-101 工作流中均未使用。Warp 内核是可微的：wp.Tape 会记录在其上下文中进行的向前内核启动，并在调用 backward() 时反向重放其伴随算子。这就是为什么团队在 Warp 中构建可微几何、CFD 和自定义物理引擎的原因，包括用于仿真和设计优化的 CAE 工作流。

Warp also supports deterministic execution, introduced in Warp 1.15: GPU atomics are scheduler-dependent by default, so repeated launches of the same kernel can differ slightly, and the opt-in deterministic modes trade some performance for reproducible ordering in simulation, validation, and regression tests. These are Warp capabilities, not guarantees of differentiability or determinism for an entire MJWarp rollout. See the Warp documentation on differentiability and deterministic execution for the details.

Warp 还支持在 1.15 版本中引入的确定性执行：GPU 原子操作默认依赖于调度程序，因此同一内核的重复启动可能会略有不同。选择性的确定性模式以牺牲部分性能为代价，换取在仿真、验证和回归测试中可重现的顺序。这些是 Warp 的功能，并不保证整个 MJWarp 部署的可微性或确定性。详情请参阅 Warp 关于可微性和确定性执行的文档。

Try Warp: pip install warp-lang (≥ 1.15 for GPU determinism), then python -m warp.examples.browse, or the tutorial notebooks.

尝试 Warp：运行 pip install warp-lang（GPU 确定性需 ≥ 1.15），然后运行 python -m warp.examples.browse，或查看教程笔记本。

What is MuJoCo Warp (MJWarp)?

什么是 MuJoCo Warp (MJWarp)？

A robot simulator repeatedly computes what happens next: given the current joint positions, velocities, controls, and contacts, it advances the scene by one small timestep. In this article, a world means one independent copy of that scene and its state.

机器人仿真器会重复计算下一步发生的情况：给定当前的关节位置、速度、控制和接触，它将场景向前推进一个小的时间步长。在本文中，一个“世界”意味着该场景及其状态的一个独立副本。