← 返回 2026-04-16

目标策略优化(TPO):用交叉熵直接拟合带锚点的指数加权目标分布 Target Policy Optimization

Jean Kaddour 📅 2026-04-07 👍 23 2026-07-13 08:36
GRPO LLM 强化学习 PPO RLVR 强化学习 目标策略优化 策略梯度

把组内候选重加权变成闭式目标分布,用交叉熵直接拟合,梯度自动消失。

前置知识

策略梯度(Policy Gradient)

强化学习中的核心方法家族,通过对策略 $\pi_\theta(a|x)$ 的对数概率乘以回报估计(advantage)来构造损失 $L = -\mathbb{E}[A \log \pi]$,再用梯度下降更新参数。REINFORCE、PPO、GRPO 都属于这一家族。

TPO 的全部讨论都建立在策略梯度范式之上,理解 REINFORCE/PPO 的 update rule 才能体会到 TPO 把标量加权换成分布拟合到底改变了什么。

GRPO(Group Relative Policy Optimization)

DeepSeek 提出的大模型 RLVR 标准方法。对同一 prompt 采样 $K$ 个回答,用组内回报 z-score 作为优势 $A_i$,做 PPO 截断 surrogate $L = -\mathbb{E}[\min(rA,\text{clip}(r)A)]$,通常叠加反向 KL 惩罚。

TPO 的核心 baseline 就是 GRPO,理解它的截断 surrogate、组内归一化和 KL 锚定是看懂 TPO 与之本质差异的前提。

目标分布与交叉熵拟合(Target Distribution & Cross-Entropy)

监督学习中的标准范式:先在某种规则下构造一个目标概率分布 $q$(如 one-hot 标签),再用 KL 散度或交叉熵 $L = -\sum q_i \log p_\theta(y_i)$ 让模型概率 $p_\theta$ 去逼近它。TPO 的核心思想就是把这个范式搬到 RL 的组内候选集合上。

TPO 的全部新颖性都建立在把组内 $K$ 个候选当做一个离散动作集、对它们构造一个指数加权目标 $q$、再用交叉熵去拟合这个观察之上。

KL 正则化策略改进(KL-Regularized Policy Improvement)

REPS、MPO 等方法通过 KL 正则化目标得到软最优分布 $q \propto \pi_{\text{old}} \exp(u/\eta)$,是 trust region 思想源头。TPO 目标分布形式与之相同,但只作用在组内 $K$ 个候选上。

理解 TPO 的目标分布是 KL 正则化问题的精确解,才能读懂 Proposition 1 和为什么 target 在闭式下存在。

稀疏奖励(Sparse Reward)与信用分配(Credit Assignment)

稀疏奖励指智能体在长序列末尾才能获得一次非零回报(如完全匹配的 reward = 1,否则 0),中间步无信号。信用分配问题就是如何把这一份末梢信号合理地分摊到序列中每一步的决策上。

TPO 的论文标题虽覆盖广,但实验设计、对比和为什么会赢的论证几乎全部围绕稀疏奖励场景——理解这一困境才能体会 TPO 的实际价值。

研究动机

现有组内策略梯度方法(PPO、GRPO、DG 等)把想要怎样重新分配概率和参数如何动起来实现这一分配两件事纠缠在同一步梯度更新里,这带来三个具体痛点。第一,在稀疏奖励场景下几乎所有候选动作回报方差为 0,这些 all-fail 组既无信号也仍在贡献梯度噪声或不一致的优势估计,导致 GRPO 在 token reversal H=10 任务上 terminal reward 实验中退化到 50.4% 错误率,PPO 在 H=9 直接掉到 90.6%、DG 在 H=9 之后完全无法学习(>95%)。第二,标量加权的策略梯度缺乏不动点:即便策略已经达到目标分布附近,$\nabla L$ 仍可能非零,因此学习率、clip 范围、epoch 数等优化器超参会显著影响最终表现——GRPO 的 epoch ablation 显示在 epoch=2 时错误率爆炸到 37.6%,而 epoch=16 时才回落到 1.1%,极不单调。第三,组内信息在标量加权下被压扁——错误集中在一个混淆类别上的样本无法用一维 advantage 表达其结构化信息,限制了学习的精细度。

本文的目标是本文提出 Target Policy Optimization(TPO),其目标是把上述两个问题解耦:先用一个与奖励相关的规则在组内候选集合上构造一个明确的目标分布 $q_i \propto p_i^{\text{old}} \exp(u_i/\eta)$,其中 $u_i$ 是组内标准化的奖励分数;然后用交叉熵损失 $L_{\text{TPO}} = -\sum_i q_i \log p_\theta(y_i|x)$ 把策略 $\pi_\theta$ 拟到这个固定目标上。设计目标具体包括:让梯度在 $p_\theta = q$ 时自动消失、避免 importance ratio 和 clipping、在不引入 critic 的前提下让多 epoch 复用稳定,以及让组内的全失败组在 rollout snapshot 下贡献严格为零梯度。

与已有工作不同的是,此前的 KL 正则化方法(REPS、MPO、V-MPO)虽然求解的是相同形式的最优分布 $q \propto \pi_{\text{old}} \exp(u/\eta)$,但它们依赖 learned Q-function、需要对完整动作空间做受约束优化,工程开销大;AWR 同样用指数权重但把它当作 log-likelihood 上的固定标量,梯度不会自我消失;PMPO 划分 accepted/rejected 两堆并用冻结的 $\pi_{\text{ref}}$ 做正则,结构上更接近 DPO 风格。TPO 的独特切入角度是:只对组内 $K$ 个已评分候选做闭式目标构造(无需 critic、无需对完整动作空间优化)、只锚定到 $\pi_{\text{old}}$ 自身、用交叉熵一次性拟合整个组而非标量加权,从而把组结构信息完整保留下来。这与 Osband (2026) 的 DG 互补:DG 解决跨上下文梯度分配,TPO 解决组内结构化更新。

核心方法

TPO 的整体思路可以一句话总结:把 RL 中对一组候选做重加权这个动作从为每个候选算一个标量 advantage 再加到 log-prob 上,改成在组内候选集上构造一个指数加权分布、再用监督学习的交叉熵去拟合它。直觉上,这与语言模型做 SFT 时把 one-hot 标签当 target 几乎完全一样,只是 target 不再是 one-hot,而是用旧策略概率和奖励共同决定的软目标。技术路线上,TPO 先冻结当前策略为 $\pi_{\text{old}}$,对每个 prompt 采样 $K$ 个候选 $y_1, \dots, y_K$,用打分器 $S$ 给出原始分数 $s_i$ 并在组内做 z-score 标准化得到 $u_i$,然后按 $q_i \propto p_i^{\text{old}} \exp(u_i/\eta)$ 计算闭式目标,最后用 $-q \log p_\theta$ 形式的交叉熵做一步或多步梯度更新。

TPO 的核心创新点可拆为三件事。其一是把组内 $K$ 个候选看成一个独立的离散动作集,并在该集合上定义 KL 正则化策略改进问题 $\max_r \sum_i r_i u_i - \eta \, \text{KL}(r \| p^{\text{old}})$,其闭式解恰好就是目标分布 $q_i$,这意味着 TPO 既有理论保证(Proposition 1)又无需引入 critic;其二是 loss 的梯度是 $\nabla L = p_\theta - q$,形式上极简且 self-extinguishing——当 $p_\theta$ 匹配 $q$ 时梯度严格为零,从而摆脱 PPO 的 clipping、importance ratio、超参敏感性问题;其三是标准化的 $u_i$ 加上 $p^{\text{old}}$ 的锚定让 all-fail 组($u=0$)在 rollout snapshot 下贡献严格为零,天然消除稀疏奖励下大量无信号样本带来的噪声,这与 PMPO 锚定到独立 $\pi_{\text{ref}}$ 的做法有本质区别。

方法步骤详情

TPO 六步流程。Step 1:$\pi_{\text{old}} \leftarrow \pi_\theta$ 并 detach。Step 2:从 $\pi_{\text{old}}$ 采 $K$ 个候选。Step 3:$s_i = S(x, y_i)$。Step 4:组内 z-score 得 $u_i$,零方差映射 $u=0$。Step 5:$q_i \propto p_i^{\text{old}} \exp(u_i/\eta)$,$\eta=1$。Step 6:最小化 $L_{\text{TPO}} = -\sum_i q_i \log p_\theta(y_i|x)$。代码 10 行:`q = softmax(log_softmax(log_scores) + u/eta).detach()` 后用 `-(q * log_p).sum(-1).mean()` 算 loss。Prop 1:$q$ 是 KL 正则化在 $K$-simplex 的唯一极大点;固定 $q$ 时 $\nabla L = p_\theta - q$,故 $p_\theta = q$ 是不动点。

技术新颖性

TPO 的技术新颖性可从三方面评估。算法层面,它是第一个把 REPS/MPO 风格的指数加权目标分布直接搬到组内候选集上、并用纯交叉熵拟合的方法,省去了 critic 和对完整动作空间的优化;这一思路此前在 AWR 中以对 log-prob 乘以 $\exp(A/\beta)$ 的形式出现过,但 TPO 用 target distribution + cross-entropy 替换了它,因此 gradient $\nabla L = p_\theta - q$ 才具备 self-extinguishing 性质。理论层面,Proposition 1 同时给出 $q$ 的最优性(KL 正则化闭式解)和 $p_\theta = q$ 的不动点性,这两点共同保证了方法在不动点附近的稳定性。系统层面,作者设计了一系列诊断实验(gradient L2 norm 衰减曲线、all-fail 比例、per-candidate weight proxy)把为什么会赢拆解为梯度自我消失、信号自动聚焦于有信息组、多 epoch 复用稳定三个独立贡献,这是论文实验部分最值得借鉴的论证范式。

Implementation sketch: JAX 与 PyTorch 各五行实现 TPO 目标与损失。
Figure 2: Implementation sketch: JAX 与 PyTorch 各五行实现 TPO 目标与损失。

实验结果

实验覆盖七类场景。多上下文 bandit(图 4)TPO 对齐 CE oracle 方向最紧;$\beta_{TPO}$ 在 $p_n=0.1$ 时为 0.73,远高于 $\beta_{DG}=0.09$、$\beta_{GRPO}=0.33$。MNIST(图 5)TPO 1600 步到 5%、最终 2.9% 最低。Token reversal V=16 时 TPOtoken 102 步达 1% error,GRPOtoken 148、PPO 259、DG 393(表 1)。Terminal reward H=10(表 3):TPO 7.4% vs GRPO 50.4% vs PPO/DG >95%。LLM(图 10):GSM8K ~85–87% 但 TPO 早收敛 10 步;graph coloring TPO ~0.96、GRPO ~0。机理(图 11–16):TPO 梯度收敛后衰减到 0;K sweep TPO 从 K=4 的 8.9% 单调到 K=64 的 0.36%,GRPO K=64 反弹到 5.6%。

Steps to 1% error on Token Reversal (bag-of-tokens reward, K=8).
Table 1: Steps to 1% error on Token Reversal (bag-of-tokens reward, K=8).
Steps to 1% error on task variations (K=8 token candidates).
Table 2: Steps to 1% error on task variations (K=8 token candidates).
Exact-match error (%) on terminal reward, prompt-matched vs interaction-matched.
Table 3: Exact-match error (%) on terminal reward, prompt-matched vs interaction-matched.
Comparison of policy optimization methods: Group / Critic / Fixed ref. 三列属性。
Table 4: Comparison of policy optimization methods: Group / Critic / Fixed ref. 三列属性。
Single-context symmetric bandit (K=100, B=100): TPO 与 DG 收敛最快,GRPO 与 PG 早期 plateau。
Figure 3: Single-context symmetric bandit (K=100, B=100): TPO 与 DG 收敛最快,GRPO 与 PG 早期 plateau。
Multi-context bandit (N=100, K=10, exact gradients): 所有方法最终收敛,CE oracle 最快,TPO 最贴近 CE。
Figure 4: Multi-context bandit (N=100, K=10, exact gradients): 所有方法最终收敛,CE oracle 最快,TPO 最贴近 CE。
MNIST contextual bandit: TPO 收敛最快、最终错误率最低;高混淆样本上 TPO 获得最大超额增益。
Figure 5: MNIST contextual bandit: TPO 收敛最快、最终错误率最低;高混淆样本上 TPO 获得最大超额增益。
Token Reversal (bag-of-tokens reward, K=8 token candidates), V ∈ {2,4,8,16}。
Figure 6: Token Reversal (bag-of-tokens reward, K=8 token candidates), V ∈ {2,4,8,16}。
Task variations, prompt- and interaction-matched.
Figure 7: Task variations, prompt- and interaction-matched.
Terminal reward, prompt- and interaction-matched, H ∈ {7,8,9,10}。
Figure 8: Terminal reward, prompt- and interaction-matched, H ∈ {7,8,9,10}。
Anchor and target-matching ablations (H ∈ {7,8,10}, V=2, K=8, 20 seeds)。
Figure 9: Anchor and target-matching ablations (H ∈ {7,8,10}, V=2, K=8, 20 seeds)。
LLM RLVR on Qwen3-1.7B (top) and DeepSeek-R1-Distill-Qwen-1.5B (bottom).
Figure 10: LLM RLVR on Qwen3-1.7B (top) and DeepSeek-R1-Distill-Qwen-1.5B (bottom).
TPO 梯度自我消失 vs GRPO 持续不衰减 (H=8, V=2, K=32)。
Figure 11: TPO 梯度自我消失 vs GRPO 持续不衰减 (H=8, V=2, K=32)。
稀疏奖励下大部分组不含信号,TPO 最快消除全失败组 (H=8, V=2, K=32)。
Figure 12: 稀疏奖励下大部分组不含信号,TPO 最快消除全失败组 (H=8, V=2, K=32)。
Group-size sensitivity sweep (H=8, V=2, epochs=4, K ∈ {4,8,16,32,64}, 30 seeds)。
Figure 13: Group-size sensitivity sweep (H=8, V=2, epochs=4, K ∈ {4,8,16,32,64}, 30 seeds)。
Zero-variance masking (H=8, V=2, K=32, epochs=4, 30 seeds)。
Figure 14: Zero-variance masking (H=8, V=2, K=32, epochs=4, 30 seeds)。
Multi-epoch extraction (H=8, V=2, K=32).
Figure 15: Multi-epoch extraction (H=8, V=2, K=32).
Epoch-count ablation (H=8, V=2, K=32, 30 seeds).
Figure 16: Epoch-count ablation (H=8, V=2, K=32, 30 seeds).
查看结构化数据
任务指标本文基线提升
Token Reversal V=16 (token-level grouping, K=8) 达到 1% 错误率所需步数(越少越好) TPOtoken = 102 步 GRPOtoken=148 / PPO=259 / DG=393 TPO 是最快基线 GRPO 的 1.45×,是 DG 的 3.85×
Terminal Reward Token Reversal H=10 (prompt-matched) Exact-match 错误率(%) TPO = 7.4% GRPO=50.4% / PPO>95% / DG>95% / GRPO(no KL) 无法学习 TPO 相对最优 baseline 降低 6.8× 错误率
Multi-context bandit (N=100, K=10) 达到 CE oracle 方向对齐的余弦相似度 TPO 余弦偏差接近 0(全训练过程) GRPO/DG 在训练中段与 CE 方向显著偏离 TPO 的更新方向在所有 RL 风格方法中最贴近期望方向
MNIST Contextual Bandit (20 seeds) Classification 错误率(%) TPO = 2.9%(最低) DG ≈ 3.x% / Group PG = 7.2% / GRPO = 5.9% / PG = 5.3% TPO 相对次优降低约 0.5–0.7 个绝对百分点
LLM RLVR: Reasoning Gym Graph Coloring (Qwen3-1.7B) Train mean score(0–1) TPO ≈ 0.96 GRPO ≈ 0.0(300 步内不学习) TPO 在 GRPO 完全失败的任务上达到接近满分的训练分数
Group-size sweep (H=8, V=2, K=64, 30 seeds) Final exact-match 错误率(%) TPO = 0.36% GRPO = 5.6%(K=32 时 4.4%,K=64 反而回升) TPO 在 K=64 上比 GRPO 低 15.6× 错误率且单调下降

局限与改进

作者在 Section 6 坦诚讨论了 TPO 的边界。候选质量依赖:TPO 只能在给定的 $K$ 个候选上重分配概率,采样多样性差则 target 没信息量,与 GRPO 共享此约束。组大小成本:在无法单次 forward 拿到所有动作评分的场景,TPO 与 GRPO 同样需要 $K$ 倍 rollout,不减少 environment interaction。标准化放大小差异:组内 reward 方差极小时 z-score 把微差放大成接近 one-hot 的 target,触发与 GRPO 相同的 difficulty bias(Dr. GRPO 问题),是明确的开放方向。规模有限:LLM 评估仅 1.5–1.7B 与 3 个任务(GSM8K、graph coloring、Knights & Knaves),MATH/AIME/7B+ 留作 future work。off-policy 不直接适用:依赖 target $q$ 在每步 detach,真正 off-policy 时需要 Retrace/V-trace 风格修正,作者明确指出这是未来工作。

独立分析的弱点

独立审视后 TPO 仍有可改进之处。其一,标准化稳定性边界未严格刻画:作者承认 z-score 在极小方差下放大小差异,只用一个 setting 做 masking,没给出 reward 方差多大时 target 不可信的判定准则;改进方向是引入可学习 target temperature 或对 $u$ 做 clipping/tanh 限幅。其二,目标分布依赖打分器 calibration:$p^{\text{old}}$ 极偏时 $q \approx p^{\text{old}}$、训练停滞;可叠加 prompt-level 全局 baseline 或 RLOO 风格 leave-one-out 修正。其三,没有跨 prompt 的 KL 锚点:TPO 只锚定 $\pi_{\text{old}}$,在 IcePop 类异步场景下可能不稳定;可允许在 target 上叠加弱的 reference-anchor 项。其四,缺少 $K$ vs batch size 的 compute-matched ablation,未来需更系统评估。

未来方向

作者明确提出的方向:把 TPO 扩展到 7B+ 大模型与 MATH/AIME 等更难的 RLVR 基准;在 off-policy/异步引擎场景下引入 Retrace 或 V-trace 风格的修正;把 TPO 与 IcePop、ScaleRL 等 off-policy 稳定技术正交组合;探索对低方差组更鲁棒的目标分布构造。基于成果可延伸的方向还包括:把 target distribution 的闭式解思路推广到 multi-reward 场景(GDPO 走 objective 修正路线,TPO 可在 group 层面对 reward 向量构造 joint target);与 MaxRL 的高阶目标修正正交组合(MaxRL 改优化什么,TPO 改怎么优化);探索 token-level 与 sequence-level target 的层次化构造;用 TPO 替代 GRPO 的先 trust region 后 PPO clip 思路,看能否在 LLM long-CoT RL 中带来更稳定的 entropy 控制。

复现评估

可复现性较高。代码开源:仓库 https://github.com/JeanKaddour/tpo 提供 token reversal、MNIST bandit、tabular bandit 等小实验。LLM 栈:verl,模型 Qwen3-1.7B 与 DeepSeek-R1-Distill-Qwen-1.5B,任务 GSM8K + Reasoning Gym 的 graph coloring 和 Knights & Knaves,K=16 rollouts/prompt。小规模:transformer Optax Muon(lr=$10^{-3}$, B=100);tabular bandit $\alpha=0.1$;MNIST Adam(lr=$10^{-3}$, B=100) 训练 10000 步。种子:20/30 seeds。算力:toy 单卡即可,LLM 需多卡(具体 GPU 数与 wall-clock 未给出),对资源有限研究者 LLM 部分较难完整复现。算法本身只依赖 10 行 PyTorch/JAX 实现,中等经验 RL 研究者 1–2 周可复现非 LLM 实验。