PaperlayerView exact revision on arXiv

Source-linked AI summary

Target Policy Optimization

Jean Kaddour

arXiv:2604.06159v1cs.LG

TL;DR

RL methods must decide both which sampled completions deserve more probability and how parameters should realize that redistribution, a coupling that can be fragile under sparse reward. TPO separates these decisions by constructing an old-policy-anchored target and fitting the policy to it by cross-entropy. Across bandits, transformer tasks, and billion-parameter LLM RLVR, it matches established baselines on easier or dense-reward settings and substantially outperforms them under sparse reward.

  • Problem

    Standard policy-gradient methods entangle desired probability redistribution with optimizer mechanics, making learning fragile especially when reward is sparse.

  • Method

    TPO constructs a target distribution from old-policy probabilities and standardized scores, then fits the policy to that target by cross-entropy.

  • Results

    Across tabular bandits, transformer sequence tasks, and LLM RLVR, TPO matches baselines on easier tasks and substantially outperforms them under sparse reward.

  • Takeaways & Limitations

    Separating desired redistribution from optimizer realization can make updates more robust across dense- and sparse-reward settings.

  • Takeaways & Limitations

    LLM-scale experiments use 1.5–1.7B-parameter models on three tasks; larger models and harder benchmarks remain future work.

Abstract

from arXiv · show

In RL, given a prompt, we sample a group of completions from a model and score them. Two questions follow: which completions should gain probability mass, and how should the parameters move to realize that change? Standard policy-gradient methods answer both at once, so the update can overshoot or undershoot depending on the learning rate, clipping, and other optimizer choices. We introduce \emph{Target Policy Optimization} (TPO), which separates the two questions. Given scored completions, TPO constructs a target distribution $q_i \propto p_i^{\,\mathrm{old}} \exp(u_i)$ and fits the policy to it by cross-entropy. The loss gradient on sampled-completion logits is $p^θ- q$, which vanishes once the policy matches the target. On tabular bandits, transformer sequence tasks, and billion-parameter LLM RLVR, TPO matches PG, PPO, GRPO, and DG on easy tasks and substantially outperforms them under sparse reward. Code is available at https://github.com/JeanKaddour/tpo.

1 Introduction

Target Policy Optimization separates choosing the desired probability redistribution from fitting the policy to it. It constructs a closed-form target over sampled completions and evaluates this approach across increasingly difficult RL settings.

  • Standard policy-gradient methods couple probability redistribution with optimizer mechanics, making learning fragile under sparse reward.
  • TPO applies reweight-then-fit to finite candidate sets sampled in group-based RL.
  • The target distribution is available in closed form without a critic or dual optimization.
  • TPO constructs q_i ∝ p_i^old exp(u_i) from behavior-policy probabilities and standardized scores, then fits the policy by cross-entropy.
  • The cross-entropy gradient vanishes exactly when the policy matches the target.
  • Across tabular bandits, contextual bandits, transformer tasks, and LLM RLVR, TPO matches policy-gradient baselines on easier tasks and outperforms them under sparse reward.

2 Target Policy Optimization

TPO samples and scores candidate completions, standardizes within-group scores, constructs a temperature-controlled target anchored to the old policy, and fits the current policy by cross-entropy. Its update has a unique stationary distribution at the target, while standardization makes the redistribution depend on relative within-group performance rather than arbitrary score units.

  • For each context, TPO samples K candidates from π_old and scores them with a scalar scorer.
  • TPO standardizes scores within each group, mapping the zero-variance case to u = 0.
  • The target tilts the frozen old-policy distribution toward higher-scoring candidates using temperature η, set to 1 in the experiments.
  • The policy is fitted to the target by minimizing cross-entropy, with q detached from the computation graph in the on-policy implementation.
  • Standardization prevents identical rankings with different score scales from producing drastically different targets and largely removes the need to tune η.
  • The target is the unique maximizer of the KL-regularized objective when p_i^old > 0 for every sampled candidate.
  • The logit gradient is p_θ − q, so the unique stationary distribution over sampled candidates is p_θ = q.

3 Experiments

Across tabular bandits, neural contextual bandits, transformer sequence tasks, and LLM RLVR, TPO generally matches baselines on easier settings and gains its clearest advantages as reward becomes sparse or tasks become harder. Its target-matching update and old-policy anchor support faster convergence, better cross-context allocation, and stronger sequence-level performance.

  • Single-context bandit: TPO and DG converge fastest in the single-context bandit, while PG and GRPO plateau at higher error as policies concentrate.TPO remains closest to the oracle policy-gradient direction throughout training.
  • Multi-context bandit: TPO stays closest to the CE oracle direction in the multi-context bandit and finishes with the lowest error among RL updates after an early transient.Its coefficient allocates more normalized update budget to hard contexts, with βTPO=0.73 at pn=0.1 versus 0.09 for DG and 0.33 for GRPO.
  • Neural policy learning: TPO reaches 5% error at step 1,600 versus 2,200 for DG and finishes at 2.9% error on the MNIST contextual bandit.At step 2,000, its surplus true-class gain reaches 0.073 in the highest wrong-class-concentration bin, while DG remains slightly negative.
  • Dense sequence reward: At V=16, TPOtoken reaches 1% error at step 102, ahead of GRPOtoken at 148, PPO at 259, and DG at 393.The methods use the same environment interactions per step, while TPOtoken and GRPOtoken exploit within-prefix candidate groups.
  • Task and reward variants: TPOtoken reaches 1% error first on all eight bag-of-tokens variants, 2–6× faster than the runner-up, and reaches it on all four sequential-reward tasks within budget.Under sequential reward, DG converges more slowly, while GRPOtoken and PPO fail to converge on any task.
  • Ablations and sparse credit assignment: At H=10, TPO reaches 7.4% error while every ablation exceeds 99%, and the gaps widen with sequence length.Removing the old-policy anchor, replacing target matching with scalar weighting, or removing GRPO’s KL penalty is consistently harmful.
  • LLM RLVR: On LLM RLVR, TPO learns about 10 steps earlier than GRPO to reach 50% GSM8K accuracy, while harder Reasoning Gym tasks show larger final-score gaps.On graph coloring, TPO reaches ∼0.96 while GRPO remains near zero; on R1-Distill-1.5B, TPO reaches ∼0.96 versus ∼0.81.

4 What explains TPO’s gains under sparse reward?

TPO’s sparse-reward gains reflect several reinforcing properties: its gradient approaches a fixed point, it concentrates updates on informative groups, and it supports stable multi-epoch reuse. These properties yield smoother sensitivity to group size and epoch count than GRPO.

  • Gradient self-extinguishing: TPO’s gradient decays to near zero after convergence, whereas GRPO maintains persistent gradients even after its error plateaus at 12.7%.TPO converges near episode 300; GRPO’s policy continues moving rather than settling near a fixed point.
  • Informative-group allocation: Roughly 90% of groups are all-fail initially, but TPO makes these groups exactly neutral on the rollout snapshot when score variance is zero.Their standardized scores are u = 0, so q = pold and the first-epoch grouped-loss contribution is zero.
  • Informative-group allocation: TPO rapidly reduces the all-fail fraction, concentrating nontrivial target structure on groups containing at least one successful candidate.GRPO leaves a larger residual all-fail fraction, while TPO drives it near zero quickly.
  • Group-size sensitivity: TPO’s final error improves smoothly from 8.9% at K=4 to 0.36% at K=64, while GRPO reaches 4.4% at K=32 before worsening to 5.6% at K=64.The sweep jointly changes candidate coverage and grouped-signal sharpness.
  • Zero-variance masking: Masking zero-variance groups raises GRPO’s final error from 6.3% to 29.7%, while TPO reaches 0.05% without masking.The result indicates that zero-variance groups can later provide an anchor toward the rollout snapshot in multi-epoch training.
  • Multi-epoch reuse: TPO remains stable across epoch counts with final error below 2.3%, whereas GRPO is strongly non-monotonic, reaching 37.6% at 2 epochs and 1.1% at 16.With four epochs, TPO reaches 0.2% error at episode 400 versus 1.1% with one epoch, while both eventually reach below 0.1%.

5 Related work

TPO belongs to a family of target-matching and policy-optimization methods, but differs in how it constructs and fits targets over candidate groups. Related methods vary in their use of critics, scalar weights, reference policies, regression losses, and objective corrections.

  • Target-matching methods: REPS, MPO, and V-MPO use exponential tilting but require a critic or value estimate, whereas TPO obtains its target in closed form on sampled candidates.TPO’s target solves a KL-regularized improvement problem restricted to the candidate set.
  • Target-matching methods: AWR uses fixed exp(A/β) weights on log-likelihood, while TPO’s cross-entropy gradient self-extinguishes when the policy matches its target.The distinction is between scalar weighting and fitting a target distribution.
  • Group-based policy gradients: RLOO and GRPO score multiple candidates but use scalar-weighted policy gradients, whereas TPO fits a distribution over the candidate simplex.Recent GRPO variants remain scalar-weighted policy-gradient methods while addressing specific failure modes.
  • Regression and preference methods: REBEL uses squared loss on paired log-probability ratios, while TPO uses cross-entropy over a full candidate group.Both construct targets from rewards and the behavior policy but differ in loss and structure.
  • Regression and preference methods: PMPO partitions candidates and anchors to a frozen reference policy, whereas TPO keeps one soft target over the group and anchors only to πold.Offline pairwise methods are more distant because TPO is online, setwise, and scorer-agnostic.
  • Objective-level corrections: GDPO and MT-GRPO correct GRPO’s objective for multi-reward and multi-task settings, while TPO changes how within-context signals become updates.The paper characterizes this distinction as orthogonal to those objective-level corrections.

6 Limitations

TPO’s limitations concern candidate quality, rollout cost, score standardization, and evaluation scale. Its relative gains remain untested on larger models and harder benchmarks.

  • Low-diversity or uniformly poor candidate groups yield uninformative targets, limiting TPO’s ability to redistribute probability effectively.
  • Sequence settings without a critic require K rollouts per context, so TPO does not remove the group-based rollout cost.More aggressive rollout reuse would create an off-policy regime that may require Retrace- or V-trace-style corrections.
  • Tiny within-group score variance can make z-scoring amplify numerical differences and produce excessively sharp targets.A group scored (0.001, 0, 0) is given as an example.
  • The LLM experiments use 1.5–1.7B-parameter models on three tasks, leaving 7B+ models and harder benchmarks such as MATH and AIME for future work.Whether TPO’s relative gains persist at larger scale remains open.

7 Conclusion

TPO separates target redistribution from optimizer realization by fitting the policy to a target distribution over scored candidates. Across tested settings, it matches baselines on dense-reward tasks and substantially outperforms them under sparse reward.

  • TPO builds a target distribution on scored candidates and fits the policy to it by cross-entropy instead of using scalar-weighted policy gradients.
  • Across tabular bandits, neural bandits, transformers, and billion-parameter LLM RLVR, TPO matches PG, PPO, DG, and GRPO on dense-reward tasks.
  • TPO substantially outperforms PG, PPO, DG, and GRPO under sparse reward across the tested settings.
  • The authors plan to test TPO on larger models.

B Multi-context tabular weighting derivation

In the one-hot multi-context tabular case, TPO’s target-matching update shares the correct-action direction with other methods but uses a distinct cross-context coefficient. The coefficient is flatter across current correctness probabilities than DG and GRPO.

  • For each context, the tabular derivation uses the correct action, current policy probability, and one-hot correct-action vector to express updates in logit space.
  • All exact updates share the within-context direction e_y − π and differ only in the scalar coefficient multiplying it.For TPO, this scalar form emerges after constructing the target and simplifying q_n − π_n in the one-hot setting.
  • TPO forms a target that multiplies correct-versus-incorrect odds by a fixed factor; for A = 10, λ = exp(10/3) ≈ 28.
  • At p_n = 0.1 and A = 10, β_TPO = 0.73, versus 0.09 for DG and 0.33 for GRPO.TPO’s coefficient is much flatter in p_n and therefore closer to CE’s equal-weight allocation.

C MNIST single-example logit updates

The MNIST derivation compares how single-example updates preserve information from sampled actions. TPO can use the detailed wrong-class distribution, whereas PG and related scalar methods generally reduce the signal to correct-versus-incorrect geometry.

  • The MNIST analysis considers one labeled example, its 10-class policy, correct-class probability, and one-hot class vectors in logit space.
  • Expected updates are taken over sampled actions, with surrogate-loss coefficients treated as stop-gradient constants.
  • Batch-standardized MNIST REINFORCE couples examples through minibatch mean and standard deviation but introduces no new within-example geometry.
  • TPO standardizes a one-hot sampled score so a positive sample has u_a = 3 and other classes have u_i = −1/3 for K = 10.After standardization, only the sign of the sampled score matters.
  • Unlike PG, GRPO, and DG, TPO directly suppresses a sampled wrong class after failure, redistributing that mass across the remaining logits.
  • PG, conditional single-sample GRPO, and Group PG preserve only a scalar correct-versus-incorrect signal in expectation, while TPO retains detailed wrong-class information.When wrong classes are nearly symmetric, TPO reduces to a scalar multiple of e_y − π.

D Temperature robustness

TPO’s temperature controls the sharpness of its target distribution, with strong performance across a broad range of values. Only the highest tested temperature materially slows convergence.

  • Score standardization uses an effective temperature of η=1 in q_i ∝ p_i^old · exp(u_i/η).
  • η ∈ {0.25, 0.5, 1, 2, 4} was evaluated on the token-reversal task across 10 seeds.
  • TPO performance is robust across a 16× temperature range, with only η=4 meaningfully slower.
  • All temperatures from 0.25 to 2.0 reached 1% error within 141 episodes, while η=4 degraded substantially.

E Multi-epoch DG instability

TPO, PPO, and GRPO constrain policy movement for multi-epoch rollout reuse, whereas DG does not bound per-step policy shifts. Consequently, repeated DG epochs are highly sensitive and usually worsen sparse-task performance.

  • TPO’s KL-anchored target, PPO clipping, and GRPO’s KL penalty stabilize multi-epoch reuse and improve extraction from rollout batches.
  • Four-epoch DG reached 48.3% final error versus 2.0% for the standard one-epoch update on reverse-copy transformer RLVR.
  • Across eight prompt-matched token-reversal variants, four-epoch DG was worse in 7 of 8 settings.
  • DG’s largest sequential-task regressions were flip from 0.07% to 4.56% and reverse flip from 0.00% to 0.82%.
  • The experiments use one DG gradient epoch per rollout batch, the most favorable setting for DG.
  • Removing GRPO’s KL penalty caused collapse under sparse terminal reward, with error increasing during training.

G LLM RLVR implementation details

The LLM RLVR experiments use a common verl-based training setup, with task-specific rewards and matched TPO and GRPO runs differing only in policy loss.

  • All LLM RLVR experiments use verl with AdamW at learning rate 10^-5, batch size 16, and 4×A100-80GB GPUs.
  • GSM8K uses exact-match rewards, graph coloring uses quasi-binary native scores, and Knights & Knaves uses partial-credit scores.
  • GSM8K adds LoRA with rank 32 and a KL penalty λ_KL=10^-3 to both TPO and GRPO.
  • Paired TPO and GRPO runs are otherwise identical and differ only in policy loss.
Loading 2604.06159v1…