把 启元机器人Q1 的训练搬到昇腾 910B：从 MuJoCo 到 PPO 的一次完整迁移
把 Q1 的训练搬到昇腾 910B
从物理仿真到 PPO 的完整拆解，以及还没做完的部分
2026-10-30 · GewisLab
目录
手上有台机器，两张 Ascend 910B，没有 NVIDIA 显卡。我们想训练启元 Q1 的人形动作跟踪，而官方那套 mjlab 训练链是绑死在 CUDA 上的。所以这件事得从最底层拆起，先搞清楚启元做了什么，再讲我们怎么设计，最后说清楚还没做完的部分。
先把结论放在前面。仿真吞吐方面，我们用原生 MuJoCo 多进程跑到 25.7 万物理步每秒，而 warp 的 CPU 后端只有 5.3k，差 49 倍。训练侧是 2304 个环境、96 个 worker，约 2.57 万 env-step 每秒，2.1 秒一个 iteration，策略跑在 910B 上，训练时 npu-smi 显示 AICore 大约 185%。与官方实现做开环对标，第 10 步的奖励差 2.7%，观测 RMSE 0.31。训练结果上，官方采样（随机起始帧加关节扰动）下，平均回合回报从 -50 涨到 150 左右，平均回合长度从 20 步涨到 80 步（1.6 秒）；从动作首帧做确定性评估是 18 步。这个成绩不算好，原因放在第三节。
最值得说的一句是：昇腾缺的其实不是算力，是物理仿真这条生态。把物理留在 CPU、把策略网络放到 NPU，这条路能走，而且比"给 Warp 补一个 NPU 后端"现实得多。
一、启元这套东西是怎么搭的
1.1 硬件与控制栈
Q1 是人形机器人，22 个自由度，高 88 厘米，整机 16 公斤。硬件上是双板结构：运控板用 RK 系芯片，大脑板是 Jetson Orin。软件层用 ROS2 加 FastDDS，走的是去中心化架构。这样设计的好处是，二次开发的程序放在 PC、运控板或者大脑板上，都能通过同一套 ROS2 接口透明访问机器人的全部能力，不用关心底层协议。
这套设计其实挺聪明的，二次开发门槛低。但有个隐含前提：训练侧默认你有 NVIDIA 显卡。
1.2 仓库都放了什么
启元把工程拆得比较细，几个仓库各管一段：
| 仓库 | 作用 |
|---|---|
q1-sdk | ROS2 接口，32 个服务、18 个话题，负责部署和控制 |
q1-urdf | v3 的 URDF 和三套 MJCF（v3 / full_range / mc） |
mjlab | 官方 fork 的物理引擎接口层，基于 mujoco_warp |
primebot_rl_mjlab | Q1 动作跟踪训练链，真正干活的地方 |
GMR | 动作重定向，把人类动作转成机器人可用的轨迹 |
SuperMimic | V2 版本全身动作跟踪，但所需资产没有公开 |
| 工具链 | arm、x86 的原生和交叉编译环境 |
需要看清一点：SDK 只管部署，不管训练。训练全部落在 primebot_rl_mjlab 和 mjlab 这两个仓库里。也就是说，想改训练链路，战场就在这两个地方，跟 SDK 关系不大。
1.3 训练链路怎么串起来
mjlab 用的是 Isaac Lab 那套 manager API，物理后端是 mujoco_warp，再往下是 warp。问题就出在 warp 这里：它的 GPU 后端只有 CUDA，wp.get_device("npu:0") 会直接报错。
启元还把 mujoco-warp 锁定在一个具体的 git revision 上：
mujoco-warp = { git = "https://github.com/google-deepmind/mujoco_warp", rev = "ea7f05b" }
我们专门去查了这个 ea7f05b，它就是官方 3.7.0.1 的 release 提交，不是私改分支。这一点确认之后，"上游是不是藏了定制 kernel"的担忧就排除了，后面的版本组合也才有得谈。
另一个细节挺关键。启元的 mjlab fork 要求 torch>=2.7.0，而上游原版要求 2.14。这个差异反而帮了忙，因为 torch_npu 2.7.1 刚好能对得上。换句话说，启元锁定的版本区间，天然就落在昇腾可支持的范围里。
1.4 环境设计比想象中讲究
这部分是我们对着源码一项项抠出来的，也是后面做对标的基础。
MJCF 里配置的其实是力矩电机，不是位置伺服。官方在模型外面包了一层 PD 控制，四组增益分别是：手臂 20/0.8/21，腿和腰 yaw 30/1.5/36，腰 roll 80/3.5/36，踝 20/0.8/21（kp/kd/力矩上限）。默认姿态来自 HOME_KEYFRAME，因为 MJCF 里其实没有 keyframe，nkey=0。
奖励和观测只用了 14 个 body，不是全部 26 个；参考帧是跟着机器人 yaw 重锚定的，不是全球坐标；motion_anchor_ori_b 是参考 anchor 系里的 3 维重力，不是旋转矩阵前两行。这些细节叠在一起，actor 观测是 119 维。
这些东西任何一项不对齐，训出来的就是另一个东西。所以后面我们对标，也是按这份清单逐条来的。
二、我们的设想
2.1 起点和方案
手上只有 910B，没有 N 卡。官方 mjlab 依赖 mujoco_warp，而 warp 的 GPU 后端只有 CUDA，物理这一层在昇腾上直接跑不起来。
我们的做法很直接：物理仿真留在 CPU，用原生 MuJoCo（C 库）做多进程批量仿真；策略网络放到 NPU 上训练。这个形态和华为 CANN 官方具身智能配方里的 LQC 样例是一致的：CPU C 库 MuJoCo 跑物理，torch_npu 跑 PPO。
2.2 先跑官方栈，看看卡在哪
一开始没有急着重写。我们先把官方 mjlab 在纯 CPU 上跑通了，能跑，但很慢。做了一组微基准之后原因就清楚了：warp 的 CPU 后端实际上是串行的。100 万个线程、每个线程 2000 次 sin，总共 20 亿次运算跑了 12.97 秒，折合 0.15 Gops/s，进程 CPU 占用 90.9%。单核跑满，40 核的配额在旁边空着。1024 个环境只有 5.3k 步每秒。
这就是整个改造的起点。
2.3 自己写仿真层
思路很直接。每个进程持有一批 MjData，直接调用原生 mujoco.mj_step；进程之间用共享内存传 obs、奖励和动作，用 Barrier 同步，不走 pickle。核心步进大概是这样：
# 单进程批量步进（简化）
for i, d in enumerate(self.datas):
q_des = self.default_joint + actions[i] * 0.25
for _ in range(10): # decimation = 10
tau = self.kp * (q_des - q) - self.kd * qd
d.ctrl[:] = np.clip(tau, -self.effort, self.effort)
mujoco.mj_step(self.model, d)
进程间的同步协议是每步两次 rendezvous：
def _run(self, cmd: int) -> None:
self.cmd[0] = cmd
self.barrier.wait() # 放行 worker
self.barrier.wait() # 等 worker 写完
实测下来，40 核配额里进程数开到 96、每进程 24 个环境吞吐最高。进程不是越多越差，因为物理步进之间有等待，超配一点反而能把核喂饱。
| 配置 | env-step/秒（1 个 env-step = 10 个物理子步） |
|---|---|
| 36 workers × 32 env | 18,305 |
| 48 workers × 32 env | 22,520 |
| 64 workers × 32 env | 25,536 |
| 96 workers × 24 env | 26,516 |
2.4 快没用，得跟官方算得一样
这是最花时间的部分。我们把官方源码当参照，一项一项对齐。
第一件是执行器。MJCF 里的 actuator 是力矩电机，直接往 ctrl 写位置目标，物理就错了。所以按官方那四组参数自己算力矩，再按 effort 截断。
第二件是默认关节角。MJCF 里没有 keyframe，默认姿态在官方 q1_v3_constants.py 的 InitialStateCfg 里。动作目标是 default 加 action 乘 0.25，关节观测也要减这个偏置。
第三件是 body 子集。官方 motion_cmd.body_names 只有 14 个，包括 pelvis、髋膝踝、torso、肩 roll、肘和腕，不是全部 26 个。
第四件是参考帧重锚定，这是 BeyondMimic 那类跟踪的关键：
delta_pos_w = robot_anchor_pos_w
delta_pos_w[..., 2] = anchor_pos_w[..., 2] # z 用参考的
delta_ori_w = yaw_quat(quat_mul(robot_anchor_quat, quat_inv(anchor_quat)))
body_pos_relative_w = delta_pos_w + quat_apply(delta_ori_w, body_pos_w - anchor_pos_w)
如果按全球坐标去算 body 位置奖励，训练会直接跑偏。
第五件是观测维度。motion_anchor_ori_b 是参考 anchor 系里的重力 3 维，不是旋转矩阵前两行，所以最终是 44 + 3 + 3 + 3 + 22 + 22 + 22，共 119 维。
对齐之后，用同一初始帧（与官方 play 模式一致）、同一零动作序列，一步一步和官方比：
| 控制步 | 官方奖励 | 本实现 | 绝对差 | 观测 RMSE |
|---|---|---|---|---|
| 0 | 2.60 | 4.31 | 1.71 | 1.18 |
| 4 | 3.33 | 4.11 | 0.77 | 0.42 |
| 8 | 3.47 | 3.52 | 0.05 | 0.32 |
| 9 | 3.44 | 3.35 | 0.09（2.7%） | 0.31 |
前几步的差异主要来自两处还没对齐的细节：官方执行器有 0 到 10 个 lag 的随机延迟，观测项也注入了高斯噪声，这两项目前都还没建模。第 8 步以后两侧收敛到 0.1 以内。
2.5 NPU 这边，把 PPO 按官方配置复刻
物理在 CPU，网络在 NPU，中间每步搬两次张量。超参直接对着官方 rl_cfg.py 抄：网络 512-256-128 ELU，观测经验归一化，init_std 1.0，5 epochs、4 minibatches，lr 1e-3 配 adaptive 调度（desired_kl 0.01），gamma 0.99、lam 0.95，clip 0.2 加值函数裁剪，entropy_coef 0.005，steps_per_env 24。基本是一比一。
观测归一化这段代码不复杂，但很关键：
@torch.no_grad()
def __call__(self, obs):
batch_mean, batch_var = obs.mean(0), obs.var(0, unbiased=False)
delta = batch_mean - self.mean
new_count = self.count + obs.shape[0]
self.mean = self.mean + delta * obs.shape[0] / new_count
m2 = self.var * self.count + batch_var * obs.shape[0] \
+ delta.pow(2) * self.count * obs.shape[0] / new_count
self.var = m2 / new_count
self.count = new_count
return (obs - self.mean) / torch.sqrt(self.var + self.eps)
自适应学习率和 KL 早停也补上了：
if losses["kl"] > 2.0 * cfg.desired_kl: # 漂太大，减速
for g in opt.param_groups:
g["lr"] = max(g["lr"] / 1.5, cfg.lr_min)
elif losses["kl"] < 0.5 * cfg.desired_kl: # 太保守，加速
for g in opt.param_groups:
g["lr"] = min(g["lr"] * 1.5, cfg.lr_max)
这里踩了一个很贵的坑。checkpoint 一开始只存了权重，没存归一化统计量。训练时完全看不出来，一切正常；到了离线评估，喂进去的是原始 obs，策略 9 步就崩了。这种坑不自己踩一遍很难想到。修法就是 checkpoint 把 mean、var、count 一起存下来，--resume 和评估都加载。
2.6 整体架构
左边是仿真：32 到 64 个进程，每个进程一批 MjData，直接调原生 MuJoCo；进程间用共享内存传数据，Barrier 同步。右边是训练：PPO 的 rollout 缓冲、GAE、裁剪更新都在 910B 上。每次控制步只搬两次数据，obs 上卡，action 下卡。
2.7 跑出来的结果
回合回报从 -50 涨到 150 左右，回合长度从 20 步涨到 80 步，参考动作全长 350 步。KL 稳定在 0.008 到 0.013 之间，目标是 0.01。训练期间用 npu-smi 看，AICore 大约 185%，NPU 是真实在算的。
确定性评估（从动作首帧开始，不加动作噪声）大概 18 步、回报 66；按策略分布采样跑 20 集，平均 18 步、回报 34。主要卡在手腕的 z 偏差超过 0.15 米阈值，说明手臂还没学会跟随。
三、还没做完的部分
代码已经跑起来了，NPU 也已经忙起来了，仿真里的 Q1 也学会走了……就是小编还没见过真正的 Q1。
这部分想写清楚，免得看着像已经收工了。
首先，动作只跟到 1.6 秒左右。官方采样是随机起始帧，我们的策略平均只能撑 80 步；从首帧开始更是只有 18 步，终止时是手腕 z 偏差超过 0.15 米。原因有两个：一是样本预算，官方是 4096 个环境训 30001 个 iteration，我们大约是人家的二十分之一；二是难度，固定首帧的对照训练能跟到 5.1 秒（约 256 步），换成随机起始后同样的预算只到 1.6 秒。所以瓶颈在采样覆盖和预算，不在环境或训练链路本身。
其次，有两个对齐项还没做：执行器延迟和观测噪声。官方是 0 到 10 个 lag 随机，每个观测项还加了高斯噪声，我们都没建模。这两块会影响 sim2real 的真实度，可能得排到下一轮。
第三，多卡还没上。现在策略跑单卡，另一张卡拿来跑别的实验。HCCL 是通的，all_reduce 验证过，但 PPO 的 DDP 还没接，也就是数据并行加梯度 all_reduce。工程量不大，但得等单卡这条先稳。
第四，Triton kernel 还在样例级。triton-ascend 装了，也跑通了自写的 vec_add，但训练热路径里还没用上。现在的想法是先把 obs 归一化或者奖励计算里的大核换掉，得先 profile 清楚，不能为了用而用。
第五，sim2real 没上真机。导出 ONNX 和 play、sim2sim 脚本官方都有，我们还没跑。端侧这块我们的原则是不动，官方用什么我们用什么的，所以这一步其实是验证，不是改造。
第六，上游 PR 还没提。只提了 issue 没提 PR，因为之前定过规矩，干完验证完再提。现在 issue 在讨论，代码在 fork 里，等这条链稳定了再整理成 PR。
最后，T1 没碰。T1 开源的是四足 V2，人形没放，训练侧也没有官方链。不过我们的栈是配置化的，关节名从 MJCF 读，真要接 T1，大概就是加配置、加奖励权重、再用 GMR 重定向动作，不是重写，是工程量。
四、复现
# 1. 环境（版本组合见仓库 docs/VERSION_MATRIX.md）
pip3 install -r ascend/requirements-ascend.txt
. /usr/local/Ascend/ascend-toolkit/set_env.sh
# 2. 冒烟
cd ascend/native_sim && python3 test_vec_env.py
# 3. 训练
OMP_NUM_THREADS=1 python3 scripts/train.py --num-envs 2304 --workers 96 \
--iterations 3000 --logdir logs/training
# 4. 评估
python3 scripts/evaluate.py --ckpt logs/training/<run>/model_best.pt \
--episodes 20 --stochastic
五、最后
这一轮证伪了三件事。给 warp 补 NPU 后端就能跑 mjlab，不成立，参考 AMD 的数据，非 CUDA 后端还不如 CPU；MJX 替代也不成立，xla-npu 现在还是 beta 推理；“官方配方让 MuJoCo 提速 2 倍”，是 8 卡对 1 卡，不是对 GPU。
证实的事更简单。原生 MuJoCo 多进程在 40 核配额上比 warp CPU 后端快 49 倍；跟官方实现能对到个位数百分比；物理留 CPU、策略上 NPU 这条路能规模化、能复现、能对标。
端侧我们完全没动，训练产物按官方方式导出，真机和推理还是启元自己的 SDK。
后面有进展再补。
项目地址：atomgit.com/GewisLab/primebot-ascend ｜ 上游反馈：q1-sdk#3、primebot_rl_mjlab#1、SuperMimic#1
https://atomgit.com/PrimebotTech
至于 sim2real……奈何小编目前还没能把 Q1 请到实验室，暂时只能在数字世界里反复折腾。真机验证这一步，就先留个 TODO 了。🥲
作为“人工智能6S店”的官方数字引擎，为AI开发者与企业提供一个覆盖软硬件全栈、一站式门户。
更多推荐
0
0- 0
已为社区贡献2条内容
所有评论(0)