AMP（Adversarial Motion Priors）Runner 是本项目中实现模仿学习与强化学习协同训练的核心执行器。它在标准 On-Policy Runner 的基础上，引入了一个对抗性判别器网络（Discriminator），通过比较智能体运动轨迹与参考示范动作（Demonstration）之间的风格差异，生成额外的 Style Reward，并将其与任务奖励（Task Reward）融合，从而驱动策略学习出既满足任务目标、又具备参考运动风格的控制策略。本章面向具备深度强化学习基础的高级开发者，系统解析 AMP Runner 的架构设计、训练循环、联合优化机制以及配置调优方法。

AMP 架构全景

AMP Runner 的架构可以概括为**"双层继承 + 双缓冲存储 + 双损失联合优化"**。AMPRunner 继承自 OnPolicyRunner，仅扩展了与判别器相关的生命周期管理；PPOAMP 继承自 PPO，在每个环境交互步骤中插入样式奖励计算，并在策略更新阶段并行训练判别器；AMPDiscriminator 作为独立的神经网络模块，负责从观测特征中提取运动风格表征并输出判别分数。

flowchart TB
    subgraph Runner层
        A[AMPRunner] -->|继承| B[OnPolicyRunner]
    end
    
    subgraph 算法层
        C[PPOAMP] -->|继承| D[PPO]
        C --> E[AMPDiscriminator]
    end
    
    subgraph 存储层
        F[RolloutStorage<br/>PPO轨迹数据]
        G[CircularBuffer<br/>智能体Disc Obs]
        H[CircularBuffer<br/>示范Disc Obs]
    end
    
    subgraph 环境层
        I[VecEnv] -->|obs, reward| C
        I -->|disc_demo| C
    end
    
    A -->|管理| C
    C -->|读写| F
    C -->|读写| G
    C -->|读写| H
    E -->|predict_style_reward| C

在整个训练管线中，obs_groups 的配置起到了数据路由的关键作用。与标准 PPO 仅需 policy 和 critic 两组观测不同，AMP 必须额外声明 discriminator（智能体当前状态的判别器输入特征）和 discriminator_demonstration（来自参考运动数据的对应特征）两个观测组。resolve_amp_config 函数会在初始化时自动扫描这些观测组的维度，推断出判别器的输入尺寸和历史步长。环境在每一步返回的 TensorDict 中同时包含这两组特征，算法层通过 get_disc_obs 和 get_disc_demo_obs 分别提取并送入循环缓冲区。Sources: amp_runner.py, amp.py, atom01_amp_agent_cfg.py

AMPDiscriminator：风格判别的三种数学范式

AMPDiscriminator 是 AMP 训练范式的核心模块，其网络结构为简单的 MLP 编码器（disc_trunk）加单输出线性层（disc_linear），输入维度为 disc_obs_dim * disc_obs_steps，即展平后的多步观测历史。该模块的关键设计在于支持三种不同的对抗损失范式，每种范式对判别器输出到样式奖励的映射关系有着不同的数学定义。

GAN 范式采用标准二分类逻辑，判别器输出经 Sigmoid 转换为智能体轨迹被判别为"假"的概率，样式奖励定义为负对数似然的变形：rew = -log(max(1-prob, ε))。这种方法理论上收敛到 JS 散度，但在实际机器人训练中容易出现梯度饱和问题。LSGAN 范式将判别问题转化为最小二乘回归，要求智能体样本的输出逼近 -1、示范样本逼近 +1，奖励通过钳位平方误差计算：rew = clamp(1 - 0.25 * (disc_score - 1)^2, min=0)。这是本项目 ATOM01 AMP 配置中的默认选择，因其训练稳定性在足式机器人模仿任务中表现更优。WGAN 范式则直接以判别器输出的均值差异近似 Wasserstein 距离，需配合判别器输出的经验归一化（disc_output_normalizer）使用，奖励即为归一化后的原始判别分数。Sources: amp.py, amp.py, atom01_amp_agent_cfg.py

三种范式共享同一套梯度惩罚机制。compute_grad_penalty 对示范数据输入计算判别器输出的梯度范数，并以 (grad_norm - 0)^2 * scale 的形式施加惩罚，默认 scale=10.0。这一设计源于 WGAN-GP 的思想，但在 GAN 和 LSGAN 模式下同样起到平滑判别器决策边界、防止过度自信的作用。Sources: amp.py, ppo_amp.py

PPOAMP：联合训练的数据流与控制流

PPOAMP 的训练流程可以拆解为环境交互阶段和参数更新阶段两个部分，每个阶段都体现了 AMP 对标准 PPO 的侵入式扩展。

在环境交互阶段，process_env_step 方法首先调用判别器的 get_disc_obs 和 get_disc_demo_obs，从当前环境观测中提取用于判别器的特征张量，形状均为 [num_envs, disc_obs_steps, disc_obs_dim]。随后通过 predict_style_reward 计算样式奖励 style_rewards，该方法在 torch.no_grad() 上下文中执行前向传播，确保样式奖励不影响策略网络的历史梯度。核心融合逻辑由 lerp_reward 实现：rewards_lerp = task_style_lerp * task_reward + (1.0 - task_style_lerp) * style_reward。当 task_style_lerp=0.0 时，策略完全由样式奖励驱动；当 task_style_lerp=1.0 时，AMP 退化为纯 PPO。提取的 disc_obs 和 disc_demo_obs 分别追加到两个 CircularBuffer 实例中，而融合后的 rewards_lerp 则传入父类 PPO.process_env_step，存入标准的 RolloutStorage。Sources: ppo_amp.py, amp.py

在参数更新阶段，update 方法同时从三个数据源采样：标准 PPO 的 RolloutStorage、智能体的 disc_obs_buffer 和示范的 disc_demo_obs_buffer。三个生成器通过 zip 同步迭代，确保每个 mini-batch 内的策略样本与判别器样本一一对应。判别器的损失计算严格区分离三种范式：GAN 使用 BCEWithLogitsLoss 将智能体标签设为 0、示范标签设为 1；LSGAN 使用 MSELoss 将目标值分别设为 -1 和 +1；WGAN 则直接最大化示范与智能体输出的均值差。判别器总损失为 disc_loss + grad_penalty。梯度回传阶段，PPO 损失、RND 损失（如启用）和判别器损失分别通过独立的优化器执行 backward() 和 step()：策略参数由 self.optimizer 更新，判别器参数由 self.disc_optimizer 更新。这种解耦优化器的设计至关重要，因为策略和判别器的学习率往往相差两个数量级（例如策略 1e-4，判别器 1e-5），且梯度裁剪阈值也各自独立控制。Sources: ppo_amp.py

sequenceDiagram
    participant Env as VecEnv
    participant Alg as PPOAMP
    participant Disc as AMPDiscriminator
    participant BufA as disc_obs_buffer
    participant BufD as disc_demo_obs_buffer
    participant Roll as RolloutStorage

    loop num_steps_per_env
        Env->>Alg: obs, task_reward, dones, extras
        Alg->>Disc: get_disc_obs(obs)
        Alg->>Disc: get_disc_demo_obs(obs)
        Disc-->>Alg: style_reward (no grad)
        Alg->>Alg: lerp_reward(task, style)
        Alg->>BufA: append(disc_obs)
        Alg->>BufD: append(disc_demo_obs)
        Alg->>Roll: process_env_step(obs, lerp_reward, ...)
    end

    loop num_epochs x num_mini_batches
        Roll-->>Alg: PPO batch
        BufA-->>Alg: disc_obs batch
        BufD-->>Alg: disc_demo_obs batch
        Alg->>Disc: disc_score, disc_demo_score
        Alg->>Alg: PPO loss.backward()
        Alg->>Alg: disc_loss.backward()
        Alg->>Alg: optimizer.step()
        Alg->>Alg: disc_optimizer.step()
    end

CircularBuffer：面向时序特征的环形采样器

AMP 训练需要维护两个额外的经验缓冲区，但它们的语义与 PPO 的 RolloutStorage 有着本质区别。RolloutStorage 存储的是单步转移（状态-动作-奖励-值），用于计算广义优势估计（GAE）；而 AMP 的 CircularBuffer 存储的是多步判别器观测窗口，用于提供判别器训练的负样本（智能体轨迹）和正样本（示范轨迹）。

CircularBuffer 内部以 (max_len, batch_size, ...) 的环形张量布局存储数据，支持按环境实例独立重置。其 mini_batch_generator 方法实现了与 PPO 同步的采样节奏：在每个 epoch 中，从当前缓冲区内的所有有效数据中随机抽取 epoch_batch_size = batch_size * fetch_length 个样本，再均分为 num_mini_batches 份。这里 fetch_length 固定为 storage.num_transitions_per_env，即每个环境在一次 rollout 中产生的步数，确保判别器的 batch size 与策略更新保持一致。这种设计允许判别器在每个 PPO mini-batch 中都能见到 freshest 的智能体轨迹，避免传统经验回放中数据过时的问题。Sources: circular_buffer.py, amp_runner.py

AMP Runner 的生命周期管理

AMPRunner 对基类 OnPolicyRunner 的重写集中在四个维度：日志记录、模型持久化、模式切换和算法构造。

在日志维度，AMPRunner 使用 LoggerAMP 替代标准 Logger，额外追踪 style_rewards（样式奖励的原始值）和 total_rewards（融合后的任务+样式奖励）。每次环境步进时，process_env_step 会将这两项累加到 episode 级别的缓冲区，并在 episode 结束时向 TensorBoard/WandB 输出 AMP/mean_style_reward 和 AMP/mean_total_reward。Sources: amp_runner.py, amp_logger.py

在持久化维度，save 方法扩展了状态字典，必须包含 amp_discriminator_state_dict、amp_discriminator_normalizer_state_dict 和 amp_discriminator_optimizer_state_dict；对应的 load 方法则按相反顺序恢复这些参数。这意味着 AMP 模型的检查点体积显著大于纯 PPO 模型，且恢复训练时必须完整加载判别器及其归一化统计量，否则奖励分布会发生漂移。Sources: amp_runner.py

在模式切换维度，train_mode 和 eval_mode 除了管理策略网络和 RND 模块外，还必须显式切换判别器网络及其观测归一器（disc_obs_normalizer）的训练/评估状态。值得注意的是，predict_style_reward 内部会临时将判别器设为 eval()，并在完成后恢复原始模式，因此 runner 级别的模式切换主要影响归一器的统计量更新行为。Sources: amp_runner.py

在算法构造维度，_construct_algorithm 在基类逻辑之前插入了 resolve_amp_config 调用，自动解析 obs_groups 中 discriminator 相关组的形状，推断出 disc_obs_dim 和 disc_obs_steps，并注入环境的时间步长 step_dt 用于奖励缩放。随后初始化两个 CircularBuffer，最后将缓冲区和 AMP 配置一并传入 PPOAMP 构造函数。Sources: amp_runner.py

关键配置参数与调优建议

ATOM01 的 AMP 训练配置位于 robolab/tasks/manager_based/amp/agents/atom01_amp_agent_cfg.py，以下参数对训练动态有着最直接的影响：

特别需要关注的是 obs_groups 的定义。discriminator 和 discriminator_demonstration 必须包含相同维度、相同历史长度的观测组。环境开发者需要确保 disc_demo 组的数据来自运动重定向模块的真实参考轨迹，而非当前状态的简单复制，否则判别器将无法学到有意义的风格差异。Sources: atom01_amp_agent_cfg.py

与标准 PPO Runner 的差异速查

下表对比了 AMP Runner 与 On-Policy Runner 训练循环 在核心行为上的差异，便于已有 PPO 训练经验的开发者快速迁移：

总结与延伸阅读

AMP Runner 通过将对抗性运动先验嵌入标准 PPO 训练循环，实现了无需显式设计奖励函数的模仿学习。其核心工程挑战在于维持策略网络与判别器网络训练的动态平衡：判别器太强会导致样式奖励稀疏、策略探索困难；判别器太弱则无法提供有效的风格引导。实践中建议从 LSGAN 损失、适中的 style_reward_scale（如 1.0~2.0）以及较保守的 task_style_lerp（如 0.3）开始调优，逐步增大模仿权重。

若你尚未了解运动数据如何准备和重定向，建议先阅读 运动数据重定向与数据集准备；若希望深入理解判别器背后的对抗学习数学原理，可参考 AMP 对抗运动先验算法。完成 AMP 训练后，模型的部署流程与标准策略一致，详见 策略测试与 Sim2Sim 部署 和 MuJoCo Sim2Sim 部署与真机迁移。