# STAIR: Addressing Stage Misalignment through Temporal-Aligned Preference Reinforcement Learning

arXiv: https://arxiv.org/abs/2509.23802 · HTML: https://mlsanity.com/paper/2509.23802/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Reward modeling and RLHF, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 35 of its 48 references that are in the map.
- First posted on arXiv 28 Sep 2025.

## Placed in

- LLM alignment / RLHF / preference optimisation / [Reward modeling and RLHF](https://mlsanity.com/alignment/reward-modeling-and-rlhf/)

## Builds on (most cited first, 30 of 35)

- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [On Variational Bounds of Mutual Information](https://mlsanity.com/paper/1905.06922/) (2019)
- [Time-Contrastive Networks: Self-Supervised Learning from Video](https://mlsanity.com/paper/1704.06888/) (2017)
- [Reward learning from human preferences and demonstrations in Atari](https://mlsanity.com/paper/1811.06521/) (2018)
- [PEBBLE: Feedback-Efficient Interactive Reinforcement Learning via Relabeling Experience and Unsupervised Pre-training](https://mlsanity.com/paper/2106.05091/) (2021)
- [Automatic Goal Generation for Reinforcement Learning Agents](https://mlsanity.com/paper/1705.06366/) (2017)
- [Contrastive Learning as Goal-Conditioned Reinforcement Learning](https://mlsanity.com/paper/2206.07568/) (2022)
- [VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning](https://mlsanity.com/paper/1910.08348/) (2019)
- [Mastering the Game of Stratego with Model-Free Multiagent Reinforcement Learning](https://mlsanity.com/paper/2206.15378/) (2022)
- [Optimal Goal-Reaching Reinforcement Learning via Quasimetric Learning](https://mlsanity.com/paper/2304.01203/) (2023)
- [Towards Human-Level Bimanual Dexterous Manipulation with Reinforcement Learning](https://mlsanity.com/paper/2206.08686/) (2022)
- [B-Pref: Benchmarking Preference-Based Reinforcement Learning](https://mlsanity.com/paper/2111.03026/) (2021)
- [SURF: Semi-supervised Reward Learning with Data Augmentation for Feedback-efficient Preference-based Reinforcement Learning](https://mlsanity.com/paper/2203.10050/) (2022)
- [Few-Shot Preference Learning for Human-in-the-Loop RL](https://mlsanity.com/paper/2212.03363/) (2022)
- [METRA: Scalable Unsupervised RL with Metric-Aware Abstraction](https://mlsanity.com/paper/2310.08887/) (2023)
- [The Laplacian in RL: Learning Representations with Efficient Approximations](https://mlsanity.com/paper/1810.04586/) (2018)
- [Preference Transformer: Modeling Human Preferences using Transformers for RL](https://mlsanity.com/paper/2303.00957/) (2023)
- [Active Preference-Based Gaussian Process Regression for Reward Learning](https://mlsanity.com/paper/2005.02575/) (2020)
- [Reward Uncertainty for Exploration in Preference-based Reinforcement Learning](https://mlsanity.com/paper/2205.12401/) (2022)
- [Benchmarks and Algorithms for Offline Preference-Based Reward Learning](https://mlsanity.com/paper/2301.01392/) (2023)
- [Learning Temporal Distances: Contrastive Successor Features Can Provide a Metric Structure for Decision-Making](https://mlsanity.com/paper/2406.17098/) (2024)
- [RIME: Robust Preference-based Reinforcement Learning with Noisy Preferences](https://mlsanity.com/paper/2402.17257/) (2024)
- [Episodic Novelty Through Temporal Distance](https://mlsanity.com/paper/2501.15418/) (2025)
- [Query-Policy Misalignment in Preference-Based Reinforcement Learning](https://mlsanity.com/paper/2305.17400/) (2023)
- [Efficient Preference-based Reinforcement Learning via Aligned Experience Estimation](https://mlsanity.com/paper/2405.18688/) (2024)

## Built on by (most cited first, 3 of 3)

- [COLLIE: Guiding Skill Discovery in Semantically Coherent Latent Space](https://mlsanity.com/paper/2606.00950/) (2026)
- [From Reward-Free Representations to Preferences: Rethinking Offline Preference-Based Reinforcement Learning](https://mlsanity.com/paper/2606.01123/) (2026)
- [Multi-Agent Reinforcement Learning via Agent-Specific Preference](https://mlsanity.com/paper/2608.08604/) (2026)

## Abstract (arXiv)

Preference-based reinforcement learning (PbRL) bypasses complex reward engineering by learning rewards directly from human preferences, enabling better alignment with human intentions. However, its effectiveness in multi-stage tasks, where agents sequentially perform sub-tasks (e.g., navigation, grasping), is limited by stage misalignment: Comparing segments from mismatched stages, such as movement versus manipulation, results in uninformative feedback, thus hindering policy learning. In this paper, we validate the stage misalignment issue through theoretical analysis and empirical experiments. To address this issue, we propose STage-AlIgned Reward learning (STAIR), which first learns a stage approximation based on temporal distance, then prioritizes comparisons within the same stage. Temporal distance is learned via contrastive learning, which groups temporally close states into coherent stages, without predefined task knowledge, and adapts dynamically to policy changes. Extensive experiments demonstrate STAIR's superiority in multi-stage tasks and competitive performance in single-stage tasks. Furthermore, human studies show that stages approximated by STAIR are consistent with human cognition, confirming its effectiveness in mitigating stage misalignment.