斯坦福大学 | CS336 | 从零开始构建语言模型 | Spring 2025 | 笔记 | Assignment 5: GRPO

目录

• • 前言
  • [1. Primer on Policy Gradients](#1. Primer on Policy Gradients)
  • • [1.1 Language Models as Policies](#1.1 Language Models as Policies)
    • [1.2 Trajectories](#1.2 Trajectories)
    • [1.3 Rewards and Return](#1.3 Rewards and Return)
    • [1.4 Vanilla Policy Gradient](#1.4 Vanilla Policy Gradient)
    • [1.5 Policy Gradient Baselines](#1.5 Policy Gradient Baselines)
    • [1.6 Off-Policy Policy Gradient](#1.6 Off-Policy Policy Gradient)
  • [2. Group Relative Policy Optimization](#2. Group Relative Policy Optimization)
  • • [2.1 GRPO Algorithm](#2.1 GRPO Algorithm)
    • [2.2 Implement](#2.2 Implement)
    • • [Problem (compute_group_normalized_rewards): Group normalization (2 points)](#Problem (compute_group_normalized_rewards): Group normalization (2 points))
      • [Problem (compute_naive_policy_gradient_loss): Naive policy gradient (1 point)](#Problem (compute_naive_policy_gradient_loss): Naive policy gradient (1 point))
      • [Problem (compute_grpo_clip_loss): GRPO-Clip loss (2 points)](#Problem (compute_grpo_clip_loss): GRPO-Clip loss (2 points))
      • [Problem (compute_policy_gradient_loss): Policy-gradient wrapper (1 point)](#Problem (compute_policy_gradient_loss): Policy-gradient wrapper (1 point))
      • [Problem (masked_mean): Masked mean (1 point)](#Problem (masked_mean): Masked mean (1 point))
      • [Problem (grpo_microbatch_train_step): Microbatch train step (3 points)](#Problem (grpo_microbatch_train_step): Microbatch train step (3 points))
      • [Problem (grpo_train_loop): GRPO train loop (5 points)](#Problem (grpo_train_loop): GRPO train loop (5 points))
  • [3. GRPO Experiments](#3. GRPO Experiments)
  • • • [Problem (grpo_learning_rate): Tune the learning rate (2 points) (6 H100 hrs)](#Problem (grpo_learning_rate): Tune the learning rate (2 points) (6 H100 hrs))
      • [Problem (grpo_baselines): Effect of baselining (2 points) (2 H100 hrs)](#Problem (grpo_baselines): Effect of baselining (2 points) (2 H100 hrs))
      • [Problem (think_about_length_normalization): Think about length normalization (1 point)](#Problem (think_about_length_normalization): Think about length normalization (1 point))
      • [Problem (grpo_length_normalization): Effect of length normalization (2 points) (2 H100 hrs)](#Problem (grpo_length_normalization): Effect of length normalization (2 points) (2 H100 hrs))
      • [Problem (grpo_group_standard_deviation): Effect of standard deviation normalization (2 points) (2 H100 hrs)](#Problem (grpo_group_standard_deviation): Effect of standard deviation normalization (2 points) (2 H100 hrs))
      • [Problem (grpo_off_policy): Implement off-policy GRPO](#Problem (grpo_off_policy): Implement off-policy GRPO)
      • [Problem (grpo_off_policy_sweep): Off-policy GRPO hyperparameter sweep (4 points)(12 H100 hrs)](#Problem (grpo_off_policy_sweep): Off-policy GRPO hyperparameter sweep (4 points)(12 H100 hrs))
      • [Problem (grpo_off_policy_clip_ablation): Off-policy GRPO-Clip ablation (2 points) (2 H100 hrs)](#Problem (grpo_off_policy_clip_ablation): Off-policy GRPO-Clip ablation (2 points) (2 H100 hrs))
      • [Problem (grpo_prompt_ablation): Prompt ablation (2 points) (2 H100 hrs)](#Problem (grpo_prompt_ablation): Prompt ablation (2 points) (2 H100 hrs))
  • [4. Leaderboard: GRPO on MATH](#4. Leaderboard: GRPO on MATH)
  • • • [Problem (leaderboard): Leaderboard (16 points) (16 H100 hrs)](#Problem (leaderboard): Leaderboard (16 points) (16 H100 hrs))
  • [5. Epilogue](#5. Epilogue)
  • 结语
  • 参考

前言

本篇文章记录 CS336 作业 Assignment 5: Alignment 中的 GRPO 作业要求，仅供自己参考😄

Assignment 5 ：https://github.com/stanford-cs336/assignment5-alignment

reference ：https://chatgpt.com/

1. Primer on Policy Gradients

以下内容均翻译自 cs336_spring2025_assignment5_alignment.pdf，请大家查看原文档获取更详细的内容

语言模型研究中的一个令人振奋的新发现是：针对 可验证的奖励（verified rewards） 进行强化学习（RL），并结合强大的基础模型，可以显著提升模型的推理能力和整体性能 [OpenAI+ 2024] [DeepSeek-AI+ 2025]

目前最强的开源推理模型，例如 DeepSeek-R1 和 Kimi k1.5 [Kimi Team+ 2025]，都是通过 策略梯度（policy gradients） 进行训练的，这是一种能够优化任意奖励函数的强大强化学习算法。

下面我们将对语言模型上的策略梯度强化学习做一个简要介绍，我们的讲解主要基于几份优秀的资料，这些资料对相关概念进行了更深入的阐述，包括：

• OpenAI 的《Spinning Up in DeepRL》[Achiam, 2018a]
• Nathan Lambert 的《Reinforcement Learning from Human Feedback (RLHF)》书籍 [Lambert, 2024]

1.1 Language Models as Policies

一个带有参数 θ \theta θ 的因果语言模型（LM），在给定当前文本前缀 s t s_t st（即状态 / 观测）时，会定义下一个 token a t ∈ V a_t \in \mathcal{V} at∈V 的概率分布。在强化学习（RL）的语境下：

• 下一个 token a t a_t at 被视为一个 "动作"（action）
• 当前文本前缀 s t s_t st 被视为 "状态"（state）

因此，语言模型可以看作是一个 离散的随机策略（categorical stochastic policy）：

a t ∼ π θ ( ⋅ ∣ s t ) , π θ ( a t ∣ s t ) = [ softmax ( f θ ( s t ) ) ] a t . (3) a_t \sim \pi_\theta(\cdot \mid s_t), \quad \pi_\theta(a_t \mid s_t) = [\text{softmax}(f_\theta(s_t))]_{a_t}. \tag{3} at∼πθ(⋅∣st),πθ(at∣st)=[softmax(fθ(st))]at.(3)

在使用策略梯度优化策略时，需要两个基本操作：

1. 从策略中采样（Sampling from the policy） 即从上述类别分布中抽取一个动作 a t a_t at。

2. 计算动作的对数似然（Scoring the log-likelihood of an action） 即计算 log ⁡ π θ ( a t ∣ s t ) \log \pi_\theta(a_t \mid s_t) logπθ(at∣st)。

通常，在使用 LLM 进行强化学习时 s t s_t st 是当前已经生成的部分解答（partial completion / solution），每个 a t a_t at 是解答中的下一个 token，当生成结束标记时，一个 episode 结束，例如 <|end_of_text|> 或在我们的 r1_zero prompt 中使用的 </answer>。

1.2 Trajectories

一个（有限时域）轨迹是智能体在交互过程中经历的状态与动作交替组成的序列：

τ = ( s 0 , a 0 , s 1 , a 1 , ... , s T , a T ) (4) \tau = (s_0, a_0, s_1, a_1, \dots, s_T, a_T) \tag{4} τ=(s0,a0,s1,a1,...,sT,aT)(4)

其中 T T T 是轨迹的长度，即：

• 要么 a T a_T aT 是一个文本结束 token
• 要么已经达到了最大生成 token 数限制

初始状态来自起始分布，即 s 0 ∼ ρ 0 ( s 0 ) s_0 \sim \rho_0(s_0) s0∼ρ0(s0)，在 LLM 强化学习的情境下， ρ 0 ( s 0 ) \rho_0(s_0) ρ0(s0) 表示对 格式化 prompt 的分布 。在一般情况下，状态转移遵循某种环境动力学 s t + 1 ∼ P ( ⋅ ∣ s t , a t ) s_{t+1} \sim P(\cdot \mid s_t, a_t) st+1∼P(⋅∣st,at)，而在基于 LLM 的强化学习中，环境是 确定性的 ：下一个状态等于旧的前缀加上新生成的 token 即 s t + 1 = s t ∣ ∣ a t s_{t+1} = s_t || a_t st+1=st∣∣at

轨迹（trajectories）也常被称为 episode 或 rollout，这些术语在后文中会交替使用。

1.3 Rewards and Return

一个标量奖励 r t = R ( s t , a t ) r_t = R(s_t, a_t) rt=R(st,at) 用于评估系统在状态 s t s_t st 下采取动作的即时质量。在具有可验证结果的强化学习任务中，通常的做法是：对中间步骤赋予 零奖励 ，对最终动作赋予一个 验证奖励（verified reward）

具体而言：

r T = R ( s T , a T ) : = { 1 , 如果轨迹 s T ∣ ∣ a T 根据奖励函数与真实答案匹配 0 , 否则 r_T = R(s_T, a_T) := \begin{cases} 1, & \text{如果轨迹 } s_T || a_T \text{ 根据奖励函数与真实答案匹配} \\ 0, & \text{否则} \end{cases} rT=R(sT,aT):={1,0,如果轨迹 sT∣∣aT 根据奖励函数与真实答案匹配否则

轨迹的回报 R ( τ ) R(\tau) R(τ) 是对整条轨迹上所有奖励的累积。常见的两种定义方式包括：

1. 有限时域无折扣回报

R ( τ ) : = ∑ t = 0 T r t (5) R(\tau) := \sum_{t=0}^{T} r_t \tag{5} R(τ):=t=0∑Trt(5)

2. 无限时域折扣回报

R ( τ ) : = ∑ t = 0 ∞ γ t r t , 0 < γ < 1 (6) R(\tau) := \sum_{t=0}^{\infty} \gamma^t r_t,\quad 0 < \gamma < 1 \tag{6} R(τ):=t=0∑∞γtrt,0<γ<1(6)

在我们的设置中，将采样 无折扣回报 形式，因为每个 episode 都具有自然的终止条件（如生成结束 token 或达到最大生成长度）。

智能体（Agent）的目标是最大化期望回报：

J ( θ ) = E τ ∼ π θ [ R ( τ ) ] (7) J(\theta) = \mathbb{E}{\tau \sim \pi\theta} [R(\tau)] \tag{7} J(θ)=Eτ∼πθ[R(τ)](7)

从而得到优化目标：

θ ∗ = arg ⁡ max ⁡ θ J ( θ ) (8) \theta^* = \arg\max_\theta J(\theta) \tag{8} θ∗=argθmaxJ(θ)(8)

1.4 Vanilla Policy Gradient

接下来，我们尝试通过对期望回报进行 梯度上升（gradient ascent） 来学习策略参数 θ \theta θ：

θ k + 1 = θ k + α ∇ θ J ( θ k ) (9) \theta_{k+1} = \theta_k + \alpha \nabla_\theta J(\theta_k) \tag{9} θk+1=θk+α∇θJ(θk)(9)

实现这一目标的核心公式是 REINFORCE 策略梯度，如下所示：

∇ θ J ( π θ ) = E τ ∼ π θ [ ∑ t = 0 T ∇ θ log ⁡ π θ ( a t ∣ s t ) R ( τ ) ] (10) \nabla_\theta J(\pi_\theta) = \mathbb{E}{\tau \sim \pi\theta} \left[ \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t \mid s_t) R(\tau) \right] \tag{10} ∇θJ(πθ)=Eτ∼πθ[t=0∑T∇θlogπθ(at∣st)R(τ)](10)

Deriving the policy gradient.

我们是如何得到这个公式的呢？为完整起见，下面给出该公式的推导思路，我们将使用基本恒等式：

1. 轨迹的概率

一条轨迹的概率为：

P ( τ ∣ θ ) = ρ 0 ( s 0 ) ∏ t = 0 T P ( s t + 1 ∣ s t , a t ) π θ ( a t ∣ s t ) (11) P(\tau \mid \theta)=\rho_0(s_0) \prod_{t=0}^{T} P(s_{t+1} \mid s_t, a_t)\pi_\theta(a_t \mid s_t) \tag{11} P(τ∣θ)=ρ0(s0)t=0∏TP(st+1∣st,at)πθ(at∣st)(11)

因此，轨迹的对数概率为：

log ⁡ P ( τ ∣ θ ) = log ⁡ ρ 0 ( s 0 ) + ∑ t = 0 T [ log ⁡ P ( s t + 1 ∣ s t , a t ) + log ⁡ π θ ( a t ∣ s t ) ] (12) \log P(\tau \mid \theta)=\log \rho_0(s_0)+\sum_{t=0}^{T} \left[\log P(s_{t+1} \mid s_t, a_t)+\log \pi_\theta(a_t \mid s_t)\right] \tag{12} logP(τ∣θ)=logρ0(s0)+t=0∑T[logP(st+1∣st,at)+logπθ(at∣st)](12)

2. 对数求导技巧

∇ θ P = P ∇ θ log ⁡ P (13) \nabla_\theta P = P \nabla_\theta \log P \tag{13} ∇θP=P∇θlogP(13)

3. 环境项与参数无关

环境相关项在参数 θ \theta θ 下是常数， ρ 0 , P ( ⋅ ∣ ⋅ ) , R ( τ ) \rho_0, P(\cdot \mid \cdot), R(\tau) ρ0,P(⋅∣⋅),R(τ) 都不依赖于策略参数，因此：

∇ θ ρ 0 = ∇ θ P = ∇ θ R ( τ ) = 0 (14) \nabla_\theta \rho_0=\nabla_\theta P=\nabla_\theta R(\tau)=0 \tag{14} ∇θρ0=∇θP=∇θR(τ)=0(14)

应用上述结论，我们可以得到：

∇ θ J ( θ ) = ∇ θ E τ ∼ π θ [ R ( τ ) ] = ∇ θ ∑ τ P ( τ ∣ θ ) R ( τ ) = ∑ τ ∇ θ P ( τ ∣ θ ) R ( τ ) = ∑ τ P ( τ ∣ θ ) ∇ θ log ⁡ P ( τ ∣ θ ) R ( τ ) (Log-derivative trick) = E τ ∼ π θ [ ∇ θ log ⁡ P ( τ ∣ θ ) R ( τ ) ] \begin{align} \nabla_\theta J(\theta) &= \nabla_\theta \mathbb{E}{\tau \sim \pi\theta}[R(\tau)] \tag{15} \\ &= \nabla_\theta \sum_{\tau} P(\tau|\theta) R(\tau) \tag{16} \\ &= \sum_{\tau} \nabla_\theta P(\tau|\theta) R(\tau) \tag{17} \\ &= \sum_{\tau} P(\tau|\theta) \nabla_\theta \log P(\tau|\theta) R(\tau) \quad \text{(Log-derivative trick)} \tag{18} \\ &= \mathbb{E}{\tau \sim \pi\theta} \left[ \nabla_\theta \log P(\tau|\theta) R(\tau) \right] \tag{19} \end{align} ∇θJ(θ)=∇θEτ∼πθ[R(τ)]=∇θτ∑P(τ∣θ)R(τ)=τ∑∇θP(τ∣θ)R(τ)=τ∑P(τ∣θ)∇θlogP(τ∣θ)R(τ)(Log-derivative trick)=Eτ∼πθ[∇θlogP(τ∣θ)R(τ)](15)(16)(17)(18)(19)

因此，将轨迹的对数概率表达式代入，并利用环境项在 θ \theta θ 下为常数的事实，就可以得到 基础策略梯度（REINFORCE）公式：

∇ θ J ( π θ ) = E τ ∼ π θ [ ∑ t = 0 T ∇ θ log ⁡ π θ ( a t ∣ s t ) R ( τ ) ] (20) \nabla_\theta J(\pi_\theta)=\mathbb{E}{\tau \sim \pi\theta} \left[\sum_{t=0}^{T}\nabla_\theta \log \pi_\theta(a_t \mid s_t)R(\tau)\right] \tag{20} ∇θJ(πθ)=Eτ∼πθ[t=0∑T∇θlogπθ(at∣st)R(τ)](20)

从直觉上看，这个梯度会：提高高回报轨迹中每个动作的对数概率，降低低回报轨迹中动作的概率。

Sample estimate of the gradient.

给定一批通过采样获得的 rollout D = { τ ( i ) } i = 1 N \mathcal{D}=\{\tau^{(i)}\}{i=1}^{N} D={τ(i)}i=1N，其中初始状态来自 s 0 ( i ) ∼ ρ 0 ( s 0 ) s_0^{(i)} \sim \rho_0(s_0) s0(i)∼ρ0(s0) 然后策略 π θ \pi\theta πθ 在环境中运行，我们可以构造一个 无偏梯度估计：

g ^ = 1 N ∑ i = 1 N ∑ t = 0 T ∇ θ log ⁡ π θ ( a t ( i ) ∣ s t ( i ) ) R ( τ ( i ) ) (21) \hat{g}= \frac{1}{N} \sum_{i=1}^{N} \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t^{(i)} \mid s_t^{(i)}) R(\tau^{(i)})\tag{21} g^=N1i=1∑Nt=0∑T∇θlogπθ(at(i)∣st(i))R(τ(i))(21)

该向量可用于梯度上升更新即 θ ← θ + α g ^ \theta \leftarrow \theta + \alpha \hat{g} θ←θ+αg^

1.5 Policy Gradient Baselines

基础策略梯度（vanilla policy gradient）的主要为问题在于其 梯度估计的方差较高 。一种常见的缓解方法是：从奖励中减去一个只依赖于状态的 基线函数（baseline） b b b。这是一种 控制变量（control variate） 技术 [Ross 2022]，其核心思想是：通过减去一个与回报 R R R 相关的项来降低估计方差，同时不会引入偏差（bias）。

我们可以定义带基线的策略梯度为：

B = E τ ∼ π θ [ ∑ t = 0 T ∇ θ log ⁡ π θ ( a t ∣ s t ) ( R ( τ ) − b ( s t ) ) ] (22) B=\mathbb{E}{\tau \sim \pi\theta} \left[ \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t \mid s_t) \big(R(\tau) - b(s_t)\big) \right] \tag{22} B=Eτ∼πθ[t=0∑T∇θlogπθ(at∣st)(R(τ)−b(st))](22)

例如，一个合理的基线选择是 在策略上的价值函数（on-policy value function） V π ( s ) = E τ ∼ π θ [ R ( τ ) ∣ s t = s ] V^\pi(s)=\mathbb{E}{\tau \sim \pi\theta}[R(\tau) \mid s_t = s] Vπ(s)=Eτ∼πθ[R(τ)∣st=s]，即从状态 s t = s s_t=s st=s 出发并按照策略 π θ \pi_\theta πθ 运行时的 期望回报 ，那么 R ( τ ) − V π ( s t ) R(\tau) - V^\pi(s_t) R(τ)−Vπ(st) 可以直观理解为当前轨迹相对于预期表现 "好多少" 或 "差多少"。

只要基线函数只依赖于状态，那么带基线的策略梯度就是 无偏的。我们可以通过重写带基线的策略梯度来说明这一点：

B = E τ ∼ π θ [ ∑ t = 0 T ∇ θ log ⁡ π θ ( a t ∣ s t ) R ( τ ) ] − E τ ∼ π θ [ ∑ t = 0 T ∇ θ log ⁡ π θ ( a t ∣ s t ) b ( s t ) ] (23) B = \mathbb{E}{\tau \sim \pi\theta} \left[ \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t \mid s_t) R(\tau) \right]-\mathbb{E}{\tau \sim \pi\theta} \left[ \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t \mid s_t) b(s_t) \right] \tag{23} B=Eτ∼πθ[t=0∑T∇θlogπθ(at∣st)R(τ)]−Eτ∼πθ[t=0∑T∇θlogπθ(at∣st)b(st)](23)

关注基线项，可以得到：

E τ ∼ π θ [ ∑ t = 0 T ∇ θ log ⁡ π θ ( a t ∣ s t ) b ( s t ) ] = ∑ t = 0 T E s t [ b ( s t ) E a t ∼ π θ ( ⋅ ∣ s t ) [ ∇ θ log ⁡ π θ ( a t ∣ s t ) ] ] (24) \mathbb{E}{\tau \sim \pi\theta} \left[ \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t \mid s_t) b(s_t) \right]= \sum_{t=0}^{T} \mathbb{E}{s_t} \left[ b(s_t) \mathbb{E}{a_t \sim \pi_\theta(\cdot \mid s_t)} \left[ \nabla_\theta \log \pi_\theta(a_t \mid s_t) \right] \right] \tag{24} Eτ∼πθ[t=0∑T∇θlogπθ(at∣st)b(st)]=t=0∑TEst[b(st)Eat∼πθ(⋅∣st)[∇θlogπθ(at∣st)]](24)

一般而言，score function 的期望为零 E x ∼ p θ [ ∇ θ log ⁡ p θ ( x ) ] = 0 \mathbb{E}{x \sim p\theta}\left[\nabla_\theta \log p_\theta(x)\right]= 0 Ex∼pθ[∇θlogpθ(x)]=0，因此，上式为零，从而可以得到：

B = E τ ∼ π θ [ ∑ t = 0 T ∇ θ log ⁡ π θ ( a t ∣ s t ) R ( τ ) ] ∇ θ J ( π θ ) (25) B = \mathbb{E}{\tau \sim \pi\theta} \left[ \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t \mid s_t) R(\tau) \right]\nabla_\theta J(\pi_\theta) \tag{25} B=Eτ∼πθ[t=0∑T∇θlogπθ(at∣st)R(τ)]∇θJ(πθ)(25)

由此可见带基线的策略梯度仍然是 无偏的，我们将在后续实验观察使用基线是否能够提升下游性能。

A note on policy gradient "losses."

在 PyTorch 等框架中实现策略梯度方法时，我们通常会定义一个所谓的策略梯度损失 pg_loss，调用 pg_loss.backward() 就会将近似策略梯度 g ^ \hat{g} g^ 写入模型参数的梯度缓冲区，在数学形式上，它可以表示为：

p g _ l o s s = 1 N ∑ i = 1 N ∑ t = 0 T log ⁡ π θ ( a t ( i ) ∣ s t ( i ) ) ( R ( τ ( i ) ) − b ( s t ( i ) ) ) (26) pg\loss = \frac{1}{N} \sum{i=1}^{N} \sum_{t=0}^{T} \log \pi_\theta(a_t^{(i)} \mid s_t^{(i)}) \big(R(\tau^{(i)}) - b(s_t^{(i)})\big) \tag{26} pg_loss=N1i=1∑Nt=0∑Tlogπθ(at(i)∣st(i))(R(τ(i))−b(st(i)))(26)

需要注意的是 pg_loss 并不是传统意义上的损失函数 ，它在训练集或验证集上报告 pg_loss 并没有实际意义，验证集上的 pg_loss 也无法反映模型是否具有良好的泛化能力。它只是一个用于反向传播的标量，使得 pg_loss.backward() 得到的梯度等价于近似策略梯度 g ^ \hat{g} g^。

在强化学习中，更重要的评估指标是 训练奖励（train rewards）和验证奖励（validation rewards），这些才是真正 "有意义" 的指标，也是策略梯度方法试图优化的目标。

1.6 Off-Policy Policy Gradient

REINFORCE 是一种 在策略（on-policy）算法：训练数据由当前正在优化的同一个策略收集得到，为什么说明这一点，我们可以将 REINFORCE 算法写为以下步骤：

1. 从当前策略 π θ \pi_\theta πθ 中采样一批轨迹（rollouts） { τ ( i ) } i = 1 N \{\tau^{(i)}\}_{i=1}^{N} {τ(i)}i=1N

2. 用这些轨迹来近似策略梯度 ∇ θ J ( π θ ) ≈ g ^ 1 N ∑ i = 1 N ∑ t = 0 T ∇ θ log ⁡ π θ ( a t ( i ) ∣ s t ( i ) ) R ( τ ( i ) ) \nabla_\theta J(\pi_\theta) \approx \hat{g}\frac{1}{N}\sum_{i=1}^{N}\sum_{t=0}^{T}\nabla_\theta \log \pi_\theta(a_t^{(i)} \mid s_t^{(i)}) R(\tau^{(i)}) ∇θJ(πθ)≈g^N1∑i=1N∑t=0T∇θlogπθ(at(i)∣st(i))R(τ(i))

3. 使用该梯度更新策略参数 θ ← θ + α g ^ \theta \leftarrow \theta + \alpha \hat{g} θ←θ+αg^

为了获得一小步梯度更新，我们需要进行大量推理来采样新的轨迹批次，然而，语言模型的行为通常不会在一次更新中发生显著变化，因此这种 on-policy 方法效率非常低。

Off-policy policy gradient.

在离策略学习中，我们使用的轨迹（rollouts）并不是来自当前正在优化的策略，而是来自 另一个策略 。像 PPO 和 GRPO 这样的主流策略梯度算法的离策略变体，会使用旧策略 π old \pi_{\text{old}} πold 生成的轨迹来优化当前策略 π θ \pi_{\theta} πθ。离策略策略梯度的估计形式为：

g ^ off-policy = 1 N ∑ i = 1 N ∑ t = 0 T π θ ( a t ( i ) ∣ s t ( i ) ) π old ( a t ( i ) ∣ s t ( i ) ) ∇ θ log ⁡ π θ ( a t ( i ) ∣ s t ( i ) ) R ( τ ( i ) ) (27) \hat{g}{\text{off-policy}}= \frac{1}{N} \sum{i=1}^{N} \sum_{t=0}^{T} \frac{\pi_\theta(a_t^{(i)} \mid s_t^{(i)})} {\pi_{\text{old}}(a_t^{(i)} \mid s_t^{(i)})} \nabla_\theta \log \pi_\theta(a_t^{(i)} \mid s_t^{(i)}) R(\tau^{(i)}) \tag{27} g^off-policy=N1i=1∑Nt=0∑Tπold(at(i)∣st(i))πθ(at(i)∣st(i))∇θlogπθ(at(i)∣st(i))R(τ(i))(27)

该形式看起来类似于对基础策略梯度（vanilla policy gradient）进行 重要性采样（importance sampling） 后的版本，其中的重加权项为 π θ ( a t ( i ) ∣ s t ( i ) ) π old ( a t ( i ) ∣ s t ( i ) ) \frac{\pi_\theta(a_t^{(i)} \mid s_t^{(i)})}{\pi_{\text{old}}(a_t^{(i)} \mid s_t^{(i)})} πold(at(i)∣st(i))πθ(at(i)∣st(i))。

实际上，公式 (27) 可以通过重要性采样推导得到，并依赖于一个合理的近似假设：当前策略 π θ \pi_\theta πθ 与旧策略 π old \pi_{\text{old}} πold 不会相差过大，更多相关细节可参考 [Degris+ 2013] 的工作。

2. Group Relative Policy Optimization

以下内容均翻译自 cs336_spring2025_assignment5_alignment.pdf，请大家查看原文档获取更详细的内容

接下来，我们将介绍 组相对策略优化（GRPO），这是一种策略梯度的变体，你将在数学问题求解任务中实现并进行实验。

2.1 GRPO Algorithm

Advantage estimation.

GRPO 的核心思想是：对每个问题，从策略 π θ \pi_\theta πθ 中采样多个输出，并利用这些输出来计算一个 baseline（基线），这样做的好处是我们无需再学习一个神经价值函数 V ϕ ( s ) V_{\phi}(s) Vϕ(s)，而该函数往往难以训练，并且在系统实现上也较为复杂。

对于某个问题 q q q 及其对应的输出组 { o ( i ) } i = 1 G ∼ π θ ( ⋅ ∣ q ) \{o^{(i)}\}{i=1}^{G} \sim \pi\theta(\cdot \mid q) {o(i)}i=1G∼πθ(⋅∣q)，设第 i i i 个输出的奖励为 r ( i ) = R ( q , o ( i ) ) r^{(i)} = R(q, o^{(i)}) r(i)=R(q,o(i))，DeepSeekMath [Shao+ 2024] 与 DeepSeek R1 [DeepSeek-AI+ 2025] 使用如下方式计算 组归一化优势（group-normalized reward）：

A ( i ) = r ( i ) − mean ( r ( 1 ) , r ( 2 ) , . . . , r ( G ) ) std ( r ( 1 ) , r ( 2 ) , . . . , r ( G ) ) + advantage_eps (28) A^{(i)}=\frac{r^{(i)} - \text{mean}(r^{(1)}, r^{(2)}, ..., r^{(G)})}{\text{std}(r^{(1)}, r^{(2)}, ..., r^{(G)})+ \text{advantage\_eps}} \tag{28} A(i)=std(r(1),r(2),...,r(G))+advantage_epsr(i)−mean(r(1),r(2),...,r(G))(28)

其中 advantage_eps 是一个小常数，用于防止除零。此外，该优势值在整个 response 的所有 token 上共享即 A t ( i ) = A ( i ) , ∀ t ∈ 1 , . . . , ∣ o ( i ) ∣ A_t^{(i)} = A^{(i)}, \quad \forall t \in {1, ..., |o^{(i)}|} At(i)=A(i),∀t∈1,...,∣o(i)∣，因此后续省略时间下标 t t t。

High-level algorithm.

在深入 GRPO 目标函数之前，我们先通过 [Shao+ 2024] 中 Algorithm 3 的算法直观理解训练流程

Note：这里采用的是 DeepSeekMath GRPO 的一个特例：使用 verified reward、不包含 KL 项、不进行参考模型或奖励模型的迭代更新。

GRPO objective.

GRPO 的目标函数融合了三种思想：

1. 离策略策略梯度（Off-policy policy gradient）（如公式 27）

2. 组归一化优势计算（如公式 28）

3. 裁剪机制（Clipping） （类似 PPO [Schulman+ 2017]）

裁剪机制的作用是在对同一批轨迹执行多次梯度更新时保持训练稳定性，同时防止策略 π θ \pi_\theta πθ 偏离旧策略过远。

我们首先写出完整的 GRPO-Clip 目标函数，然后再直观理解其中的 clipping（裁剪）机制：

J GRPO-Clip ( θ ) = E q ∼ D , { o ( i ) } i = 1 G ∼ π θ ( ⋅ ∣ q ) [ 1 G ∑ i = 1 G 1 ∣ o ( i ) ∣ ∑ t = 1 ∣ o ( i ) ∣ min ⁡ ( π θ ( o t ( i ) ∣ q , o < t ( i ) ) π θ old ( o t ( i ) ∣ q , o < t ( i ) ) A ( i ) ,    clip ( π θ ( o t ( i ) ∣ q , o < t ( i ) ) π θ old ( o t ( i ) ∣ q , o < t ( i ) ) ,    1 − ϵ ,    1 + ϵ ) A ( i ) ) ⏟ per-token objective ] J_{\text{GRPO-Clip}}(\theta)= \mathbb{E}{q \sim \mathcal{D}, \{o^{(i)}\}{i=1}^{G} \sim \pi_\theta(\cdot \mid q)} \left[ \frac{1}{G} \sum_{i=1}^{G} \frac{1}{\left| o^{(i)} \right|} \sum_{t=1}^{\left| o^{(i)} \right|} \underbrace{ \min \left( \frac{ \pi_\theta \left( o^{(i)}t \mid q, o^{(i)}{<t} \right) }{ \pi_{\theta_{\text{old}}} \left( o^{(i)}t \mid q, o^{(i)}{<t} \right) } A^{(i)}, \; \text{clip} \left( \frac{ \pi_\theta \left( o^{(i)}t \mid q, o^{(i)}{<t} \right) }{ \pi_{\theta_{\text{old}}} \left( o^{(i)}t \mid q, o^{(i)}{<t} \right) }, \; 1-\epsilon, \; 1+\epsilon \right) A^{(i)} \right) }_{\text{per-token objective}} \right] JGRPO-Clip(θ)=Eq∼D,{o(i)}i=1G∼πθ(⋅∣q) G1i=1∑G o(i) 1t=1∑∣o(i)∣per-token objective min πθold(ot(i)∣q,o<t(i))πθ(ot(i)∣q,o<t(i))A(i),clip πθold(ot(i)∣q,o<t(i))πθ(ot(i)∣q,o<t(i)),1−ϵ,1+ϵ A(i)

其中，超参数 ϵ > 0 \epsilon > 0 ϵ>0 控制策略允许变化的幅度。

为了更直观地理解这一点，我们可以按照 [Achiam 2018b] 的方式重写 per-token 目标函数，定义函数：

g ( ϵ , A ( i ) ) = { ( 1 + ϵ ) A ( i ) , A ( i ) ≥ 0 ( 1 − ϵ ) A ( i ) , A ( i ) < 0 (30) g(\epsilon, A^{(i)}) = \begin{cases} (1+\epsilon)A^{(i)}, & A^{(i)} \ge 0 \\ (1-\epsilon)A^{(i)}, & A^{(i)} < 0 \end{cases} \tag{30} g(ϵ,A(i))={(1+ϵ)A(i),(1−ϵ)A(i),A(i)≥0A(i)<0(30)

于是 per-token 目标函数可以写为：

per-token objective = min ⁡ ( π θ ( o t ( i ) ∣ q , o < t ( i ) ) π θ old ( o t ( i ) ∣ q , o < t ( i ) ) A ( i ) , g ( ϵ , A ( i ) ) ) \text{per-token objective} = \min \left( \frac{\pi_\theta(o_t^{(i)} \mid q, o_{<t}^{(i)})} {\pi_{\theta_{\text{old}}}(o_t^{(i)} \mid q, o_{<t}^{(i)})} A^{(i)}, g(\epsilon, A^{(i)}) \right) per-token objective=min(πθold(ot(i)∣q,o<t(i))πθ(ot(i)∣q,o<t(i))A(i),g(ϵ,A(i)))

接下来可以分情况理解：

情况二：优势 A ( i ) > 0 A^{(i)}>0 A(i)>0

此时目标函数简化为：

per-token objective = min ⁡ ( π θ ( o t ( i ) ∣ q , o < t ( i ) ) π θ old ( o t ( i ) ∣ q , o < t ( i ) ) , 1 + ϵ ) A ( i ) \text{per-token objective}= \min \left( \frac{\pi_\theta(o_t^{(i)} \mid q, o_{<t}^{(i)})} {\pi_{\theta_{\text{old}}}(o_t^{(i)} \mid q, o_{<t}^{(i)})}, 1+\epsilon \right) A^{(i)} per-token objective=min(πθold(ot(i)∣q,o<t(i))πθ(ot(i)∣q,o<t(i)),1+ϵ)A(i)

由于 A ( i ) > 0 A^{(i)}>0 A(i)>0，当 π θ ( o t ( i ) ∣ q , o < t ( i ) ) \pi_\theta(o_t^{(i)} \mid q, o_{<t}^{(i)}) πθ(ot(i)∣q,o<t(i)) 变大时（即该动作在新策略下更可能发生），目标函数会增加，但由于存在 min 裁剪项 ，目标函数的增长被限制，一旦 π θ ( o t ( i ) ∣ q , o < t ( i ) ) > ( 1 + ϵ ) π θ old ( o t ( i ) ∣ q , o < t ( i ) ) \pi_\theta(o_t^{(i)} \mid q, o_{<t}^{(i)})>(1+\epsilon)\pi_{\theta_{\text{old}}}(o_t^{(i)} \mid q, o_{<t}^{(i)}) πθ(ot(i)∣q,o<t(i))>(1+ϵ)πθold(ot(i)∣q,o<t(i)) 目标函数就会达到最大值 ( 1 + ϵ ) A ( i ) (1+\epsilon)A^{(i)} (1+ϵ)A(i)，因此，新策略不会被鼓励偏离旧策略太远。

情况一：优势 A ( i ) < 0 A^{(i)}<0 A(i)<0

类似地，当优势为负时，模型希望降低该动作的概率，但 clipping 会阻止其下降得过多，策略不会被鼓励使概率低于 ( 1 − ϵ ) π θ old ( o t ( i ) ∣ q , o < t ( i ) ) (1-\epsilon)\pi_{\theta_{\text{old}}}(o_t^{(i)} \mid q, o_{<t}^{(i)}) (1−ϵ)πθold(ot(i)∣q,o<t(i))（完整论证请参考 [Achiam 2018b]）

2.2 Implement

现在我们已经对 GRPO 的训练流程和目标函数有了高层理解，接下来将开始实现其中的各个组成部分。在 SFT 和 EI 部分中实现的许多组件也会在 GRPO 中复用。

Problem (compute_group_normalized_rewards): Group normalization (2 points)

Computing advantages (group-normalized rewards).

首先，我们将实现用于计算 rollout 批次中每个样本优势（advantages）的逻辑，即 组归一化奖励（group-normalized rewards）。我们将考虑两种获取组归一化奖励的方式：

1. 上文公式 (28) 中介绍的方法

2. 一种较新的简化方法

Dr. GRPO [Liu+ 2025] 指出，使用 std ( r ( 1 ) , r ( 2 ) , . . . , r ( G ) ) \text{std}(r^{(1)}, r^{(2)}, ..., r^{(G)}) std(r(1),r(2),...,r(G)) 进行归一化，可能会对那些答案正确率变化较小（方差较低）的批次产生不理想的影响，因此，他们提出去除标准差归一化步骤，仅计算：

A ( i ) = r ( i ) − mean ( r ( 1 ) , r ( 2 ) , . . . , r ( G ) ) (31) A^{(i)} =r^{(i)} - \text{mean}(r^{(1)}, r^{(2)}, ..., r^{(G)}) \tag{31} A(i)=r(i)−mean(r(1),r(2),...,r(G))(31)

我们将在本次作业中实现这两种变体，并在后续对比它们的性能表现。

Deliverable ：实现一个方法 compute_group_normalized_rewards，用于：

• 计算每个 rollout 响应的原始奖励（raw rewards）
• 在各自的组内进行归一化
• 返回归一化奖励和原始奖励
• 以及任何你认为有用的元数据（metadata）

推荐接口如下：

def compute_group_normalized_rewards(
    reward_fn,
    rollout_responses,
    repeated_ground_truths,
    group_size,
    advantage_eps,
    normalize_by_std,
):

该函数需要对每一组 rollout 响应计算奖励，并按照组大小进行归一化。

参数说明：

• reward_fn: Callable[[str, str], dict[str, float]]：用于对 rollout 响应与对应 ground truth 进行打分，返回一个包含以下键的字典："reward"、"format_reward"、"answer_reward"
• rollout_responses: list[str]：来自策略模型的 rollout 结果列表，该列表长度为 rollout_batch_size = n_prompts_per_rollout_batch x group_size
• repeated_ground_truths: list[str]：对应样本的 ground truth 列表，该列表长度等于 rollout_batch_size，因为每个样本的 ground truth 会重复 group_size 次
• group_size: int：每个问题对应的响应数量（即每组的大小）
• advantage_eps: float：用于归一化时防止除零的小常数
• normalize_by_std: bool：若为 True，则按组内标准差进行归一化；否则，仅减去组均值（不除以标准差）

返回值：tuple[torch.Tensor, torch.Tensor, dict[str, float]]

包含：

• advantages：形状为 (rollout_batch_size,)，每个 rollout 响应对应的组归一化奖励
• raw_rewards：形状为 (rollout_batch_size,)，每个 rollout 响应对应的未归一化奖励
• metadata：你选择记录的其他统计信息，例如：reward 的均值、标准差、最大值 / 最小值等

为了测试你的实现，请实现 [adapters.run_compute_group_normalized_rewards] 然后运行：

uv run pytest -k test_compute_group_normalized_rewards

确保测试通过。

Problem (compute_naive_policy_gradient_loss): Naive policy gradient (1 point)

Naive policy gradient loss.

接下来，我们将实现一些用于计算 "损失" 的方法，需要提醒 / 说明的是：这些并不是严格意义上的 "损失函数"，因此不应作为评估指标进行汇报，在强化学习中，应重点跟踪训练和验证奖励等指标（详见第 1.5 节的讨论）。

我们首先从 朴素策略梯度损失 开始，该损失形式非常简单：它将优势（advantage）与动作的对数概率相乘（并取负号）。给定一个问题 q q q、一个响应 o o o 以及响应中的某个 token o t o_t ot，朴素的 per-token 策略梯度损失为：

− A t ⋅ log ⁡ p θ ( o t ∣ q , o < t ) (32) -A_t \cdot \log p_\theta(o_t \mid q, o_{<t}) \tag{32} −At⋅logpθ(ot∣q,o<t)(32)

Deliverable ：实现一个方法 compute_naive_policy_gradient_loss，用于使用原始奖励或预先计算好的优势（advantage）来计算 逐 token 的策略梯度损失。

推荐接口如下：

def compute_naive_policy_gradient_loss(
    raw_rewards_or_advantages: torch.Tensor,
    policy_log_probs: torch.Tensor,
) -> torch.Tensor:

该函数需要在每个 token 上计算策略梯度损失，其中 raw_rewards_or_advantages 可以是原始奖励（raw reward）或已经归一化的优势（advantage）

参数说明：

• raw_rewards_or_advantages: torch.Tensor：形状 (batch_size, 1)，表示每个 rollout 响应对应的标量奖励 / 优势
• policy_log_probs: torch.Tensor：形状 (batch_size, sequence_length)，表示每个 token 的对数概率（log-probabilities）

返回值：

• torch.Tensor：形状 (batch_size, sequence_length) 表示逐 token 的策略梯度损失（在训练循环不会沿 batch 和 sequence 维度进行聚合）

Tips ：需要将 raw_rewards_or_advantages 在 sequence_length 维度上进行 broadcast

为测试你的实现，请实现 [adapters.run_compute_naive_policy_gradient_loss]，然后运行：

uv run pytest -k test_compute_naive_policy_gradient_loss

确保测试通过。

Problem (compute_grpo_clip_loss): GRPO-Clip loss (2 points)

GRPO-Clip loss.

接下来，我们将实现更为有趣的 GRPO-Clip 损失。

逐 token 的 GRPO-Clip 损失为：

− min ⁡ ( π θ ( o t ∣ q , o < t ) π θ old ( o t ∣ q , o < t ) A t , clip ( π θ ( o t ∣ q , o < t ) π θ old ( o t ∣ q , o < t ) , 1 − ϵ , 1 + ϵ ) A t ) . (33) -\min \Bigg( \frac{\pi_\theta(o_t \mid q, o_{<t})}{\pi_{\theta_{\text{old}}}(o_t \mid q, o_{<t})} A_t, \text{clip}\left( \frac{\pi_\theta(o_t \mid q, o_{<t})}{\pi_{\theta_{\text{old}}}(o_t \mid q, o_{<t})}, 1-\epsilon, 1+\epsilon \right) A_t \Bigg). \tag{33} −min(πθold(ot∣q,o<t)πθ(ot∣q,o<t)At,clip(πθold(ot∣q,o<t)πθ(ot∣q,o<t),1−ϵ,1+ϵ)At).(33)

Deliverable ：实现一个方法 compute_grpo_clip_loss，用于计算逐 token 的 GRPO-Clip 损失。

推荐接口如下：

def compute_grpo_clip_loss(
    advantages: torch.Tensor,
    policy_log_probs: torch.Tensor,
    old_log_probs: torch.Tensor,
    cliprange: float,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:

参数说明：

• advantages: torch.Tensor：形状 (batch_size, 1)，表示每个样本的优势 A A A
• policy_log_probs: torch.Tensor：形状 (batch_size, sequence_length)，表示当前正在训练的策略模型输出的逐 token 对数概率
• old_log_probs: torch.Tensor：形状 (batch_size, sequence_length)，表示旧策略模型输出的逐 token 对数概率
• cliprange: float：裁剪参数 ϵ \epsilon ϵ（例如 0.2）

返回值：tuple[torch.Tensor, dict[str, torch.Tensor]]

包括：

• loss：形状为 (batch_size, sequence_length) 的张量，表示逐 token 的裁剪损失
• metadata：一个字典，用于记录你希望日志中输出的内容。建议记录每个 token 是否被裁剪，例如：是否右侧（clipped 版本）的策略梯度损失在 min 运算中是否小于左侧（未裁剪版本）

Tips ：需要将 advantages 在 sequence_length 维度上进行 broadcast。

为测试你的实现，请实现 [adapters.run_compute_grpo_clip_loss]，然后运行：

uv run pytest -k test_compute_grpo_clip_loss

确保测试通过。

Problem (compute_policy_gradient_loss): Policy-gradient wrapper (1 point)

Policy gradient loss wrapper.

我们将进行消融实验（ablation），比较三种不同形式的策略梯度：

• (a) no_baseline ：不使用 baseline 的朴素策略梯度损失，也就是说，优势（advantage）直接等于原始奖励即 A = R ( q , o ) A=R(q,o) A=R(q,o)。
• (b) reinforce_with_baseline ：带 baseline 的朴素策略梯度损失，这里我们使用 分组归一化后的奖励 作为优势（advantage）。如果 r ˉ \bar{r} rˉ 是通过 compute_group_normalized_rewards 得到的分组归一化奖励（可能做了或未做组内标准差归一化），那么 A = r ˉ A = \bar{r} A=rˉ。
• (c) grpo_clip：GRPO-Clip 损失。

为了方便，我们将实现一个封装（warpper），使我们可以在这三种策略梯度损失之间轻松切换

Deliverable ：实现 compute_policy_gradient_loss，这是一个便捷的封装函数，用于根据指定类型分发到对应的损失函数（no_baseline、reinforce_with_baseline 或 grpo_clip），并返回逐 token 的损失以及相关辅助统计信息。

推荐接口如下：

def compute_policy_gradient_loss(
    policy_log_probs: torch.Tensor,
    loss_type: Literal["no_baseline", "reinforce_with_baseline", "grpo_clip"],
    raw_rewards: torch.Tensor | None = None,
    advantages: torch.Tensor | None = None,
    old_log_probs: torch.Tensor | None = None,
    cliprange: float | None = None,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:

该函数需要根据 loss_type 选择并计算对应的策略梯度损失。

参数说明：

• policy_log_probs：形状 (batch_size, sequence_length)，当前正在训练的策略模型输出的逐 token 对数概率
• loss_type：三者取值之一："no_baseline"、"reinforce_with_baseline"、"grpo_clip"
• raw_rewards：当 loss_type == "no_baseline" 时必需，形状 (batch_size, 1)
• advantages：当 loss_type == "reinforce_with_baseline" 或 grpo_clip 时必需，形状 (batch_size, 1)
• old_log_probs：当 loss_type == "grpo_clip" 时必需，形状 (batch_size, sequence_length)
• cliprange：当 loss_type == "grpo_clip" 时必需，标量 ϵ \epsilon ϵ，用于裁剪

返回值：tuple[torch.Tensor, dict[str, torch.Tensor]]

包括：

• loss：形状 (batch_size, sequence_length)，逐 token 的损失
• metadata：字典，包含底层损失函数返回的统计信息（例如 GRPO-Clip 中的 clip fraction）

Tips ：调用 compute_naive_policy_gradient_loss 或 compute_grpo_clip_loss 对参数进行检查（参照前面使用的 assertion 模式），将子函数返回的 metadata 汇总为一个统一字典

为测试你的实现，请实现 [adapters.run_compute_policy_gradient_loss]，然后运行：

uv run pytest -k test_compute_policy_gradient_loss

确保测试通过。

Problem (masked_mean): Masked mean (1 point)

Masked mean.

到目前为止，我们已经实现了计算优势（advantages）、对数概率（log probabilities）、逐 token 损失以及一些有用统计量（例如逐 token 熵和 clip fraction）所需的逻辑。

为了将形状为 (batch_size, sequence_length) 的逐 token 损失张量压缩为一个损失张量（每个样本一个标量），我们将在序列维度上对损失取平均，但只在对应响应部分的索引上进行（即仅在 mask[i, j] == 1 的 token 位置上计算）。

在许多基于 LLM 的强化学习代码库中，通常会按序列长度进行归一化，然而，这是否是正确做法其实并不显然 - 你可以注意到，在我们给出的策略梯度估计式（公式 21）中，并没有出现类似 1 T ( i ) \frac{1}{T^{(i)}} T(i)1 的归一化因子。我们将首先采用这种标准操作（称为 masked_mean），随后也会测试使用我们在 SFT 中实现的 masked_normalize 方法。

我们将允许用户指定在哪个维度上计算均值，如果 dim 等于 None，则对所有被 mask 的元素计算整体平均。这在计算如下指标时可能会很有用：

• 响应 token 上的平均逐 token 熵
• clip fraction 等统计量

Deliverable ：实现一个方法 masked_mean，用于在考虑布尔掩码的情况下对张量元素进行平均。

推荐接口如下：

def masked_mean(
    tensor: torch.Tensor,
    mask: torch.Tensor,
    dim: int | None = None,
) -> torch.Tensor:

该函数需要在指定维度上计算张量的均值，仅对 mask == 1 的元素进行平均。

参数说明：

• tensor: torch.Tensor：需要进行平均计算的数据张量
• mask: torch.Tensor：与 tensor 形状相同，值为 1 的位置会参与均值计算
• dim: int | None：指定在哪个维度上进行平均，如果为 None，则对所有被 mask 的元素整体计算均值

返回值：

• torch.Tensor：带掩码的均值结果，其形状遵循 tensor.mean(dim) 的语义规则

为测试你的实现，请实现 [adapters.run_masked_mean]，然后运行：

uv run pytest -k test_masked_mean

确保测试通过。

Problem (grpo_microbatch_train_step): Microbatch train step (3 points)

GRPO microbatch train step.

现在，我们已经准备好实现 GRPO 的单次微批（microbatch）训练步骤了，请记住，对于一个完整的训练 mini-batch，如果 gradient_accumulation_steps > 1，我们会迭代多个 microbatch。

具体来说，在给定原始奖励或优势（advantages）以及对数概率（log probs）的情况下，我们将：

1. 计算逐 token 的损失；

2. 使用 masked_mean 将其聚合为每个样本一个标量损失；

3. 在 batch 维度上取平均；

4. 根据梯度累积步数进行调整；

5. 最后执行反向传播。

Deliverable：实现一次 GRPO 的单个 micro-batch 更新步骤，包括：

• 策略梯度损失计算
• 使用掩码进行平均
• 梯度缩放（gradient scaling）

推荐接口如下：

def grpo_microbatch_train_step(
    policy_log_probs: torch.Tensor,
    response_mask: torch.Tensor,
    gradient_accumulation_steps: int,
    loss_type: Literal["no_baseline", "reinforce_with_baseline", "grpo_clip"],
    raw_rewards: torch.Tensor | None = None,
    advantages: torch.Tensor | None = None,
    old_log_probs: torch.Tensor | None = None,
    cliprange: float | None = None,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:

该函数需要完成在一个 microbatch 上执行一次前向与反向传播。

参数说明：

• policy_log_probs：形状 (batch_size, sequence_length)，当前训练策略模型输出的逐 token 对数概率
• response_mask：形状 (batch_size, sequence_length)，响应 token 位置为 1，prompt 或 padding 位置为 0
• gradient_accumulation_steps：每次优化器更新所包含的 microbatch 数量
• loss_type：三者取值之一："no_baseline"、"reinforce_with_baseline"、"grpo_clip"
• raw_rewards：当 loss_type == "no_baseline" 时必需，形状 (batch_size, 1)
• advantages：当 loss_type != "no_baseline" 时必需，形状 (batch_size, 1)
• old_log_probs：当使用 GRPO-Clip 时必需，形状 (batch_size, sequence_length)
• cliprange：GRPO-Clip 中的裁剪参数 ϵ \epsilon ϵ

返回值：tuple[torch.Tensor, dict[str, torch.Tensor]]

包括：

• loss：标量张量，表示经过梯度累积调整后的 microbatch 损失，我们返回该值以便记录日志
• metadata：字典，包含底层损失函数返回的统计信息，以便其他你希望记录的指标

Tips ：需要在该函数中调用 loss.backward()，确保根据 gradient_accumulation_steps 对损失进行适当缩放

为测试你的实现，请实现 [adapters.run_grpo_microbatch_train_step]，然后运行：

uv run pytest -k test_grpo_microbatch_train_step

确保测试通过。

Problem (grpo_train_loop): GRPO train loop (5 points)

Putting it all together: GRPO train loop.

现在，我们将整个各个模块，构建一个完整的 GRPO 训练循环，整体结构可以参考第 2.1 节中的算法流程，并在适当位置使用我们已经实现的方法。

下面提供了一组示例超参数，如果你的实现是正确的，使用这些超参数应该能够得到合理的训练结果：

n_grpo_steps: int = 200
learning_rate: float = 1e-5

advantage_eps: float = 1e-6
rollout_batch_size: int = 256
group_size: int = 8
sampling_temperature: float = 1.0
sampling_min_tokens: int = 4  # As in Expiter, disallow empty string responses
sampling_max_tokens: int = 1024
epochs_per_rollout_batch: int = 1  # On-policy
train_batch_size: int = 256  # On-policy
gradient_accumulation_steps: int = 128  # microbatch size is 2, will fit on H100
gpu_memory_utilization: float = 0.85
loss_type: Literal[
    "no_baseline",
    "reinforce_with_baseline",
    "grpo_clip",
] = "reinforce_with_baseline"
use_std_normalization: bool = True

optimizer = torch.optim.AdamW(
    policy.parameters(),
    lr=learning_rate,
    weight_decay=0.0,
    betas=(0.9, 0.95),
)

这些默认超参数会使你从 on-policy 设置开始训练 - 对于每一个 rollout batch，我们只执行一次梯度更新，从超参数的角度来看，这意味着 train_batch_size 等于 rollout_batch_size，epochs_per_rollout_batch 等于 1

下面是一些用于 sanity check 的断言和常量设置，可以帮助你避免一些边界情况，并指引你朝正确方向实现：

assert train_batch_size % gradient_accumulation_steps == 0, (
    "train_batch_size must be divisible by gradient_accumulation_steps"
)

micro_train_batch_size = train_batch_size // gradient_accumulation_steps

assert rollout_batch_size % group_size == 0, (
    "rollout_batch_size must be divisible by group_size"
)

n_prompts_per_rollout_batch = rollout_batch_size // group_size

assert train_batch_size >= group_size, (
    "train_batch_size must be greater than or equal to group_size"
)

n_microbatches_per_rollout_batch = rollout_batch_size // micro_train_batch_size

另外还有一些额外建议：

• 记得使用 r1_zero prompt，并像之前实验一样，让 vLLM 在第二个答案标签 处停止生成。
• 建议使用 typer 进行参数解析。
• 使用梯度裁剪（clip value 设为 1.0）。
• 应定期记录验证奖励（例如每 5 或 10 步一次）。建议至少在 1024 个验证样本上进行评估，以便比较不同超参数（因为 CoT/RL 评估可能存在较大噪声）。
• 在我们当前的损失实现中，GRPO-Clip 仅应在 off-policy 情况下使用（因为它需要旧策略的对数概率）。
• 在 off-policy 设置中，如果每个 rollout batch 进行多轮（multiple epochs）梯度更新，那么每一轮都重新计算旧对数概率会比较浪费。更合理的做法是：只计算一次旧对数概率，然后在后续 epoch 中复用。
• 不要对旧对数概率（old log-probabilities）进行梯度计算。
• 每次优化器更新时，应记录以下部分或全部指标：
  • 损失（loss）
  • 梯度范数（gradient norm）
  • Token 熵（token entropy）
  • Clip fraction（若为 off-policy）
  • 训练奖励（总奖励、格式奖励、答案奖励）
  • 以及任何你认为对调试有帮助的其他指标。

Deliverable：实现一个完整的 GRPO 训练循环。开始在 MATH 数据集上训练策略模型，并确认验证奖励（validation rewards）随着训练步数增加而提升；随着时间推移，生成的 rollout 结果呈现出合理变化。

此外，请提供：一张展示验证奖励随训练步数变化的曲线图以及若干随时间变化的示例 rollout 输出。

3. GRPO Experiments

以下内容均翻译自 cs336_spring2025_assignment5_alignment.pdf，请大家查看原文档获取更详细的内容

现在，我们可以基于 GRPO 训练循环开始进行实验，尝试不同的超参数和算法改动，每个实验需要使用 2 块 GPU：一块用于 vLLM 实例，另一块用于策略模型。

Note on stopping runs early.

如果在 200 个 GRPO step 之前，你已经看到不同超参数之间存在显著差异（例如某个配置发散，或明显效果较差），可以提前停止该实验，以节省时间和计算资源用于后续实验，下面提到的 GPU 时长仅为粗略估计。

Problem (grpo_learning_rate): Tune the learning rate (2 points) (6 H100 hrs)

从上面给出的建议超参数开始，对学习率进行 sweep，并报告最终的验证答案奖励（如果优化器发散，请注明）。

Deliverable：

• 不同学习率对应的验证奖励曲线；
• 一个在 MATH 上达到至少 25% 验证准确率的模型；
• 用 2 句话简要讨论你在日志指标中观察到的其他趋势。

在后续实验中，可以使用本次 sweep 中表现最好的学习率。

Problem (grpo_baselines): Effect of baselining (2 points) (2 H100 hrs)

Effect of baselines.

在上述超参数设置（使用调好的学习率）下，我们将研究 baseline 的影响，当前处于 on-policy 设置，因此我们将比较以下损失类型：

• no_baseline
• reinforce_with_baseline

注意：默认超参数中 use_std_normalization = True。

请分别使用 reinforce_with_baseline 和 no_baseline 训练策略模型。

Deliverable：

• 每种损失类型对应的验证奖励曲线；
• 用 2 句话简要讨论你在其他日志指标中观察到的趋势。

在接下来的几个实验中，应使用上述实验中表现最好的损失类型。

Problem (think_about_length_normalization): Think about length normalization (1 point)

Length normalization.

正如我们在实现 masked_mean 时提到的那样，在序列长度维度上对损失取平均并不是必要的，甚至未必是正确的。如何对损失进行求和，是一个重要的超参数选择，因为它会影响策略对各个动作的 credit attribution（归因分配）。

我们参考 [Lambert 2024] 中的一个示例来说明这一点：

在 GRPO 的训练步骤中，我们首先会计算逐 token 的策略梯度损失（暂时忽略 clipping）：

advantages  # (batch_size, 1)
per_token_probability_ratios  # (batch_size, sequence_length)
per_token_loss = -advantages * per_token_probability_ratios

这里，我们已经将优势（advantages）在序列长度维度上进行了 broadcast。

接下来，我们比较两种对这些逐 token 损失进行聚合的方法：

• 使用我们实现的 masked_mean，即在每个序列中仅对未被 mask 的 token 取平均；
• 对每个序列中未被 mask 的 token 进行求和，然后除以一个固定的常数标量，这一方式可以通过我们实现的 masked_normalize 方法支持，例如设置 constant_normalizer = 1.0，参考 [Liu+ 2025] [Yu+ 2025]

我们考虑一个示例：假设 batch size 为 2，第一个响应有 4 个 token，第二个响应有 7 个 token，通过这个例子，我们可以观察不同归一化方式如何影响梯度：

from your_utils import masked_mean, masked_normalize

ratio = torch.tensor([
    [1, 1, 1, 1, 1, 1, 1],
    [1, 1, 1, 1, 1, 1, 1],
], requires_grad=True)

advs = torch.tensor([
    [2, 2, 2, 2, 2, 2, 2],
    [2, 2, 2, 2, 2, 2, 2],
])

masks = torch.tensor([
    # generation 1: 4 tokens
    [1, 1, 1, 1, 0, 0, 0],
    # generation 2: 7 tokens
    [1, 1, 1, 1, 1, 1, 1],
])

# Normalize with each approach
max_gen_len = 7
masked_mean_result = masked_mean(ratio * advs, masks, dim=1)
masked_normalize_result = masked_normalize(
    ratio * advs, masks, dim=1, constant_normalizer=max_gen_len
)

print("masked_mean", masked_mean_result)
print("masked_normalize", masked_normalize_result)

# masked_mean tensor([2., 2.], grad_fn=<DivBackward0>)
# masked_normalize tensor([1.1429, 2.0000], grad_fn=<DivBackward0>)

masked_mean_result.mean().backward()
print("ratio.grad", ratio.grad)

# ratio.grad:
# tensor([[0.2500, 0.2500, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000],
#         [0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429]])

ratio.grad.zero_()

masked_normalize_result.mean().backward()
print("ratio.grad", ratio.grad)

# ratio.grad:
# tensor([[0.1429, 0.1429, 0.1429, 0.1429, 0.0000, 0.0000, 0.0000],
#         [0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429]])

Deliverable：在不运行实验的前提下，对这两种方法进行比较。

• 各自的优缺点是什么？
• 是否存在在某些特定场景或示例，使其中一种方法更优？

接下来，我们将从实验角度比较 masked_mean 和 masked_normalize。

Problem (grpo_length_normalization): Effect of length normalization (2 points) (2 H100 hrs)

Deliverable ：在端到端 GRPO 训练中，分别使用 masked_mean 和 masked_normalize 进行归一化，并进行比较。

• 报告验证答案奖励曲线；
• 对实验结果进行分析，包括是否有其他指标（例如梯度范数）表现出明显趋势。

Hint：可以关注与训练稳定性相关的指标，例如梯度范数（gradient norm）

在后续实验中，请固定使用表现更好的长度归一化方法。

Problem (grpo_group_standard_deviation): Effect of standard deviation normalization (2 points) (2 H100 hrs)

Normalization with group standard deviation.

回顾我们在 compute_group_normalized_rewards 的标准实现（基于 [Shao+ 2024] [DeepSeek-AI+ 2025]），我们会按照组内标准差进行归一化。

[Liu+ 2025] 指出，用组内标准差进行除法归一化，可能会在训练过程中引入不希望的偏置，例如：

• 标准差较低的问题（例如过于简单或过于困难的问题，其所有奖励几乎都是 1 或 0），
• 在训练中会被赋予更高的权重。

为了解决这一问题，[Liu+ 2025] 提出去掉按标准差归一化这一步骤，我们已经在 compute_group_normalized_rewards 中实现了这一变体，现在将进行测试。

Deliverable：比较以下两种设置的表现：

• use_std_normalization == True
• use_std_normalization == False

请：

• 报告验证答案奖励曲线；
• 对实验结果进行分析，包括是否有其他指标表现出明显趋势。

Hint：可以关注与训练稳定性相关的指标，例如梯度范数（gradient norm）。

在后续实验中，请固定使用表现更好的组归一化方式。

Problem (grpo_off_policy): Implement off-policy GRPO

Off-policy versus on-policy.

到目前为止，我们尝试的所有超参数设置都是 on-policy 的：每个 rollout batch 只进行一次梯度更新，因此我们几乎是在使用对策略梯度的 "原则性" 估计 g ^ \hat{g} g^（除了前面提到的长度归一化和优势归一化之外）。

虽然这种方法在理论上是合理且稳定的，但它效率较低。Rollout 需要策略模型进行较慢的生成过程，因此是 GRPO 的主要计算成本；而每个 rollout batch 只进行一次梯度更新，可能不足以显著改变策略的行为，因此显得比较浪费。

因此，接下来我们将尝试 off-policy 训练，即在每个 rollout batch 上进行多次梯度更新（甚至多个 epoch）。

Deliverable：实现 off-policy 的 GRPO 训练。

根据你在上面实现的完整 GRPO 训练循环，你可能已经具备了相关基础设施，如果没有，则需要实现以下内容：

• 允许在每个 rollout batch 上进行多个 epoch 的梯度更新，其中：
  • 每个 rollout batch 的 epoch 数由 epochs_per_rollout_batch 控制；
  • 每个 rollout batch 内的优化器更新次数由 rollout_batch_size 和 train_batch_size 控制。
• 修改主训练循环：在每个 rollout batch 生成阶段之后、在进入内部梯度更新循环之前，先计算策略的响应 log-probs，并将其保存为 old_log_probs。建议使用 torch.inference_mode()
• 使用 "GRPO-Clip" 作为损失类型。

Problem (grpo_off_policy_sweep): Off-policy GRPO hyperparameter sweep (4 points)(12 H100 hrs)

现在，我们可以通过：

• 每个 rollout batch 的 epoch 数
• 每个 rollout batch 的优化器更新次数

来控制训练的 off-policy 程度。

Deliverable ：固定 rollout_batch_size = 256，选择一组 epochs_per_rollout_batch 和 train_batch_size 的取值范围进行超参数扫描。

建议流程：

1. 先进行一次 粗略扫描（GRPO 步数较少，例如 <50），以大致了解性能表现的分布情况；

2. 然后进行一次 更精细的扫描（例如 200 个 GRPO step）。

请提供一段简要的实验说明，解释你选择扫描范围的原因。

将结果与 on-policy 的设置进行比较：

• epochs_per_rollout_batch = 1
• train_batch_size = 256

并分别给出以验证步数（number of validation steps）为横轴的曲线图以及以 wall-clock time 为横轴的曲线图。

请报告验证答案奖励曲线，并对实验结果进行分析，包括：

• 是否有其他指标表现出明显趋势（例如熵 entropy、响应长度 response length）；
• 比较模型在训练过程中熵的变化，与之前 EI 实验中观察到的情况进行对比。

Hint ：为了保持显存使用量不变，你需要调整 gradient_accumulation_steps。

Problem (grpo_off_policy_clip_ablation): Off-policy GRPO-Clip ablation (2 points) (2 H100 hrs)

Ablating clipping in the off-policy setting.

回顾 GRPO-Clip 中 clipping 的目的：当在单个 rollout batch 上进行多次梯度更新时，防止策略偏离旧策略过远。

接下来，我们将在 off-policy 设置中对 clipping 进行消融实验，以测试它在多大程度上是必要的。换句话说，我们将使用如下逐 token 损失：

− π θ ( o t ∣ q , o < t ) π θ old ( o t ∣ q , o < t ) A t -\frac{\pi_\theta(o_t \mid q, o_{<t})}{\pi_{\theta_{\text{old}}}(o_t \mid q, o_{<t})} A_t −πθold(ot∣q,o<t)πθ(ot∣q,o<t)At

Deliverable ：实现未裁剪（unclipped）的逐 token 损失，作为一种新的损失类型 "GRPO-No-Clip"。

使用你在上一题中找到的表现最佳的 off-policy 超参数配置，运行该未裁剪版本的损失，请：

• 报告验证答案奖励曲线；
• 对比 GRPO-Clip 的结果进行分析；
• 评论实验发现，包括是否有其他指标呈现明显趋势，例如：
  • 熵（entropy）
  • 响应长度（response length）
  • 梯度范数（gradient norm）

Problem (grpo_prompt_ablation): Prompt ablation (2 points) (2 H100 hrs)

Effect of prompt.

作为最后一个消融实验，我们将研究一个颇为出人意料的现象：在强化学习（RL）过程中所使用的 prompt，可能会对模型性能产生显著影响，这种影响往往取决于模型的预训练方式。

我们不再使用 cs336_alignment/prompts/r1_zero.prompt 中的 R1-Zero prompt，而是改为使用一个非常简单的 prompt，位于 cs336_alignment/prompts/question_only.prompt

内容为：

{question}

你需要在训练和验证阶段都使用该 prompt，并将奖励函数（训练和验证均使用）更改为 question_only_reward_fn，其定义位于 cs336_alignment/drgrpo_grader.py

Deliverable：报告以下两种 prompt 的验证答案奖励曲线：

• R1-Zero prompt
• question-only prompt

请比较它们的指标表现，包括是否有其他指标呈现明显趋势，例如：

• 熵（entropy）
• 响应长度（response length）
• 梯度范数（gradient norm）

并尝试解释你观察到的现象。

4. Leaderboard: GRPO on MATH

以下内容均翻译自 cs336_spring2025_assignment5_alignment.pdf，请大家查看原文档获取更详细的内容

作为本次（必做）作业的最后一部分，你需要在 2 块 H100 GPU 上、4 小时训练时间内，探索各种方法，以获得尽可能高的验证奖励。

Model.

继续使用 Qwen 2.5 Math 1.5B Base 模型。

Dataset.

继续使用集群中提供的 MATH 训练与验证数据集：

• 训练集：/data/a5-alignment/MATH/train.jsonl
• 验证集：/data/a5-alignment/MATH/validation.jsonl

限制条件如下：

• 不允许使用任何其他数据；
• 不允许使用更强模型生成的推理链（reasoning chains）进行 SFT；
• 必须在 完整验证集（全部 5K 个样本）上报告验证准确率；
• 验证时必须使用给定的采样超参数：
  • temperature = 1.0
  • max_tokens = 1024
• 可以自行对训练集进行过滤，或设计课程学习（curriculum）。
• 验证阶段必须使用 R1-Zero prompt；
• 验证阶段必须严格使用 starter code 中提供的 r1_zero_reward_fn 奖励函数（但训练阶段可以自行设计奖励函数）。

Algorithm.

你可以自由调节超参数，甚至完全改变训练算法，只要：

• 不使用任何额外数据；
• 不使用其他模型（可以使用同一模型的多个副本）。

Systems optimizations.

你可能会注意到，在我们上面给出的简易 GRPO 实现中，至少有一块 GPU 始终处于空闲状态。通过改进系统层面的实现，你很可能会获得显著性能提升。

例如，你可以考虑：

• 在 rollout 或训练中使用低精度（mixed precision / lower precision）
• 使用 torch.compile
• 以及其他系统优化手段

此外，你完全不必拘泥于将 vLLM 放在一张 GPU 上、将训练策略放在另一张 GPU 上的简单分工方式。我们鼓励你思考更优的并行化方案。

Ideas.

如果你希望参考一些可能的优化方向，可以查看以下仓库：

• veRL
• trl
• torchtune
• oat

On KL divergence.

我们还需要指出，在前面的实验中，我们并没有加入相对于某个参考模型的 KL 散度项（通常这个参考模型是一个冻结的 SFT 模型或预训练 checkpoint）。

在我们的实验以及部分文献结果（[Liu+ 2025] [NTT123 2025]）中发现：

• 去掉 KL 项对性能没有明显影响
• 同时还能节省 GPU 显存（因为不需要存储参考模型）

不过，许多 GRPO 相关仓库默认都会加入 KL 项。我们鼓励你在实验中尝试加入 KL 或其他形式的正则化，只要你使用的是 Qwen 2.5 Math 1.5B Base，或者是通过你自己的算法训练得到的模型即可。

Problem (leaderboard): Leaderboard (16 points) (16 H100 hrs)

Deliverable ：请报告在 2 块 H100 GPU 上、4 小时训练时间内 获得的验证准确率，并提交一张 验证准确率随实际运行时间（wall-clock time）变化的截图（横轴需截至 ≤ 4 小时）。

提醒：我们对评估过程有如下约束：

1. 验证准确率 必须是在整个 MATH 验证集（全部 5000 个样本）上的平均准确率。

2. 验证阶段必须使用 R1-Zero prompt。

3. 评估时必须使用 vLLM，并设置：

• temperature = 1.0
• max_tokens = 1024

4. 验证准确率必须通过对 starter code 中提供的 r1_zero_reward_fn 奖励函数输出的答案奖励进行平均计算得到。

5. Epilogue

恭喜你完成了本课程的最后一次作业！你理应为自己的努力感到自豪。

我们希望你通过从零开始构建语言模型的核心组件，真正理解了现代语言模型背后的基础原理，并在这个过程中收获了成长与乐趣。

结语

本篇文章系统梳理并拆解了 CS336 Assignment 5 中关于策略梯度与 GRPO 的核心内容。从语言模型作为随机策略的基本视角出发，我们推导了 REINFORCE 的策略梯度公式，引入 baseline 以降低方差，并进一步理解 on-policy 与 off-policy 之间的效率差异。在此基础上，我们深入分析了 GRPO 的设计思想---通过组归一化优势替代显式价值函数、通过 clipping 保持策略更新稳定，并将这些思想具体落实为可运行的训练循环。

下篇文章我们就来一起看看 GRPO 具体该如何实现，敬请期待🤗

参考

• https://github.com/stanford-cs336/assignment5-alignment