Uniform-Correct Policy Optimization: Breaking RLVR’s Indifference to Diversity

Paper: arXiv:2605.00365 Code: AnamikaLochab/UCPO Code reference: main @ 93ce5064 (2026-05-04)

1. Motivation (研究动机)

RLVR / GRPO 在数学推理上的成功，主要体现在单次采样能否答对，也就是 Pass@1 的提升。但很多数学题并不是“只有一个唯一解轨迹”的任务：同一个最终答案，往往可以来自不同的分步推导、不同的中间变形路径、不同的公式拆解顺序。论文抓住的问题是，现有 RLVR 目标把这些“都能过 verifier 的正确轨迹”统统视为同一个桶里的正样本，却没有规定模型应该怎样在这些正确轨迹之间分配概率质量。这会带来一个很具体但过去经常被忽略的现象：Pass@1 上升，不代表 Pass@K 一定也上升。如果策略逐渐把概率集中到某一条已经常见、已经容易采到的正确轨迹上，那么多次采样时模型虽然可能更稳定地产出这一条解，但对其他同样正确的推理路径覆盖会越来越差。

对只看 single-shot accuracy 的设置来说，这似乎不是大问题；但对 test-time sampling、self-consistency、多轨迹 reranking、以及需要从多个正确候选中再做下游筛选的场景来说，这种塌缩会直接伤害可用性。作者把这个现象命名为 diversity collapse，并且强调它不是“训练偶然噪声”或“某个实现细节不稳定”这么简单。论文的核心立场是：塌缩来源于 RLVR 目标本身的结构性缺口。只要目标只关心“落在正确集合里”，却不关心“落在正确集合里的哪一部分、以什么分布落进去”，那么训练动态就可能自发把概率质量推向少数高频正确解，而把其余正确解边缘化。

1.1 现有 RLVR 优化的盲点

在 verifier-based RL 中，最自然的目标是最大化 $J (π) = E_{x \sim D, y \sim π (\cdot ∣ x)} [R (x, y)],$ 其中 $R (x, y) \in {0, 1}$ 通常表示 response 是否被规则或判题器接受。对单个 prompt $x$ 而言，只要定义正确解集合 $Y^{+} (x)$ ，这个目标等价于尽可能提高 $Z_{π} (x) = \sum_{y \in Y^{+} (x)} π (y ∣ x) .$ 这里的问题非常尖锐：只要概率质量都落在 $Y^{+} (x)$ 上，目标值就已经最优。至于这部分质量是均匀分在十条正确解上，还是 99% 压在一条正确解上，原始 RLVR 目标完全不区分。

这意味着传统 RLVR 在目标层面只有“correctness 约束”，没有“correct-set geometry 约束”。作者并不是说 correctness 不重要，而是说：当任务天然存在多个正确轨迹时，仅仅把错误质量推走是不够的，正确集合内部的分布形状同样决定了最终的采样表现。

1.2 为什么单看 Pass@1 会掩盖问题

如果把 RLVR 看成一种“把模型输出推向 verifier 接受区域”的机制，那么 Pass@1 的提升很容易掩盖内部塌缩。因为在早期训练中，一条略占优势的正确轨迹会更频繁地被采样到，继而获得更多正梯度；这条轨迹的概率升得越快，后续越容易继续被采样，从而形成正反馈闭环。对 Pass@1 而言，这种“快速锁定一条对的路”甚至可能看起来是件好事。

但 Pass@K 衡量的是另一件事：给模型 $K$ 次机会，它能否覆盖更多不同但仍然正确的候选。一旦正确集合内部的概率分布过于尖锐，多次采样得到的就只会是“同一种正确答案的不同表面改写”或者干脆是同一条轨迹的重复。此时：

单样本正确率可能仍在上升。
多样正确轨迹覆盖率却在下降。
reranking / majority voting 能利用的信息量变少。
验证器允许的“多个可行思路”被训练过程自己抹平。

因此作者把论文问题明确成：如何在不牺牲 RLVR 对 correctness 的追求前提下，显式约束正确集合内部的概率分布，使其不要塌缩到少数模式上。

1.3 论文要解决的不是“探索不足”，而是“正确集内部的探索失衡”

很多读者第一反应会把这个问题理解成一般性的 exploration problem，但论文的切入点更细。它并不是说模型没有探索错误和正确之间的边界，也不是说采样温度太低、全局熵太小。UCPO 关注的是一个条件化后的问题：

当我们已经只看 verifier 认可的 response 时，这些正确 response 之间是否仍然维持了足够均匀的概率分布？

这个条件化视角非常关键。因为如果直接在全输出空间上鼓励熵，模型也可能把概率撒向明显错误的轨迹；这会伤害 correctness mass。UCPO 试图做到的是更精细的修正：

错误集合仍然应该被压低。
正确集合的总质量仍然应该增大。
但正确集合内部不能被单一模式垄断。

也就是说，这篇论文研究的不是“要不要 diversity”，而是“在正确的前提下，要不要保留多样性”。这比全局熵正则更贴近 reasoning RL 的实际目标。

1.4 为什么这个问题值得研究

这个问题一旦解决，收益并不只局限于数学竞赛题上的 Pass@64 数值更高。更广义地看，它会影响所有依赖多次采样的推理场景：

self-consistency 需要多条互相独立的正确推理轨迹，而不是同一路径重复采样。
best-of- $K$ 或 verifier reranking 依赖候选池覆盖更多解法模式。
需要可解释性或人类审阅时，多样正确轨迹能展示模型是否真正掌握了问题，而不是只记住一种答法。
多解问题上的稳健性，往往来自对正确解集的分散建模，而不是对单一高分解的极端 sharpening。

因此，UCPO 实际上在问一个更基础的问题：RLVR 训练出的“会做题”模型，究竟是在学一个正确解，还是在学一整个正确解集合？

1.5 论文的理论起点

作者没有停留在经验观察层面，而是进一步指出：如果把 RLVR 目标写成“正确集合上的总概率质量”，那么任何支持集完全落在 $Y^{+} (x)$ 内的策略，都是这个目标的 stationary point。形式上，这解释了为什么 collapsed-correct policy 和 uniform-correct policy 在原始目标下都同样“合法”。这一步很重要，因为它把经验上的 “GRPO 似乎会塌缩” 升级为理论上的 “GRPO 没有理由不塌缩”。既然目标函数本身对正确集合内部的概率分布 indifferent，那么最终落到哪一种正确分布形状，就只能由采样偏置与梯度动态决定，而不是由目标显式约束。也正因为如此，作者后续提出 UCPO 时，不是试图给 GRPO 打一个经验补丁，而是试图把“正确集合内部应更均匀”这件事直接写进优化目标。

2. Idea (核心思想)

核心洞察是：RLVR 失败的根源不是它不会奖励正确答案，而是它对“多个正确答案之间应该如何共享概率质量”完全无感。 UCPO 的做法不是在整个输出空间盲目加熵，而是把分布整形精确限制在“已经被 verifier 判对的集合”内部，让这些正确解的条件分布尽量接近均匀。

换句话说，UCPO 把原来的优化目标从“只要压到正确集合里就行”改成了“既要压到正确集合里，也要避免在正确集合内部过度偏置到少数模式上”。这一步听起来像小改动，但它改变的是目标的最优解结构：collapsed-correct 不再和 uniform-correct 等价。

2.1 一句话概括论文的新意

如果用一句话概括，这篇论文的新意是：

在 verifier-based RL 中，把 diversity 约束从全局输出空间收缩到正确集合内部，并把它做成一个与 correctness 兼容的条件均匀化目标。

这和常见的 entropy regularization、temperature sampling、trajectory-level exploration 的最大区别在于，UCPO 并不把“多样性”当作单独目标，而是把它理解为“正确集合上的分布形状问题”。

2.2 为什么是 Uniform-Correct，而不是 High-Entropy

论文并没有直接主张“熵越大越好”。如果策略在整个输出空间上提升熵，那么错误答案的概率也可能一起升高，这显然和 RLVR 的目标冲突。作者真正想要的是一个条件化的目标：

在错误集合上，概率应继续被压制。
在正确集合上，概率总质量应尽可能高。
在正确集合内部，分布应尽量均匀而不是塌缩。

因此最理想的结构不是 global high-entropy policy，而是 Uniform-Correct Policy： $π (y ∣ x) = ⎩ ⎨ ⎧ \frac{1}{∣ Y ^{+} ( x ) ∣}, 0, y \in Y^{+} (x) y \in / Y^{+} (x)$ 这个 policy 既保证所有质量都在正确集合内，又保证正确集合内部没有模式垄断。作者在第 5 节正是围绕这一点展开理论分析，并从 robustness 与 entropy-regularized optimality 两个角度说明它为什么是理想结构。

2.3 与已有方法的根本区别

论文对比的对象并不只是 vanilla GRPO，还包括两类常见缓解思路。第一类是 全局 token / sequence entropy regularization。这类方法通过鼓励更平滑的输出分布来避免过度 sharpening，但问题是它没有区分正确与错误 token / trajectory，因而可能把概率重新分给错误解。UCPO 则只在正确集合上整形，因此目标更“窄”、更精准。第二类是 trajectory-level exploration 或 reward reshaping，例如面向 Pass@K 的奖励、协方差控制、或基于 baseline 的 advantage 改写。这些方法的共同点是想让多样性在训练中“顺便出现”。而 UCPO 的不同点是：它先明确指出最优 policy 结构应该是什么，再反推一个使该结构成为唯一最优解的目标。因此 UCPO 和现有方法的根本差异不是“又加了一个正则项”，而是：

现有方法多在全空间上调探索强度。
UCPO 在正确集合这个条件空间上直接规定目标几何。
现有方法通常只经验性抑制 collapse。
UCPO 先理论刻画 collapse，再构造对应的修正目标。

2.4 关键创新到底在哪里

UCPO 的关键创新可以拆成三层。第一层是 问题重新表述。作者把 RLVR 中的 diversity 问题从“采样更随机一点”改写成“条件分布是否均匀”。这一表述更贴合 verifier-based RL，因为 verifier 自然把输出空间切成“正确 / 错误”两块。第二层是 最优结构刻画。论文没有直接从工程上试 τ、试正则，而是证明 Uniform-Correct Policy 在 robustness 和小 τ 的 entropy-regularized objective 下都具有特殊地位。这让后面的算法不是拍脑袋设计。第三层是 实现方式足够轻量。UCPO 并不需要额外 reward model，也不需要离线枚举所有正确解，更不需要大改 RL pipeline。它复用 GRPO 已经会收集的 verifier-accepted rollouts，在这些样本上做 advantage reweighting，就能把梯度更多地分给低概率但仍正确的响应。

2.5 直觉上为什么它应该有效

如果把 GRPO 想象成“给所有正确样本统一发奖金”，那么最容易被采到的正确样本会因为出现次数更多而拿走更多总奖金。UCPO 的做法更像是“奖金总额不变，但对经常出现的正确样本打折，对稀有但正确的样本加权补贴”。这样做的结果不是放松 correctness，而是把正确集合内部的更新预算重新均摊。

这个直觉也解释了为什么 UCPO 的工程代价不大。它没有增加新的训练阶段，也没有引入昂贵的新监督信号，而是改变同一批正确 rollout 在反向传播中的相对权重。论文多次强调：UCPO preserves the total advantage mass of GRPO but reallocates it toward underrepresented correct responses。真正改变的是“梯度流向”，不是“训练框架骨架”。

2.6 与 GRPO 的关系

UCPO 不是要替换 GRPO 的一切，而是对 GRPO 做一个有方向性的目标修正。把它理解为 “GRPO + conditional uniformity shaping” 最合适：

当 $τ = 0$ 时，UCPO 退化回 GRPO。
当 $τ > 0$ 时，正确集合内部的不均匀分布会受到惩罚。
$τ$ 越大，越偏向均匀正确分布。
但训练仍然依赖 GRPO 同样的 on-policy sampling、verifier 打分和 PPO-style update。

这种设计让 UCPO 有一个很实际的优点：对现有 RLVR codebase 来说，它更像“可插拔的 update rule 变体”，而不是一套需要整体重写的训练范式。

3. Method (方法)

3.1 Overall framework

Figure 3 解读：这张图对比了 GRPO、UCPO 和全局 entropy regularization 的训练动态。GRPO 会持续把概率质量推向少数正确解；UCPO 则把更新信号重新分配到低概率但仍正确的响应上，从而在不牺牲 correctness mass 的前提下保持更高的 conditional diversity。

UCPO 的训练流程和 GRPO 基本一致：先对每个 prompt 采样一组 responses，用 compute_score 作为 verifier 判断对错，再基于这些 rollout 估计 advantage，并通过 PPO/GRPO 风格的 surrogate objective 更新策略。区别只在于，UCPO 不再把所有正确响应视为“只要是正样本就完全等价”，而是显式关心正确集合内部的条件分布形状。

把整套框架拆开看，可以分成四步：

从当前策略 $π_{θ} (\cdot ∣ x)$ 为每个 prompt 采样 $K$ 条 response。
用 verifier 将 response 切分为正确集合与错误集合，并累积 token-level outcome reward。
在同一 prompt 内，对正确 rollout 的 advantage 进行再分配，让低概率正确解得到更强的提升信号。
继续沿用现有 RLVR 的优化与并行训练骨架，不引入新的判别器或离线标注。

因此 UCPO 的“新模块”并不是一套新的采样器或新的 reward model，而是一套围绕 conditional uniformity 设计的目标与梯度估计规则。

3.2 Problem setup and notation

论文围绕单个 prompt $x$ 建立分析。令输出空间中的正确解集合为 $Y^{+} (x)$ ，错误解集合为 $Y^{-} (x)$ 。对于一个策略 $π (\cdot ∣ x)$ ，定义它落在正确集合上的总概率质量为 $Z_{π} (x) = \sum_{y \in Y^{+} (x)} π (y ∣ x) .$ 进一步定义在正确集合上的条件分布 $q_{π} (y ∣ x) = \frac{π ( y ∣ x )}{Z _{π} ( x )}, y \in Y^{+} (x) .$ 这个 $q_{π}$ 是理解整篇论文最关键的量。因为：

$Z_{π} (x)$ 决定“有多少质量被放到正确集合里”。
$q_{π} (\cdot ∣ x)$ 决定“这些正确质量在不同正确解之间如何分配”。

原始 RLVR 主要优化的是前者，UCPO 则在保留前者的同时，额外塑造后者。

3.3 RLVR objective 的结构性 indifferent

RLVR 对单个 prompt 的目标可以写成 $J_{x} (π) = E_{y \sim π (\cdot ∣ x)} [R (x, y)] = \sum_{y \in Y^{+} (x)} π (y ∣ x) = Z_{π} (x) .$ 若策略满足 $supp (π (\cdot ∣ x)) \subseteq Y^{+} (x)$ ，则 $J_{x} (π) = 1$ 已经达到最大值。论文的 Theorem 4.1 正是在形式上指出：所有这种策略都是 stationary point。它们之间的差异只体现在条件分布 $q_{π}$ 上，而原始目标对此完全 indifferent。这带来两个直接后果。第一，RLVR 目标无法区分下面两种策略：

所有正确解等概率分布的 uniform-correct policy。
几乎所有质量都压在一条正确解上的 collapsed-correct policy。

第二，最终训练轨迹会被随机初始化、on-policy sampling、以及正样本频次放大效应牵着走。也就是说，塌缩不是目标“失败”了，而是目标根本没有惩罚它。

3.2 Collapse mechanism

Figure 1 解读：作者用 toy RLVR 环境验证 collapse 的四步链条。高概率正确解更容易被采样到，采样偏置又放大了它的梯度贡献，随后更新进一步抬高它的概率，最终把其他正确解挤出支持集。

论文把这个 collapse 过程拆成一个自强化循环。

第一步：sampling bias

从当前策略采样 $K$ 个 rollout 时，概率更高的正确解自然更容易重复出现。若某个正确解 $y$ 的采样次数记为 $N_{y}$ ，则 $N_{y} \sim Binomial (K, π_{θ} (y ∣ x)), E [N_{y}] = K \cdot π_{θ} (y ∣ x) .$ 这意味着哪怕两个解都正确，只要一开始一个略高于另一个，它在 minibatch 中就更常被看见。

第二步：gradient amplification

在 outcome-only 的 GRPO 设定下，正确响应共享同一个正向 advantage，错误响应则拿到低或负的更新。于是 minibatch 中“谁出现得多”，谁的总梯度贡献就更大。论文把每个正确解的期望梯度贡献写成与 $K \cdot π_{θ} (y ∣ x)$ 成正比，这一步直接把采样频率转成了更新强度。

第三步：ratio divergence

一旦高频正确解在当前步获得更强更新，其下一轮概率会进一步抬高。作者分析了 dominant / minority 正确解之间对数概率比的演化，并指出它表现为带正漂移的过程。直观上说，领先者会因为领先而进一步领先。

第四步：irreversible pruning

当少数正确解的概率降到很低后，它们在有限 $K$ 的采样下几乎不再出现。此时即便这些解理论上仍然正确，也拿不到梯度信号，恢复机会变得极小。塌缩于是从“偏好某个模式”升级为“训练上再也看不见其他模式”。

这四步解释了为什么单纯依赖 on-policy positive update 的 RLVR 在多解任务上会天然向模式集中。

3.5 Controlled experiment 的作用

论文不仅给出理论分析，还设计了 toy 控制实验验证上面的四步链条。Figure 1 与 Figure 2 的作用不是单纯展示现象，而是逐层验证：

初始化偏置会被采样频次放大。
采样频次差异会转化成 expected update 差异。
update 差异会不断拉大正确解之间的概率比。
低概率正确解一旦跌破采样阈值，就进入近似不可恢复状态。

这一步很关键，因为它说明 UCPO 需要修正的是“采样与梯度之间的结构耦合”，而不是某个偶然的 optimizer 或 seed 问题。

3.3 GRPO 的目标与 indifference

GRPO 采用 clipped surrogate objective： $L_{GRPO} (θ) = E_{x \sim D, {y_{i}}_{i = 1}^{K} \sim π_{θ_{old}} (\cdot ∣ x)} [\frac{1}{K} \sum_{i = 1}^{K} min (r_{i} (θ) A_{i}, clip (r_{i} (θ), 1 - ϵ, 1 + ϵ) A_{i})],$ 其中 $r_{i} (θ) = \frac{π _{θ} ( y _{i} ∣ x )}{π _{θ_{old}} ( y _{i} ∣ x )} .$ 论文关心的不是 clipped PPO 形式本身，而是其中 advantage 的语义。对同一 prompt 的 outcome reward 来说，只要 rollout 被 verifier 接受，它就进入“正确样本”这一类，并共享正向更新趋势。于是 GRPO 的优化倾向变成：

优先提高已经采样到的正确样本的概率。
不会额外补偿那些“当前没被采到但理论上同样正确”的解。
采样到的频率越高，累计更新越大。

这正是 objective indifference 在采样近似下的工程化体现。

3.6 Optimal policy characterization

UCPO 不是直接从经验 loss 出发，而是先问：如果我们想让 policy 在多解任务上更稳健，最理想的结构是什么？论文给出两个角度。第一是 robustness。如果部署时面临轻微的 verifier 扰动、候选重加权、或需要在多个正确解之间稳定覆盖，那么把概率全部压在少数模式上显然更脆弱；相反，正确集合内部越均匀，越不容易因为某条路径波动而整体性能下降。第二是 entropy-regularized optimality。作者分析 $E_{y \sim π (\cdot ∣ x)} [R (x, y)] + τ H (π (\cdot ∣ x))$ 这一类熵正则目标，并证明在足够小的 $τ$ 区间里，Uniform-Correct Policy 会成为唯一最优解。这个结论的意义不在于“以后都要用全局熵”，而在于说明：从最优结构角度看，正确集合上的均匀分布并不是拍脑袋偏好，而是具有明确理论合理性的目标形态。

3.4 UCPO objective

UCPO 在 GRPO 上增加条件均匀性项： $L_{UCPO} (x) = lo g Z_{θ} (x) - τ KL (u (\cdot ∣ x) ∥ q_{θ} (\cdot ∣ x)),$ 其中 $Z_{θ} (x)$ 是策略在正确集合上的总质量， $q_{θ} (\cdot ∣ x)$ 是正确集合上的条件分布， $u (\cdot ∣ x)$ 是这个正确集合上的均匀分布。 $τ$ 控制 uniformity penalty 的强弱。

把这个目标拆开理解：

$lo g Z_{θ} (x)$ 继续鼓励“把更多总概率质量放到正确集合中”。
$KL (u ∥ q_{θ})$ 惩罚“正确集合上的条件分布偏离均匀”。
两项组合后，目标既保留 correctness，又引入 correct-set 内部的分布整形。

论文在 Theorem 6.1 中证明：对于任意 $τ > 0$ ，Uniform-Correct Policy 是该目标的唯一最优解。这正是 UCPO 与原始 RLVR / GRPO 的本质区别。原始目标把 collapsed-correct 和 uniform-correct 一视同仁；UCPO 则显式打破这种 indifferent。

3.7 为什么是反向 KL： $KL (u ∥ q_{θ})$

这里使用的是 $KL (u ∥ q_{θ})$ ，而不是 $KL (q_{θ} ∥ u)$ 。直觉上，前者更像“要求策略覆盖住所有正确解”，因为如果某个正确解在 $q_{θ}$ 下概率极小，就会受到更强惩罚。它强调的是 coverage，而不只是平均熵变大。

这与论文的问题设定高度一致：作者不是只想让 correct set 内部“更平滑”，而是希望低概率正确解不要被剪枝。用 $u$ 作为参照分布，正好把“每个正确解都应被留住”编码进目标。

3.5 Gradient reweighting

Figure 4 解读：UCPO 本质上是在 correct set 内做重加权。图里高概率正确解的权重被压低，低概率正确解的权重被抬高；所以它不是增加总梯度，而是重排相同的 advantage budget。

对这个目标求梯度后，论文得到的更新不是“额外再做一次采样”，而是对同一批正确样本施加不同权重。其核心思想可概括为：高概率正确解的更新系数被压低，低概率正确解的更新系数被放大，从而把同一个 prompt 的正梯度预算重新分配给 underrepresented correct responses。

论文给出的权重形态可以写成一种“从均匀分布与当前条件分布之间插值”的形式。直觉上，更新不再只看某条正确轨迹是否被 verifier 接受，还看它在当前正确集合内部是否已经占了太多概率。于是：

已经占优的正确轨迹仍然得到正更新，但边际收益变小。
稀有正确轨迹得到更强补偿，防止继续掉出支持集。
错误轨迹仍然不会因为 uniformity penalty 被主动鼓励。

这也是它和 global entropy regularization 的核心差异。全局 entropy 会把质量推向整个输出空间，容易误伤 incorrect outputs；UCPO 只在 verifier-accepted samples 上工作，因此更贴近 RLVR 的目标。

3.8 与 released code 的对应关系

论文层面把 UCPO 写成条件均匀性正则目标，而 released code 采用的是 advantage estimator 分支的实现。仓库中 scripts_c/run_ucpo_1.5b.sh 与 scripts_c/run_ucpo_7b.sh 都把

algorithm.adv_estimator=iq
algorithm.tau=0.2
algorithm.alpha=0.0

作为 UCPO 训练入口。随后在 [ucpo/trainer/ray_trainer.py](/Users/bytedance/Library/CloudStorage/OneDrive-个人/paper_notes/notes/LLM & VLM/RL & Post-Training/Uniform-Correct Policy Optimization.md) 对应的 compute_advantage(...) 分支中，当 adv_estimator == AdvantageEstimator.Inv_q 时，会调用 core_algos.compute_dgrpo_advantage_outcome_invq_l(...)，并将 tau、alpha 传入。

这说明 released code 的实现重点不在显式构造一个单独的 KL loss 模块，而是在 advantage 估计阶段完成“correct set 内重加权”。从工程视角看，这和论文的目标是等价取向：

论文表达的是 objective-level 的优化目标。
代码实现的是 gradient estimator / advantage reweighting。
两者都在做一件事：把正梯度从高频正确解转移给低频正确解。

3.9 Verifier and reward path

仓库中的 ucpo/trainer/fn_score.py 给出了最直接的 verifier 入口。它用 math_verify.parse 与 verify 判断模型输出和 ground-truth 是否一致；若正确则返回 {"score": 1, "acc": 1}，否则返回 {"score": 0, "acc": 0}。

这意味着 UCPO 并没有依赖额外 learned reward model。它仍然属于典型的 RLVR 设置：

reward 是可验证的、离散的 outcome reward；
训练重点在于如何利用这些 outcome signal；
UCPO 的新意完全落在“如何组织正确 rollout 的更新结构”上。

这也解释了为什么论文能把讨论聚焦到正确集合内部，而不用处理 reward model calibration 带来的额外不确定性。

3.10 End-to-end training flow

结合论文与 released code，可以把 UCPO 的训练流程总结为：

def ucpo_train_step(policy, prompt_batch, verifier, tau, alpha):
    responses = sample_k_responses(policy, prompt_batch, k=8)
    scores = [verifier(resp) for resp in responses]
    correct_mask = [s["score"] == 1 for s in scores]
    old_log_probs = compute_old_log_probs(policy, responses)
    advantages = compute_dgrpo_advantage_outcome_invq_l(
        token_level_rewards=scores,
        old_log_prob=old_log_probs,
        tau=tau,
        alpha=alpha,
        correct_mask=correct_mask,
    )
    loss = clipped_grpo_surrogate(policy, responses, advantages)
    loss.backward()
    optimizer.step()
    return loss

这段伪代码的重点不是逐字符还原源码，而是说明 UCPO 没有改动整个 RL skeleton：采样、验证、估计 advantage、再走 PPO/GRPO 更新，主干都与已有 pipeline 一致。

3.11 Key components by function

如果按模块功能拆分，UCPO 方法里真正重要的组件有四个。

组件 A：Verifier-defined correct set

输入是 prompt 与 response，输出是 0/1 正确性标签。它定义了后续所有条件分布、uniformity penalty 与 advantage reweighting 的作用范围。没有这一层，correct-set uniformity 就无从谈起。

组件 B：On-policy grouped sampling

同一 prompt 采 $K$ 条 response，是识别 collapse 的必要条件。因为只有把同一个 prompt 的多个正确轨迹放在一起，才能观察“正确集合内部是否被某一模式主导”。如果只做单样本更新，correct-set geometry 的信息会直接丢失。

组件 C：Conditional uniformity shaping

这是论文的核心模块。它不新增 reward，而是改变对正确 rollout 的相对 weighting。若去掉这一项，训练就会退回 GRPO 的模式频次放大逻辑。

组件 D：GRPO-compatible policy update

UCPO 的实用价值之一，是它保留了已有 GRPO 训练管线的主体结构。若这一层被改写成完全不同的优化器或额外外环，方法的迁移成本会明显上升；论文显然刻意避免了这一点。

3.12 What would fail if a component were removed

把每个组件拿掉后，预期失败模式其实很清楚。

去掉 verifier：就失去 correct set 定义，无法区分“正确内部均匀化”和“全局熵扩散”。
去掉 grouped sampling：无法估计同一 prompt 下不同正确轨迹的相对覆盖，collapse 结构不可见。
去掉 uniformity shaping：训练退回 GRPO，自强化塌缩机制重新出现。
去掉 GRPO-compatible update：方法虽可能仍成立，但会失去论文强调的低成本、可插拔优势。

3.13 Code-to-paper mapping

Code reference: main @ 93ce5064 (2026-05-04) — pseudocode and mapping based on this commit

Paper Concept	Source File	Key Class/Function
Hydra entry and training bootstrap	`ucpo/main_run.py`	`main()`, `run_ppo()`, `get_custom_reward_fn()`
Verifier-based reward	`ucpo/trainer/fn_score.py`	`compute_score()`
UCPO/GRPO advantage dispatch	`ucpo/trainer/ray_trainer.py`	`compute_advantage()`
UCPO advantage branch	`ucpo/trainer/ray_trainer.py`	`AdvantageEstimator.Inv_q`, `core_algos.compute_dgrpo_advantage_outcome_invq_l()`
Default hyperparameter base	`ucpo/config/ppo_trainer.yaml`	`algorithm.tau`, `algorithm.alpha`, `custom_reward_function`, rollout defaults
1.5B experiment override	`scripts_c/run_ucpo_1.5b.sh`	`algorithm.adv_estimator=iq`, `tau=0.2`, `lr=5e-6`, `trainer.n_gpus_per_node=4`
7B experiment override	`scripts_c/run_ucpo_7b.sh`	`algorithm.adv_estimator=iq`, `tau=0.2`, `lr=2e-6`, `trainer.n_gpus_per_node=8`

这张映射表也说明了一点：论文中许多训练数字不能直接从 ucpo/config/ppo_trainer.yaml 默认值读取，因为实际实验是由 scripts_c/run_ucpo_1.5b.sh 与 scripts_c/run_ucpo_7b.sh 显式 override 的。例如 YAML 默认 data.train_batch_size=1024、max_prompt_length=512、max_response_length=512、algorithm.tau=1.0，而实验脚本把它们覆盖成 64 / 1024 / 3072 / 0.2。因此本笔记中所有训练配置数字都应以 launch script 为准，而不是以基线 YAML 默认值为准。

3.14 Pseudocode for key components

Verifier:

def compute_score(sequence: str, ground_truth: str) -> dict:
    answer = parse(sequence)
    target = ground_truth if "\\boxed" in ground_truth else f"\\boxed{{{ground_truth}}}"
    target = parse(target)
    if verify(answer, target):
        return {"score": 1, "acc": 1}
    return {"score": 0, "acc": 0}

GRPO-style grouped rollout:

def sample_group(policy, prompt, k: int = 8):
    responses = [policy.sample(prompt) for _ in range(k)]
    log_probs = [policy.log_prob(prompt, resp) for resp in responses]
    return responses, log_probs

UCPO advantage branch:

def compute_ucpo_advantage(batch, tau: float, alpha: float):
    return core_algos.compute_dgrpo_advantage_outcome_invq_l(
        token_level_rewards=batch["token_level_rewards"],
        response_mask=batch["response_mask"],
        old_log_prob=batch["old_log_probs"],
        index=batch["uid"],
        tau=tau,
        alpha=alpha,
    )

Policy update:

def update_policy_with_ucpo(policy, batch, tau: float, alpha: float):
    advantages, returns, stats = compute_ucpo_advantage(batch, tau=tau, alpha=alpha)
    loss = clipped_surrogate_loss(
        policy=policy,
        responses=batch["responses"],
        advantages=advantages,
        old_log_probs=batch["old_log_probs"],
    )
    loss.backward()
    optimizer.step()
    return stats

3.15 方法直觉总结

如果只用一句话总结 UCPO 的方法直觉，那就是：

让“正确但稀有”的轨迹在训练中不要因为一开始采样频次低，就永远失去被继续学习的机会。

GRPO 的问题不是它不会奖励正确答案，而是它默认“被采到多少次”就近似等于“值得被学多少”。UCPO 通过 correct-set conditional uniformity 打破这个近似，把正确解内部的 long-tail 模式重新拉回训练视野。这也是为什么它能提升 Pass@K 与 equation-level diversity，同时又不必显著牺牲 Pass@1。

3.16 论文公式与 released code 实现差异

论文把 UCPO 叙述为“在正确集合上加条件均匀性惩罚”，而 released code 里真正落地的是 algorithm.adv_estimator=iq 分支，并通过 core_algos.compute_dgrpo_advantage_outcome_invq_l() 来实现 advantage reweighting。也就是说：

论文写的是 objective-level 形式： $lo g Z_{θ} (x) - τ KL (u ∥ q_{θ})$ 。
代码写的是 estimator-level 形式：在 grouped correct rollout 上做重加权。

两者目标一致，但表述层级不同。本笔记保留这一差异，是因为它直接关系到“读论文的人如何理解方法”和“看源码的人如何定位实现”。

3.6 Source code mapping

Code reference: main @ 93ce5064 (2026-05-04) — pseudocode and mapping based on this commit

Paper Concept	Source File	Key Class/Function
Verifier-based reward	`ucpo/trainer/fn_score.py`	`compute_score()`
Hydra entrypoint	`ucpo/main_run.py`	`main()` / `run_ppo()` / `get_custom_reward_fn()`
UCPO / IQ advantage branch	`ucpo/trainer/ray_trainer.py`	`AdvantageEstimator.Inv_q` → `compute_dgrpo_advantage_outcome_invq_l()`
Default training config	`ucpo/config/ppo_trainer.yaml`	`algorithm.tau`, `algorithm.alpha`, `custom_reward_function`
1.5B launch script	`scripts_c/run_ucpo_1.5b.sh`	`algorithm.adv_estimator=iq`, `tau=0.2`, `alpha=0.0`
7B launch script	`scripts_c/run_ucpo_7b.sh`	`algorithm.adv_estimator=iq`, `tau=0.2`, `alpha=0.0`

伪代码（按仓库实现抽象）：

def compute_score(sequences: str, ground_truth: str) -> dict:
    answer = parse(sequences)
    gt = parse(ground_truth if "\boxed" in ground_truth else f"\boxed{{{ground_truth}}}")
    return {"score": 1, "acc": 1} if verify(answer, gt) else {"score": 0, "acc": 0}

def run_ppo(config):
    compute_score = get_custom_reward_fn(config)
    reward_fn = reward_manager_cls(tokenizer=tokenizer, num_examine=0, compute_score=compute_score)
    val_reward_fn = reward_manager_cls(tokenizer=tokenizer, num_examine=1, compute_score=compute_score)
    trainer = RayPPOTrainer(..., reward_fn=reward_fn, val_reward_fn=val_reward_fn)
    trainer.fit()

def select_advantage_branch(batch, config):
    if config.algorithm.adv_estimator == "iq":
        advantages, returns, stats = compute_dgrpo_advantage_outcome_invq_l(
            token_level_rewards=batch["token_level_rewards"],
            response_mask=batch["response_mask"],
            old_log_prob=batch["old_log_probs"],
            index=batch["uid"],
            tau=config.algorithm.tau,
            alpha=config.algorithm.alpha,
        )
    return advantages, returns, stats

3.7 论文公式与 released code 实现差异

论文把 UCPO 叙述为“在正确集合上加条件均匀性惩罚”，但 released code 里真正走的是 algorithm.adv_estimator=iq 分支，并通过 compute_dgrpo_advantage_outcome_invq_l() 实现加权。换句话说，论文层面是 objective 级描述，代码层面是 advantage reweighting 的实现；两者目标一致，但表述层次不同。

4. Experimental Setup (实验设置)

数据集：训练集是 train_data_10k.parquet，验证集是 valid_data.parquet；论文正文说明训练池来自 DeepScaleR / MATH / AIME / AMC / Omni-MATH 的 10K competition-style math problems。
基线：GRPO、KL-Cov、Pass@K Training、FGRPO、Ent-Reg、Ent-Adv。
指标：Pass@1 和 Pass@64；Pass@K 表示采样 $K$ 次后至少一个样本被 verifier 接受的比例。另一个指标是 equation-level diversity，衡量 correct rollouts 内数学表达式的去重程度。
训练配置：1.5B 用 4 张 GPU，7B 用 8 张 H100；train_batch_size=64，max_prompt_length=1024，max_response_length=3072，rollout.n=8，ppo_mini_batch_size=512，1.5B 学习率 5e-6，7B 学习率 2e-6，tau=0.2，alpha=0.0，entropy_coeff=0。

5. Experimental Results (实验结果)

5.1 Main results

Figure 7 解读：UCPO 在所有模型和 benchmark 上都把 Pass@K 曲线整体抬高，尤其在 K 增大后优势更明显，说明它恢复的是“多样正确解覆盖率”，而不是只优化单样本命中率。

Figure/Table 1 解读：主表展示了 3 个模型族在 5 个数学 benchmark 上的 Pass@1 / Pass@64。论文结论是：UCPO 在保持 Pass@1 竞争力的同时，平均 Pass@64 持续优于 GRPO，并在 AIME24 上取得最高 +10 点级别提升。

Figure 2 解读：该图表明 collapse 并不是偶然噪声，而是从初始化偏置开始经采样偏置和梯度放大逐步积累出来的结构性现象。也因此，单纯加 global entropy 不足以稳定修复。

5.2 Diversity

Table 2 的 equation-level diversity 显示，UCPO 在 AIME / MATH / Olympiad 上都最高，AIME 相比 GRPO 有约 45% 的相对提升，说明它确实恢复了正确解集合内部的覆盖率。

具体数值上，UCPO 的 diversity 分别达到 AIME 0.218、AMC 0.217、MATH 0.196、Olympiad 0.228，平均值 0.215；而 GRPO 分别是 0.150 / 0.197 / 0.173 / 0.208，平均 0.182。这说明 UCPO 的收益不是来自把更多概率撒向错误答案，而是让 verifier 认可的正确轨迹本身更丰富。

5.3 Ablation and sensitivity

Figure 5 解读：全局 entropy regularization 的收益很小而且不稳定，说明“在整个输出空间鼓励熵”并不能精确修复正确集合内的塌缩。

Figure 6 解读：UCPO 的 $τ$ 只要在中等范围内就能稳定改善 diversity，同时维持 correctness；这比全局 entropy 的超参更稳。

Figure 8 解读：UCPO 和 GRPO 的每步训练耗时几乎一致，说明这项改动在工程上几乎是“零额外成本”的。

5.4 Limitations

论文没有提出明显的安全或通用性限制，但它的实验范围仍然主要是数学 reasoning；因此它更像是一个针对 RLVR 多解覆盖问题的专门修正，而不是通用 RL 目标的替代品。

Paper Notes

探索

Uniform-Correct Policy Optimization: Breaking RLVR's Indifference to Diversity

Uniform-Correct Policy Optimization: Breaking RLVR’s Indifference to Diversity

1. Motivation (研究动机)

1.1 现有 RLVR 优化的盲点

1.2 为什么单看 Pass@1 会掩盖问题

1.3 论文要解决的不是“探索不足”，而是“正确集内部的探索失衡”

1.4 为什么这个问题值得研究

1.5 论文的理论起点

2. Idea (核心思想)

2.1 一句话概括论文的新意

2.2 为什么是 Uniform-Correct，而不是 High-Entropy

2.3 与已有方法的根本区别

2.4 关键创新到底在哪里

2.5 直觉上为什么它应该有效

2.6 与 GRPO 的关系

3. Method (方法)

3.1 Overall framework

3.2 Problem setup and notation

3.3 RLVR objective 的结构性 indifferent

3.2 Collapse mechanism

第一步：sampling bias

第二步：gradient amplification

第三步：ratio divergence

第四步：irreversible pruning

3.5 Controlled experiment 的作用

3.3 GRPO 的目标与 indifference

3.6 Optimal policy characterization

3.4 UCPO objective

3.7 为什么是反向 KL：KL(u∥qθ​)

3.5 Gradient reweighting

3.8 与 released code 的对应关系

3.9 Verifier and reward path

3.10 End-to-end training flow

3.11 Key components by function

组件 A：Verifier-defined correct set

组件 B：On-policy grouped sampling

组件 C：Conditional uniformity shaping

组件 D：GRPO-compatible policy update

3.12 What would fail if a component were removed

3.13 Code-to-paper mapping

3.14 Pseudocode for key components

3.15 方法直觉总结

3.16 论文公式与 released code 实现差异

3.6 Source code mapping

3.7 论文公式与 released code 实现差异

4. Experimental Setup (实验设置)

5. Experimental Results (实验结果)

5.1 Main results

5.2 Diversity

5.3 Ablation and sensitivity

5.4 Limitations

目录

3.7 为什么是反向 KL： $KL (u ∥ q_{θ})$