Reward Learning through Ranking Mean Squared Error

TL;DR

This work tackles reward design bottlenecks in reinforcement learning by learning rewards from ordinal human ratings rather than handcrafting them. It introduces Ranked Return Regression for RL (R4), a rating-based method that optimizes a novel ranking mean squared error (rMSE) loss using differentiable soft ranks to align predicted trajectory returns with teacher ratings. The authors prove theoretical guarantees—the rMSE solution is complete and minimal (with relaxed versions under bounded ranking error)—and demonstrate empirically that R4 outperforms rating-based and preference-based baselines in offline and online robotic locomotion tasks, often with substantially less feedback. The approach offers a principled, scalable path to robust reward learning with practical impact for real-world RL deployment.

Abstract

Reward design remains a significant bottleneck in applying reinforcement learning (RL) to real-world problems. A popular alternative is reward learning, where reward functions are inferred from human feedback rather than manually specified. Recent work has proposed learning reward functions from human feedback in the form of ratings, rather than traditional binary preferences, enabling richer and potentially less cognitively demanding supervision. Building on this paradigm, we introduce a new rating-based RL method, Ranked Return Regression for RL (R4). At its core, R4 employs a novel ranking mean squared error (rMSE) loss, which treats teacher-provided ratings as ordinal targets. Our approach learns from a dataset of trajectory-rating pairs, where each trajectory is labeled with a discrete rating (e.g., "bad," "neutral," "good"). At each training step, we sample a set of trajectories, predict their returns, and rank them using a differentiable sorting operator (soft ranks). We then optimize a mean squared error loss between the resulting soft ranks and the teacher's ratings. Unlike prior rating-based approaches, R4 offers formal guarantees: its solution set is provably minimal and complete under mild assumptions. Empirically, using simulated human feedback, we demonstrate that R4 consistently matches or outperforms existing rating and preference-based RL methods on robotic locomotion benchmarks from OpenAI Gym and the DeepMind Control Suite, while requiring significantly less feedback.

Figures (15)

• Figure 1: Illustration of the R4 learning process: Given a dataset of trajectory–rating pairs, we compute the predicted return for each trajectory under $\hat{r}_\theta$ and apply a differentiable sorting algorithm to obtain soft ranks. Then, we minimize the MSE between the soft ranks and the original ratings.
• Figure 2: Performance of a SAC agent trained with (1) R4, (2) RbRL, and (3) the environment reward. Budget denotes the number of labeled trajectories used to learn the offline reward function.
• Figure 3: This shows online SAC performance evaluated with the true environment reward, using either the environment reward or learned rewards from different rating/preference-based algorithms.
• Figure 4: We evaluate R4 under ablations of our implementation choices: stratified sampling and the dynamic feedback schedule. We find that both features can improve the base R4 method.
• Figure 5: Combined performance of policies trained on reward models learned from five human raters. Light curves show the average performance across five SAC training runs for each rater; the dark curve shows the overall mean across raters. R4 consistently outperforms RbRL despite substantial variation in individual rating behavior. Furthermore, on average, R4 with human ratings performs similarly to R4 with perfect simulated ratings.

Theorems & Definitions (6)

• Definition 1
• Proposition 1: Consistency
• Theorem 1: Completeness and Minimality
• Definition 2
• Theorem 2: Completeness and Minimality under Bounded Ranking Error - Relaxed Theorem \ref{['thm:minimality']}
• proof : Proof of Theorem \ref{['thm:relaxedthm1']}