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

Dear Reviewer,

Thanks for your valuable and detailed comments. We hope the following statement clears your concern.

W1, Q1: Novelty concerns.

A for W1, Q1: While 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.

• 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.
• 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.
• 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.

W2, Q2: Human experiments lack SOTA baselines and statistical significance analysis.

A for W2, Q2: To 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. Since evaluating these new baselines requires comparisons with STAIR, we updated the original results to ensure consistent human preference assessments across all methods.

As 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.

• Confidence Intervals. 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.
• Two-Sample t-Tests. 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.

W3, Q3: Code dependency.

A for W3, Q3: We appreciate the reviewer's concern regarding the reproducibility and long-term usability of our code. To 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. We 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.

L1: Detailed discussion of limitations.

A for L1: The quadrilateral distance is currently defined to evaluate pairwise segment differences, which are then used to compute scores for segment pairs. Therefore, more general feedback formats, such as listwise or scalar feedback, which contain multiple segments in a single query, are not directly supported. However, STAIR can be extended to handle these feedback types by converting them into pairwise preferences:

• Listwise or Ordinal Feedback: When humans rank a set of samples, pairs can be sampled from the ranked list to construct pairwise feedback.
• Scalar Scores: When humans provide scores for samples, pairs can be generated by comparing the scores to create pairwise preferences.

Extending STAIR to explicitly handle more preference formats is a promising direction for future research. We will consider that in future work.

References:

[1] Lee et al. Pebble: Feedback-efficient interactive reinforcement learning via relabeling experience and unsupervised pre-training.

[2] Cheng et al. RIME: Robust Preference-based Reinforcement Learning with Noisy Preferences.

[3] Liang et al. Reward uncertainty for exploration in preference-based reinforcement learning.

We would like to thank the reviewer for standing by a positive evaluation. We also appreciate the valuable comments, which helped us significantly improve the paper's strengths.

Dear Reviewer,

Thanks for your valuable and detailed comments. We hope the following statement clears your concern.

W1: Explanation for performance in single-stage tasks.

A for W1: 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 a fall (which has a small temporal distance), helping the agent learn stability. As training progresses and the policy improves (with the quadruped becoming more stable), the temporal distance between such segments (before and after a 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-learned behaviors like maintaining stability.

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 to focus on 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 [1-2].

We will explore a theoretical explanation of the connection between curriculum learning and STAIR in future work.

W2, L1: Extension to general feedback formats.

A for W2, L1: While 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:

• 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.
• 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.

Moreover, though human feedback in real-world scenarios can often be more expressive, obtaining such detailed feedback is typically challenging [3]. In contrast, binary preference feedback remains a common approach in practice [4-6]. Therefore, we prioritize binary preference and leave non-binary preference for future work.

W3: Assumption on stage reward bias.

A for W3: While 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. Instead, Proposition 3 highlights a specific scenario where our algorithm exhibits greater performance advantages. Moreover, for more general cases (without the assumption of stage reward bias), Proposition 2 shows that our method remains effective.

W4, L3: Deployment cost due to update frequency ($K_{SD}$).

A for W4, L3: We acknowledge that the performance of STAIR requires frequent contrastive learning updates (small $K_{SD}$), which introduces additional computational cost. However, despite this frequent update, the overall complexity of STAIR does not increase significantly. The 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:

• 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.
• 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.

Q1: Sensitivity and extension to environment reversibility.

A for Q1: The 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:

• Irreversible tasks: The temporal distance learning algorithm proposed in [7] does not assume environment reversibility. Therefore, STAIR can be extended to irreversible tasks.
• 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.

Q2: Performance on noisy/inconsistent feedback.

A for Q2: As 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":

1. Error teacher: A teacher with a random error rate $\epsilon=0.1$, resulting in 10% incorrect feedback.
2. 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$.

Results 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.

L2: Temporal embedding lacks supervision.

A for L2: We acknowledge that STAIR depends on the quality of contrastively learned temporal embeddings. However, 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. Therefore, contrastive learning enables effective learning of this relative ordering without requiring ground truth supervision, consistent with prior works adopting similar techniques [9].

We sincerely thank the reviewer again for the timely and valuable comments. We hope that our response has cleared most of your concerns.

References:

[1] Florensa, Carlos, et al. Automatic Goal Generation for Reinforcement Learning Agents.

[2] Racaniere, Sebastien, et al. Automated curriculum generation through setter-solver interactions.

[3] Choi et al. Listwise Reward Estimation for Offline Preference-based Reinforcement Learning.

[4] Lee et al. Pebble: Feedback-efficient interactive reinforcement learning via relabeling experience and unsupervised pre-training.

[5] Cheng et al. RIME: Robust Preference-based Reinforcement Learning with Noisy Preferences.

[6] Liang et al. Reward uncertainty for exploration in preference-based reinforcement learning.

[7] Myers, Vivek, et al. Learning temporal distances: Contrastive successor features can provide a metric structure for decision-making.

[8] Lee, Kimin, et al. B-pref: Benchmarking preference-based reinforcement learning.

[9] Jiang, Yuhua, et al. Episodic novelty through temporal distance.

Dear reviewer,

We were wondering if our response and revision have cleared all your concerns. In the previous responses, we have tried to address all the points you have raised. In the remaining days of the rebuttal period, we would appreciate it if you could kindly let us know whether you have any other questions, so that we can still have time to respond and address them. We are looking forward to discussions that can further improve our current manuscript. Thanks!

Best regards,

The Authors

We would like to thank the reviewer for standing by a positive score. We also appreciate the valuable comments, which helped us significantly improve the paper's strengths.

Dear Reviewer,

Thanks for your valuable and detailed comments.

W4: Performance on noisy feedback.

A for W4: As 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":

1. Error teacher: A teacher with a random error rate ε=0.1, resulting in 10% incorrect feedback.
2. 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.

Results 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.

W5: Ambiguity of stage-aligned queries.

A for W5: We thank the reviewer for highlighting this important aspect. The 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.

W2, Q1, Q2, Q3: Clarification on theory

• Q1: Eq.3: We would like to clarify that Eq.3 does not state $G^\tau(i)\le G^\tau(i-1)$. Instead, 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.

• Q2: Proposition 1: Yes, 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.

• Q3: Appendix D.1 and Algorithm 3:

  • (1) We briefly introduce the outline of Appendix D.1 and Algorithm 3, respectively.

    • 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.
    • 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.

    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.

  • (2) Yes, both $r_\textrm{stage}$ and $\bar{r}_\textrm{bias}$ are environment parameters. They are sampled from predefined distributions.

Q4, Q8: Clarity of the statement.

• Q4: $T$ in $i=(0,...,100,T)$ denotes terminate, and $w_T$ is a terminated state in the abstract MDP.
• Q8: We conduct 5 experiments (with different seeds) for each result.

Q5: The effect of changing $|\Omega|$ and $|\Upsilon|$.

A for Q5: As 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. We report the normalized episode reward for each method. Results show that stage-aligned PbRL consistently outperforms conventional PbRL. Furthermore, 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|$.

Q6: Does learning with $r_\textrm{bias}=0$ decay into plain PbRL.

A for Q6: Learning 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$.

Q7: How is $d_\textrm{state}$ computed:

A for Q7: $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. Specifically, $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.
It 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].

W1, Q9: Human experiments.

A for W1, Q9: To address your concern, we conducted additional experiments.

(1) Recruitment of participants and mitigation of bias.

As suggested, we invite 8 colleagues from our department to participate in the human experiment. To avoid potential bias, queries generated by STAIR and baselines were randomly shuffled, ensuring labelers were unaware of the source algorithm. Please note that the original result in section 5.3 is also derived under this mechanism, which ensures fairness, as detailed in Appendix E.1.

(2) Statistical significance analysis.

As 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.

• Confidence Intervals. Due to the limited number of samples, the results in Section 5.3 initially had large error bars. With 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.
• Two-Sample t-Tests. 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.

Q10: Ablation on Eq.8.

A for Q10: As 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.

Q11: Query efficiency experiments.

A for Q11: As 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. However, with $N_{total}=500$, STAIR begins to show strong performance, outperforming all baselines, which shows its effectiveness.

Q12-14, PFC: Minor questions, Paper Formatting Concerns: Thank you for your keen attention to detail! We have corrected them in the revised version.

We sincerely thank the reviewer again for the timely and valuable comments. We hope that our response has cleared most of your concerns.

References:

[1] Kurby et al. Segmentation in the perception and memory of events.

[2] Zacks et al. Event perception: a mind-brain perspective.

[3] Lee, Kimin et al. Pebble: Feedback-efficient interactive reinforcement learning via relabeling experience and unsupervised pre-training.

[4] Liang, Xinran, et al. Reward uncertainty for exploration in preference-based reinforcement learning.

[5] Lee, Kimin, et al. B-pref: Benchmarking preference-based reinforcement learning.

Dear Reviewer,

We greatly appreciate your willingness to improve your score. We also appreciate the valuable comments, which helped us significantly improve the paper's strengths. We hope the following statement clears your remaining concern.

Q3: Interpretation for Eq.3

A for Q3:

Thank you for pointing out this issue. The inequality $G^\tau(i) - 1 \le G^\tau(i)$ is always true (as $-1 \leq 0$), and it indeed holds no specific meaning in this context. Our original intention was to represent two key constraints in a more compact form: $G^\tau(i) - 1 \le G^\tau(i-1)$ (ensuring sequential stages) and $G^\tau(i-1) \le G^\tau(i)$ (ensuring no going back). To avoid potential confusion, we will present these two inequalities separately in the revised version.

Q7: Clarity on the calculation of $d_\text{state}$

A for Q7:

Thank you for the helpful feedback. We will revise the manuscript to include additional details and references as suggested.

Q9: Concerns about data independence and multiple comparisons in human experiments.

A for Q9:

We sincerely thank the reviewer for pointing out this important aspect.

1. You are correct that we conducted pairwise t-tests between STAIR and each baseline individually (i.e., STAIR vs. PEBBLE, STAIR vs. MRN, and STAIR vs. RUNE). We will explicitly indicate it in the revised version.

2. Statistical Analysis: We acknowledge that the p-values should indeed be adjusted for multiple comparisons, and we appreciate the reviewer for highlighting this oversight.

   • To address this concern, we applied the suggested Bonferroni correction to the p-values, and the corrected values are provided in the updated table below. After the Bonferroni correction, STAIR remains statistically significant in 7 out of 9 experiments.
   • To further demonstrate the effectiveness of our method, we also calculated the p-values using the Benjamini-Hochberg procedure. As shown in the table, under the Benjamini-Hochberg correction, STAIR achieves statistical significance across all 9 experiments.

These updated results reconfirm that STAIR achieves significant improvements over the baselines in learning human-recognized stages. We hope this clarification resolves your concerns.

Notes for the table:

• p-value: The original p-value from the statistical test.
• Bonf-p-value: The p-value adjusted using the Bonferroni correction.
• Bonf-Significance: The statistical significance determined based on the Bonferroni-adjusted p-value (Bonf-p-value).
• BH-p-value: The p-value adjusted using the Benjamini-Hochberg correction.
• BH-Significance: The statistical significance determined based on the Benjamini-Hochberg-adjusted p-value (BH-p-value).

We sincerely thank the reviewer again for the thoughtful and constructive feedback. We hope that our responses and additional experimental results have addressed your concerns.

Best,

The Authors

Thank you once again for your detailed response, as well as your patience while I considered my review.

In particular, I appreciate the new statistical analysis of the alignment-to-human-stages experiments, it is interesting that not only the results are statistically significant, but also the higher success rate over the baselines.

In light of our discussion, I have decided to increase my score to a 4.

Dear Reviewer,

Thanks for your valuable and detailed comments. We hope the following statement clears your concern.

W1, W3, Q1: Robustness on noisy or inconsistent feedback.

A for W1, W3, Q1: As 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":

1. Error teacher: A teacher with a random error rate $\epsilon=0.1$, resulting in 10% incorrect feedback.
2. 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$.

Results 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.

W2: Complexity concerns.

A for W2: To address your concern, we analyze the complexity of contrastive learning and the complexity of the extra hyperparameters, respectively.

(1) As for the complexity, STAIR introduces additional components, but the overall complexity does not increase significantly. We evaluate this in terms of both training and inference:

• 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.
• 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.

(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.

Q2: Applicability in non-linear stage structures.

A for Q2: We believe STAIR can be adapted to handle branching or non-linear stage structures. STAIR's core mechanism relies on using temporal distance to measure stage differences. The temporal distance used in our work [2], defined as

which quantifies the transition probabilities between states under a given policy and does not depend on linear or sequential task structures. This 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. Extending 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.

Q3: Extension to general feedback formats.

A for Q3: STAIR can generalize to broader feedback formats by extracting pairwise feedback from them:

• Ranked Lists: Given a set of samples ranked by human preference, we can sample pairwise comparisons from the ranked list to construct pairwise feedback.
• 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.

We thank the reviewer for this insightful suggestion, which we will consider in our future work.

Thanks 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.

References:

[1] Lee, Kimin, et al. B-pref: Benchmarking preference-based reinforcement learning.

[2] Myers, Vivek, et al. Learning temporal distances: Contrastive successor features can provide a metric structure for decision-making.

Dear reviewer,

We were wondering if our response and revision have cleared all your concerns. In the previous responses, we have tried to address all the points you have raised. In the remaining days of the rebuttal period, we would appreciate it if you could kindly let us know whether you have any other questions, so that we can still have time to respond and address them. We are looking forward to discussions that can further improve our current manuscript. Thanks!

Best regards,

The Authors