🤖 machine learning

Value-Free Policy Optimization via Reward Partitioning

This paper introduces Reward Partition Optimization (RPO), a simple and scalable single-trajectory preference learning method that eliminates the need for value function estimation by normalizing rewards through prompt-level distributions, thereby achieving superior alignment, diversity, and stability compared to existing baselines like DRO and KTO.

Original authors: Bilal Faye, Hanane Azzag, Mustapha Lebbah

Published 2026-06-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Bilal Faye, Hanane Azzag, Mustapha Lebbah

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot chef how to cook the perfect meal. In the past, there were two main ways to do this, and both had some major headaches.

The Old Ways: The "Taste Test" and the "Gourmet Judge"

The "Taste Test" (Pairwise Preferences):
Traditionally, to teach the chef, you'd give them two dishes for the same order (e.g., two different lasagnas) and ask a human, "Which one is better?" The robot learns by comparing the two.
- The Problem: As the robot gets better, it becomes incredibly hard for humans to tell the difference between two good lasagnas. It's like asking a food critic to choose between two 5-star meals; it's exhausting, expensive, and slow.
The "Gourmet Judge" (Reward Models + Reinforcement Learning):
Another method involves hiring a "Gourmet Judge" (a separate AI) to taste the food and give it a score from 1 to 10. The robot then tries to cook dishes that get the highest score from this Judge.
- The Problem: The Judge can be wrong, or the robot might try to "game" the Judge by making weird, toxic, or nonsensical dishes just to get a high score. It's also very computationally expensive and unstable, like trying to balance a house of cards in a windstorm.

The New Method: RPO (Reward Partitioning Optimization)

The authors of this paper introduce a new, simpler way called RPO. Instead of comparing two dishes or hiring a separate Judge, RPO looks at all the dishes the robot has ever made for a single order and figures out the "average quality" of that specific order.

Here is how RPO works, using a simple analogy:

The "Classroom Grade" Analogy

Imagine a teacher (the robot) is grading essays.

The Old Way (DRO): The teacher tries to guess the "perfect score" for every single essay they write before they even finish writing it. They have to guess this "perfect score" while simultaneously writing the essay. This is confusing and leads to mistakes.
The RPO Way: The teacher looks at all the essays written for a specific prompt (e.g., "Write a story about a cat").
- Some essays are terrible (score 2).
- Some are okay (score 5).
- Some are amazing (score 9).
- RPO calculates the average quality of all those cat stories.
- If the robot writes a story that is better than the average for that prompt, it gets a "thumbs up" and learns to do it again.
- If it writes a story worse than the average, it gets a "thumbs down" and learns to change.

The Magic Trick: RPO doesn't need a separate "Judge" AI to tell it what the average is. It just looks at the data it already has (the prompt, the response, and the score) and does a quick math calculation to find that average. This makes the training process much faster, more stable, and less likely to go crazy.

What Did They Find?

The researchers tested this new method on many different types of AI models (both ones that read and write, and ones that just write). Here is what happened:

Better Results: The RPO-trained robots wrote responses that were more helpful, more diverse (less repetitive), and safer (less toxic) than robots trained with the old methods.
Stability: While the old methods sometimes made the robot's behavior swing wildly (like a car losing control), RPO kept the robot on a smooth, steady path.
Speed: It took less time to train the robots using RPO.
Robustness: Even if the scores given to the essays were a little bit "noisy" or imperfect (like a human grader having a bad day), RPO still worked well. It didn't break.

The Catch (Limitations)

The paper is honest about where RPO might struggle:

You need a crowd: To calculate that "average," you need to see multiple different responses for the same prompt. If you only have one response for a prompt, RPO can't do its magic math.
Garbage In, Garbage Out: If the scores (rewards) are completely wrong or corrupted, the method will still struggle, though it handles small mistakes better than the old ways.

The Bottom Line

The paper proposes RPO as a smarter, simpler way to teach AI to be helpful. Instead of asking humans to compare two options or building a complex system to guess "perfect scores," RPO simply looks at the whole group of answers for a question, finds the average, and teaches the AI to aim higher than that average. It's a more stable, efficient, and effective way to align AI with human preferences.

Technical Summary: Value-Free Policy Optimization via Reward Partitioning

Problem Statement

Aligning Large Language Models (LLMs) with human preferences typically relies on Reinforcement Learning from Human Feedback (RLHF) or Direct Preference Optimization (DPO). However, these paradigms face significant scalability and stability challenges:

RLHF requires costly human annotation of pairwise preferences, training a separate Reward Model (RM), and unstable reinforcement learning loops (e.g., PPO) that are sample-inefficient and sensitive to hyperparameters.
Direct Preference Optimization (DPO-style) methods eliminate the explicit RM and RL loop but require curated pairwise preference data. As models improve, distinguishing between high-quality responses becomes increasingly difficult, driving up annotation costs.
Single-Trajectory Learning offers a practical alternative using $(prompt, response, scalar\_reward)$ tuples, which are easier to annotate (e.g., thumbs up/down). However, existing methods for this setting, such as Direct Reward Optimization (DRO), rely on jointly learning a policy and a value function. This joint optimization introduces variance, instability, and sensitivity to off-policy data, as errors in value estimation can bias the policy training.

The core problem addressed is the need for a stable, scalable, and simple method to optimize policies directly from scalar reward data without the complexity and instability of learning an auxiliary value function.

Methodology: Reward Partitioning Optimization (RPO)

The authors propose Reward Partitioning Optimization (RPO), a value-free objective that leverages the analytical form of the optimal policy under KL regularization to eliminate the need for a learned value function.

Theoretical Foundation

Starting from the optimal policy formulation in KL-regularized reward maximization:
$\pi^*(y|x) = \pi_{ref}(y|x) \exp\left(\frac{1}{\tau} (r(x, y) - V^*(x))\right)$
where $V^*(x)$ is the soft value function acting as a partition function to normalize rewards. Traditional methods like DRO approximate $V^*(x)$ using a separate neural network, leading to coupled optimization of policy ( $\pi_\theta$ ) and value ( $V_\phi$ ).

RPO bypasses this by empirically estimating the partition function directly from the data. For a given prompt $x$ , the method aggregates all available responses $y_j$ in the dataset to compute an empirical partition function $\hat{Z}(x)$ :
$\hat{Z}(x) = \sum_{j \in I_x} \pi_{ref}(y_j|x) \exp\left(\frac{1}{\tau} r(x, y_j)\right)$
where $I_x$ is the set of indices for samples with prompt $x$ . The empirical value is then $\hat{V}(x) = \tau \log \hat{Z}(x)$ .

Optimization Objective

RPO formulates the learning task as a direct regression on the log-ratio between the learned policy and the reference policy. The loss function minimizes the squared error between the predicted log-ratio and the centered, scaled reward:
$\mathcal{L}_{RPO}(\theta) = \frac{1}{2} \mathbb{E}_{(x,y,r) \sim D} \left[ \left( \log \frac{\pi_\theta(y|x)}{\pi_{ref}(y|x)} - \frac{1}{\tau}(r(x, y) - \hat{V}(x)) \right)^2 \right]$

Key characteristics of the method include:

Value-Free: No auxiliary value network is trained; $\hat{V}(x)$ is computed on-the-fly from the reference policy and observed rewards.
Stable Gradients: The gradient update depends only on the policy parameters, avoiding backpropagation through a learned value estimator.
Offline Compatibility: The method is fully offline and scales to large datasets without requiring online RL loops.

Key Contributions

Proposal of RPO: A new single-trajectory policy optimization method that replaces value estimation with empirical reward normalization.
Analysis of DRO: A theoretical and practical critique of Direct Reward Optimization, highlighting its instability due to joint policy-value optimization.
Formal Derivation: A proof showing that RPO approximates the optimal policy structure by leveraging the partition function formulation.
Extensive Evaluation: Comprehensive experiments across encoder-decoder (Flan-T5) and decoder-only (Mistral, LLaMA, Qwen) architectures, demonstrating consistent superiority over SFT, KTO, and DRO.

Experimental Results

The authors evaluated RPO on the UltraFeedback dataset and across multiple benchmarks (AlpacaEval, MT-Bench, IFEval, GSM8K).

In-Domain Performance (UltraFeedback)

Metrics: RPO consistently outperformed SFT, KTO, and DRO on automatic metrics (BERTScore, ROUGE-L, Distinct-2) and safety (Toxicity).
LLM-as-a-Judge: RPO achieved higher win rates against baselines when evaluated by GPT-4o and Claude-3.5.
Stability: KL divergence analysis showed RPO maintains smooth, controlled trajectories, whereas DRO exhibited strong divergence growth (especially in decoder-only models) and KTO showed oscillations.
Efficiency: RPO demonstrated faster convergence and lower training time compared to DRO and KTO across all model scales.

Out-of-Distribution (OOD) Generalization

RPO showed robust but task-dependent improvements. Gains were most pronounced in instruction-following benchmarks (AlpacaEval, IFEval).
Performance on reasoning tasks (GSM8K) and multi-turn conversation (MT-Bench) showed consistent but more moderate improvements, varying by model family.

Ablation Studies

Temperature ( $\tau$ ): The method is robust to moderate smoothing ( $\tau=5.0$ ) but suffers significant performance degradation with very sharp scaling ( $\tau=0.1$ ), leading to unstable optimization.
Empirical Normalization ( $\hat{V}(x)$ ): Removing the prompt-level reward normalization caused consistent degradation in semantic similarity, diversity, and safety, confirming its critical role in stabilizing reward scales.
Dataset Density: Performance degrades as the number of completions per prompt decreases, with a 25% density setting causing substantial drops in quality. This highlights the method's reliance on diverse reward distributions per prompt for accurate partition estimation.
Noise Robustness: RPO remains stable under moderate reward noise ( $\sigma=0.1, 0.5$ ) but performance declines as noise intensity increases, though it does not collapse even under strong corruption ( $\sigma=1.0$ ).

Significance and Limitations

The paper claims that RPO offers a simple, scalable, and stable alternative to existing alignment methods. By eliminating the need for value function learning, it reduces optimization complexity and variance while providing absolute reward supervision. The authors position RPO as a promising direction for scalable preference optimization, particularly in scenarios where pairwise data is scarce but scalar feedback is available.

Limitations acknowledged by the authors:

Data Requirements: RPO relies on grouped data where multiple completions exist per prompt to estimate the partition function effectively. Performance degrades if prompt-level diversity is insufficient.
Reward Quality: While robust to moderate noise, the method remains dependent on the quality of reward annotations; severe corruption negatively impacts alignment.
Evaluation Scope: The current evaluation is primarily focused on instruction-following and alignment benchmarks. The authors note that broader domains, multilingual settings, and long-horizon reasoning tasks require further study to fully assess generalization.