ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

从Policy Gradient到异步Agent RL:统一框架与公式推导

从Policy Gradient到异步Agent RL:统一框架与公式推导 1. 从策略梯度到异步 Agent RL为什么需要一套统一框架搞大模型强化学习的人大概率都经历过这样一个阶段一开始看 Policy Gradient 的推导觉得数学上挺干净REINFORCE 的梯度估计、基线方差缩减、优势函数一步步推下来逻辑自洽。但一旦把场景换成 LLM Agent尤其是那种需要多轮工具调用、环境反馈延迟、多个 Agent 并行跑任务的异步场景原来那套公式就有点不够用了。不是公式错了而是它假设的“一个回合、一个策略、同步采样”的前提在 Agent 场景里几乎全被打破。我自己在这个方向上踩了不少坑。最早做单轮 RLHF 的时候PPO 那套东西跑得挺顺reward model 打分、advantage 估计、clip 更新流程清晰。后来开始做带工具调用的 Agent问题就来了一个任务可能涉及十几次环境交互每次交互的延迟不一样有的工具秒回有的要等好几秒多个 Agent 并行跑的时候轨迹长度参差不齐有的早早结束有的还在等环境返回。这时候如果还按同步 PPO 的方式去攒 batchGPU 利用率低得可怜训练效率直接崩掉。所以这篇文章想做的事情很明确把从最基础的 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_{t0}^{T-1} \pi_\theta(a_t|s_t) P(s_{t1}|s_t,a_t)$其中环境转移 $P$ 和初始分布 $\rho$ 都不依赖 $\theta$所以$$\nabla_\theta \log \pi_\theta(\tau) \sum_{t0}^{T-1} \nabla_\theta \log \pi_\theta(a_t|s_t)$$代入后得到 REINFORCE 的梯度估计$$\nabla_\theta J(\theta) \mathbb{E}{\tau \sim \pi\theta}\left[\sum_{t0}^{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_{tt}^{T-1} \gamma^{t-t} r_{t}$。这样梯度变成$$\nabla_\theta J(\theta) \mathbb{E}\left[\sum_{t0}^{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_{t1}) - 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_{t1}^{T} \pi_\theta(y_t | x, y_{t})$$对应的 log 概率是 $\sum_{t1}^{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 里还常用GAEGeneralized Advantage Estimation。GAE 通过引入参数 $\lambda$ 来平衡偏差和方差$$A_t^{GAE(\gamma,\lambda)} \sum_{l0}^{\infty} (\gamma \lambda)^l \delta_{tl}$$其中 $\delta_t r_t \gamma V(s_{t1}) - V(s_t)$。当 $\lambda0$ 时退化成单步 TD 误差偏差大方差小当 $\lambda1$ 时退化成蒙特卡洛偏差小方差大。实际用的时候$\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 hackingbatch 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]$ 内时$w1$超出范围时梯度被截断。而在异步场景里可以根据样本的“年龄”来调整 $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 值调大学习率、放松 KLreward 震荡学习率太大、batch 太小看 reward 曲线、梯度方差降学习率、增大 batchKL 飙升策略跑偏、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 设计、不同的采样策略、不同的超参效果可能差好几倍。所以不要迷信某个算法多试、多调、多记录才是正道。
返回列表