本文档介绍ATOM01足式机器人基于NVIDIA IsaacLab的仿真强化学习环境搭建与训练工作流，涵盖从依赖安装、环境注册、训练启动到策略导出的完整操作路径。阅读本文前，建议先完成URDF运动学与动力学描述和MuJoCo与Isaac Sim模型配置的相关配置，以确保机器人模型文件已就位。

Sources: README_CN.md

技术架构与组件关系

ATOM01的训练系统采用隔离式扩展架构：核心仿真引擎IsaacSim与IsaacLab作为底层依赖独立安装，项目代码（robolab环境与rsl_rl算法库）以独立仓库形式存在，通过Python包注册机制与IsaacLab解耦。这种设计避免了直接修改IsaacLab源码，降低了版本升级时的重构成本。

robolab负责任务环境定义（场景、奖励、观测、地形），rsl_rl负责策略优化（PPO、AMP、Distillation等算法）。两者通过gymnasium注册表与IsaacLab的DirectRLEnv/ManagerBasedRLEnv接口衔接。

flowchart TB
    subgraph 上层工具层
        A[训练脚本 train.py]
        B[测试脚本 play.py]
        C[Sim2Sim脚本 sim2sim_atom01.py]
    end

    subgraph 算法层
        D[rsl_rl<br/>OnPolicyRunner / AMPRunner / DistillationRunner]
    end

    subgraph 环境层
        E[robolab.tasks.direct<br/>BaseEnv / AttnEncEnv / InterruptEnv]
        F[robolab.tasks.manager_based<br/>AmpEnv / BeyondMimicEnv]
    end

    subgraph 底层引擎层
        G[IsaacLab 2.3.2]
        H[IsaacSim 5.1.0]
    end

    A --> D
    B --> D
    D --> E
    D --> F
    E --> G
    F --> G
    G --> H

Direct RL Env与Manager-Based RL Env是IsaacLab提供的两种环境范式。本项目中的基础运动（Locomotion）、注意力编码（AttnEnc）和中断恢复（Interrupt）采用DirectRLEnv实现，代码透明度高，便于精细控制物理步进与观测计算；AMP模仿学习与BeyondMimic多技能策略采用ManagerBasedRLEnv实现，利用IsaacLab的Manager系统模块化组合奖励、终止条件与命令生成器。

Sources: README_CN.md, init.py

前置条件与系统要求

在开始安装前，请确保软硬件环境满足以下最低要求。

Sources: README_CN.md

环境安装步骤

安装遵循先底层后上层的顺序：先通过conda/pip部署IsaacLab，再将本项目作为独立扩展包安装。

flowchart LR
    A[安装IsaacLab<br/>conda环境] --> B[克隆 atom01_train<br/>仓库]
    B --> C[初始化子模块<br/>git submodule update]
    C --> D[安装 robolab<br/>pip install -e .]
    D --> E[安装 rsl_rl<br/>pip install -e .]
    E --> F[验证环境列表<br/>python list_envs.py]

步骤一：安装IsaacLab。按照IsaacLab官方文档完成conda安装。推荐使用conda方案，因为它简化了终端调用Python脚本的路径管理，并自动处理CUDA与PyTorch版本对齐。

步骤二：克隆仓库（需在IsaacLab目录之外执行，保持隔离性）：

git clone https://github.com/Roboparty/atom01_train.git
cd atom01_train
git submodule update --init --recursive

步骤三：安装robolab扩展包：

cd robolab
pip install -e .

setup.py中声明了robolab的额外依赖，包括mujoco==3.3.3、mujoco-python-viewer、psutil、joblib和pynput，用于后续Sim2Sim验证与工具脚本。

Sources: setup.py

步骤四：安装rsl_rl算法库：

cd ../rsl_rl
pip install -e .

该子模块基于ETH Zurich的rsl_rl进行扩展，pyproject.toml中指定了torch>=2.6.0、tensordict>=0.7.0、onnx等依赖，用于支持策略网络训练与导出。

Sources: pyproject.toml

步骤五：验证安装。运行以下命令打印所有已注册的ATOM01环境：

python robolab/scripts/tools/list_envs.py

若输出包含Atom01-Flat、Atom01-Rough等任务名称，说明扩展包已正确挂载到gymnasium注册表。

Sources: list_envs.py

已注册环境总览

本项目通过gymnasium.register在导入robolab.tasks时自动注册以下环境。环境ID可直接作为--task参数传入训练脚本。

各环境的entry_point指向具体的环境类（如BaseEnv、AmpEnv），env_cfg_entry_point指向配置类（如ATOM01FlatEnvCfg），rsl_rl_cfg_entry_point指向智能体配置类（如ATOM01FlatAgentCfg）。这种分离设计允许用户在不修改环境代码的情况下，通过Hydra配置或CLI参数覆盖超参数。

Sources: init.py, init.py, init.py, init.py, init.py

训练工作流

启动训练

以Atom01-Flat环境为例，基础训练命令如下：

python robolab/scripts/rsl_rl/train.py \
    --task=Atom01-Flat \
    --headless \
    --logger=tensorboard \
    --num_envs=8192

train.py脚本的工作流程如下：

sequenceDiagram
    participant User
    participant train.py
    participant AppLauncher
    participant gym
    participant RslRlVecEnvWrapper
    participant Runner

    User->>train.py: 传入 --task 等参数
    train.py->>AppLauncher: 解析并启动IsaacSim
    AppLauncher-->>train.py: 返回 simulation_app
    train.py->>gym: gym.make(task, cfg=env_cfg)
    gym-->>train.py: 返回 IsaacLab 环境实例
    train.py->>RslRlVecEnvWrapper: 包装环境（动作裁剪）
    RslRlVecEnvWrapper-->>train.py: 返回 VecEnv
    train.py->>Runner: 实例化 OnPolicyRunner / AMPRunner
    Runner-->>train.py: 返回 runner
    train.py->>Runner: runner.learn(max_iterations)
    loop 每次迭代
        Runner->>RslRlVecEnvWrapper: env.step(actions)
        RslRlVecEnvWrapper->>gym: 执行物理仿真
        gym-->>RslRlVecEnvWrapper: 返回 obs, reward, done
        RslRlVecEnvWrapper-->>Runner: 返回转移数据
        Runner->>Runner: PPO/AMP 参数更新
    end
    Runner-->>train.py: 训练结束
    train.py->>AppLauncher: simulation_app.close()

Sources: train.py

关键CLI参数说明

RSL-RL专用参数通过cli_args.py统一管理，支持覆盖实验名称（--experiment_name）、运行名后缀（--run_name）和日志项目名（--log_project_name）。

Sources: cli_args.py

算法Runner类型选择

train.py根据agent_cfg.class_name自动实例化对应的Runner：

• OnPolicyRunner：标准PPO训练，适用于Atom01-Flat、Atom01-Rough、Atom01-AttnEnc、Atom01-Interrupt。
• AMPRunner：Adversarial Motion Prior训练，适用于Atom01-AMP。
• DistillationRunner：策略蒸馏，适用于Atom01-Interrupt等需要蒸馏的场景。

实例化前，环境会被RslRlVecEnvWrapper包装，统一动作裁剪（clip_actions）并适配RSL-RL的VecEnv接口。

Sources: train.py

恢复训练与多GPU配置

恢复训练：添加--resume即可从logs/rsl_rl/<experiment_name>/下最近的检查点继续。也可通过--load_run和--checkpoint精确定位。

多GPU训练：添加--distributed后，脚本会自动为每个进程分配独立的cuda:{local_rank}设备，并基于local_rank偏移随机种子，保证各进程探索多样性。注意多GPU训练不支持CPU设备（--device cpu），必须显式指定GPU。

Sources: train.py

策略验证与导出

Play模式

训练完成后，使用play.py加载策略进行可视化验证：

python robolab/scripts/rsl_rl/play.py \
    --task=Atom01-Flat \
    --num_envs=1

Play脚本支持以下附加选项：

• --plane：强制使用平面地形，便于观察步态。
• --push_robot：在运行中施加推力，测试鲁棒性。
• --real-time：尽可能按实时速率渲染。
• --disable_fabric：禁用Fabric缓存，回退到标准USD I/O。

Sources: play.py

模型导出

Play脚本内置策略导出功能。加载检查点后，脚本会自动将策略导出为Torch JIT（.pt）和ONNX（.onnx）两种格式，存放于检查点同级目录的exported/文件夹下。对于标准ActorCritic策略，调用export_policy_as_jit与export_policy_as_onnx完成导出；对于注意力编码策略（AttnEnc），脚本内嵌了TorchAttnEncPolicyExporter和OnnxAttnEncPolicyExporter，将encoder与actor组装为统一的前向图，确保导出后的模型与Sim2Sim、部署阶段的输入输出维度一致。

导出的ONNX模型将在后续推理节点与ONNX策略部署章节中用于ROS2推理节点部署。

Sources: play.py

Sim2Sim验证（简要）

训练导出的.pt策略可通过MuJoCo进行Sim2Sim验证：

python robolab/scripts/mujoco/sim2sim_atom01.py \
    --load_model /path/to/exported/policy.pt

项目为不同算法准备了对应的Sim2Sim入口：sim2sim_atom01.py（基础）、sim2sim_atom01_amp.py（AMP）、sim2sim_atom01_attn_enc.py（注意力编码）、sim2sim_atom01_bm.py（BeyondMimic）、sim2sim_atom01_interrupt.py（中断恢复）。Sim2Sim的详细操作流程与验证指标请参见Sim2Sim迁移与MuJoCo验证。

Sources: README_CN.md

数据集准备工具（AMP / BeyondMimic）

对于AMP和BeyondMimic训练，需要先准备动作数据集。项目提供scripts/tools/retarget/下的重定向工具链，将GMR（参考GMR仓库）生成的动作数据转换为IsaacLab可用的格式。由于GMR数据集的关节顺序与IsaacLab内部顺序不同，需事先准备关节映射YAML（如scripts/tools/retarget/config/atom01.yaml），然后运行dataset_retarget.py完成批量转换。

转换后的数据集存放在data/motions/目录下，按算法分子目录（atom01_gmr、atom01_lab、atom01_bm）。

Sources: README_CN.md

故障排查

下一步

完成环境搭建并成功运行基础训练后，建议按以下路径继续深入：

• 如需理解基础运动策略的奖励设计与训练技巧，阅读基础Locomotion策略训练。
• 如需使用AMP进行模仿学习，阅读AMP模仿学习与动作数据集。
• 如需训练多技能与注意力策略，阅读BeyondMimic多技能与注意力策略。
• 如需将策略迁移到MuJoCo验证，阅读Sim2Sim迁移与MuJoCo验证。