Source-linked AI summary
Principled Reinforcement Learning with Human Feedback from Pairwise or $K$-wise Comparisons
Banghua Zhu, Jiantao Jiao, Michael I. Jordan
TL;DR
The paper studies how reliably RLHF can learn linear rewards from pairwise or K-wise human comparisons and whether accurate reward estimation yields good policies. It analyzes MLE and pessimistic MLE under BTL and PL models, showing that pessimism protects policy performance under coverage assumptions while true MLE is more efficient for K-wise comparisons.
Problem
The paper examines the potential sub-optimality of MLE for downstream policy training in RLHF, beyond its reward-estimation error.
Method
The paper analyzes MLE and pessimistic MLE for linear reward models learned from pairwise or K-wise comparisons under BTL and PL models.
Results
MLE converges for reward estimation, but MLE can fail for induced-policy performance while pessimistic MLE achieves a near-optimal rate under coverage assumptions; true MLE is asymptotically more efficient for K-wise comparisons.
Takeaways & Limitations
The results support introducing pessimism in RLHF reward learning or policy training and provide theoretical guarantees connecting RLHF with max-entropy IRL.
Takeaways & Limitations
The theory assumes greedy policies with respect to learned rewards rather than the full practical pipeline of pre-training, reward learning, and PPO or policy-gradient fine-tuning.
Abstract
from arXiv · showhide
We provide a theoretical framework for Reinforcement Learning with Human Feedback (RLHF). Our analysis shows that when the true reward function is linear, the widely used maximum likelihood estimator (MLE) converges under both the Bradley-Terry-Luce (BTL) model and the Plackett-Luce (PL) model. However, we show that when training a policy based on the learned reward model, MLE fails while a pessimistic MLE provides policies with improved performance under certain coverage assumptions. Additionally, we demonstrate that under the PL model, the true MLE and an alternative MLE that splits the $K$-wise comparison into pairwise comparisons both converge. Moreover, the true MLE is asymptotically more efficient. Our results validate the empirical success of existing RLHF algorithms in InstructGPT and provide new insights for algorithm design. Furthermore, our results unify the problem of RLHF and max-entropy Inverse Reinforcement Learning (IRL), and provide the first sample complexity bound for max-entropy IRL.
1 Introduction
This paper develops a theoretical RLHF framework for linear rewards learned from pairwise or K-wise human comparisons. It shows that reward estimation can converge while policy optimization still requires pessimism under limited coverage, and that true MLE is more efficient for K-wise comparisons.
- 1 Introduction: RLHF trains a reward model from human rankings of responses sampled for prompts, then fine-tunes the language model with PPO.The analysis models prompts as contextual-bandit states and responses as actions, with a fixed linear reward representation.
- 1 Introduction: The framework studies sample complexity for learning linear rewards from pairwise or K-wise comparison data under the BTL and PL models.When K = 2, the PL setting reduces to pairwise BTL comparisons used in existing RLHF algorithms.
- 1.1 Main Results: MLE converges in a query-dependent semi-norm under regularity conditions, but its induced greedy policy can be suboptimal.The paper distinguishes reward-parameter estimation from downstream policy performance and studies the performance gap against the optimal policy.
- 1.1 Main Results: Pessimistic MLE achieves a near-optimal induced-policy rate under coverage assumptions, whereas MLE can fail under the same assumptions.Pessimism discounts actions that are less represented in the observed dataset.
- 1.1 Main Results: For K-wise comparisons, both the true MLE and the pairwise-splitting estimator converge and achieve near-optimal policies with pessimism, while true MLE has smaller asymptotic variance.The pairwise-splitting estimator corresponds to the algorithm analyzed for InstructGPT.
- 1.1 Main Results: The results extend to MDPs and IRL, provide guarantees for max-entropy IRL, and motivate pessimism or regularization in RLHF policy training.The paper connects RLHF with reward learning from pre-collected data and offline reinforcement learning.
2 Preliminaries
The paper formulates RLHF as learning linear rewards from pairwise or K-wise human rankings, beginning with contextual bandits and extending to MDPs. It specifies the MDP objects, offline sampling procedure, and Plackett-Luce model, whose K=2 case recovers BTL comparisons.
- Scope: The paper studies reward-model sample complexity from pairwise or K-wise comparisons and later extends the analysis to MDPs, IRL, and experiments.The stated framework includes BTL and PL models, with MLE used for reward training.
- Reward and environment model: A finite-horizon MDP consists of states, actions, horizon H, transition kernels, reward distributions, and an initial state distribution.With H = 1 and no transitions, the model reduces to a contextual bandit.
- Reward and environment model: Policies map states to actions, while value functions measure expected cumulative rewards along trajectories generated by the transition and reward processes.The paper uses V⋆ and Q⋆ for optimal value and action-value functions and defines state and state-action occupancy measures.
- Sampling and comparison model: The Plackett-Luce model governs K-wise rankings; when K = 2, it reduces to the pairwise Bradley-Terry-Luce comparison model.Pairwise observations can be represented as Bernoulli outcomes whose parameter depends on the two actions’ true rewards.
3 Learning from Pairwise Comparison
The pairwise-comparison analysis first bounds MLE estimation error, then constructs a pessimistic MLE using confidence sets and lower confidence bounds for policy learning. Under coverage assumptions, pessimistic MLE achieves near-optimal policy performance, whereas MLE can fail.
- 3.1 Algorithms: MLE and Pessimistic MLE: MLE minimizes the negative log likelihood of observed pairwise preferences, using reward differences between the compared actions.The estimator may be nonunique, in which case any minimizer is selected.
- 3.1 Algorithms: MLE and Pessimistic MLE: Pessimistic MLE forms a confidence set around MLE and uses a lower confidence bound as the reward estimate for policy optimization.The construction includes a regularization parameter, a reference vector v, and a state distribution q.
- 3.1 Algorithms: MLE and Pessimistic MLE: The estimation analysis provides a high-probability confidence bound for the parameter error under the pairwise BTL model.The bound depends on the data covariance and regularization, with curvature controlled by γ.
- 3.1 Algorithms: MLE and Pessimistic MLE: A bounded concentratability coefficient represents coverage of the target policy’s feature vector by the observed dataset.The guarantee can also use a reference policy in place of the optimal policy.
- 3.1 Algorithms: MLE and Pessimistic MLE: Choosing the reference vector v appropriately can reduce the concentratability coefficient, especially when feature differences are covered although the features themselves are not.The paper suggests using a frequently occurring feature vector as a rule of thumb.
- 3.1 Algorithms: MLE and Pessimistic MLE: Pessimism can be approximated for neural-network rewards through conservative offline-RL methods or policy-training regularization.The cited policy regularization keeps the policy close to the original policy and within pretraining-data coverage.
- 3.2 Failure of MLE and Lower Bounds: MLE can fail on induced policy performance, while pessimistic MLE achieves a near-optimal rate by discounting poorly represented actions.The result separates MLE from pessimistic MLE when the concentratability coefficient is bounded.
- 3.2 Failure of MLE and Lower Bounds: Pessimistic MLE is minimax-optimal up to constant factors for policy suboptimality under bounded concentratability.The paper provides both an upper-bound comparison and a matching lower-bound construction.
4 Learning from K-wise comparisons
The paper analyzes direct and pairwise-split MLEs for K-wise Plackett-Luce comparisons. Both estimators converge, while the direct MLE is asymptotically more efficient despite the split estimator’s favorable nonasymptotic dependence on K.
- Estimators: The direct estimator MLEK maximizes the joint Plackett-Luce likelihood, whereas MLE2 splits each K-wise ranking into K(K−1)/2 pairwise comparisons and applies pairwise MLE.The split construction replaces the joint ranking probability with a product of pairwise marginals.
- Setting: K-wise comparisons assume a complete ranking of K items, with K restricted to O(1) because humans have limited capacity for comparing many options.The paper contrasts complete rankings with prior work that observes only the maximum or provides expectation bounds.
- MLEK: Theorem 4.1 provides a high-probability estimation guarantee for MLEK under the K-wise Plackett-Luce model.The guarantee holds for any λ > 0 with probability at least 1 − δ.
- Comparison: Relative to pairwise comparisons, the K-wise covariance aggregates feature differences across all action pairs, introducing additional K-dependence in the bound.The paper considers this dependence likely loose but near-optimal when K = O(1), given pairwise minimax lower bounds.
- MLE2: Theorem 4.2 provides a high-probability guarantee for MLE2 and validates its empirical performance in Ouyang et al. (2022).The analysis reports a seemingly better nonasymptotic dependence on K for MLE2, while noting that the MLEK analysis may be loose.
- Comparison: Although both estimators are unbiased and converge, MLEK is asymptotically more efficient because MLE2 is an M-estimator with larger asymptotic variance.The paper also empirically verifies the performances of both estimators.
5 Extension to MDPs
The paper extends its RLHF analysis from contextual bandits to MDPs with action-based and trajectory-based comparisons. The resulting guarantees adapt the comparison model to Q-values or cumulative trajectory rewards, while policy learning still depends on coverage and pessimism.
- MDP comparison regimes: The MDP extension distinguishes action-based comparisons, based on expected cumulative returns from state-action pairs, from trajectory-based comparisons, based on cumulative rewards of entire trajectories.Both regimes begin from sampled initial states or states and generate comparison observations under the MDP dynamics.
- Trajectory-based comparison: For trajectory comparisons, the likelihood exponent changes from a single reward to the difference between cumulative rewards of two trajectories.The corresponding covariance uses differences between cumulative trajectory features and reduces to the contextual-bandit form when H = 1.
- Trajectory-based comparison: Lemma 5.1 gives a high-probability MLE estimation guarantee for trajectory-based comparisons under bounded features.The assumption is ∥ϕ(·, ·)∥∞ ≤ L, and the guarantee holds for any λ > 0 with probability at least 1 − δ.
- Trajectory-based comparison: Theorem 5.2 gives a high-probability guarantee for the induced policy in the trajectory-comparison setting.The result uses the policy-output procedure with the MLE reward estimate and extends to K-wise comparisons by the same argument.
- Action-based comparison: For action-based RLHF, MLE estimates a linear optimal Q-function from Bernoulli comparisons between two actions sampled at each state.The resulting estimator is analyzed through the same approach as the contextual-bandit MLE.
- Action-based comparison: Action-based policy guarantees require good coverage of the covariance directions; without it, constructing lower confidence bounds for every policy can be difficult.When the covariance is invertible and covers all directions well, confidence bounds for Q⋆ can support a good greedy policy without pessimism.
6 Connection with Inverse Reinforcement Learning
The paper connects RLHF with inverse reinforcement learning by viewing human behavior or comparisons through Plackett-Luce models. It derives estimation and policy guarantees for max-entropy IRL and characterizes a related limitation for action-based IRL.
- Connection to IRL: RLHF is framed as reward learning from pre-collected data, linking it to inverse reinforcement learning and offline reinforcement learning.The same techniques provide theoretical guarantees for max-entropy IRL and action-based IRL.
- Max-entropy IRL: In max-entropy IRL, human-selected trajectories follow a Plackett-Luce model, with trajectory sets containing all possible trajectories from an initial state.When the normalization denominator is exactly computable, maximum-entropy IRL reduces to MLE.
- Max-entropy IRL: Although enumerating all trajectories is exponentially difficult in horizon H, expected state frequencies provide an efficient way to compute the MLE gradient.This is the computational approach attributed to Ziebart et al. (2008).
- Max-entropy IRL: Lemma 6.1 gives a high-probability estimation guarantee for max-entropy IRL under the Plackett-Luce model.The guarantee holds for any λ > 0 with probability at least 1 − δ.
- Max-entropy IRL: Combining IRL with pessimism yields a policy guarantee for max-entropy IRL.Theorem 6.2 analyzes the policy produced using the MLE reward estimate in the paper’s policy-learning procedure.
- Limitations: The IRL bound depends on sups |T(s)|, which can be much larger than the feature dimension d.Splitting one observation into pairwise comparisons may improve this dependence in the current analysis.
- Action-based IRL: Action-based IRL models human behavior through an optimal Q-function and observes one human action per state rather than a pairwise comparison.The paper states that MLE still provides a similar performance guarantee in this single-action observation setting.
- Limitations: A provable lower-confidence-bound algorithm for policy learning remains an open problem in action-based IRL.The paper identifies this as analogous to the unresolved issue in action-based RLHF.
7 Experiments
The experiments support the theoretical distinction between parameter estimation and policy quality: MLE converges under the semi-norm but can yield poor policies, whereas pessimistic MLE achieves vanishing sub-optimality. For K-wise comparisons, both estimators converge, with the true MLE showing lower estimation error, especially as K increases.
- Contextual bandit: Figure 1 compares semi-norm convergence on the left with MLE and pessimistic MLE under sub-optimality on the right.Both methods use the same estimated parameter; they differ in how the induced policy is trained.
- Contextual bandit: MLE converges under the semi-norm but fails to give good policy performance, while pessimistic MLE achieves vanishing sub-optimality.The contextual-bandit experiment uses sample sizes from 10 to 500, with 100 repetitions per sample size.
- K-wise comparisons: As sample size increases, both MLE2 and MLEK converge on K-wise comparisons, while MLEK has smaller estimation error.The comparison uses K = 4 and K = 9, with samples ranging from 10 to 500.
- K-wise comparisons: The estimation-error gap between MLEK and MLE2 grows with K, matching the prediction that MLEK is asymptotically more efficient.Figure 2 reports the comparison for K = 4 on the left and K = 9 on the right.
8 Conclusion
The paper’s conclusion identifies two central findings: pessimism is important for guaranteeing good policies, and both estimators converge for K-wise comparisons while the true MLE is asymptotically more efficient. It also outlines scope boundaries involving policy training, behavioral models, and changing features.
- The analysis concludes that pessimism is important for guaranteeing a good policy.
- For K-wise comparisons, both MLEK and MLE2 converge, while MLEK is asymptotically more efficient.
- Scope boundaries: The analysis assumes a greedy policy trained from the learned reward rather than the full pre-training, reward-learning, and PPO or policy-gradient procedure.
- Scope boundaries: The theory focuses on BTL and PL models, leaving extensions to Thurstone and cardinal models for future work.
- Scope boundaries: The constructed confidence bound assumes a fixed feature function, whereas practical fine-tuning may involve slowly changing features.
A Analysis for nonlinear rθ
This analysis extends the reward-learning treatment beyond linear rewards to nonlinear, potentially non-convex reward functions. Under boundedness and smoothness assumptions it gives a finite-sample estimation guarantee, but a nonvanishing dependence remains when the nonlinear smoothness parameter is positive.
- Assumptions: The nonlinear analysis assumes boundedness and smoothness of rθ over parameters, states, and distinct action pairs.
- Assumptions: The linear reward model satisfies the nonlinear assumptions with α0 = LB, α1 = L, and α2 = 0.
- Guarantee: Theorem A.2 provides a high-probability estimation-error guarantee for nonlinear rewards under the stated assumptions.The theorem is stated for any λ > 0 with probability at least 1 −δ.
- Limitation: When α2 > 0, the bound retains dependence on α2 as n grows, so the analysis does not establish a vanishing rate for general nonlinear rewards.
- Extensions: The same argument can be extended to K-wise comparisons and MDPs, and pessimistic MLE can be designed from the resulting confidence bound.
B.3 Proof of Theorem 3.9
The appendix proves the pairwise-comparison policy result and develops the estimation-error argument for K-wise comparisons. The proof combines a constructed counterexample for MLE with coverage-based control for pessimistic MLE and concentration and convexity arguments for estimation.
- Pairwise MLE counterexample: A four-action construction shows that pairwise MLE can select a1 or a2 instead of the optimal action a4 under specified comparison events.
- Pairwise MLE counterexample: The constructed event yields a positive expected suboptimality for the policy induced by MLE.
- Pessimistic MLE: The example’s coverage calculation and Theorem 3.2 imply that pessimistic MLE achieves vanishing error.
- K-wise estimation: For K-wise comparisons, the proof bounds MLE estimation error using strong convexity of the negative log likelihood with respect to the data-dependent semi-norm.
- K-wise estimation: Under the K-wise PL model, the proof establishes zero-mean sub-Gaussian gradient terms and applies concentration inequalities to control the estimation error.
B.6 Proof of Theorem 4.2
The proof bounds the K-wise MLE’s estimation error by combining likelihood optimality, strong convexity, and concentration of the score under the PL model. It then transfers the resulting bound to the theorem’s sub-optimality guarantee.
- Score concentration: Under the K-wise PL model, Lemma B.3 provides the high-probability bound needed for the regularized estimation analysis.The matrix comparison used in the proof relates M to X(ΣD + λI)−1Xᵀ/n².
- Estimation error: The proof begins by bounding the K-wise MLE estimation error in the regularized covariance seminorm.Likelihood optimality yields an inequality involving the error vector and the score at the true parameter.
- Score concentration: The score term is controlled through a dual-norm inequality and a quadratic-form concentration bound.The proof represents the score norm as VᵀMV and applies Bernstein’s inequality to sub-Gaussian variables.
- Conclusion: The resulting inequality is solved for the estimation error, after which the sub-optimality argument follows the same reasoning as Theorem 3.2.The proof explicitly reuses the earlier sub-optimality analysis once the estimation bound is obtained.
B.8 Proof of Theorem 5.2
The proof uses the confidence set from Lemma 5.1 to compare the pessimistic policy’s estimated and true values. Combining this containment with the policy optimality relation yields the theorem’s performance bound.
- Confidence-set argument: Lemma 5.1 places the true parameter inside the confidence set around the MLE with high probability.This containment is the starting point for evaluating the pessimistic policy.
- Policy comparison: Because ˆπPE optimizes the estimated expected value, the proof controls the second value difference using its optimality.The argument compares the pessimistic policy under estimated and true objectives.
- Policy comparison: The confidence-set definition implies a nonpositive estimated-to-auxiliary value difference for ˆπPE.The proof states that ˆJ(ˆπPE) − J′(ˆπPE) ≤ 0 when the true parameter lies in the confidence set.
- Conclusion: Combining these inequalities gives the final performance guarantee with probability at least 1 −2δ.The last step uses the confidence-set control and the policy optimality relation together.
B.9 Proof of Theorem 6.2
The proof of Theorem 6.2 extends the MLE estimation analysis to structured trajectory comparisons. It establishes strong convexity and concentration under the sampling assumptions, then reuses the earlier theorem’s sub-optimality argument.
- Conclusion: The remaining sub-optimality proof follows the same argument as Theorem 5.2.Theorem 6.2 is presented as a direct corollary after Lemma 6.1 is established.
- Strong convexity: The proof first establishes strong convexity of the negative log-likelihood in the data-induced seminorm.The curvature parameter is set to γ = exp(−4LB)/2 under the stated boundedness conditions.
- Estimation error: Solving the resulting inequality yields the MLE estimation bound under the structured comparison model.The proof obtains the bound after combining likelihood optimality, strong convexity, and score control.
- Concentration: Independent sampling makes the concatenated comparison vectors independent, mean-zero, and 1-sub-Gaussian.These properties support concentration of the score quadratic form.
- Concentration: The score norm is written as VᵀMV and bounded using matrix domination plus Bernstein’s inequality.The matrix M is compared with X(ΣD + λI)−1Xᵀ/n² before applying quadratic-form concentration.