解读 32 · export.py 与 infer_policy.py:导出与回放

解读对象:scripts/export.py(11 行)+ src/mjlab_microduck/export.py(311 行)、scripts/infer_policy.py(1524 行)
需要的前置:解读 27–解读 31(被导出的训练配置)、总纲

在做什么

部署件的「最后一公里」两端。export 侧:把训练 checkpoint 变成 ONNX——scripts/export.py 是 11 行薄壳(export.py:8-11),真正逻辑在 mjlab_microduck/export.py 的 run_export,发布流水线 uv run publish 直接调用同一函数,发布物不可能绕过这条路。infer 侧:scripts/infer_policy.py 是 1500 行的 CPU 排练场——用 onnxruntime 在 MuJoCo 里跑导出的策略,50 Hz 控制(decimation 4 × timestep 0.005,infer_policy.py:1149),键盘当遥操,把真机要发生的一切先在笔记本上过一遍。

真实符号:导出的关键一击

训练配置里 obs_normalization=True,即观测先过 EmpiricalNormalization(用运行统计量把每维观测归一化的层)再进网络。导出时 runner.export_policy_to_onnx(path, filename)(mjlab_microduck/export.py:264)把这个层作为子模块烘焙进计算图,产出 actor(normalizer(obs))——机器人上跑的图和训练看到的逐位一致。文件头的警告值得抄在墙上:仿真内 play 是自己在外面临时套归一化的,会掩盖一个手工转换、忘了归一化的 checkpoint——「never convert by hand」(:1-9)。随后元数据经 attach_metadata_to_onnx(:266-267)写进 ONNX 的 custom_metadata_map,infer_policy 读 gait_period 走的正是这个通道(infer_policy.py:246-252)。checkpoint 支持本地 logs/rsl_rl/<experiment_name> 或按 wandb_run_path 下载,迭代号用正则 model_(\d+)\.pt 解析(:65-69)。

真实符号:回放侧的 BAM 复刻

infer_policy 最有分量的部分,是把训练用的 BAM 执行器(执行器摩擦/电压的参数化模型)在 CPU MuJoCo 里复刻:文件头手工镜像训练常量——BAM_MOTOR_NAME = "xl330"、BAM_MODEL = "m6"、BAM_KP_FW = 200.0、BAM_VIN_RANGE = (6.5, 8.2)、BAM_STIFF_SOLREF_FRICTION = (-5.0e4, -2.0e2)(:33-45)。刻意不 import 训练侧常量模块(那会拖进 mjlab/torch/warp,约 16 秒),改由 tests/test_infer_policy_bam.py 锁死一致(:28-32)。load_mujoco_with_bam(:58)把 XML 位置执行器换成电压受限的力矩电机、清零关节 damping/frictionloss,交给 MujocoController 每步重写。

观测拼装 get_observations(:655)就是 27-31 篇反复出现的 61 维(旧策略 51 维):角速度 3 + 重力 3 + joint_pos/vel 各 14 + last_action 14 + 命令 13(--new-cmd-obs)。多策略共存是骨架:_update_policy_session(:514)按速度命令模长与 switch_threshold = 0.05 在 walking/standing 间热切换;trigger_behavior(:731)把 K/L/R 键变成 kick/roulade 的纯会话替换——这些策略训练时命令全零,触发即换 ONNX,计时到点(kick_duration = 3.0、roulade_duration = 2.0)交还行走策略。踢球的球由 _place_ball(:764)按训练同款 yaw 系偏移瞬移;ground_pick 的相位钟走到 φ≥0.7 自动收工(:721),而训练侧 RISE_END = 0.80 大于它——ground_pick 配置里留有「待验证」警告,窗口对不上时起身段播不完。

设计取舍

这个脚本的立场是:回放器必须和训练「失真得一样」。它复制训练的每个失真(执行器模型、--delay 观测延迟、命令布局),又提供「调向真机」的旋钮——--foot-friction 2.0、--foot-solref 0.04 模拟 PU 鞋底(:942-949),复现「真机跑着跑着向前摔」这类训练分布外现象。键盘从终端读(TerminalInput,:143)而非 viewer 窗口,因为窗口按键会触发 MuJoCo 内置可视化快捷键。

你带走的收获

• 归一化层必须烘焙进 ONNX 图;任何「手工转换再在外面补」的路径都会掩盖不一致。
• 回放脚本手工镜像训练常量时,用测试锁死同步,而不是共享 import。
• 多策略运行时 = 会话指针切换 + 每个策略一个命令填充规则;零命令即怠速让切换无感。
• 回放器的价值在「可调的分布外」:摩擦、接触柔度、电压、延迟都是复现真机差异的变量。

延伸

• 真机侧对应物是 duck-control 的 policy.rs(第一辑 03);策略来自:解读 27
• 本目录:microduck-replica/upstream/microduck_rl/scripts/export.py、.../scripts/infer_policy.py、.../src/mjlab_microduck/export.py