Skip to content

l2/rl/ — L2 来源 #2:训练 PPO + rollout 录制

在 pipeline 里的位置

L2 数据生成层中等难度的一支:用 RSL-RL PPO 在 IsaacLab 自带的 Isaac-Lift-Cube-Franka-v0 (joint_pos task)上学一个 policy,然后让它在 sim 里 rollout,把 (obs, action) 录到 HDF5。

数据金字塔角色:量级中等(10³–10⁴ 条),主要价值是自然产生 recovery / 长尾轨迹——RL 探索期 + 收敛后的失败案例都是 scripted 永远生成不出来的数据分布。

关键设计:训练 joint_pos / 录制时换 IK-abs

l2/scripted/ 比起来,这里有个反直觉的点:训练时 PPO 学的是 9d joint_pos action, 但写到 HDF5 的是 8d IK-abs 等价 action(对齐 scripted/teleop 的 schema)。

为什么这么做

  • IsaacLab 只对 Isaac-Lift-Cube-Franka-v0(joint_pos)注册了 PPO/SKRL/SB3 cfg; Isaac-Lift-Cube-Franka-IK-Abs-v0 没有 RL agent cfg。
  • PPO 直接学 7d EE pose(含 4d unit quat)的 box action space 不收敛。改 6D rotation 要写 自定义 ActionTerm + 自定义 reward,工程量大且教学价值低。
  • 工程最干净、教学含金量最高的方案是直接复用 IsaacLab 验证过的 joint_pos PPO 配方,把 action 空间对齐挪到 rollout 阶段当数据后处理做

怎么实现的

  • recorders.py 里两个自定义 RecorderTerm:
    • IKAbsActionRecorder:每步从 ee_frame sensor 当前 EE pose + finger joint 状态 二值化拼成 8d IK-abs action,覆写默认的 actions 列。
    • PolicySplitObsRecorder:rollout 时 concatenate_terms=True (PPO 网络硬要求 1-D 输入),本 term 按 ObservationManager.group_obs_term_dim["policy"] 把 1-D obs 切回 dict,同时用上面的 8d action 替换 actions term(原 9d joint_pos last_action)。
  • env_cfg.pyRolloutRecorderManagerCfg 把默认的 record_pre_step_actionsrecord_pre_step_flat_policy_observations 设为 None 禁掉,让自定义的两个 term 接管。

详细推理见 obs/action 对齐 末尾"l2/rl 实现细节"。

怎么跑

bash
# 1. 训 PPO(默认 1024 envs × 600 iters,A100 上 ~10–20 分钟到 sanity 收敛)
uv run --project l2 python l2/rl/train.py
# 输出:l2/rl/checkpoints/lift_cube/<时间戳>/model_<iter>.pt
#       l2/rl/checkpoints/lift_cube/latest.pt -> 上面那个

# 2. 用 checkpoint rollout 录数据
uv run --project l2 python l2/rl/rollout_and_record.py \
    --checkpoint l2/rl/checkpoints/lift_cube/latest.pt \
    --num_demos 100 \
    --num_envs 32
# 输出:data/demos/rl_lift_cube.hdf5

跑完之后跟 scripted 同样过一遍 schema 检查:

bash
uv run --project l2 python tools/h5_inspect.py data/demos/rl_lift_cube.hdf5

并丢进 L3 转换器,合并多来源 dataset:

bash
cd l3
uv run python convert.py ../data/demos/ ../data/lerobot/lift_cube_v0_mixed/
# convert.py 按文件名前缀(scripted_* / rl_*)自动给每条 episode 打 source 字段

关键约定

  • 录制基础设施沿用 IsaacLab RecorderManager——直接在 env_cfg 上挂自定义的 RolloutRecorderManagerCfg,不要自己写 HDF5 buffer。
  • success 由 terminations.success DoneTerm 自动判定——加进 env_cfg(同 l2/scripted), RecorderManager.record_pre_reset 在每个 reset env 上读它的值写到 episode attr。 EXPORT_SUCCEEDED_ONLY 模式只保留 success=True 的 episode。
  • 训练 / rollout 共用同一个 PPO 网络结构(LiftCubePPORunnerCfg),都保 concatenate_terms=True。rollout 时录制器自己把 1-D obs 切回 dict 落盘,不影响 PPO 推理。
  • HDF5 schema 与 l2/scripted 完全对齐——8d action,5 个 obs term,L3 convert.py 不需要 按 source 分支处理。

文件清单

文件用途
train.pyRSL-RL PPO 训练入口(精简版,不走 IsaacLab hydra 流程)
rollout_and_record.py加载 checkpoint + 挂 RecorderManager + rollout 录数据
env_cfg.py训练 / rollout 两个 env_cfg 工厂(共享同一 task,差异在 obs concat 与 recorders)
recorders.py自定义 IKAbsActionRecorder / PolicySplitObsRecorder
__init__.pypackage 入口(让 import l2/rl.recorders 工作)
checkpoints/训练产物(gitignored,具体规则见仓库根 .gitignore)

注意事项

  • 训练耗 GPU 时(默认 1024 envs × 600 iters,A100 ~10–20 分钟到 sanity 收敛)。本任务相对简单, 到 100% success rate 通常 1500–2000 iters,要更高质量 demo 适当加 --max_iterations
  • rollout_and_record.py--num_envs 32 是经验值——更高吞吐但显存敏感;BC 训练只关心 episode 数和质量,不关心吞吐。
  • 推荐先跑通 l2/scripted 再来这步——scripted 提供 sanity check 的 baseline,且两边产物的 HDF5 schema 必须对齐(本仓库的合并多来源能力依赖这一点)。

与 IsaacLab 自带 train.py / play.py 的差异

  • 不走 hydra,直接 hardcode task_id + 构造 cfg。本仓库范围内不需要 hydra 的多任务覆盖能力。
  • 不导出 jit/onnx——下游 rollout_and_record.py 直接用 OnPolicyRunner.get_inference_policy 拿到 nn.Module 就够了,jit 主要给 ROS / 部署用。
  • rollout_and_record.py 自带 RecorderManager(IsaacLab 自带 play.py 没有),并且把 9d joint_pos action 重写成 8d IK-abs schema(IsaacLab 没有这个步骤)。

参考:L2 三种轨迹来源 §2、 obs/action 对齐 "l2/rl 实现细节"。