本项目在 IsaacLab 框架之上实现了 Adversarial Motion Priors (AMP) 模仿学习管线，使 ATOM01 双足机器人能够从人类动作捕捉数据中学习自然、流畅的运动风格。与纯强化学习策略相比，AMP 通过对抗性判别器为策略提供隐式的动作先验约束，在不牺牲任务完成度的前提下显著改善运动的生物拟真度。本页涵盖动作数据集的完整制备流程、AMP 环境架构、对抗训练算法核心以及训练配置参数。

AMP 核心原理与系统架构

AMP 的核心思想源自生成对抗网络（GAN）：一个 判别器（Discriminator） 学习区分真实的人类示范动作与策略生成的动作，而 策略网络（Policy） 则通过最大化判别器的"欺骗得分"来获得风格奖励（Style Reward）。该风格奖励与任务奖励（Task Reward）通过线性插值融合，驱动策略既完成任务指令又模仿人类动作特征。

系统整体数据流遵循"环境产出观测 → 判别器计算风格奖励 → 联合优化策略与判别器"的闭环。具体而言，AnimationEnv 在常规物理仿真的基础上增加了动作数据管理器（MotionDataManager）与动画管理器（AnimationManager），前者加载并采样 .pkl 格式的动作片段，后者按时间步推进参考动画并填充观测缓存。AmpEnv 进一步重写了 step 方法，确保在环境重置前保留用于判别器训练的观测历史。

flowchart TB
    subgraph 动作数据集层
        GMR[GMR原始动作数据 .pkl] --> RETARGET[Retarget工具链<br/>gmr_to_lab.py]
        RETARGET --> LAB[Lab格式动作数据<br/>root_pos / root_rot / dof_pos / key_body_pos]
    end

    subgraph 环境层
        LAB --> MDM[MotionDataManager<br/>加载 / 采样 / 插值]
        MDM --> AM[AnimationManager<br/>时间推进 / 观测缓存]
        AM --> AMP_ENV[AmpEnv<br/>step重写 / 观测保留]
    end

    subgraph 训练层
        AMP_ENV --> POLICY[Policy网络<br/>Actor-Critic]
        AMP_ENV --> DISC[AMPDiscriminator<br/>风格判别]
        POLICY --> PPOAMP[PPOAMP算法<br/>PPO + AMP联合更新]
        DISC --> PPOAMP
        PPOAMP --> AMP_ENV
    end

    style PPOAMP fill:#e1f5fe
    style DISC fill:#fff3e0

Sources: amp_env.py, animation_env.py

动作数据集制备与重定向

AMP 训练的前提是高质量、与机器人本体匹配的动作数据集。本项目的数据来源为 GMR 动作捕捉管线，输出格式包含 fps、root_pos、root_rot（xyzw 四元数）、dof_pos 等字段。由于 GMR 的关节顺序与 IsaacLab 的机器人模型不一致，项目提供了完整的重定向（Retarget）工具链将原始数据转换为 Lab 格式。

重定向流程分为两个主要步骤：关节重排序与关键体位置计算。gmr_to_lab.py 中的 extract_gmr_data 函数依据 YAML 配置文件中的 gmr_dof_names 到 lab_dof_names 的映射表对自由度数组进行重排；随后 run_simulator 在 Isaac Sim 中逐帧回放动作，驱动机器人模型运动并记录 key_body_pos（关键连杆在世界坐标系下的位置）。输出字典还包含 loop_mode 字段（0 为 CLAMP，1 为 WRAP），用于控制动画在末端是冻结还是循环。

flowchart LR
    A[GMR .pkl] --> B{YAML配置<br/>atom01.yaml}
    B --> C[extract_gmr_data<br/>关节映射 / 裁剪]
    C --> D[run_simulator<br/>物理回放 / 关键体记录]
    D --> E[Lab .pkl<br/>wxyz四元数 / key_body_pos]

项目中已预置的动作数据集位于 data/motions/atom01_lab/，涵盖行走、跑步、转向、站立等 16 组动作片段，数据来源包括 CMU、ACCAD 以及 GVHMR：

Sources: gmr_to_lab.py, dataset_retarget.py, atom01.yaml

批量重定向命令

对于批量处理，使用 dataset_retarget.py 可以一次性转换整个目录：

python scripts/tools/retarget/dataset_retarget.py \
    --robot atom01 \
    --input_dir robolab/data/motions/atom01_gmr \
    --output_dir robolab/data/motions/atom01_lab \
    --config_file robolab/scripts/tools/retarget/config/atom01.yaml \
    --loop clamp

Sources: dataset_retarget.py

AMP 环境架构详解

AMP 环境采用三层继承结构：ManagerBasedRLEnv → AnimationEnv → AmpEnv。每层分别负责基础强化学习逻辑、参考动画驱动、以及 AMP 特有的观测保留机制。

AnimationEnv 与动画管理器

AnimationEnv 在 load_managers 阶段实例化了两个专属管理器：MotionDataManager 负责从磁盘加载 .pkl 动作文件并进行帧级插值；AnimationManager 则为每个并行环境维护一条参考动画时间线，按 step_dt 推进并缓存最近 num_steps_to_use 步的参考观测。当环境触发重置（reset_buf）时，AnimationManager.reset 会依据动作权重分布重新采样 motion_id，并随机初始化时间相位，从而保证训练初期的状态多样性。

MotionDataTerm 是底层的数据加载单元。它通过 joblib 读取动作文件，在初始化阶段预计算 root 速度与关节速度的数值微分，并将所有动作张量沿帧维度拼接存储。运行时通过 get_motion_state 执行帧索引查找与线性插值：先调用 calc_phase 将时间映射到 [0, 1] 相位（支持 WRAP 循环模式），再转换为帧索引并做 lerp 插值。关键体位置从世界系转换到机体系，使用 root 四元数的逆旋转完成。

Sources: motion_data_manager.py, animation_manager.py

AmpEnv 的 Step 重写

标准 ManagerBasedRLEnv 在环境终止后会先执行重置再计算观测，这会导致 AMP 判别器丢失终止前一帧的观测历史。AmpEnv.step 保留了父类的大部分逻辑，但将观测计算延后到重置完成之后，确保 observation_manager.compute(update_history=True) 捕获到正确的时序数据。该设计对于依赖历史窗口（history_length=3 或 10）的判别器观测至关重要。

Sources: amp_env.py

观测组配置

AMP 环境定义了四组观测：policy、critic、disc（判别器策略观测）、disc_demo（判别器示范观测）。策略组包含机身角速度、投影重力、速度指令、关节位置与速度、上一动作，维度精简以利于实时推理。判别器组则捕捉更具风格辨识度的特征，如关节位置、关节速度、机身角速度，历史长度默认为 10 步。示范观测组 disc_demo 从 AnimationManager 的缓存中读取对应时刻的参考值，确保判别器在相同状态空间下比较策略与示范。

flowchart LR
    subgraph PolicyObs["policy (history_length=1)"]
        P1[base_ang_vel]
        P2[projected_gravity]
        P3[velocity_commands]
        P4[joint_pos_rel]
        P5[joint_vel_rel]
        P6[last_action]
    end

    subgraph DiscObs["disc / disc_demo (history_length=10)"]
        D1[base_ang_vel]
        D2[joint_pos]
        D3[joint_vel]
    end

Sources: amp_env_cfg.py

对抗训练算法实现

项目的 AMP 训练后端基于 rsl_rl 库扩展，核心组件包括 PPOAMP 算法、AMPDiscriminator 判别器、AMPRunner 训练运行器以及 CircularBuffer 环形缓存。

AMPDiscriminator 判别器网络

AMPDiscriminator 是一个 MLP 二分类器，输入为展平后的时序观测（disc_obs_dim × disc_obs_steps），经过若干隐藏层后输出标量判别分数。网络支持三种损失模式：

判别器包含一个 EmpiricalNormalization 输入归一化层，在训练过程中在线更新均值与方差。compute_grad_penalty 对示范数据施加梯度惩罚（默认缩放系数 10.0），促使判别器在真实数据附近的梯度趋近于零，有效抑制模式崩溃。

Sources: amp.py

PPOAMP 联合优化

PPOAMP 继承自标准 PPO，在 process_env_step 阶段执行以下操作：首先提取当前策略观测 disc_obs 与参考动画观测 disc_demo_obs，调用判别器计算 style_rewards，再将任务奖励与风格奖励按 task_style_lerp 插值得到最终奖励 rewards_lerp。两组观测分别追加到独立的 CircularBuffer 中，供后续判别器训练批次采样。

update 方法实现了策略与判别器的交替梯度更新。每轮 mini-batch 迭代中，PPO 损失（surrogate + value + entropy）与判别器损失独立计算、独立反向传播。判别器损失由策略数据的分类损失、示范数据的分类损失以及梯度惩罚项组成。优化器层面，策略使用 self.optimizer，判别器使用独立的 disc_optimizer（Adam），两者的学习率与梯度裁剪阈值分别配置。

Sources: ppo_amp.py

AMPRunner 与模型持久化

AMPRunner 继承自 OnPolicyRunner，重写了 learn、save、load、train_mode、eval_mode 五个方法。learn 在 rollout 阶段额外提取 style_rewards 与 total_rewards 并送入 LoggerAMP 记录。save 将判别器网络状态、判别器归一化状态以及判别器优化器状态一并写入 checkpoint，保证训练中断后可完整恢复。load 则对称地恢复上述全部状态。

Sources: amp_runner.py, amp_logger.py

CircularBuffer 环形缓存

CircularBuffer 为每个并行环境维护一个固定长度的时序数据缓存，支持按 mini-batch 随机采样。其内部以 (max_len, batch_size, ...) 的张量存储数据，通过 _pointer 循环写入。mini_batch_generator 方法在每个 epoch 内将缓存数据随机打乱后切分为多个 mini-batch，满足判别器对非时序独立同分布采样的需求。

Sources: circular_buffer.py

训练配置参数解析

ATOM01 的 AMP 训练配置集中在 atom01_amp_agent_cfg.py 中，分为网络结构、PPO 超参数、AMP 超参数三个层次。

网络与算法配置

policy = RslRlPpoActorCriticCfg(
    init_noise_std=1.0,
    actor_hidden_dims=[512, 256, 128],
    critic_hidden_dims=[512, 256, 128],
    activation="elu",
)
algorithm = RslRlPpoAmpAlgorithmCfg(
    class_name="PPOAMP",
    learning_rate=1.0e-4,
    num_learning_epochs=5,
    num_mini_batches=4,
    schedule="adaptive",
    desired_kl=0.01,
    symmetry_cfg=...,
    amp_cfg=...,
)

策略与价值网络均采用 [512, 256, 128] 的三层 MLP，激活函数为 elu。学习率启用自适应调度：当 KL 散度超过 0.02 时除以 1.5，低于 0.005 时乘以 1.5。对称增强（Symmetry）开启数据增强与镜像损失，镜像损失系数为 0.2，通过 atom01.compute_symmetric_states 实现左右肢体的状态镜像。

AMP 专属超参数

判别器网络结构为 [1024, 512]，大于策略网络，以提供充足的表示能力来捕捉动作风格的细微差异。task_style_lerp=0.3 意味着最终奖励中 30% 来自风格奖励、70% 来自任务奖励，该平衡在跟踪指令与自然步态之间取得折中。

Sources: atom01_amp_agent_cfg.py

动作数据权重与课程

在 Atom01AmpEnvCfg 中，通过 motion_data_weights 字典为每个动作片段分配采样权重。所有权重在初始化时归一化为概率分布，供 MotionDataTerm.sample_motions 进行多项式采样。当前配置中 CMU、ACCAD、GVHMR 三类数据权重均为 1，实现均匀混合。如需侧重某一类动作风格，可直接调整对应权重值。

环境奖励函数在 AMP 模式下进行了简化：任务跟踪奖励（track_lin_vel_xy_exp、track_ang_vel_z_exp）的权重设为 0，因为风格奖励已隐式编码了速度模式；保留的惩罚项包括非期望接触（undesired_contacts）与存活奖励（alive），用于维持基本的物理可行性。

Sources: atom01_amp_env_cfg.py

训练与推理

启动 AMP 训练

cd modules/atom01_train/robolab
python scripts/rsl_rl/train.py --task=Atom01-Amp-v0 --headless --logger=tensorboard --num_envs=8192

训练默认使用 8192 个并行环境，单步环境步长 24，总迭代次数 5000。每 100 次迭代保存一次模型 checkpoint。

播放与导出策略

# 可视化策略
python scripts/rsl_rl/play_amp.py --task=Atom01-Amp-v0 --num_envs=1 --load_run <run_dir> --checkpoint <model.pt>

# Sim2Sim 迁移到 MuJoCo
python scripts/mujoco/sim2sim_atom01_amp.py --load_model <exported/policy.pt>

play_amp.py 支持将策略导出为 ONNX 格式，便于后续在 OrangePi 主控板上通过 ONNX Runtime 部署。sim2sim_atom01_amp.py 专门用于验证 AMP 策略在 MuJoCo 物理引擎上的零样本迁移能力。

Sources: play_amp.py, sim2sim_atom01_amp.py

进阶话题与调参建议

判别器输入设计

判别器观测的选择直接决定 AMP 能捕捉到的风格维度。当前配置使用 base_ang_vel、joint_pos、joint_vel，舍弃了 root 线速度与关键体位置。若发现策略无法复现特定的躯干倾斜或摆臂风格，可在 amp_env_cfg.py 中取消注释 base_lin_vel 或 key_body_pos_b 项，并同步更新 disc_obs_dim。

风格奖励与任务奖励的平衡

task_style_lerp 是控制策略"服从指令"还是"模仿风格"的关键旋钮。当 lerp=0 时，策略完全依赖任务奖励，退化为普通 PPO；当 lerp=1 时，策略完全模仿动作数据，可能忽视速度指令。对于需要精确跟踪速度指令的足式机器人，建议保持在 0.2 ~ 0.4 区间。

梯度惩罚与训练稳定性

若观察到判别器损失剧烈震荡或风格奖励收敛到极值，可尝试：1) 增大 grad_penalty_scale 到 20.0；2) 降低 disc_learning_rate 到 5.0e-5；3) 切换为 WGAN 损失并启用 disc_output_normalizer。LSGAN 在本项目中表现最为稳定，是推荐的默认配置。

阅读衔接

完成 AMP 训练后，若希望进一步扩展至多技能切换与注意力机制，请参阅 BeyondMimic多技能与注意力策略。若需将训练好的策略迁移至 MuJoCo 并进行硬件部署前的验证，请参阅 Sim2Sim迁移与MuJoCo验证。AMP 环境依赖 IsaacLab 的 Manager-based 架构，若对环境配置原理尚不熟悉，建议先回顾 IsaacLab环境搭建与训练流程 与 基础Locomotion策略训练。