mjbatch：在 CPU 上批量并行推进 MuJoCo 仿真

一个模型，几千个仿真，不需要加速卡

机器人仿真里真正有意思的工作，很少是「跑一条轨迹」，几乎都是「跑一群轨迹」。轨迹优化器要靠一束 rollout 去有限差分出梯度；采样式 MPC 每个控制步要评估上千条候选动作序列；强化学习要上千个并行环境，否则策略更新的信号会被方差淹没；系统辨识要把同一段录制动作对七十个候选惯性向量各重放一遍。这些负载有一个共同形状：在仿真之间是彻底可并行的，在每个仿真内部是彻底串行的。这恰好是多核 CPU 擅长、而单线程 Python 循环极其难堪的形状。

mjbatch 是 Kevin Zakka 写的一个 Python + C++ 小库，正对着这个缺口：在 CPU 上并行推进几千个 MuJoCo 仿真，对外只暴露一个 Python 对象，物理调用期间释放 GIL。2026 年 9 月 10 日以 0.1.0 版本、Apache-2.0 协议发布，锁定 mujoco==3.11.0，提供 CPython 3.10 到 3.14 的 wheel，包含 free-threading（无 GIL）构建。整个库只有两个 C++ 头文件加一个绑定文件，外面再包一层很薄的 Python 类做命名访问器。

整个 API 就是一个对象

模型在构造时被复制，所以要在交给它之前改完 MjModel，之后不要再改。调用是串行的：一次只有一个 step，不允许并发重入。

import mujoco, numpy as np
from mjbatch import Batch

model = mujoco.MjModel.from_xml_path("scene.xml")
batch = Batch(model, num_sims=4096)          # 线程数默认用满所有逻辑 CPU
qpos, ctrl = batch.bind("qpos"), batch.bind("ctrl")
batch.expand("geom_friction")[:, :, 0] = np.random.uniform(0.4, 1.2, (4096, 1))
for _ in range(1000):
    ctrl[:] = policy(qpos)                   # 你的控制器，一次算完 4096 个
    batch.step()                             # 并行推进；qpos 原地更新

num_threads=0 表示用满每一个逻辑 CPU，并且夹到不超过 num_sims。这个默认值是有意为之，源码里把理由写成了注释：MuJoCo 的线性代数规模小、矩阵密，卡在访存延迟上而不是把核心算力吃满，所以每个物理核跑两个线程实测有 1.3 到 1.5 倍吞吐（在 Threadripper 7960X 上测得）。另一个构造参数 forward=True 会在每次 step 结束时补一次 mj_forward，让派生字段与状态同步而不是落后一个子步，代价是每个仿真每次调用多一次 forward。

方法	签名	行为
bind	bind(name, dtype=None) -> NDArray	按 MjData 自己的布局，返回一个 (N, ...) 的活视图。而 bind("state") 返回的是原始的 (N, nstate) 积分状态行。
expand	expand(name, dtype=None) -> NDArray	每个仿真各自的 mjModel 或 mjOption 取值，从模型播种，在每次物理调用前应用。
step	step(ids=None, nstep=1, history=None)	每个仿真在一个 worker 上跑 nstep 次 mj_step。ids 用来选子集，可以是排好序的唯一整数，也可以是布尔掩码。
forward / reset	forward(ids=None)、reset(ids=None, keyframe=-1)	reset 先 mj_resetData（或 mj_resetDataKeyframe），再套用待写入的字段，最后 mj_forward，所以派生字段立刻可用。
set_const	set_const(ids=None)	逐仿真跑 mj_setConst，并把它改过的字段全部 expand，于是派生常量也是逐仿真独立的。

C++ 绑定之上，Python 层加了与 MjData 同名的访问器：sensor("name")、joint("name")、actuator("name")、body("name")、site("name")，每个都返回从绑定字段里切出来的活视图。joint 知道自由关节在 qpos 里占 7 列、在 qvel 里占 6 列，而球关节是 4 和 3，所以 batch.joint("floating_base").qpos 的形状天然就对，不用自己数偏移量。

内存随线程数走，不随仿真数走

这是让 4096 个仿真能塞进一台笔记本的那个设计决定，值得说准确，因为它并不是最直觉的那个做法。每个仿真只保存自己的 mjSTATE_INTEGRATION 向量加上一组 warning 计数器。mjData 是每个 worker 线程一份，不是每个仿真一份。每次调用前，worker 把分给它的那些仿真的状态装进自己的 mjData，推进，再把状态写回去。

结论是：内存开销是线程数的函数（一个很小的定值），而不是仿真数的函数（你自己选的数）。测试套件用一个子进程量 ru_maxrss 的增长把这条钉死了：对测试模型而言，给 4096 个仿真各分配一个 mjData 会让常驻内存涨 712 MB，而 4096 个状态向量只有几 MB，断言是这个 batch 必须留在 256 MB 以下。开启了 sleep 的模型在构造时直接拒绝，因为 sleep 的记账不在 mjtState 里，会静默失同步。

同一个不对称也解释了基准测试里「批量子步」那一栏的结果。每次调用跑 10 个子步，等于把状态装载和回写的开销摊到 10 个物理步上而不是 1 个，所以 24 线程那一列从 88,111 跳到 676,019 sim-substeps/s。如果你的求解器能容忍子步之间不回 Python，nstep 就是最便宜的加速手段。examples/hello.py 把这条路推到极限：4096 个单摆从 4096 个不同角度释放，用一次 batch.step(nstep=1000) 调用推进 1000 步，整个过程 GIL 全程释放。

活数组与拷贝纪律

bind 返回的是覆盖在 batch 缓冲区上的真 numpy 数组，不是快照。写它就是设状态和控制量，step 之后读它就是观测。库为输入字段保留了一份「上次由我写入」的镜像，所以物理调用前的拷入是逐元素的，只动自上次写入以来变过的部分。每个绑定字段在调用后都会被拷出；在两次调用之间才绑定的派生字段，会在某个仿真的下一次调用时被填上。写 bind("state") 的某一行，等于设定该仿真在下一次调用时的状态，并且待写入的字段改动会叠加在其上，在那之前字段视图是陈旧的。reset 把两者一起丢弃，就像它丢弃一次字段写入那样。

dtype 规则是机械的：mjtNum 字段是 float64（若链接的是 float32 版 libmujoco 则是 float32），float 是 float32，int 是 int32，mjtByte 是 uint8，mjtBool 是 bool；任何 mjtNum 字段都可以显式要 float32。在 float32 构建下这条路径就是一次纯 memcpy。

expand 是把「一堆相同的仿真」变成「一堆不同的仿真」的那个开关。它覆盖 geom_friction、body_mass、dof_armature 这类 mjModel 数组，也覆盖 mjOption：标量给成 (N,)，向量给成 (N, size)。所以重力、时间步长、积分器、求解器设置全都是逐仿真独立的，这让域随机化和逐环境物理参数变成一行赋值。绑定层的文档里写了两个坑。把 iterations 或 ls_iterations 调高、或把 cone 换成 elliptic，可能让某个仿真需要比模板给 worker 的 mjData 所预留的 arena 更多的空间，这会以常规的 MuJoCo 捕获错误形式报出来并指名是谁。另外 enableflags 不能打开 sleep，理由和构造函数拒绝 sleep 模型是同一个。派生常量只在你调用 set_const 之后才跟随 expand 的输入变化，这与对单个模型调用 mj_setConst 的语义一致；mj_setConst 会写的那些 mjModel 标量（如 flags、stat）同样是逐仿真保存的。

数字到底说了什么

基准脚本 benchmarks/scaling.py 对基线交代得很老实：比较对象是一个 mjData 在普通 Python 循环里被推进，先 20 个子步预热再计时 200 个，所有加速比都以此为分母。它接受场景 XML、MuJoCo Menagerie 的模型名，或者一个 mesh 由 Menagerie 提供的 XML（配 --assets）。下表是发布时在 Unitree G1 场景上、256 个仿真、机器报告 12 核的那次运行结果。

配置	每次调用 1 个子步	加速比	每次调用 10 个子步
串行 mj_step 循环（单个 mjData）	8,523	基线	不适用
threads=1	8,530	1.0x	85,429
threads=8	58,217	6.8x	487,360
threads=24	88,111	10.3x	676,019

这里有两个读数值得注意。第一，threads=1 时批量路径与串行循环的差距在千分之一以内，也就是说在无可并行的场合，这层抽象不收你任何费用。第二，从 8 线程到 24 线程只换来 1.5 倍而不是 3 倍，这正是超线程小矩阵负载应有的形状：前 8 个线程各占一个物理核，后 16 个与它们共享。真正的奖品是单线程到 24 线程之间那 7.7 倍空间，而批量子步带来的那约 7.7 倍是叠在它之上的。

工程量真正所在的地方是线程池。它是一个持久线程池，提供阻塞式 parallel-for 和「粘性切片」：第 i 个条目属于 worker i * T / n 的切片，这让同一个仿真在多次调用之间留在同一个核上，从而让它那份 mjData 一直是缓存热的。跑完自己切片的 worker 会通过一个原子计数器去认领别人剩下的条目，所以没有人会卡在等最慢的那一片上。同一时刻只允许有一个 Run 在跑，而且 worker 函数不许抛异常，这也是为什么 MuJoCo 错误是被「捕获」而不是作为 C++ 异常向上传播的。

错误处理值得单独说一句，因为它会改变你写求解器的方式。任一 worker 上的 MuJoCo 错误都会抛出 RuntimeError，并指名第一个失败的仿真；其余仿真照常跑完，失败的那一个保持调用前的状态，它的待写入改动仍然挂着。这个捕获机制是 import 时装的一个 MuJoCo 日志 handler，之后再装别的 handler 会让它失效。

六个求解器，每个都自成一档

这些 example 本身就是「光有这个原语就够了」的论据。每一个都是单文件里的完整求解器，不是某个框架的包装壳，方法学覆盖四大类。

G1 后空翻是这组里最吃力的一个。它跟踪一段重定向后动捕片段的 200 帧，代价覆盖 12 个 body，位置和姿态误差按每平方米、每弧度加权，超过阈值后做类似 Huber 的软化，让大位置误差线性增长而不是平方增长；根节点权重向量在 x 和 y 上刻意放松，在 z 上紧 1000 倍。这段片段在 200 帧之后脚踝会翻转，所以跟踪就停在那里。

Go1 训练器是一个写在单文件里的完整 PPO：1024 个环境、24 步 horizon、600 轮迭代、每回合 500 步，时间步长 4 ms 配 decimation 5 得到 50 Hz 控制，PD 增益 35 与 0.5，并把 Go1 电机真实转子惯量 0.000111842 kg m² 作为 armature 加进去；观测 50 维、动作 12 维。域随机化每 5 秒重抽一次地面摩擦（0.4 到 1.0），奖励混合了速度跟踪、转向跟踪、2 Hz 小跑步态项（按对角腿配对相位）、姿态、倾斜、弹跳、晃动、关节限位与动作变化率。网络在有 CUDA 或 MPS 时放上去，但在 MPS 下 actor 仍留在 CPU，因为这么小的网络，数据搬运比计算更贵。

两个 cartpole 文件展示的是同一个任务的两种解法。cartpole_swingup.py 跑 iLQR，用解析的代价梯度与沿轨迹的 Hessian，并对杆垂在支点下方时出现的负对角项做正则化。cartpole_mpc.py 按 arXiv 2212.00541 的方法跑预测采样：25 个 knot 的时域上放 1024 条带噪 rollout，挑最好的一条，提交前 4 个子步，如此重复 150 个无窗口控制步。因为滚动时域永远到不了终端 knot，所以每个 knot 都用运行代价。第二个文件最清楚地说明了这个库为什么存在：每个控制步 1024 条 rollout，那是一个 batch 维度，不是一个循环。

抛球机械臂做的是硬件协同设计而不是控制。它固定 1 m 工作半径和 0.045 m 球半径，积分器设为 implicit-fast、时间步长 2 ms、求解器迭代 40 次，碰撞只保留球与地面，然后在形态和控制上同时搜索。电机模型是真实的：堵转转矩 6 与 3 N m、空载转速 24 rad/s、转子惯量 0.0006 kg m²，释放截止 0.65 s、每个电机 4 个 knot。没有逐仿真模型参数的话，每个候选都要重新编译一次模型；有了 expand，它就是一次列赋值。

系统辨识这个例子对「batch 即 Jacobian」这个思路依赖最重。它先录 1000 个 100 Hz 的指令、每个 5 个物理子步，驱动方式是带真实 PD 增益的多正弦激励，各关节幅值在 0.08 到 0.3 弧度之间、频率在 0.37 到 1.07 Hz 之间。然后用步长 1e-6 的有限差分对这段录制拟合 70 个参数，这意味着 Jacobian 的每一列都是另一个仿真。每步用 50 个样本的窗口，让 batch 保持小规模、迭代保持高频。

打包、可移植性与老实交代的边界

构建用 scikit-build-core 加 nanobind，而且打包决定是写在注释里的，不是留给你猜。wheel 覆盖 cp310 到 cp314 再加 cp314t（free-threading 构建），这个信号有意义：一个全部职责就是在物理调用期间释放 GIL 的库，是少数几个在无 GIL 解释器里会变得严格更有用的库。Windows 和 musllinux 被跳过，理由很具体：mujoco 的 wheel 里 mujoco.dll 没有配套的 import library，而扩展是靠 rpath 找到 libmujoco 的，Windows 没有对等机制。Linux 上 auditwheel repair 排除 libmujoco.so.3.11.0，macOS 上 delocate-wheel 排除对应的 dylib。把 mujoco==3.11.0 精确锁死，正是让这个排除安全的前提。

测试套件是一个文件里的 29 个测试，覆盖的正是本文一直在描述的语义：带变更追踪的拷入、派生字段的陈旧性、bind("state") 与返回数组共享内存、reset 丢弃待写入改动、按 ids 与按布尔掩码选子集，以及「内存不随仿真数增长」这条性质。开发工具是 pytest、ruff、pyright，两空格缩进、100 列行宽。

边界也要点名。调用是串行的，所以没有办法让一个对象里的两个 batch 并发推进。模型在构造时被复制，因此你之后再改 MjModel 不会传进去；流程是先改、再构造，凡是应当逐仿真变化的东西都走 expand。开了 sleep 的模型直接拒绝。调高求解器迭代次数或换用 elliptic 锥，可能超出模板为 worker 的 mjData 预留的 arena，而这会以被捕获的 MuJoCo 错误形式出现，不是静默变慢。没有 GPU 路径：这是一个按设计就是 CPU 的库，对标的应该是 MJX 而不是取代它。换来的东西是：你的整条工具链仍然是 numpy，你的调试器仍然是调试器，一台没有加速卡的笔记本能在一分钟内训出一个四足行走策略。

怎么跑起来

# 在 Menagerie 模型上跑吞吐量基准
uv run python benchmarks/scaling.py unitree_g1 --num-sims 256 --threads 1,8,24

# 4096 个单摆，1000 步，一次调用
uv run examples/hello.py

# 需要开窗口的 example 得有显示器；加 --headless 就不需要
uv sync --group examples
uv run examples/go1_joystick.py
uv run examples/g1_flip.py --play        # 回放已保存的求解结果

这些 example 会把求解结果落盘，所以 g1_flip.py --play 是加载 g1_flip_plan.npz 而不是重跑优化器。对昂贵的那几个来说这个分离很重要：求解一次，然后在可视化上反复迭代。