🤖 machine learning

Best Policy Learning from Trajectory Preference Feedback

This paper introduces Posterior Sampling for Preference Learning (PSPL), a novel algorithm that combines offline preference data with online pure exploration to achieve the first Bayesian simple regret guarantees for best policy identification in Preference-based Reinforcement Learning, demonstrating superior performance over existing baselines on various benchmarks.

Original authors: Akhil Agnihotri, Rahul Jain, Deepak Ramachandran, Zheng Wen

Published 2026-04-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Akhil Agnihotri, Rahul Jain, Deepak Ramachandran, Zheng Wen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to drive a car, but you don't have a manual, and you can't ask the robot "what is the reward for this action?" because the robot doesn't understand numbers. Instead, you have to act like a parent: you watch the robot drive two different routes, and you simply say, "I liked route A better than route B."

This is the core challenge of Reinforcement Learning from Human Feedback (RLHF), a technology used to train AI models (like chatbots or image generators) to behave the way humans want.

Here is a simple breakdown of the paper's story, using everyday analogies.

1. The Problem: The "Bad Teacher" and the "Broken Map"

Currently, AI training often relies on a two-step process that is prone to errors:

The Reward Model: First, humans try to write a "scorecard" (a mathematical formula) that predicts how good a human will like a specific output.
The Optimization: The AI tries to maximize that score.

The Catch: Humans are bad at writing perfect scorecards. Sometimes the AI finds a loophole to get a high score without actually being helpful (this is called "reward hacking"). Also, the data used to build that scorecard often comes from a specific group of people or a specific time, making it biased. If the "teacher" (the human rater) isn't an expert, the AI learns bad habits.

2. The Solution: "Posterior Sampling for Preference Learning" (PSPL)

The authors propose a smarter way to learn called PSPL. Instead of trying to build a perfect scorecard first, they let the AI learn directly from the "A vs. B" comparisons, while keeping a "mental map" of what it knows and what it doesn't.

Think of PSPL as a Detective with a Crystal Ball:

The Offline Clues (The Old Case Files): Before the detective starts investigating, they are given a box of old case files (an offline dataset). These files contain past comparisons (e.g., "In 2023, people preferred this route").
- The Twist: The detective knows these files might be from a "subpar" detective (a rater with low competence). Maybe the old detective was tired or didn't know the city well.
The Online Investigation (The New Patrol): The detective goes out on the street to gather fresh evidence. They try two different routes, ask a fresh set of people which is better, and update their map.
The Crystal Ball (Posterior Sampling): The detective doesn't just guess one single "best route." Instead, they imagine two different versions of reality simultaneously.
- Reality A: "What if the old files were mostly right, and the city is mostly safe?"
- Reality B: "What if the old files were wrong, and the city is actually dangerous?"
- The detective tries to drive a route that works well in both realities. This forces them to explore areas they are unsure about, rather than just sticking to what they think they know.

3. The Secret Sauce: Knowing How "Competent" the Rater Is

The paper introduces a clever concept called Rater Competence.

Imagine you are learning to cook.

Scenario A: You get feedback from a Michelin-star chef. Their opinion is gold.
Scenario B: You get feedback from a toddler who just likes food that is spicy.

The PSPL algorithm is smart enough to ask: "How much should I trust this feedback?"

If the rater is an expert (high competence), the AI leans heavily on their feedback.
If the rater is clueless (low competence), the AI treats their feedback as "noise" and relies more on its own exploration.

The paper proves mathematically that even if the offline data comes from a "mediocre" rater, the AI can still learn the best policy (the best route) by combining that imperfect data with smart online exploration.

4. The "Bootstrapped" Version (Making it Practical)

The perfect version of this algorithm (PSPL) is mathematically beautiful but computationally heavy—like trying to solve a Rubik's cube while juggling chainsaws.

To make it usable in the real world, the authors created Bootstrapped PSPL.

The Analogy: Instead of solving the Rubik's cube perfectly, they take a quick snapshot of the current state, shake the cube slightly (add random noise), and solve that slightly different version. They do this many times.
The Result: This creates a "good enough" approximation that runs fast on computers but still keeps the magic of exploring the unknown.

5. The Results: Does it Work?

The authors tested this on:

Video Games: Like Mountain Car (getting a car up a hill) and RiverSwim (a fish trying to swim upstream against a current).
Image Generation: Teaching an AI to generate better pictures based on human preferences (using the "Pick-a-Pic" dataset).

The Outcome:
In every test, PSPL found the "best policy" (the best way to drive, swim, or generate images) faster and with fewer mistakes than existing methods. It was particularly good at handling situations where the initial data was biased or the human raters weren't perfect experts.

Summary

This paper is about teaching AI to learn from comparisons ("I like A better than B") rather than scores. It solves the problem of "bad data" by using a smart statistical trick (sampling two realities) to figure out how much to trust the data. It's like having a student who can learn from a flawed textbook and a real-world teacher, knowing exactly how much to listen to each, to become the ultimate expert.

1. Problem Statement

The paper addresses the Best Policy Identification (BPI) problem within the framework of Preference-based Reinforcement Learning (PbRL), specifically in settings where the agent must learn from a combination of offline preference datasets and online pure exploration.

Context: Standard Reinforcement Learning from Human Feedback (RLHF) relies on learning a scalar reward model from human annotations, which is prone to misspecification and reward hacking. PbRL offers a robust alternative by using binary trajectory comparisons (e.g., "Trajectory A is better than B").
The Challenge: In practice, PbRL often starts with an offline dataset of trajectory preferences collected from a "rater" (human or automated) who may have subpar competence (noisy, biased, or out-of-distribution). The goal is to leverage this imperfect offline data to "warm-start" an online learning algorithm that performs pure exploration (focusing on finding the optimal policy rather than maximizing cumulative reward during learning) to identify the best policy with minimal simple regret.
Objective: Given an unknown episodic Markov Decision Process (MDP) with unknown transition dynamics and reward functions, and an offline dataset of trajectory pairs with binary preferences, design an algorithm that minimizes the simple Bayesian regret (the gap between the optimal policy's reward and the learned policy's reward) after $K$ online episodes.

2. Methodology: Posterior Sampling for Preference Learning (PSPL)

The authors propose PSPL, a novel algorithm inspired by Top-Two Thompson Sampling. Unlike standard RLHF pipelines that train a separate reward model, PSPL maintains Bayesian posteriors over both the reward model and the transition dynamics.

A. Theoretical Framework

Modeling:
- MDP: Defined as $M = (P_\eta, S, A, H, r_\theta, \rho)$ , where $\eta$ parameterizes transitions and $\theta$ parameterizes the reward.
- Preference Model: Uses a Bradley-Terry model where the probability of preferring trajectory $\tau_1$ $τ_{1}$ over $\tau_0$ $τ_{0}$ depends on a rater's "competence" parameters $(\lambda, \beta)$ $(λ, β)$ .
  - $\beta$ : Deliberateness/surety of the rater (higher $\beta$ means less random noise).
  - $\lambda$ : Knowledgeability of the rater regarding the true reward vector $\theta$ (higher $\lambda$ means the rater's estimate $\vartheta$ is closer to $\theta$ ).
- Trajectory Embedding: Trajectories are mapped to a feature space via $\phi(\tau)$ .
Informed Priors:
- The algorithm constructs informed priors for $\theta$ and $\eta$ by conditioning on the offline dataset $D_0$ .
- The prior over $\theta$ incorporates the likelihood of the observed preferences given the rater's competence model.
- The prior over $\eta$ incorporates the transition counts observed in the offline trajectories.
Algorithm Steps (Algorithm 1):
- Initialization: Construct informed priors $\nu_1(\theta)$ and $\chi_1(\eta)$ from $D_0$ .
- Loop (for $k=1 \dots K$ ):
  1. Sample: Draw two independent samples of parameters $(\hat{\theta}^{(0)}, \hat{\eta}^{(0)})$ and $(\hat{\theta}^{(1)}, \hat{\eta}^{(1)})$ from the current posteriors.
  2. Compute Policies: Derive two policies $\pi^{(0)}$ and $\pi^{(1)}$ based on these sampled parameters (using value iteration or linear programming).
  3. Explore: Execute both policies to generate two trajectories $\tau^{(0)}$ and $\tau^{(1)}$ .
  4. Feedback: Query the rater for a binary preference $Y_k$ between the two trajectories.
  5. Update: Append the new data to the dataset and update the posteriors $\nu_{k+1}(\theta)$ and $\chi_{k+1}(\eta)$ .
- Output: Return the policy derived from the Maximum A Posteriori (MAP) estimate after $K$ episodes.

B. Practical Approximation: Bootstrapped PSPL

Since exact posterior updates are intractable due to the loss of conjugacy in the Bradley-Terry likelihood, the authors introduce Bootstrapped PSPL (Algorithm 2).

Technique: Uses Bayesian Bootstrapping to approximate posterior samples.
Mechanism: Instead of solving the posterior analytically, it perturbs the surrogate loss function (composed of online loss, offline loss, and prior loss) with random weights (Bernoulli for data, Gaussian for priors).
Optimization: Solves the perturbed loss to find a point estimate (MAP), which serves as a sample from the approximate posterior distribution. This allows the algorithm to be computationally feasible for high-dimensional problems like image generation.

3. Key Contributions

First BPI Framework for PbRL with Offline Data: The paper establishes the first formal framework for Best Policy Identification in unknown MDPs using a mix of offline trajectory preferences and online pure exploration.
Theoretical Guarantees:
- Provides the first Bayesian simple regret bounds for PbRL.
- Derives a regret bound that explicitly depends on the size of the offline dataset ( $N$ ) and the rater's competence ( $\lambda, \beta$ ).
- Shows that as $N \to \infty$ and the rater approaches an expert ( $\lambda, \beta \to \infty$ ), the simple regret converges to zero exponentially fast.
Algorithmic Innovation: Introduces PSPL and its bootstrapped approximation, which effectively integrates offline data quality (rater competence) into the online exploration strategy.
Empirical Validation: Demonstrates superior performance over state-of-the-art baselines (DPO, IPO, LPbRL, DPS) in both control tasks and complex generative tasks.

4. Experimental Results

The authors evaluated Bootstrapped PSPL on two types of benchmarks:

Control Environments (MountainCar, RiverSwim):
- Comparison: PSPL was compared against Logistic Preference-based RL (LPbRL) and Dueling Posterior Sampling (DPS).
- Findings:
  - PSPL consistently achieved lower simple regret than baselines.
  - Rater Competence: Performance improved significantly as the rater's competence ( $\lambda, \beta$ ) increased.
  - Offline Data Size: Increasing the size of the offline dataset ( $N$ ) reduced regret, even with a "mediocre" expert.
  - Robustness: The algorithm remained robust even when the assumed rater competence parameters were misspecified.
Image Generation (Pick-a-Pic Dataset):
- Setup: Applied to text-to-image generation (diffusion models) using the Pick-a-Pic dataset (500k+ human preferences).
- Method: Treated the denoising chain as a trajectory. Used an ensemble of automatic reward models (ImageReward, Aesthetic) to evaluate final performance.
- Results: PSPL outperformed DPS in generating images with higher final rewards. The bootstrapped approach successfully handled the high-dimensional state space of diffusion models where exact posterior sampling is impossible.

5. Significance and Impact

Bridging Offline and Online Learning: The paper solves a critical gap in RLHF by providing a principled way to use large, potentially noisy offline preference datasets to bootstrap online learning, rather than treating them as static training data for a reward model.
Focus on Simple Regret: By targeting simple regret (finding the best policy) rather than cumulative regret, the work aligns more closely with real-world RLHF objectives where the final model performance matters more than the intermediate exploration cost.
Handling Imperfect Feedback: The explicit modeling of rater competence ( $\lambda, \beta$ ) allows the algorithm to weigh offline data appropriately, making it robust to biased or low-quality human feedback, a common issue in large-scale alignment.
Scalability: The bootstrapped approximation makes the theoretically sound Bayesian approach applicable to modern generative AI tasks (like diffusion models), moving beyond simple tabular or linear bandit settings.

In summary, this work provides a rigorous theoretical and practical foundation for hybrid offline-online preference learning, demonstrating that incorporating offline data with known (or estimated) rater competence significantly accelerates the identification of optimal policies in complex, unknown environments.