本文档聚焦于 rsl_rl 中 PPO（Proximal Policy Optimization） 算法的工程实现与端到端训练流程。我们将从策略网络的数据交互出发，逐步拆解环境交互、优势估计、小批量策略更新等核心阶段，帮助你在掌握 PPO 数学原理的基础上，理解代码中的时序编排、损失计算与扩展机制。如需了解策略网络的具体架构，可参阅 Actor-Critic 基础架构设计；关于训练循环的更高层生命周期管理，请参阅 训练运行器生命周期管理。

整体架构与数据流

在 rsl_rl 中，PPO 并非孤立存在，而是由 OnPolicyRunner 进行生命周期编排、RolloutStorage 承载轨迹数据、ActorCritic 提供策略与价值输出。三者的协作关系可概括为：Runner 驱动循环，Storage 缓存经验，PPO 计算梯度。

flowchart TD
    A[OnPolicyRunner.learn] -->|获取观测| B[PPO.act]
    B -->|采样动作| C[VecEnv.step]
    C -->|obs, rewards, dones, extras| D[PPO.process_env_step]
    D -->|写入 transition| E[RolloutStorage]
    E -->|单轮 rollouts 结束| F[PPO.compute_returns]
    F -->|GAE 计算| E
    E -->|生成 mini-batch| G[PPO.update]
    G -->|Adam 更新参数| H[ActorCritic]
    H -->|下一轮| A

PPO 类本身不直接实例化网络，而是通过构造函数接收一个已构建的 ActorCritic（或其变体）实例与 RolloutStorage 实例。这种依赖注入的设计使得策略架构与算法解耦，你可以在不变更 PPO 核心逻辑的前提下，替换为循环网络或 CNN 编码器。OnPolicyRunner 在 _construct_algorithm 中完成这一装配过程，并负责将环境观测的 TensorDict 结构解析为算法所需的观测分组。

Sources: ppo.py, on_policy_runner.py

环境交互：动作采样与过渡记录

每一轮策略更新的起点是数据收集。PPO.act 接收当前观测 TensorDict，调用策略网络的 act 方法采样动作，同时缓存价值估计、动作对数概率、分布均值与标准差等中间量到 self.transition。如果策略是循环网络，还会额外保存隐藏状态。这些字段被组织在 RolloutStorage.Transition 容器中，作为单步经验的原子单元。

当环境执行 step 返回新的观测、奖励和终止信号后，PPO.process_env_step 负责完成三项工作：一是更新观测归一化器的运行统计量；二是处理超时引导（timeout bootstrapping）——若 extras 中包含 time_outs，则将被截断回合的价值估计按折扣率加回到奖励中，避免固定长度回合过早截断带来的价值估计偏差；三是将填充好的 transition 写入 RolloutStorage，并清空过渡容器以备下一步使用。若启用了 RND（Random Network Distillation），本阶段还会计算内在奖励并叠加到外在奖励上。

Sources: ppo.py, rollout_storage.py, vec_env.py

回报估计：GAE-Lambda 的后向传播

当 num_steps_per_env 步的轨迹收集完毕后，PPO.compute_returns 采用 广义优势估计（GAE-Lambda） 计算回报与优势。该方法从轨迹末端向前回溯：对于每一步，先计算 TD 残差 delta = r_t + gamma * V(s_{t+1}) * (1 - done) - V(s_t)，再通过指数加权移动平均累加优势 advantage = delta + gamma * lambda * (1 - done) * advantage。最终的回报值为 returns = advantages + values。

graph LR
    A[last_values] -->|bootstrap| B[step N-1]
    B -->|delta, advantage| C[step N-2]
    C -->|...| D[step 0]
    D -->|advantages + values| E[returns]

完成 GAE 后，若未启用按 mini-batch 的局部优势归一化，则会对全部优势做全局标准化 (advantages - mean) / (std + 1e-8)。这种归一化有助于稳定策略梯度，尤其在奖励尺度变化剧烈的任务中效果显著。

Sources: ppo.py

策略更新： clipped surrogate 与多损失组合

PPO.update 是算法实现的核心，它通过多次遍历存储的经验数据来优化策略。更新阶段首先根据策略类型选择小批量生成器：对于前馈网络使用 mini_batch_generator，将时间步与环境数两个维度展平后随机打乱索引，再按固定 mini-batch 大小切分；对于循环网络则使用 recurrent_mini_batch_generator，按回合边界对轨迹做填充与掩码处理，以保证隐藏状态在有效序列上传播。随后进入双重循环：num_learning_epochs 控制外层 epoch 数，num_mini_batches 控制每轮切分数。

在每个 mini-batch 上，算法依次执行以下计算：

1. 优势归一化（可选）：若开启 normalize_advantage_per_mini_batch，则仅对当前 batch 内的优势做标准化。
2. 对称增强（可选）：若开启对称性数据增强，会通过外部函数对观测与动作做镜像扩充，batch 规模相应倍增。
3. 重新评估当前策略：对 batch 中的观测再次前向传播，得到新的动作对数概率 actions_log_prob、价值估计 value 和分布参数 mu, sigma。这一步是 PPO 的“on-policy”约束体现——虽然数据来自旧策略，但目标函数使用当前策略的概率比进行评估。
4. 自适应 KL 学习率调度：若配置为 adaptive，会根据新旧策略间的 KL 散度调整学习率。代码中采用解析公式直接计算高斯分布的 KL，而非蒙特卡洛采样估计。当 KL 超过 desired_kl * 2.0 时降低学习率，低于 desired_kl / 2.0 且大于 0 时提高学习率。
5. Surrogate Loss：计算概率比 ratio = exp(log_prob_new - log_prob_old)，并取 clipped 与 unclipped 目标的最大值： loss = max(-A * ratio, -A * clamp(ratio, 1-clip, 1+clip))。符号取负是因为 PyTorch 默认执行梯度下降，而 PPO 目标是最大化。
6. Value Loss：支持两种模式。若启用 use_clipped_value_loss，则对价值更新量做裁剪 value_clipped = target_values + clamp(value - target, -clip, +clip)，并取 (value - returns)^2 与 (value_clipped - returns)^2 的较大者；否则使用标准的均方误差。
7. 总损失合成： loss = surrogate + value_loss_coef * value_loss - entropy_coef * entropy 若启用辅助损失（aux loss）或对称镜像损失（mirror loss），也会以对应系数叠加到总损失中。

梯度计算完成后，调用 clip_grad_norm_ 进行梯度裁剪，随后通过 Adam 更新策略参数。若启用了 RND，则同步更新 RND 预测器的参数。

Sources: ppo.py, rollout_storage.py, actor_critic.py

训练编排：OnPolicyRunner 的主循环

虽然 PPO 类实现了算法的数学核心，但训练的节奏控制由 OnPolicyRunner.learn 负责。该方法遵循严格的“采样-计算-更新-日志”四阶段节拍：

1. 初始化：将策略设为训练模式，获取初始观测，并在分布式场景下广播参数以确保各 GPU 起点一致。
2. Rollout 阶段：在 torch.inference_mode() 上下文中循环 num_steps_per_env 次，逐与环境交互并调用 alg.act 与 alg.process_env_step。此阶段不记录计算图，显著降低显存占用。
3. 回报计算：调用 alg.compute_returns 完成 GAE。
4. 策略更新：退出 inference 模式，调用 alg.update 执行多 epoch 优化。
5. 日志与保存：记录收集耗时、学习耗时、各项损失及当前学习率；按 save_interval 周期保存模型检查点。

这种编排方式确保了 on-policy 的严格性——每次策略更新后立即丢弃旧数据，下一轮 rollout 必须使用更新后的策略重新采样。对于希望深入理解运行器状态转换的读者，建议继续阅读 训练运行器生命周期管理。

Sources: on_policy_runner.py

分布式训练与梯度同步

当环境变量 WORLD_SIZE > 1 时，OnPolicyRunner 会自动进入分布式模式，并将配置通过 multi_gpu_cfg 传递给 PPO。PPO.broadcast_parameters 在训练启动时由 Runner 调用，将 rank 0 上的模型状态广播至所有进程；PPO.reduce_parameters 则在每次 update 的反向传播后被调用，使用 torch.distributed.all_reduce 收集并平均所有 GPU 上的梯度，再写回各参数。这种数据并行策略在保持逻辑简单的同时，有效扩展了 batch size。

Sources: ppo.py, on_policy_runner.py

核心参数速查

Sources: ppo.py

后续阅读建议

掌握 PPO 的训练流程后，你可以根据兴趣选择以下方向深入：

• 若希望了解策略网络的具体 MLP 构造、观测归一化与动作分布参数化，请阅读 Actor-Critic 基础架构设计。
• 若关注循环网络与注意力编码器在 PPO 中的接入方式，请阅读 循环与注意力策略变体。
• 若对 rollout 数据的内存布局、Transition 结构与小批量生成器的内部实现感兴趣，请阅读 Rollout 数据存储与 Transition 与 经验采样与小批量生成。
• 若需要引入对抗动作先验（AMP）进行模仿学习，请阅读 AMP 对抗动作先验算法。
• 若希望利用 RND 提升探索效率，请阅读 RND 随机网络蒸馏探索。
• 若需要利用对称性先验约束策略，请阅读 对称性增强与镜像损失。