1. 从策略梯度到异步 Agent RL:为什么需要一套统一框架
搞大模型强化学习的人,大概率都经历过这样一个阶段:一开始看 Policy Gradient 的推导,觉得数学上挺干净,REINFORCE 的梯度估计、基线方差缩减、优势函数,一步步推下来逻辑自洽。但一旦把场景换成 LLM Agent,尤其是那种需要多轮工具调用、环境反馈延迟、多个 Agent 并行跑任务的异步场景,原来那套公式就有点不够用了。不是公式错了,而是它假设的“一个回合、一个策略、同步采样”的前提,在 Agent 场景里几乎全被打破。
我自己在这个方向上踩了不少坑。最早做单轮 RLHF 的时候,PPO 那套东西跑得挺顺,reward model 打分、advantage 估计、clip 更新,流程清晰。后来开始做带工具调用的 Agent,问题就来了:一个任务可能涉及十几次环境交互,每次交互的延迟不一样,有的工具秒回,有的要等好几秒;多个 Agent 并行跑的时候,轨迹长度参差不齐,有的早早结束,有的还在等环境返回。这时候如果还按同步 PPO 的方式去攒 batch,GPU 利用率低得可怜,训练效率直接崩掉。
所以这篇文章想做的事情很明确:把从最基础的 Policy Gradient 到异步 Agent RL 的整条推导链路捋清楚,并且给出一个能统一描述这些方法的框架。不是纯理论综述,而是带着“我要把它跑起来”的视角去推。核心关键词就几个:LLM 强化学习、Policy Gradient、Agent RL、公式推导、统一框架。适合谁看?如果你已经了解基本的强化学习概念,做过或者正准备做 LLM 的 RL 训练,尤其是涉及 Agent 场景的,那这篇内容应该能帮你省掉不少重新推公式的时间。
我个人的习惯是,任何一套 RL 方法,先把它拆成三个问题:策略怎么参数化、梯度怎么估计、数据怎么采样。这三个问题回答清楚了,方法的基本形态就定了。后面所有的变体,本质上都是在这三个维度上做取舍。下面我就按这个思路,从 Policy Gradient 开始,一步步推到异步 Agent RL。
2. Policy Gradient 的公式推导与核心直觉
2.1 从目标函数到梯度估计:REINFORCE 的完整推导
Policy Gradient 的出发点很朴素:我们有一个参数化的策略 $\pi_\theta(a|s)$,希望最大化期望回报 $J(\theta) = \mathbb{E}{\tau \sim \pi\theta}[R(\tau)]$。其中 $\tau = (s_0, a_0, s_1, a_1, \ldots)$ 是一条完整轨迹,$R(\tau)$ 是这条轨迹的总回报。
直接对 $J(\theta)$ 求梯度,难点在于期望的分布本身依赖 $\theta$。这里用到一个关键技巧——log-derivative trick,也叫 score function 技巧:
$$\nabla_\theta \pi_\theta(\tau) = \pi_\theta(\tau) \nabla_\theta \log \pi_\theta(\tau)$$
这个等式的来源很简单,就是 $\nabla \log f = \nabla f / f$ 的变形。有了它,梯度就可以写成:
$$\nabla_\theta J(\theta) = \mathbb{E}{\tau \sim \pi\theta}[\nabla_\theta \log \pi_\theta(\tau) \cdot R(\tau)]$$
再把 $\log \pi_\theta(\tau)$ 展开。因为轨迹概率是 $\pi_\theta(\tau) = \rho(s_0) \prod_{t=0}^{T-1} \pi_\theta(a_t|s_t) P(s_{t+1}|s_t,a_t)$,其中环境转移 $P$ 和初始分布 $\rho$ 都不依赖 $\theta$,所以:
$$\nabla_\theta \log \pi_\theta(\tau) = \sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t|s_t)$$
代入后得到 REINFORCE 的梯度估计:
$$\nabla_\theta J(\theta) = \mathbb{E}{\tau \sim \pi\theta}\left[\sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t|s_t) \cdot R(\tau)\right]$$
这个公式的直觉是:如果一条轨迹的总回报高,就提高这条轨迹上所有动作的概率;如果回报低,就降低它们的概率。但这里有个明显的问题——用整条轨迹的回报 $R(\tau)$ 去乘每个动作的 log 概率,方差非常大。因为一个动作的好坏,不应该由整条轨迹的回报来决定,而应该由这个动作之后的回报来决定。
2.2 因果性与基线:方差缩减的两个关键手段
第一个改进是引入因果性。在时刻 $t$ 采取的动作 $a_t$,不可能影响时刻 $t$ 之前的奖励。所以 $R(\tau)$ 可以替换成从 $t$ 开始的累积回报 $G_t = \sum_{t'=t}^{T-1} \gamma^{t'-t} r_{t'}$。这样梯度变成:
$$\nabla_\theta J(\theta) = \mathbb{E}\left[\sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t|s_t) \cdot G_t\right]$$
第二个改进是引入基线$b(s_t)$。因为 $\mathbb{E}{a \sim \pi\theta}[\nabla_\theta \log \pi_\theta(a|s) \cdot b(s)] = 0$,所以减去一个只依赖状态的基线不会改变梯度的期望,但能显著降低方差。最常用的基线是状态价值函数 $V(s_t)$,于是 $G_t - V(s_t)$ 就成了优势函数$A_t$ 的一个估计。
这一步在 LLM 场景里特别重要。因为 LLM 的 action space 是词表大小,动辄几万到十几万,方差本来就大。如果不做基线,训练很容易震荡。我实测下来,加不加 baseline,收敛速度能差两三倍。
2.3 从 REINFORCE 到 Actor-Critic:引入价值函数
REINFORCE 是蒙特卡洛方法,必须等一条轨迹跑完才能更新。在 LLM 场景里,这意味着要等整个生成过程结束,拿到 reward 之后才能算梯度。如果序列很长,或者 Agent 需要多轮交互,这个等待成本很高。
Actor-Critic 的思路是:用一个价值网络 $V_\phi(s)$ 来估计状态价值,然后用 TD 误差 $\delta_t = r_t + \gamma V_\phi(s_{t+1}) - V_\phi(s_t)$ 作为优势的估计。这样就不需要等整条轨迹结束,可以做到单步更新或者多步更新。Actor 负责策略 $\pi_\theta$,Critic 负责价值 $V_\phi$,两者交替更新。
在 LLM 的 RLHF 里,PPO 就是典型的 Actor-Critic 结构。Actor 是语言模型本身,Critic 通常是另一个同规模或者稍小的模型,输出每个 token 位置的价值估计。这里有个工程上的坑:Critic 的输入和 Actor 一样,但输出是一个标量,所以 Critic 的最后一层需要改。而且 Critic 的训练往往比 Actor 更不稳定,因为价值函数的尺度会随着 reward 的变化而漂移。
3. LLM 场景下 Policy Gradient 的适配与改造
3.1 动作空间与序列生成:token 级策略的梯度计算
LLM 的强化学习和传统 RL 最大的区别在于动作空间。传统 RL 里,动作通常是离散的有限集合或者连续向量;而 LLM 的每个“动作”是在词表上选一个 token,整个序列生成过程就是一系列 token 决策的串联。
假设 prompt 是 $x$,生成的回复是 $y = (y_1, y_2, \ldots, y_T)$,那么策略就是:
$$\pi_\theta(y|x) = \prod_{t=1}^{T} \pi_\theta(y_t | x, y_{<t})$$
对应的 log 概率是 $\sum_{t=1}^{T} \log \pi_\theta(y_t | x, y_{<t})$。如果整条回复只有一个最终 reward $r(x,y)$,那么 REINFORCE 的梯度就是:
$$\nabla_\theta J = \mathbb{E}\left[\nabla_\theta \log \pi_\theta(y|x) \cdot (r(x,y) - b)\right]$$
这里 $b$ 通常是一个 baseline,可以是当前 batch 的平均 reward,也可以是价值网络的估计。注意,这里 reward 是序列级别的,不是 token 级别的。所以同一个序列里所有 token 共享同一个 reward 信号。这在实际操作中会带来一个问题:如果一条回复整体得分高,但中间某个 token 其实选得不好,它也会被一起加强。这是序列级 reward 的固有局限。
3.2 奖励设计与优势估计:从稀疏到稠密
在 LLM 的 Agent 场景里,reward 往往更复杂。一个任务可能包含多个子目标,比如“先查天气,再根据天气推荐穿搭,最后生成一段文案”。如果只在最后给一个总分,那就是典型的稀疏奖励,学习效率很低。
常见的做法是设计过程奖励或者中间奖励。比如每个工具调用成功给一个小奖励,格式正确给一个小奖励,最终任务完成给一个大奖励。这样 reward 信号更稠密,梯度估计的方差也会小很多。
但这里有个权衡:过程奖励设计得不好,容易导致 reward hacking。模型可能会学会刷那些容易拿的小奖励,而忽略最终目标。我自己的经验是,过程奖励的权重不要太高,通常控制在最终奖励的 10% 到 20% 左右比较合适。而且最好用优势函数来做归一化,而不是直接用原始 reward。
优势估计方面,除了前面说的 $G_t - V(s_t)$,在 LLM 里还常用GAE(Generalized Advantage Estimation)。GAE 通过引入参数 $\lambda$ 来平衡偏差和方差:
$$A_t^{GAE(\gamma,\lambda)} = \sum_{l=0}^{\infty} (\gamma \lambda)^l \delta_{t+l}$$
其中 $\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$。当 $\lambda=0$ 时,退化成单步 TD 误差,偏差大方差小;当 $\lambda=1$ 时,退化成蒙特卡洛,偏差小方差大。实际用的时候,$\lambda$ 通常取 0.95 左右,$\gamma$ 取 1.0(因为 LLM 的回合通常有明确终止)。
3.3 PPO 在 LLM 训练中的关键参数与实操细节
PPO 是目前 LLM RL 训练里最常用的算法,核心是** clipped surrogate objective**:
$$L^{CLIP}(\theta) = \mathbb{E}\left[\min\left(\rho_t(\theta) A_t, \text{clip}(\rho_t(\theta), 1-\epsilon, 1+\epsilon) A_t\right)\right]$$
其中 $\rho_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}$ 是重要性采样比率。clip 的作用是限制每次更新的幅度,防止策略跑偏。
在 LLM 场景里,有几个参数特别关键:
| 参数 | 典型值 | 作用 | 踩坑经验 |
|---|---|---|---|
| clip range $\epsilon$ | 0.1 ~ 0.2 | 限制策略更新幅度 | 太大容易崩,太小学习慢 |
| KL 系数 | 0.01 ~ 0.1 | 约束新策略不偏离旧策略太远 | 太大导致学不动,太小导致 reward hacking |
| batch size | 64 ~ 512 | 每次更新的样本数 | 太小方差大,太大显存吃不消 |
| mini-batch | 8 ~ 64 | 梯度累积的粒度 | 影响训练稳定性 |
| 学习率 | 1e-6 ~ 5e-6 | Actor 的学习率 | 比预训练小一到两个数量级 |
KL 惩罚这一项在 LLM 里几乎是必须的。因为语言模型的输出空间太大,如果不加约束,策略很容易退化到重复输出某些高 reward 的模板,也就是常说的 mode collapse。我一般会把 KL 系数设成动态的,根据实际 KL 散度来调整,让 KL 维持在一个目标值附近。
还有一个实操细节:reward 归一化。不同任务的 reward 尺度可能差很多,有的任务满分是 1,有的是 10。如果不做归一化,梯度尺度会很不稳定。我通常会在 batch 内做标准化,减去均值除以标准差,然后再乘一个缩放系数。
4. 异步 Agent RL 的挑战与统一框架构建
4.1 异步采样带来的分布偏移问题
到了 Agent 场景,事情变得更复杂。一个 Agent 任务通常包含多轮交互:模型生成一个动作(比如调用某个工具),环境返回一个观察,模型再生成下一个动作,直到任务结束。如果多个 Agent 并行跑,每个任务的进度不一样,有的已经跑到第 10 步,有的还在第 2 步。
异步采样的核心问题是分布偏移。当我们在时刻 $t$ 用策略 $\pi_{\theta_{old}}$ 采样了一个动作,但等到这个样本被用于更新时,策略已经变成了 $\pi_{\theta_{new}}$。如果 $\theta_{new}$ 和 $\theta_{old}$ 差太多,重要性采样比率就会很大或者很小,导致梯度估计不准。
在同步 PPO 里,这个问题通过 clip 和 KL 约束来缓解。但在异步场景里,样本的“年龄”可能差很多,有的样本是几步之前的策略采的,有的是几十步之前的。如果统一用一个 clip range,老样本的比率可能早就超出范围了。
4.2 统一框架:把同步与异步 RL 放进同一个数学表达
我一直在想,能不能用一个统一的框架来描述这些方法。后来发现,关键是把采样策略和更新策略分开看。定义 $\mu$ 为采样策略,$\pi_\theta$ 为当前要优化的策略,那么重要性采样比率就是:
$$\rho = \frac{\pi_\theta(a|s)}{\mu(a|s)}$$
同步 PPO 里,$\mu = \pi_{\theta_{old}}$,而且 $\theta_{old}$ 和 $\theta$ 只差一次更新。异步场景里,$\mu$ 可能是多个历史策略的混合,每个样本的 $\mu$ 都不一样。
统一的目标函数可以写成:
$$J(\theta) = \mathbb{E}_{(s,a) \sim \mu}\left[\rho \cdot A(s,a) \cdot w(\rho)\right]$$
其中 $w(\rho)$ 是一个权重函数,用来控制重要性采样的影响。PPO 的 clip 本质上就是一种 $w(\rho)$:当 $\rho$ 在 $[1-\epsilon, 1+\epsilon]$ 内时,$w=1$;超出范围时,梯度被截断。而在异步场景里,可以根据样本的“年龄”来调整 $w$,老样本的权重低一些,新样本的权重高一些。
这个框架的好处是,它把同步和异步统一起来了。同步只是异步的一个特例,即所有样本的 $\mu$ 都相同且足够新。基于这个框架,可以设计更灵活的异步算法,比如按样本年龄加权或者按策略差异加权。
4.3 工程实现:异步 Agent RL 的系统架构要点
理论框架清楚了,工程实现还有一堆坑。异步 Agent RL 的系统通常包含几个模块:推理服务、环境服务、经验回放、训练服务。推理服务负责用当前策略生成动作,环境服务负责执行动作并返回观察,经验回放负责存储轨迹,训练服务负责从回放里采样并更新策略。
这里最大的挑战是吞吐量匹配。推理和环境的交互速度往往比训练慢很多,如果训练服务等数据,GPU 利用率就上不去。常见的做法是让推理服务持续跑,把经验不断写入回放池,训练服务从池子里采样。但这样又引入了样本陈旧的问题。
我自己的经验是,回放池的容量不要太大,通常控制在几千条轨迹左右。太大了样本太旧,太小了又不够训练用。另外,可以给每条轨迹打一个时间戳,训练的时候优先采样新的轨迹,老轨迹按一定概率丢弃。这样能在吞吐量和样本新鲜度之间找到一个平衡。
还有一个细节是环境调用的容错。Agent 调用工具的时候,工具可能超时、返回错误、或者返回格式不对。如果不处理,整个训练流程可能卡住。我一般会设置超时重试,重试几次还不行就返回一个默认观察,并给一个负奖励。这样至少不会让训练中断。
5. 实操中的常见问题与排查技巧
5.1 训练不收敛或 reward 震荡的排查思路
LLM 的 RL 训练不收敛,原因可能有很多。我一般按下面的顺序排查:
第一,看KL 散度。如果 KL 突然飙升,说明策略跑偏了,可能是学习率太大或者 KL 系数太小。这时候先把学习率降下来,或者把 KL 系数调大。
第二,看reward 分布。如果 reward 的方差特别大,或者大部分样本的 reward 都差不多,那梯度信号就很弱。这时候要检查 reward 设计是否合理,或者考虑做 reward 归一化。
第三,看advantage 的尺度。如果 advantage 的绝对值特别大,梯度就会很大,容易导致更新过猛。通常会对 advantage 做标准化,减去均值除以标准差。
第四,看clip fraction。如果大部分样本的比率都被 clip 了,说明新旧策略差异太大,可能是采样和更新之间的间隔太长。这时候要减少每次更新的样本数,或者增加采样频率。
5.2 异步场景下的样本陈旧与重要性采样修正
异步场景里,样本陈旧是绕不开的问题。我试过几种修正方法:
一种是按年龄衰减。给每个样本一个权重 $w = \gamma^{age}$,age 是样本从采集到被使用的步数。这样老样本的权重自然降低。
另一种是按策略差异修正。计算采样策略和当前策略的 KL 散度,如果差异太大,就降低这个样本的权重,或者直接丢弃。
还有一种更激进的做法是只保留最新样本。回放池只保留最近 N 步的样本,老的直接扔掉。这样样本新鲜度有保证,但可能会浪费一些数据。
实测下来,按年龄衰减比较稳,实现也简单。通常 $\gamma$ 取 0.99 左右,也就是每过 100 步,权重降到原来的 0.37。这样既利用了老样本,又不会让它们主导更新。
5.3 常见问题速查表
| 问题现象 | 可能原因 | 排查方法 | 解决思路 |
|---|---|---|---|
| reward 不涨 | 学习率太小、KL 约束太强 | 看梯度范数、KL 值 | 调大学习率、放松 KL |
| reward 震荡 | 学习率太大、batch 太小 | 看 reward 曲线、梯度方差 | 降学习率、增大 batch |
| KL 飙升 | 策略跑偏、reward hacking | 看 KL 曲线、输出样本 | 增大 KL 系数、检查 reward |
| 输出重复 | mode collapse | 看生成样本多样性 | 增大 KL 系数、加熵正则 |
| 训练速度慢 | 环境调用阻塞、回放池太小 | 看 GPU 利用率、数据吞吐 | 异步化、增大回放池 |
| 显存溢出 | batch 太大、序列太长 | 看显存占用 | 减小 batch、梯度累积 |
5.4 几个容易被忽略的实操细节
第一个是随机种子。LLM 的 RL 训练对随机种子很敏感,不同的种子可能导致完全不同的结果。我一般会跑至少三个种子,取平均来看效果。如果只有一个种子的结果,很难判断是算法好还是运气好。
第二个是评估频率。训练过程中要定期在验证集上评估,不能只看训练 reward。因为训练 reward 可能会被 reward hacking 抬高,但实际效果并没有变好。评估的时候要用独立的 reward 或者人工评估。
第三个是检查点管理。RL 训练很容易崩,所以一定要定期保存检查点。我一般每几百步就存一次,而且会保留最近几个检查点,方便回滚。
第四个是日志记录。除了 reward 和 loss,还要记录 KL、clip fraction、advantage 均值方差、梯度范数、学习率等。这些指标在排查问题的时候非常有用。我习惯用 TensorBoard 或者 WandB 来可视化,一眼就能看出哪里不对。
6. 从公式到落地:一套可复现的异步 Agent RL 方案
6.1 整体流程与模块划分
把前面所有东西串起来,一套完整的异步 Agent RL 方案大概是这样:
- 策略模型:LLM,负责生成动作
- 价值模型:可以是同规模模型,也可以是轻量级模型,负责估计状态价值
- 推理服务:用当前策略生成轨迹,支持多 Agent 并行
- 环境服务:执行工具调用,返回观察和奖励
- 经验回放池:存储轨迹,支持按时间戳采样
- 训练服务:从回放池采样,计算 GAE,更新策略和价值网络
整个流程是异步的:推理服务持续生成轨迹写入回放池,训练服务持续从回放池采样更新。两者通过回放池解耦,互不阻塞。
6.2 关键参数配置与调参经验
下面是我在实际项目中用过的一套配置,供参考:
# 策略优化 learning_rate: 2e-6 clip_range: 0.2 kl_coef: 0.05 entropy_coef: 0.01 gamma: 1.0 lambda_gae: 0.95 # 采样 batch_size: 128 mini_batch_size: 16 rollout_steps: 512 replay_buffer_size: 4096 sample_age_decay: 0.99 # 训练 num_epochs: 4 max_grad_norm: 1.0 value_loss_coef: 0.5 target_kl: 0.02调参的时候,我一般先固定其他参数,只调学习率和 KL 系数。学习率从 1e-6 开始试,如果 reward 不涨就慢慢加大,如果震荡就减小。KL 系数从 0.05 开始,如果 KL 经常超过目标值就加大,如果 KL 一直很低就减小。
6.3 效果验证与迭代方向
验证的时候,不能只看训练 reward。我一般会看几个指标:任务完成率、平均交互步数、工具调用成功率、输出多样性。任务完成率是最直接的,但有时候模型会学会一些捷径,所以还要看交互步数和工具调用成功率,确保模型是真的在解决问题,而不是在刷分。
迭代方向的话,我觉得有几个值得探索的点。一个是更精细的信用分配,现在序列级 reward 还是太粗,如果能做到 token 级或者步骤级的信用分配,学习效率会高很多。另一个是多 Agent 协同,现在大部分工作还是单 Agent,多 Agent 场景下的策略优化和信用分配是另一个难题。还有就是离线与在线的结合,用离线数据预训练,再用在线数据微调,可能比纯在线效率更高。
我个人在实际操作中的体会是,LLM 的 RL 训练,算法本身只是一部分,工程实现和数据质量同样重要。同样的算法,不同的 reward 设计、不同的采样策略、不同的超参,效果可能差好几倍。所以不要迷信某个算法,多试、多调、多记录,才是正道。