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

Ni Mu (Tsinghua University) · Yiqin Yang (Institute of automation, Chinese academy of science, Chinese Academy of Sciences) · Bo Xu (Wuhan University) · Yao Luan (Tsinghua University) · Qing-Shan Jia (Tsinghua University, Tsinghua University)

coherent stagescontrastive learning dynamic adaptation empirical experiments human cognition human preferences multi-stage tasks policy learning preference-based reinforcement learning reward engineeringstage approximationstage misalignmentstairtemporal distance theoretical analysis

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 **ST**age-**A**l**I**gned **R**eward 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.