PAWS: Preference Learning with Advantage-Weighted Segments

Technical Summary: PAWS (Preference Learning with Advantage-Weighted Segments)

Problem Statement
Preference-based Reinforcement Learning (PbRL) aims to learn policies from human trajectory-level comparisons rather than explicit reward functions. However, existing methods suffer from a fundamental training-inference distribution shift. Typically, PbRL methods train a utility function (e.g., a reward or advantage model) on segment-level or trajectory-level preferences. Yet, during policy optimization, these models are queried on individual state-action pairs to compute per-step learning signals.

This mismatch creates a temporal credit assignment problem. As illustrated in the paper, many distinct per-step advantage assignments can sum to the same segment-level preference label. Consequently, the utility model is underconstrained at the individual step level. When the policy optimizer queries this model on single steps (inference) despite it being trained on segments, the resulting learning signals are ambiguous and unreliable, severely degrading policy learning, particularly in low-data regimes.

Methodology: PAWS
The authors propose PAWS (Preference Learning with Advantage-Weighted Segments) to resolve this distribution shift by aligning utility training with policy optimization at the segment level.

1. Advantage Learning:

   • Instead of learning a reward function, PAWS learns an advantage function $A_\phi$ directly from segment-level preference pairs $(\tau^+, \tau^-)$.
   • The preference likelihood is modeled using a Bradley-Terry formulation based on the sum of advantages within a segment: $P_{A_\phi}[\tau^+ \succ \tau^-] = \frac{\exp(A_\phi(\tau^+))}{\exp(A_\phi(\tau^+)) + \exp(A_\phi(\tau^-))}$.
   • The model can be implemented using either a Transformer encoder or a Multi-Layer Perceptron (MLP).

2. Segment-Based Policy Optimization:

   • Rather than inferring per-step utilities, PAWS performs policy updates directly on trajectory segments.
   • The method formulates a constrained optimization problem to find an optimal segment distribution $p^*(\tau)$ that maximizes the expected advantage while staying within a trust region (KL-divergence bound $\epsilon$) of the original data distribution $p_D(\tau)$.
   • The solution yields a reweighted segment distribution: $p^*(\tau) \propto p_D(\tau) \exp(\frac{1}{\lambda} A_\phi(\tau))$.
   • The policy $\pi_\theta$ is then updated by maximizing a weighted maximum likelihood objective, where the weights are the exponential of the segment advantages:
     $$\mathcal{L}(\theta) = \sum_{\tau \in D} \sum_{(s_t, a_t) \in \tau} \exp\left(\frac{A_\phi(\tau)}{\lambda}\right) \log \pi_\theta(a_t|s_t)$$
   • Crucially, the importance weight is applied to the entire segment, ensuring that the policy update respects the segment-level preference information.

3. Hyperparameter Control via Effective Sample Size:

   • To avoid manual tuning of the KL-bound $\epsilon$ (which depends on action dimensionality and dataset size), PAWS introduces a data-driven strategy to control the Lagrange multiplier $\lambda$.
   • The authors propose setting $\lambda$ to achieve a desired effective sample size ($n_{eff}$), calculated from the importance weights. This allows intuitive control over how many segments meaningfully contribute to the update, balancing exploitation of high-advantage segments with stability.

Key Contributions

• Analysis of Distribution Shift: The paper identifies the mismatch between segment-level utility training and step-level inference as a core limitation of existing PbRL methods, causing ambiguous temporal credit assignment.
• PAWS Algorithm: A novel segment-based preference learning method that performs policy updates using advantage-weighted segments, preserving trajectory-level preference information and avoiding unreliable per-step signals.
• Data-Driven Hyperparameter Tuning: An intuitive strategy for setting policy optimization hyperparameters based on the effective sample size of preference-weighted data, reducing reliance on manual tuning.
• Empirical Validation: Comprehensive evaluation on simulated robotic manipulation (Meta-World) and locomotion (D4RL) tasks, demonstrating consistent improvements over established baselines.

Experimental Results
The authors evaluated PAWS on 10 Meta-World manipulation tasks and 4 D4RL locomotion tasks under two preference budgets (50 and 500 pairs).

• Performance: PAWS (both Transformer and MLP variants) consistently outperformed baselines including Behavior Cloning (BC), P-IQL, CPL, CPL+KL, Preference Transformer, and IPL.
• Low-Data Regime: In the 50-preference setting, many baselines degraded to perform worse than Behavior Cloning, whereas PAWS maintained robust performance, highlighting its efficiency in data-scarce scenarios.
• Ablation Studies:
  • Update Granularity: Replacing PAWS's segment-level updates with state-action-level updates (using the same learned advantage function) resulted in a significant performance drop, confirming that the distribution shift is the primary cause of failure in standard approaches.
  • Segment Length: Increasing segment length exacerbated performance degradation for state-action-based updates but had minimal negative impact on PAWS, which operates natively at the segment level.
  • Human Feedback: When tested on preferences collected from 10 non-author human labelers, PAWS achieved the highest success rates on "Button Press" and "Door Open" tasks, validating its robustness to real human noise.
  • Credit Assignment: PAWS demonstrated a higher Spearman's rank correlation with expert rankings compared to step-based updates, indicating superior preservation of temporal preference structures.

Significance and Claims
The paper claims that the performance gains of PAWS stem primarily from aligning the training and inference distributions of the learned utility function, rather than from specific architectural choices (as both MLP and Transformer variants outperformed their step-level counterparts).

The authors emphasize that PAWS offers a principled solution to the temporal credit assignment problem in PbRL by avoiding the "unreliable per-step learning signals" inherent in methods that train on segments but infer on steps. The method is presented as particularly effective in low-data regimes where step-level utility estimates are noisiest. While the evaluation is restricted to simulated robotic tasks with oracle and small-scale human data, the authors suggest the segment-level perspective is directly applicable to preference fine-tuning of large generative models (e.g., RLHF), where comparisons are typically made over entire responses rather than individual tokens.

Limitations
The authors acknowledge that PAWS assigns uniform importance weights to all state-action pairs within a segment, which may be suboptimal if a segment contains a mix of high- and low-quality actions. Additionally, the choice of effective sample size ($n_{eff}$) requires selection, though it is noted to be more intuitive than tuning the KL bound directly. Finally, the current evaluation is limited to simulation and does not yet address large-scale, noisy human feedback.