☰
从Policy Gradient到异步Agent RL:统一框架与公式推导
2026/10/7 13:28:35 网站建设 项目流程

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 size64 ~ 512每次更新的样本数太小方差大,太大显存吃不消
mini-batch8 ~ 64梯度累积的粒度影响训练稳定性
学习率1e-6 ~ 5e-6Actor 的学习率比预训练小一到两个数量级

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 设计、不同的采样策略、不同的超参,效果可能差好几倍。所以不要迷信某个算法,多试、多调、多记录,才是正道。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询