AMP(Adversarial Motion Prior)是一种将对抗模仿学习嵌入策略优化框架的算法,其核心目标是在保留任务导向能力的同时,让智能体的动作风格逼近专家示范轨迹。与纯行为克隆不同,AMP 不直接拟合专家动作分布,而是训练一个判别器来区分"策略生成的轨迹片段"与"专家示范片段",并将判别器的置信度转化为额外的风格奖励信号,与任务奖励线性混合后驱动 PPO 策略学习。这种设计允许策略在对抗博弈中自发地探索与专家风格一致但又能完成特定任务的动作模式。
Sources: modules/amp.py, algorithms/ppo_amp.py
核心设计思想
AMP 的底层逻辑是一个双人极小极大博弈:判别器试图准确区分智能体轨迹与专家轨迹,而策略则试图生成让判别器无法区分的轨迹。与传统的 GAIL 不同,AMP 将判别器的输入限制在低维的运动学观测空间(如关节位置、速度、根姿态等),而非完整的策略观测空间。这意味着判别器只评判"动作风格"是否像专家,而不关心任务相关的上下文信息(如目标位置、地形感知等)。策略网络仍然通过标准的 Actor-Critic 架构接收完整任务观测并输出动作,但它在优化目标中额外获得了一项来自判别器的风格奖励。
Sources: modules/amp.py
系统架构与模块交互
下图展示了 AMP 训练流程中各核心模块的交互关系。环境在每个时间步输出两组观测:discriminator(智能体当前运动学状态的历史窗口)和 discriminator_demonstration(来自专家数据集的对齐窗口)。AMPDiscriminator 对这两组输入分别打分,PPOAMP 将分数转换为风格奖励并与任务奖励插值,最终通过标准 PPO 更新策略参数,同时以对抗损失独立更新判别器参数。
flowchart TD
Env[向量化环境 VecEnv] -->|obs: discriminator / discriminator_demonstration| AMP[AMPDiscriminator]
Env -->|task obs, rewards, dones| PPO[PPOAMP 策略更新]
AMP -->|disc_score| SR[风格奖励 predict_style_reward]
SR -->|style_reward| LR[奖励混合 lerp_reward]
Env -->|task_reward| LR
LR -->|rewards_lerp| PPO
PPO -->|actions| Env
Buffer1[CircularBuffer<br/>disc_obs_buffer] -.->|mini-batch| AMP
Buffer2[CircularBuffer<br/>disc_demo_obs_buffer] -.->|mini-batch| AMP
AMP -->|disc_loss + grad_penalty| DiscOpt[判别器优化器 disc_optimizer]
PPO -->|surrogate + value + entropy| PolOpt[策略优化器 optimizer]
Sources: runners/amp_runner.py, algorithms/ppo_amp.py
判别器网络与观测组设计
AMPDiscriminator 是一个浅层 MLP 分类器,其输入维度为 disc_obs_dim × disc_obs_steps,即把历史时间步内的运动学观测拼接成一个扁平向量。网络结构由三层隐藏层(默认 [256, 256, 256])和一个单神经元输出层组成。为了稳定训练,判别器对输入观测使用 EmpiricalNormalization 进行在线均值方差归一化,该归一化模块会在训练过程中持续累积统计量直到达到预设的上限(默认 1e8 个样本)。
观测组配置要求环境同时提供 discriminator 和 discriminator_demonstration 两个键,且它们的历史长度和拼接后总维度必须严格一致。resolve_amp_config 函数会在初始化时自动从环境观测张量中推断这些维度,并将其写入配置字典,避免用户手动计算。
Sources: modules/amp.py, modules/amp.py
三种对抗损失变体
框架内置了三种 GAN 损失变体,通过 LossType 枚举选择。不同变体在判别器的输出解释、奖励计算方式和训练稳定性上存在显著差异。
| 损失类型 | 判别器目标 | 风格奖励公式 | 输出归一化 | 适用场景 |
|---|---|---|---|---|
| GAN | 标准二分类 BCE | -log(1 - σ(disc_score)) |
无 | 经典实现,但可能面临梯度饱和 |
| LSGAN | 最小二乘 MSE | clamp(1 - 0.25·(disc_score - 1)², min=0) |
无 | 训练更稳定,梯度更平滑 |
| WGAN | Wasserstein 距离 | norm(disc_score) |
有(EmpiricalNormalization) | 需配合梯度惩罚,理论上避免模式崩溃 |
在 update 阶段,判别器对智能体轨迹的期望输出为负(GAN 中标签为 0,LSGAN 中目标为 -1),对专家轨迹的期望输出为正。WGAN 则直接最大化专家分数与策略分数之间的间隔。三种变体均额外施加梯度惩罚项 compute_grad_penalty,该项在示范数据上计算梯度范数与零的平方误差,以约束判别器的 Lipschitz 常数。
Sources: modules/amp.py, modules/amp.py, algorithms/ppo_amp.py
风格奖励计算与任务奖励混合
风格奖励的计算发生在环境交互阶段(process_env_step),而非策略更新阶段。判别器以推理模式接收智能体的 disc_obs,经过归一化后前向传播得到 disc_score,再根据选定的 LossType 映射为原始奖励值。最终风格奖励会乘以时间步长 dt 和缩放系数 style_reward_scale,以消除帧率差异对奖励幅度的影响。
任务奖励与风格奖励通过线性插值混合:rewards_lerp = task_style_lerp × task_reward + (1 - task_style_lerp) × style_reward。当 task_style_lerp = 0 时,策略完全依赖风格奖励模仿专家动作;当 task_style_lerp = 1 时,AMP 退化为标准 PPO。该参数为用户提供了在"任务完成度"与"动作逼真度"之间进行显式权衡的单一旋钮。
Sources: modules/amp.py, modules/amp.py, algorithms/ppo_amp.py
数据流与环形缓冲区
由于判别器需要同时访问智能体轨迹和专家轨迹,而专家数据通常以预录数据集形式存在,AMP 使用两个独立的 CircularBuffer 分别缓存 disc_obs 和 disc_demo_obs。这两个缓冲区的容量由配置中的 disc_obs_buffer_size 决定。在每次 process_env_step 时,当前步的判别器观测被追加到对应缓冲区;在 update 时,缓冲区通过 mini_batch_generator 按随机排列的线性索引采样,生成与 PPO rollout 数据对齐的小批量。
CircularBuffer 的设计支持多环境并行:每个环境拥有独立的写入指针和推送计数器,采样时从所有环境的有效历史中进行无放回随机抽取,确保每个小批量都包含多样化的环境-时间步组合。
Sources: storage/circular_buffer.py, algorithms/ppo_amp.py, algorithms/ppo_amp.py
训练更新流程
PPOAMP.update() 在标准 PPO 的内循环基础上扩展了判别器训练步骤。具体而言,每个 epoch 的每个 mini-batch 都会执行以下操作序列:首先计算 PPO 的 surrogate loss、value loss 和 entropy loss;随后对判别器小批量进行归一化并前向传播,分别得到智能体轨迹分数与专家轨迹分数;接着根据所选损失类型计算 disc_loss,并叠加在专家数据上计算的梯度惩罚 disc_grad_penalty;最后,策略参数和判别器参数分别通过各自的优化器进行梯度下降。判别器的梯度裁剪阈值由 disc_max_grad_norm 控制,默认通常为 0.5,比策略的 1.0 更为保守。
值得注意的是,判别器的观测归一化统计量仅在策略反向传播之后、参数更新之前通过 update_normalization 进行增量更新,以确保当前批次使用的归一化参数与梯度计算时一致。
Sources: algorithms/ppo_amp.py
Runner 生命周期与存储管理
AMPRunner 继承自 OnPolicyRunner,重写了算法构造、模式切换、检查点保存与加载逻辑。在 _construct_algorithm 中,Runner 除了初始化策略和 RolloutStorage 外,还额外实例化两个 CircularBuffer 并注入 PPOAMP 的构造函数。train_mode 和 eval_mode 被扩展以同步切换判别器及其归一化模块的训练/评估状态。
保存检查点时,save 方法将判别器的网络状态、归一化状态以及独立优化器状态一并序列化;load 方法则按对称逻辑恢复。这保证了训练中断后可以从任意迭代点无缝续训,而不会丢失判别器的对抗学习进度。
Sources: runners/amp_runner.py, runners/amp_runner.py
日志与监控
LoggerAMP 继承基础 Logger,专门追踪 AMP 相关的训练指标。在每个环境步中,它累积风格奖励和混合后的总奖励;当 episode 结束时,计算每秒钟的平均风格奖励并写入 Episode_Reward/style。此外,训练日志还包含判别器损失(amp/disc_loss)、梯度惩罚(amp/disc_grad_penalty)、智能体轨迹平均分数(amp/disc_score)和专家轨迹平均分数(amp/disc_demo_score)。通过监控智能体分数与专家分数的趋近程度,开发者可以直观判断对抗博弈是否达到纳什均衡。
Sources: utils/amp_logger.py
AMP 与标准 PPO 的关键差异
| 维度 | 标准 PPO (OnPolicyRunner) |
AMP (AMPRunner) |
|---|---|---|
| 算法类 | PPO |
PPOAMP(继承 PPO) |
| 额外模块 | 无 | AMPDiscriminator + 两个 CircularBuffer |
| 奖励信号 | 仅任务奖励 | 任务奖励 + 风格奖励(插值混合) |
| 优化器数量 | 1(策略) | 2(策略 + 判别器) |
| 观测组要求 | critic,可选 rnd_state |
额外要求 discriminator 和 discriminator_demonstration |
| 保存内容 | 策略、优化器、RND(如有) | 额外保存判别器、判别器归一化、判别器优化器 |
| 日志指标 | 标准 PPO 损失与奖励 | 额外增加风格奖励、判别器分数、判别器损失 |
Sources: algorithms/ppo.py, algorithms/ppo_amp.py, runners/on_policy_runner.py
配置要点与常见注意事项
AMP 的配置字典 amp_cfg 必须在算法配置中显式提供,否则 PPOAMP 初始化将抛出异常。关键配置项包括:loss_type(选择 GAN/LSGAN/WGAN)、disc_obs_buffer_size(环形缓冲区容量)、step_dt(由 resolve_amp_config 自动从环境提取)、disc_learning_rate(通常低于策略学习率)、grad_penalty_scale(梯度惩罚系数)以及 task_style_lerp(任务-风格奖励插值权重)。此外,判别器主干和输出层可以分别配置 weight_decay,以对不同深度的参数施加差异化的 L2 正则化。
需要特别注意的是,环境必须保证 discriminator 和 discriminator_demonstration 两组观测具有完全相同的历史步长和特征维度,否则 resolve_amp_config 会在初始化阶段断言失败。
Sources: modules/amp.py, algorithms/ppo_amp.py
总结
AMP 算法通过将对抗模仿学习嵌入 PPO 框架,实现了任务目标与动作风格的双目标优化。其核心创新在于利用一个轻量级判别器为策略提供密集的风格奖励信号,同时保持策略网络对完整任务观测的访问能力。框架通过观测组分离、环形缓冲区和独立优化器的设计,确保了对抗训练与策略训练的稳定耦合。对于需要在仿真中复现真实运动捕捉数据风格(如人形机器人行走、跑酷)的场景,AMP 提供了一种比纯行为克隆更具泛化性的解决方案。
下一步阅读建议
理解了 AMP 的对抗训练机制后,建议继续阅读以下页面以建立完整的训练图景:
- 若希望深入理解策略网络如何接收和切分多组观测,请参阅 Actor-Critic 基础架构设计。
- 若对环形缓冲区的采样机制与 rollout 数据的组织方式感兴趣,请参阅 Rollout 数据存储与 Transition 和 经验采样与小批量生成。
- 若需对比另一种利用示范数据的训练范式,请参阅 策略蒸馏与师生框架。
- 若关注训练运行器的完整生命周期与分布式支持,请参阅 训练运行器生命周期管理。