Wasserstein Distributionally Robust Regret Optimization for Reinforcement Learning from Human Feedback

Abstract: Reinforcement learning from human feedback (RLHF) has become a core post-training step for aligning LLMs, yet the reward signal used in RLHF is only a learned proxy for true human utility. From an operations research perspective, this creates a decision problem under objective misspecification: the policy is optimized against an estimated reward, while deployment performance is determined by an unobserved objective. The resulting gap leads to reward over-optimization, or Goodharting, where proxy reward continues to improve even after true quality deteriorates. Existing mitigations address this problem through uncertainty penalties, pessimistic rewards, or conservative constraints, but they can be computationally burdensome and overly pessimistic. We propose Wasserstein distributionally robust regret optimization (DRRO) for RLHF. Instead of pessimizing worst-case value as in standard DRO, DRRO pessimizes worst-case regret relative to the best policy under the same plausible reward perturbation. We study the promptwise problem through a simplex allocation model and show that, under an $\ell_1$ ambiguity set, the inner worst-case regret admits an exact solution and the optimal policy has a water-filling structure. These results lead to a practical policy-gradient algorithm with a simple sampled-bonus interpretation and only minor changes to PPO/GRPO-style RLHF training. The framework also clarifies theoretically why DRRO is less pessimistic than DRO, and our experiments show that DRRO mitigates over-optimization more effectively than existing baselines while standard DRO is systematically over-pessimistic.

• The paper introduces DRRO—a method that optimizes robust regret in RLHF to counteract reward over-optimization and objective misspecification.
• It employs a closed-form solution on a simplex allocation model and dynamic ambiguity calibration to distribute robust bonuses for practical policy updates.
• Empirical comparisons show that DRRO outperforms traditional DRO by balancing robustness and policy improvement over the widest useful KL range.

Wasserstein Distributionally Robust Regret Optimization in RLHF

Motivation and Problem Formulation

Reinforcement Learning from Human Feedback (RLHF) is a critical step in aligning LLMs, yet its reward signal is a learned proxy rather than true human utility. This creates an objective misspecification problem: policies are optimized against a proxy reward, but their deployment performance is dictated by an unobserved, underlying objective. Empirical work has demonstrated reward over-optimization, or Goodharting, where the proxy reward improves, but gold-standard metrics of quality deteriorate as policy diverges from its initial state (see (Figure 1)).

Figure 1: Reward over-optimization in RLHF: proxy reward rises as policy moves away from initialization, but true evaluator begins to fall.

The paper formalizes RLHF as a robust decision problem with endogenous objective uncertainty, highlighting the tension between methods that hedge against reward misspecification (uncertainty penalties, pessimistic rewards, conservative constraints) and the risk that robustness mechanisms themselves become over-pessimistic, hindering policy improvement.

DRRO: Formulation and Theoretical Foundations

Standard Distributionally Robust Optimization (DRO) protects against worst-case value under reward perturbation, but this is an excessively conservative deployment benchmark. The paper introduces Distributionally Robust Regret Optimization (DRRO), which pessimizes the worst-case regret relative to the best policy under the same reward perturbation. Formulating RLHF as an ex-ante DRRO problem, the paper proves that Wasserstein ambiguity over reward functions reduces to deterministic ambiguity in their mean reward, specifically using an $\ell_1$ norm ambiguity set.

For a fixed prompt, the problem space is reduced to a simplex allocation model where the inner worst-case regret admits a closed-form solution. The adversary optimally concentrates its ambiguity budget on a single response and the minimizer exhibits a water-filling structure, assigning probability mass only as needed to equalize uncovered rewards above a threshold.

Algorithmic Translation and Practical Policy Optimization

The promptwise closed-form motivates a practical policy-gradient algorithm. The robust correction acts as a reward-shaping bonus that is introduced within PPO or GRPO-style RLHF updates. A hard robust bonus can be realized by assigning the full ambiguity budget to the most threatening sampled response; a softmax relaxation distributes the bonus across sampled responses proportional to their adversarial threat level. SNIS estimators provide a tractable means of estimating the robust bonus from partial rollouts, and the ambiguity budget is made dynamic via a sequence-level KL statistic (calibrated using Donsker–Varadhan variational bounds), which models the growth in reward-model error as policy drifts.

Empirical Comparison of DRRO vs DRO and Baselines

Extensive experiments compare five RLHF update rules—PPO, GRPO, DRO, DRRO-hard, and DRRO-soft/dynamic—on a unified codebase with standardized prompt distribution, reward models, and evaluation metrics. The main benchmark demonstrates that DRRO achieves both the peak and most durable improvement in held-out gold reward, resisting over-optimization and maintaining strong performance across the widest useful KL range (see (Figure 2)).

Figure 2: Main benchmark: Proxy reward rises monotonically, but gold reward deteriorates for several baselines; DRRO maintains strongest gold trajectory over largest KL range.

Direct comparison between DRO and DRRO highlights that DRO is systematically over-pessimistic, flattening the policy to avoid worst-case value, while DRRO allows more concentrated optimization on responses that retain attractiveness under plausible reward perturbations (see (Figure 3)).

Figure 3: DRO vs DRRO: DRRO achieves superior reward-vs-conservatism tradeoff, DRO is excessively pessimistic throughout training.

Internal ablations show that the dynamic ambiguity budget and the soft assignment of robust bonuses are vital for DRRO's practical performance; hard fixed budgets underperform throughout most of the KL range (see (Figure 4)).

Figure 4: DRRO ablation: Dynamic budget and soft assignment yield best gold reward, fixed-budget/hard-assignment variant weakest across useful KL.

Theoretical Comparisons and Coverage-Based Perspective

Analytically, DRRO dominates DRO under rank-preserving proxy rewards due to its concentration-preserving allocation. DRO equalizes exposure, penalizing all uncertainty indiscriminately, while DRRO only hedges relative regret against plausible competitors. Further, coverage-based certificates show DRO pays for absolute coverage of the chosen policy, while DRRO pays for relative coverage versus plausible alternative policies, formalizing why DRRO is less pessimistic and more efficient in its use of the robustness budget.

Implications and Future Directions

The DRRO framework clarifies the nature of robustness in RLHF, providing stronger mitigation of reward over-optimization with minimal computational overhead and direct integration into mature PPO/GRPO pipelines. By robustifying regret rather than value, DRRO allows optimization to proceed unless improvement is fundamentally unattainable under plausible reward misspecification.

Practically, DRRO improves the reward-vs-KL frontier throughout training, supporting greater exploration and optimization in policy space. Theoretically, its tractable structure and coverage-based guarantees motivate further study in robust policy optimization under uncertain rewards. Future work may generalize to sequential settings, investigate alternative ambiguity geometries, and extend dynamic budget calibration to broader uncertainty metrics.

Conclusion

Wasserstein DRRO constitutes a principled, tractable, and empirically effective method for mitigating reward over-optimization in RLHF. It achieves superior balance between robustness and optimization by pessimizing relative regret, demonstrating clear advantages in empirical benchmarks and theoretical comparisons with DRO. The adoption of DRRO in large-scale RLHF pipelines is expected to yield safer, more reliable policy optimization in the presence of reward model misspecification (2605.00155).