AMP（Adversarial Motion Priors）是一种将生成对抗网络与近端策略优化相结合的模仿学习框架，其核心目标是在不依赖显式时间对齐的前提下，让强化学习策略自动习得参考运动数据中的自然运动风格。传统行为克隆要求参考轨迹与策略 rollout 在时序上严格对齐，而 AMP 通过判别器学习运动片段的分布特征，将“风格相似度”转化为可优化的标量奖励，从而支持异步、变长度的运动模仿。在本项目中，AMP 模块以 PPOAMP 算法类为中心，配合 AMPDiscriminator 判别器、AMPRunner 训练循环以及环境端的观察组配置，构成一条完整的对抗运动先验训练管线。

Sources: ppo_amp.py, amp.py

算法原理：对抗训练与策略优化的共生

AMP 的理论基础源于对抗性模仿学习。系统同时维护两个互相对抗的优化目标：策略网络（Actor-Critic）试图最大化由任务奖励和风格奖励共同构成的复合回报；判别器则试图准确区分策略生成的运动状态与来自运动捕捉数据集的参考状态。这种对抗关系被嵌入到标准的 PPO 更新循环中，每一轮迭代都包含策略梯度更新和判别器梯度更新两个独立步骤。与传统 GAIL 不同，AMP 的判别器输入不是完整状态-动作对，而是经过设计的“运动特征观察”（motion-specific observations），这降低了对策略分布的过度敏感性，同时允许环境保留独立的任务奖励结构。本项目进一步引入了 task_style_lerp 参数，支持在线性插值空间内连续调节任务奖励与风格奖励的权重，使得开发者可以在“严格模仿”与“任务优先”之间平滑过渡。

Sources: ppo_amp.py, amp.py

系统架构与数据流

AMP 系统的数据流可概括为四层：环境层负责提供并行的四类观察（policy、critic、discriminator、discriminator_demonstration）；Runner 层通过 AMPRunner 调度 rollout 与日志；算法层在 PPOAMP 中完成奖励混合与损失聚合；网络层则包含策略网络与判别器网络各自的独立优化器。以下架构图展示了从环境步进到判别器更新的完整闭环。

graph TD
    Env[环境 AnimationEnv<br/>输出 policy / critic / disc / disc_demo] -->|obs| Runner[AMPRunner]
    Runner -->|actions| Env
    Runner -->|learn| Alg[PPOAMP Algorithm]
    Alg -->|act / evaluate| Policy[Actor-Critic Policy]
    Policy -->|actions| Runner
    Alg -->|get_disc_obs / predict_style_reward| Disc[AMPDiscriminator]
    Disc -->|style_rewards<br/>rewards_lerp| Alg
    Dataset[运动捕捉数据集<br/>Motion Dataset] -->|参考运动| Env
    Alg -->|lerp reward| Storage[RolloutStorage]
    Buf1[CircularBuffer<br/>disc_obs_buffer] -->|agent states| Disc
    Buf2[CircularBuffer<br/>disc_demo_obs_buffer] -->|demo states| Disc
    Disc -->|disc_loss + grad_penalty| DiscOpt[Disc Optimizer<br/>Adam]
    Policy -->|surrogate + value + entropy| PolicyOpt[Policy Optimizer<br/>Adam]

环境每一步都会输出 disc 和 disc_demo 两组观察，前者描述智能体当前的运动学状态，后者描述参考动画在相同时刻的对应状态。PPOAMP.process_env_step 调用判别器将这两组观察转化为风格奖励，再通过 lerp_reward 与原始任务奖励混合，最终存入 RolloutStorage。在更新阶段，两组观察分别通过各自的 CircularBuffer 采样，供判别器计算对抗损失。

Sources: amp_runner.py, ppo_amp.py

核心组件详解

PPOAMP：融合 AMP 的 PPO 算法

PPOAMP 继承自标准 PPO 算法类，在保留所有原生 PPO 功能（包括自适应学习率、RND 好奇心奖励、对称性增强）的基础上，扩展了判别器的初始化、风格奖励计算和联合优化逻辑。构造函数接收两个额外的 CircularBuffer 实例——disc_obs_buffer 与 disc_demo_obs_buffer——分别缓存智能体轨迹与演示轨迹的判别器观察。算法通过 amp_cfg 字典配置判别器的网络维度、损失类型、学习率及梯度惩罚系数。

在 process_env_step 中，算法首先调用 AMPDiscriminator.get_disc_obs 与 get_disc_demo_obs 从环境观察字典中提取对应观察组，随后通过 predict_style_reward 计算风格奖励，并使用 lerp_reward 进行线性插值混合。混合后的奖励被传递给父类的 process_env_step，从而无缝融入 PPO 的经验回放。在 update 方法中，PPO 损失与判别器损失分别计算、分别反向传播，但共享同一轮 mini-batch 循环。判别器使用独立的 Adam 优化器，其参数分为 disc_trunk 和 disc_linear 两组，分别施加不同的 L2 权重衰减。

Sources: ppo_amp.py, ppo_amp.py

AMPDiscriminator：运动风格判别器

AMPDiscriminator 是一个标准的多层感知机，输入维度为 disc_obs_dim × disc_obs_steps，即把历史观察在时间维度上展平后送入网络。主干网络 disc_trunk 由若干全连接层与激活函数构成，最终通过 disc_linear 输出一个标量分数。该模块支持三种对抗损失模式：标准 GAN、最小二乘 GAN（LSGAN）和 Wasserstein GAN（WGAN）。模块内部维护一个 EmpiricalNormalization 观察归一器，用于稳定判别器的输入分布；当使用 WGAN 时，还会额外维护一个输出归一器。

判别器的核心接口包括：

• get_disc_obs / get_disc_demo_obs：从环境观察字典中按 obs_groups 提取对应观察组，并做维度校验。
• normalize_disc_obs：将三维观察 (num_envs, steps, dim) reshape 后送入归一器，再恢复原形状。
• compute_grad_penalty：对演示数据计算 R1 风格的梯度惩罚，即要求演示样本处的梯度范数趋近于 0，从而增强判别器的平滑性。
• predict_style_reward：根据当前损失类型将判别分数映射为非负的风格奖励，并乘以 style_reward_scale 和步长 dt。

classDiagram
    class AMPDiscriminator {
        +nn.Sequential disc_trunk
        +nn.Linear disc_linear
        +EmpiricalNormalization disc_obs_normalizer
        +nn.Identity/EmpiricalNormalization disc_output_normalizer
        +LossType loss_type
        +forward(x) Tensor
        +get_disc_obs(obs, flatten) Tensor
        +get_disc_demo_obs(obs, flatten) Tensor
        +normalize_disc_obs(disc_obs) Tensor
        +compute_grad_penalty(demo_data, scale) Tensor
        +predict_style_reward(disc_obs, dt) Tuple[Tensor, Tensor]
        +lerp_reward(task_reward, style_reward) Tensor
    }

Sources: amp.py, amp.py, amp.py

AMPRunner：AMP 训练循环管理器

AMPRunner 继承自 OnPolicyRunner，负责实例化 AMP 所需的额外存储与日志系统。在 _construct_algorithm 中，Runner 首先调用 resolve_amp_config 自动推断判别器观察的维度与历史长度，随后初始化两个 CircularBuffer 作为判别器的经验缓存，最后构造 PPOAMP 实例。learn 方法在标准 rollout 循环的基础上，将 style_rewards 与 total_rewards 提取出来传递给 LoggerAMP，以便在 TensorBoard 或 WandB 中监控风格奖励的演化趋势。模型的保存与加载逻辑也被扩展为包含判别器状态、判别器归一器状态以及判别器优化器状态，确保训练可以完整恢复。

Sources: amp_runner.py, amp_runner.py, amp_runner.py

CircularBuffer：判别器观察的循环缓存

CircularBuffer 是一个专为多环境并行设计的循环缓冲区，形状为 (max_len, batch_size, ...)，支持按环境 ID 独立重置。与普通 rollout storage 不同，判别器缓冲区需要在多个训练 epoch 内持续采样，因此其 mini_batch_generator 采用随机索引组合策略：先从当前有效长度与 batch size 的笛卡尔积中随机选取 epoch_batch_size 个线性索引，再映射为 buffer 的维度索引，从而实现跨时间与跨环境的均匀采样。该设计保证即使 buffer 长度远大于 num_steps_per_env，判别器仍能在每个 PPO epoch 中获得充足且多样的训练样本。

Sources: circular_buffer.py, circular_buffer.py

环境侧配置：观察组与运动数据

AMP 的训练高度依赖环境端对观察组的精确划分。在 obs_groups 配置中，必须定义 discriminator 和 discriminator_demonstration 两个键，分别映射到环境观察字典中的对应字段。以 ATOM01 AMP 任务为例，disc 观察组包含智能体的基础角速度、关节位置和关节速度，而 disc_demo 观察组包含参考动画的根角速度、关节位置和关节速度。两者在历史长度和拼接维度上保持一致，确保判别器输入维度完全匹配。

graph LR
    subgraph 环境观察输出
        A[policy] -->|用于 Actor-Critic| B[策略网络]
        C[critic] -->|用于 Value| D[价值网络]
        E[disc] -->|智能体运动特征| F[AMPDiscriminator]
        G[disc_demo] -->|参考运动特征| F
    end

在 AmpEnvCfg 中，DiscriminatorCfg 与 DiscriminatorDemoCfg 均继承自 ObservationGroupCfg，支持通过 history_length 指定时序窗口长度。环境在每一步会自动将 flatten_steps_dim=False 的观察保留为三维张量 (num_envs, history_length, obs_dim)，这与判别器期望的输入格式完全一致。resolve_amp_config 函数在训练启动时会遍历这些观察组，自动累加各 term 的维度并校验历史长度一致性，从而避免手动计算输入维度的繁琐与出错风险。

Sources: amp_env_cfg.py, atom01_amp_agent_cfg.py, amp.py

训练流程与损失计算

AMP 的训练更新可分为三个并行计算的损失分支：策略损失、判别器损失以及可选的辅助损失（RND、对称性）。以下流程图展示了 PPOAMP.update 的核心逻辑。

flowchart TD
    Start[开始 Update] --> Sample1[从 RolloutStorage<br/>采样 PPO batch]
    Start --> Sample2[从 CircularBuffer<br/>采样 disc & disc_demo]
    Sample1 --> PolicyForward[Policy 前向<br/>计算 log_prob / value / entropy]
    Sample2 --> DiscForward[Discriminator 前向<br/>计算 disc_score & disc_demo_score]
    PolicyForward --> PPO Loss[Surrogate + Value + Entropy]
    DiscForward --> LossType{LossType}
    LossType -- GAN --> BCE[BCEWithLogitsLoss]
    LossType -- LSGAN --> MSE[MSELoss]
    LossType -- WGAN --> Wass[-mean(demo) + mean(agent)]
    BCE --> GP[梯度惩罚<br/>compute_grad_penalty]
    MSE --> GP
    Wass --> GP
    PPO Loss --> Back1[Policy Optimizer<br/>zero_grad / backward / step]
    GP --> Back2[Disc Optimizer<br/>zero_grad / backward / step]
    Back1 --> End[返回 loss_dict]
    Back2 --> End

判别器损失的具体形式取决于 loss_type 配置。在 GAN 模式下，智能体样本被标注为 0、演示样本被标注为 1，使用二元交叉熵；风格奖励通过 -log(1 - sigmoid(score)) 计算，物理意义为“被误判为演示的惊讶度”。LSGAN 将二分类问题转化为回归问题，目标值分别为 -1 和 +1，风格奖励为 clamp(1 - 0.25*(score-1)^2, min=0)，其梯度在训练初期更稳定，因此作为本项目的默认配置。WGAN 则直接使用分数差值作为近似 Wasserstein 距离，风格奖励为经过在线归一化的原始分数，适用于需要精细控制奖励尺度的场景。无论采用哪种模式，判别器总损失都会加上一个对演示数据计算的梯度惩罚项，其数学形式为 scale * (||∇D(x)||_2 - 0)^2，强制判别器在真实数据附近保持局部线性，有效抑制模式崩溃。

Sources: ppo_amp.py, amp.py, amp.py

关键超参数与调优指南

AMP 引入的超参数主要集中在判别器结构、优化器设置以及奖励混合策略三个方面。下表汇总了核心参数的工程语义与推荐配置范围。

判别器与策略的学习率配比是一个关键调参维度。若判别器学习率过高，它会快速过拟合到当前策略分布，导致风格奖励信号稀疏，策略难以获得有效梯度；反之，若判别器过弱，则风格奖励趋于随机，失去引导意义。通常建议将判别器学习率设为策略学习率的十分之一到五分之一，并监控 disc_score 与 disc_demo_score 的均值差异，理想情况下两者应维持在 0.5~2.0 的合理间隔内。

Sources: atom01_amp_agent_cfg.py, ppo_amp.py

监控指标与调试

LoggerAMP 在标准 PPO 日志的基础上，额外追踪三类 AMP 专属指标，并写入 TensorBoard 的 AMP/ 命名空间下：

• AMP/mean_style_reward：每个 episode 内风格奖励的时间均值，直接反映策略与参考运动的风格接近程度。
• AMP/mean_total_reward：混合奖励（rewards_lerp）的 episode 均值，用于评估任务与风格的综合表现。
• Loss/amp/disc_score 与 Loss/amp/disc_demo_score：判别器对智能体样本和演示样本的打分均值。两者差距过大意味着判别器过度自信，可能需要降低学习率或增大梯度惩罚；差距过小则意味着判别器无法区分，可能需要增加网络容量或降低正则。

在控制台输出中，Mean amp/disc_loss 和 Mean amp/disc_grad_penalty 可以帮助快速判断判别器的训练健康度。若 disc_grad_penalty 持续高于 disc_loss 一个数量级，说明梯度惩罚项主导了优化方向，应适当下调 grad_penalty_scale。此外，Episode_Reward/style 以每秒均值的形式被附加到 episode extras 中，便于与任务奖励曲线进行归一化对比。

Sources: amp_logger.py, amp_logger.py

进阶话题

梯度惩罚的数学动机

本项目采用的梯度惩罚并非 WGAN-GP 中常见的插值样本梯度惩罚，而是针对演示数据单独施加的 R1 正则化。其理论基础在于：在真实数据分布的支撑集上，最优判别器的梯度应为零（当生成器完全匹配真实分布时）。通过在演示样本上最小化梯度范数的平方，判别器被约束为在真实数据附近局部平滑的函数，这有效降低了策略因判别器突变而接收到的噪声奖励信号。该设计对 LSGAN 尤为关键，因为最小二乘损失的判别器在远离目标值时会产生二次增长的梯度，R1 正则化可抑制这种发散趋势。

Sources: amp.py

对称增强与 AMP 的协同

PPOAMP 完全兼容对称性增强模块。在 update 中，若启用了对称数据增强，观察与动作批次会先经过镜像变换扩充，随后 PPO 损失在扩充后的批次上计算；而判别器损失仅在原始批次上计算，因为判别器关注的是物理运动特征而非策略分布的对称性。这种解耦设计保证了对称增强只影响策略优化，不干扰对抗训练的分布匹配。开发者可以在 AMP 训练中同时启用 symmetry_cfg 和 rnd_cfg，实现风格模仿、探索激励与物理对称约束的三重目标优化。

Sources: ppo_amp.py, ppo_amp.py

关联阅读

掌握 AMP 算法后，建议继续阅读以下页面以理解完整的训练管线与部署链路：

• 若需深入了解 PPO 的基础 clipped surrogate loss、GAE 以及自适应学习率调度机制，请参阅 PPO 核心算法与超参数。
• 若需配置 ATOM01 的关节参数、传感器设置及运动数据重定向流程，请参阅 ATOM01 机器人模型配置 与 运动数据重定向与数据集准备。
• 若需了解 AMP 训练的完整启动流程、命令行参数及 checkpoint 管理，请参阅 AMP Runner 与模仿学习训练。
• 若计划将训练好的 AMP 策略迁移到 MuJoCo 或真机，请参阅 MuJoCo Sim2Sim 部署与真机迁移。