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. 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.
图 1：MJWarp 如何把 Python 连到 GPU 仿真——MuJoCo 加载并编译 MJCF 模型；MJWarp 用 NVIDIA Warp 实现物理，Warp 编译 CUDA 内核在 NVIDIA GPU 上推进仿真状态。
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.
这是「Physical AI 仿真现状」系列的第二篇：第一篇绘制了机器人仿真的版图；本篇准备并规模化仿真环境（不训练策略）；后续 Newton 与 Isaac Lab 两篇覆盖再上一层的集成。
Putting it together · 串起来
Decision shortcut:
决策捷径（按场景选工具）：
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 框架：开发者用 Python 写静态类型的 kernel，Warp 把它编译成 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.
这个面向机器人的小内核在重力下推进点的位置。一个逻辑线程管一个点，于是同一段代码从 2 个点扩到数百万个点，也不需要往控制流里塞任何 GPU 术语。
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())
（代码原样，不译）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())
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 管线应改用框架适配器或 DLPack 共享）；**可组合的内核启动**——程序可以依次启动一组各司其职的内核，并把支持的 CUDA 工作捕获成图以减少重复派发开销——图回放是对既有缓冲重放启动，不等于任意内核融合。
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 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.
**可微与确定性**：还有两个 Warp 能力值得知道（本文 SO-101 流程都没用到）。Warp 内核可微——wp.Tape 记录其上下文里的前向内核启动，backward() 时按逆序回放伴随（adjoint），所以团队用它做可微几何、CFD 与定制物理，包括仿真与设计优化的 CAE 工作流。Warp 1.15 起还支持确定性执行——GPU 原子操作默认依赖调度器，同一内核反复启动结果可能略有差异；可选的确定性模式用一点性能换仿真/验证/回归测试里可复现的顺序。注意这是 Warp 的能力，不等于整个 MJWarp rollout 的可微或确定性保证。
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，或教程 notebook。
What is MuJoCo Warp (MJWarp)? · MuJoCo Warp 是什么
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. One world might contain the SO-101 arm reaching for a cube; another can contain the same arm starting from a slightly different pose.
机器人仿真器反复计算「接下来发生什么」：给定当前关节位置、速度、控制与接触，把场景推进一个小时间步。本文里，一个世界（world）指该场景及其状态的一份独立副本——一个世界里是 SO-101 臂去够方块，另一个世界里同一条臂从稍不同的姿态出发。
MuJoCo and MJWarp can run the same compatible robot and task, but they organize the work differently. MuJoCo naturally suits developing and inspecting one or a few CPU worlds. MJWarp is a NVIDIA Warp implementation of MuJoCo's physics pipeline that places the model and a batch of independent states on NVIDIA GPUs; one call to mjw.step advances the entire batch.
MuJoCo 与 MJWarp 能跑同一个兼容机器人与任务，但组织工作的方式不同：MuJoCo 天然适合开发与检视 CPU 上的一两个世界；MJWarp 是 MuJoCo 物理管线的 NVIDIA Warp 实现，把模型和一批独立状态放上 GPU——一次 mjw.step 推进整批。
MJWarp's value is not necessarily a faster step for one world. It is the ability to advance hundreds or thousands together, giving the GPU enough parallel work to improve aggregate throughput, the total world-steps completed per second. That favors reinforcement learning and large-scale sampling, where collecting experience matters more than minimizing one environment's latency.
MJWarp 的价值不一定是单世界更快的一步，而是同时推进成百上千个世界、给 GPU 足够的并行工作量以提升聚合吞吐（aggregate throughput，每秒完成的总世界步数）。这利好强化学习与大规模采样——在这些场景里「收经验」比「压低单环境延迟」重要。
This blog covers the following: validate one MuJoCo world, move it to MJWarp, form a batch, verify it, and measure it correctly. Solver tuning, Jacobian representation, and specialized multi-GPU or determinism topics are not required for this migration and can be covered separately.
本文覆盖：验证一个 MuJoCo 世界 → 搬进 MJWarp 成批 → 验证 → 正确测量。求解器调参、雅可比表示、多 GPU 与确定性等专题不在本次迁移必需范围内，另行覆盖。
Then, the distinction is precise: Latency is wall-clock time for one simulation step. Aggregate throughput is the total number of world-steps completed per measured wall-clock second.
两个词的区分要精确：**延迟** = 一步仿真的墙上时间；**聚合吞吐** = 每实测墙钟秒完成的总世界步数。
### Basic usage: structs, batch sizes, and a minimal step · 基本用法：结构体、批大小与最小步进 The core API transition is small: Use mjw.make_data() when default/fresh state is intended. Use mjw.put_data() when the exact initialized MuJoCo state must cross the migration boundary. Allocating batched resources requires defining batch-size parameters (nworld / nconmax / njmax).
核心 API 迁移面很小：想要默认/全新状态用 mjw.make_data()；要把「精确初始化过的 MuJoCo 状态」原样带过迁移边界用 mjw.put_data()。分配批式资源需要定义批大小参数（nworld / nconmax / njmax）。
### Performance tuning · 性能调优 Tuning 1 — CUDA graph capture: mjw.step is many kernel launches; capture once, replay often:
调优 1：CUDA 图捕获——mjw.step 是一大串内核启动，捕获一次、反复回放：
with wp.ScopedCapture() as capture: mjw.step(m, d) wp.capture_launch(capture.graph)
（代码原样，不译）with wp.ScopedCapture() as capture: mjw.step(m, d) wp.capture_launch(capture.graph)
Tuning 2 — Size nconmax / naconmax / njmax tightly: memory and work scale with them. Tune with mjwarp-testspeed: --measure_alloc and watch overflows in mjwarp-viewer.
调优 2：把 nconmax / naconmax / njmax 配紧——内存与工作量随它们涨；用 mjwarp-testspeed --measure_alloc 调，并在 mjwarp-viewer 里盯溢出。
Additional tuning considerations. After sizing contact and constraint buffers, test solver iteration limits without changing task behavior. Meshes and CCD settings can increase memory use; nccdmax / naccdmax can reduce CCD buffer allocation when the measured contact counts allow it. MJWarp's compact solver uses MuJoCo's Newton constraint solver and sleeping, not the separate Newton physics-engine framework.
其他调优项：接触/约束缓冲配好后，在不改变任务行为的前提下测求解器迭代上限；网格与 CCD 设置会加大内存，实测接触数允许时可用 nccdmax / naccdmax 缩 CCD 缓冲。注意 MJWarp 的紧凑求解器用的是 MuJoCo 的 Newton 约束求解器与睡眠机制，不是另一个同名的 Newton 物理引擎框架。
To train policies on MJWarp physics: Isaac Lab via Newton; mjlab (manager API directly on MJWarp + PyTorch); MuJoCo Playground via MJX (impl='warp'). Install / try: pip install mujoco-warp · mjwarp-viewer path/to/scene.xml · Colab tutorial.
在 MJWarp 物理上训策略的三条路：Isaac Lab（经 Newton）；mjlab（manager API 直接建在 MJWarp + PyTorch 上）；MuJoCo Playground 走 MJX（impl='warp'）。安装试用：pip install mujoco-warp、mjwarp-viewer 查看场景、Colab 教程。
Workflow to migrate a MuJoCo scene to MjWarp · 把 MuJoCo 场景迁到 MJWarp 的工作流
Step 1: Establish a MuJoCo CPU baseline. The scene: nothing MJWarp-specific yet — an SO-101 arm, a table, and two cubes to stack, written as ordinary MJCF.
第 1 步：建 CPU 基线。场景：还没任何 MJWarp 特有东西——SO-101 臂、一张桌子、两块要摞起来的方块，普通 MJCF 写成。
Figure 2. SO-101 pick-and-place scene, rendered from the MuJoCo CPU simulation. The task is to grasp the red 44 mm cube and stack it on the blue cube; the same robot and scene are used for MJWarp validation.
图 2：SO-101 抓放场景（MuJoCo CPU 仿真渲染）：抓起红色 44 mm 方块摞到蓝色方块上；同一机器人与场景用于 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"> <freejoint name="red_cube_joint"/> <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> <body name="blue_cube" pos="0.33 0.06 0.046"> <freejoint name="blue_cube_joint"/> <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"/> </body> </worldbody> </mujoco>
（代码原样，不译）<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"> <freejoint name="red_cube_joint"/> <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> <body name="blue_cube" pos="0.33 0.06 0.046"> <freejoint name="blue_cube_joint"/> <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"/> </body> </worldbody> </mujoco>
In the companion repository this file is generated rather than hand-written: resolve_pick_place_scene() copies the Menagerie arm into .generated/, fills the table and cube coordinates from a robot profile, and writes scene_pick_place.xml.
配套仓库里这个文件是生成的而非手写：resolve_pick_place_scene() 把 Menagerie 的臂拷进 .generated/、按机器人档案填桌面与方块坐标、写出 scene_pick_place.xml。本演练用 SO-101 档案。
import mujoco mjm = mujoco.MjModel.from_xml_path("scene_pick_place.xml") mjd = mujoco.MjData(mjm) fps = 50 # controller rate sim_substeps = 10 # physics steps per control frame frame_dt = 1.0 / fps mjm.opt.timestep = frame_dt / sim_substeps controller = PickPlaceController(spec=spec) # waypoints + damped-least-squares IK for _ in range(600): # 600 control frames ctrl = controller.step(mjm, mjd, frame_dt) for _ in range(sim_substeps): mjd.ctrl[: mjm.nu] = ctrl mujoco.mj_step(mjm, mjd)
（代码原样，不译）import mujoco mjm = mujoco.MjModel.from_xml_path("scene_pick_place.xml") mjd = mujoco.MjData(mjm) fps = 50 # controller rate sim_substeps = 10 # physics steps per control frame frame_dt = 1.0 / fps mjm.opt.timestep = frame_dt / sim_substeps controller = PickPlaceController(spec=spec) # waypoints + damped-least-squares IK for _ in range(600): # 600 control frames ctrl = controller.step(mjm, mjd, frame_dt) for _ in range(sim_substeps): mjd.ctrl[: mjm.nu] = ctrl mujoco.mj_step(mjm, mjd)
Match the simulation and control rates. At 50 control frames per second and 10 physics substeps per frame, use a physics timestep of 0.002 seconds. Set it before the CPU rollout and before uploading the model with mjw.put_model so both backends advance the same simulated time. Without that line, every later measurement inherits the mismatch: parity comparisons, throughput numbers quoted as "simulated seconds," and any learned policy whose action rate no longer matches deployment.
对齐仿真与控制频率：每秒 50 控制帧 × 每帧 10 物理子步 → 物理时间步取 0.002 秒，且必须在 CPU rollout 之前和 mjw.put_model 上传之前都设好，两个后端才推进同样的模拟时间。少了这一行，后面所有测量都继承这个错位：parity 对比、以「模拟秒」计的吞吐数、以及任何动作率从此对不上部署的已学策略。
Check whether the cubes are stacked successfully. With 44 mm cubes, success becomes two measurable conditions: a horizontal center error of xy_err ≤ 0.015 m (measured between the cube centers) and a vertical separation of 0.035 m ≤ dz ≤ 0.055 m between cube centers (one cube edge, with slack for settling). Evaluate both conditions after the cubes have settled; a successful process exit alone does not establish task success.
摞方块是否成功要量化：44 mm 方块下，成功 = 两个可测条件——水平中心误差 xy_err ≤ 0.015 m（方块中心间测量）且垂直间距 0.035 m ≤ dz ≤ 0.055 m（一个方块边长加沉降余量）。两条件都要在方块停稳后评估；进程正常退出本身不构成任务成功。
Run the CPU task from the companion checkout. The run ends by printing the two numbers above (stack check: xy_err=… dz=…), which is the assertion the rest of the article compares against.
从配套仓库跑 CPU 任务（clone → uv 建 Python 3.12 环境 → 装依赖 → 跑 600 帧无头 + debug）。运行结束打印上面两个数（stack check: xy_err=… dz=…）——这正是后文所有对比所参照的断言。
The arm comes straight from MuJoCo Menagerie pinned to a known-good commit, since Menagerie assets change, so treat the scene as a template. Optional reBot variant: the companion code also exposes --robot rebot with a separate profile for the scene layout, gripper, and capacity limits (nconmax=256, njmax=500). This walkthrough uses SO-101. Validate the reBot asset and task separately before reporting its results.
臂直接取自 MuJoCo Menagerie 并钉在已知良好 commit（Menagerie 资产会变，场景当模板用）。可选 reBot 变体：配套代码另有 --robot rebot，用独立档案管理场景布局、夹爪与容量上限（nconmax=256、njmax=500）；本文用 SO-101，reBot 的资产与任务须单独验证后再报数。
Step 2: Validate one-world MJWarp parity. Run one world on the GPU first, with the host still in the loop, so you can watch the same task in the same viewer and compare the same two numbers. Upload the model, allocate batched state, seed it from the initialized host state, and run one forward pass before stepping.
第 2 步：单世界 parity 验证——先在 GPU 上跑一个世界、宿主留在环里，同屏同任务对同样的两个数。上传模型、分配批式状态、用初始化过的宿主状态播种、步进前先跑一次 forward。
wp.init() import mujoco_warp as mjw device = wp.get_device() m = mjw.put_model(mjm) d = mjw.make_data(mjm, nworld=1, nconmax=spec.nconmax, njmax=spec.njmax) wp.copy(d.qpos, wp.array(mjd.qpos[None, :], dtype=wp.float32, device=device)) wp.copy(d.qvel, wp.array(mjd.qvel[None, :], dtype=wp.float32, device=device)) wp.copy(d.ctrl, wp.array(mjd.ctrl[None, :], dtype=wp.float32, device=device)) mjw.forward(m, d)
（代码原样，不译）wp.init() import mujoco_warp as mjw device = wp.get_device() m = mjw.put_model(mjm) d = mjw.make_data(mjm, nworld=1, nconmax=spec.nconmax, njmax=spec.njmax) wp.copy(d.qpos, wp.array(mjd.qpos[None, :], dtype=wp.float32, device=device)) wp.copy(d.qvel, wp.array(mjd.qvel[None, :], dtype=wp.float32, device=device)) wp.copy(d.ctrl, wp.array(mjd.ctrl[None, :], dtype=wp.float32, device=device)) mjw.forward(m, d)
Seeding the three fields explicitly is the transparent option, and it makes clear exactly what crosses to the device; mjw.put_data(mjm, mjd, nworld=…) carries the whole initialized struct over in one call instead.
显式播种这三个字段是更透明的做法——什么过了设备一清二楚；mjw.put_data(mjm, mjd, nworld=…) 则一把把整个初始化结构带过去。
def simulate_frame() -> None: ctrl = controller.step(mjm, mjd, frame_dt) for _ in range(sim_substeps): mjd.ctrl[: mjm.nu] = ctrl wp.copy(d.ctrl, wp.array(mjd.ctrl[None, :], dtype=wp.float32, device=device)) mjw.step(m, d) mjd.qpos[:] = d.qpos.numpy()[0] mjd.qvel[:] = d.qvel.numpy()[0] mujoco.mj_forward(mjm, mjd)
（代码原样，不译）def simulate_frame() -> None: ctrl = controller.step(mjm, mjd, frame_dt) for _ in range(sim_substeps): mjd.ctrl[: mjm.nu] = ctrl wp.copy(d.ctrl, wp.array(mjd.ctrl[None, :], dtype=wp.float32, device=device)) mjw.step(m, d) mjd.qpos[:] = d.qpos.numpy()[0] mjd.qvel[:] = d.qvel.numpy()[0] mujoco.mj_forward(mjm, mjd)
Step 3: Size contact and constraint capacity. MJWarp allocates contact and constraint buffers before stepping. Exceeding those capacities invalidates the affected rollout for verification or benchmarking, even when execution continues with an overflow warning rather than an exception. Increase the relevant limit and rerun the task. Larger buffers use more GPU memory, so verify capacity over the full task before tightening the allocation.
第 3 步：配接触与约束容量。MJWarp 步进前分配接触/约束缓冲；超容会使该 rollout 的验证或基准作废——即使执行只打溢出警告不抛异常。该提上限就提、重跑任务；缓冲越大越吃 GPU 显存，收紧分配前先在全任务上验容量。
Set contact and constraint limits for the robot and task being simulated. The SO-101 profile uses nconmax=128 and njmax=300 as starting capacities. Check that these limits are sufficient during the most contact-heavy part of the task — for pick-and-place, the instant both jaws and the table touch a cube, not the arm hovering in free space. An overflow is reported rather than raised: MJWarp prints the budget to increase ("narrowphase overflow - please increase nconmax to …") to the terminal, and flags the affected worlds in Data.overflow for you to read back after a step. Only mjw.put_data raises an error outright. mjwarp-testspeed --measure_alloc reports the contacts and constraints a scene actually consumed, and it aborts the rollout with the offending world IDs as soon as any world overflows. Treat those reports as failures: raise the limit and re-run before trusting either the trajectory or the benchmark, then tighten again whenever the model, collision geometry, or task changes.
按机器人与任务设限：SO-101 档案以 nconmax=128、njmax=300 起步，并在任务接触最重的时刻验够不够——抓放任务是双爪和桌面同时碰方块那一瞬，不是臂悬空时。溢出是「报告」不是「抛错」：MJWarp 在终端打出该提的预算（"narrowphase overflow - please increase nconmax to …"）并在 Data.overflow 里标记受影响世界供步后回读；只有 mjw.put_data 会直接抛错（它手里有可对照的 MuJoCo 状态）。mjwarp-testspeed --measure_alloc 报告场景实际消耗的接触与约束，任一世界溢出即带肇事 world ID 中止。把这些报告当失败对待：先提限重跑再信轨迹或基准；模型、碰撞几何或任务一变就重新收紧。
Step 4: Scale to 2,048 worlds. Once one-world parity passes, reallocate at the target size and replicate the initialized state across the batch. Two things change relative to Gate 2: nworld, and the fact that nothing crosses the PCIe bus per step.
第 4 步：扩到 2,048 个世界。单世界 parity 过关后，按目标规模重分配、把初始化状态复制满整批。相对第 2 关只变两件事：nworld，以及每步不再有任何数据跨 PCIe 总线。
nworld = 2_048 d = mjw.make_data(mjm, nworld=nworld, nconmax=spec.nconmax, njmax=spec.njmax) wp.copy(d.qpos, wp.array(np.tile(mjd.qpos, (nworld, 1)), dtype=wp.float32, device=device)) wp.copy(d.qvel, wp.array(np.tile(mjd.qvel, (nworld, 1)), dtype=wp.float32, device=device)) wp.copy(d.ctrl, wp.array(np.tile(mjd.ctrl, (nworld, 1)), dtype=wp.float32, device=device)) mjw.forward(m, d) with wp.ScopedCapture() as capture: mjw.step(m, d) step_graph = capture.graph np.tile gives every world the same starting state, which is the right baseline for a throughput measurement; per-world randomization would instead write different rows of d.qpos on the device.
（代码原样，不译）nworld = 2_048 d = mjw.make_data(mjm, nworld=nworld, nconmax=spec.nconmax, njmax=spec.njmax) wp.copy(d.qpos, wp.array(np.tile(mjd.qpos, (nworld, 1)), dtype=wp.float32, device=device)) wp.copy(d.qvel, wp.array(np.tile(mjd.qvel, (nworld, 1)), dtype=wp.float32, device=device)) wp.copy(d.ctrl, wp.array(np.tile(mjd.ctrl, (nworld, 1)), dtype=wp.float32, device=device)) mjw.forward(m, d) with wp.ScopedCapture() as capture: mjw.step(m, d) step_graph = capture.graph np.tile gives every world the same starting state, which is the right baseline for a throughput measurement; per-world randomization would instead write different rows of d.qpos on the device.
Figure 3. Scaling the SO-101 task from one CPU world to 2,048 independent GPU states using the same compatible model. A single MJWarp step advances the full batch. This conceptual illustration highlights aggregate throughput, measured as world-steps per wall-clock second.
图 3：同一个兼容模型，把 SO-101 任务从一个 CPU 世界扩到 2,048 个独立 GPU 状态；单次 MJWarp 步进推进整批。此概念图强调聚合吞吐——每墙钟秒的世界步数。
Step 5: Verify, then measure. GPU launches are asynchronous, so a naive timer measures how fast Python queued work, not how fast the GPU finished it. Warm up first — the first launches pay kernel compilation and allocation — then synchronize immediately before and after the timed region. Report both aggregate world-steps per second and milliseconds per batched step, together with the batch size. Use the measured curve to identify where additional worlds improve throughput and where memory or compute limits reduce the benefit. Results depend on the scene, simulation settings, and hardware; a one-world latency comparison does not establish batched throughput.
第 5 步：先验证再测量。GPU 启动是异步的，naive 计时器量的是 Python 排队的速度、不是 GPU 干完活的速度。先预热（首批启动要付内核编译与分配成本），再在计时区前后立刻各同步一次。报告口径：聚合「每秒世界步数」与「每批步毫秒」一起报、并附批大小；用实测曲线找「加世界还涨吞吐」与「内存/算力开始限制收益」的拐点。结果取决于场景、仿真设置与硬件——单世界延迟对比不能推出批吞吐。
import time for _ in range(10): # warm-up: compilation, allocation, caches wp.capture_launch(step_graph) wp.synchronize() t0 = time.perf_counter() for _ in range(200): wp.capture_launch(step_graph) wp.synchronize() # without this you time the queue, not the work elapsed = time.perf_counter() - t0 total = 200 * nworld print(f"{total / elapsed:,.0f} world-steps/second")
（代码原样，不译）import time for _ in range(10): # warm-up: compilation, allocation, caches wp.capture_launch(step_graph) wp.synchronize() t0 = time.perf_counter() for _ in range(200): wp.capture_launch(step_graph) wp.synchronize() # without this you time the queue, not the work elapsed = time.perf_counter() - t0 total = 200 * nworld print(f"{total / elapsed:,.0f} world-steps/second")
To see that curve on your own hardware, scaling_study.py sweeps the batch size and prints ms/step alongside throughput and speedup:
想在自己的硬件上看这条曲线：scaling_study.py 扫批大小，打印 ms/step 与吞吐、加速比——
cd /tutorials/sim2real-blogs/notebooks/mujoco/part2 python solutions/so101_mjwarp_solution.py --headless-steps 600 # parity, needs CUDA python scaling_study.py --worlds 1 64 1024 2048 8192 --steps 100
（代码原样，不译）cd /tutorials/sim2real-blogs/notebooks/mujoco/part2 python solutions/so101_mjwarp_solution.py --headless-steps 600 # parity, needs CUDA python scaling_study.py --worlds 1 64 1024 2048 8192 --steps 100
Get started · 上手
Warp (kernel layer): pip install warp-lang → python -m warp.examples.browse → docs · GitHub. MJWarp (GPU MuJoCo): pip install mujoco-warp → mjwarp-viewer benchmarks/humanoid/humanoid.xml → docs · Colab tutorial. SO-101 context: sim-to-real course · Physical AI learning paths. Train on top of MJWarp: mjlab · MuJoCo Playground · Isaac Lab + Newton (upcoming posts).
上手路径：Warp（内核层）pip install warp-lang、示例浏览器、文档与 GitHub；MJWarp（GPU 版 MuJoCo）pip install mujoco-warp、用 mjwarp-viewer 看 humanoid 示例、文档/Colab；SO-101 语境：sim-to-real 课程与 Physical AI 学习路径；在 MJWarp 上训练：mjlab、MuJoCo Playground、Isaac Lab + Newton（后续文章）。
What's next · 下一步
This post covered raw Warp → MJWarp: GPU kernels, batched stepping, and an SO-101 scene using mjw.step. Next, we will port the same MJCF environment into Newton, using MuJoCo Warp as its rigid-body solver (newton.solvers.SolverMuJoCo). Newton will manage the model, state, controls, and contacts, while MJWarp runs underneath. You will also see what Newton adds: multi-format assets, swappable solvers, sensors/IK helpers, and an Isaac Lab path. If you build something with Warp or MJWarp, open an issue on the linked repositories or find us on Discord NVIDIA Omniverse.
本篇覆盖 raw Warp → MJWarp：GPU 内核、批步进与用 mjw.step 的 SO-101 场景。下一篇把同一 MJCF 环境移植进 Newton——用 MuJoCo Warp 当它的刚体求解器（newton.solvers.SolverMuJoCo）：Newton 管模型、状态、控制与接触，MJWarp 在底下跑；还会看到 Newton 加了什么：多格式资产、可换求解器、传感器/IK 辅助与 Isaac Lab 路径。用 Warp 或 MJWarp 做了东西，欢迎到相关仓库开 issue 或在 Discord NVIDIA Omniverse 找我们。
References · 参考
Blog 1: The State of Simulation for Physical AI (series post 1); NVIDIA Warp GitHub · Documentation · v1.15.0 release (GPU determinism) · deterministic execution guide; MuJoCo Warp GitHub · docs; Warp differentiable physics article; Warp 1.5.0 tile-based programming; mjlab (arXiv:2601.22074); MuJoCo Playground; NVIDIA SO-101 sim-to-real course; Newton next post.
系列第一篇《The State of Simulation for Physical AI》；NVIDIA Warp 仓库/文档/v1.15.0（GPU 确定性）/确定性执行指南；MuJoCo Warp 仓库与文档；Warp 可微物理文章；Warp 1.5.0 tile 编程；mjlab（arXiv:2601.22074）；MuJoCo Playground；NVIDIA SO-101 sim-to-real 课程；Newton 下一篇。