[CS336-17] Policy Gradient 与 GRPO 细节

这篇是 Stanford CS336 Spring 2025 第 17 讲 Course Materials 的学习笔记。它继续深入 policy gradient 和 GRPO 的机制。

1. 把语言模型看成 RL policy#

在语言模型 RL 里，state 可以理解成 prompt 加上已经生成的 response prefix，action 是下一个 token，policy 就是语言模型给出的 next-token distribution。

一次 rollout 是从 prompt 开始生成完整回答，最后得到 outcome reward。对于数学或代码任务，reward 可以是答案是否正确、测试是否通过。

这和机器人 RL 不同：状态转移基本由字符串拼接决定，reward 通常在最后才出现，而且可以对同一 prompt 反复采样多个回答。

2. Policy gradient 的基本推导#

目标是最大化期望 reward：

E[R]

通过 log-derivative trick，可以得到 policy gradient 形式：

gradient ≈ grad log pi(a | s) * R(s, a)

直觉是：高 reward 的回答应该提高概率，低 reward 的回答应该降低或不提高概率。

如果 reward 是 0/1，naive policy gradient 只会从正确回答中学习。这在 sparse reward 场景下方差很高，训练不稳定。

3. Baseline 和 advantage#

为了降低方差，可以从 reward 中减去 baseline：

grad log pi(a | s) * (R(s, a) - b(s))

只要 baseline 不依赖 action，就不会改变期望梯度，但可以显著降低噪声。

最自然的 baseline 是同一 state 下的平均 reward。于是得到 advantage：

A(s, a) = Q(s, a) - V(s)

它表示某个 action 比当前 state 下的平均水平好多少。

4. GRPO 的 group baseline#

GRPO 利用了语言模型任务中的 group structure：对同一个 prompt 采样多条回答。

对每个 prompt 的一组 responses，计算 reward 后做中心化或标准化。高于组内平均的回答得到正 advantage，低于组内平均的回答得到负 advantage。

这避免了单独训练 value model，也让 baseline 更贴近当前 prompt。

5. Reward normalization 的影响#

Lecture 17 用一个排序任务展示了不同 reward 处理方式的影响。

直接使用 raw reward 时，模型可能学习慢、容易卡住。使用 centered reward 后，低于平均的样本会产生负向更新，同一组都一样好或一样差时不更新。再除以标准差可以进一步归一化不同 prompt 的难度。

但标准差归一化也可能带来长度偏差或过度放大容易/困难样本，因此后续 GRPO 变体会调整这些细节。

6. KL penalty 和 reference model#

RL 更新可能让模型偏离原始能力，尤其在窄 reward 上过拟合。显式 KL penalty 可以约束新 policy 不要离 reference policy 太远。

这在想给模型加入新能力、又不希望它忘掉原有语言能力时很重要。

7. takeaway#

第 17 讲把 GRPO 从“一个热门算法名”拆成几个基本部件：

• policy gradient 用 log probability 加权 reward 更新模型；
• baseline 不改变期望梯度，但能降低方差；
• advantage 表示相对当前 prompt 的好坏；
• GRPO 用同 prompt 多回答的组内统计替代 value model；
• reward 归一化和 KL penalty 是稳定训练的关键细节。

因此 GRPO 的核心并不神秘：它是适配语言模型 rollout 结构的简化 policy gradient。

参考#

• Stanford CS336 Spring 2025 Course Materials: lecture_17.py