A5 · Alignment and Reasoning RL（让模型学会做数学题）
年份 2026（Spring 2026）—— 往年（2025 / 2024）题面不同，见「作业说明」
分值 未记录 ✗（官方主页没有公布各项作业的比重 ✓ 不推断 ✗）
发布 / 截止 5 月 20 日发布 / 6 月 3 日截止（L19）（官方日程标注 tentative / 暂定 ✓）
算力 训练需 GPU（SFT 与 RLVR 都要真跑 ✓）
配哪几讲 候选（机器按术语匹配，供参考）：第15讲 §RLHF 算法：从 PPO 到 DPO；第15讲 §SFT 的方法：其实很无聊；第16讲 §三、核心算法（下）：GRPO；第16讲 §二、核心算法（上）：PPO —— 通吃但难驯；第17讲 §2. 回到 2021：CLIP 的对比学习；第17讲 §3. SigLIP：把损失函数改得更简单、更高效
代码仓库 https://github.com/stanford-cs336/assignment5-alignment
官方题面 PDF（本地原件在 _local/A5/）

官方原文：Apply supervised finetuning and reinforcement learning to train LMs to reason when solving math problems. Optional Part 2: implement and apply safety alignment methods such as DPO.

做什么

题面来自官方材料（本地原件在 _local/，官方链接见「材料」）。 编号是原件的编号（notebook 是 cell 序号、PDF 是页码），方便跟官方题面对号。 本项目不提供答案（D53）。

第 1 页

CS336 Assignment 5 (alignment): Reasoning RL

Version 26.0.0

CS336 Staff

Spring 2026

1 Assignment Overview（作业概览）

在本作业中，你将获得一些实践经验：训练一个 language model 进行推理并解决下游任务。

What you will implement（你将实现什么）

  1. Zero-shot、few-shot 和 chain-of-thought prompting。
  2. Group Relative Policy Optimization (GRPO)，一种在给定外部奖励的情况下提升模型性能的 reinforcement learning 算法。
  3. Policy gradient estimation 的变体，用来探索 variance reduction 和 importance weight clipping 策略。

What you will run（你将运行什么）

  1. 在 GSM8K 上测量 OLMo-2-0425-1B 的 prompting 性能。
  2. 在 OLMo-2-0425-1B 上运行 on-policy GRPO，以提升其在 GSM8K 上的性能。
  3. 运行 RL 变体，包括 RFT、Dr. GRPO 和 MaxRL，以探索 RL 中的算法选择。
  4. 运行 off-policy GRPO，以加速训练并探索各种 clipping 策略。

What the code looks like（代码结构）

所有作业代码以及这份 writeup 都可在 GitHub 上获取：github.com/stanford-cs336/assignment5-alignment。请 git clone 该仓库。如果有任何更新，我们会通知你，你可以 git pull 获取最新版本。

  1. cs336_alignment/*：你将在这里编写 assignment 5 的代码。注意，这里没有代码（除了下面描述的 starter code），因此你可以从零开始按自己的想法实现。
  2. cs336_alignment/vllm_utils.py：该文件包含运行 vLLM server 的代码，该 server 可生成 generations 并同步 weights。
  3. cs336_alignment/drgrpo_grader.py：该文件包含对数学问题的模型输出进行评分的代码。
  4. cs336_alignment/prompts/*：为了方便你，我们提供了包含 prompts 的文本文件。
  5. tests/*.py：该目录包含必需的 GRPO 测试以及可选的 alignment/safety supplement 测试。

tests/test_grpo.py 中必需的测试会调用 tests/adapters.py 中定义的 hooks。你将实现 adapters，把你的代码连接到测试。编写更多测试和/或修改测试代码可能有助于调试你的代码，但你的实现应通过原先提供的 test suite。

  1. README.md：该文件包含一些关于设置 environment 的基本说明。

How to submit（如何提交）

你将把以下文件提交到 Gradescope：

第 2 页

2 Introduction（引言）

2.1 Background（背景）

在本课程的前四次作业中，我们学习了如何预训练一个 base model。现在，我们准备学习 post-training：一旦有了 base model，如何把它变成一个能解决下游任务的有用工具？

post-training 的一个组成部分是 alignment。Pre-training 由于数据混合和训练目标，会产生这样的效果：赋予模型广泛的知识和行为。但当我们要求模型解决任务时，我们希望得到一种特定的行为，即一个 helpful and harmless 的 assistant。把一个广泛的 base model 变成专注的 chat model 的过程称为“instruction tuning”或“alignment”，我们将在可选的 Assignment 5 supplement 中介绍这些技术。

post-training 的另一个组成部分是 reinforcement learning。在 pre-training 中，因为我们的目标是赋予模型广泛的知识基础，所以我们使用广泛的数据（网络上的文本）。在 post-training RL 中，我们的目标变窄了：我们希望模型在特定任务上达到高准确率，比如解决数学问题。这个更窄的目标至少在两个方面不同于 pre-training：(a) 我们不再有那么多数据；(b) 我们的训练目标从“coverage”（拥有广泛的知识基础）变为“precision”（产生准确的回答）。由于这些差异，我们需要一种新技术，即 reinforcement learning (RL)。

在 RL 中，我们给定一个包含问题的 dataset，以及一个 scoring function，用来判断某个 response 是否正确解决了该问题。例如，在 coding 场景中，我们可能会给定一个 coding problem（例如“编写一个反转列表的 Python 函数”），以及一个由一组 test cases 组成的 scoring function（例如 assert f([0, 1, 2]) == [2, 1, 0]）。或者在数学场景中，我们可能会给定一个数学推理 word problem（例如“如果 Tom 四年前的年龄是 John 的一半，而 John 现在二十岁，Tom 现在多大？”），以及一个 scoring function，它解析最终答案并检查其是否正确（12）。值得注意的是，与 pretraining 不同，我们不会得到一个用于模仿的 responses dataset：在 coding 例子中，我们不会得到一个能解决该任务的正确 Python 程序；在数学推理例子中，我们不会得到一条能产生答案的正确推理链。相反，在 RL 中，我们直接以模型准确率作为目标函数来取 gradient steps。在高层次上，RL 会涉及从模型采样 responses，用 scoring function 给它们打分，然后对正确的 responses 加权提升。

在本作业中，我们将从数学和实证两方面学习 RL。事实证明，RL 很难，因为它很慢且不稳定。做 RL 的科学也很难，因为 RL 运行在不同 random seeds 之间具有高方差，而且实现中看似很小的细节也很重要。本作业将介绍 LLM RL，并探索其中的一些挑战。

2.2 Model and dataset（模型与数据集）

在本作业中，我们将使用 OLMo-2-0425-1B，这是一个 base model，先在 OLMo-mix-1124 上 pre-trained，再在 Dolmino-mix-1124 上 mid-trained，总共使用 4 trillion tokens。OLMo-mix-1124 主要由 DCLM-Baseline [J. Li et al., 2024] 组成，这是我们在课堂上讲过的一个 dataset；而 Dolmino-mix-1124 是一个更聚焦的数据混合，包含大约 50% DCLM 和 50% 的 instruction following、math、code、STEM paper 以及 wiki data。其他 OLMo-2-0425-1B 的细节可以在 OLMo 2 技术报告 [T. OLMo et al., 2024] 中找到。

第 3 页

在课程的这个阶段，你应该已经具备理解他们每个设计决策所需的全部背景知识。

作为我们的下游任务，我们将使用 GSM8K dataset [K. Cobbe et al., 2021]，它可在作业仓库中的 data/gsm8k/train.jsonl 和 data/gsm8k/test.jsonl 中找到，也可在线获取。该 dataset 包含一组相对简单的小学数学推理 word problems。下面是一个例子：

{
  "question": "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?",
  "answer": "Natalia sold 48/2 = <<48/2=24>>24 clips in May.\nNatalia sold 48+24 = <<48+24=72>>72 clips altogether in April and May.\n#### 72"
}

在 RL 期间，我们的模型将学习为这种风格的问题生成推理链，从而提高其解决数学问题的能力。

关于模型和数据集选择的说明：由于资源限制，我们只选择了一个小规模的 model-dataset 组合作为我们的 RL testbed。具体来说，我们选择了一个相当有能力的小模型，因为它是在非常大量的 tokens 上训练的（是参数数量的 4000 倍！），这使我们能够在现实 dataset 上以较小规模看到合理的 RL 结果。不幸的是，我们在本课程中训练的模型太弱，无法解决数学问题。如果你在完成本作业后感到好奇，也欢迎拿你在本课程中训练的模型，在一个更简单的任务上运行 RL。由于 RL 训练动态非常依赖模型和 dataset，你应该记住，我们在本作业中看到的结果可能无法迁移到其他 model-dataset 组合。

2.3 Notation（符号）

本作业会涉及一些数学，因此在下表中我们提供一些稍后会使用的 notation。Language modeling 和 reinforcement learning 经常对同一个对象使用不同术语，所以表中包含两套术语。如果你忘了某个 notation 指的是什么，可以随时回来看这张表。

Notation LM terminology RL terminology 含义
$\rho$ prompt distribution or dataset initial state distribution 对 prompts/problems 的分布。
$x$ prompt; question initial state 从 $\rho$ 中采样得到的 prompt/problem。
$y$ response; completion; generation; sample rollout; trajectory; sampled action sequence 对 prompt $x$ 的一个采样答案。
$y_t$ token action $y$ 中第 $t$ 个生成的 token。
$y_{ prefix 位置 $t$ 之前的所有生成 tokens，或 $y_1, \ldots, y_{t-1}$。
$\pi_\theta$ model policy 参数为 $\theta$ 的模型，它对给定 prompts $x$ 的 responses $y$ 赋予概率 $\pi_\theta(y \mid x)$。

第 4 页

符号 名称 说明
$\pi_\theta(y_t \mid x, y_{ next-token distribution / timestep $t$ 的 policy 给定 prompt 与之前生成的 tokens 时，token $y_t$ 的条件概率。
$r(y \mid x)$ reward（奖励） 表示采样得到的 response 正确性的标量分数（本作业中为 0 或 1）。
$B$ 每个 batch 的 prompt 数量 每个 inference batch 中的 prompt 数量。
$G$ 每个 prompt 的 generation 数量 / group size 为每个 prompt 采样的 response 数量。
$\text{len}(y)$ 或 $L$ response 长度 / horizon 一个 response 中生成的 token 数量。
$A(i,j)$ advantage（优势） 经过 baseline 和 normalization 后，分配给 prompt $i$ 的 response $j$ 的权重。

3 Prompting（提示）

我们将预训练 base model 应用到下游任务的第一步，是给它提供 prompt（提示）。Base models 在广泛的行为上经过预训练，而 prompting 是一种轻量级方式，可将其行为塑造成任务求解模式。在作业后面，我们还会看到 prompt 的选择会塑造 RL 动态与探索。

最基本的 prompting 方式是为模型提供问题，并从其 next token distribution 中采样以生成答案；我们将此策略称为 question_only。我们会将该策略与 r1_zero prompt 比较；后者同时包含问题以及让模型进行 chain-of-thought reasoning 的指令 [DeepSeek-AI et al., 2025]。

3.1 Using vLLM for inference（使用 vLLM 进行推理）

为了从模型生成 responses，我们需要一个 inference engine。实现 inference engine 超出本作业范围，因此我们将使用 vLLM inference engine [W. Kwon et al., 2023]；它实现了多种优化，包括快速 CUDA kernels、用于高效 attention KV caching 的 PagedAttention 等。用于启动 vLLM server 并产生 generations 的代码已提供在 cs336_alignment/vllm_utils.py 中，其接口如下：

@dataclass
class VLLMCompletion:
    text: str
    token_ids: list[int]
    finish_reason: str | None

@dataclass
class VLLMServer:
    model_id: str
    gpu: int = 0
    seed: int = 0
    gpu_memory_utilization: float = 0.9

第 5 页

    def start(self) -> None: ...
    def generate_completions(
        self,
        prompts: list[str],
        sampling_params: dict,
        batch_size: int | None = None,
    ) -> list[VLLMCompletion]: ...

3.2 Zero-shot、few-shot 与 chain-of-thought prompting

除非另有说明，在 GSM8K 上的实验我们将使用来自 DeepSeek R1-Zero 模型的以下 prompt [DeepSeek-AI et al., 2025]。我们称之为 r1_zero prompt：

A conversation between User and Assistant. The User asks a question, and the Assistant solves it. The Assistant first thinks about the reasoning process in the mind and then provides the User with the answer. The reasoning process is enclosed within <think> </think> and the answer is enclosed within <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think> <answer> answer here </answer>.
User: {question}
Assistant: <think>

r1_zero prompt 位于文本文件 cs336_alignment/prompts/r1_zero.prompt 中。

在 prompt 中，question 指我们插入的某个问题（例如，Natalia 在 4 月向 48 位朋友卖出了发夹，然后在 5 月卖出的数量是 4 月的一半。Natalia 在 4 月和 5 月总共卖出了多少个发夹？）。预期是模型扮演 assistant 的角色，开始生成思考过程（因为我们已经包含了一个开头的 <think> 标签），用 </think> 结束思考过程，然后在 answer 标签内生成最终符号答案，例如 <answer> 4x + 10 </answer>。让模型生成 <answer> </answer> 这类标签的目的，是让我们可以轻松解析模型的输出并与 ground truth answer 比较，并且在看到结束的 answer 标签 </answer> 时可以停止 response generation。

另一种 prompting 方法称为 few-shot prompting：在向模型提供实际问题之前，先在前面加上几个 question-answer 对。r1_zero prompt 的 few-shot 版本如下：

A conversation between User and Assistant. The User asks a question, and the Assistant solves it. The Assistant first thinks about the reasoning process in the mind and then provides the User with the answer. The reasoning process is enclosed within <think> </think> and the answer is enclosed within <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think> <answer> answer here </answer>.
User: {question-1}
Assistant: <think> {reasoning-1} </think> <answer> {answer-1} </answer>
User: {question-2}
Assistant: <think> {reasoning-2} </think> <answer> {answer-2} </answer>
User: {question-3}
Assistant: <think> {reasoning-3} </think> <answer> {answer-3} </answer>
User: {question}
Assistant: <think>

第 6 页

Few-shot prompting 通过给模型几个它要解决的任务示例来提升模型性能。开放 base model 的已报告 benchmark 数字常常是 few-shot 的：例如，OLMo-2-0425-1B 的 model card 上为 GSM8K 报告的结果使用了 8-shot prompting。我们在 cs336_alignment/prompts/r1_zero_three_shot_gsm8k.prompt 中包含了一个 3-shot 版本的 r1_zero prompt，其中的示例取自 OLMES repository。

最后，作为 baseline，我们将包含 question_only prompt，位于 cs336_alignment/prompts/question_only.prompt：

{question} Please put your final answer within \\boxed{{}}.

注意，虽然该 prompt 名为 question_only，但我们要求将最终答案放在 boxed 中，以便 grader 更容易从 response 中解析答案，如下一节所述。

3.3 Grading function（评分函数）

一旦模型生成 response，我们需要检查它是否正确。我们的数学题包含 ground truth answers（例如 0.5），但模型可以用许多方式给出正确答案：例如，它可以回答 <answer> 1/2 </answer> 或 The answer is 0.5。因此，为了正确评分模型输出，我们需要一个 answer parsing function，它接收模型输出和已知的 ground-truth 作为输入，并返回一个 boolean，指示模型输出是否正确。

在我们的实验中，我们将使用近期 reasoning RL 工作中使用的一个快速且相当准确的 answer parser [Z. Liu et al., 2025]。对于 r1_zero prompts，该 reward function 实现在 cs336_alignment.drgrpo_grader.r1_zero_reward_fn。question_only prompt 不要求模型使用 <think> 和 <answer> 标签，因此应改用同一 grader 文件中的 cs336_alignment.drgrpo_grader.question_only_reward_fn 进行评估。注意，question_only parser 会查找 \boxed{} 中的最终答案。这些 reward functions 返回一个 total reward，以及单独的 format reward 和 answer reward，分别指示 response 是否遵循预期格式，以及解析出的答案是否正确。一些工作会对格式正确但答案错误的 responses 给予 partial credit（非零 total reward）。在我们的实验中，我们不会给予 partial credit，因此 total reward 将只是 answer reward，我们只会将 format reward 用于 logging。

注意，这些 reward functions 期望 ground_truth 参数只包含答案。为了从 GSM8K 数据集的 ground truth responses 中提取答案，这些 responses 包含 {rationale} #### {answer}，你应当按 #### 分割并去除首尾空白。

3.4 Experiments（实验）

现在我们可以评估 base model 的 prompting 性能了。

Generation hyperparameters. 生成 responses 时，我们将使用 temperature 1.0、top-p 1.0、max generation length 512 进行采样。r1_zero prompts 要求模型以字符串 </answer> 结束其答案，因此在使用这些 prompts 时，我们可以让 vLLM 在模型输出该字符串时停止：

# Based on Dr. GRPO: stop when the model completes its answer
# https://github.com/sail-sg/understand-r1-zero/blob/
#   c18804602b85da9e88b4aeeb6c43e2f08c594fbc/train_zero_math.py#L167
sampling_params['stop'] = ["</answer>"]
sampling_params['include_stop_str_in_output'] = True

仅对 r1_zero 和 r1_zero_three_shot prompts 使用这个 stop string，不要用于 question_only prompt。

第 7 页

Problem (prompting_baselines)：在 GSM8K 上运行 OLMo-2-0425-1B（5 分）

(a) 编写一个脚本，用 zero-shot 的 question_only、zero-shot 的 r1_zero 和 few-shot 的 r1_zero_three_shot prompt 评估 OLMo-2-0425-1B 在 GSM8K 上的表现。然后运行脚本并观察输出。对每个 prompt，有多少条模型生成分别落入以下类别：(1) 答案正确，且格式奖励和正确性奖励都为 1；(2) 格式奖励为 1、正确性奖励为 0；(3) 格式奖励为 0、正确性奖励为 0？观察至少十个类别 2 的样例，其中有多少模型输出实际上是正确的，只是没有被正确解析？类别 3 呢？

Deliverable: 几句评述、评估指标，以及若干 prompt 和 response 示例。

(b) 观察模型输出，描述模型在每个 prompt 下的行为特征。例如，如果我们希望模型回答问题，仅仅提供问题是否足够，还是模型会表现出除了回答问题之外的其他行为？zero-shot 的 r1_zero 和 few-shot 的 r1_zero_three_shot prompt 如何塑造模型的行为？

Deliverable: 几句带支撑示例的评述。

4 Group Relative Policy Optimization（组相对策略优化）

现在我们已经测量了模型仅靠 prompting 的表现，下一步将通过训练提升这一表现。具体来说，我们希望优化模型的准确率，或等价地，优化其期望奖励：

$$J_\theta = \mathbb{E}_{x\sim\rho}\mathbb{E}_{y\sim\pi_\theta(y\mid x)}[r(y\mid x)], \tag{1}$$

其中 $\rho$ 是我们对 prompt/问题 $x$ 的任务分布，$\pi_\theta$ 是我们的模型，$y$ 是针对 $x$ 采样得到的 response/解答，$r(y\mid x)$ 表示 $y$ 是否为 $x$ 的正确答案（我们也会把 $r$ 称为“reward function”）。注意，这个目标与我们在课上到目前为止学习的预训练目标相当不同：与从数据集 $\mathcal{D}$ 中采样得到的 cross entropy loss $\mathbb{E}_{x\sim\mathcal{D}}[-\log \pi_\theta(x)]$ 不同，我们的准确率目标涉及对来自模型自身样本的期望。因此，我们需要新工具来优化这个目标：特别地，我们需要使用 reinforcement learning (RL)。从高层看，RL 会涉及从模型采样 response、用评分函数给它们打分，并强化其中正确的那些。RL 通常很慢且具有挑战性，所以本作业的很大一部分将探索让它更快、更稳定的工具。

RL 的另一个方面是：由于它涉及从模型采样，其动态会受到采样时所使用 prompt 的影响。特别地，语言模型研究中一个令人兴奋的发现是：当使用 chain-of-thought 推理 prompt 时，在强 base model 上执行 RL 可以显著提升其推理能力 [OpenAI et al., 2024; DeepSeek-AI et al., 2025]。因此，在本作业中，除了探索 RL 的核心算法思想，我们还会探索它与 chain-of-thought prompting 等语言建模技术的交互。

4.1 推导 on-policy GRPO

我们将首先探索 RL 背后的数学，逐步构建到 Group Relative Policy Optimization (GRPO) 算法——训练语言模型时常用的一种 RL 算法 [Z. Shao et al., 2024]。

第 8 页

我们的讲解紧密基于几份很好的资料，它们更深入地讲解了这些概念：OpenAI 的 Spinning Up in Deep RL [J. Achiam, 2018] 和 Nathan Lambert 的 Reinforcement Learning from Human Feedback (RLHF) Book [N. Lambert, 2024]。

4.1.1 Language models as policies

传统上，在 reinforcement learning 中，我们考虑一个 policy：它接收状态 $s_t$ 并输出动作 $a_t$。这个动作随后带来某个奖励 $r_t = r(s_t, a_t)$，以及从某个 next state distribution 中抽取的后续状态 $s_{t+1} \sim \text{next\_state}(s_t, a_t)$。然后，reinforcement learning 涉及优化该 policy 以获得高奖励。

一个参数为 $\theta$ 的 causal language model (LM) $\pi_\theta$ 定义了在给定当前文本前缀 $y_{

$$a_t \sim \pi_\theta(\cdot\mid s_t), \quad \pi_\theta(a_t\mid s_t) = [\mathrm{softmax}(f_\theta(s_t))]_{a_t}. \tag{2}$$

用 policy gradients 优化 policy 时，我们需要两个基本操作： 1. 从 policy 采样：从上述 categorical distribution 中抽取一个动作 $a_t$； 2. 给动作的 log-likelihood 打分：计算 $\log \pi_\theta(a_t\mid s_t)$。

一般来说，在用 LLM 做 RL 时，$s_t$ 是到目前为止产生的部分完成/解答，每个 $a_t$ 是解答的下一个 token；当发出 end-of-text token 时，episode 结束，例如 <|end_of_text|>，或者在我们的 r1_zero prompt 情形下是 </answer>。

4.1.2 Trajectories

在 reinforcement learning 中，我们从一个初始状态开始，该状态从某个起始分布 $s_0 \sim \rho$ 中抽取。然后我们从 policy $\pi_\theta$ 中采样一个动作 $a_t$，根据某个 next state distribution $s_{t+1} \sim \text{next\_state}(s_t, a_t)$ 转移到下一个状态，并重复。这一系列状态和动作随后形成一条（finite-horizon）trajectory：

$$\tau = (s_0, a_0, s_1, a_1, \ldots, s_T, a_T), \tag{3}$$

其中 $T$ 是 trajectory 的长度，即 $a_T$ 是一个 end-of-text token，或者我们已经达到最大 token 生成预算。

在我们的情形中，$s_0 \sim \rho$ 是我们的 prompt $x$，其中包含我们要解决的数学题（可能应用了 CoT prompting 或 few-shot）。从这个前缀开始，我们随后逐个采样 token，下一个状态就是旧前缀与已发出 token 的拼接，即 $s_{t+1} = (s_t, a_t)$。Trajectories 也称为 episodes 或 rollouts；我们会互换使用这些术语。

4.1.3 Rewards and return

在 RL 中，标量奖励 $r_t = r(s_t, a_t)$ 评判在状态 $s_t$ 处所采取动作的即时质量。对于可验证领域上的 RL，比如我们的数学应用题任务，中间推理步骤的奖励为零，直到我们输出答案。然后我们在终止动作上收到一个经过验证的奖励

$$r_T = r(s_T, a_T) := \begin{cases}

1 & \text{if trajectory } \tau \text{ matches the ground-truth according to our reward function}\ 0 & \text{otherwise}

\end{cases}. \tag{4}$$

在我们的情形中，如果最终答案正确，$r_T$ 为 1，否则为 0。Return $R(\tau)$ 聚合沿 trajectory 的奖励，而在我们的情形中它就是 $r_T$。

第 9 页

Agent（智能体）的目标是最大化期望 return

$$J_\theta = \mathbb{E}_{\tau\sim\pi_\theta}[r(\tau)], \tag{5}$$

其中 $\tau\sim\pi_\theta$ 是 trajectory 上的分布：我们先采样 $s_0 \sim \rho$，然后从 policy 采样动作、从 next state distribution 采样下一个状态。在我们的情形中，这意味着先从数据集中采样数学问题 $x\sim\rho$，然后采样 token 直到到达 response 的末尾；这个过程可以记为 $y\sim\pi_\theta(y\mid x)$。这个目标引出优化问题

$$\theta^* = \arg\max_\theta J_\theta. \tag{6}$$

4.1.4 Policy gradients

我们现在已经具备学习一种名为 policy gradients 的 RL 算法所需的记号与定义。在作业剩余部分，我们将坚持使用语言模型记号 $(x, y)$ 来表示 prompt 和 response，而不是 RL 中状态 $s_t$ 和动作 $a_t$ 的记号。

回想一下，我们希望优化模型的期望奖励（即准确率），其由下式给出

$$J_\theta = \mathbb{E}_{x\sim\rho}\mathbb{E}_{y\sim\pi_\theta(y\mid x)}[r(y\mid x)], \tag{7}$$

其中 $r(y\mid x)$ 表示 response $y$ 对 prompt $x$ 是否正确。我们的方法是对目标做梯度上升，即

$$\theta_{k+1} = \theta_k + \alpha \nabla_\theta J_{\theta_k}, \tag{8}$$

但要做到这一点，我们需要能够从样本中估计梯度 $\nabla_\theta J_{\theta_k}$。把这个梯度改写为对样本的期望，会得到一个称为“REINFORCE policy gradient”的表达式，名字来自 REINFORCE 论文 [R. J. Williams, 1992]。特别地，我们会看到可以如下改写梯度：

$$\nabla_\theta \mathbb{E}_{x\sim\rho}\mathbb{E}_{y\sim\pi_\theta(y\mid x)}[r(y\mid x)] = \mathbb{E}_{x\sim\rho}\mathbb{E}_{y\sim\pi_\theta(y\mid x)}[r(y\mid x)\nabla_\theta \log \pi_\theta(y\mid x)]. \tag{9}$$

在走完推导之前，请注意这个表达式暗示了一个直观的算法（REINFORCE）：我们可以从数据集中采样问题 $x$，从模型中采样 response $y$，在这些样本上计算 $r(y\mid x)\nabla_\theta \log \pi_\theta(y\mid x)$ 的平均值，然后走一步梯度。表达式 $r(y\mid x)\nabla_\theta \log \pi_\theta(y\mid x)$ 说的是，我们应该提高高奖励 $r(y\mid x)$ 的 response $y$ 的 log probability。如果我们重复这个过程并走多步梯度，就得到一个基本的 reinforcement learning 训练循环。

为了推导 REINFORCE policy gradient，我们可以使用 log-derivative trick（回想 $\log f(x)$ 的导数是 $\frac{f'(x)}{f(x)}$）：

$$\nabla_\theta \pi_\theta(y\mid x) = \pi_\theta(y\mid x)\nabla_\theta \log \pi_\theta(y\mid x). \tag{10}$$

直接应用这个恒等式就得到所需结果：

$$\nabla_\theta \mathbb{E}_{y\sim\pi_\theta(y\mid x)}[r(y\mid x)] = \sum_y \nabla_\theta \pi_\theta(y\mid x) r(y\mid x) \tag{11}$$

$$= \sum_y \pi_\theta(y\mid x)\nabla_\theta \log \pi_\theta(y\mid x) r(y\mid x) \tag{12}$$

$$= \mathbb{E}_{y\sim\pi_\theta(y\mid x)}[r(y\mid x)\nabla_\theta \log \pi_\theta(y\mid x)]. \tag{13}$$

注意，在这些等式中，我们省略了对 prompt 的期望 $\mathbb{E}_{x\sim\rho}$，因为 $\nabla_\theta$ 与 $\mathbb{E}_{x\sim\rho}$ 可以交换。

第 10 页

因此，给定一批 prompt $x^{(1)}, \ldots, x^{(B)} \sim^{\text{iid}} \rho$ 以及每个 prompt 的 $G$ 个 response $y^{(i,j)} \sim^{\text{iid}} \pi_\theta(y \mid x^{(i)})$，我们得到 policy gradient 的如下样本估计：

$$ \hat{g} \leftarrow \frac{1}{BG}\sum_{i=1}^{B}\sum_{j=1}^{G} r(y^{(i,j)} \mid x^{(i)}) \nabla_\theta \log \pi_\theta(y^{(i,j)} \mid x^{(i)}), \tag{14} $$

它会提高高 reward 的 response 的权重。为了推导 GRPO，我们将从这个基本的 policy gradient estimator 出发，然后逐一应用修改。

4.1.5 Baselines

虽然基本的 REINFORCE policy gradient estimator 的期望 $\nabla_\theta J_\theta$ 是正确的，但它的 variance 很高，使 RL 不稳定。为了稳定训练并降低 variance，人们使用的一个工具是 baselines。虽然 baselines 可能变得相当复杂，但最简单的 baseline 是从 reward 中减去一个常数 $b$：

$$ \mathbb{E}_{x\sim\rho}\mathbb{E}_{y\sim\pi_\theta(y\mid x)}\left[(r(y\mid x)-b)\nabla_\theta \log \pi_\theta(y\mid x)\right]. \tag{15} $$

例如，如果我们减去 $b=0.5$，并且 reward $r$ 是二值的（正确为 1，错误为 0），那么使用这个 estimator 意味着我们从模型采样 response，提高正确 response 的权重（权重为 0.5），并降低错误 response 的权重（权重为 $-0.5$）。这与我们原来的 gradient estimator 相反：原来的 estimator 以权重 1 提高正确 response，而对错误 response 不做任何事（权重为 0）。

本节的关键数学事实告诉我们，减去 baseline $b$ 会保持 estimator 的期望不变，只要 $b$ 不依赖于 action/response $y$（例如它可以是常数，或 $x$ 的函数，或 policy 的函数）。只要 $b$ 不依赖于 $y$，在 estimator 中减去它就会保持其期望不变：

$$ \mathbb{E}_{x\sim\rho}\mathbb{E}_{y\sim\pi_\theta(y\mid x)}\left[(r(y\mid x)-b)\nabla_\theta\log\pi_\theta(y\mid x)\right] = \mathbb{E}_{x\sim\rho}\mathbb{E}_{y\sim\pi_\theta(y\mid x)}\left[r(y\mid x)\nabla_\theta\log\pi_\theta(y\mid x)\right]. \tag{16} $$

这个事实直接来自以下恒等式：

$$ \mathbb{E}_{y\sim\pi_\theta(y\mid x)}\left[\nabla_\theta\log\pi_\theta(y\mid x)\right] = \nabla_\theta \mathbb{E}_{y\sim\pi_\theta(y\mid x)}[1] = 0, \tag{17} $$

其中第一步来自反向应用我们之前的 log-derivative trick。

虽然 baselines 保持 estimator 的期望不变，但它们可能增加或降低 variance。用记号来说，“variance”指的是以下内容：把 policy gradient estimator 写成样本均值 $\frac{1}{n}\sum_{i=1}^{n} Z_i$，其中每个 $Z_i$ 是某个样本 $(x,y)$ 的表达式 $(r(y\mid x)-b)\nabla_\theta \log \pi_\theta(y\mid x)$，那么 estimator 的 variance 是 $\frac{\mathbb{E}[Z_i^2]-\mathbb{E}[Z_i]^2}{n}$。我们将在下面的问题中探讨 baselines 何时增加或降低 variance。

Problem (baseline_calcs)：计算 policy gradient estimator 的 variance（5 分）

令 $\pi_\theta$ 定义二元 action 空间 $\mathcal{A}=\{0,1\}$ 上的 policy，且 $\pi_\theta(A=1)=p=\sigma(\theta)$，其中 $\sigma$ 表示 sigmoid function $\sigma(\theta)=\frac{1}{1+e^{-\theta}}$。我们的二元 reward function 会给 action $A=1$ 分配 reward 1，否则分配 reward 0，或者写成 $r(A)=\mathbf{1}\{A=1\}$。

(a) 令 policy gradient estimator 由下式给出

$$ \frac{1}{n}\sum_{i=1}^{n} r(A_i)\nabla_\theta\log\pi_\theta(A_i) \tag{18} $$

第 11 页

给定 $n$ 个样本 $A_i \sim^{\text{iid}} \pi_\theta$。这个 estimator 的 variance 是多少？

Deliverable：一个用 $n$ 和 $p$ 表示的表达式，并附上推导。

(b) 令 baseline-adjusted policy gradient estimator 由下式给出

$$ \frac{1}{n}\sum_{i=1}^{n} (r(A_i)-b)\nabla_\theta\log\pi_\theta(A_i) \tag{19} $$

给定 $n$ 个样本 $A_i \sim^{\text{iid}} \pi_\theta$。这个 estimator 的 variance 是多少？

Deliverable：一个用 $n$、$b$ 和 $p$ 表示的表达式，并附上推导和一些讨论。

(c) 如果我们代入“population mean” baseline $b=p$，得到的 variance 是多少？将这个 variance 与未调整的 policy gradient estimator 的 variance 进行比较：它总是更低、总是更高，还是有时更高有时更低，取决于 $p$？

Baselines in GRPO

GRPO 的第一个组件将是应用“group mean” baseline：

$$ \hat{g}\leftarrow \frac{1}{BG}\sum_{i=1}^{B}\sum_{j=1}^{G}(r(y^{(i,j)}\mid x^{(i)})-\mu_i)\nabla_\theta\log\pi_\theta(y^{(i,j)}\mid x^{(i)}) \tag{20} $$

其中 $\mu_i=\frac{1}{G}\sum_{j=1}^{G} r(y^{(i,j)}\mid x^{(i)})$，即模型在 prompt $x^{(i)}$ 上的平均 reward。注意，即使 $\mu_i$ 依赖于 response $y$，它仍然保持期望不变，最多相差一个缩放因子 $\frac{G-1}{G}$：

$$ \begin{aligned} &\mathbb{E}_{y^{(1)},\ldots,y^{(G)}\sim\pi_\theta(y\mid x)}\left[\frac{1}{G}\sum_{j=1}^{G}(r(y^{(j)}\mid x)-\mu_i)\nabla_\theta\log\pi_\theta(y^{(j)}\mid x)\right] && (21)\\ &= \frac{1}{G}\sum_{j=1}^{G}\mathbb{E}_{y^{(2)},\ldots,y^{(G)}\sim\pi_\theta(y\mid x)}\left[\mathbb{E}_{y^{(1)}\sim\pi_\theta(y\mid x)}\left[(r(y^{(1)}\mid x)-\mu_i)\nabla_\theta\log\pi_\theta(y^{(1)}\mid x)\right]\right] && (22)\\ &= \frac{1}{G}\sum_{j=1}^{G}\mathbb{E}_{y^{(1)}\sim\pi_\theta(y\mid x)}\left[\left(r(y^{(1)}\mid x)-\frac{1}{G}r(y^{(1)}\mid x)\right)\nabla_\theta\log\pi_\theta(y^{(1)}\mid x)\right] && (23)\\ &= \frac{G-1}{G}\mathbb{E}_{y\sim\pi_\theta(y\mid x)}\left[r(y\mid x)\nabla_\theta\log\pi_\theta(y\mid x)\right]. && (24) \end{aligned} $$

这些经过 group mean 调整的 reward 通常称为 advantages：这一项不是绝对 reward，而是表示该 rollout 相对于平均 reward 的“advantage”。

4.1.6 Advantage normalization

在减去 group mean 之后，GRPO 的下一个修改是除以 group standard deviation：

$$ \hat{g}\leftarrow \frac{1}{BG}\sum_{i=1}^{B}\sum_{j=1}^{G}\frac{r(y^{(i,j)}\mid x^{(i)})-\mu_i}{\text{std}_i}\nabla_\theta\log\pi_\theta(y^{(i,j)}\mid x^{(i)}), \tag{25} $$

其中 $\text{std}_i=\sqrt{\frac{1}{G}\sum_{j=1}^{G}(r(y^{(i,j)}\mid x^{(i)})-\mu_i)^2}$（注意，默认的 torch.std 函数使用略有不同的 sample std 表达式来做 bias correction；在你的实现中应使用默认的 torch.std）。除以 standard deviation 不再保持 estimator 的期望，因此我们不再是在原始目标 $J_\theta$（expected reward）上做 gradient ascent。理解这个修改的一种方式是把它看作一个稳定性技巧：如果我们假设每个 gradient vector 都是 iid Gaussian，

第 12 页

除以 group std 能确保每个 group 的 gradient update 在范数上大致相等。我们将在下一节更深入地讨论 advantage normalization。

4.1.7 Sequence normalization

GRPO 中还有一个细节似乎继承自早期的 LLM-PPO 实现，称为 sequence normalization。回忆一下，我们的 policy gradient estimator（到目前为止）如下（这里我们把 $\log \pi_\theta(y^{(i,j)}\mid x^{(i)})$ 展开成对时间步的求和）：

$$ \hat{g}\leftarrow \frac{1}{BG}\sum_{i=1}^{B}\sum_{j=1}^{G}\frac{r(y^{(i,j)}\mid x^{(i)})-\mu_i}{\text{std}_i}\nabla_\theta\sum_{t=1}^{\text{len}(y^{(i,j)})}\log\pi_\theta(y_t^{(i,j)}\mid x^{(i)}, y_{

GRPO 增加了一个额外的 sequence length normalization 项：

$$ \hat{g}\leftarrow \frac{1}{BG}\sum_{i=1}^{B}\sum_{j=1}^{G}\frac{r(y^{(i,j)}\mid x^{(i)})-\mu_i}{\text{std}_i}\nabla_\theta\left(\frac{1}{\text{len}(y^{(i,j)})}\sum_{t=1}^{\text{len}(y^{(i,j)})}\log\pi_\theta(y_t^{(i,j)}\mid x^{(i)}, y_{

注意，这个修改进一步改变了期望：不再是在所有 sequence 中同等地加权所有 token，较长 sequence 中的 token 相对于较短 sequence 中的 token 被降权。下一节我们会讨论这个选择，并做 ablations 来看它是否是一个好的修改。但就目前而言，我们只会实现原始论文中给出的标准 GRPO 算法，所以我们会包含 sequence normalization。

4.1.8 Putting everything together

GRPO 中还有一个组件我们尚未讨论，即 clipped importance reweighting。我们将在作业后面讨论 off-policy RL 时讨论这个组件。简而言之，off-policy RL 指的是通过在每个 inference batch 上采取多个 gradient step 来加速 RL 的方法，与之相对的是 on-policy RL，即每个 inference batch 只采取一个 gradient step。“off-policy”一词指的是：从第二个 gradient step 开始，inference batch 中的样本现在是 stale 的，因为它们来自模型的一个更旧版本。Importance reweighting 则是对 gradient estimator 的一种修改，使其能够使用 stale 样本。目前我们只会做 on-policy RL，所以还不需要 importance reweighting。

我们现在已经拥有实现 on-policy GRPO 所需的所有组件。算法见 Algorithm 1。

第 13 页

Algorithm 1: On-policy Group Relative Policy Optimization (GRPO)

Input 初始 policy model $\pi_{\theta_0}$；reward function $r$；task distribution（或 dataset）$\rho$；learning rate $\alpha$

Output $\pi_\theta$

1 policy model $\pi_\theta \leftarrow \pi_{\theta_0}$

2 对于 step = 1, …, n_grpo_steps

3 采样一批问题 $x^{(1)},\ldots,x^{(B)} \sim^{\text{iid}} \rho$

4 每个问题采样 $G$ 个输出：$y^{(i,1)},\ldots,y^{(i,G)} \sim^{\text{iid}} \pi_\theta(y\mid x^{(i)})$

5 计算 on-policy GRPO policy gradient estimator

$$ \hat{g}\leftarrow \frac{1}{BG}\sum_{i=1}^{B}\sum_{j=1}^{G}\frac{1}{\text{len}(y^{(i,j)})}\sum_{t=1}^{\text{len}(y^{(i,j)})}\frac{r(y^{(i,j)}\mid x^{(i)})-\mu_i}{\text{std}_i}\nabla_\theta\log\pi_\theta(y_t^{(i,j)}\mid x^{(i)}, y_{

其中

$$ \mu_i=\frac{1}{G}\sum_{j=1}^{G}r(y^{(i,j)}\mid x^{(i)}) \tag{29} $$

表示 group mean，

$$ \text{std}_i=\sqrt{\frac{1}{G}\sum_{j=1}^{G}(r(y^{(i,j)}\mid x^{(i)})-\mu_i)^2} \tag{30} $$

表示 group standard deviation。

6 用你最喜欢的 optimizer 更新 $\theta$（我们在下面给出 stochastic gradient ascent update）：

$$ \theta \leftarrow \theta + \alpha \hat{g}. \tag{31} $$

4.2 Implementing on-policy GRPO

4.2.1 Using Hugging Face models

虽然之前的作业使用我们在 cs336_basics 中自己的 language model 实现，但在这次作业中，我们将直接使用 Hugging Face transformers 库来加载预训练的 base model。如果你愿意，也欢迎实现你自己的 transformer 和预训练权重加载器；只需确保 architecture 与 OLMo-2-0425-1B 匹配。

要加载一个 Hugging Face model 和 tokenizer（使用 bfloat16，并用 FlashAttention-2 节省内存），你可以使用以下 starter code，它位于 cs336_alignment/checkpoint.py：

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def get_model_and_tokenizer(model_id_or_dir: str, device: str):
    model = AutoModelForCausalLM.from_pretrained(
        model_id_or_dir,
        device_map=device,
        torch_dtype=torch.bfloat16,
        attn_implementation="eager" if device=='cpu' else "flash_attention_2",
    )

第 14 页

tokenizer = AutoTokenizer.from_pretrained(model_id_or_dir)
    return model, tokenizer

model_id_or_directory 可以是一个模型名，例如 allenai/OLMo-2-0425-1B，也可以是某个目录的路径。这个目录通常会是调用 save_pretrained 的结果；你可以调用它来保存训练好的模型：

# Save the model weights
model.save_pretrained(save_directory=output_dir)
tokenizer.save_pretrained(save_directory=output_dir)

首先，让我们实现一个 helper function，它使用预训练的 tokenizer 来 tokenize 输入 prompts 和 responses。分别 tokenize prompt 和 response，不添加 special tokens，然后直接把它们拼接为 prompt_ids + response_ids，二者之间没有 EOS、BOS 或分隔 token。我们还需要构造一个 response_mask：一个 boolean mask，与 shifted labels 对齐，当且仅当对应的 label token 来自 response 时为 True，对于 prompt 和 padding 位置为 False。我们将在 training loop 中使用这个 mask，以确保只在 response tokens 上计算 loss。

Problem (tokenize_prompt_and_output)：Prompt 与 output 的 tokenization（1 分）

Deliverable: 实现一个方法 tokenize_prompt_and_output，它分别 tokenize prompt 和 output，在不插入 special tokens 的情况下把它们拼接起来，并构造一个 response_mask。推荐使用以下接口：

def tokenize_prompt_and_output(
    prompt_strs: list[str], 
    output_strs: list[str], 
    tokenizer: PreTrainedTokenizer,
) -> dict[str, torch.Tensor]:

Tokenize prompt 和 output 字符串，并构造一个与 labels 对齐的 mask：response tokens 为 1，其他 tokens（prompt 或 padding）为 0。

Args: - prompt_strs: list[str] prompt 字符串列表。 - output_strs: list[str] output 字符串列表。 - tokenizer: PreTrainedTokenizer 用于 tokenization 的 tokenizer。

Returns: - dict[str, torch.Tensor]。令 prompt_and_output_lens 为一个列表，包含拼接后的 tokenized prompt 和 output 字符串的长度。那么返回的 dictionary 应包含以下 keys： - input_ids：torch.Tensor，shape 为 (batch_size, max(prompt_and_output_lens) - 1)：tokenized 的 prompt 和 output 字符串，并切掉最后一个 token。 - labels：torch.Tensor，shape 为 (batch_size, max(prompt_and_output_lens) - 1)：shifted input ids，即去掉第一个 token 后的 input ids。

第 15 页

要测试你的代码，实现 [adapters.run_tokenize_prompt_and_output]。然后运行以下测试

uv run pytest -k test_tokenize_prompt_and_output

并确保你的实现通过测试。

一旦我们有了 tokenized input，就可以按如下方式把它传给 model：

input_ids = train_batch["input_ids"].to(device)
labels = train_batch["labels"].to(device)
logits = model(input_ids).logits

有了这个语法，请实现下面的方法，它计算 response 在 model 下的 per-token log-probabilities；这是计算 policy gradient 所需的一个基础原语。在 RL 中，记录 per-token entropies 通常也很有用，因此这个函数还会包含一个 return_token_entropy 选项，你也需要实现它。

Problem (get_response_log_probs)：Response log-probs（与 entropy）（1 分）

Deliverable: 实现一个方法 get_response_log_probs，它从 causal language model 获取 per-token conditional log-probabilities（给定之前的 tokens），并可选地返回 model 的 next-token distribution 的 entropy。

推荐使用以下接口：

def get_response_log_probs(
    model: PreTrainedModel,
    input_ids: torch.Tensor,
    labels: torch.Tensor,
    return_token_entropy: bool = False,
) -> dict[str, torch.Tensor]:

Args: - model: PreTrainedModel 用于 scoring 的 HuggingFace model（放在正确的 device 上；如果不应计算 gradients，则置于 inference mode）。 - input_ids: torch.Tensor shape (batch_size, sequence_length)，由你的 tokenization 方法产生的拼接后的 prompt + response tokens。 - labels: torch.Tensor shape (batch_size, sequence_length)，由你的 tokenization 方法产生的 labels。 - return_token_entropy: bool 如果为 True，还返回 per-token entropy。

Returns: - dict[str, torch.Tensor]。 - "log_probs"：shape (batch_size, sequence_length)，conditional log-probabilities $\log p_\theta(x_t \mid x_{"token_entropy"：可选，shape (batch_size, sequence_length)，每个位置的 per-token entropy（仅当 return_token_entropy=True 时存在）。

第 16 页

要测试你的代码，实现 [adapters.run_get_response_log_probs]。然后运行

uv run pytest -k test_get_response_log_probs

并确保测试通过。

4.2.2 在 reinforcement learning loop 中使用 vLLM

在我们的 RL loop 中，还需要生成 rollouts。我们将使用的具体配置是：把一个 Hugging Face model 和 optimizer 放在一块 GPU 上用于 training，把 vLLM（其中包含 model 和 kv cache）放在另一块 GPU 上。因此，除了上一节描述的 vLLM initialization 和 generation 函数之外，我们还需要在每次 inference step 之前同步两个 device 之间的 weights。下面描述 weight sync 代码，它位于 cs336_alignment/vllm_utils.py 中。

@dataclass
class VLLMServer:
    gpu: int = 1 # Run training on gpu 0 and inference on gpu 1
    # Create the NCCL weight-transfer group between the training GPU and vLLM.
    def init_weight_sync(self, policy_device: str): ...
    # Copy the current Hugging Face policy weights into the vLLM server and
    # reset vLLM caches that depended on the old weights.
    def sync_policy_weights(self, policy: torch.nn.Module) -> None: ...
    # Generate rollouts from the current vLLM weights.
    def generate_completions(
        self,
        prompts: list[str],
        sampling_params: dict,
        batch_size: int | None = None,
    ) -> list[VLLMCompletion]: ...

与我们的 prompting 实验一样，我们将以 temperature 1.0、top-p 1.0、max generation length 512 进行采样。Prompt 要求 model 以字符串 </answer> 结束其回答，因此我们可以让 vLLM 在 model 输出这个字符串时停止：

# Based on Dr. GRPO: stop when the model completes its answer
# https://github.com/sail-sg/understand-r1-zero/blob/
#   c18804602b85da9e88b4aeeb6c43e2f08c594fbc/train_zero_math.py#L167
sampling_params['stop'] = ["</answer>"]
sampling_params['include_stop_str_in_output'] = True

4.2.3 GRPO components

接下来，让我们实现计算 GRPO loss 各部分的 components。注意，在作业后面，我们会实现变体，例如不同的 advantage normalizers 或 importance reweighting 方法。下面的 components 经过设计，使得我们可以在不同变体之间切换，而只需最少的代码重叠，因此我们提供了带有参数、用于指定使用哪个变体的接口。对于作业的这一部分，你只需要实现标准的 GRPO 变体。

我们的第一步是实现一个 helper function，用于计算 responses 的 rewards。

Problem (compute_rollout_rewards)：Computing the rewards of rollouts（1 分）

第 17 页

Deliverable: 实现方法 compute_rollout_rewards，为每个 rollout response 计算 raw reward。

推荐使用如下接口：

def compute_rollout_rewards(
    reward_fn: Callable[[str, str], dict[str, float]],
    rollout_responses: list[str],
    repeated_ground_truths: list[str],
) -> tuple[torch.Tensor, dict[str, float]]:

为一系列 rollout responses 计算 reward，并返回 reward 各组件的 metadata。

Args: - reward_fn: Callable[[str, str], dict[str, float]] 根据 ground truths 给 rollout responses 打分，产生一个字典，键为 "reward"、"format_reward" 和 "answer_reward"。 - rollout_responses: list[str] 来自 policy 的 rollouts。该列表长度为 rollout_batch_size = n_prompts_per_rollout_batch * group_size。 - repeated_ground_truths: list[str] 各示例的 ground truths。该列表长度为 rollout_batch_size，因为每个示例的 ground truth 会被重复 group_size 次。

Returns: - tuple[torch.Tensor, dict[str, float]]。 - raw_rewards 形状 (rollout_batch_size,)。每个 rollout response 的未归一化 reward。 - metadata 用于记录的 reward 统计量。至少包括整个 rollout batch 上 total reward 和 format reward 的均值。

要测试你的代码，请实现 [adapters.run_compute_rollout_rewards]。然后运行测试：

uv run pytest -k compute_rollout_rewards

并确保你的实现通过该测试。

接下来，我们来实现 GRPO 的核心数学部分，它会对这些 reward 做归一化，把它们变成 advantages。注意，我们上面的函数输出的 rewards 是展平的，因此需要接收 group size，把它 reshape 回各个 group 以进行归一化。为了通过测试，你需要支持 baseline = "mean" 与 advantage_normalizer = "std"。为避免除以零，你还应在除以每个 group 的 std 之前，把 advantage_eps 加到该 std 上。

Problem (compute_group_normalized_rewards_grpo)：Group normalization（1 point）

Deliverable: 实现方法 compute_group_normalized_rewards，在其所属 group 内对 raw rewards 做归一化，并返回归一化后的 rewards 以及你认为有用的 metadata。

目前，你只需要支持 baseline = "mean" 和 advantage_normalizer = "std"。对于不支持的输入，可以随意抛出 NotImplementedError。在作业后续部分，我们会实现其他选项并运行 ablations。记住要把 advantage_eps 加到 normalizer 上以避免除以零。

推荐使用如下接口：

第 18 页

def compute_group_normalized_rewards(
    raw_rewards: torch.Tensor,
    group_size: int,
    baseline: Literal["mean", "none"] = "mean",
    advantage_eps: float = 1e-6,
    advantage_normalizer: Literal["std", "none", "mean"] = "std",
):

计算 advantages：在每个 group 内应用所要求的 baseline 和 normalization。

Args: - raw_rewards: torch.Tensor 形状 (rollout_batch_size,)。每个 rollout response 的未归一化 reward，其中 rollout_batch_size = n_prompts_per_rollout_batch * group_size。 - group_size: int 每个问题（group）对应的 response 数。 - baseline: Literal["mean", "none"] 对本题，支持 mean，即减去每个 group 的平均 reward。之后，none 表示不减去 baseline。 - advantage_eps: float 用于归一化时避免除以零的小常数。 - advantage_normalizer: Literal["std", "none", "mean"] 对本题，支持 std，即除以每个 group 的标准差。之后，none 表示不归一化，mean 表示除以每个 group 的平均 reward。

Returns: - tuple[torch.Tensor, dict[str, float]]。 - advantages 形状 (rollout_batch_size,)。每个 rollout response 的 group-normalized reward。 - metadata 由你选择要记录的其他统计量（例如 reward 的 mean、std、max/min）。

要测试你的代码，请实现 [adapters.run_compute_group_normalized_rewards]。然后运行测试：

uv run pytest -k compute_group_normalized_rewards_grpo

并确保你的实现通过该测试。

下一步是实现一个函数，用来计算 per-token policy gradient loss。注意：Algorithm 1 中的公式写的是 gradient estimator，但本题要求你写一个 per-token loss，使得对它求梯度会产生 GRPO gradient estimator 的每一项（之后我们会在 token 上聚合）。这个表达式并不是传统意义上的 loss——在传统意义中我们期望它在训练过程中下降。它只是一个表达式，使得我们求梯度时能得到 policy gradient。

你应实现的具体 loss 是

$$ J_{\theta}^{\text{GRPO-on-policy}} = \frac{1}{BG}\sum_{i=1}^{B}\sum_{j=1}^{G} \frac{1}{\operatorname{len}(y^{(i,j)})} \sum_{t=1}^{\operatorname{len}(y^{(i,j)})} \frac{r(y^{(i,j)} \mid x^{(i)})-\mu_i}{\operatorname{std}_i} \log \pi_\theta\bigl(y_t^{(i,j)} \mid x^{(i)}, y_{

你可以验证，对它求导会产生 Algorithm 1 中的梯度项。由于 PyTorch 优化器执行梯度下降，你的实现应返回该目标的负值作为 loss。compute_policy_gradient_loss 会计算每个 token 和每个 sequence 的 loss 项，而 aggregate_loss_across_microbatch_sequence 会在 tokens 和 sequences 上对 per-token losses 取平均。

第 19 页

Problem (compute_policy_gradient_loss_on_policy)：On-policy policy gradient（1 point）

Deliverable: 实现方法 compute_policy_gradient_loss，在给定 raw rewards 或预计算好的 advantages 时，计算 per-token policy-gradient loss。

目前，我们假设所有 rollouts 都是 on-policy，所以你只需要支持 importance_reweighting_method = "none"，并且可以忽略 old-log-prob 和 clipping 参数。对于不支持的输入，可以随意抛出 NotImplementedError。我们会在作业后续部分实现 clipping。

推荐使用如下接口：

def compute_policy_gradient_loss(
    raw_rewards_or_advantages: torch.Tensor,
    policy_log_probs: torch.Tensor,
    importance_reweighting_method: Literal["none", "noclip", "grpo", "gspo"] = "none",
    old_log_probs: torch.Tensor | None = None,
    cliprange: float | None = None,
    response_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:

在每个 token 上计算 policy-gradient loss，其中 raw_rewards_or_advantages 要么是 raw reward，要么是已经归一化好的 advantage。

Args: - raw_rewards_or_advantages: torch.Tensor 形状 (batch_size,) 或 (batch_size, 1)，每个 rollout response 的标量 reward/advantage。 - policy_log_probs: torch.Tensor 形状 (batch_size, sequence_length)，每个 token 的 logprobs。 - importance_reweighting_method: Literal["none", "noclip", "grpo", "gspo"] "none"：不做 importance reweighting；"noclip"：应用 importance reweighting 但不裁剪；"grpo"：做 PPO/GRPO 风格的 token-level reweighting 和 clipping；"gspo"：做 GSPO 风格的 sequence-level reweighting 和 clipping。 - old_log_probs: torch.Tensor | None 除非 importance_reweighting_method = "none"，否则为必需；形状 (batch_size, sequence_length)。 - cliprange: float | None = None 裁剪参数 $\varepsilon$，当 importance_reweighting_method 为 "grpo" 或 "gspo" 时必需。 - response_mask: torch.Tensor | None = None 可选的形状 (batch_size, sequence_length) 的 response tokens mask。对于只在 response tokens 上平均 sequence-level log-ratio 的 GSPO 实现，此参数为必需。

Returns: - tuple[torch.Tensor, dict[str, torch.Tensor]]。 - per_token_policy_gradient_loss 形状 (batch_size, sequence_length)，per-token policy-gradient loss（将在训练循环中跨 batch 和 sequence 维度聚合）。 - metadata 来自底层 loss 调用的统计量，例如 clip-fraction 各分量。

第 20 页

要测试你的代码，实现 [adapters.run_compute_policy_gradient_loss]。然后运行

uv run pytest -k test_compute_policy_gradient_loss_on_policy

并确保测试通过。

最后，让我们来实现对 token 和 sequence 的聚合函数；在标准 GRPO 中，这涉及先对每条 sequence 内的 token 求平均，然后再对 sequence 求平均。稍后我们还会实现另一种做法，即直接除以一个常数；我们在推导中已经看到，这种做法更忠实于我们试图估计的原始 policy gradient。

Problem (aggregate_loss_across_microbatch_sequence)：跨 token 和 sequence 聚合 loss（0.5 分）

Deliverable：实现一个方法 aggregate_loss_across_microbatch，接收逐 token 的 policy-gradient loss 和 response mask，输出平均 loss。目前我们采用标准 GRPO 聚合方式，即先在每条 sequence 内部求平均，再在 sequence 之间求平均，因此你可以假定 loss_normalization = "sequence"。对不支持的输入，尽可抛出 NotImplementedError。

推荐的接口如下：

def aggregate_loss_across_microbatch(
    per_token_policy_gradient_loss: torch.Tensor,
    mask: torch.Tensor,
    loss_normalization: Literal["sequence", "constant"] = "sequence",
    normalization_constant: int | None = None,
) -> torch.Tensor:

按 response mask 和 loss 归一化策略聚合逐 token 的 policy-gradient loss。

Args：

Returns：

要测试你的代码，实现 [adapters.run_aggregate_loss_across_microbatch]。然后运行

uv run pytest -k test_aggregate_loss_across_microbatch_sequence

并确保测试通过。

第 21 页

4.2.4 GRPO training step

现在我们可以把这些组件拼成一个完整的训练步（training step）了。

Gradient accumulation

我们需要较大的 batch size，才能在 inference 时获得良好的利用率。因此在 on-policy RL 中（我们把 train batch size 设为与 inference batch size 相同），GPU 的内存不足以一次性对整个 batch 计算梯度。所以我们需要把 batch 切成一系列 microbatch，并在这些 microbatch 上累积梯度。主要难点在于正确处理归一化，以确保 microbatch 累积得到的梯度等价于在整个 batch 上计算出的梯度。

在 PyTorch 中实现 gradient accumulation 很直接。回忆一下，每个权重 tensor 都有一个 .grad 属性用来存梯度。在我们的代码里，调用 loss.backward() 之前，.grad 属性是 None；调用 loss.backward() 之后，.grad 属性就包含了梯度。通常我们会先做一次 optimizer step，然后用 optimizer.zero_grad() 清零梯度，它会把权重 tensor 的 .grad 字段重置：

# Forward pass.
logits = model(inputs)
loss = loss_fn(logits, labels)
# Backward pass.
loss.backward()
# Update weights.
optimizer.step()
# Zero gradients in preparation for next iteration.
optimizer.zero_grad()

要实现 gradient accumulation，我们会把 batch 切成 $k$ 个 microbatch，并且只调用一次 optimizer.step() 和 optimizer.zero_grad()。假定采用 sequence normalization，即先对每条 sequence 求平均 loss、再在 sequence 之间求平均，那么在算出每个 microbatch 的平均 loss 之后，我们还需要按 sequence 数量重新加权。

gradient_accumulation_steps = 4
microbatch_size = len(inputs) // gradient_accumulation_steps
for i in range(0, len(inputs), microbatch_size):
    inputs_microbatch = inputs[i:i+microbatch_size]
    labels_microbatch = labels[i:i+microbatch_size]
    # Forward pass.
    logits = model(inputs_microbatch)
    loss = loss_fn(logits, labels_microbatch) * (len(inputs_microbatch) / len(inputs))
    # Backward pass.
    loss.backward()
# Update weights once across entire batch.
optimizer.step()
# Zero gradients once across entire batch.
optimizer.zero_grad()

注意，后面当我们改用 constant normalization 而不是 sequence normalization 时，microbatch 的代码会略有变化，以适配不同的归一化方式。

第 22 页

Implementing the train step

在下面的函数中，我们最终实现一个 GRPO 训练步（training step），输入是某一批 rollouts。该函数接收 model、tokenizer、optimizer、reward function、prompts 和 rollouts，以及各种超参数；它会累积梯度并执行一次 optimizer step。你需要按上文所述实现 gradient accumulation，以避免显存不足。在 optimizer step 之前，该函数还应把梯度范数裁剪到 max_grad_norm。函数随后应返回该 batch 上的 training loss 以及 metadata，两者都用于 logging。请在你的 train step 中至少记录以下指标：

也欢迎记录其他指标：例如记录 $k$ 在 1 到 group_size 之间的 pass@k 可能很有用（即在前 $k$ 个 rollout 中至少有一个正确答案的 prompt 所占比例）。

Problem (grpo_train_step_standard_on_policy)：GRPO train step（5 分）

Deliverable：给定 model、tokenizer 和 rollouts，实现 policy gradient 的单 batch 更新。本次作业的这一部分，你只需要在 on-policy 设定下为标准 GRPO 实现该函数，即 baseline = "mean"、advantage_normalizer = "std"、importance_reweighting_method = "none"、loss_normalization = "sequence"。对不支持的输入，尽可抛出 NotImplementedError。

推荐的接口如下：

def grpo_train_step(
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizer,
    optimizer: Optimizer,
    gradient_accumulation_steps: int,
    max_grad_norm: float | None,
    reward_fn: Callable[[str, str], dict[str, float]],
    repeated_prompts: list[str],
    rollout_responses: list[str],
    repeated_ground_truths: list[str],
    group_size: int,
    # Reward normalization
    baseline: Literal["mean", "none"] = "mean",
    advantage_eps: float = 1e-6,
    advantage_normalizer: Literal["std", "none", "mean"] = "std",
    # Importance reweighting and clipping
    importance_reweighting_method: Literal["none", "noclip", "grpo", "gspo"] = "none",
    old_log_probs: torch.Tensor | None = None,
    cliprange: float | None = None,
    # Loss normalization
    loss_normalization: Literal["sequence", "constant"] = "sequence",
    normalization_constant: int | None = None,
) -> tuple[torch.Tensor, dict[str, torch.Tensor | float]]:

第 23 页

执行前向与反向传播，使用 gradient_accumulation_steps 个 microbatches。

Args（参数）：

Returns（返回值）：

第 24 页

要测试你的代码，请实现 [adapters.run_grpo_train_step]。然后运行 uv run pytest -k test_grpo_train_step_standard_on_policy，并确认它通过。

4.3 Experiments（实验）

我们现在可以搭建完整的 GRPO 训练循环了，它应当：

下面给出了一些建议的 hyperparameters。如果你的训练脚本是正确的，那么用这些 hyperparameters 训练时，对大多数随机种子，你应该会看到 validation reward 随时间上升。

n_train_examples = 6400
n_val_examples = 1024
num_rollout_steps = 200
learning_rate = 1e-5
rollout_batch_size = train_batch_size = 256
group_size = 8
gradient_accumulation_steps = 32
sampling_temperature = 1.0
sampling_max_tokens = 512
max_grad_norm = 1.0
optimizer = torch.optim.AdamW(
    policy.parameters(), lr=learning_rate, betas=(0.9, 0.95), weight_decay=0.0
)

注意 rollout_batch_size 和 train_batch_size 统计的是 responses 而不是 prompts。所以 rollout_batch_size = train_batch_size = 256 意味着 32 个 prompts、每个 8 个 rollouts。

作为一个合理的默认做法，你可以每 10 个 rollout batches 在 validation set 上评估一次。你应确保至少在 1024 个 validation examples 上评估，因为 CoT/RL 的评估可能有较大噪声。除了记录上一小节描述的 metrics 之外，记录 rollouts 也很有用，可以定性地了解模型在做什么。一个合理的默认做法是每 40 个 rollout batches 记录一次当前 batch 的 training rollouts。可以随意使用 data/gsm8k/train.jsonl 作为 training examples，data/gsm8k/test.jsonl 作为 validation examples。

由于 policy gradient 估计器的方差可能很大，而 RL 又会自我强化，所以 RL 训练运行的轨迹往往方差很高。因此，一旦你确信脚本是正确的，我们会要求你用 4 个 random seeds 运行你的 RL 循环。脚本在固定 seed 下不完全可复现也没关系；你也可以干脆把实验跑 4 次。在你的图中，请画出带有跨 run 方差信息的 metrics（即不要只画均值）。一些合理的选项包括：

第 25 页

Problem (grpo_experiments_standard_on_policy)：使用 GRPO 提升 OLMo-2-0425-1B 在 GSM8K 上的表现（2 B200 hrs）（10 分）

(a) 写一个脚本，给定模型名、prompt、训练集文件路径、验证集文件路径、sampling hyperparameters 和 training hyperparameters，运行 GRPO 训练循环。在高层次上，脚本会先初始化 vLLM server、wandb logging、datasets、模型和 optimizer。然后训练循环反复做以下事情：与 vLLM server 同步权重、产生 training rollouts、在 training batches 上取 policy gradients，并定期检查在 validation set 上的表现（同时记录 generations）。

Deliverable: 一个在任务 GSM8K、模型 OLMo-2-0425-1B 上运行 GRPO（standard, on-policy）的脚本。

(b) 把脚本运行大约 50 步，确认你看到 validation rewards 上升，并且随时间产生合理的 rollouts。Note：reward 一开始接近零是正常的，取决于随机种子，可能需要若干步才能采到非零 reward 的 rollouts，模型才开始改进。

Deliverable: 能让你相信脚本是正确的证据（例如每步 validation rewards 上升、看起来合理的 rollouts）。

(c) 一旦你确信脚本是正确的，就用上面给出的 hyperparameters、4 个 random seeds，对 OLMo-2-0425-1B 在 GSM8K 上运行它。你应当使用 zero-shot 的 r1_zero prompt。如果代码正确，你应该会看到 validation rewards 随着模型被训练而上升。

记录下列 metrics，并把它们随时间的变化画成图，同时注明不同 run 之间的方差（例如画出跨 run 的 mean/std 或 min/max）。

还请定期记录 rollouts 并观察它们。responses 是否因训练而改善？

Deliverable: 每个 metric 几句评论加一张图，描述它在训练过程中如何变化，以及不同 run 之间有多少方差。

Deliverable: 几个训练前后模型 rollouts 的例子。

Deliverable: 一套平均跨 random seeds 达到至少 25% 最终 validation accuracy 的流程。

现在我们已经有了能跑的脚本，可以开始调 hyperparameters 和做 ablations 了。我们先从调 learning rate 开始，它通常是影响训练最大的 hyperparameter。

第 26 页

Problem (grpo_learning_rate)：调学习率（4 B200 hrs）（3 分）

从建议的超参数出发，对学习率进行扫描（至少包含一个小于推荐默认值的学习率，以及一个大于它的学习率），并报告最终验证奖励（如果优化器发散，则注明发散）。根据你在上一部分观察到的方差大小，你应该决定每次训练运行要使用多少个随机种子。在作业剩余部分，如果你愿意，可以随意使用你调好的学习率，而不使用建议的默认值。

Deliverable（交付物）： 一张最终验证奖励随学习率变化的图，并附几句评论。

另一个影响 RL 动态的选择是 prompt：由于 RL 涉及提高模型 rollout 的权重，prompt 可以塑造模型在训练期间探索哪些 rollout。下面我们将探索 prompt 选择的影响。

Problem (grpo_prompt_ablation)：Prompt 消融（4 B200 hrs）（3 分）

使用 question_only prompt 和 r1_zero_three_shot prompt，配合几个随机种子运行 GRPO 脚本。与上文 zero-shot r1_zero prompt 的结果相比，这些 prompt 表现如何？哪些 prompt 的平均奖励最好、方差最低？在其他记录指标中是否还有其他系统性差异？基于不同运行之间的方差，你对你的发现有多大信心？

Deliverable（交付物）： 评论，并用所引用指标的图作为支撑。

5 RL algorithm variants（RL 算法变体）

GRPO 是语言模型 RL 的一个流行选择，但它的算法选择在多种论文中存在争议。在本作业中，我们将学习这些选择背后的一些理论论证，并进行受控实验，亲自看看哪些算法最好（至少对于我们模型和数据集的组合而言）。本节将涵盖 on-policy 设置，下一节将涵盖 off-policy 算法。

5.1 Dr. GRPO

在 GRPO 的推导过程中，你可能注意到了一系列选择，这些选择导致 GRPO 策略梯度估计器不再具有“正确”的期望 $\nabla_\theta J_\theta$，其中 $J_\theta$ 是模型的期望奖励。这两个选择是标准差优势归一化和序列长度归一化。我们要考虑的第一个消融（Dr. GRPO 论文 [Z. Liu et al., 2025] 中主张的）是撤销这些选择：去掉标准差归一化，并将总损失除以一个常数，而不是先按序列长度对每条序列归一化。将这个常数归一化项记为 $Z$，则 Dr. GRPO 的梯度估计器为

$$ \hat{g} \leftarrow \frac{1}{Z}\sum_{i=1}^{B}\sum_{j=1}^{G}\sum_{t=1}^{\operatorname{len}(y^{(i,j)})}\left(r(y^{(i,j)}\mid x^{(i)})-\mu_i\right)\nabla_\theta \log \pi_\theta\left(y_t^{(i,j)}\mid x^{(i)}, y_{

第 27 页

其中 $Z$ 通常设为训练 batch size 乘以最大生成长度（即 $Z = BGL$，其中 $L$ 是最大生成长度（在我们的情况中为 512））。

Problem (think_about_length_normalization)：思考长度归一化（1 分）

在运行任何实验之前，思考按序列长度对每条序列归一化，与用同一个常数对所有序列归一化之间的区别。每种方法各有什么优缺点？是否存在某些特定设置或例子，其中一种方法看起来更好？

Deliverable（交付物）： 几句讨论。

Problem (compute_group_normalized_rewards_drgrpo)：Dr. GRPO 组归一化（0.5 分）

Deliverable（交付物）： 更新你的 compute_group_normalized_rewards 方法，使其支持 advantage_normalizer = "none"。虽然 Dr. GRPO 使用 baseline = "mean"，但在这个问题中，也请添加对 baseline = "none" 的支持。

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

uv run pytest -k compute_group_normalized_rewards_drgrpo

并确保你的实现通过它。

Problem (aggregate_loss_across_microbatch_constant)：Dr. GRPO 损失聚合（0.5 分）

Deliverable（交付物）： 更新你的 aggregate_loss_across_microbatch 方法，使其支持 loss_normalization = "constant"。

为了测试你的代码，实现 [adapters.run_aggregate_loss_across_microbatch]。然后运行测试：

uv run pytest -k test_aggregate_loss_across_microbatch_constant

并确保测试通过。

5.2 Rejection fine tuning（拒绝微调）

我们要看的下一个消融是一个简单得多的算法，通常称为 “rejection fine tuning”（RFT）或 “expert iteration”（EI）。顾名思义，该算法包括采样一批 rollout，保留其中正确的那些，然后对这些正确的 rollout 做监督微调（即最小化下一词预测对数损失）。具体来说，我们的 RFT 梯度为

$$ \hat{g} \leftarrow \nabla_\theta \left[ \frac{1}{Z}\sum_{i=1}^{B}\sum_{j=1}^{G}\mathbf{1}\{r(y^{(i,j)}\mid x^{(i)}) = 1\} \sum_{t=1}^{\operatorname{len}(y^{(i,j)})}\log \pi_\theta\left(y_t^{(i,j)}\mid x^{(i)}, y_{

其中 $Z$ 是常数归一化项（如同在 Dr. GRPO 中一样），$\mathbf{1}\{r(y^{(i,j)}\mid x^{(i)}) = 1\}$ 是一个指示函数，只保留正确的回答。在接下来的问题中，我们将思考 RFT 梯度是否实际上是一个 policy gradient（即它是否真的优化期望奖励 $J_\theta$），以及它与 GRPO 有何关系。和之前一样，在将其实现为 PyTorch 损失时，请返回负目标，以便最小化损失会对该估计器执行梯度上升。

第 28 页

Problem (think_about_rft)：思考 RFT（2 分）

回忆一下，RFT 目标（使用常数归一化）为

$$ J_\theta = \frac{1}{Z}\sum_{x}\sum_{j=1}^{G}\mathbf{1}\{r(y^{(j)}\mid x)=1\}\log \pi_\theta(y^{(j)}\mid x), \tag{35} $$

其中 $x$ 是 prompt，每个 $y^{(j)}$ 表示一个回答 $y^{(j)} \overset{\text{iid}}{\sim} \pi_\theta(\cdot\mid x)$，$G$ 是我们从策略中采样的生成数量，$Z$ 是我们的常数归一化项，$r$ 是奖励函数。RFT 梯度则由 $\nabla_\theta J_\theta$ 给出。

相比之下，on-policy 的 Dr. GRPO 策略梯度估计器（即使用常数归一化且不做标准差归一化的 GRPO）为

$$ \frac{1}{Z}\sum_{x}\sum_{j=1}^{G}\left(r(y^{(j)}\mid x)-\mu\right)\nabla_\theta\log \pi_\theta(y^{(j)}\mid x), \tag{36} $$

其中 $\mu$ 表示组均值 $\mu = \frac{1}{G}\sum_{j=1}^{G} r(y^{(j)}\mid x)$。假设我们的奖励函数 $r$ 是二值的，比较并对比这些目标。它们具有相同的期望吗？你认为哪一个会有更低的方差？直观地说，是否存在你会更偏好其中一个的情形？

Deliverable（交付物）： 几句讨论。

5.3 MaxRL

最后，我们将探索一种近期方法，称为 Maximum Likelihood Reinforcement Learning（最大似然强化学习），也称为 MaxRL [F. Tajwar et al., 2026]。MaxRL 不是像 GRPO 那样用组标准差对优势进行归一化，也不是像 Dr. GRPO 那样完全不归一化，而是提出除以组均值。这一选择得到以下梯度估计器：

$$ \hat{g} \leftarrow \frac{1}{Z}\sum_{i=1}^{B}\sum_{j=1}^{G}\sum_{t=1}^{\operatorname{len}(y^{(i,j)})} \frac{r(y^{(i,j)}\mid x^{(i)})-\mu_i}{\mu_i} \nabla_\theta \log \pi_\theta\left(y_t^{(i,j)}\mid x^{(i)}, y_{

注意，在原始的 MaxRL 论文中，他们不是使用常数归一化项 $Z$，而是除以 batch 中 token 的总数 $\sum_{i=1}^{B}\sum_{j=1}^{G}\operatorname{len}(y^{(i,j)})$（现在这依赖 batch，而非常数）。在本作业中，我们改用常数归一化项，这样相对于其他 baseline 有更少的变动因素，从而更容易分离出按 $\mu_i$ 而非 $\mathrm{std}_i$ 归一化所带来的影响。

在接下来的问题中，我们将做一些基础数学，以探索不同优势归一化器的影响。直观地说，除以平均奖励 $\mu_i$ 平均而言会提高更困难 prompt 的权重，因为这些 prompt 相比更容易的 prompt 更可能具有较低的 $\mu_i$。我们可以通过定义期望奖励的一个重新加权版本来形式化这一直觉：

$$ J_{\theta,w} = \mathbb{E}_{x\sim \rho}\left[w(x, \operatorname{stopgrad}(\pi_\theta))\mathbb{E}_{y\sim \pi_\theta} r(y\mid x)\right], \tag{38} $$

其中 $w(x, \operatorname{stopgrad}(\pi_\theta))$ 对 prompt $x$ 重新加权，例如根据它们的难度 $\eta(x) = \mathbb{E}_{y\sim \pi_\theta} r(y\mid x)$ 来加权，而该难度由当前策略在 $x$ 上的期望奖励刻画；$\operatorname{stopgrad}$ 表示我们不对重新加权函数求梯度。在接下来的问题中，我们将推导到目前为止见过的每个优势归一化器的难度重加权。

第 29 页

Problem (derive_difficulty_reweightings): 推导由不同 advantage normalizers 引起的 difficulty reweightings（6 分）

回顾：在 prompts 上聚合时，标准 policy gradient 为

$$\nabla_\theta J_\theta = \nabla_\theta \mathbb{E}_{x\sim\rho}\left[\mathbb{E}_{y\sim\pi_\theta} r(y\mid x)\right], \tag{39}$$

其中 $\rho$ 是我们在 prompts $x$ 上的 prompt 分布，$\pi_\theta$ 表示我们的 policy，$r(y\mid x)$ 表示 $y$ 是否为对 $x$ 的正确回答。在这道题中，我们将推导 GRPO 各变体所优化的 surrogate objectives，这些 surrogate 具有如下形式

$$\nabla_\theta J_{\theta,w} = \nabla_\theta \mathbb{E}_{x\sim\rho}\left[w(x, \operatorname{stopgrad}(\pi_\theta))\mathbb{E}_{y\sim\pi_\theta} r(y\mid x)\right] \tag{40}$$

其中 $w$ 是对 prompts 的某种重加权（我们不对 $w$ 求导）。

(a) 回顾 Dr. GRPO 的 policy gradient estimator 为

$$\mathbb{E}_{x\sim\rho}\left[\frac{1}{Z}\sum_{j=1}^{G}\left(r(y^{(j)}\mid x)-\mu\right)\nabla_\theta\log \pi_\theta(y^{(j)}\mid x)\right] \tag{41}$$

其中 $\mu$ 表示组均值 $\mu = \frac{1}{G}\sum_{j=1}^{G} r(y^{(j)}\mid x)$。令常数 normalizer $Z=G$，并令组大小 $G\to\infty$，该 estimator 在什么重加权函数 $w$ 下等价于优化 surrogate objective $J_{\theta,w}$？

Deliverable: 用问题参数的一个子集给出表达式，并附上几句论证。

(b) 在采用 constant normalization 的情况下，GRPO estimator 与 Dr. GRPO 的不同之处在于它除以标准差：

$$\mathbb{E}_{x\sim\rho}\left[\frac{1}{Z}\sum_{j=1}^{G}\frac{r(y^{(j)}\mid x)-\mu}{\operatorname{std}}\nabla_\theta\log \pi_\theta(y^{(j)}\mid x)\right], \tag{42}$$

其中 $\operatorname{std}$ 表示组标准差。令常数 normalizer $Z=G$，并令组大小 $G\to\infty$，该 estimator 在什么重加权函数 $w$ 下等价于优化 surrogate objective $J_{\theta,w}$？

Deliverable: 用问题参数的一个子集给出表达式，并附上几句论证。

(c) MaxRL 则改为除以组均值，即

$$\mathbb{E}_{x\sim\rho}\left[\frac{1}{Z}\sum_{j=1}^{G}\frac{r(y^{(j)}\mid x)-\mu}{\mu}\nabla_\theta\log \pi_\theta(y^{(j)}\mid x)\right]. \tag{43}$$

令常数 normalizer $Z=G$，并令组大小 $G\to\infty$，该 estimator 在什么重加权函数 $w$ 下等价于优化 surrogate objective $J_{\theta,w}$？

Deliverable: 用问题参数的一个子集给出表达式，并附上几句论证。

Problem (think_about_advantage_normalization): 思考 advantage normalization（2 分）

第 30 页

在运行任何实验之前，思考一下将组 advantages 分别按组 std、组 mean 归一化，或者不做 advantage normalization，这三者之间有什么区别。每种方法各有什么优点和缺点？是否存在某些特定设置或例子，其中某一种方法看起来更好？

Deliverable: 几句讨论。

Problem (compute_group_normalized_rewards_maxrl): MaxRL 的组归一化（0.5 分）

Deliverable: 更新你的方法 compute_group_normalized_rewards，使其支持 advantage_normalizer = "mean"。与 advantage_normalizer = "std" 一样，你应在 normalizer 中加上 advantage_eps，以避免除以零。

要测试你的代码，请实现 [adapters.run_compute_group_normalized_rewards]。然后使用以下命令运行测试，并确保你的实现通过：

uv run pytest -k compute_group_normalized_rewards_maxrl

5.4 Experiments

现在可以运行实验来探索这些不同的 policy gradient estimators 了。首先，我们需要更新 grpo_train_step 方法以支持这些变体；其各个组成部分我们已在上文实现。具体设置如下：

为加速训练，请注意：一旦我们为每条序列计算了归一化后的 advantages，advantage 为零的序列就无需再传入模型，因为它们对梯度的贡献为零。由于我们的奖励是二值的，序列会在两种情况下具有零 advantage：

一旦剪掉零 advantage 的序列，我们还可以按相同比例减少 gradient_accumulation_steps（即可以保持 microbatch_size 不变）：例如，如果一半序列的 advantage 为零，我们可以执行 $\frac{k}{2}$ 个 grad accum steps，而不是 $k$ 个 grad accum steps。这个优化对 RFT 尤其有帮助，因为 RFT 会有相当多的零奖励序列。但你应在数学推导和实现上小心，确保剪枝后的版本计算出的梯度与未剪枝版本相同。

Problem (grpo_train_step_variants_on_policy): GRPO train step 变体（2.5 分）

第 31 页

Deliverable: 更新你的 grpo_train_step 方法，使其支持全部 on-policy 变体（仍然使用 importance_reweighting_method = "none"）。这些选项包括 baseline: Literal["mean", "none"]、advantage_normalizer: Literal["std", "none", "mean"] 和 loss_normalization: Literal["sequence", "constant"]。此外，为加速训练，你的方法应避免将 advantage 为零的序列传入模型。

注意，对于 baseline = "none"，错误样本获得零奖励，因此在 loss 中的权重为零，所以没有必要将它们传入模型。你可以利用这一点来调整实现并加速训练。

要测试你的代码，请实现 [adapters.run_grpo_train_step]。然后运行以下命令并确认其通过：

uv run pytest -k test_grpo_train_step_variants_on_policy

现在训练步已经实现，我们可以运行实验来观察各方法的性能，并判断它们是否优于标准 GRPO。你应使用与之前 runs 相同的 hyperparameters，以使 runs 之间具有可比性。

关于调 hyperparameters 的说明： 如果我们的目标是判断一种方法是否优于另一种方法，那么实验上最可靠的做法是为每种方法分别调 hyperparameters。然而，hyperparameter tuning 开销很大，而这份作业没有那么多算力。另一种开销较小的选择是：为 baseline 调 hyperparameters，然后将这些 hyperparameters 用于新方法（例如，在我们的情况中，我们在标准 GRPO 上调 learning rate，并把该 learning rate 用于后续实验）。在这种情况下，如果我们的新方法优于 baseline，我们仍然可以声称它更好，因为这个次优 learning rate 会给新方法在充分调参后的性能提供一个下界。但反过来，如果新方法表现更差，我们不能声称它比 baseline 更差，因为我们没有正确调它的 hyperparameters。

Problem (grpo_experiments_variants_on_policy): 比较不同 RL 算法的性能（8 B200 hrs）（10 分）

在与你的标准 GRPO 训练 runs 保持 hyperparameters 固定（就你所选择的 learning rate 而言）的情况下，运行以下变体，每种使用 4 个随机种子。你应使用 zero-shot 的 r1_zero prompt。

与你之前运行的标准 GRPO runs 相比，这些方法表现如何？哪些方法的平均性能最好，哪些方法的方差最低？你认为是否有方法会从更多 hyperparameter tuning 中受益？基于不同 runs 之间的方差，你对你的发现有多大信心？

第 32 页

Deliverable: 针对所引用指标的文字评述（commentary），并配以图表支持。

6 Off-policy RL

我们至今的实验和算法都是完全 on-policy 的，也就是说，每个梯度估计器所用的样本都直接抽自正在被更新的模型（即：每个 inference batch 只做 1 个训练步，且训练 batch size 被设为与 inference batch size 相等）。本节我们将探讨 off-policy RL：它对每个 inference batch 做多个训练步，有潜力加快训练速度，代价是可能出现不稳定、算法复杂度上升。

6.1 Importance reweighting

在 policy gradient 的朴素版本中，我们采样一批 responses，然后在这批 responses 上做一个大步的 batch 梯度更新。但如果我们把这一批拆成若干 minibatch、每个 minibatch 做一次训练步，模型或许能训练得更快。这种做法称为 “off-policy RL”，与我们原先的 “on-policy RL”（每个 inference batch 只做一个训练步）相对。

这种做法之所以叫 off-policy RL，是因为在第一个 minibatch step 之后，当前 policy 已经不同于 inference 时所用的 policy。此时样本是 “off-policy” 的、即 “stale” 的：它们并非来自当前 policy，而我们想计算的恰恰是当前 policy 的 policy gradient。用记号来说，若把 inference policy 记为 $\pi_0$、当前 policy 记为 $\pi_\theta$，则有以下关系：

$$\mathbb{E}_{x\sim\rho}\mathbb{E}_{y\sim\pi_0(y\mid x)}\left[r(y\mid x)\nabla_\theta\log\pi_\theta(y\mid x)\right]\neq\mathbb{E}_{x\sim\rho}\mathbb{E}_{y\sim\pi_\theta(y\mid x)}\left[r(y\mid x)\nabla_\theta\log\pi_\theta(y\mid x)\right]. \tag{44}$$

左端是我们的 “naive” off-policy 估计器：我们假装 responses $y$ 来自当前 policy，机械地套用同样的步骤去算 policy gradient。这个式子说明，该 naive 估计器的期望是不对的。

纠正这一偏差的一种方法是使用 importance reweighting，得到如下估计器：

$$\mathbb{E}_{x\sim\rho}\mathbb{E}_{y\sim\pi_0(y\mid x)}\left[\frac{\pi_\theta(y\mid x)}{\pi_0(y\mid x)}r(y\mid x)\nabla_\theta\log\pi_\theta(y\mid x)\right], \tag{45}$$

其中我们施加了 importance weights $\frac{\pi_\theta(y\mid x)}{\pi_0(y\mid x)}$。这些权重的作用是：对仍然与当前 policy $\pi_\theta$ 相关的 $y$ 上调权重，而对 $\pi_\theta(y\mid x)$ 很低的 “stale” 样本下调权重。不难看出，这个修改后的表达式具有正确的期望：

$$\mathbb{E}_{x\sim\rho}\mathbb{E}_{y\sim\pi_0(y\mid x)}\left[\frac{\pi_\theta(y\mid x)}{\pi_0(y\mid x)}r(y\mid x)\nabla_\theta\log\pi_\theta(y\mid x)\right]=\mathbb{E}_{x\sim\rho}\mathbb{E}_{y\sim\pi_\theta(y\mid x)}\left[r(y\mid x)\nabla_\theta\log\pi_\theta(y\mid x)\right]. \tag{46}$$

缺点是 importance reweighting 会增大估计器的方差：如果 importance weight $\frac{\pi_\theta(y\mid x)}{\pi_0(y\mid x)}$ 可以大到 $C$，我们的方差也可能膨胀到 $C$ 倍。这一问题使得 naive importance reweighting 无法用于 language model —— 它的 importance weight 是 $\text{len}(y)$ 个项的连乘，随 response 长度呈指数增长：

$$\frac{\pi_\theta(y\mid x)}{\pi_0(y\mid x)}=\prod_{t=1}^{\text{len}(y)}\frac{\pi_\theta(y_t\mid x,y_{

因此，与 “naive” 的不加权 off-policy 估计器（高偏差，但没有 reweighting 引入的额外方差）相比，这个 reweighted 估计器零偏差、高方差。下面的方法介于两者之间，提出各种启发式方法来权衡这一 bias-variance tradeoff。

第 33 页

6.2 PPO/GRPO-style importance reweighting and clipping

6.2.1 Token-level reweighting

标准做法最早由 PPO [J. Schulman et al., 2017] 提出，GRPO [Z. Shao et al., 2024] 也采用了它：做 token 级的 importance reweighting，而不是 sequence 级的 reweighting。这一选择得到如下 “未截断的 token 级 reweighting” 估计器：

$$\mathbb{E}_{x\sim\rho}\mathbb{E}_{y\sim\pi_0(y\mid x)}\left[r(y\mid x)\sum_{t=1}^{\text{len}(y)}\frac{\pi_\theta(y_t\mid x,y_{

注意这个表达式与我们上面看到的 “principled” sequence 级估计器并不相同，后者可以改写成下面这样（按时间步展开）：

$$\mathbb{E}_{x\sim\rho}\mathbb{E}_{y\sim\pi_0(y\mid x)}\left[r(y\mid x)\left(\prod_{t=1}^{\text{len}(y)}\frac{\pi_\theta(y_t\mid x,y_{

在深入讨论这个修改后的表达式在做什么之前，先注意 reweighting 项不再是 $\text{len}(y)$ 的指数级，这大幅降低了方差。

与作业前面几节一样，为了更好理解这个估计器，我们将推导它所优化的 surrogate objective [T. Degris et al., 2013]。具体地，令 $\tilde{\pi}_t$ 表示这样的 surrogate policy：所有时间步都从旧 policy $\pi_0$ 采样，唯独第 $t$ 步从 $\pi_\theta$ 采样：

$$\tilde{\pi}_t(y\mid x)=\left(\prod_{s=1}^{t-1}\pi_0(y_s\mid x,y_{

那么，如果假设所有 response 长度相同、均为 $L$，可以证明 token 级 reweighting 优化的正是每个 $\tilde{\pi}_t$ 下的期望 reward，即

$$J_\theta^{\text{token}}=\mathbb{E}_x\left[\sum_{t=1}^{L}\mathbb{E}_{y\sim\tilde{\pi}_t(y\mid x)}\left[r(y\mid x)\right]\right]. \tag{51}$$

这一事实直接来自用 importance reweighting 项改写该 surrogate objective：

$$\begin{aligned}\nabla_\theta J_\theta^{\text{token}}&=\nabla_\theta\mathbb{E}_x\left[\sum_{t=1}^{L}\mathbb{E}_{y\sim\tilde{\pi}_t(y\mid x)}\left[r(y\mid x)\right]\right]\qquad(52)\\&=\nabla_\theta\mathbb{E}_x\left[\sum_{t=1}^{L}\mathbb{E}_{y\sim\pi_0(y\mid x)}\left[\frac{\pi_\theta(y_t\mid x,y_{

其中最后一步来自本作业通篇使用的 log derivative trick。

这个 surrogate objective 让我们直观地看出 token 级 reweighting 引入的偏差。第一，对于第 $t$ 项，prefix $y_{t}$ 是从 $\pi_0$ 而非 $\pi_\theta$ 采样的。换言之，当该目标评估动作 $y_t$ 如何影响后续步骤时，它是基于 stale policy $\pi_0$ 的行为、而不是当前 policy $\pi_\theta$ 的行为来评估的。这两种偏差都由 sequence 级（下页续）

第 34 页

reweighting 来纠正，它对 prefix 和 suffix 都做了正确的重新加权（代价是方差增大）。直观上，$\pi_\theta$ 离 $\pi_0$ 越远，这些偏差就越严重。

Problem (derive_surrogate_objectives): Derive surrogate objectives for importance reweighting approaches (2 points)

上面我们看到，token 级 importance reweighting 优化的是这样一族 surrogate policy 下的期望 reward：除一个位置外，其余所有位置都从旧 policy 采样。现在假设我们改为考虑一种 “pairwise” importance reweighting 方法，它由下面的 policy gradient 估计器给出：

$$\sum_{t=1}^{\frac{L}{2}}\frac{\pi_\theta(y_{2t-1}\mid x,y_{<2t-1})\pi_\theta(y_{2t}\mid x,y_{<2t})}{\pi_0(y_{2t-1}\mid x,y_{<2t-1})\pi_0(y_{2t}\mid x,y_{<2t})}r(y\mid x)\nabla_\theta\left[\log\left(\pi_\theta(y_{2t-1}\mid x,y_{<2t-1})\pi_\theta(y_{2t}\mid x,y_{<2t})\right)\right] \tag{55}$$

其中 $\pi_\theta$ 是当前 policy，$\pi_0$ 是 stale 采样 policy，且 $y\sim\pi_0(y\mid x)$。这个估计器优化的是什么 surrogate objective？

Deliverable: 一个用问题的部分参数表示的表达式，并附上推导。

6.2.2 Clipping

除了 token 级 reweighting，PPO 和 GRPO 引入的另一条启发式是 importance weight clipping。如上所见，随着 $\pi_\theta$ 离 $\pi_0$ 越来越远，importance reweighting 的方差会越来越严重。我们的 surrogate objective 告诉我们，token 级 reweighting 的偏差也会越来越严重。所以，为了让 off-policy RL 保持稳定，我们需要确保 $\pi_\theta$ 始终靠近 $\pi_0$。

确保 $\pi_\theta$ 靠近 $\pi_0$ 的最简单办法，是减少每个 inference batch 的训练步数，但某些 inference batch 或许允许走更多步。为了把算法性能榨出来，我们希望在不超出 inference batch 允许的范围内尽可能多走步。为此，有一类做法是对 importance weight 过大或过小的项做 clipping。

近期论文对于 clipping 具体应如何实现意见不一，但在本作业中我们将实现 PPO 的 clipping 做法，它一直延续到 GRPO，似乎也是主流做法。首先，把 token 级 reweighting 与标准 GRPO 目标结合，得到如下 loss 函数：

$$J_\theta^{\text{GRPO-off-policy-noclip}}=\frac{1}{BG}\sum_{i=1}^{B}\sum_{j=1}^{G}\frac{1}{\text{len}(y^{(i,j)})}\frac{r(y^{(i,j)}\mid x^{(i)})-\mu_i}{\text{std}_i}\sum_{t=1}^{\text{len}(y^{(i,j)})}\frac{\pi_\theta(y_t\mid x,y_{

其中样本抽自我们 stale 的 inference policy $\pi_0$。你可以验证，对该目标求导得到的正是我们上面推导的 token-reweighted policy gradient 估计器，只是带上了 GRPO 的 advantages 和 sequence normalization。

现在我们准备在这个目标上加入 clipping。令 $A^{(i,j)}=\frac{r(y^{(i,j)}\mid x^{(i)})-\mu_i}{\text{std}_i}$ 表示 prompt $i$ 上 response $j$ 的 advantage，令 $w_t^{(i,j)}=\frac{\pi_\theta(y_t\mid x,y_{

$$J_\theta^{\text{GRPO-off-policy-clip}}=\frac{1}{BG}\sum_{i=1}^{B}\sum_{j=1}^{G}\frac{1}{\text{len}(y^{(i,j)})}\sum_{t=1}^{\text{len}(y^{(i,j)})}\min\left(A^{(i,j)}w_t^{(i,j)},\ A^{(i,j)}\text{clip}(w_t^{(i,j)},[1-\varepsilon,1+\varepsilon])\right), \tag{57}$$

第 35 页

其中 $\operatorname{clip}(w,[1-\varepsilon,1+\varepsilon])=\min(\max(w,1-\varepsilon),1+\varepsilon)$ 会把 importance weight 裁剪到区间 $[1-\varepsilon,1+\varepsilon]$。把该目标写成下面这种形式会稍微更直观一些 [J. Achiam, 2018]：

$$ J_{\text{GRPO-off-policy-clip}}^{\theta} = \frac{1}{BG}\sum_{i=1}^{B}\sum_{j=1}^{G} \frac{1}{\operatorname{len}(y^{(i,j)})} \sum_{t=1}^{\operatorname{len}(y^{(i,j)})} \begin{cases} \min\left(w_t^{(i,j)},1+\varepsilon\right)A^{(i,j)}, & A^{(i,j)}\ge 0,\\[2pt] \max\left(w_t^{(i,j)},1-\varepsilon\right)A^{(i,j)}, & A^{(i,j)}<0. \end{cases} \tag{58} $$

在这个目标下，模型会有动力给 advantage 高的动作提高权重，直到它给该动作的权重比 old policy 高 $1+\varepsilon$。到那时，$\min$ 会走 $1+\varepsilon$ 这一支，对它求导会产生零梯度。负 advantage 的动作类似：模型会有动力降低该动作的权重，直到它给该动作的权重比 old policy 低 $1-\varepsilon$，此时该动作的梯度变为零。因此求梯度后，我们得到如下 policy gradient estimator：

$$ \nabla_\theta J_{\text{GRPO-off-policy-clip}}^{\theta} = \frac{1}{BG}\sum_{i=1}^{B}\sum_{j=1}^{G} \frac{1}{\operatorname{len}(y^{(i,j)})} \sum_{t=1}^{\operatorname{len}(y^{(i,j)})} \operatorname{mask}_t^{(i,j)} w_t^{(i,j)} A^{(i,j)} \nabla_\theta \log \pi_\theta(y_t\mid x,y_{

其中 mask 函数为

$$ \operatorname{mask}_t^{(i,j)} = \begin{cases} \mathbf{1}\left\{w_t^{(i,j)}<1+\varepsilon\right\}, & A^{(i,j)}\ge 0,\\[2pt] \mathbf{1}\left\{w_t^{(i,j)}>1-\varepsilon\right\}, & A^{(i,j)}<0, \end{cases} \tag{60} $$

它对正 advantage 的大 importance weight 做 mask，对负 advantage 的小 importance weight 做 mask。

现在我们有了所需数学，可以在 loss function 中实现 importance_reweighting_method = "noclip" 和 importance_reweighting_method = "grpo"，它们应分别对应 $J_{\text{GRPO-off-policy-noclip}}^{\theta}$ 和 $J_{\text{GRPO-off-policy-clip}}^{\theta}$。和 on-policy 情形一样，从你的 loss 实现中返回负的目标值，这样 gradient descent 就会执行预期的 gradient ascent 更新。

Problem（compute_policy_gradient_loss_off_policy）：带 token 级重加权的 off-policy policy gradient（1 分）

Deliverable： 更新你的 compute_policy_gradient_loss 方法，使其支持 importance_reweighting_method = "noclip" 或 "grpo"，它们会使用 old_log_probs 和 cliprange 参数。old_log_probs 表示用于生成 rollouts 的模型下每个 token 的 log probabilities，你需要在训练脚本中加几行来计算这些值。cliprange 表示裁剪强度参数 $\varepsilon$。

要测试你的代码，请实现 [adapters.run_compute_policy_gradient_loss]。然后运行

uv run pytest -k test_compute_policy_gradient_loss_off_policy

并确保测试通过。

Note： 虽然 clipping 在 PPO 中的动机是防止当前 policy 偏离 old policy 太远，但它也可以被看作一种降低 importance reweighted estimator 方差的方法，因为我们把过大的 importance weight 项压掉了。如果我们不那么担心偏离 old policy 太远，并且希望在更新中更激进，那么一种更简单、更直接的方差控制方式，是直接在 gradient estimator 中对 importance weight 做上界裁剪：

第 36 页

$$ \hat{g}\leftarrow \frac{1}{BG}\sum_{i=1}^{B}\sum_{j=1}^{G} \frac{1}{\operatorname{len}(y^{(i,j)})} \sum_{t=1}^{\operatorname{len}(y^{(i,j)})} \min\left(w_t^{(i,j)},1+\varepsilon\right) A^{(i,j)} \nabla_\theta\log\pi_\theta(y_t\mid x,y_{

这里与 PPO/GRPO 不同：对于那些被提高权重超过 $1+\varepsilon$ 的动作，我们仍然通过它们取得非零梯度。这个更简单的裁剪过程由 CISPO [MiniMax et al., 2025] 提出（注：他们也用每个 group 的 token 数而不是每个 sequence 的长度做归一化，但为简单起见，我们写出按 sequence 归一化的目标）。可选地，欢迎你尝试 CISPO 并与其他裁剪方法比较。

6.3 GSPO

如上所见，token 级 reweighting 是有偏的，因为它忽略了 prefix 和 suffix 的 reweighting。在 GSPO 论文 [C. Zheng et al., 2025] 中，作者发现 GRPO 不稳定，并主张我们应该改为做 sequence 级 reweighting。为了绕开 sequence 级 importance weight 是 $L$ 项乘积这一事实，他们直接把这个表达式取 $\frac{1}{L}$ 次幂。换句话说，他们按整个 sequence 上的 geometric mean importance weight 来 reweight，而不是按乘积。结果得到如下 loss function：

$$ J_{\text{GRPO-off-policy-gspo}}^{\theta} = \frac{1}{BG}\sum_{i=1}^{B}\sum_{j=1}^{G} \min\left( A^{(i,j)}s^{(i,j)}, A^{(i,j)}\operatorname{clip}\left(s^{(i,j)},[1-\varepsilon,1+\varepsilon]\right) \right), \tag{62} $$

其中我们现在有一个跨时间步共享的 sequence 级 importance weight $s^{(i,j)}$：

$$ s^{(i,j)} = \left( \prod_{t=1}^{\operatorname{len}(y^{(i,j)})} \frac{\pi_\theta(y_t\mid x,y_{

注意，如果没有这个 $\frac{1}{\operatorname{len}(y^{(i,j)})}$ 次幂，也没有 clipping，这个目标就变成 sequence 级 importance-reweighted policy gradient loss。取 geometric mean 和 clipping 以偏差为代价降低方差。另一点是，在求梯度时，由于 geometric mean 的幂，我们自然会得到 sequence 长度归一化（为简洁起见，下面的公式中忽略 clipping）：

$$ \begin{aligned} \nabla_\theta J_{\text{GRPO-off-policy-gspo}}^{\theta} &= \frac{1}{BG}\sum_{i=1}^{B}\sum_{j=1}^{G} A^{(i,j)}\nabla_\theta s^{(i,j)}\tag{64}\\ &= \frac{1}{BG}\sum_{i=1}^{B}\sum_{j=1}^{G} A^{(i,j)}s^{(i,j)} \frac{1}{\operatorname{len}(y^{(i,j)})} \sum_{t=1}^{\operatorname{len}(y^{(i,j)})} \nabla_\theta\log\pi_\theta(y_t\mid x,y_{

所以，如果你想将 GSPO 应用到你最喜欢的 constant-normalized RL 算法（可选），你需要改变 sequence 级 importance weight 中的指数。

Problem（think_about_importance_reweighting）：思考 importance reweighting（2 分）

考虑 off-policy RL 的三种可能的 importance reweighting 策略：(a) 不做 importance reweighting，(b) PPO/GRPO 风格的裁剪 token 级 importance reweighting，以及 (c) 使用 geometric mean 的 GSPO 风格裁剪 sequence 级 importance reweighting。在 bias 和 variance 的权衡方面，每种方法落在谱系的哪个位置？你能想到一些情形，其中某一种方法直觉上可能比其他方法更好吗？

Deliverable： 几句评论。

第 37 页

在下面的问题中，我们将实现 GSPO loss。记住要以数值稳定的方式计算 geometric mean。

Problem（compute_policy_gradient_loss_off_policy_gspo）：带 sequence 级重加权的 off-policy policy gradient（1 分）

Deliverable： 更新你的 compute_policy_gradient_loss 方法，使其支持 importance_reweighting_method = "gspo"，它会使用 old_log_probs 和 cliprange 参数。old_log_probs 表示用于生成 rollouts 的模型下每个 token 的 log probabilities，你需要在训练脚本中加几行来计算这些值。cliprange 表示裁剪强度参数 $\varepsilon$。和上面一样，返回负的目标值，这样最小化 loss 就会执行 gradient ascent。

要测试你的代码，请实现 [adapters.run_compute_policy_gradient_loss]。然后运行

uv run pytest -k test_compute_policy_gradient_loss_off_policy_gspo

并确保测试通过。

6.4 实验

我们现在已经具备了更新 grpo_train_step 方法以支持 off-policy 训练所需的组件。

Problem（grpo_train_step_off_policy）：Off-policy GRPO 训练步（2.5 分）

Deliverable： 更新你的 grpo_train_step 方法，使其支持 off-policy 参数（importance_reweighting_method、old_log_probs、cliprange）。该函数现在应支持完整的参数范围。

要测试你的代码，请实现 [adapters.run_grpo_train_step]。然后运行

uv run pytest -k test_grpo_train_step_off_policy

并确认测试通过。

让训练变成 off-policy 真的会让它更快吗？我们会在多大程度上失去稳定性？哪一种 reweighting/clipping 方法最好？我们将在下面的实验中探索这些问题。请保持超参数与你之前的 RL 运行一致，只是现在这些运行应当是 32× off policy（也就是说，对每个 inference batch 我们取 32 个 training steps）：

rollout_batch_size = 256
train_batch_size = 8
gradient_accumulation_steps = 1

另外，作为默认值，你可以使用以下裁剪范围值（即上面推导中的 $\varepsilon$），它们取自原始论文。可选地，也欢迎你自己调这些值。

注意，你还需要在训练脚本中加几行来计算 old_log_probs，以便传入你的 reweighted loss。

第 38 页

问题（grpo_experiments_off_policy）：比较不同 off-policy 算法的性能（8 B200 hrs）（10 分）

保持超参数与你标准 GRPO 训练运行时一致（学习率由你选择），运行以下变体，每个变体使用 4 个随机种子。你应当使用 zero-shot 的 r1_zero prompt。如果你愿意，也可以从上一节中挑选你最喜欢的 RL 算法变体，而不必坚持标准版本。

除了你已经在记录的指标之外，请务必也记录 clip fraction。

与你之前跑的完全 on-policy 的 GRPO 运行相比，这些方法表现如何？变成 off-policy 会如何影响训练的稳定性以及不同 run 之间的方差？两种 clipping 方法在 clip fraction 上有何差异，哪一种看起来更稳定？有没有哪些方法你认为会从更多超参数调优中受益？根据 run 之间的方差，你对你的结论有多大信心？

Deliverable：解说（commentary），辅以所引用指标的图。

7 尝试你自己的 policy gradient 估计器

问题（try_your_own）：尝试你自己的 policy gradient 估计器（10 分）

既然你已经看过文献中提出的各种 policy gradient 估计器，以及它们背后的一些理论，现在就该你来提出自己的了！以下是一些建议：

提出你自己的 policy gradient 估计器。然后在 OLMo-2-0425-1B 上、GSM8K 这一相同的 RL 设定中运行它，并与我们已尝试过的方法作为基线进行比较；请使用多个随机种子，并保持问题参数固定，以使结果具有可比性。请相对于目前尝试过的方法只改动一处（也就是说，本作业中至少应有一个先前涵盖的方法，使得你的方法相对于它只改了一件事）。

你所提出的改动背后的理论依据或直觉是什么？你提出的改动是否胜过基线？

Deliverable：解说，辅以所引用指标的图；如果数学推导有助于论证你的方法，也可以选择附上。

第 39 页

参考文献

材料

数据 / 仓库（只给链接，不下载 ✗）

按 D50：数据集与 starter code 只给链接 ✓（starter code 用 git clone 拿 ✓，数据集太大 ✗ 而且年年变 ✓）。这些链接故意不放在「材料」小节 ✓ —— 放了就会被 --fetch 下载 ✗。

自查