人形机器人“天工”AMP-PPO算法源码解析：从PPO基础到对抗模仿学习

上一篇文章介绍了“天工”机器人开源代码 train.py 的运行流程，本篇我们将深入剖析其核心训练算法——AMP-PPO。

PPO 基本概念

强化学习非常适合训练机器人，其模式是：给定当前状态，策略网络决定执行什么动作，与环境交互后获得奖励。在众多算法中，PPO因其简单高效而广受欢迎，但它并非凭空出现，其设计思想源于两位前辈：REINFORCE 和 TRPO。

最朴素的想法是初始化一个随机策略，让机器人尝试直到失败，然后更新策略，如此循环。这就是 REINFORCE 算法。但这种方案数据利用率极低，且更新幅度难以控制，训练过程很不稳定。

TRPO 旨在提高数据利用率同时保证效果不下降。它引入了信任区域（Trust Region）的概念，限制新旧策略之间的 KL 散度变化范围。为了重复利用旧策略收集的数据来计算新策略的期望回报，TRPO 在数学上使用了重要性采样。

本质上，TRPO 是一个带约束的二阶优化问题。通常有两种求解思路：一是有限差分近似法，通过泰勒展开进行近似计算，实现简单但稳定性较差；二是共轭梯度法，需要计算精确的海森向量积。尽管现代框架如 PyTorch 支持二阶自动微分，但计算海森矩阵依然会带来巨大的计算和内存开销。

PPO 的巧妙之处在于，它用一个裁剪函数替代了复杂的 KL 散度约束，将约束条件直接写入了目标函数中，从而又能使用简单快速的一阶梯度下降进行优化了！虽然单次更新的“质量”可能不及 TRPO 精准，但其极快的速度足以用“数量”弥补“质量”的差距，整体效率更高。

了解了算法的演进思想，我们接下来聚焦代码实现。首先从不带 AMP 的经典 PPO 部分开始。

PPO

1. 概率比 (Probability Ratio)

这个比率用于比较新旧策略对同一动作的评估概率。比值 >1 表示新策略认为该动作更好，<1 则表示旧策略更优。代码实现如下：

ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch))

这里使用对数概率相减再取指数，等价于概率相除。这样做的好处是能避免原始概率值过小导致计算中的数值下溢。

2. Clipped Surrogate Objective (裁剪替代目标)

这是 PPO 的核心。优势函数 A_t 告诉我们“这个动作是好是坏”，概率比 r_t 告诉我们“策略变化了多少”。两者相乘给出了策略期望的更新方向。min 和 clip 操作将策略的变化幅度限制在一定范围内，确保了训练的稳定性。这部分代码对应 Actor 网络，目的是使其产生更优的策略。

# 未裁剪的 surrogate: r_t * A_t
surrogate = -torch.squeeze(advantages_batch) * ratio

# 裁剪后的 surrogate: clip(r_t, 1-ε, 1+ε) * A_t
surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp(
    ratio, 1.0 - self.clip_param, 1.0 + self.clip_param
)

# 取 max（因为论文是最大化目标函数，代码是最小化损失，所以取max）
surrogate_loss = torch.max(surrogate, surrogate_clipped).mean()

请注意一个关键细节：原论文是最大化目标函数，而代码实现通常是最小化损失函数。因此：

• 论文中的 min 操作，在代码中对应 max。
• 代码中在 advantage 前加了负号。

参数对应关系：

• self.clip_param = 0.2，与 PPO 论文中的设定保持一致。

3. Value Function Loss (价值函数损失)

这部分代码对应 Critic 网络，目标是让价值网络的评估更准确。 论文中只提到了均方误差（MSE）的实现方式，而代码中为了增强稳定性，还提供了 Clipped Value Loss 方法，其思路与前述的 Surrogate Loss 裁剪类似。

if self.use_clipped_value_loss:
    # Clipped value loss（额外的稳定性措施）
    value_clipped = target_values_batch + (value_batch - target_values_batch).clamp(
        -self.clip_param, self.clip_param
    )
    value_losses = (value_batch - returns_batch).pow(2)
    value_losses_clipped = (value_clipped - returns_batch).pow(2)
    value_loss = torch.max(value_losses, value_losses_clipped).mean()
else:
    # 标准 squared error loss
    value_loss = (returns_batch - value_batch).pow(2).mean()

4. 综合目标函数

最后，加上鼓励探索的熵奖励（熵值越大，策略随机性越高），就构成了完整的损失函数。

loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean()

5. Adaptive KL Penalty (自适应学习率)

这里实现的自适应 KL 惩罚机制与 PPO 原始论文中提到的并非同一回事，但利用了相似的思想：通过监控策略的 KL 散度变化来动态调整学习率。

if self.desired_kl is not None and self.schedule == "adaptive":
    with torch.inference_mode():
        # 计算 KL 散度（假设动作服从高斯分布）
        kl = torch.sum(
            torch.log(sigma_batch / old_sigma_batch + 1.0e-5)
            + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch))
            / (2.0 * torch.square(sigma_batch))
            - 0.5,
            axis=-1,
        )
        kl_mean = torch.mean(kl)

    if kl_mean > self.desired_kl * 2.0:
        self.learning_rate = max(1e-5, self.learning_rate / 1.5)
    elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
        self.learning_rate = min(1e-2, self.learning_rate * 1.5)

6. GAE (广义优势估计)

GAE 是一种平衡偏差与方差的优势估计技巧。

• 当 λ = 0 时，退化为 1-step TD，偏差较大（依赖 Critic 的准确性），但方差小。
• 当 λ = 1 时，等价于蒙特卡洛采样，无偏差（使用真实奖励），但方差极大（受环境随机性影响严重）。

def compute_returns(self, last_values, gamma, lam, normalize_advantage: bool = True):
    advantage = 0
    for step in reversed(range(self.num_transitions_per_env)):
        # if we are at the last step, bootstrap the return value
        if step == self.num_transitions_per_env - 1:
            next_values = last_values
        else:
            next_values = self.values[step + 1]
        # 1 if we are not in a terminal state, 0 otherwise
        next_is_not_terminal = 1.0 - self.dones[step].float()
        # TD error: r_t + gamma * V(s_{t+1}) - V(s_t)
        delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step]
        # Advantage: A(s_t, a_t) = delta_t + gamma * lambda * A(s_{t+1}, a_{t+1})
        advantage = delta + next_is_not_terminal * gamma * lam * advantage
        # Return: R_t = A(s_t, a_t) + V(s_t)
        self.returns[step] = advantage + self.values[step]

    # Compute the advantages
    self.advantages = self.returns - self.values
    # Normalize the advantages if flag is set
    if normalize_advantage:
        self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8)

因为 t 时刻的优势估计 A_t 依赖于之后所有时刻的优势估计，所以需要从后向前倒序计算。

看完了通用的 PPO 部分，我们再来看看为了让人形机器人跑得更“像人”而引入的特有算法。

AMP (对抗性运动先验)

AMP 的思想来源于论文《AMP: Adversarial Motion Priors for Stylized Physics-Based Character Control》。它是一种对抗性模仿学习，即通过一个判别器（Discriminator）来缩小智能体策略产生的动作与专家演示动作之间的分布差异。你可以把它想象成生成对抗网络（GAN）在机器人运动控制上的一个变体：GAN 模仿的是图像分布，而 AMP 模仿的是状态转移分布。

为什么只关注状态转移，而不直接模仿动作？因为对于学习“像人一样跑步”这个姿态任务，我们更关心最终呈现的“状态”结果（如关节角度、身体姿态），而非产生这个结果的具体肌肉控制信号（动作）。

1. 风格奖励

在 predict_amp_reward 方法中，奖励的计算公式如下：

# d 是判别器的输出 D(s, s')
d = self.amp_linear(self.trunk(torch.cat([state, next_state], dim=-1)))

# 核心公式：1 - 0.25 * (d - 1)^2，并与 0 取 max (clamp)
reward = self.amp_reward_coef * torch.clamp(1 - (1 / 4) * torch.square(d - 1), min=0)

• self.amp_reward_coef 对应权重系数 w，默认配置为 0.3。
• torch.clamp(..., min=0) 对应公式中的 [·]+（取正值）。
• 这个设计确保了当判别器输出 d ≈ 1（认为动作很真实）时，奖励接近最大值；当 d 远离 1 时，奖励会迅速衰减至 0。

2. 梯度惩罚

梯度惩罚用于约束判别器，使其决策函数更加平滑，防止训练过程中出现梯度爆炸或模式崩溃。

公式对应：λ · (||∇D(x)||₂ - 0)²。这是一种零中心梯度惩罚，强制判别器在专家数据附近的梯度范数趋向于 0。

代码实现：在 compute_grad_pen 方法中：

# 获取专家数据输入的梯度
grad = autograd.grad(
    outputs=disc, inputs=expert_data, grad_outputs=ones, ...
)[0]

# 公式：lambda * ||grad||^2 的均值
# 这里 grad.norm(2, dim=1) - 0 表示目标梯度模长为 0
grad_pen = lambda_ * (grad.norm(2, dim=1) - 0).pow(2).mean()

• lambda_ 是惩罚系数。
• grad.norm(2, dim=1).pow(2) 计算了梯度向量的 L2 范数的平方。

3. 奖励融合

我们可以通过参数 task_reward_lerp 来平衡机器人是更倾向于“完成目标任务”（比如跑得快）还是更倾向于“动作规范美观”（跑得像人）。这是一种线性插值。

def _lerp_reward(self, disc_r, task_r):
    # r = (1 - alpha) * style_reward + alpha * task_reward
    r = (1.0 - self.task_reward_lerp) * disc_r + self.task_reward_lerp * task_r
    return r

4. AMP Loss

AMP 判别器本身的损失函数由三部分构成：

1. 专家部分损失：目标是让判别器给专家数据打高分（接近1）。
2. 策略部分损失：目标是让判别器给策略产生的数据打低分（接近-1）。这是一种“对抗”训练。
3. 梯度惩罚损失：如前所述，用于稳定训练。

policy_d = self.discriminator(torch.cat([policy_state, policy_next_state], dim=-1))
expert_d = self.discriminator(torch.cat([expert_state, expert_next_state], dim=-1))
expert_loss = torch.nn.MSELoss()(expert_d, torch.ones(expert_d.size(), device=self.device))
policy_loss = torch.nn.MSELoss()(policy_d, -1 * torch.ones(policy_d.size(), device=self.device))
amp_loss = 0.5 * (expert_loss + policy_loss)
grad_pen_loss = self.discriminator.compute_grad_pen(*sample_amp_expert, lambda_=10)
loss += self.amploss_coef * amp_loss + self.amploss_coef * grad_pen_loss

除了 AMP，项目中还应用了其他辅助算法来提升训练效果。

RND (随机网络蒸馏)

这个想法来自论文《Exploration by Random Network Distillation》。其核心是奖励智能体去探索那些它未曾见过的状态。

首先，随机初始化两个结构相同但参数不同的神经网络：Predictor（预测器）和 Target（目标网络）。其中 Target 网络被固定，不参与训练。

# Create network architecture
self.predictor = self._build_mlp(num_states, predictor_hidden_dims, num_outputs, activation).to(self.device)
self.target = self._build_mlp(num_states, target_hidden_dims, num_outputs, activation).to(self.device)

# make target network not trainable
self.target.eval()

然后，在训练过程中，让 Predictor 网络去拟合 Target 网络对同一状态的特征输出。

if self.rnd:
    # predict the embedding and the target
    predicted_embedding = self.rnd.predictor(rnd_state_batch)
    target_embedding = self.rnd.target(rnd_state_batch).detach()
    # compute the loss as the mean squared error
    mseloss = torch.nn.MSELoss()
    rnd_loss = mseloss(predicted_embedding, target_embedding)

关键在于，如果策略进入了一个新颖的、之前很少访问的状态，Predictor 网络由于缺乏“经验”，其预测误差（与固定 Target 的输出之差）会很大。反之，对于熟悉的状态，预测误差就小。因此，这个预测误差本身就可以作为一种内在奖励信号，在数据收集阶段叠加到环境给出的外在奖励上，激励智能体进行探索。

# Compute the intrinsic rewards and add to extrinsic rewards
if self.rnd:
    # Obtain curiosity gates / observations from infos
    rnd_state = infos["observations"]["rnd_state"]
    # Compute the intrinsic rewards
    self.intrinsic_rewards, rnd_state = self.rnd.get_intrinsic_reward(rnd_state)
    # Add intrinsic rewards to extrinsic rewards
    self.transition.rewards += self.intrinsic_rewards
    # Record the curiosity gates
    self.transition.rnd_state = rnd_state.clone()

Symmetry Loss (对称性损失)

这个想法来自论文《On Learning Symmetric Locomotion》。对于双足、四足等具有对称结构的机器人，其控制策略理论上也应该具有对称性。具体来说，如果策略学到在状态 s 下应该采取动作 a，那么在经过镜像变换的对称状态 s' 下，就应该采取对称的动作 a'。引入对称性损失可以约束策略学习这种特性，减少需要学习的模式，提高样本效率和泛化能力。

不过，在“天工”项目当前的代码仓库中，这部分功能暂时没有启用，相关配置为 None。

algorithm = RslRlPpoAlgorithmCfg(
    class_name="AMPPPO",
    value_loss_coef=1.0,
    use_clipped_value_loss=True,
    clip_param=0.2,
    entropy_coef=0.005,
    num_learning_epochs=5,
    num_mini_batches=4,
    learning_rate=1.0e-3,
    schedule="adaptive",
    gamma=0.99,
    lam=0.95,
    desired_kl=0.01,
    max_grad_norm=1.0,
    normalize_advantage_per_mini_batch=False,
    symmetry_cfg=None,  # RslRlSymmetryCfg()
    rnd_cfg=None,  # RslRlRndCfg()
)

在 Isaac Lab 仓库的 Issues 中，可以找到一个关于 G1 机器人的配置样例，其中启用了对称性损失，可供参考：https://github.com/jloganolson/g1_23dof_locomotion_isaac/blob/ee072967a7bea4cd6c7c89d64b76a421957245da/source/g1_23dof_locomotion_isaac/g1_23dof_locomotion_isaac/tasks/manager_based/g1_23dof_locomotion_isaac/agents/rsl_rl_ppo_cfg.py。

总结

至此，我们已将“天工”人形机器人项目中用到的几个核心算法——PPO、AMP、RND 和 Symmetry Loss 的原理与代码实现梳理了一遍。通过阅读代码，可以深刻体会到强化学习的数据流比监督学习复杂得多，它涉及策略交互、数据收集、优势估计、多目标损失计算等多个循环阶段。希望这篇对开源代码的解读，能帮助你更好地理解这些前沿算法是如何在具体项目中被实现和整合的。

如果你对机器人、强化学习或其他AI技术有更多兴趣，欢迎来到云栈社区与更多开发者交流讨论。在后续的文章中，我们将进一步分析这些算法模块是如何在完整的训练流程中协同工作的。