Target Policy Optimization: Decoupling What to Learn from How to Learn in RL

Paper: 2604.06159 Authors: Jean Kaddour Categories: cs.LG

The Gap

Policy gradient methods (PG, PPO, GRPO) answer two questions simultaneously: which completions should gain probability mass, and how should parameters move to realize that change. This coupling creates a problem—the update can overshoot or undershoot depending on learning rate, clipping, and optimizer hyperparameters. When rewards are sparse, this becomes critical: you might take a tiny step toward a rare good completion, or overshoot and lose it entirely. Prior methods treat this as a hyperparameter tuning problem. This paper asks: what if we separate the two questions?

Sparse Reward Problem:
  |
  v
Hypothesis: Coupling target selection + parameter update
            causes instability
  |
  v
Method: Decouple via explicit target distribution
        q_i ∝ p_i^old * exp(u_i)
        then fit policy by cross-entropy
  |
  v
Evidence: Tabular bandits, transformers, billion-param LLMs
          TPO matches baselines on easy tasks
          substantially outperforms on sparse reward
  |
  v
Conclusion: Explicit target + supervised fitting
            is more stable than implicit PG updates

The Increment

One sentence: Before, RL updates entangled target selection with optimization dynamics; after, we construct an explicit target distribution and fit to it via supervised learning.

Core Mechanism

TPO operates in two clean phases. First, given a prompt, sample completions from the current policy and score them with a reward model. Use these scores to construct a target distribution: q_i is proportional to the old policy probability times the exponential of the utility (advantage). This target distribution explicitly encodes which completions should gain mass.

Second, fit the policy to this target via cross-entropy loss. The gradient on sampled-completion logits becomes p^θ - q, which vanishes when the policy matches the target. No clipping, no KL penalties in the loss—the target distribution already incorporates the old policy, so you’re not drifting arbitrarily far. The update is a supervised learning step, not a policy gradient step.

Phase 1: Construct Target          Phase 2: Fit Policy
                                   
prompt --> sample --> score        target q_i ∝ p_i^old * exp(u_i)
           |                                |
           v                                v
    [c1, c2, ..., cn]              minimize: -sum q_i log p_i^θ
           |                                |
           v                                v
    [u1, u2, ..., un]              gradient: p^θ - q --> 0

Think of it like adjusting a recipe. Standard policy gradient methods taste the dish (get rewards), then immediately adjust the heat and ingredients (update parameters) in one motion—easy to overshoot. TPO instead writes down the ideal flavor profile first (target distribution), then carefully adjusts the recipe to match that profile (supervised fitting). The target is your reference point; you stop when you’ve matched it, not when some learning rate schedule says to stop.

Key Concepts

• Target Distribution as Reweighted Policy: The formula q_i ∝ p_i^old ** exp(u_i) is doing something subtle. It’s not just “make good things more likely”—it’s reweighting the old policy by exponentiating advantages. If a completion had 1% probability and high advantage, it might jump to 5% in the target. If it had 40% probability and low advantage, it might drop to 20%. The target is a concrete distribution you can sample from, visualize, and debug. Standard PG methods have an implicit target that emerges from the gradient; TPO makes it explicit.

• Cross-Entropy Fitting vs Policy Gradient: Policy gradient computes ∇log p(a|s) ** A(s,a) and steps in that direction. The step size is controlled by learning rate, and you hope it lands somewhere good. Cross-entropy fitting computes p^θ - q and steps to minimize that gap. The “step size” is implicit—you stop when p^θ ≈ q. This is why TPO is more stable under sparse reward: you’re not taking blind steps scaled by a hyperparameter; you’re fitting to a fixed target.

Framework Shift

Before (Policy Gradient):           After (TPO):

sample --> reward --> gradient      sample --> reward --> target dist
              |                                    |
              v                                    v
         ∇log p * A                    q ∝ p^old * exp(u)
              |                                    |
              v                                    v
      θ += lr * gradient              fit p^θ to q via CE
              |                                    |
              v                                    v
      (hope it's good)                (stop when p ≈ q)

Implicit target, explicit step      Explicit target, implicit step

From gradient-driven exploration to target-driven fitting, the core shift is replacing a scaled gradient step with a supervised learning objective.

Expert Assessment

Problem choice: Real gap. Sparse reward instability is a known pain point in RLHF and RL fine-tuning. The coupling of target selection and optimization has been discussed informally, but rarely formalized as a design choice to revisit. This sits at the intersection of practical LLM training and foundational RL theory—good positioning.

Method maturity: Elegant insight, not brute force. The exponential reweighting q ∝ p^old ** exp(u) is borrowed from control-as-inference literature, but applying it as an explicit intermediate target is clever. The simplicity is a strength—no new hyperparameters beyond what you’d already tune in PPO. However, the paper doesn’t deeply explore failure modes: what if the target distribution is multimodal and hard to fit? What if exp(u) creates extreme weights?

Experimental integrity: Baselines are fair (PG, PPO, GRPO, DG). The billion-parameter LLM experiments are convincing—sparse reward tasks show clear wins. But the paper leans heavily on “substantially outperforms” without error bars or significance tests in some plots. The tabular bandit results are clean, but those are toy problems. I’d want to see more ablation on the target construction—what if you use q ∝ p^old ** u instead of exp(u)?

Writing quality: The abstract and intro are crisp. The method section could use a diagram showing the two-phase structure more explicitly. The related work is thin—missing connections to KL-regularized RL, DPO, and other recent RLHF methods that also construct implicit targets. The conclusion is a missed opportunity to discuss when TPO might fail or when PG is preferable.

Verdict: weak accept — Solid idea with convincing evidence on a real problem, but needs deeper analysis of failure modes and stronger statistical rigor in experiments.

Takeaways

Steal the two-phase structure: separate “what should the model do” from “how do I update parameters to get there.” This applies beyond RL—any time you’re optimizing with noisy signals, consider constructing an explicit intermediate target. The exponential reweighting trick (q ∝ p * exp(u)) is a clean way to turn advantages into a distribution without clipping or thresholding. If you’re training LLMs with sparse feedback, TPO’s cross-entropy fitting is worth trying—it’s a drop-in replacement for PPO with potentially better stability.