STAIR: Addressing Stage Misalignment through Temporal-Aligned Preference Reinforcement Learning
Authors: Yao Luan
Ni Mu
Yiqin Yang
Bo XU
Qing-Shan Jia
Venue: NeurIPS 2025 poster ()
Keywords: reinforcement learning
preference based reinforcement learning
Primary area: reinforcement_learning

Abstract: 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.

STAIR: Addressing Stage Misalignment through Temporal-Aligned Preference Reinforcement Learning
Yao Luan"!, Ni Mu Yiqin Yang??, Bo Xu', Qing-Shan Jia Beijing Key Laboratory of Embodied Intelligence Systems, Department Automation, Tsinghua University, Beijing, China The Key Laboratory of Cognition and Decision Intelligence for Complex Systems, Institute of Automation, Chinese Academy of Sciences, Beijing, China {1uany23, mnn2))mmall tsinghua. edu. cn yangBia. ac. en, jiaqs@tsinghua. edu cn
Abstract
Preference-based reinforcement learning (PbRL) bypasses complex reward en- gineering 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 (c.g., navigation, grasping). is limited by stage misalignment: Comparing segments from mismatched stages, such 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 "'age-AlIgnne Reward learning (STAIR), which first learns 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 policy changes. Extensive experiments demonstrate STAIR's su- periority in multi-stage tasks and competitive performance in single-stage tasks. Furthermore, human studies show that stages approximated by STAIR con- sistent with human cognition, confirming its effectiveness in mitigating stage misalignment. Code is available https: htt::/gittttt github.
Introduction
Reinforcement Learning (RL) has achieved significant progress in various applications, including robotics 5], gaming [39, 34, 26], and autonomous systems [2, 8, 25]. Yet, the effectiveness of RL relies heavily on well-designed reward functions, which often require substantial manual effort and expert knowledge. To address this challenge, preference-bare basedd reinforcement learning (PbRL) has emerged as promising alternative. This approach leverages human preferences among different agent behaviors as the reward signal, thereby alleviating the need for complex reward design.
However, many real-world sequential tasks exhibit multi-stage structures [41, 20], factor often overlooked in existing research. For example, in a robotic ggo-fetch-returyy task, where an agent retrieves an object from distance, the agent must ① navigate to the object, grasp and 3 transport it to target location, as illustrated in Figure 1. Current PbRL methods [22, 32] face challenges of stage misalignment when collecting human preferences for these multi-stage tasks. This issue arises when behaviors from different stages, such as navigation and grasping, are presented
"Yo Luan Ni Mu contributed equally. Correspondence to Yiqin Yang and Qing- Shan Jia.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Phase Learning a Stage Approximation
Phase Stage-Aligned Query Selection
Stage Approxmation
Implicit Stages
Stage-aligned query
Stage
Efficient for Reword Leorning
Stage
Stage-misaligned query (
Dettnggish m's true future states a from other states a sia contractive leamning
insfficient
Stage
Temporal Distance
Segment's color Indicating the approximated stage
Figure 1: An overview of the proposed STAIR. (1) Learn the temporal distance to construct an on-policy stage difference approximator. (2) Use the temporal distance to select stage-aligned queries.
to humans for comparison. It leads ambiguous feedback, labelers struggle compare behaviors in distinct subtasks, like efficient movement versus precise manipulation. This ambiguity can be clarified by event segmentation theory from cognitive science [45, 19J. The theory suggests that humans process action sequences by dividing them into discrete event boundaries. Consequently, comparisons involving behaviors that cross these natural event boundaries (from different stages) impose a higher cognitive load on labelers, thereby increasing ambiguity in their assessments. Moreover, significant differences in state distributions across stages can reduce the information gained from misaligned queries, adversely impacting policy learning.
In this paper, we systematically analyze the impact of stage misalignment both theoretically and empirically. As illustrated in Proposition 2, comparing behaviors across different stages leads to inefficient policy learning, which requires significantly more feedback. Additionally, Proposition and the experiments in Figure 3 reveal that when humans have inherent preferences for certain stages (e.g., favoring stages closer to task completion), conventional methods show quadratic growth in feedback demands, while the stage-aligned approach scales linearly. These findings underscore the importance of stage alignment for efficient preference learning in multi-stage tasks.
To address the stage misalignment issue, we focus on selecting queries with aligned stages for comparison, where key challenge is identifying these stages without prior task knowledge. We propose novel method, STAIR, which leverages temporal distance to measure stage differences. Temporal distance learned through contrastive learning, grouping closely occurring states together, while separating those with significant temporal gaps. As illustrated in Figure 1, STAIR involves two main phases: 1) utilize contrastive learning develop a temporal distance model for measuring stage differences; and 2) prioritize the comparison of segments with small stage differences. Extensive experiments show that STAIR outperforms state-of-the-arr PbRL methods, achieving higher success rates, improved feedback efficiency, and faster convergence. Moreover, human studies validate that queries selected by STAIR align with human cognition, confirming its effectiveness in addressing stage misalignment.
In summary, our contributions are threefold: (1) We identify the critical issue of stage misalignment in PbRL through theoretical analysis and human experiments. (2) We propose STAIR, a novel stage- aligned learning method that automatically approximates stage differences via temporal distance, and selects stage-aligned queries to address the stage misalignment issue. (3) Extensive experiments show that STAIR outperforms existing methods, validating our key insight that stage alignment significantly improves preference learning efficiency in multi-stage scenarios.
2 Preliminaries
Reinforcement learning. An RL problem typically builds on Markov Decision Problem (MDP) characterized by tuple (S,A, P.r. no, 7), where S and are the state space and the action space, P: x A + ((S) is the state transition, r S A IR is the reward function, d4 is the initial state distribution, and r E (0. 1) the discount factor. A policy म S - ((A) interacts with the environment by sampling action from the distribution π (als) when observing state The goal of

---

RL agent is to learn a policy v, which maximizes the expectation of a discounted cumulative reward:
[(T] E po..T 1== 7 ('r(s), a.]].
Preference-based reinforcement learning. In PbRL, the true reward function is not available and replaced by an estimated reward function fe parameterized by v, which is trained to be consistent with human preference. Specifically, given a segment pair (1),01), where a segment H
state-action sequence with fixed length H, the preference is a binary indicating the index of the preferred segment. To construct fy, we follow the the Bradley-Terry model [4, 7]. and define the preference predictor Ps. as follows:
{sk, Ok....,. # ,
exp(E (ot (p(s),a!))
Polor Øol Eie (xp(( Pol (s))),
(1)
where σi do indicates the human prefer d than du- Then, fgg can be trained by minimizing the cross-entropy loss between the true preference y and the preference estimated by Ps:
Breward (15) = 1))o) [11 y) log Poloo > ai] + yog Pa > ool]
(2)
where Do is the preference buffer storing (ao, 01, y) labeled by human.
Stage Formulation. For complex tasks with multi-stage characteristics, such as the "go-fetch- return" task in Figure 1, we reformalize the problem as a chain of stages as W1 - W2 + Wange Each stage w represents an aggregation of states, and all stages comprise the stage space = [w]..... (uig...,WNuup) ), Saage = |s]_ mapping function F(s, : S x Ω - [0, 1] assigns the probability of a state s belonging to stage w. Using F(s,w), w), we can sequentially decomposed a trajectory r = (5), .., - 1) into N, segments {aī N. (N₂ S Nauge). where all states within a segment belong to the same stage. Specifically, let GT (i) denote the index of the stage for state si, then, the segmentation of 7 is obtained by solving the following optimization problem:
max Eejer (or()))) G
S.t. (0 = -1s
This optimization problem maximizes the likelihood of the stage decomposition, quantified by F, while maintaining the sequential chain structure stage transitions. Once G' (.) is solved, segments can be derived as of = (sj,aj)c=(()==)
1, ...., .,. where N, max; GT( is the number of stages in trajectory т. N₂ = Nsage holds only if the task is completed within this trajectory. The optimal objective value of problem (3) quantifies the degree to which the trajectory 7 can be divided into stages. To evaluate the degree to which an MDP can be divided into stages, we propose using the average optimal objective value of problem (3) across trajectories generated by the optimal policy n². This measure can be formally expressed as: F f Ererr [maxor (+) Esjer Er [F(sT war (s)]]-
3 Impact of Stage Misalignment
3.1 Impact of Stage Misalignment in Multi-Stage Tasks
In this subsection, we first validate the existence of the multi-stage property in real-world tasks, then demonstrate the negative effects of stage misalignment. To verify the multi-stage property, we analyze the state distributions at each timestep, as greater differences among them suggest higher potential for problem segmentation into stages. Formally, we introduce Proposition (proof in Appendix A), which considers classifier that predicts the collected timestep for given state, with high accuracy suggesting the task's 's multi-stage property. Note that Proposition holds independently of Nauge-
Proposition 1. For an MDP and trajectories generated by its optimal policy n², consider a calibrated classifier T(s) that takes a state s as input and outputs the probability pr(s) that the state is collected at step t E (1,2....N) N.. ((s) = max, ps(s). Denote the accuracy of the classifier as acc, the multi-stage measure FF has lower bound F 2 acc?.
To empirically verify the multi-stage property in real-world scenarios, we analyze Meta World robotic manipulation tasks [44] an example, since most MetaWorld tasks are multi-stage. For example,
3

---

Prefer 0a
Prefer O2
Stage-Aligned PaRL
Conventional PbRL.
Prefer Both
window-oper
Roias 20
Roua 40
100
1100,1
V100,2
10055
Start Step of Segment
80
Egisade Reward
90
90
75
75
11100
60
60
20
45
45
V2,
V2,2
V2,5
0/2
Start Step of Segment 0
20
80
25
50 Epochs
75
100
25
Epochs
50
75
100
F11
P12
F15
Figure 3: (Left) Human preferences of segments started at different timesteps in the window-open task. Each point (fa, represents that segment do and σi are collected from steps by and by. Humans prefer segments in later timesteps, suggesting a stage reward bias, where humans' underlying reward is higher in these later stages. (Right) Normalized episode reward of various Roias in the abstract MDP model. Additional results are shown in Appendix F.1.
Figure 2: An illustration of the experiment in the abstract MDP model.
the window-open task requires the arm open the window, which can be divided into two stages: grasping the handle and pulling the handle. Specifically, we train a classifier as in Proposition I using trajectories of the expert policy. The classifier achieves an impressive 81% accuracy, indicating significant differences in state distributions across timesteps, thereby supporting the presence of the multi-stage property in these tasks. Experimental details are shown in Appendix D.I.
Then, we theoretically evaluate the impact of stage misalignment on efficiency in multi-stage tasks. We reformalize these tasks to an abstract MDP (St, (T.),-): P,r, (o, 7), where each state w; rep- resents a stage' and has [T,] actions { |T.]
j=1' Transitions follow a deterministic chain structure (w; W1+1), and the reward function F perfectly aligns with human preferences. With this for- mulation, we compare two query selection methods: (1)
selection samples segments within the same stage, and (2) conventional selection samples segments across all stages [7, 22]. As analyzed by Proposition 2, generally, the conventional PbRL requires more queries to learn an optimal policy than the stage-aligned PbRI.. The analysis leverages the similarity between PbRL. and ranking methods, where the true ordering is defined by the reward function F(w, v). The stage-aligned method learns local orderings for actions in each stage, and the conventional method learns global ordering of all stage-action pairs. Detailed assumptions and the proof are provided in Appendix A.
Proposition 2. In the worst case scenario, the conventional PbRL needs (1 log([(2( additional queries to learn the optimal policy compared the stage-aligned PhRL
3.2 Impact of Stage Misalignment Under Stage Reward Bias
While Section 3.1 analyzes the impact of stage misalignment in general multi-stage tasks, this subsection focuses specific case, where humans prefer certain stages others. Our theoretical and empirical analysis shows that stage misalignment has a more pronounced effect in these scenarios.
Assuming human preferences can be modeled by reward function, we formalize these multi-stage tasks through a reward decomposition: r(s, a) = Tstage (s) + +Ts(( a), where (s) represents the average reward of the stage containing state 8, and (sa a) denotes the reward advantage of the state-action pair within that stage. This decomposition reflects that stages differ in importance in human preference judgments. We refer to (suge(s) (s) the stage reward bias, since higher stage reward bias reflects more preferred stages. Note that this decomposition is primarily for validating the severity of stage misalignment, and is not necessary for the proposed method (Section 4).
To validate the existence of stage reward bias in practical tasks, we conduct a human preference experiment. We sample segments of 20-timestep lengths from expert trajectories, and collect human preference labels for segment pairs. Human labelers watch video renderings of each segment and
'To avoid confusion with the general MDP formulation, we use w to denote both the state and the stage for the abstract MDP, and use to denote the action.

---

select the one that is more beneficial for achieving the task objective, as detailed in Appendix D. Results Figure (Left) indicate a clear human preference for segments starting later, which empirically demonstrates the existence of stage reward bias in human preference, where certain stages receive higher valuation than others.
Further, we evaluate the benefit of stage-aligned query selection under stage reward bias. Specifically, we instantiate the abstract MDP model as in Figure 2, where wi, Vi,j denote the i-th stage and the j-th action in stage wi. In this model, [S2] = 101, [X| = 5, E {1, 100. т} (wT is a terminal state) and Fssae (w;) Uniform (0, R (w,, 02) Uniform [0, 10), with Roias controls the scale of stage reward bias. As shown in Figure 3 (Right), stage-aligned selection yields better learning efficiency than conventional methods in this model. This advantage may result from stage-miin maieemisaligged queries providing limited information for action selection in the current stage, especially under significant stage reward bias. Experimental details are provided in Appendix D.I.
We also provide a theoretical analysis on query complexity to support these findings. Proposition analyzes cases with significant stage reward bias, yielding more precise result than Proposition 2. Detailed assumptions and proof's are available in Appendix A.
Proposition 3. If the stage reward bias is sufficiently large, such that the reward ordering between (wi, UU) and (wy, (w) depends solely on Wi and we for i & i', then in the worst case scenario, conventional PbRL requires (([22 log((T[)) additional queries to learn the optimal policy compared stage-aligned PbRL.
In practice, the stage reward bias is often less significant than that assumed in Proposition 3. Con- sequently, the sample complexity for conventional PbRL ranges from [[[]][]] log(( and ([ log([T)),, which remains considerably higher than that of stage-aligned PbRL. In sum- mary, in multi-stage tasks, stage-aligned query selection enhances policy learning by eliminating less informative queries, especially when stage-specific bias is significant.
Stage-Aligned Reward Learning
To address the stage misalignment in PbRL, we focus selecting stage-aligned queries. However, two key challenges arise: First, stage definitions are often based on subjective expertise of specific tasks, limiting their generalizability across different tasks. Second, stage measurement should be adaptable to evolving policies; otherwise, stages assessed with early suboptimal policies may become incompatible with the newer, improved ones. We elaborate these challenges in Appendix C.
To tackle these challenges, we propose STage-Allgned Reward learning (STAIR). A core innovation of STAIR is the use of temporal distance to develop on on-policy stage difference approximator. First, STAIR employs contrastive learning to construct a temporal distance model, which efficiently approximates stage differences between states in an on-policy manner, as detailed in Section 4.1. Then, we design query selection method in Section 4.2 to identify stage-aligned queries. The overall framework of STAIR is illustrated in Figure and Algorithm 1.
4.1 Learning Temporal Distance
In STAIR, utilize temporal distance as measure of stage differences, as it effectively addresses the aforementioned challenges. Specifically, temporal distance quantifies the transition probabilities between states under given policy. Easily reachable state pairs have smaller temporal distances, indicating similar stages, whereas hard-to-reach state pairs show larger distances, suggesting distinct stages. This measure does not rely on any predefined task-specific stage definitions, and could adapt dynamically to policy changes, as temporal distance can be learned in an on-policy manner. Therefore, temporal distance serves as an effective stage approximator.
Let p. (8+ = y|so = x)) denote the discounted state occupancy measure of state y when starting from under policy It represents the discounted probability of reaching state y from x:
ro y =) = (1 (1-7)
ND (st = -1so = =x),
(4)
where pa (sk y|s0 = xr) is the probability of reaching y from x in exactly steps under policy Using the discounted state occupancy measure in (4), temporal distance quantifies the transition
5

---

probability between states. The successor distance [30] serves as start-of-the thattoof-hhearrt implementation of the temporal distance, which is defined as
(2 y) log (p,, y/sa = g)/p, y/su z))
(5)
The successor distance is proven be quasimetrie [30], even in stochastic MDPs, making it reliable measure of the stage difference between states.
To learn the successor distance dīp (z, y). we employ contrastive learning [30]. This method trains an energy function fa y) that assigns higher scores to state pairs (z, y) belonging to the same trajectory, which follows the joint distribution (* y/8o x); and that assigns lower scores to pairs (Ir, y) from different trajectories, which equivalently sampled from the marginal distributions Pa(a)) and p+ p+(y) & (,Po(r)p, optimized using the symmetrized infoNCE loss [31]:
yl8o æ)) respectively. The energy function fo (z, y) is
(x):
+ log exp(folI;, 31))
Lgg =
log Σ j== exp(fo(2x, U1))
jj= exp((f((z),)
[=(0==01=0==) [35]. the successor
As the optimal energy function fa- satisfies fo- (2. = log
C-p+
distance could be derived by dīpp(t (z. y) = lo (y, v) - So [o- y). Following prior work [30], we
parameterize the score function fo(2, v) = ca. (v) do (x, v): where 0 = (0e: 0a). Then after
contrastive learning, we have dip (2, v) dæz (x. v). allowing do; (x, to serve directly as the temporal distance.
4.2 Stage-Aligned Query Selection
In this subsection, we focus on addressing the stage misalignment issue via temporal distance. Specifically, we propose a stage-aligned query selection method, which mitigates stage misalignment by filtering queries identified as misaligned based temporal distance.
Measuring stage difference between segments. The query selection method requires measuring the stage difference between segments du, 01. However, the temporal distance in Section 4.1 was originally designed to measure state distances. Therefore, we propose a quadrilateral distance to adapt the temporal distance in Section 4.1 for segment distances:
daagg 01)
1 4
(d% sb) + dsu sh-1) + dso (5.. s/11) + dS st), (7)
where s and 8i denote the initial and final states of seg- ment J, This quadrilateral distance ensures that stage-aligned segments yield smaller daage values. As visualized in Fig- ure 4. the side terms, (s sb) + (5 $ s))=)). favor segments with closely aligned starting and ending points, in- dicating straightforward stage alignment. The diagonal terms, (sg: s/==) + dgp( (s) H s)), favor segments with shorter temporal spans, thereby focusing on segments concentrated within single stage. We further analyze the behaviour of dange theoretically in Appendix A.3.
(gg
Go
(5-11. f
Length measured
dy
(s ahrz)
(sb. ab)
Figure 4: An illustration of the quadrilateral distance defined in (7). The distance dsage (oo, σi) ) is the av- erage length (measured by tempo- ral distance dīp)) of the four dashed lines in the quadrilateral.
Query selection. Using the stage difference metric, we then propose the stage-aligned query selection method in Algorithm 2. This method calculates score function [(o, Ø1) for each candidate segment pair (oo. 01), and selects queries with the largest scores. To enhance the informativeness of comparisons, the score function, we integrate stage alignment with the reward model's uncertainty, a widely adopted metric of infor- mativeness [22, 38). We represent this uncertainty with (ao, J.), which calculates the variance of Po 01. value across ensemble members: daaae (0o, 01) VariP
)[==)) follow-
ing [22]. Here Ps. [oo > " denotes an ensemble of ne identical reward models with different randomly initalized parameters. Formally, the selection score is defined as:
1 01) (Caage dstago (oo, σ ()) (Cstate + datate dauate 01)),
(8)

---

SAC with task reward (oracle)
STAIR
PEBBLE
RUNE
MRN
QPA
RIME
door-open (Nura=0000 5000)
door-open 1000000000
sweep-into Wisss 100000
sweep-into 20000)
100
100
100
100
75
Success
50-
25
25
25
1.00
0.25
0.50
1.00
Window-oppn (000))2000))
window-close (Mws)=2000)
door-unlock (Nna050000
faucet-open (Nncy-33000)
100
100
100
100
75
75
75-
75
25
0.00
0.25
0.00
0.25
0.50
0.75
0.50
0.75
1.00
0.75
1,00
X10
Timestops
Timesteps
Figure 5: Learning curves on robot manipulation tasks from MetaWorld. The solid line and the shaded area represent the mean and the standard deviation of success rates (%).
SAC with task reward (oracle)
STAIR
PEBBLE
RUNE
MRN
QPA
RIME
(Nww==0000
quadruped (Merx 2000)
quadruped-run (Num== 2000)
push-green [Naery
1000
1000
1000
Success Rate (%)
100
750
750
Eplcadic
500
500
500
50
250
250
0.00
0.50
1.00
0.50
1.00
1.00
0.00
0.50
0.75
1.00
Timesteps
x10
Timesteps x10*
DMControl
RoboDesk
Figure 6: Learning curves on locomotion tasks from DMControl and robot manipulation task from RoboDesk. The solid line and the shaded area represent the mean and the standard deviation of the episodic returns (DMControl) and success rates (RoboDesk).
where dauge and dstate are normalized to [0, 1], and Catge, Cstate are hyperparameters. The first term encourages stage alignment, while the second term emphasizes queries where the reward model shows high uncertainty.
Experiments
We design our experiments answer the following questions: Ql: How does STAIR compare to other state-of-the-arr methods in multi-stage tasks? Q2: Is the stage-aligned query selection still beneficial in single-stage tasks? Q3: Does the stage approximated by temporal distance align with human cognition? Q4: What is the contribution of each of the proposed techniques in STAIR?
5.1 Setup
Domains. We evaluate STAIR on several complex robotic manipulation and locomotion tasks from MetaWorld [44], DMControl [40], RoboDesk [16]. In RoboDesk, we modify the task representation from pixel-based camera image to a robot-arm state representation, which aligns with Meta World. Details these domains are shown Appendix D.2. MetaWorld and RoboDesk focus on multi-stage tasks that achieve specific objectives, such opening window, which requires arm to first grasp handle and then pull it. Conversely, DMControl includes single-stage tasks, which focus on maximizing travel distance velocity, posing challenges in defining stages for humans. We evaluate STAIR both multi-stage and single-stage domains to show its robustness and generalizability.
Baselines and Implementation. including PEBBLE [22], RUNE [23], MRN [24], RIME [6] and QPA [13]. We also evaluate SAC [11] with ground truth reward as a performance upper bound. For PEBBLE and RUNE, we employ disagrcement-badt based query selection, it yields the best performance. Following prior works
compare STAIR with several state-of-the-an PbRL methods,
7

---

STAIR
STAIR (ISR)
Timestep
PEBBLE
STAIR (Kg==)
STAIR (Ks0=4)
STAIR (Ksp=16)
Soccess Rate (16:
100
door-open (Wem 5000)
100
Sweep-into (Msza=10000)
Success Rate (%)
100
door-open (Mrra (5000
100
sweep-into INoue 100000
75
75
75
25
25
25
25
0.00
0.25
0.50
0.75
1,00
0.00
0.25
0.50 0.75 1,00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1,00
Timesteps X10*
Timestaps X10*
Timesteps x10"
Timesieps x10'
Figure Ablation results. (Left) Performance of STAIR with various stage difference approximators.
(Right) Performance of STAIR under range of temporal distance update frequencies KSD-
[21, 22, 23], consider oracle script teacher that provides preference by comparing the total task reward two segments. Please refer to Appendix D.3 for further details.
5.2 Results on Benchmark Tasks
Multi-stage tasks in MetaWorld and RoboDesk. As shown in Figure 5 (MetaWorld) and (RoboDesk), STAIR outperforms baselines across all evaluated multi-stage tasks, achieving success rates close to 100% in most cases. Moreover, STAIR shows a faster convergence speed. For example, in door-open (5000)))000 and window-open 22000),, STAIR's performance rapidly improves over fewer time steps, and reaches a stable high performance earlier than the other baselines.
Single-stage tasks in DMControl. As shown in Figure 6, STAIR outperforms PEBBLE and is competitive with other baselines, even in single-stage stage tasks. This suggests the potential of STAIR for broader applications. This success may arise from the implicit curriculum learning induced by STAIR: Though the task is not multi stage STAIR implicitly divides the reward learning process into stages by introducing queries progressively. Later learning stages are presented only after the agent masters the earlier ones, enabling the model to focus on the complexities of the newly added stages. We provide further discussions on it in Appendix G, and will explore this further in future work.
5.3 Human Experiments
STAIR
PEBBLE
We conducted human labeling experiment to examine whether the stages approximated by temporal distance align with human cognition. Specifically, human labelers are instructed to review queries generated by STAIR and PEBBLE, and assess whether the two segments in query correspond to the same stage, as detailed in Appendix E. Figure 8 shows the ratio of queries identified by labelers stage-aligned. Additionally, Figure 13 visualizes segment pairs selected by STAIR and PEBBLE, where the pair selected by STAIR shows different behaviors in similar stages, and the one PEBBLE exhibits different stages, The results indicate that STAIR effectively selects queries recognized by humans stage-aligned, thereby facilitating human labeling.
Align Ratio (%)
75
50
docc-opes
vindon-clcae vindow open
Figure 8: The ratio of two selected seg- ments in query being at the same stage from a human perspective.
5.4 Ablation Study
Impact of the stage difference approximator. To show the efficiency of the segment stage difference approximator in (7), we design two variants of STAIR: (1) Timestep: This variant uses the collected timestep of a state as the stage approximation, then uses Interval Span Ratio (ISR) We mas min
disR (Zo, 01) = Wi.aos Wi. Wimii and Wi,max are the minimum and maximum approximated stage for states in d,. It is assumed that min S W1.min otherwise, we swap ag and σι. The ISR quantifies the alignment between two intervals by measuring their intersection if they overlap, or the gap between them otherwise, thereby serving as an indicator stage alignment. (2) STAIR (ISR): This variant uses the temporal distance from the current state to the trajectory's 's initial state the stage approximation, which serves as
to assess the stage difference between two segments dp, Ø1, where

---

Table 1: Performance of STAIR with various
Table 2: Performance of STAIR and PEBBLE on door-open using different numbers of query feedback.
Estage and Ctate values.
door-open
sweep-into
Cange
Canse
(Noual = 5000)
(Noul = 10000)
Nottl
STAIR
PEBBLE
0.1
99.91 + 0.07
87.81
1.7
500
52.01
23118
20.00
17.88
1
0.5
99.53 = 0.44
92.31
+ 2,28
2000
77.77
1: 11.67
28.79 1 17.02
2
0.1
100.00 = 0.00
91.08
3.42
5000
100.00
85.57
12.77
2
0.5
98.67 + 1.14
93.58 + 4.22
10000
99.93 = 0.06
92.53 + 6.53
the on-policy version of timestep, and also uses ISR to evaluate the stage difference between two segments.
As shown in Figure 7 (Left), STAIR outperforms all the variants. The limitation of these variants may lie in measuring stage differences between segments in one-dimensional axis. contrast, our quadrilateral distance evaluates stage differences in two-dimensional space, effectively modeling more complex relationships between segments.
Impact of temporal distance update frequency. Accurately estimating stage differences requires the stage approximator to adapt effectively to the evolving policy. Figure 7 (Right) shows that a lower frequency of temporal distance updates (larger KSD) leads to decreased performance of STAIR. This emphasizes the importance of training the stage difference approximator in an on-policy manner.
Robustness on query selection hyperparameters. Table shows that STAIR's performance remains consistent across hyperparameter configurations. This shows the robustness of STAIR to hyperparameter changes, which ensures stable and effective learning in more general settings.
Enhanced feedback efficiency. We compare the performance of STAIR and PEBBLE using different numbers of queries (Niotal on the World door-open task. Table 2 shows that STAIR consistently outperforms PEBBLE, demonstrating its effectiveness in utilizing limited feedback. This result arises from STAIR's stage-aligned query selection, which offers more informative queries for policy learning, enabling the agent to learn effectively with fewer feedback.
6 Related Work
Preference-based reinforcement learning. PbRL has emerged as a promising framework for aligning agent behaviors with human intentions, thereby alleviating the need for complex reward engineering [7. 21, 27, 28]. To enhance feedback efficiency, prior works have explored unsupervised pre-training [22], semi-supervised data augmentation [32, 1], reward uncertainty- -ased exploration [23], on-policy query selection [13], and meta-learning [12]. These methods aim to minimize dependence human feedback by selecting more informative queries based entropy [14), ensemble disagreement [22], or feature-space diversity [3]. However, they often overlook the inherent multi-stage nature of real-world tasks [41, 20], where comparisons of cross-stage segments can lead to inefficient learning. Our approach addresses this by introducing stage-aligned query selection based on temporal distance, which improves learning efficiency by focusing on segment comparisons within the same stage.
Temporal distance learning in RL. Temporal distance quantifies the expected transition steps between states under given policy, which is essential in goal-conditioned RL [9]. skill discovery [33], and state representation learning [37]. Prior works learn it through spectral decomposition of state transitions [43], constrained optimization for temporal consistency [42], and contrastive
learning that groups temporally adjacent states [30]. Recent research [30] introduced successor distance, which is formally guaranteed to be quasimetric. The quasimetric characteristic ensures the successor distance to a reliable state similarity measure, which thus has been employed to design intrinsic rewards that enhance exploration [15]. Nevertheless, these methods remain underexplored in multi-stage sequential decision problems. Our work fills this gap by employing temporal distance as stage similarity measure, enabling stage decomposition without task knowledge.

---

7 Conclusion
This paper presents STAIR, novel approach for addressing stage misalignment in multi-stage environments. STAIR first constructs stage difference approximator via temporal distance, learned through contrastive learning. Then, STAIR extends the temporal distance to assess segment distances via quadrilateral distance, enabling the selection of stage-air aoggeal queries. Experiments show that STAIR outperforms baselines multi-stage tasks and remains competitive in single-stage stage domains. Human experiments further confirm the effectiveness of the learned stage approximation.
Limitations. One limitation of STAIR is that the quadrilateral distance only assesses pairwise segment differences, limiting its applicability to other preference formats. We will explore this in future work.
Acknowledgments and Disclosure of Funding
This work is supported by NSFC (No. 62125304), the National Key Research and Development Program China 1002600)0044000,, the 111 International Collaboration Project (BP2018006) 8006), the Beijing Natural Science Foundation (L233005), and BNRist project (BNR2024TD03000),
References
[1] Fengshuo Bai, Rui Zhao, Hongming Zhang, Sijia Cui, Ying Wen, Yaodong Yang, Bo Xu, and Lei Han. Efficient preference- based reinforcement learning via aligned experience estimation. arXiv preprint arXiv:2405. 18688, 2024.
[2] Marc G Bellemare, Salvatore Candido, Pablo Samuel Castro, Jun Gong, Marlos C Machado, Subhodeep Moitra, Sameera S Ponda, and Ziyu Wang. Autonomous navigation of stratospheric balloons using reinforcement learning. Nature, 588(7836):77-82 2020.
[3] Erdem Biyik, Nicolas Huynh, Mykel Kochenderfer, and Dorsa Sadigh. Active preference-based gaussian process regression for reward learning. In Robotics: Science and Systems, 2020.
[4] Ralph Allan Bradley and Milton E Terry. Rank analysis of incomplete block designs: method of paired comparisons. Biomelrika, 59(3/4):324-445 1952.
[5] Yuanpei Chen, Tianhao Wu, Shengjie Wang. Xidong Feng, Jiechuan Jiang, Zongqing Lu, Stephen McAleer, Hao Dong, Song-Chun Zhu, and Yaodong Yang. Towards human-level bimanual dexterous manipulation with reinforcement learning. Advances in Neural Information Processing Systems, 35:5150-5166 2022.
[6] Jie Cheng, Gang Xiong, Xingyuan Dai, Qinghai Miao, Yisheng Lv, and Fui-Yue Wang. Rime: Robust preference-baseee reinforcement learning with noisy preferences. arXiv preprint arXiv:2402.172577 2024.
[7] Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. Advances in neural information processing systems, 30, 2017.
[8] Jonas Degrave, Federico Felici, Jonas Buchli, Michael Neunert, Brendan Tracey, Francesco Carpanese, Timo Ewalds, Roland Hafner, Abbas Abdolmaleki, Diego de Las Casas, et al. Mag- netic control of tokamak plasmas through deep reinforcement learning. Nature, 602(7897):444 419, 2022.
[9] Benjamin Eysenbach. Tianjun Zhang, Sergey Levine, and Ruslan Salakhutdinov. Contrastive learning as goal-conditioned reinforcement learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances Neural Information Processing Systems, 2022.
[10] Carlos Florensa, David Held, Xinyang Geng, and Pieter Abbeel. Automatic goal generation for reinforcement learning agents. International conference on machine learning, pages 1515-1528. PMLR, 2018.
10

---

[11] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforeement learning with a stochastic actor. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning. volume 80 of Proceedings of Machine Learning Research, pages 1861-1870. PMIRR, 10-15 Jul 2018.
[12] Donald Joseph Hejna and Dorsa Sadigh. Few-shot preference learning for human-in-the-loop rl. In Conference on Robot Learning, pages 2014-2025. PMLR, 2023.
[13] Xiao Hu, Jianxiong Li, Xianyuan Zhan, Qing-Shan Jia, and Ya-Qin Zhang. Query-policy misalignment in preference-based reinforcement learning. arXiv preprint 2rXv: 2023.
[14] Borja Ibarz, Jan Leike, Tobias Pohlen, Geoffrey Irving, Shane Legg, and Dario Amodei. Reward learning from human preferences and demonstrations in atari. Advances in neural information processing systems, 31, 2018.
[15] Yuhua Jiang, Qihan Liu, Yiqin Yang, Xiaoteng Ma, Dianyu Zhong, Hao Hu, Jun Yang, Bin Liang, Bo XU, Chongjie Zhang, and Qianchuan Zhao. Episodic novelty through temporal distance. In The Thirteenth International Conference on Learning Representations, 2025.
[16] Harini Kannan, Danijar Hafner, Chelsea Finn, and Dumitru Erhan. Robodesk: A multi-task reinforcement learning benchmark. https://g /tttee/////tutuu./// wwwtuuu.oooo/ooo/ 2021.
[17] Elia Kaufmann, Leonard Bauersfeld, Antonio Loquercio, Matthias Müller, Vladlen Koltun, and Davide Scaramuzza. Champion-level drone racing using deep reinforcement learning, Nature, 620(7976):982-987, 2023.
[18) Changyeon Kim, Jongjin Park, Jinwoo Shin, Honglak Lee, Pieter Abbeel, and Kimin Lee. Preference transformer: Modeling human preferences using transformers for arXiv preprint arXiv:2303.00957 2023.
[19] Christopher A Kurby and Jeffrey M Zacks. Segmentation in the perception and memory of events. Trends in cognitive sciences, 12(2):72-79, 2008
[20] Gilwoo Lee, Tomás Lozano-Pérez and Leslie Pack Kaelbling. Hierarchical planning for multi- contact non-prehensile manipulation. In 2015 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 264-271, 2015.
[21] Kimin Lee, Laura Smith, Anca Dragan, and Pieter Abbeel. B-pref: Benchmarking preference- based reinforcement learning. arXiv preprint arXiv:2111.03026, 2021.
[22] Kimin Lee, Laura M Smith, and Pieter Abbeel. Pebble: bredback-efficient interactive rein- forcement learning via relabeling experience and unsupervised pre-training. In International Conference Machine Learning, pages 6152-6163. PMLR, 2021.
[23] Xinran Liang, Katherine Shu, Kimin Lee, and Pieter Abbeel. Reward uncertainty for exploration in preference-based reinforcement learning, arXiv preprint urXiv:2205.1 120./ 2022.
[24] Runze Liu, Fengshuo Bai, Yali Du, and Yaodong Yang. Meta-reward-nett Implicitly differ- entiable reward learning for preference-basee reinforcement learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances Neural Information Processing Systems, 2022.
[25] Yao Luan, Qing-Shan Jia, Yi Xing, Zhiyu Li, and Tengfei Wang. An efficient real-time railway container yard management method based on partial decoupling. IEEE Transactions on Automation Science and Engineering, 2025.
[26] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint 1322v.331..5602, 2013.

---

[27] Ni Mu, Hao Hu, Yiqin Yang, Bo Xu, and Qing-shan Jia. Clarify; Contrastive preference rein- forcement learning for untangling ambiguous queries. Proceedings of the 42th International Conference Machine Learning. 2025.
[28] Ni Mu, Yao Luan, and Qing- Shan Jia. Preference-bast based multi-objective reinforcement learning. IEEF Transactions on Automation Science and Engineering, 2025.
[29] Ni Mu, Yao Luan, Yiqin Yang, and Qing-shan Jia. S-epoa: Overcoming the indistinguisha- bility of segments with skill-driven preference-bar reinforcement learning. arXiv preprint arXiv:2408. 12130, 2024.
[30] Vivek Myers, Chongyi Zheng, Anca Dragan, Sergey Levine, and Benjamin Eysenbach. Learn- ing Temporal Distances: Contrastive Successor Features Can Provide Metric Structure for Decision-Making. In International Conference on Machine Learning. 2024.
[31] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding, arXiv preprint arXiv: 018800003777 2018.
[32] Jongjin Park, Younggyo Seo, Jinwoo Shin, Honglak Lee, Pieter Abbeel, and Kimin Lee. Surf: Semi-supervised reward learning with data augmentation for feedhack-efffcient preference-based reinforcement learning. arXiv preprint arXiv:2203. /0050, 2022.
[33] Seohong Park, Olch Rybkin, and Sergey Levine. METRA: Scalable unsupervised RL with metric-aware abstraction. The Twelfth International Conference on Learning Representations, 2024.
[34] Julien Perolat, Bart De Vylder, Daniel Hennes, Eugene Tarassov, Florian Strub, Vincent de Boer, Paul Muller, Jerome T Connor, Neil Burch, Thomas Anthony, et al. Mastering the game of stratego with model-free multiagent reinforcement learning. Science, 378(6623):990 996, 2022.
[35] Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds of mutual information. In International Conference on Machine Learning, pages 5171-5180. PMLR, 2019.
[36] Sebastien Racaniere, Andrew Lampinen, Adam Santoro, David Reichert, Vlad Firoiu, and Timothy Lillicrap. Automated curriculum generation through setter-solver interactions. In International conference on learning representations, 2020.
[37] Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, Sergey Levine, and Google Brain. Time-contrastivv networks: Self-supervised learning from video. In 2018 IEEE international conference on robotics and automation (ICRA), pages 1134-1141. IEEE, 2018
[38] Daniel Shin, Anca D Dragan, and Daniel S Brown. Benchmarks and algorithms for offline preference-based reward learning. arXiv preprint arXiv:2301.01392, 2023.
[39] David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driess- che, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mas- tering the game of go with deep neural networks and tree search. nature, 529(7587):484-48) 2016.
[40] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprini arXiv: 180100069000 2018.
[41] Dexin Wang, Chunsheng Liu, Faliang Chang, Hengqiang Huan, and Kun Cheng. Multi-stage reinforcement learning for non-prehensile manipulation. IEEE Robotics and Automation Letters, 9(7):6712-6719, 2024.
[42] Tongzhou Wang, Antonio Torralba, Phillip Isola, and Amy Zhang Optimal goal-reaching reinforcement learning via quasimetric learning. International Conference on Machine Learning, pages 36411-36430. PMLR, 2023.
12

---

[43] Yifan Wu, George Tucker, and Ofir Nachum. The laplacian in RL: Learning representations with efficient approximations. In International Conference on Learning Representations, 2019.
[44] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pages 1094-1 100. PMLR, 2020.
[45] Jeffrey M Zacks, Nicole K Speer, Khena M Swallow, Todd Braver, and Jeremy R Reynolds. Event perception: mind-brain perspective. Psychological bulletin, 133(2):273, 2007.
[46] Dianyu Zhong, Yiqin Yang, and Qianchuan Zhao. No prior mask: Eliminate redundant action for deep reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 17078-17086, 2024.
[47] Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep rl via meta- learning. In International Conference on Learning Representations, 2020.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The main claims made in the abstract and introduction accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: See Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
14

---

Justification:
See Appendix A.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: See Appendix D and Appendix E.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
15

---

Answer: [Yes]
Justification:
We provide code in the anonymous Github linked in Appendix D.3.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /CodeSubmic / //nips. wttpccccpppccccceeececccccccccct /ttp:///giiieeeeeeeuuubiiiiiii ssisnnoillll onPolicy))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
See Appendix and Appendix E.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We present the standard deviation of the results. See Section 5.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
16

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
show in tables or
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
See Appendix D)
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: htap:::://nuuuee http://eerrrrspeeeeeee
Answer: [Yes]
Justification: The research conducted in the paper conforms, in every respect, with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
Answer: [NA]
Justification: This work studies general multi-stage tasks, and is not tied particular applications.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
17

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: See Appendix D.3.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/dataseta has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
18

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We provide code in the anonymous Github linked Appendix D.3, with README file along with it.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [Yes]
Justification: See Appendix E.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [Yes]
Justification: See Appendix E.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
19

---

16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: / /ttp///ersp////// theurips.cc/oonfeeee eeenneer///eee22//////// /LLM/ for what should or should be described.
20

---

A Proof
A.1 Proof of Proposition
Proposition 4. For MDP and trajectories generated by is optimal policy x², consider classifier ((s) that takes state as input and outputs the probability Pe(s) that the state is collected at step e (1,2..... N, T(s) max; Pt(s). Denote the accuracy of the classifier acc, the multi-stage measure F is lower bounded by acc. Er ,EE [max, Pt pr(s)]
Proof. We prove this proposition by constructing stage chain with the classifier.
First, we aggregate the timesteps into [ sets by defining
$ = (tl
i}.
(9)
T/
Then, we define an aggregated classifier / (s) that receives a state and outputs the probability Pez (8) that the given state collected at the aggregated step 1,",
Specifically, we have
Σ (s(),,
(10)
NEE
and the aggregated classifier is defined as
**(() max P.- (s).
(11)
t
The accuracy of the classifier is given by
acc = Eysx' ser [I(t" (s, T) arg max (s)))
(12)
where 1" (s) is the step of state in trajectory lower bounded by acc, since any errors made by the aggregated classifier must also be present in the original classifier, while the reverse is not true.
т. Therefore, the accuracy of the aggregated classifier
Note that the aggregated timestep shows a chain structure. We divide the stages according to these aggregated timestep sets, where the stage w; corresponds to the aggregated timestep set t," Then we model the probability of state being stage w, using the aggregated classifier:
F(s, Pc- (s).
(13)
the classifier is perfectly accurate, i.c. acc 1, the chain constraints defined in (3) are naturally satisfied since the stages determined by the aggregated time steps. In this ideal scenario, achieve a multi-stage measure represented as:
F E SE [max p. (+))-
(14)
However, in practice, acc f 1. In this case, each state s has a worst-case probability acc assigned wrong stage w where * wi " Consequently, the mapping probability of the true stage
acc of being
F F(s.w ( * arg max F(s,w;)) 2 0. 7S10
(15)
My
Therefore, we could derive a lower bound of F:
F ace Eywnt
max p (s)
acc ETart" ,ee
max Σ pr(s)
SEL
(16)
acc Erwrt .EE
maxp:(s)
which concludes the proof.
21

---

Lemma 1. For a multi-class classification problem. where &' is the input space. Y (1, 21,2.... K) } is the set of classes. Consider a classifier f X + Ay, and let fe(x) denote the predicted probability that sample belongs to class k. Suppose the data (z, v(z)) (or (2, y) for simplicity) follows the joint distribution P. Define the accuracy of the classifier as acc R( [T( = arg maxk fi [1(())) If and probability p [0. 11, we have Pr(y = k))((r)= p) p. then we have ace = Ez~p mmxx [x)].
the classifie is perfectly calibrated, i.e., for all classes Ez=p|maxa
Proof. Since the classifier is perfectly calibrated, we have
Pr(v(z) k) f() p)
Pr(v(x) = k. fx( p)
p
Pr(fx(x)=)) =))
Pr(m)Pr(((xx)= = fx(z) = p/z)
(17)
Pr(x)Pr(fx(r)) p/))
Pr(x)Pr(y(x) = -((())(f((z)) = p|x)
Pr(z)Pr(faae pla)) p/x)
Pr(v(x) k|2) pp=
where Pr(x:) represents the marginal distribution of sample x. The fourth line is because the classifier conditionally independent of the ground truth label given the sample x.
Then, we focus the accuracy:
acc (am)) [[V arg max (x))
Ex Σ Pr(yl=)(U(y arg max [x()))
Fxz Σ (w)
= arg max fa (x))
(18)
Eg Farg max, (x))
E [max (x)),
where the third equation is obtained by substituting (17). That concludes the proof.
Proposition 1. For an MDP and trajectories generated by its optimal policy x", consider a calibrated
classifier T(s) that takes a state 5 as input and outputs the probability pe(s) that the state is collected at step &. E (1, 2,..,...,, ((s) = max, (s). Denote the accuracy of the classifier as acc, the multi-stage measure F has lower bound F 2 acc?
Proof. Using Proposition 4, have
F2 acc Eyaxr eee max (x())
(19)
Using Lemma I, we have Ex ser [max; pr())] = acc, which leads to the conclusion straightfor- wardly.
A.2 Proof of Proposition 2 and? 3
The proofs are based on the abstract MDP introduced in Section 3.1. This MDP is fully stage- wise, with a discrete state space S2 and a discrete action space Y, where each state comprises stage. In addition, make the following assumptions:
Human preferences are perfectly aligned with oracle reward function F(w, v). The estimated reward function Fi (w, v) modeled with a tabular model, which is parameterized by matrix a of shape [ x |T].
22

---

The learning process is ideal, i.e., consider the PbRL problem as ranking problem where the true order is defined by the reward function F(w, v), and suppose the model can fully fit all preferences provided during training. Such an oracle learning process exists. For instance, the algorithm could utilize decision forest to comprehensively represent all provided preferences, where each parent node is more preferred than its child nodes.
Lemma 2. Recover the order of nt elements needs C(n) = logz (n!) = OO log n) queries.
Consider a set of n arbitrary elements, where the total number of possible sorting outcomes is n!. The ranking process can be modeled with a decision tree, where each Ieaf node corresponds to a distinct sorted outcome, and each branch represents comparison between elements. The height of this decision tree reflects the minimum number of comparisons required to arrive at a correctly ordered result, i.e., the number of comparisons required can be expressed C(n) loga (n!). Then, it is intuitive to check loga (n!) = On n), which concludes the proof.
Proof.
Proposition 2. In the worst case scenario, the conventional PbRL needs O([[[TT) log(( additional queries to learn the optimal policy compared to the stage-aligned PbRL.
Proof. We consider the worst-case scenario, where the optimal solution can only be determined after recovering the total order relation of action v for each stage w. i.e., (wh, Uy), (we, ui)..... (Wa, U(Ta)) is ordered for each wi-
Stage-aligned PbRL ranks all actions from each stage, which requires e₁ | (() queries to learn the optimal policy.
|2 C(M)
Conventional PbRL. ranks all stage-action pairs (w, v). If the [S2] groups of actions are ordered, the optimal policy derived. Similar to Lemma 2, there are ([2||T]) possible choices, and ([Ω)) of them are correct. Therefore, conventional PbRL requires c log((({[[Ti)) - log((15())) queries to learn the optimal policy.
Compare and C22 we have
(122 TI) exp (c2 a1)
[2 ||
(20)
e el (IT)
22x
|| -(2|+|TD)
where the approximation is based Stirling's ' formula, ie., n! 2 (3) simplify the above equation to
As |22], [T| > c,
1 r Mill
1
[22 ||
[
exp (C2 - -1)=
> 1.
(21)
2m
22m
Then, we have
22
ci --
(22)
Which concludes the proof.
Proposition 3. If the stage reward bias is sufficiently large, such that the reward ordering between U5)) and (wy, Uzy) depends solely on W; and we for i + i' then in the worst case scenario, conventional PbRL requires (( log((T() additional queries to learn the optimal policy compared to slage-aligned PbRI
Proof. In scenarios with large stage bias, queries with # i' do not contribute to the ranking (in fact, only [S2] of them do some contribution, as to determine the order of w;).
The probability to sample wi, s.t. i = i! in all stage-action pair is
[ Ciri | 1) = |2
(23)
As stage-aligned PbRL ranks all actions from each stage, which requires C(M)) = [9]] [())( queries to learn the optimal policy, the conventional PbRL. requires 0 (r log([T)) queries to derive the optimal policy.
|
23

---

A.3 Theoretical Analysis for the Quadrilateral Distance
Proposition 5. Let Q be the set of stages and dsn x R be quasimetric temporal distance
function. Assume:
(1) (Ideal stage partition) For each state E S, there exists a unique stage w E Q such that (ss,w)= =1, denoted as E w.
(2) (Stage compactness) For any stage w St there exists 0 2 S.t. Vs, s' w, dsp(s, s') S 6,
(3) (Cross-stage separation) There exists o7 > s.t. if E w and of E w' with w # w. then
dsp(s, s') 2 6-
(4) (Boundedness) There exists A 2 8 S.I. Vs, E ES, S')
For two segments do = [+); - and 0s = = (...ft11) we define the quadrilateral distance d 01) in (7). As visualized in Figure 4, d Muuge(ar σ1) is calculated the average value of temporal distances between the two start and end points of the two segments:
dstage 01) = [dsp[=0] ii) + dsp(2 ;
1 ; + den s1 - dsn H-1, [9)]]
Then, datage satisfies the upper and lower bounds as in Table 3, and the upper bounds and lower bounds are strictly increasing across each alignment case, indicating smaller distance values for stage-aligned queries.
Table 3: The upper and lower bounds of quadrilateral distance daaae go. σ1) in alignment cases.
Case
Stage Alignment Condition
Bounds of dsuge ( F1)
All start and end points E w $0 of w, $0 --1 s # E w' d sf E w, others & w All start and end points are distinct stages
0 S datuge (an, 01)
I 6++
B
5 /2 S dsugee (oo, 01) (5+
A+)
C
36 /4 S dstago (oo, ②1) S (0,
+ +3A )/4
d S dstagel (Tp, σ1)
SA-
Proof. Case A (Complete Alignment): If all endpoints belong the same stage w, the stage compactness implies all terms in dstuge (oo, σ1) to be smaller than d,t. Thus,
dauug (00, S 1 (40t) = -5
(24)
Non-negativity of dsD ensures dsuage 2 0.
Case (Endpoint Alignment): For the aligned start points sy. of E w and the aligned end H-1 E w', the cross-stage terms in dougolooo σ1) satisfy dsp($g. " 2 6, and points s0 .81
(s -1 5) 2 8- Thus, the lower bound and upper bound of dssuge (a σ1) are
1 10+0+0 0 +65)
5c
dsiage
(25)
8 +
dstago
18. 4
+ + + + A+)
(26)
2
Case C (Partial Alignment): With only s6. l4 E w, three terms in datgg (go, 01) involve cross-stage pairs. Thus, the lower bound and upper bound of daage (0o, σ1) are
+ 8. = datage 2 1 (0 00 5-
35
(27)
4
1 4 10 + + A
8
+3
daage
(28)
Case D (No Alignment): In this case, all pairs are cross-stage. Thus, the lower bound and upper bound of daage (o 1 are
8 < dstage S A
(29)
24

---

Priority Ordering: From A > o, we derive:
+
A
8
+
3
A
(30)
2
4
30
0
6.
(31)
2
Thus, strict inequality A < << <<C+ D holds for both bounds.
B
Algorithm Implementation
We illustrate the full process of STAIR as in Algorithm I andd 2.
Algorithm STAIR
Require:
Feedback frequency K, number of queries per feedback session M, Total feedback N temporal distance update frequency KSD
1: Initialize replay buffer D, DSD, feedback buffer D
2: Initialize the policy [a)ss) with unsupervised pretraining [22]
3: for each iteration do
4: Rollout with (alss) s) and store (s,arr,s' s') into D, DSD
5: if iteration % = 0 and 'D7 < NNotal then
6: Select { 01) ) J=z
DD using stage-aligned query selection (see Section 4.2) Query the teacher for preference {y] 8: Store preference {(ov, 01, 1) into D end if
7:
9:
10:
Update the reward model fg with D using (2) if iteration % KSD = then
12:
Update the temporal distance model with DSD using (6)
13:
Set D
14:
end if
15:
Relabel D with fv
16: Update the policy = (als) using D
17: end for
Algorithm 2 STAGE- QUERY SELECTION
Require: Number of candidate queries Ng, number of queries per feedback session M
1: Sample Ne segment pairs {(o), )
2: Initialize query selection vector of shape No with zeros: i i= (0.......00) 0).
3: for each segment pair (a), of) do
4: Calculate selection score and store it in i: i(i) + (o of)
5: end for
6: Select M queries with the largest selection score î
25

---

C Challenges in Stage-Aligned Reward Learning
Challenge 1: Stage definitions are often based on subjective expertise and lack generalizabilit across tasks.
Challenge 2: Stage measurement should remain valid and reliable as the policy evolves. [Take timestep as an example)
Stage distribution along the trajectory
Go to the box
Go to the box
(BEB)
Policy in early epochs: falls reach the box
timestep
Fetch: Before obtaining the box
10
Fetch the box
(0
Fetch: After obtaining the
Policy in later epochs: complete the entire task
Stage is timestep
Return with the box
Return with the box
Aligning the timestep does not align the stage since the policy evolves!
Two reasonable stage definitions for the go-fetch-return task.
Figure 9: Illustrations of the two challenges that emerged in the design of the stage-aligned reward
learning method.
Stage definitions are often based on subjective expertise and lack generalizability across tasks. The definition of the stage often requires expert insights and subjective judgment, and can vary significantly across domains and individuals. For instance, in the go-fetch-return example shown in Figure I, one might opt to further subdivide the retrieval stage into two parts: before and after obtaining the box. Therefore, there is need for an alternative framework that approximates stages without requiring explicit stage definitions. Meta-learning approaches learning task embeddings for segments [47] or clustering methods [46], could potentially address this challenge. However, these methods are typically complex and demand large amounts of data for effective learning, which is unavailable in PbRL.
Stage measurement should remain valid and reliable as the policy evolves. Take an intuitive approximation, timestep, as an example. For a fixed policy, the timestep serves as a continuous approximation of the stage, states that are close together stage. However, the timestep relies on the sampling policy, which cannot be relabeled once collected. As the policy evolves during the training process and given that state-of-the-an PbRL methods off-policy for query reuse, segments collected in the early training phases may confuse the query selection in later phases. Therefore, it is essential develop stage approximation that is tied to specific policy and can be easily relabeled according to the trained policy.
trajectory tend belong to the same
26

---

D Experimental Details
D.I Analysis of the Impact of Stage Misalignment
This section provides the setups and implementation details of experiments in Section 3.
Details the classification experiment in Section 3.1. We conduct the classification experiment in Section 3.1 on Metaworld door-open, drawer-open, and window-open tasks. For each task, we first train policy using SAC [11] with the ground truth task reward the expert policy, and then collect le6 transitions (s, r, t) with the trained policy, where is the timestep. The SAC is implemented with the code provided by BPref [21], and is consistent with the SAC used in Section 5. The detailed hyperparameters of SAC are shown in Table 7.
Then, to verify the multi-stage property, we train a classifier which outputs the estimated timestep f of the input state Among the collected 1 10th transitions, we randomly select 5 x 10 transitions for training the classifier, and report the result that evaluated on the other 5 x 105 transitions. The detailed hyperparameters of the classifier are shown in Table 4.
Table 4: Hyperparameters of the timestep classifier.
Hyperparameter
Value
Size of hidden layers
256
Number of hidden layers
3
Train epochs
20
Batch size
256
Learning rate
1 10-1
Optimizer
Adam
Details the human segment preferences experiment in Section 3.2. To show the existence of the stage reward bias, where the reward here refers to the underlying reward of humans, we sample segments of length 20 from the SAC trajectories, compose them into queries randomly, and then let humans label preferences. Details about the human experiments are shown in Appendix E.
Details the abstract MDP experiment in Section 3.2. To show the benefits of stage-aligned query selection in multi-stage problems, we instantiate the abstract MDP introduced in Section 3.2 with [ 101, | = 5, (1.... 100, T). as shown in Figure 2. The reward function is F(w, v) = Fin (w, v) + Fstage (w), where (u denotes the stage reward bias. Faage(w;) ~ Uniform (0, Riins] (w,,, Uy) Uniform(0, 10). Roins indicates the strength of the stage reward bias. Additionally, we normalize the reward to ensure fixed scale across all Roias values, eliminating the need tune hyperparameters for policy training. Specifically, we normalize the reward function as follows:
F(w, v) 1 Ene min, F(w', e') 13
F' (w.
(32)
64 Σn, max, F(w', u') [& Ew min u()
ie., normalize the reward such that the performance of an arbitrary policy is in [0, [2[]. We assume that human preference could be perfectly modeled by this underlying reward function.
In this problem, we use tabular reward function (s, a) parameterized by a matrix 1a of shape [2] x |.). The conventional PbRL and the stage-aligned PbRL are implemented as described in Algorithm 3, which details the training process of this reward model based on the Bradley-Tery model. Specifically, lines 3~4 describe the query collection process of the two methods: conventional sampling uniformly samples state-action pairs from the entire state space, while stage-aligned sampling ensures that both state-action pairs come from the same stage. Then, in line 5, the reward model is trained by optimizing the cross-entropy loss.
27

---

Algorithm 3 Tabular PbRL Algorithms for the Abstract MDP
Require: Number of queries in one epoch M.
1: Initialize tabular reward function Ty (8, a) by setting vb to zero matrix.
2: for each epoch do (For conventional PbRL) Randomly sample 2M state-action pairs {((3), a:) ] composing
M queries D {(($2i-1, 1: 022-1); (22)) )= i=m--
4: (For stage-aligned PbRL) Randomly sample M states {8 and 2M actions J
comprising M queries D {((8;, 02;-1); (8s, (2:)) T==1 )
5: Update & by conducting gradient descent with loss function (2) with D.
6: end for
Hyperparameters of the tabular PbRL algorithms are shown in Table
Table 5: Hyperparameters for policy training the abstract MDP.
Hyperparameter
Value
Train epochs
100
Number of queries in one epoch
200
Learning rate Optimizer
0.05
Adam
D.2 Tasks
The locomotion tasks from DMControl [40] and robotic manipulation tasks from Meta World [44] used in our experiments are shown in Figure 10.
(a) door-open
(b) sweep-into
(c) window-open
(d) window-close
(c) door-unlock
(f) faucet-open
(g) walker
(h) quadruped
(i)
push-green
Figure 10: Rendered images tasks from MetaWorld (a-f), DMControl (g-h), and RoboDesk (i).
Robotic manipulation tasks in Meta World. MetaWorld [44] provides diverse high-dimensional robotic manipulation tasks. We choose the following complex tasks this work.
Door-open: An agent controls robotic arm to open door with revolving joint. The door is placed at a random position.
Sweep-into: An agent controls robotic
sweep block into the hole. The block is a random position. The hole is placed at a fixed location.
Window-open: An agent controls a robotic arm to open a window, The window is placed at random position.
28

---

Window-close:
An agent controls a robotic arm to close window. The window is placed at random position.
Door-unlock: An agent controls robotic arm unlock door by rotating the lock counter- clockwise. The door is placed at a random position.
An agent controls a robotic arm to open a faucet by rotating the faucet coumter-clockwise. The faucet is placed at random position.
Faucet-open:
Please refer to [44] for detailed descriptions of the state space, action space, and the ground truth rewards.
Locomotion tasks DMControl suite. DMControl suite [40] provides diverse high-dimensional locomotion tasks. We choose the following complex tasks in this work.
Walker-rum: A planar walker is trained to control its body and walk on the ground.
-uadruped-walk: A four-legged ant is trained to control its body and limbs, and crawl slowly on the ground.
Quadruped-rum: A four-legged ant is trained to control its body and limbs, and crawl fast on the ground.
The ground truth reward incorporates components designed to promote forward velocity across all tasks. Additionally, for tasks such as Walker-run, Quadruped-wasd walk, and Cheetah-run, supplementary terms are included to encourage upright torso posture.
Robotic manipulation tasks in RoboDesk. RoboDesk [16] provides diverse robotic manipulation tasks that test diverse behaviors within the same environment. We choose the following complex tasks in this work.
Push-green: A robotic arm is trained to push the green button to turn the green light
We modify the state representation from the original implementation [16], while keeping the action and ground truth rewards unchanged. Specifically, the state is represented 68-dimensional vector that combines current and previous timestep information, including the end-effector coordinates (3 dimensions), the wrist joint angle (I dimension), the gripper finger position dimension), and the positions of all objects in the scenario (29 dimensions).
D.3 Implementation Details
The implementation of STAIR is based on BPref [21]. Code is available
https://githes httpa///ggthuu...omiiiiiiiiie com/iiiiiii////A.
For the implementation of SAC [11], PEBBLE [22], RIME [6], RUNE [23], and QPA [13], refer to their corresponding official repositories, shown in Table 6.
Table 6: Source codes of baselines.
Algorithm
Url
License
SAC. PEBBLE
MIT
RUNE
MIT
MRN
http:///gghhbbbccoomyaaaauu////p
MIT
RIME
httpe://gggtauocccmm////leiiitne ICML2024
MIT
QPA
MIT
SAC serves as a performance upper bound because uses the ground-truth reward function, which unavailable in PbRL settings for training. The detailed hyperparameters of SAC are shown in Table 7. PEBBLE's settings remain consistent with its original implementation, and the specifics are detailed in Table 8, For RUNE, MRN, RIME, QPA and STAIR, most hyperparameters are the same those
29

---

of PEBBLE, and other hyperparameters are detailed in Table 9, 10, 11, 12 and 13, respectively. For RIME, employ an oracle script teacher BPref by setting eps_skip=O, replacing the noisy teacher, as non-ideal feedback falls outside the scope of this paper. For QPA, we remove the data augmentation to ensure fair comparison, as none of the other baselines incorporates this technique, and the authors of QPA have noted that data augmentation does not have significant effect on performance in the Meta World environment the official repository QPA. The total amount of feedback and feedback amount per session are detailed in Table 14.
The experiments are conducted on a server with Intel(R) Xeon(R) Platinum 8352V CPU, 512 GB RAM, NVIDIA RTX 4090 GPU, and Ubuntu 20.04 LTS. For all baselines and our method, we run different seeds, and report the mean performance and the standard deviation.
Table 7: Hyperparameters of SAC.
Hyperparameter
Value
Number of layers
(DMControl), (MetaWorld
Hidden units per layer
1024 (DMControl), 256 (MetaWorld
Activation function
ReLU
Optimizer
Adam
Learning rate
0.0005 (DMControl), 0.0001 (MetaWorld)
Initial temperature
0.2
Critic target update freq
Critic EMA
0.01
Batch Size
1024 (DMControl), 512 (MetaWorld
(31, B2) Discount
(0.9, 0.999)
0.99
Table 8: Hyperparameters of PEBBLE.
Hyperparameter
Value
Segment length
50
Learning rate
0.0005 (DMControl), 0.0001 (MetaWorld)
Feedback frequency
20000 (DMControl), 5000 (MetaWorld)
Num of reward ensembles
Reward model activator Unsupervised pretraining steps
tanh
9000
Table 9: Hyperparameters RUNE.
Hyperparameter
Value
Initial weight of intrinsic reward 3g
0.05
Decay rate p
0.001
30

---

Table 10: Hyperparameters of MRN.
Hyperparameter
Value
Bi-level updating frequency
5000 (Cheetah, Hammer, Button Press), 1000 (Walker)
3000 (Quadruped), 10000 (Sweep Into)
Table 11: Hyperparameters RIME.
Hyperparameter
Value
Coefficient in the lower bound Nlower
0.5
Minimum weight 8min
Maximum weight Smax
3
Decay rate k
1/30 (DMControl), 1/300 (MetaWorid)
Upper bound Tupper
3 In
for the intrinsic reward
1 10
Steps of unsupervised pre-training
9000
Table 12: Hyperparameters of QPA
Hyperparameter
Value
Learning rate
0.0005 (walker-run),
0.0001 (quadruped-ward -wak,, quadruped-run, Meta Morld)))d)
Size of alliny- buffer N
30 (door-unlock), 60 (door-open),
10 (Other tasks)
Data augmentation ratio 7
20
Hybrid experience replay sample ratio
0.5
Min/Max length of subsampled snippets
[35, 45]
Table 13: Hyperparameters of STAIR.
Hy perpaarrreee
Value
State coefficient of the quadrilateral distance Gstate
0.1
Stage coefficient of the quadrilateral distance Cstage
2
The frequency of temporal distance update/sD
1
Learning rate of the temporal distance
3 x 10-4
Number of layers of Temporal distance
3
Hidden units per layer
256
Energy function in Temporal distance Contrastive loss function in Temporal distance
MRN-POT [15]
InfoNCE Symmetric [15]
Table 14:
Feedback amount in each environment. The "value" column refers to the feedback amount
in total / per session.
Environment
Value
Environment
Value
walker-run
1000 /100
window-open
2000/50
quadruped-wal walk
2000/200
window-close
2000/50
quadruped-run
2000/200
door-unlock
5000/50
door-open
5000/50
faucet-open
3000/50
10000/50
push-green
1000/50
sweep-into
10000/50
20000/50
31

---

E Human Experiments
E.1 Preference collection
In human experiments, we collect feedback from human labelers (the authors) familiar with the tasks.
Task 1. For the human experiment in Section 1, the labelers provide 20 pieces of feedback to each task. Since the task is completed in 100 steps, we only sample segments that begin before the 100th step. Each segment about second long, which has 20 timesteps. The labelers are instructed to watch a video rendering each segment and determine which one performs better in achieving the specified objective. For each query, the labelers are presented with three options: (I) dg is better, (2) 01 is better, and (3) the two segments are indistinguishable.
Task For the human experiment in Section 5.3, the labelers provide 50 pieces of feedback to each task. Each segment is about 2 seconds long, which has 50 timesteps. To ensure faimess, we shuffle the queries generated by the algorithms so that the labelers do not know the algorithm that generates the queries. The labelers are instructed to watch video rendering each segment and determine whether the two segments are the same stage of this task from their perspective. For each query, the labelers are presented with two options: (1) the two segments are in the same stage, (2) the two segments are in different stages.
F,2 Guidance human labelers
Below, we provide the instructions we provided to the human labelers. The instructions are inspired by [7, 18. 29].
Door-open.
The target behavior is that the robot arm smoothly rotates the door until it stays fully open clearly visible angle. (For task only) If the arm moves abnormally, lower your priority for the segment.
The target behavior is that the drawer is fully extended to its final position with controlled, direct pushing. (For task only) If the arm moves abnormally, lower your priority for the segment.
Drawer-open.
Window-open.
The target behavior is that the window slides horizontally to clearly open position with coordinated gripper guidance. (For task 1 only) If the arm moves abnormally, lower your priority for the segment.
The target behavior that the window slides horizontally to a clearly closed position with coordinated gripper guidance. (For task 1 only) If the arm moves abnormally, lower your priority for the segment.
Window-close.
32

---

More Experimental Results
F.1 Multi-Stage Property and Stage Misalignment
In this section, we provide additional results for Section 3. Figure 11 provides additional human experiments validating the existence of stage reward bias in MetaWorld domains, which serves example for practical tasks. Figure 12 provides the episode reward and reward estimation error of different Roiss in the abstract MDP model.
Prefer do
Prefer
Prefer Both
door-open
draw trawer-open
window-open
100
Start Step Segment
60
40
20
20
20
0
0
20
40
60
80
100
20
40
60
80
100
0
20 40
60
80
100
Step Segment
Segment
Start Step Segment
Figure II: Human Preferences of segments started at different timesteps in the MetaWorld Each
point (ta, by) represents that segment 0 and segment are collected from steps t and by respectively. Humans prefer segments in later timesteps, suggesting stage reward bias where the humans' underlying reward is higher in these later stages.
Stage-Wligned PbRL
Conventional PbRL
Roua
Rai = 10
Rava 20
Ravs 30
Rav 40
Foleate Bewoord
90
M0
90
90
90
60
60
60
45.
Reaned Cap
100
75
50
50
100
100
Fpochs
Pporhs
Sports
Spoctts
Ppochs
Figure 12: Episode reward (upper row) and reward estimation error (lower row) of different Roins in the abstract MDP model. The solid line and the shaded area represent the mean and the standard deviation of the corresponding value.
F.2 Query Visualization
In Figure 13, we visualize the queries selected by PEBBLE and STAIR, respectively. The segment pair (0o, 01) selected by PEBBLE is in different stages: the window is already closed in oo. while the arm is in the process of closing the window d1e Comparing these two segments does not provide the with sufficient information learn how perform the task of closing the window. In contrast, for the segment pair (o6- of) selected by STAIR, the arm is closing the window in two different ways, and the two segments within the same stage. This comparison directly provides the arm with actionable information on the mechanics of closing the window, making it more beneficial for learning.
E33 Performance noisy feedback
To evaluate STAIR's robustness to noisy feedback, which mimics imperfect and inconsistent human feedback, we conduct experiments considering two types of "scripted teachers". following prior work
33

---

Query selected by PEBBLE: (op, oi)
G
Query selected by STAIR: (og, oj
Figure 13: Visualization of segment pairs selected by (a) PEBBLE (disagreement-basee based query selection) and (b) STAIR (stage-aligned query selection), under the window-open task.
[21]. Specifically, consider two types of "scripted teachers": (1) Error teacher: A teacher with random error rate c = 0.1, resulting in 10% incorrect feedback. (2) Inconsistent teacher: Feedback is randomly sampled from mixture of two sources: myopic teacher with discounted factor y 0.9, and an error teacher with e = 0.2.
We show the results door-open (Nout 5000) and sweep-into (Nocal 10000) in Table 15. As shown in the table, STAIR consistently outperforms baselines under both conditions, highlighting its robustness to non-ideal feedback.
Table 15: Performance on noisy feedback.
Door-open
Sweep-into
Teacher
Error
Inconsistent
Error
Inconsistent
STAIR
99.89 0.99
98.53 136
49.12 1471
56.67 11.18
PEBBLE
91.41 = 661
88.83 + 7.29
29.64 = 11:6
29.86 + 14.1
RUNE
63.15 13.50
75.97 13.28
11.82 5.88
10.66 8.66
34

---

G Further Discussions
Explanation for performance in single-stage tasks. In single-stage tasks, the performance of STAIR primarily comes from the induced implicit curriculum learning mechanism, where the method adaptively adjusts the learning focus based on the evolving policy.
To explain how the curriculum learning works, we use the Quadruped task as an example. In the quadruped task, early training with STAIR might prioritize selecting segments before and after fall (which has a small temporal distance), helping the agent learn stability. As training progresses and the policy improves (with the quadruped becoming such segments (before and after fall) increases. At this point, STAIR shifts its focus to segments where the quadruped shows different movement behaviors, rather than emphasizing stability-related segments. This gradual shift enables the agent to learn better movement behaviors while avoiding excessive focus on already-lean behaviors like maintaining stability.
stable), the temporal distance between
This induced automatic curriculum learning mechanism implicitly divides the reward learning process into stages by introducing queries progressively. In this way, later learning stages (e.g., learning how to walk faster) are presented only after the agent masters the earlier ones (e.g., ensuring stability), enabling the model focus the complexities of the newly added stages. Recent works have demonstrated the effectiveness of automatic curriculum learning, which guides the agent with tasks that align with its current capabilities [10, 36].
35


Reviews:
[{"review_id": "U90lrVncYG", "reviewer": "Reviewer_iPSk", "summary": "This paper identifies and addresses the problem of stage misalignment in preference-based reinforcement learning (PbRL), where comparisons between behavior segments from different stages of a multi-stage task lead to uninformative or misleading feedback. The authors propose STAIR (STage-AlIgned Reward learning), which learns an on-policy, task-agnostic temporal distance metric via contrastive learning. This metric is used to filter for stage-aligned segment pairs when querying preferences. A novel “quadrilateral distance” measures stage similarity between segments. STAIR significantly improves feedback efficiency and policy performance on multi-stage tasks (MetaWorld, RoboDesk) and remains competitive on single-stage benchmarks (DMControl). A small-scale human study confirms that STAIR-selected queries align well with human perception of stages.", "questions": "1. How does STAIR perform with noisy or inconsistent human feedback? Injecting noise into oracle labels could help assess robustness.\n2. Can STAIR handle branching or non-linear stage structures? Most experiments assume linear sequential tasks.\n3. How could STAIR extend to more general feedback formats (e.g., ranked lists or scalar scores)?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. Strong theoretical justification (Propositions 2–3) and rigorous evaluation.\n2. STAIR consistently outperforms PbRL baselines on multi-stage tasks with far fewer queries.\n3. The paper is well-written and structured. Key concepts (temporal distance, stage alignment) are clearly explained with useful diagrams.\n4. The paper tackles an important underexplored issue in PbRL with implications for real-world learning from humans.\n5. The paper enjoys strong potential to improve sample efficiency for long-horizon tasks.\n\nWeaknesses:\n1. Human feedback is simulated via an oracle—true robustness to real-world human noise remains untested.\n2. The proposed algorithm adds considerable complexity (contrastive learning, extra hyperparameters), which may hinder adoption.\n3. The limited demonstration with noisy or real human-in-the-loop training slightly makes real-world application less practical.", "was_revised": true, "final_justification": "I have read the author response and have no further questions. I'll maintain my score.", "rebuttal": "Dear Reviewer,\n\nThanks for your valuable and detailed comments. We hope the following statement clears your concern. \n\n**W1, W3, Q1: Robustness on noisy or inconsistent feedback.**\n\n**A for W1, W3, Q1:**\nAs suggested, we conduct experiments to evaluate STAIR's robustness to noisy feedback, which mimics imperfect and inconsistent human feedback. Following prior work [1], we consider two types of \"scripted teachers\": \n1. Error teacher: A teacher with a random error rate $\\epsilon=0.1$, resulting in 10\\% incorrect feedback. \n2. Inconsistent teacher: Feedback is randomly sampled from a mixture of two sources: a myopic teacher with discounted factor $\\gamma=0.9$, and an error teacher with $\\epsilon=0.2$. \n\nResults on door-open ($N_\\text{total}=5000$) and sweep-into ($N_\\text{total}=10000$) show that STAIR consistently outperforms baselines under both conditions, highlighting its robustness to non-ideal feedback. \n\n||Door-open||Sweep-into||\n|:-:|:-:|:-:|:-:|:-:|\n|Teacher|Error|Inconsistent|Error|Inconsistent|\n|STAIR|**99.89±0.09**|**98.53±1.36**|**49.12±14.71**|**56.67±11.18**|\n|PEBBLE|91.41±6.61|88.83±7.29|29.64±11.66|29.86±14.41|\n|RUNE|63.15±13.50|75.97±13.28|11.82±5.88|10.66±8.76|\n\n\n**W2: Complexity concerns.**\n\n**A for W2:**\nTo address your concern, we analyze the complexity of contrastive learning and the complexity of the extra hyperparameters, respectively.\n\n(1) \n**As for the complexity**, STAIR introduces additional components, but the overall complexity does not increase significantly. \nWe evaluate this in terms of both training and inference:\n- Training: STAIR's training time is approximately 1.27 times that of PEBBLE, which remains manageable. This is because the temporal distance model is a lightweight three-layer MLP, and its update frequency does not exceed the policy. \n- Inference: During inference, STAIR relies exclusively on the SAC policy, with no additional components applied. Therefore, its inference time matches that of existing methods like PEBBLE or SAC, ensuring no increase in runtime complexity during deployment.\n\n(2) **Regarding hyperparameters**, the additional hyperparameters introduced by STAIR are consistent across all environments. Besides, ablation studies in Section 5.4 show the robustness of STAIR to query selection parameters, and provide guidance for the selection of temporal distance update frequency, which can facilitate more effective applications. \n\n\n\n**Q2: Applicability in non-linear stage structures.**\n\n**A for Q2:**\nWe believe STAIR can be adapted to handle branching or non-linear stage structures. \nSTAIR's core mechanism relies on using temporal distance to measure stage differences. \nThe temporal distance used in our work [2], defined as \n$$\nd^{\\pi}_{\\text{SD}}(x, y) = \\log \\left( {p^{\\pi} _\\gamma (s _+ = y | s_0 = y)}/{p^{\\pi} _\\gamma (s _+ = y | s_0 = x)}\\right),\n$$\nwhich quantifies the transition probabilities between states under a given policy and does not depend on linear or sequential task structures. \nThis independence makes it inherently adaptable to more complex scenarios, such as branching or non-linear stage structures, as long as temporal distance appropriately reflects stage differences in these settings.\nExtending STAIR to explicitly handle branching or non-linear tasks, along with analyzing its performance in such contexts, could be a promising direction for future research.\n\n\n\n**Q3: Extension to general feedback formats.**\n\n**A for Q3:**\nSTAIR can generalize to broader feedback formats by extracting pairwise feedback from them:\n- Ranked Lists: Given a set of samples ranked by human preference, we can sample pairwise comparisons from the ranked list to construct pairwise feedback. \n- Scalar Scores: When humans assign scalar scores to samples, we can sample pairs from the dataset and use the scores to derive pairwise feedback by comparing their relative magnitudes.\n\nWe thank the reviewer for this insightful suggestion, which we will consider in our future work.\n\n\nThanks again for the valuable comments. We sincerely hope our additional experimental results and explanation have cleared the concern. More comments on further improving the presentation are also very much welcome. \n\n**References:**\n\n[1] Lee, Kimin, et al. B-pref: Benchmarking preference-based reinforcement learning.\n\n[2] Myers, Vivek, et al. Learning temporal distances: Contrastive successor features can provide a metric structure for decision-making."}, {"review_id": "vgGhCI23KJ", "reviewer": "Reviewer_sEyy", "summary": "This paper presents STAIR, a preference-based reinforcement learning (PbRL) algorithm which leverages temporal distance between trajectory segments to select more informative queries for preference labelling. STAIR was specifically designed for multi-stage environments where an agent has to do multiple subtasks, and the preferences between trajectories across tasks are not very informative. Still, empirical results show that STAIR does not negatively affect the performance of single-stage tasks (such as walker-walk in DMControl).\n\nTo estimate the temporal distance of trajectory segments, STAIR leverages a quadrilateral distance based on the successor distance, which is itself learnt through contrastive learning.\n\nExperimental results show that STAIR outperforms many recent PbRL algorithms in multi-stage environments (particularly MetaWorld).", "questions": "**Main questions**\n\n_Answering the following questions would help address the weaknesses identified above, and raise the paper rating_.\n\n*Q1*: in Eq 3., why does $G^\\tau(i) \\le G^\\tau(i -1)$? I understand $G^\\tau(i -1) \\le G^\\tau(i)$ (ie. stages must be in order), and $G^\\tau(i) -1 \\le G^\\tau(i)$ (once a stage changes, we cannot go back to the previous one).\n\n*Q2*: Proposition 1 hold independently of $|\\Omega|$ ($N_\\textrm{stage}$), correct? I believe that is the case based on Appendix A. Either way I would recommend to indicate this in the main text.\n\n*Q3*: I was not able to follow Appendix D.1 and Algorithm 3, was exactly is being trained? It would seem that $r_\\textrm{stage}$ and $\\bar{r}_\\textrm{bias}$ are fixed and sampled independently?\n\n*Q4*: In the Abstract MDP formulation, what does T represent in the expression $i = (0, ..., 100, T)$. \n\n*Q5*:  What is the effect of changing $|\\Omega|$ and $|\\Upsilon|$?\n\n*Q6*: Does learning with $r_\\textrm{bias} = 0$ decay into plain PbRL?\n\n*Q7*: How is $d_\\textrm{state}$ computed? Is it simply the variance of the current reward function across all states in the buffer D?\n\n*Q8*: How many experiment repeats were carried out to compute the standard deviations in Figs 5, 6, and 7?\n\n*Q9*: The results in section 5.3 are problematic for two reasons: a) No statistical significance analysis is carried out and the error bars are quite large (therefore it is possible there is actually only a sampling difference between STAIR and PEBBLE), b) the labellers where the authors themselves, which likely have an unconscious (and natural) bias to agree with STAIR on the task segmentation. The easiest fix is to redo the experiment with unrelated participants (ideally enough to be able to carry out a statistical significance test).\n\n*Q10*: An ablation for eq.8  with only $d_\\textrm{stage}$ and only $d_\\textrm{state}$ would help clarify the contribution of each part of STAIR. \n\n*Q11*: In Table 2, how does STAIR perform with 100 and 50 total feedback? Would it be possible to compare to some of the PbRL methods already included in the main results (many of those improved upon PEBBLE on feedback efficiency).\n\n\n**Minor questions**\n\n\n_These questions likely will not affect the review rating, but would help with the paper readability_.\n\n*Q12*: In Sec. 3.2, line 311, I would not use _indicate_ but something like _chosen to_ or _reflect_.\n\n*Q13*: Following [24], $d^\\pi_\\textrm{SD}$ should be defined as $d^\\pi_\\textrm{SD} = f_\\theta(\\mathbf{y}, \\mathbf{y}) - f_\\theta(x, y)$, otherwise the derivations that follow do not work out. I would additionally mention this is the one-step formulation of the successor distance.\n\n*Q14*: Sec 5.4, line 296 should read _higher_ $K_\\mathrm{SD}$ rather than _smaller_.", "limitations": "Limitations are adequately addressed.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths** \n\n* The research questions are interesting, and address the claims of the paper (namely that STAIR helps with multi-stage tasks, that it does not negatively affect the performance of single-stage task, and that the quadrilateral distance is an adequate form of trajectory clustering).\n* STAIR outperforms many recent PbRL algorithms (PEBBLE, RUNE, MRN QPA, and RIME) in MetaWorld.\n* Figure 1 clearly conveys the main ideas behind STAIR.\n* The theoretical foundations of STAIR are interesting and (mostly) clear (though see questions below).\n\n\n**Weaknesses**\n\n* _W1_: The human evaluation of alignment between stages of STAIR and human judgement is problematic, calling into question whether the observed benefits from STAIR come from stage clustering, rather than a form of curriculum learning (as the paper hypothesises may be happening for single stage tasks). See questions below for specifics.\n* _W2_: Some theoretical aspects of the paper (Figure 3 (right), section 3.2, and $d_\\textrm{state}$) need clarification. See questions below.\n* _W4_: There is no evaluation of an agent based on actual human feedback, instead of an oracle. Experiments showcasing the robustness of STAIR to human feedback (or at least oracle noise) would increase the significance of the paper.\n* _W5_: No citation is provided for \"This issue arises when behaviors from different stages, such as navigation and grasping, are presented to humans for comparison. It leads to ambiguous feedback, as labellers struggle to compare behaviours in distinct subtasks, like efficient movement versus precise manipulation\", a statement which is key in justifying STAIR.", "was_revised": true, "final_justification": "The authors have mostly address the concerns I raised in my review.\n\n* _W1_: The human evaluation of STAIR stages -> authors redid this experiment with new evaluators (unrelated to the paper), and found that STAIR is significantly better at aligning stages with human judgement than the other baselines.\n\n* _W2_: The theoretical questions about the paper were clarified during the rebuttal.\n\n* _W4_: Use of oracle as human-proxy -> authors added experiments with noisy scripted teachers. No evaluation with actual human feedback was provided (this is the main reason I have not increased my score further).\n\n* _W5_: Missing citation -> This was clarified during rebuttal.", "rebuttal": "Dear Reviewer,\n\nThanks for your valuable and detailed comments. \n\n**W4: Performance on noisy feedback.**\n\n**A for W4:**\nAs suggested, we conduct experiments to evaluate STAIR's robustness to noisy feedback, which mimics imperfect and inconsistent human feedback. Following prior work [5], we consider two types of \"scripted teachers\": \n1. Error teacher: A teacher with a random error rate ε=0.1, resulting in 10\\% incorrect feedback. \n2. Inconsistent teacher: Feedback is randomly sampled from a mixture of two sources: a myopic teacher with discounted factor γ=0.9, and an error teacher with ε=0.2. \n\nResults on door-open ($N_{total}=5000$) and sweep-into ($N_{total}=10000$) show that STAIR consistently outperforms baselines under both conditions, highlighting its robustness to non-ideal feedback. \n\n||Door-open||Sweep-into||\n|:-:|:-:|:-:|:-:|:-:|\n|Teacher|Error|Inconsistent|Error|Inconsistent|\n|STAIR|99.89±0.09|98.53±1.36|49.12±14.71|56.67±11.18|\n|PEBBLE|91.41±6.61|88.83±7.29|29.64±11.66|29.86±14.41|\n|RUNE|63.15±13.50|75.97±13.28|11.82±5.88|10.66±8.76|\n\n**W5: Ambiguity of stage-aligned queries.**\n\n**A for W5:**\nWe thank the reviewer for highlighting this important aspect.\nThe ambiguity in human comparisons across different stages can be understood through the lens of Event Segmentation Theory in cognitive sciences [1-2]. This theory suggests that humans naturally perceive continuous actions as segmented into event boundaries. Comparisons that span these boundaries (i.e., comparisons across different stages) significantly increase cognitive load, thereby leading to ambiguous queries. \n\n**W2, Q1, Q2, Q3: Clarification on theory**\n\n- **Q1: Eq.3:**\nWe would like to clarify that Eq.3 does not state $G^\\tau(i)\\le G^\\tau(i-1)$. \nInstead, it states $G^\\tau(i)-1\\le G^\\tau(i-1)\\le G^\\tau(i)$, where the first inequality ensures that the stage changes at most by 1 at each transition, i.e., no stage can be skipped, and the second ensures that stages progress in order without going back. \n\n- **Q2: Proposition 1:**\nYes, Proposition 1 applies to arbitrary $|\\Omega|\\le T$, where $T$ is the number of timesteps of the trajectory. We will indicate it in the main text. \n\n- **Q3: Appendix D.1 and Algorithm 3:**\n  - (1) We briefly introduce the outline of Appendix D.1 and Algorithm 3, respectively.\n    - **Appendix D.1** provides the setups and implementation details of experiments in Section 3.  Specifically, it describes the training process of the classifier in Section 3.1, which demonstrates the multi-stage property by predicting the timestep of a given state. It also introduces the source of human-labeled samples in Figure 3 (left) as well as the details of reward function learning in Figure 3 (right). These experiments reveal the existence of stage reward bias in MetaWorld tasks and analyze its impact by comparing conventional sampling with stage-aligned sampling on the abstract MDP. \n    - **Algorithm 3** details the training process of this reward model based on the Bradley-Terry model. Specifically, lines 3-4 describe the query collection process of the two methods: conventional sampling uniformly samples state-action pairs from the entire state space, while stage-aligned sampling ensures that both state-action pairs come from the same stage. Then, in line 5, the reward model is trained by optimizing the cross-entropy loss.\n\n    In the revised paper, we will add some summary texts as stated above in Appendix D.1 and add some high-level explanations near Algorithm 3 to enhance clarity.\n  - (2) Yes, both $r_\\textrm{stage}$ and $\\bar{r}_\\textrm{bias}$ are environment parameters. They are sampled from predefined distributions.\n\n**Q4, Q8: Clarity of the statement.**\n- **Q4:** $T$ in $i=(0,...,100,T)$ denotes terminate, and $w_T$ is a terminated state in the abstract MDP.\n- **Q8:** We conduct 5 experiments (with different seeds) for each result.\n\n**Q5: The effect of changing $|\\Omega|$ and $|\\Upsilon|$.**\n\n**A for Q5:**\nAs suggested, we conduct experiments using different $|\\Omega|$ and $|\\Upsilon|$ on the abstract MDP, with a fixed stage reward bias of 20, and rewards normalized to be at most 100 across all environments.\nWe report the normalized episode reward for each method. \nResults show that stage-aligned PbRL consistently outperforms conventional PbRL. \nFurthermore, as $|\\Omega|$ or $|\\Upsilon|$ increases, the performance delta in normalized episode reward also increases. This matches Propositions 2 and 3, which state that the additional queries required by conventional PbRL to learn the optimal policy scale with $|\\Omega|$ or $|\\Upsilon|$.\n\n|$\\vert\\Omega\\vert$|$\\vert\\Upsilon\\vert$|Stage-Aligned PbRL|Conventional PbRL|Delta|\n|:-:|:-:|:-:|:-:|:-:|\n|100|5|90.74±1.89|85.65±1.34|5.09|\n|50|5|94.92±2.44|90.22±4.86|4.69|\n|200|5|83.48±1.02|76.55±0.80|6.93|\n|100|3|95.62±1.69|93.06±1.08|2.55|\n|100|8|82.73±2.47|74.40±2.91|8.32|\n\n**Q6: Does learning with $r_\\textrm{bias}=0$ decay into plain PbRL.**\n\n**A for Q6:**\nLearning with $r_\\textrm{bias}=0$ does not decay into plain PbRL, as it still selects stage-aligned queries in multi-stage tasks. Theoretically, the $r_\\textrm{bias}=0$ case is analyzed in Proposition 2, where our method shows a feedback efficiency advantage over plain PbRL. Experimentally, this is supported by the results in Appendix F.1 (Figure 12), where our method outperforms plain PbRL when $r_\\textrm{bias}=0$.\n\n**Q7: How is $d_\\textrm{state}$ computed:**\n\n**A for Q7:**\n$d_\\textrm{state}(\\sigma_0,\\sigma_1)$ is computed as the variance of $P_\\psi[\\sigma_1\\succ\\sigma_0]$ predicted by an ensemble of reward model. \nSpecifically, $d_{\\text{state}}(\\sigma_0,\\sigma_1)=\\text{Var}[P_{\\psi_i}[\\sigma_0\\succ\\sigma_1]^3_{i=1}]$, where $\\{P_{\\psi_i}[\\sigma_0\\succ\\sigma_1]\\}_{i=1}^3$ are identical reward models with different randomly initalized parameters.  \nIt is used to identify queries where the ensemble shows higher disagreement, prioritizing more informative queries. This calculation aligns with works in this field, like [3-4]. \n\n**W1, Q9: Human experiments.**\n\n**A for W1, Q9:**\nTo address your concern, we conducted additional experiments. \n\n(1) Recruitment of participants and mitigation of bias.\n\nAs suggested, we invite 8 colleagues from our department to participate in the human experiment.\nTo avoid potential bias, queries generated by STAIR and baselines were randomly shuffled, ensuring labelers were unaware of the source algorithm.\nPlease note that the original result in section 5.3 is also derived under this mechanism, which ensures fairness, as detailed in Appendix E.1.\n\n(2) Statistical significance analysis.\n\nAs suggested, we perform statistical analyses, including confidence interval calculations and independent two-sample t-tests, to confirm the significance of differences between STAIR and baselines.\n- Confidence Intervals.\n  Due to the limited number of samples, the results in Section 5.3 initially had large error bars. \nWith additional participants, we collect more feedback, which reduces the standard deviation and narrows the confidence intervals (CIs), confirming STAIR's consistent performance in learning human-recognized stages.\n- Two-Sample t-Tests.\n  We conduct independent two-sample t-tests for each task, showing that STAIR achieves statistically significant improvements (p<0.05) in learning human-recognized stages.\n\n|Task|Method|Success Rate|95\\% CI|t-value|p-value|Significance|\n|:-:|:-:|:-:|:-:|:-:|:-:|:-:|\n|door-open|STAIR|0.78±0.11|0.78±0.09|-|-|-|\n||PEBBLE|0.51±0.16|0.51±0.14|3.93|<0.01|✔|\n||MRN|0.5±0.07|0.5±0.06|6.08|<0.01|✔|\n||RUNE|0.55±0.13|0.55±0.11|3.8|<0.01|✔|\n|window-close|STAIR|0.73±0.14|0.73±0.12|-|-|-|\n||PEBBLE|0.55±0.2|0.55±0.17|1.97|0.03|✔|\n||MRN|0.45±0.09|0.45±0.08|4.54|<0.01|✔|\n||RUNE|0.38±0.17|0.38±0.14|4.34|<0.01|✔|\n|window-open|STAIR|0.72±0.19|0.72±0.16|-|-|-|\n||PEBBLE|0.40±0.14|0.40±0.12|3.72|<0.01|✔|\n||MRN|0.53±0.1|0.53±0.09|2.36|0.02|✔|\n||RUNE|0.48±0.16|0.48±0.14|2.6|0.01|✔|\n\n**Q10: Ablation on Eq.8.**\n\n**A for Q10:**\nAs suggested, we conduct ablation studies on $d_\\textrm{stage}, d_\\textrm{state}$ in Eq. 8. The following results show that removing either term degrades performance, demonstrating the effectiveness of the proposed Eq. 8. \n\n||door-open ($N_{total}=5000$)|sweep-into ($N_{total}=10000$)|\n|:-:|:-:|:-:|\n|STAIR|100.00±0.00|91.08±3.42|\n|STAIR w/o $d_{stage}$|85.57±12.77|59.17±21.62|\n|STAIR w/o $d_{state}$|74.48±15.21|34.86±15.09|\n\n**Q11: Query efficiency experiments.**\n\n**A for Q11:**\nAs suggested, we compare STAIR and baselines with 50 and 100 feedback on Metaworld's door-open task. The results show that when $N_{total}\\leq 100$, all algorithms fail to perform, likely because such a small amount of feedback is insufficient to learn a reliable reward function.\nHowever, with $N_{total}=500$, STAIR begins to show strong performance, outperforming all baselines, which shows its effectiveness.\n\n|$N_{total}$|STAIR|PEBBLE|MRN|RUNE|\n|:-:|:-:|:-:|:-:|:-:|\n|50|0±0|0±0|0±0|0±0|\n|100|0±0|0±0|0±0|0±0|\n|500|52.01±23.18|20.00±17.88|19.33±17.29|0.94±0.84|\n|2000|77.77±11.67|28.79±17.02|58.82±21.84|72.03±16.92|\n|5000|100.00±0.00|85.57±12.77|97.40±1.46|72.15±16.45|\n|10000|99.93±0.06|92.53±6.53|93.40±5.46|78.84±17.64|\n\n**Q12-14, PFC: Minor questions, Paper Formatting Concerns:**\nThank you for your keen attention to detail! We have corrected them in the revised version.\n\nWe sincerely thank the reviewer again for the timely and valuable comments. We hope that our response has cleared most of your concerns.\n\n**References:**\n\n[1] Kurby et al. Segmentation in the perception and memory of events.\n\n[2] Zacks et al. Event perception: a mind-brain perspective.\n\n[3] Lee, Kimin et al. Pebble: Feedback-efficient interactive reinforcement learning via relabeling experience and unsupervised pre-training.\n\n[4] Liang, Xinran, et al. Reward uncertainty for exploration in preference-based reinforcement learning.\n\n[5] Lee, Kimin, et al. B-pref: Benchmarking preference-based reinforcement learning."}, {"review_id": "Iw60TJAsnZ", "reviewer": "Reviewer_AkSW", "summary": "This paper proposes STAIR (STage-AlIgned Reward learning), a method designed to address the \"stage misalignment\" issue in preference-based reinforcement learning (PbRL) for multi-stage tasks. Traditional PbRL methods often compare trajectory segments from different subtasks (e.g., navigation vs. grasping), resulting in ambiguous human feedback and reduced policy learning efficiency. STAIR learns a temporal distance via contrastive learning to approximate implicit task stages and prioritizes feedback collection from stage-aligned segment pairs. The authors provide theoretical analysis demonstrating how misaligned queries increase sample complexity and show, through both synthetic and real-world robotic manipulation tasks, that STAIR improves learning efficiency and feedback utilization. Human experiments further validate that STAIR's stage approximation aligns with human cognitive perception.", "questions": "1.\tIs the temporal distance learning process sensitive to environment reversibility? Can it generalize to irreversible domains or single-shot decision tasks?\n\n\n2.    Does STAIR still work under noisy or inconsistent manual feedback (e.g., label flipping, annotator disagreement)? Please give a reasonable explanation or more experimental demonstration.", "limitations": "1.\tThe current method only supports binary preference labels, while in real-world settings, human feedback is often more nuanced (e.g., graded or ranked).\n\n2.\tThe method heavily depends on the quality of contrastively learned temporal embeddings, yet this step lacks ground truth supervision or calibration.\n\n3.\tThe need to compute quadrilateral distances for many segment pairs during inference could become a bottleneck in large-scale or real-time applications.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n1.\tThe paper clearly identifies stage misalignment as a fundamental but underexplored challenge in PbRL for multi-stage tasks.\n\n2.\tThe authors rigorously derive complexity gaps between stage-aligned and conventional query selection through abstract MDP constructions (Propositions 2 & 3).\n\n3.\tThe use of contrastive learning for temporal distance estimation and the novel quadrilateral distance for measuring segment differences are elegant and principled.\n\n4.\tThe evaluation spans both multi-stage (MetaWorld, RoboDesk) and single-stage (DMControl) tasks, demonstrating robustness, generalizability, and superior performance.\n\n5.\tSTAIR does not rely on task-specific stage definitions or priors, making it easy to integrate into existing PbRL frameworks.\n\nWeaknesses\n\n1.\tWhile STAIR shows advantages in single-stage tasks, it lacks a deeper theoretical explanation for why stage alignment is still beneficial when explicit stages are absent.\n\n2.\tSTAIR focuses solely on pairwise preference queries and does not extend to listwise, ordinal, or scalar feedback, which are often more expressive.\n\n3.\tProposition 3 assumes a strong stage reward bias (i.e., human preferences depend primarily on stage position), which may be unrealistic or task-dependent.\n\n4.\tThough contrastive learning and quadrilateral distance are tractable, their performance is sensitive to the update frequency (K_SD), which could increase deployment cost in real systems.", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer,\n\nThanks for your valuable and detailed comments. We hope the following statement clears your concern.\n\n\n**W1: Explanation for performance in single-stage tasks.**\n\n**A for W1:**\nIn single-stage tasks, the performance of STAIR primarily comes from the induced implicit curriculum learning mechanism, where the method adaptively adjusts the learning focus based on the evolving policy.\n\nTo explain how the curriculum learning works, we use the quadruped task as an example.\nIn the quadruped task, early training with STAIR might prioritize selecting segments before and after a fall (which has a small temporal distance), helping the agent learn stability. \nAs training progresses and the policy improves (with the quadruped becoming more stable), the temporal distance between such segments (before and after a fall) increases.\nAt this point, STAIR shifts its focus to segments where the quadruped shows different movement behaviors, rather than emphasizing stability-related segments. \nThis gradual shift enables the agent to learn better movement behaviors while avoiding excessive focus on already-learned behaviors like maintaining stability.\n\nThis induced automatic curriculum learning mechanism implicitly divides the reward learning process into stages by introducing queries progressively. In this way, later learning stages (e.g., learning how to walk faster) are presented only after the agent masters the earlier ones (e.g., ensuring stability), enabling the model to focus on the complexities of the newly added stages. \nRecent works have demonstrated the effectiveness of automatic curriculum learning, which guides the agent with tasks that align with its current capabilities [1-2].\n\nWe will explore a theoretical explanation of the connection between curriculum learning and STAIR in future work.\n\n\n\n**W2, L1: Extension to general feedback formats.**\n\n**A for W2, L1:**\nWhile we agree that STAIR focuses only on pairwise feedback (as noted in the limitations section), other feedback formats can be converted into binary preferences, enabling potential extension of STAIR:\n- Listwise or ordinal feedback: When humans rank a set of samples based on preference, we can sample pairs from the ranked list to construct binary feedback.\n- Scalar scores: When humans assign scores to samples based on certain rules or subjective opinions, we can sample pairs from the scored dataset and construct binary preferences by comparing the scores.\n\nMoreover, though human feedback in real-world scenarios can often be more expressive, obtaining such detailed feedback is typically challenging [3]. \nIn contrast, binary preference feedback remains a common approach in practice [4-6]. \nTherefore, we prioritize binary preference and leave non-binary preference for future work.\n\n\n**W3: Assumption on stage reward bias.**\n\n**A for W3:**\nWhile Proposition 3 does assume a strong stage reward bias, we would like to clarify that our method and experiments do not rely on this assumption. \nInstead, Proposition 3 highlights a specific scenario where our algorithm exhibits greater performance advantages.\nMoreover, for more general cases (without the assumption of stage reward bias), Proposition 2 shows that our method remains effective.\n\n \n**W4, L3: Deployment cost due to update frequency ($K_{SD}$).**\n\n**A for W4, L3:**\nWe acknowledge that the performance of STAIR requires frequent contrastive learning updates (small $K_{SD}$), which introduces additional computational cost. \nHowever, despite this frequent update, the overall complexity of STAIR does not increase significantly. \nThe primary computational burden of contrastive learning update and quadrilateral distance calculation lies in the temporal distance model, and we analyze its impact during training and inference, respectively:  \n- **Training**: STAIR's training time is approximately **1.27 times that of PEBBLE**, which is a manageable increase. This is because the temporal distance model is a three-layer MLP, and its update frequency does not exceed the policy. Additionally, the frequency at which STAIR computes quadrilateral distances for segment pairs is comparable to the frequency at which PEBBLE computes disagreement for segment pairs.  \n- **Inference**: During inference, STAIR relies exclusively on the SAC policy, with no additional components applied. Therefore, its inference time matches that of existing methods like PEBBLE or SAC, ensuring no increase in runtime complexity during deployment.\n\n\n**Q1: Sensitivity and extension to environment reversibility.**\n\n**A for Q1:**\nThe core mechanism of STAIR is based on the learning of temporal distance [7]. Its applicability depends on whether the temporal distance is well-defined for the given domain. Below, we consider the two specific cases you mentioned:\n- Irreversible tasks: The temporal distance learning algorithm proposed in [7] does not assume environment reversibility. Therefore, STAIR can be extended to irreversible tasks.\n- Single-shot decision tasks: Single-shot decision tasks involve only single-step transitions. Temporal distance and stage concepts cannot be meaningfully defined in these tasks. Consequently, STAIR is not applicable in such scenarios.\n\n\n**Q2: Performance on noisy/inconsistent feedback.**\n\n**A for Q2:**\nAs suggested, we conduct experiments to evaluate STAIR's robustness to noisy feedback, which mimics imperfect and inconsistent human feedback. Following prior work [8], we consider two types of \"scripted teachers\": \n1. Error teacher: A teacher with a random error rate $\\epsilon=0.1$, resulting in 10\\% incorrect feedback. \n2. Inconsistent teacher: Feedback is randomly sampled from a mixture of two sources: a myopic teacher with discounted factor $\\gamma=0.9$, and an error teacher with $\\epsilon=0.2$. \n\nResults on door-open ($N_\\text{total}=5000$) and sweep-into ($N_\\text{total}=10000$) show that STAIR consistently outperforms baselines under both conditions, highlighting its robustness to non-ideal feedback. \n\n||Door-open||Sweep-into||\n|:-:|:-:|:-:|:-:|:-:|\n|Teacher|Error|Inconsistent|Error|Inconsistent|\n|STAIR|**99.89±0.09**|**98.53±1.36**|**49.12±14.71**|**56.67±11.18**|\n|PEBBLE|91.41±6.61|88.83±7.29|29.64±11.66|29.86±14.41|\n|RUNE|63.15±13.50|75.97±13.28|11.82±5.88|10.66±8.76|\n\n\n**L2: Temporal embedding lacks supervision.**\n\n**A for L2:**\nWe acknowledge that STAIR depends on the quality of contrastively learned temporal embeddings.\nHowever, similar to existing methods leveraging temporal distance [7], STAIR only requires the relative ordering of temporal distances between states rather than their exact values, as temporal distance is used to rank segments. \nTherefore, contrastive learning enables effective learning of this relative ordering without requiring ground truth supervision, consistent with prior works adopting similar techniques [9].\n\nWe sincerely thank the reviewer again for the timely and valuable comments. We hope that our response has cleared most of your concerns.\n\n**References:**\n\n[1] Florensa, Carlos, et al. Automatic Goal Generation for Reinforcement Learning Agents.\n\n[2] Racaniere, Sebastien, et al. Automated curriculum generation through setter-solver interactions.\n\n[3] Choi et al. Listwise Reward Estimation for Offline Preference-based Reinforcement Learning.\n\n[4] Lee et al. Pebble: Feedback-efficient interactive reinforcement learning via relabeling experience and unsupervised pre-training.\n\n[5] Cheng et al. RIME: Robust Preference-based Reinforcement Learning with Noisy Preferences.\n\n[6] Liang et al. Reward uncertainty for exploration in preference-based reinforcement learning.\n\n[7] Myers, Vivek, et al. Learning temporal distances: Contrastive successor features can provide a metric structure for decision-making.\n\n[8] Lee, Kimin, et al. B-pref: Benchmarking preference-based reinforcement learning.\n\n[9] Jiang, Yuhua, et al. Episodic novelty through temporal distance."}, {"review_id": "z55QtbaZr4", "reviewer": "Reviewer_HXsa", "summary": "The paper addresses a critical challenge in preference-based reinforcement learning (PbRL): stage misalignment in multi-stage tasks. PbRL learns reward functions from human preferences, but when applied to sequential tasks with distinct sub-tasks, comparing segments from different stages leads to uninformative feedback and hinders policy learning. The authors first validate the negative effects of stage misalignment through theoretical analysis and empirical experiments, demonstrating that humans tend to prefer segments from later stages (stage reward bias), and that conventional PbRL methods require significantly more feedback to converge compared to stage-aligned approaches. The authors then propose STAIR, which leverages contrastive learning to approximate stages via temporal distance, and introduces a novel quadrilateral distance for selecting stage-aligned segment comparisons. Extensive experiments on robotic manipulation and locomotion benchmarks show that STAIR outperforms existing PbRL methods in multi-stage tasks and remains competitive in single-stage settings. Human studies further confirm that STAIR's stage approximations align with human cognition, and that STAIR selects stage-aligned queries more effectively than baseline methods.", "questions": "1. **Novelty of technical contributions:** Could the authors more clearly distinguish their technical contributions from existing methods, particularly regarding the adaptation of successor distance and contrastive learning to segment distances?\n2. **Scope of human experiments:** Would the results hold if STAIR were compared to other state-of-the-art PbRL methods beyond PEBBLE? Could the authors provide statistical significance or effect size measures for the human preference results?\n3. **Code and reproducibility:** Could the authors update their code to use current, maintained libraries to ensure long-term reproducibility and adoption?", "limitations": "The authors briefly mention that the quadrilateral distance only assesses pairwise segment differences, limiting its applicability to other preference formats. The discussion of limitations is quite brief and could be expanded.", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 2, "significance": 3, "originality": 4, "strengths_and_weaknesses": "STRENGTHS:\n\n* **Quantification of stage misalignment:** The paper provides clear empirical and theoretical evidence for the negative impact of stage misalignment in PbRL.\n* **Human preference analysis:** The inclusion of human experiments demonstrates that humans prefer segments from later stages, validating the existence of stage reward bias and supporting the motivation for stage-aligned methods.\n* **Novel algorithm:** STAIR adapts successor distance and contrastive learning to segment distance, and introduces a new quadrilateral distance for temporal alignment. This combination demonstrates improved performance over existing methods.\n* **Performance and generalisation:** STAIR outperforms previous models in multi-stage tasks and shows faster convergence even in single-stage settings, indicating robust generalisation.\n* **Human-aligned query selection:** Human studies show that STAIR selects stage-aligned queries more effectively than PEBBLE, suggesting that its stage approximation aligns well with human cognition.\n\nWEAKNESSES:\n\n* **Novelty concerns:** STAIR's adaptation of successor distance and contrastive learning to segment distance raises questions about the novelty of its technical contributions, as these techniques have been introduced in prior work.\n* **Limited human experiment comparisons:** Human experiments only compare STAIR to PEBBLE and do not include other state-of-the-art PbRL methods. The statistical significance of the results is not explicitly reported or quantified.\n* **Code dependency:** The code relies on deprecated packages (such as old versions of Mujoco), which may limit reproducibility and long-term adoption.", "was_revised": true, "final_justification": "I thank the authors for their response which addressed my questions, specifically regarding the issue with experimental comparison with more state-of-the-art PbRL methods. After reviewing all the responses, I stand by my positive evaluation of the work overall.", "rebuttal": "Dear Reviewer,\n\nThanks for your valuable and detailed comments. We hope the following statement clears your concern.\n\n\n**W1, Q1: Novelty concerns.**\n\n**A for W1, Q1:**\nWhile the successor distance is an established concept, our work introduces significant novelty in the problem setting, motivation, and the proposed method, distinguishing it from prior approaches.\n- **Problem Setting**: Our core contribution lies in adapting the successor distance to estimate stage differences and solve the critical problem of stage misalignment in PbRL, a problem not previously addressed. \n- **Motivation**: Our use of successor distance is well-motivated by the challenges outlined in Appendix C, particularly the need for a dynamic, policy-adaptive stage approximation that avoids task-specific assumptions. \n- **Method**: While the successor distance is defined between individual states, it cannot be directly applied to query selection, which requires a measure between segment pairs. To address it, we tailored the successor distance to our scenario. Specifically, we propose the quadrilateral distance, which extends successor distance to capture temporal and behavioral alignment between segments. This ensures stable and efficient stage-aligned learning in our framework, going beyond the original formulation of successor distance.\n\n\n**W2, Q2: Human experiments lack SOTA baselines and statistical significance analysis.**\n\n**A for W2, Q2:**\nTo address your concern, we conduct additional experiments with more state-of-the-art PbRL methods (MRN and RUNE) as baselines, and report the statistical significance analysis result.\nSince evaluating these new baselines requires comparisons with STAIR, we updated the original results to ensure consistent human preference assessments across all methods.\n\nAs suggested, we perform statistical significance analyses on the results, including confidence interval calculations and independent two-sample t-tests, to confirm the significance of differences between STAIR and baselines.\n- Confidence Intervals.\n  We report the 95\\% confidence interval in Table 1. The narrow intervals confirm that STAIR’s performance in learning human-recognized stages is consistent and tightly clustered around the mean.\n- Two-Sample t-Tests.\n  We conduct independent two-sample t-tests for each task, showing that STAIR achieves statistically significant improvements (p<0.05) in learning human-recognized stages.\n\n|Task|Method|Success Rate|95\\% CI|t-value|p-value|Significance|\n|:-:|:-:|:-:|:-:|:-:|:-:|:-:|\n|door-open|STAIR|0.78±0.11|0.78±0.09|-|-|-|\n||PEBBLE|0.51±0.16|0.51±0.14|3.93|<0.01|✔|\n||MRN|0.5±0.07|0.5±0.06|6.08|<0.01|✔|\n||RUNE|0.55±0.13|0.55±0.11|3.8|<0.01|✔|\n|window-close|STAIR|0.73±0.14|0.73±0.12|-|-|-|\n||PEBBLE|0.55±0.2|0.55±0.17|1.97|0.03|✔|\n||MRN|0.45±0.09|0.45±0.08|4.54|<0.01|✔|\n||RUNE|0.38±0.17|0.38±0.14|4.34|<0.01|✔|\n|window-open|STAIR|0.72±0.19|0.72±0.16|-|-|-|\n||PEBBLE|0.40±0.14|0.40±0.12|3.72|<0.01|✔|\n||MRN|0.53±0.1|0.53±0.09|2.36|0.02|✔|\n||RUNE|0.48±0.16|0.48±0.14|2.6|0.01|✔|\n\n\n**W3, Q3: Code dependency.**\n\n**A for W3, Q3:**\nWe appreciate the reviewer's concern regarding the reproducibility and long-term usability of our code.\nTo ensure fair comparisons, our current environment setup aligns with state-of-the-art methods [1-3], and we have provided detailed environment configurations in the code repository to support reproducibility.\nWe are committed to improving the usability of our code and will consider updating it to use current, maintained libraries in future revisions or after acceptance to better support long-term adoption.\n\n\n**L1: Detailed discussion of limitations.**\n\n**A for L1:**\nThe quadrilateral distance is currently defined to evaluate pairwise segment differences, which are then used to compute scores for segment pairs.\nTherefore, more general feedback formats, such as listwise or scalar feedback, which contain multiple segments in a single query, are not directly supported. \nHowever, STAIR can be extended to handle these feedback types by converting them into pairwise preferences:\n- Listwise or Ordinal Feedback: When humans rank a set of samples, pairs can be sampled from the ranked list to construct pairwise feedback.\n- Scalar Scores: When humans provide scores for samples, pairs can be generated by comparing the scores to create pairwise preferences.\n\nExtending STAIR to explicitly handle more preference formats is a promising direction for future research. We will consider that in future work. \n\n**References:**\n\n[1] Lee et al. Pebble: Feedback-efficient interactive reinforcement learning via relabeling experience and unsupervised pre-training.\n\n[2] Cheng et al. RIME: Robust Preference-based Reinforcement Learning with Noisy Preferences.\n\n[3] Liang et al. Reward uncertainty for exploration in preference-based reinforcement learning."}]

Author rebuttal:

Meta-review: This work identifies a fundamental and underexplored issue in PbRL, and theoretically demonstrates the gap in multi-stage tasks. It suggests a natural but nontrivial resolution and theoretically demonstrates its benefit as well as 
well as perform empirical evaluation which spans both multi-stage (MetaWorld, RoboDesk) and single-stage (DMControl) tasks, demonstrating robustness, generalizability, and superior performance. 
Finally, the paper is very well written.
From these reasons the paper should certainly be accepted.
Decision: Accept (poster)