Target Policy Optimization
| Authors | Jean Kaddour |
| Year | 2026 |
| HF Upvotes | 22 |
| arXiv | 2604.06159 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
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 Target Policy Optimization (TPO), which separates the two questions. Given scored completions, TPO constructs a target distribution q_i propto p_i^{,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.
Engineering Breakdown
Plain English
Target Policy Optimization (TPO) addresses a fundamental problem in reinforcement learning from human feedback (RLHF): when you have scored completions, how do you update the policy without overshooting or undershooting the target? Standard policy-gradient methods conflate two separate decisions—which completions should get more probability and how much to move the parameters—into one step, making the update sensitive to hyperparameters like learning rate and clipping. TPO decouples these by first constructing an explicit target distribution based on the old policy and reward scores, then fitting the new policy to this target via cross-entropy loss. The method was validated on tabular bandits, transformer sequence tasks, and billion-parameter LLMs, showing that the gradient p^θ - q elegantly vanishes once the policy matches the target.
Core Technical Contribution
The core novelty is the separation of policy update into two independent, interpretable steps: target construction and policy fitting. Rather than combining reward optimization and policy movement into a single gradient step (as in PPO or DPO), TPO first computes a target distribution q_i proportional to p_i^old * exp(u_i), where u_i are the completion rewards and p_i^old is the old policy. The policy is then optimized to match this target via standard cross-entropy loss, yielding a gradient of p^θ - q that naturally decays to zero at convergence. This decomposition removes the tight coupling between the step size, clipping mechanisms, and optimization dynamics that plague existing methods, making the training process more stable and interpretable.
How It Works
The algorithm takes a batch of completions scored by a reward model and proceeds in two phases. First, for each completion i, compute the target probability q_i = (p_i^old * exp(u_i)) / Z, where p_i^old is the old policy's probability under the language model, u_i is the normalized reward, and Z is a partition function for normalization. Second, optimize the new policy parameters θ to minimize cross-entropy between the model's output distribution p^θ and the target q: L = -Σ q_i log(p^θ_i). The loss gradient with respect to logits is simply p^θ - q, which means the update is proportional to how far the current policy is from the target. Once p^θ converges to q, the gradient vanishes completely, providing a natural stopping point without requiring explicit clipping or learning-rate tuning to prevent overshoot.
Production Impact
For engineers deploying RLHF pipelines, TPO offers three practical advantages: (1) reduced hyperparameter sensitivity—you no longer need to carefully tune PPO's clipping ranges and learning rates to avoid divergence, (2) clearer interpretability—the target distribution is explicit and can be logged and monitored, making it easier to debug training failures, and (3) simpler implementation—the cross-entropy loss is standard and gradient computation is straightforward. In a production LLM fine-tuning system, this could reduce iteration time during hyperparameter search and make it easier to A/B test different reward models or scoring functions without worrying about downstream optimizer instability. The trade-off is computational: you must maintain and compute the old policy's log-probabilities alongside the new policy, adding memory overhead and inference cost, and the method has only been tested up to billion-parameter scale, so scaling properties to trillion-parameter models remain unclear.
Limitations and When Not to Use This
TPO assumes access to well-calibrated reward scores u_i and relies on the old policy p_i^old being meaningful, which breaks down early in training or when the reward model is poorly trained. The method's stability on very large models (10B+ parameters) with sparse reward signals or noisy human feedback annotations has not been demonstrated. Additionally, the construction of q_i via exponential weighting is sensitive to reward scale and may require careful normalization; if rewards are poorly normalized, the target distribution can collapse to a few high-reward completions or spread too uniformly. The paper does not address how to handle distribution shift when the old policy is stale or when using data from multiple reward models, which is common in practice. Finally, the gradient vanishing at convergence, while clean theoretically, means you may need additional regularization (e.g., KL divergence to base model) to prevent mode collapse or catastrophic forgetting.
Research Context
TPO builds on the RLHF pipeline established by InstructGPT/ChatGPT and improves upon recent policy-gradient variants like PPO (Schulman et al., 2017) and DPO (Rafailov et al., 2023) by removing the need for clipping and explicit KL penalties. It sits in the broader trend of making RL for LLMs more stable and interpretable: prior work like ILQL and GRPO attempted similar decoupling but TPO is cleaner and more theoretically justified. The explicit target distribution connects to older ideas in batch RL and offline RL (e.g., conservative Q-learning) where constructing a stationary target avoids extrapolation error. This work opens the door to studying the sample complexity and convergence rates of policy fitting in the LLM regime, and whether other losses (e.g., Jensen-Shannon divergence, Hellinger distance) could serve as alternatives to cross-entropy for target matching.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::