---
title: Ranked Return Regression for RL (R4)
url: https://www.emergentmind.com/topics/ranked-return-regression-for-rl-r4
type: topic
---

# Ranked Return Regression for RL (R4)

Ranked Return Regression for RL (R4) is a method for reward learning in reinforcement learning (RL) that addresses the challenge of inferring reward functions from human ratings rather than explicit reward specification or binary preferences. R4 centers on a novel ranking mean squared error (rMSE) loss, treating trajectory ratings as ordinal targets, and utilizes differentiable sorting to match predicted trajectory rankings to teacher-provided ratings. Formal guarantees on minimality and completeness distinguish R4 within the class of rating-based reward learning algorithms, and empirical results demonstrate competitive or superior policy performance and sample efficiency in robotic locomotion benchmarks [2601.09236].

## 1. Problem Setting and Notation

R4 operates in a Markov Decision Process (MDP) without a built-in reward, supplemented by teacher ratings. The core elements are:

- State space $\mathcal{S}$, action space $\mathcal{A}$, transition kernel $T(s'|s, a)$, initial distribution $p_0(s)$, and discount $\gamma \in [0, 1)$.
- Full trajectories $T = (s_0, a_0, \dots, s_H, a_H)$ are labeled by a teacher with an ordinal rating $c(T) \in \{0, 1, \dots, n-1\}$, with $0$ denoting "worst" and $n-1$ denoting "best."
- The dataset $D = \bigcup_{k=0}^{n-1} D_k$, where $D_k = \{T : c(T) = k\}$.
- Reward function $f_\theta: \mathcal{S} \times \mathcal{A} \to \mathbb{R}$ with predicted trajectory return $G_\theta(T) = \sum_{t=0}^{H} \gamma^t f_\theta(s_t, a_t)$.
- Objectives:
  1. Learn $\theta$ such that $G_\theta$ respects ordinal classes $c(T)$.
  2. Use $f_\theta$ as the reward function for policy optimization, maximizing $\mathbb{E}[\sum \gamma^t f_\theta(s_t, a_t)]$ in place of a hand-crafted reward.

## 2. Ranking Mean Squared Error (rMSE) Loss

R4 employs the rMSE loss to regress predicted soft ranks to ordinal ratings:

- In each gradient step, sample one trajectory $T_i$ from each class $D_{c_i}$ for $i=0,\dots,n-1$ so $c_i$ is strictly increasing.
- Compute predicted returns $x_i = G_\theta(T_i)$.
- Compute soft ranks $R_i = R^\varepsilon_i(x) \in [0, n-1]$ as a differentiable proxy for the integer rank of $x_i$.
- The rMSE loss is defined as:
  $$
  L_{\mathrm{rMSE}}(\theta) = \frac{1}{n} \sum_{i=0}^{n-1}(R_i - c_i)^2
  $$
  This penalizes deviations between the continuous soft ranks and discrete ratings, encouraging correct trajectory ordering without imposing interval constraints within each class.

## 3. Differentiable Sorting: Soft Ranks

R4 utilizes a differentiable sorting operator for ranking:

- Let $x \in \mathbb{R}^n$ represent predicted returns $[x_0, \ldots, x_{n-1}]$.
- Define the cost matrix $C(x)$ by $C_{ij} = |x_i - x_j|$.
- Soft permutation matrix $P^\varepsilon \in \mathbb{R}^{n \times n}$ is obtained by solving the entropic optimal transport problem:
  $$
  P^\varepsilon(x) = \arg\min_{P \in U_n} \langle P, C(x) \rangle - \varepsilon H(P)
  $$
  where $U_n = \{P \geq 0 : P1 = 1, P^T 1 = 1\}$ (Birkhoff polytope), and entropy $H(P) = -\sum_{ij} P_{ij} \log P_{ij}$.
- After Sinkhorn iterations, compute the soft-rank vector:
  $$
  R^\varepsilon(x) = P^\varepsilon(x) \cdot [0, 1, ..., n-1]^T
  $$
- As $\varepsilon \to 0$, the operator recovers hard sorting; for $\varepsilon > 0$, gradients propagate through $x_i = G_\theta(T_i)$ to $\theta$.

## 4. Training Algorithm

The R4 algorithm proceeds in two phases—reward learning and policy optimization—often cycling between the two:

- **Initialization**:
  - Reward network $f_\theta$.
  - Policy $\pi$ (e.g., Soft Actor-Critic).
  - Replay buffer $\mathcal{B}$.
- **Training Loop**:
  1. Collect environment transitions via $\pi$, add to $\mathcal{B}$.
  2. Every $K_{\text{env}}$ steps (or offline), run $B$ reward-learning updates:
     - For $k = 0,\dots,n-1$, sample $T^{(k)}$ from $D_k$ (or from $\mathcal{B}$).
     - $c_k = k$, $x_k = G_\theta(T^{(k)})$.
     - Compute $R = R^\varepsilon([x_k])$.
     - Compute $L_{\mathrm{rMSE}} = \frac{1}{n} \sum_k (R_k - c_k)^2$.
     - Update $\theta \leftarrow \theta - \eta_\theta \nabla_\theta(L_{\mathrm{rMSE}} + \lambda \|\theta\|^2)$.
     - Relabel all rewards in $\mathcal{B}$ using $f_\theta$.
  3. Run standard RL policy update (e.g., SAC) on $\mathcal{B}$ with $f_\theta$ rewards.

## 5. Theoretical Guarantees: Minimality and Completeness

R4 provides formal guarantees on the resulting learned reward functions under mild assumptions:

- **Deterministic Realizability**: A ground-truth $r^* \in \mathcal{H}$ exists, such that $G^*(T) = \sum \gamma^t r^*(s_t, a_t)$.
- **Binning**: Teacher ratings result from thresholding $G^*(T)$ into $n$ ordinal bins.
- **Model Realizability**: $r^*$ is in the class $\mathcal{H}$ of $f_\theta$.
- **Exact Differentiable Sorting**: The soft-rank operator recovers exact ranks when returns are strictly ordered.
- **Feasible Set**:
  $$
  R = \{ r \in \mathcal{H} : c(T_i) < c(T_j) \implies G_r(T_i) < G_r(T_j) \ \forall T_i, T_j \in D \}
  $$
- **Consistency**: $r^*$ minimizes $L_{\mathrm{rMSE}}$.
- **Completeness and Minimality**: $\arg\min_\theta L_{\mathrm{rMSE}}(\theta) = R$—i.e., the solution set contains all and only reward functions inducing the teacher’s ordering.
- Relaxation to bounded soft-rank error ($|\hat{R}_i - \text{true rank}| \leq \varepsilon$) preserves guarantees if $\varepsilon$ is small relative to $n$.

## 6. Empirical Evaluation

R4 has been empirically evaluated in simulated human feedback scenarios, comparing both offline and online feedback regimes against rating-based RL (RbRL) and preference-based methods (PEBBLE, SURF, QPA). The central performance metric is the downstream policy's undiscounted environment return.

- **Offline Feedback**:
  - **Domains**: OpenAI Gym Reacher, Inverted Double Pendulum, HalfCheetah.
  - **Dataset**: 50–100 trajectories per class, labeled via ground-truth return thresholds.
  - **Results**: R4-reward policies learn statistically faster ($p < 0.05$) and reach higher final returns than SAC trained on ground-truth reward and RbRL in Reacher.
- **Online Feedback**:
  - **Domains**: DeepMind Control Suite (Walker-walk/stand, Cheetah-run, Quadruped-walk/run, Humanoid-stand).
  - **Budget**: Fixed rated trajectories (100–200); preference methods with equivalent feedback count.
  - **Baselines**: RbRL, PEBBLE, SURF, QPA with optimal sampling.
  - **Results**: R4 matches or outperforms baselines, achieving faster learning or higher final return in 4/6 tasks (Bonferroni-corrected $p < 0.0125$).

Across both regimes, R4 models yield higher Trajectory Alignment Coefficient, maintain robustness to noisy labels (up to 80% noise), and are insensitive to the choice of $n$.

## 7. Practical Considerations

R4 presents practical advantages and implementation choices:

- **Computational Cost**: Each rMSE update involves $B$ network forward passes, a Sinkhorn solve (typically 5–20 iterations, size $n \times n$, worst-case $O(n^3)$), and $B$ backpropagations. Standard $n \leq 10$, $B = 4$–8, so computational overhead is modest.
- **Hyperparameters**:
  - Soft-sort regularization $\varepsilon \approx 0.01$–$1.0$, robust over wide ranges.
  - Policy learning batch size per SAC conventions (256).
  - Reward-network architecture: MLP with 2–4 hidden layers, ReLU activations; no dedicated output scaling.
  - Optional $\ell_2$ weight decay $(\lambda \approx 10^{-3}$–$10^{-2})$.
- **Online Feedback Enhancements**:
  - Dynamic feedback schedule: high labeling frequency early, decaying over time.
  - Stratified trajectory sampling: combine top-return and low-return trajectories from buffer of last 50, randomly sample sub-segments of length 8.
  - Ensemble of 3–5 reward networks for prediction stability.
- **Limitations & Open Problems**:
  - Accuracy depends on soft-sort approximation; large noise or tied returns require careful $\varepsilon$ tuning.
  - Theoretical results depend on deterministic returns and model realizability; generalization to stochastic rewards and mis-specified $\mathcal{H}$ is open.
  - Human-in-the-loop usability: preliminary pilots show robustness to varied labeling, but large-scale user studies are pending.

R4 is characterized as a simple, hyperparameter-lean method for learning reward models from multi-class trajectory ratings by regressing soft ranks to ordinal ratings under a differentiable sorting operator, inheriting formal guarantees and demonstrating empirically strong sample efficiency and policy performance in both offline and online RL applications [2601.09236].

Source: https://www.emergentmind.com/topics/ranked-return-regression-for-rl-r4