Wasserstein Distributionally Robust Regret Optimization for Reinforcement Learning from Human Feedback

TL;DR

This paper introduces Wasserstein distributionally robust regret optimization (DRRO) for reinforcement learning from human feedback, effectively reducing reward over-optimization and improving alignment with true human utility.

Contribution

It proposes a novel DRRO framework that optimizes worst-case regret, with a practical policy-gradient algorithm and theoretical insights into its advantages over standard DRO.

Findings

DRRO mitigates over-optimization more effectively than existing methods.

The inner worst-case regret admits an exact solution with a water-filling structure.

Experiments demonstrate DRRO's superior performance in alignment tasks.

Abstract

Reinforcement learning from human feedback (RLHF) has become a core post-training step for aligning large language models, 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$-ground-cost Wasserstein 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 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.