SAO: 面向智能体强化学习的单轨迹异步优化

论文: Single-Rollout Asynchronous Optimization for Agentic Reinforcement Learning
作者: Zhenyu Hou, Yujiang Li, Jie Tang, Yuxiao Dong (智源研究院/清华)
arXiv: 2607.07508

背景与挑战

在大语言模型（LLM）的后训练阶段，强化学习（RL）已成为提升模型推理能力和智能体能力的关键技术。当前主流的 RL 方法如 GRPO (Group Relative Policy Optimization) 采用同步批处理的训练范式：每个 prompt 采样 G 条轨迹，等待全部完成后计算组内相对优势进行更新。

然而，这种同步模式在长程智能体任务中面临严峻挑战：

1. 延迟驱动的 off-policy 问题：组内最慢的轨迹决定了整个批次的等待时间，导致其他轨迹的采样策略与当前训练策略严重偏离
2. 训练不稳定：GRPO 在约 160 步后常出现训练崩溃
3. 硬件利用率低：同步屏障导致 GPU 空闲等待

虽然异步 RL 系统能够提升吞吐量，但现有方案在训练稳定性和任务有效性方面探索不足。

SAO 核心方法

SAO (Single-rollout Asynchronous Optimization) 提出了一套面向智能体 RL 的异步优化框架，核心创新包括三个方面：

1. 单轨迹采样 (Single-Rollout Sampling)

SAO 摒弃了 GRPO 的组采样策略，改为每个 prompt 仅采样一条轨迹：

特性	GRPO	SAO
采样方式	组采样 (G samples/prompt)	单轨迹采样
优势估计	组内相对基线	Value model GAE
同步策略	等待全部组成员	单轨迹完成即训练

这种设计的直接好处是：轨迹一旦生成立即参与训练，无需等待最慢样本，从根本上消除了延迟导致的 off-policy 效应。

优势估计采用 GAE (Generalized Advantage Estimation)：

\[ \hat{A}_t^{\text{GAE}} = \sum_{l=0}^{|y|-t-1}(\gamma\lambda)^l \delta_{t+l} \]

其中 \(\delta_t = r_t + \gamma V_\phi(s_{t+1}) - V_\phi(s_t)\)

2. Value Model 训练设计

单轨迹采样需要可靠的 Value model 提供优势估计，SAO 提出三项实用技术：

更快的 Value 更新 (K=2)
每次策略更新时，Value model 进行 2 次更新，加速价值估计收敛。

Frozen-Attention 训练
冻结注意力层参数，仅训练 MoE 层。这一设计显著降低计算开销，同时保持价值估计质量。

Skip-Observation GAE
针对智能体任务中环境反馈的随机性，SAO 提出跨观测跳跃的 GAE 计算：

\[ \hat{A}(a_{i,N}) = \delta + \gamma\lambda\hat{A}(a_{i+1,0}) \] \[ \delta = r_t + \gamma V(a_{i+1,0}) - V(a_{i,N}) \]

其中 \(a_{i,N}\) 表示第 \(i\) 轮的第 \(N\) 个动作，\(a_{i+1,0}\) 为下一轮初始动作。这种设计将动作价值直接连接，过滤环境反馈的噪声。

3. 双边 Token-Level 裁剪 (DIS)

DIS (Direct Importance Sampling) 机制是 SAO 的训练稳定器：

直接使用 rollout 策略 \(\pi_{\text{rollout}}\) 作为行为策略，当前策略 \(\pi_\theta\) 进行重要性采样：

\[ r_t(\theta) = \exp\left(\log\pi_\theta(a_t|s_t) - \log\pi_{\text{rollout}}(a_t|s_t)\right) \]

校准函数 实现双边裁剪：

\[ f(x;\epsilon_\ell,\epsilon_h) = \begin{cases} x & \text{if } 1-\epsilon_\ell < x < 1+\epsilon_h \\ 0 & \text{otherwise} \end{cases} \]

与 PPO 的关键区别：超出信任区域的 token 被完全屏蔽（梯度=0），而非简单裁剪。这使得 SAO 可以采用更激进的裁剪参数 (\(\epsilon_{\text{low}}=0.3, \epsilon_{\text{high}}=5.0\)) 而不损失稳定性。

完整训练目标：

\[ \mathcal{L}(\theta) = \hat{\mathbb{E}}_t\left[f(r_t(\theta),\epsilon_\ell,\epsilon_h)\hat{A}_t\log\pi_\theta(a_t|s_t)\right] \]

算法流程

SAO 训练循环:
1. 初始化: 策略 π_θ, Value model V_φ
2. 异步 Rollout: π_rollout 持续生成轨迹 τ_i
3. 轨迹完成时立即:
   a. 计算重要性比 r_t(θ) = π_θ/π_rollout
   b. 屏蔽满足 r_t ∉ [1-ε_ℓ, 1+ε_h] 的 token
   c. 计算 Skip-Observation GAE 优势 Â_t
   d. 策略更新: L(θ) = E[f(r_t,ε_ℓ,ε_h)·Â_t·log π_θ]
   e. Value 更新: K=2 步 Frozen-Attention 训练
4. 无需等待批次完成，重复步骤 3

实验结果

基准测试

基准	类型	结果
SWE-Bench Verified	智能体编程	超越 GRPO 及其变体
BeyondAIME	数学推理	稳定提升
IMOAnswerBench	数学推理	稳定提升

关键发现

1. 训练稳定性：SAO 可持续训练 1000+ 步 而不崩溃，GRPO 约在 160 步后出现崩溃

2. 性能分化：SAO 与 GRPO (w/ DIS) 在初期表现相近，但约 400 步后 出现显著性能分化

3. 在线学习优势：在模拟的在线学习场景中（环境持续变化），单轨迹 RL 表现出特别强的适应性

实际部署

SAO 已成功应用于 GLM-5.2 (750B-A40B) 模型的训练，验证了方法在大规模生产环境的有效性。

方法对比总结

维度	GRPO	SAO
采样粒度	组级别 (G条)	单轨迹
同步方式	批次同步	完全异步
Off-policy 处理	旧策略检查点	直接存储 rollout log-probs
优势估计	组相对基线	Value-based GAE
训练稳定性	~160步崩溃	1000+步稳定
适用场景	短程推理	长程智能体任务

核心洞察

SAO 的设计体现了几个关键洞察：

1. 延迟是异步 RL 的头号敌人：单轨迹采样消除了同步等待，将 off-policy 程度降至最低

2. 稳定性来自严格的信任区域：DIS 的双边屏蔽策略比传统裁剪更能约束策略更新

3. Value model 需要特殊设计：Frozen-Attention 和 Skip-Observation GAE 是针对 LLM 智能体任务的实用优化

4. 异步 ≠ 不稳定：通过精心设计的采样策略和裁剪机制，异步训练可以达到甚至超越同步训练的稳定性

结语

SAO 为 LLM 后训练阶段的强化学习提供了一种新的范式，特别适用于需要长程交互的智能体任务。其核心贡献在于证明了：通过单轨迹采样配合严格的双边裁剪，异步 RL 可以同时实现高吞吐量和高稳定性。这一方法已被验证可支撑 750B 参数规模模型的训练，为下一代智能体模型的 RL 后训练提供了可靠的基础设施。

---

本文是对 arXiv:2607.07508 的技术解读，欢迎关注基模后训练领域的最新进展。