本文档聚焦于 RoboLab 项目中用于训练管线操控、资产导出、运动数据处理、策略可视化以及训练监控的各类工具脚本。掌握这些工具能够显著提升开发迭代效率——从快速验证环境注册是否正确,到将训练好的策略导出为可在 MuJoCo 或真机上运行的格式,再到利用键盘进行交互式调试。建议读者在已完成 环境搭建与快速运行 并熟悉 项目结构导读 的基础上阅读本文。
Sources: robolab/scripts/tools/list_envs.py, robolab/scripts/rsl_rl/cli_args.py
核心脚本速查与职责边界
项目中的脚本主要分布在 robolab/scripts/rsl_rl/(训练与推理)、robolab/scripts/tools/(环境资产与数据处理)以及 rsl_rl/rsl_rl/utils/(算法库内部工具)三个目录下。下表总结了各核心入口脚本的用途和典型调用场景,帮助开发者快速定位所需工具。
| 脚本路径 | 核心职责 | 典型使用场景 |
|---|---|---|
robolab/scripts/rsl_rl/train.py |
启动 PPO/AMP/Distillation 训练 | 标准训练、分布式多卡训练、断点续训 |
robolab/scripts/rsl_rl/play.py |
标准策略回放与导出 | 验证策略表现、导出 JIT/ONNX、注意力可视化 |
robolab/scripts/rsl_rl/play_amp.py |
AMP 策略回放与导出 | 验证模仿学习策略、导出含 LSTM 的 ONNX |
robolab/scripts/rsl_rl/play_bm.py |
BeyondMimic 策略回放 | 验证相位控制策略、键盘交互控制 |
robolab/scripts/rsl_rl/cli_args.py |
统一命令行参数解析 | 被所有训练/推理脚本复用 |
robolab/scripts/tools/list_envs.py |
枚举已注册环境 | 确认新任务是否成功注册到 Gymnasium |
robolab/scripts/tools/export_terrain.py |
地形网格/高度图导出 | 生成 MuJoCo 可用的地形资产 |
robolab/scripts/tools/beyondmimic/csv_to_npz.py |
CSV 运动数据转 NPZ | 将动捕数据转换为训练数据集 |
robolab/scripts/tools/retarget/dataset_retarget.py |
批量运动重定向 | 将 GMR 格式数据转换为 Isaac Lab 格式 |
Sources: robolab/scripts/rsl_rl/train.py, robolab/scripts/rsl_rl/play.py
训练脚本详解与命令行配置
训练入口 train.py 是整个项目中最常使用的脚本之一。它通过 cli_args.py 提供了一套统一的命令行接口,允许开发者在不修改 Hydra 配置文件的前提下,快速覆盖关键超参数。下图展示了从命令行输入到训练启动的完整流程。
flowchart TD
A[命令行输入] --> B[AppLauncher 解析<br>--headless / --device]
B --> C[cli_args 解析<br>--experiment_name / --resume]
C --> D[Hydra 加载<br>env_cfg + agent_cfg]
D --> E[覆盖配置<br>num_envs / seed / max_iterations]
E --> F[创建 VecEnv 与 Runner]
F --> G[runner.learn 启动训练循环]
G --> H[Logger 写入<br>TensorBoard / W&B / Neptune]
cli_args.py 中注册的参数分为实验管理、断点加载和日志后端三类。通过 --resume 配合 --load_run 与 --checkpoint 可实现无缝断点续训;而 --logger wandb 或 --logger neptune 则能在不改变代码的情况下切换日志后端。下表列出了最常用的命令行参数及其作用。
| 参数 | 类型/默认值 | 说明 |
|---|---|---|
--task |
str / None |
要训练的任务名,例如 Isaac-Velocity-Flat-Atom01-v0 |
--num_envs |
int / None |
覆盖环境并行数量,常用于快速调试时减小显存占用 |
--seed |
int / None |
随机种子;设为 -1 时会在 [0, 10000] 范围内随机采样 |
--max_iterations |
int / None |
覆盖训练总迭代次数 |
--distributed |
flag / False |
启用多 GPU 分布式训练,详见 分布式多卡训练配置 |
--resume |
flag / False |
从已有 checkpoint 恢复训练 |
--load_run |
str / None |
指定要恢复的 run 文件夹名 |
--checkpoint |
str / None |
指定具体的 checkpoint 文件名 |
--logger |
str / None |
日志后端,可选 tensorboard、wandb、neptune |
--video |
flag / False |
训练过程中自动录制视频 |
--video_interval |
int / 2000 |
视频录制间隔(步数) |
Sources: robolab/scripts/rsl_rl/cli_args.py, robolab/scripts/rsl_rl/train.py
分布式训练的调试注意事项
当使用 --distributed 时,train.py 会自动根据 app_launcher.local_rank 为每个进程分配独立的 CUDA 设备,并在种子基础上加上 local_rank 以保证各进程数据采样的多样性。需要特别注意的是,CPU 设备与分布式训练不兼容——脚本会在检测到 --device cpu 与 --distributed 同时出现时抛出 ValueError。此外,日志记录仅在 gpu_global_rank == 0 的主进程上启用,避免多进程重复写入。开发者在调试分布式问题时,应首先确认各进程的 local_rank 和 device 分配是否符合预期。
Sources: robolab/scripts/rsl_rl/train.py
策略回放、导出与可视化
训练完成后,需要通过 play.py(标准策略)、play_amp.py(AMP 策略)或 play_bm.py(BeyondMimic 策略)进行验证和导出。这些脚本的共性流程是:加载配置 → 恢复 checkpoint → 获取推理策略 → 执行仿真循环 → 可选导出模型。
JIT 与 ONNX 策略导出
所有 play 脚本均支持将训练好的策略导出为 TorchScript (JIT) 和 ONNX 格式,以便在 MuJoCo Sim2Sim 或真机控制器中部署。play.py 针对标准 Actor-Critic 策略直接调用 export_policy_as_jit 和 export_policy_as_onnx,而针对**注意力编码器(Attention Encoder)**策略则提供了专门的 TorchAttnEncPolicyExporter 和 OnnxAttnEncPolicyExporter。这两个自定义导出器会完整保留 actor、encoder、estimator 和 actor_obs_encoder 的前向链路,确保感知模块在部署端的行为与训练端一致。导出文件默认存放在 checkpoint 同级目录的 exported/ 文件夹下。
Sources: robolab/scripts/rsl_rl/play.py
注意力热力图可视化
对于配置了注意力编码器的任务,play.py 在推理阶段支持实时可视化注意力权重。当策略返回 output_attn 时,脚本会将高度扫描点云 height_scan 从机器人局部坐标系转换到世界坐标系,并根据注意力分数对点云进行颜色分级(0–10 级),通过 VisualizationMarkers 在 Isaac Sim 视口中渲染出来。这一功能无需额外配置,只需确保环境配置中包含 attn_enc 字段即可自动启用。下表对比了标准策略回放与注意力策略回放的关键差异。
| 特性 | 标准策略 (play.py 普通模式) |
注意力策略 (play.py attn_enc 模式) |
|---|---|---|
| 策略类 | OnPolicyRunner / DistillationRunner |
OnPolicyRunner |
| 导出器 | 通用 export_policy_as_jit/onnx |
TorchAttnEncPolicyExporter / OnnxAttnEncPolicyExporter |
| 可视化 | 无 | 实时注意力点云热力图 |
| 输入维度 | num_actor_obs |
num_actor_obs + map_size |
| 键盘控制 | 可选 | 可选 |
Sources: robolab/scripts/rsl_rl/play.py
交互式调试:键盘控制
在策略验证阶段,开发者经常需要手动改变机器人的速度指令以测试策略的鲁棒性。robolab.utils.keyboard 模块提供了一个轻量级的键盘控制器,集成在 play.py 和 play_bm.py 中。该控制器通过 Omniverse 的 carb.input 接口订阅键盘事件,将按键映射为 SE(2) 速度指令并直接写入环境的 command_generator。默认按键映射如下:
| 按键 | 动作 | 增量 |
|---|---|---|
W / S |
前进 / 后退 (lin_vel_x) |
±0.05 m/s |
A / D |
左移 / 右移 (lin_vel_y) |
±0.05 m/s |
Q / E |
左转 / 右转 (ang_vel) |
±0.05 rad/s |
X |
清零所有速度 | 立即置零 |
R |
重置环境 | 触发 episode_length_buf = 1e6 |
在 play.py 中,键盘控制器仅在非 headless 模式下自动实例化;而在 play_bm.py 中,可通过 --keyboard 参数显式启用。使用时需注意,该控制器直接修改 env.command_generator.command 张量,因此会覆盖环境原有的随机指令生成逻辑。
Sources: robolab/robolab/utils/keyboard.py
运动数据处理流水线
项目提供了完整的运动数据工具链,覆盖从原始动捕数据到可用于 AMP 或 BeyondMimic 训练的 NPZ/PKL 格式的全链路转换。
CSV 转 NPZ:帧率重采样与速度计算
csv_to_npz.py 负责读取原始 CSV 运动文件(格式为 base_pos[3] | base_rot[4](xyzw) | dof_pos[N]),通过 MotionLoader 类完成三项核心操作:帧率插值(支持任意输入/输出 FPS)、SO(3) 角速度计算(基于四元数差分)以及 PyTorch 自动微分线速度计算。插值采用线性插值(LERP)处理平移和关节角,球面线性插值(SLERP)处理旋转四元数,保证旋转过渡的光滑性。最终输出的 NPZ 文件包含 fps、joint_pos、joint_vel、body_pos_w、body_quat_w、body_lin_vel_w、body_ang_vel_w 等字段,可直接被 MotionLoader(位于 MDP 中)加载用于模仿学习训练。
Sources: robolab/scripts/tools/beyondmimic/csv_to_npz.py
运动重定向:GMR 到 Isaac Lab
当运动数据来源为 GMR(Generative Motion Replay)格式时,需要使用 retarget 目录下的工具进行坐标系和关节顺序转换。核心模块 gmr_to_lab.py 定义了 extract_gmr_data 函数,负责将 GMR 的 pickle 文件(含 root_pos、root_rot(xyzw)、dof_pos)转换为 Isaac Lab 格式(四元数转为 wxyz),并根据 YAML 配置文件进行自由度重排序。single_retarget.py 用于处理单个文件并可在 GUI 中可视化验证,而 dataset_retarget.py 则面向批量处理——它会读取整个输入目录下的所有 .pkl 文件,启动一次多环境仿真,并行计算所有运动的关键刚体位置,最后将结果写回输出目录。
Sources: robolab/scripts/tools/retarget/dataset_retarget.py, robolab/scripts/tools/retarget/single_retarget.py
NPZ 运动回放
replay_npz.py 是一个轻量级的可视化脚本,用于在 Isaac Sim 中回放已转换好的 NPZ 运动文件。它跳过物理仿真(仅调用 sim.render() 而不调用 sim.step()),直接将运动数据写入机器人的根状态和关节状态,适用于快速检查重定向或插值结果是否正确。脚本会自动循环播放,并在每次重置时将相机视角对准机器人重心。
Sources: robolab/scripts/tools/beyondmimic/replay_npz.py
地形导出与资产工具
export_terrain.py 是连接 Isaac Lab 地形生成器与 MuJoCo 仿真器的桥梁脚本。它基于 TerrainGenerator 实例化地形,并支持三种导出模式:
--export_mesh:导出整体合并网格为.obj,并通过 CoACD 算法将其分解为凸包集合,自动生成 MuJoCo XML 片段;--export_meshes:按地形子块分别导出,适用于需要单独碰撞体的大型场景;--export_hfield:基于射线检测采样地形表面高度,生成 16-bit PNG 高度图,并输出 MuJoCohfield的 XML 配置。
CoACD 分解的参数可通过 --coacd_threshold 调整(默认 0.01),阈值越小分解越精细、凸包数量越多。导出的 XML 片段中会包含必要的 asset 和 worldbody 节点,开发者可直接复制到 MuJoCo 模型中使用。需要特别注意的是,当使用高度图时,MuJoCo 中机器人 base link 的 Z 轴位置需要加上 abs(h_min) 的偏移量以避免陷入地下。
Sources: robolab/scripts/tools/export_terrain.py
环境注册诊断
在开发新任务或扩展环境时,首要步骤是确认环境已正确注册到 Gymnasium。list_envs.py 脚本会以无头模式启动 Isaac Sim,遍历 gym.registry 中所有 ID 包含 "Atom" 的任务,并以表格形式输出任务名、入口点和配置文件路径。如果新创建的任务未出现在列表中,通常是由于对应模块未被导入——脚本通过 import robolab.tasks 触发注册,因此请确保任务模块位于 robolab/tasks/ 下且已被包的 __init__.py 正确引用。
Sources: robolab/scripts/tools/list_envs.py
训练日志与监控系统
RSL-RL 算法库内置了多后端日志系统,统一封装在 rsl_rl.utils.logger 中。Logger 类在初始化时会根据配置自动选择 tensorboard、wandb 或 neptune 作为底层写入器,同时保存代码的 git diff 状态到 logs/{experiment}/{run}/git/ 目录,便于后续实验复现。
日志指标结构
无论是标准 PPO 还是 AMP 训练,日志指标均按层级组织:
Loss/{key}— 各类损失值(value loss、policy loss、entropy loss 等)Policy/mean_noise_std— 策略动作噪声标准差Perf/{fps,collection_time,learning_time}— 采样与训练性能Train/mean_reward与Train/mean_episode_length— 回合级统计Episode/{key}— 环境 extras 中自定义的回合信息
对于 AMP 训练,LoggerAMP 在基类基础上额外追踪 AMP/mean_total_reward 和 AMP/mean_style_reward,并在每个回合结束时将风格奖励的均值写入 Episode_Reward/style。若启用了 RND 好奇心探索,则还会记录 Rnd/mean_extrinsic_reward、Rnd/mean_intrinsic_reward 和 Rnd/weight。
Sources: rsl_rl/rsl_rl/utils/logger.py, rsl_rl/rsl_rl/utils/amp_logger.py
外部日志后端配置
使用 Weights & Biases 时,需要在 runner 配置中指定 wandb_project,并在环境变量中设置 WANDB_USERNAME。WandbSummaryWriter 继承自 SummaryWriter,因此即使在不安装 wandb 的环境中,TensorBoard 的本地日志功能仍然可用。Neptune 的配置类似,但需要额外提供 NEPTUNE_API_TOKEN 和 NEPTUNE_USERNAME。两个后端均支持自动上传模型 checkpoint 和 git diff 文件。
Sources: rsl_rl/rsl_rl/utils/wandb_utils.py, rsl_rl/rsl_rl/utils/neptune_utils.py
算法库通用工具
rsl_rl.utils.utils 提供了一批不依赖具体算法的底层工具函数,在自定义网络或扩展训练逻辑时非常实用:
resolve_nn_activation(name)— 将字符串(如"elu"、"swish")解析为对应的torch.nn.Module,支持 12 种常见激活函数;resolve_optimizer(name)— 将字符串解析为优化器类,支持 Adam、AdamW、SGD、RMSprop;resolve_callable(name)— 支持通过字符串动态导入类或函数,格式包括直接引用、module.path:ClassName以及简写名(在rsl_rl包内查找);split_and_pad_trajectories/unpad_trajectories— 在 done 索引处切分轨迹,填充至等长并返回掩码,常用于 AMP 的判别器批次重组或序列模型处理。
Sources: rsl_rl/rsl_rl/utils/utils.py
常见问题排查
下表汇总了在使用上述工具脚本时可能遇到的典型问题及其排查方法。
| 现象 | 可能原因 | 排查与解决 |
|---|---|---|
list_envs.py 中看不到新任务 |
任务模块未被导入或注册 | 检查 robolab/tasks/__init__.py 是否导入了任务模块;确认 gym.register 的 ID 包含 "Atom" |
训练时 CUDA out of memory |
--num_envs 过大或视频录制开启 |
减小 --num_envs;关闭 --video;使用 --headless |
| 断点续训加载失败 | --load_run 或 --checkpoint 路径错误 |
检查 logs/rsl_rl/{experiment_name}/ 下的文件夹名;或直接传入 --checkpoint 的绝对路径 |
| 键盘控制无响应 | 运行在 headless 模式或焦点不在视口 | play.py 需去掉 --headless;play_bm.py 需添加 --keyboard |
| ONNX 导出后在 MuJoCo 中结果不一致 | 注意力编码器策略使用了普通导出器 | 确认使用了 TorchAttnEncPolicyExporter / OnnxAttnEncPolicyExporter,它们会保留 encoder 和 estimator 链路 |
export_terrain.py 的 CoACD 分解极慢 |
--coacd_threshold 过小或 resolution 过高 |
适当增大 threshold 到 0.05;降低 resolution 以减少面数 |
| W&B 日志未上传 | 环境变量未设置或项目名错误 | 检查 WANDB_USERNAME;确认 runner config 中 wandb_project 已配置 |
分布式训练报错 CPU device not supported |
--device cpu 与 --distributed 同时传入 |
分布式训练必须使用 GPU,改为 --device cuda |
阅读进阶
掌握工具脚本的使用后,建议继续深入以下主题以完善对整个训练与部署管线的理解:
- 若需修改或新增训练任务,请参考 自定义任务与环境扩展指南。
- 若需理解训练循环的内部机制,请参阅 On-Policy Runner 训练循环。
- 若准备将策略迁移到 MuJoCo 或真机,请参阅 MuJoCo Sim2Sim 部署与真机迁移。
- 若对 AMP 模仿学习的训练流程感兴趣,请参阅 AMP Runner 与模仿学习训练。