案例分析：NablaEngine（GPU 手写解析伴随）vs MuJoCo CPU——1000× 梯度差距从哪里来，各方案适用于什么场景

[TATP-233]

我们近期研究了 NablaEngine（一个面向 Unitree G1 的融合 CUDA 可微物理引擎，∇Mimic 的底层仿真器），并与 MuJoCo C 的批量仿真能力（mjbatch）做了对比，目的是理解可微仿真的设计空间。以下是我们验证过的事实总结，性能数据来自 NablaEngine 自己发布的 benchmark（docs/results/paper_rtx6000/paper_bench.md），可能对 UniLab 用户选择物理后端有参考价值。

---

NablaEngine 是什么

NablaEngine 不是 MuJoCo 的封装——它把 MuJoCo 的物理语义（Newton 求解器、implicitfast 积分器、约束模型、碰撞原语）用手写的融合 CUDA kernel 重新实现了一遍，专门面向一个机器人（mjlab G1）。它以 JAX 或 PyTorch 的可微算子形式接入上层框架。

核心架构选择：

• 一个 CUDA warp = 一个仿真世界，每个 world 的全部状态在共享内存（约 10 KB），模型常量在编译时烘焙进 constant memory
• 正向 tape：正向 kernel 记录接触表、离散决策（clamp 掩码、约束区域位）和 Cholesky 因子（Newton 的 L_H、implicitfast 的 L_Q），反向无需重算即可重放
• 解析伴随：手推 VJP 公式直接写成 CUDA 设备函数（adjoint.py: 1494 行；adjoint_contact.py: 510 行），遵循 MuJoCo Warp PR #1535 的约定

与 MuJoCo 的关系

组件	来源	性质
碰撞原语	MuJoCo Warp collision_primitive_core.py	翻译为 CUDA C++（Apache-2.0）
VJP 符号约定（正负号、命名）	PR #1535 collision_adjoint.py, adjoint_util.py	沿用，未重新定义
四元数积分 VJP	PR #1535 _adj_quat_integrate	沿用
Newton 求解器物理	MuJoCo 的定义	重新实现
implicitfast 积分器	MuJoCo 的定义	重新实现
伴随覆盖范围	PR #1535：11 种碰撞原语，condim 1/3/4/6，elliptic + pyramidal	严格子集：5 种，condim 1/3，仅 pyramidal

NablaEngine 自己的 THIRD_PARTY_NOTICES.md 明确写道其伴随 "follows that adjoint's conventions"（遵循 PR #1535 的约定）。数学覆盖范围比 PR #1535 更窄；贡献主要是工程性能（融合核、共享内存、编译时特化）。

梯度来自哪里：手写的解析公式，不是自动微分

梯度来自手推 VJP 公式直接写成 CUDA 设备函数——不经过 NVIDIA Warp、不经过任何自动微分工具、不经过有限差分。jax.custom_vjp / torch.autograd.Function 只负责把写好的 backward kernel 暴露给上层框架。

与有限差分（mjd_transitionFD）的本质算法区别：

	有限差分（正向探测）	解析伴随（反向模式）
对 f 的知识	黑盒	白盒（知道计算图结构）
99 个输入、1 个标量输出	99 × cost(forward)	1 × cost(forward)
64 步 BPTT × 99 个输入	6336 × cost(forward)	64 × cost(forward)
精度	ε 截断误差 + 浮点舍入	与正向同级

这不是工程效率差距，而是算法复杂度差距：FD 对每个输入独立探测（O(n)）；反向模式通过链式法则组合局部导数，从输出端自动扇出到所有输入（O(1)）。

实测性能（NablaEngine benchmark，RTX PRO 6000）

操作	4096 worlds	来源
NablaEngine 正向	1.21 ms	实测
NablaEngine 正向 + 解析伴随（vjp）	1.78 ms	实测
MuJoCo C 正向（16 线程，float64）	34.7 ms	实测
MuJoCo C FD Jacobian（前向差分，99 次探测）	2750 ms	实测
MuJoCo C FD Jacobian（中心差分，198 次探测）	5253 ms	实测
MJX 正向（JAX, GPU）	255 ms	实测
MJX 反向	不支持	—
MuJoCo Warp 正向	6.15 ms	实测
MuJoCo Warp + PR #1535 伴随（vjp）	21 ms	实测

重要说明：MuJoCo C 没有解析伴随。2750 ms 是有限差分的结果（每步 99 次正向仿真）。如果有人为 CPU MuJoCo 手写解析伴随，预期代价约为正向的 1–2 倍（合计约 35–70 ms），但目前没有人实现和测量过。

1000× 差距在 ∇Mimic 场景中的意义（4096 envs × 64 步 BPTT）

后端	每次 update 的梯度计算	3 小时能跑多少 update	可行？
NablaEngine（GPU 解析伴随）	0.11 s	~100,000	✅（论文结果：0.875 return）
MuJoCo C FD（CPU，前向差分）	220 s	49	❌ 远远不够
MuJoCo C FD（缩小到 64 envs × 8 steps）	0.86 s	12,500	⚠️ 训练动态完全不同
MuJoCo C + 假想的解析伴随（4096 × 64）	~4.5 s	~2,400	⚠️ 比所需少 ~40 倍

根本原因：∇Mimic 每次 update 需要 4096 × 64 = 262K 个精确梯度。FD 每步需要 99 次探测 → 每次 update 1.04 亿次 mj_step 调用 → 每次 update 在 CPU 上需要数分钟。CPU 线程池优化无法弥补 99× 的算法差距。

能否在 CPU MuJoCo 上手写解析伴随？

可以，而且大部分数学推导已经存在：

• MJX：RNE、CRB、运动学的 VJP（JAX，float64）
• MuJoCo Warp PR test(superdex): integrated rollout, checkpoint sim2sim and support evidence (#1533) #1535：碰撞原语（11 种）、Newton 求解器、implicitfast 的 VJP
• NablaEngine：CUDA 实现（5 种碰撞原语的子集）

CPU C 实现需要：

• Tape 机制（约 300 行 C）
• VJP 翻译（从 PR test(superdex): integrated rollout, checkpoint sim2sim and support evidence (#1533) #1535 / MJX 到 C，约 2000 行）
• 稀疏矩阵处理（MuJoCo 用 CSR）
• 对比 mjd_transitionFD 验证（约 1000 行测试）
• 批量并行分发（复用 mjbatch 的线程池）

估计工作量：约 4000–5000 行 C，1–2 个月。预期性能：16 CPU 线程上 4096-world vjp 约 70 ms（比 GPU 慢约 40×，但比 CPU FD 快约 39×）。这将填补一个真实的空白：通用的、float64 的、CPU 上的可微 MuJoCo。

手写解析梯度的局限性

1. 离散决策处梯度不唯一：伴随取正向走过的分支的单侧导数，给出的是次梯度，不是真梯度
2. 覆盖范围有限：每种新原语类型（tendon、equality、ball joint、elliptic cone）需要新的人工推导
3. 人工推导出错风险：VJP 公式是手推的；符号搞反、漏链式法则项会产生静默的错误梯度（NablaEngine 用大量 FD 验证来检查）
4. bitwise 重放耦合：伴随必须重放正向的精确 float32 舍入路径；改正向不改伴随就破坏正确性
5. 隐函数定理在奇异点失效：Hessian 接近奇异 → 梯度爆炸；需要正则化或梯度裁剪
6. 求解器截断偏差：梯度假设 Newton 精确收敛；10 次迭代时存在系统性偏差
7. 参数梯度受限：不返回 ∂/∂(质量、摩擦系数、增益、时间步) → 不能做可微系统辨识

对后端选择的启示

使用场景	推荐后端	原因
大规模可微 RL（BPTT 穿过接触密集物理）	GPU 解析伴随（NablaEngine）	1000× 吞吐优势是算法层面的
PPO / SAC / 无导数 RL	CPU 批量正向（mjbatch）	不需要物理梯度
MPC / iLQR / 轨迹优化（1–64 rollouts）	CPU + 解析伴随（如已实现）	低延迟、float64 精度
系统辨识	CPU float64 + 解析伴随	精度是硬需求
原型 / 模型探索	MJX 或 MuJoCo Warp（autodiff）	自动覆盖、无需推导
生产训练、模型已冻结、性能关键	手写解析	达到所需吞吐的唯一路径

对 UniLab 的讨论点

1. UniLab 的 SimBackend 接口是否应该包含统一的梯度 API（例如 step_with_vjp(cotangents) -> gradients）？目前各后端的梯度能力根本不同：MJX 仅有前向模式；NablaEngine 有反向模式；MuJoCo C 仅有有限差分；MuJoCo Warp 原生无梯度。

2. 是否有兴趣发起社区项目，为 CPU MuJoCo C 手写解析伴随？ 数学已经发表（PR test(superdex): integrated rollout, checkpoint sim2sim and support evidence (#1533) #1535、MJX）；工程上可行（约 4000–5000 行 C）。

3. 对于"超越 GPU 主导"的范式：FD vs 解析的差距（每步 99×）意味着 CPU 在正向吞吐上可以竞争（mjbatch 的线程扩展很好），但在梯度吞吐上没有解析伴随就无法竞争。异构架构可能需要明确承认这个不对称性。

---

所有 NablaEngine 性能数据来自其公开的 docs/results/paper_rtx6000/paper_bench.md（RTX PRO 6000, nvcc 12.8, 4096 worlds, 10 Newton iterations, 4×5 ms substeps）。CPU FD 数据来自其 MuJoCo C 3.17 adapter（16 线程，float64, mjd_transitionFD）。

交叉引用：NablaEngine · ∇Mimic · mjbatch · MuJoCo Warp PR #1535