策略测试与 Sim2Sim 部署是连接训练管线与真机迁移的关键环节。在 Isaac Lab 中完成强化学习训练后,你需要在相同的物理仿真环境中验证策略的鲁棒性,随后将其导出为可移植格式,并在 MuJoCo 中执行跨仿真器验证(Sim2Sim)。本章面向中级开发者,系统讲解从 Checkpoint 回放到 MuJoCo 部署的完整工作流,涵盖标准 PPO、AMP、注意力编码器以及 BeyondMimic 等多种算法变体的测试与迁移方法。
Sources: README_CN.md
从训练到部署的整体架构
整个管线遵循“训练 → 验证 → 导出 → 跨仿真器部署”的四阶段模式。在 Isaac Lab 内部,play.py 系列脚本负责加载训练保存的 Checkpoint、恢复 Runner 状态、执行推理并自动导出 TorchScript 与 ONNX 两种格式。随后,MuJoCo 侧的 sim2sim_atom01.py 系列脚本加载导出的 .pt 文件,在独立的物理引擎中重建观测构造与 PD 控制循环,从而验证策略是否存在仿真器过拟合。
flowchart TD
A[训练完成<br/>logs/rsl_rl/model_XXX.pt] --> B[策略测试<br/>play.py / play_amp.py / play_bm.py]
B --> C{自动导出}
C --> D[JIT模型<br/>exported/policy.pt]
C --> E[ONNX模型<br/>exported/policy.onnx]
D --> F[Sim2Sim部署<br/>sim2sim_atom01*.py]
E --> G[其他推理框架]
F --> H[MuJoCo仿真<br/>视频/绘图/键盘控制]
Sources: play.py, sim2sim_atom01.py
Isaac Lab 策略测试
测试脚本与适用范围
仓库为不同算法变体提供了独立的测试入口,避免在 play 脚本中引入不必要的依赖耦合。下表列出了各脚本的功能定位与适用场景。
| 测试脚本 | 支持的 Runner | 核心特性 | 典型使用场景 |
|---|---|---|---|
play.py |
OnPolicyRunner、DistillationRunner | 支持注意力编码器导出、平面地形切换、推搡测试、注意力可视化 | 标准 PPO、带感知编码器、蒸馏策略的验证 |
play_amp.py |
OnPolicyRunner、AMPRunner、DistillationRunner | 标准 JIT/ONNX 导出、AMP 判别器加载 | AMP 对抗运动先验策略的验证 |
play_bm.py |
OnPolicyRunner、DistillationRunner | 支持键盘实时控制指令 | BeyondMimic 模仿学习策略的验证 |
Sources: play.py, play_amp.py, play_bm.py
测试脚本的执行流程
所有 play 脚本遵循一致的主流程,差异仅体现在环境配置覆盖项与导出器类型上。以 play.py 为例,核心逻辑可分为环境准备、模型加载、自动导出和推理循环四个阶段。首先,脚本通过 hydra_task_config 解析任务配置与 Agent 配置,随后用 CLI 参数覆盖环境并行数、随机种子、地形类型等关键选项。值得注意的是,测试模式会强制关闭观测噪声 (env_cfg.noise.add_noise = False) 并默认禁用推搡事件,以确保评估结果的可复现性;若需要验证抗扰动能力,可显式传入 --push_robot 开启。
Sources: play.py
模型恢复阶段,脚本优先从 logs/rsl_rl/{experiment_name} 目录自动定位最新 Checkpoint,也可通过 --checkpoint 直接指定绝对路径。随后根据 agent_cfg.class_name 实例化对应的 Runner,并调用 runner.load(resume_path) 恢复策略网络与优化器状态。紧接着,runner.get_inference_policy() 将策略切换为评估模式并返回 act_inference 函数句柄,该句柄在推理循环中被直接调用以生成动作。
Sources: play.py, on_policy_runner.py
策略自动导出
测试脚本在运行推理前会自动将策略导出到 Checkpoint 同级目录下的 exported/ 文件夹中。导出过程需提取 Runner 内部的神经网络模块 (policy_nn) 以及观测归一化器 (normalizer),确保 Sim2Sim 阶段的输入预处理与训练时完全一致。
对于标准策略,isaaclab_rl.rsl_rl 提供的 export_policy_as_jit 和 export_policy_as_onnx 会分别生成 TorchScript 和 ONNX 文件。若任务启用了注意力编码器(attn_enc),play.py 则使用自定义的 TorchAttnEncPolicyExporter 与 OnnxAttnEncPolicyExporter,在导出前将编码器、actor_obs_encoder、estimator 等子模块与 Actor 网络按前向逻辑重新组装,保证 JIT trace 能够捕获完整的前向计算图。
推理循环与实时评估
推理循环以 simulation_app.is_running() 作为持续条件,内部通过 torch.inference_mode() 禁用梯度计算。策略接收环境观测后输出动作,环境执行步进并返回新观测。若策略使用了循环网络(如 LSTM),需要在 dones 为 True 时调用 policy_nn.reset(dones) 清零隐状态。当传入 --real-time 时,脚本会根据环境步长 dt 与实际耗时计算补偿睡眠,使可视化回放接近真实时间流速。
Sources: play.py
常用 CLI 参数速查
以下参数在测试阶段最为常用,可与各 play 脚本组合使用。
| 参数 | 类型/默认值 | 说明 |
|---|---|---|
--task |
str / None |
任务名称,必须指定 |
--num_envs |
int / None |
并行环境数,测试时建议设为 1 |
--checkpoint |
str / None |
指定 Checkpoint 文件路径 |
--load_run |
str / None |
加载指定运行目录 |
--plane |
flag / False |
切换为平面地形(仅 play.py) |
--push_robot |
flag / False |
启用机器人推搡事件(仅 play.py) |
--video |
flag / False |
录制视频并保存到 videos/play/ |
--video_length |
int / 200 |
视频录制的步数长度 |
--real-time |
flag / False |
尽可能按真实时间运行 |
--keyboard |
flag / False |
启用键盘控制(仅 play_bm.py) |
Sources: cli_args.py, play.py, play_bm.py
MuJoCo Sim2Sim 部署
Sim2Sim 部署的目标是验证导出的策略模型在另一个物理引擎中的行为一致性,同时为真机部署提供与仿真解耦的推理基准。仓库中的 sim2sim_atom01*.py 系列脚本负责在 MuJoCo 中加载 JIT 模型、重建观测与控制器闭环。
Sim2Sim 脚本矩阵
根据算法特性与验证需求,MuJoCo 侧提供了多个部署变体。
| 部署脚本 | 关键差异 | 附加功能 |
|---|---|---|
sim2sim_atom01.py |
基础版本 | 关节位置/基座速度绘图、headless 视频录制 |
sim2sim_atom01_amp.py |
与基础版一致,增加交互 | 键盘实时调速(8/2/4/6/7/9)、相机跟随、实时同步 |
sim2sim_atom01_attn_enc.py |
增加地形高度扫描 | MuJoCo ray casting 感知、射线可视化 |
sim2sim_atom01_interrupt.py |
增加中断注入 | 在指定时段注入 is_interrupt 标志并覆盖右臂目标位姿 |
sim2sim_atom01_bm.py |
BeyondMimic 观测构造 | 从 motion npz 加载参考运动作为观测输入 |
Sources: sim2sim_atom01.py, sim2sim_atom01_amp.py, sim2sim_atom01_attn_enc.py, sim2sim_atom01_interrupt.py, sim2sim_atom01_bm.py
Sim2Sim 控制循环架构
MuJoCo 侧的控制循环采用“高频物理 + 低频策略”的分层架构。物理仿真以固定步长 dt(通常 0.001s,即 1000Hz)推进,而策略推理频率由 decimation 控制。例如当 decimation = 20 时,策略以 50Hz 的频率输出目标关节位置,中间帧通过 PD 控制器维持力矩跟踪。
sequenceDiagram
participant M as MuJoCo Physics (1000Hz)
participant O as get_obs
participant P as Policy Network (50Hz)
participant C as PD Control
loop Every simulation step
M->>O: q, dq, quat, v, omega, gvec
alt count_lowlevel % decimation == 0
O->>P: 构造 policy_input (含 frame_stack)
P->>C: target_q = action * action_scale
C->>M: tau = PD(target_q, q) & clip
else
O->>C: 复用上一步 target_q
C->>M: tau = PD(target_q, q) & clip
end
M->>M: mj_step
end
Sources: sim2sim_atom01.py
观测构造与关节映射
观测构造是 Sim2Sim 与训练环境保持一致的核心。get_obs 函数从 MuJoCo 的 data 结构中读取本体状态,并通过 scipy.spatial.transform.Rotation 将线速度、重力向量等转换到机体坐标系。随后,脚本通过 usd2urdf 映射数组对关节顺序进行重排——这是因为 Isaac Lab 中的 USD/Articulation 关节顺序与 MuJoCo XML 的关节顺序并不相同。以 ATOM01 为例,usd2urdf 将 Isaac Lab 的 23 维关节向量映射到 MuJoCo 的关节排列,确保策略接收到的 q_obs 和 dq_obs 与训练时的分布一致。
Sources: sim2sim_atom01.py, sim2sim_atom01.py
标准版本的单帧观测维度为 78,其排列如下表所示。BeyondMimic 版本在此基础上前置了 46 维参考运动状态(关节位置+速度),总维度扩展至 121。
| 索引范围 | 物理量 | 维度 | 说明 |
|---|---|---|---|
| 0 – 2 | 角速度 omega |
3 | 机体坐标系,来自 IMU sensor |
| 3 – 5 | 重力向量 gvec |
3 | 机体坐标系下 [0, 0, -1] 的旋转 |
| 6 – 8 | 速度指令 cmd |
3 | vx, vy, dyaw |
| 9 – 31 | 关节位置偏差 q_obs |
23 | q - default_pos,经 usd2urdf 映射 |
| 32 – 54 | 关节速度 dq_obs |
23 | 经 usd2urdf 映射 |
| 55 – 77 | 上一帧动作 action |
23 | 策略输出的上一时刻动作 |
Sources: sim2sim_atom01.py
历史帧堆叠(Frame Stack)
为了赋予策略时序记忆能力,Sim2Sim 脚本实现了与训练环境对齐的历史观测缓冲区 hist_obs。在第一帧时,脚本通过 np.tile 将初始观测复制 frame_stack 份填满缓冲区;后续每一步则将新观测拼接到末尾并丢弃最旧的一帧。最终输入网络的 policy_input 是 [1, frame_stack * num_single_obs] 的展平向量。不同算法变体对 frame_stack 的默认值不同:基础版与中断版为 10,注意力编码器版为 5,BeyondMimic 版为 1。
Sources: sim2sim_atom01.py
注意力编码器版本的感知重建
当部署包含地形感知编码器的策略时,sim2sim_atom01_attn_enc.py 需要在 MuJoCo 中重建高度扫描观测。脚本通过 get_rays 函数沿机体下方垂直向下投射射线(ray_vec = [0, 0, -1]),采集网格化地形高度信息,并裁减到 [-1.0, 1.0] 区间后与本体观测拼接。在可视化模式下,draw_rays 会在场景中绘制红色球体标记射线落点,帮助开发者直观检查感知输入的正确性。
Sources: sim2sim_atom01_attn_enc.py, sim2sim_atom01_attn_enc.py
数据记录与后处理
除实时可视化外,Sim2Sim 脚本会在运行结束后自动生成两组分析图表:一是各关节的指令位置与实际位置对比,二是基座线速度与角速度的指令值与实际值对比。这些图表以 PNG 格式保存于工作目录,便于定量分析 PD 跟踪精度与策略响应延迟。若使用 --headless 模式运行,脚本还会通过 OpenCV 将渲染帧写入 simulation.mp4 视频文件。
Sources: sim2sim_atom01.py
端到端操作示例
1. 标准 PPO 策略测试与导出
在 Isaac Lab 环境中加载训练好的策略,执行单环境推理并自动导出模型。
python robolab/scripts/rsl_rl/play.py \
--task=Isaac-Velocity-Flat-Atom01-v0 \
--num_envs=1 \
--video \
--video_length=500
执行后,logs/rsl_rl/{experiment_name}/{run}/exported/ 目录下会出现 policy.pt 与 policy.onnx。
Sources: README_CN.md
2. MuJoCo Sim2Sim 验证
将上一步导出的 policy.pt 加载到 MuJoCo 中进行跨仿真器验证。
python robolab/scripts/mujoco/sim2sim_atom01.py \
--load_model /path/to/exported/policy.pt
若需要地形场景与无头视频录制,可追加参数:
python robolab/scripts/mujoco/sim2sim_atom01.py \
--load_model /path/to/exported/policy.pt \
--terrain \
--headless
Sources: README_CN.md
3. AMP 策略的 Sim2Sim 与键盘控制
AMP 版本支持在 Sim2Sim 阶段通过小键盘实时调整速度指令,适用于交互式行为调试。
python robolab/scripts/mujoco/sim2sim_atom01_amp.py \
--load_model /path/to/exported/policy.pt
运行后,按 8/2 调节前进/后退速度,4/6 调节左右速度,7/9 调节偏航速率,F 切换相机跟随,0 重置机器人状态。
Sources: sim2sim_atom01_amp.py
关键配置对照表
Sim2Sim 脚本通过内嵌的 Sim2simCfg 类硬编码机器人控制参数。以下列出 ATOM01 平台在不同脚本中的核心配置差异,便于你在定制自己的机器人时参考修改。
| 配置项 | 基础/中断/AMP | 注意力编码器 | BeyondMimic |
|---|---|---|---|
dt |
0.001 | 0.001 | 0.005 |
decimation |
20 | 20 | 4 |
frame_stack |
10 | 5 | 1 |
num_single_obs |
78 | 78 | 121 |
action_scale |
0.25 | 0.25 | 0.25 |
num_actions |
23 | 23 | 23 |
sim_duration |
10.0 | 10.0 | 1000.0 |
Sources: sim2sim_atom01.py, sim2sim_atom01_attn_enc.py, sim2sim_atom01_bm.py
常见问题与故障排查
| 现象 | 可能原因 | 排查方法 |
|---|---|---|
| Sim2Sim 中机器人瞬间摔倒 | usd2urdf 关节映射错误 |
核对 default_pos 在 MuJoCo 中的初始化姿态是否与 Isaac Lab 一致 |
| 策略输出为 0 或常数 | 加载了错误的 Checkpoint 或 normalizer 未导出 | 确认 exported/policy.pt 与测试时的 normalizer 匹配;检查 JIT trace 是否包含完整前向图 |
| 注意力编码器版本仿真异常 | 高度扫描维度与训练时不一致 | 检查 map_size 和 resolution 是否与环境配置中的 height_scanner 匹配 |
| AMP Sim2Sim 动作僵硬 | 参考运动数据未加载或观测构造遗漏 | 确认 BeyondMimic 版本的 motion_file 路径正确,且 m_input 拼接顺序与训练一致 |
| headless 模式报错 GLX | 渲染后端未正确指定 | 确保环境变量 MUJOCO_GL=glfw 与 __GLX_VENDOR_LIBRARY_NAME=nvidia 已设置 |
Sources: sim2sim_atom01.py, sim2sim_atom01_bm.py
下一步
完成策略测试与 Sim2Sim 验证后,你已经确认策略在跨仿真器环境中具备基本的鲁棒性。若要深入了解训练过程的内部机制,可阅读 On-Policy Runner 训练循环 与 AMP Runner 与模仿学习训练;若准备将策略迁移到真机,请参考 MuJoCo Sim2Sim 部署与真机迁移。对于需要自定义任务或扩展观测空间的开发者,自定义任务与环境扩展指南 提供了环境注册与配置注入的详细说明。