cs.LGMar 1, 2025

Distributionally Robust Reinforcement Learning with Human Feedback

Authors: Debmalya MandalPaulius SasnauskasGoran Radanovic

Abstract

Reinforcement learning from human feedback (RLHF) has evolved to be one of the main methods for fine-tuning large language models (LLMs). However, existing RLHF methods are non-robust, and their performance deteriorates if the downstream task differs significantly from the preference dataset used in fine-tuning. In order to mitigate this problem, we introduce a distributionally robust RLHF for fine-tuning LLMs. In particular, our goal is to ensure that a fine-tuned model retains its performance even when the distribution of prompts significantly differs from the distribution encountered during fine-tuning. We formulate distributionally robust optimization (DRO) version of two popular fine-tuning methods -- (1) reward-based RLHF and (2) reward-free DPO (direct preference optimization). We propose a minibatch gradient descent based algorithms for both of them, and theoretically prove convergence guarantees for the algorithms. Subsequently, we evaluate our algorithms on an out-of-distribution (OOD) task by first training the model on the Unified-Feedback dataset and evaluating its performance on two different datasets. The experimental results show that our robust training improves the accuracy of the learned reward models on average, and markedly on some tasks, such as reasoning. Furthermore, we show that the robust versions of policy optimization methods, similarly improve performance on OOD tasks.

Explore similar work

Apr 30, 2026cs.LG

Wasserstein Distributionally Robust Regret Optimization for Reinforcement Learning from Human Feedback

Reinforcement learning from human feedback (RLHF) is a central post-training tool for aligning large language models, but its training reward is only a learned proxy for true human utility. This creates a decision problem under objective misspecification: the policy is optimized against an estimated reward, while deployment performance is governed by an unobserved population preference. The resulting gap leads to reward over-optimization, where proxy reward keeps improving after true quality deteriorates. We propose distributionally robust regret optimization (DRRO) for RLHF with a Wasserstein ambiguity set over reward laws, using promptwise p\ell_p distances between reward vectors as transport costs. Unlike standard distributionally robust optimization, which pessimizes worst-case value, DRRO pessimizes worst-case regret relative to the best policy under the same plausible reward perturbation. We show that the expressive-policy problem decomposes into promptwise regret problems. For each prompt, the inner adversary has a dual-norm closed form; under the 1\ell_1 transport cost used by our algorithm, the optimizer has a water-filling structure. These results lead to a practical policy-gradient algorithm that adds a simple sampled bonus to GRPO-style training. Theory and experiments both show that DRRO is less over-pessimistic than standard DRO and mitigates over-optimization more effectively than existing baselines.
Yikai Wang, Shang Liu, Jose Blanchet
Jul 28, 2026cs.LG

Meta-Learned Reward Shaping for Reinforcement Learning from Human Feedback

Reinforcement Learning from Human Feedback (RLHF) is the standard approach for aligning large language models with human preferences, but its quality is limited by static, task-agnostic reward models. This mismatch leads to sparse learning signals and suboptimal alignment. We introduce MeRLa (Meta-Learned Reward Shaping), a principled framework that meta-learns a task-aware shaping function Φ(x,y;φ)Φ(x,y;φ) across auxiliary tasks before RLHF training. The learned shaping produces a composite reward that preserves policy optimality while providing task-specific learning signals. Our meta-objective combines task discrimination, entropy regularization, and potential-based conservation for stable convergence. We provide theoretical guarantees for policy invariance, analyze representation drift sensitivity, and formally address incentive misalignment from entropy maximization. Experiments on LLaMA-3-8B across four benchmarks show consistent improvements over PPO, DPO, GRPO, and DAPO, achieving a 90.8% length-controlled win rate on AlpacaEval 2.0 and a score of 9.14 on MT-Bench, with 41% less training instability. MeRLa retains its benefits when combined with process-based and rubric-based enhanced rewards.
Yunpeng Chu
Jun 2, 2026cs.LG

RL Excursions during Pre-Training: Re-examining Policy Optimization for LLM training

The standard LLM training pipeline applies reinforcement learning (RL) only after pre-training and supervised fine-tuning (SFT). We question this status quo by training a LLM from scratch and applying RL, SFT, and SFT followed by RL directly to intermediate pre-training checkpoints. We find that RL is effective very early, and often matches the full SFT\toRL pipeline early as well. Through experiments on harder problems, we find that targeted pre-training data composition is a strong lever for RL effectiveness, even more so than model scale. Beyond reasoning accuracy, applying RL directly to base checkpoints expands the model's distribution; the sharpening effect reported in recent work arises only when RL follows SFT. The general capabilities of the model remain essentially unchanged by RL, while they degrade following SFT. Finally, we merge RL and SFT objectives by parallel averaging, which outperforms across all other training methods discussed, across metrics, while preserving general capabilities. Together, these results suggest that LLM training might benefit from an expanded use of RL.
Rachit Bansal, Clara Mohri, Tian Qin +2