Subtask-Aware Visual Reward Learning from Segmented Demonstrations

Dear Reviewer 7C5k,

We sincerely appreciate your valuable comments, which were extremely helpful in improving our draft. Below, we address each comment in detail.

[W1] Reliance on pre-trained visual and text encoders

We appreciate the reviewer’s insightful comment. We hypothesize that pre-trained models leveraging large-scale robotic datasets [1, 2] could improve reward quality by capturing subtle visual changes more effectively. However, to the best of our knowledge, no such open-sourced model is currently available. We consider this a promising avenue for future work and plan to investigate these representations as they become available.

[W2] Effect of segmented demonstrations

Thank you for your comment. While REDS requires demonstrations with subtask segmentations, we designed it to mitigate dependence on high-quality annotations. The iterative training process leverages suboptimal demonstrations with automatic segmentation, reducing the need for manual effort (Section 4.4). Moreover, as shown in Figure 9d, REDS performs robustly even with fewer demonstrations, indicating its adaptability to varying data quality.

Reference
[1] Open X-Embodiment: Robotic Learning Datasets and RT-X Models, ICRA 2024.
[2] DROID: A Large-Scale In-the-Wild Robot Manipulation Dataset, RSS 2024.
[3] Robotic Control via Embodied Chain-of-Thought Reasoning, CoRL 2024.

Dear Reviewer 7C5k,

Thank you again for your time and efforts in reviewing our paper.

As the discussion period draws close, we kindly remind you that two days remain for further comments or questions. We would appreciate the opportunity to address any additional concerns you may have before the discussion phase ends.

Thank you very much!

Many thanks,
Authors

[Q7] Influence of the form of language instructions

We observe that the format of language instructions has minimal impact on the performance of downstream RL agents, as the model primarily relies on the semantic content rather than specific syntactic structure. To ensure compatibility and consistency, we adopt the format used in CLIP, which is widely validated for language-vision tasks.

[Q8] Explanation of Figure 10

Figure 10 illustrates the limitations of VIPER, DrS, and ORIL in providing effective reward signals. ORIL suffers from mode collapse and fails to distinguish visual changes in the robot's state, resulting in improper reward signals. The sudden drop in DrS performance (right figure) is caused by misleading subtask annotations from the pre-defined script; we have corrected this issue and updated the graph in the revised draft. VIPER generates distinct reward scales when transitioning between subtasks, assigning lower rewards in later phases, which causes the agent to become stuck. In contrast, REDS, guided by the EPIC loss, produces consistent and context-aware reward signals, enabling the agent to progress effectively through subsequent subtasks.

Reference
[1] VIP: Towards Universal Visual Reward and Representation via Value-Implicit Pre-Training, ICLR 2023.
[2] XIRL: Cross-embodiment Inverse Reinforcement Learning, CoRL 2021.
[3] Unsupervised Perceptual Rewards for Imitation Learning, RSS 2017.
[4] MimicGen: A Data Generation System for Scalable Robot Learning using Human Demonstrations, CoRL 2023.
[5] Learning multi-stage tasks with one demonstration via self-replay, CoRL 2021.
[6] Visual Language Maps for Robot Navigation, ICRA 2023.
[7] Octopus: Embodied Vision-Language Programmer from Environmental Feedback, ECCV 2024.
[8] Generative Agents: Interactive Simulacra of Human Behavior, UIST 2023.
[9] Mobile-Agent-v2: Mobile Device Operation Assistant with Effective Navigation via Multi-Agent Collaboration, NeurIPS 2024.
[10] Language-Grounded Dynamic Scene Graphs for Interactive Object Search with Mobile Manipulation, RA-L 2024.
[11] BUMBLE: Unifying Reasoning and Acting with Vision-Language Models for Building-wide Mobile Manipulation, ArXiv 2024.
[12] FurnitureBench: Reproducible Real-World Benchmark for Long-Horizon Complex Manipulation, RSS 2023.
[13] Diffusion Reward: Learning Rewards via Conditional Video Diffusion, ECCV 2024.

Dear Reviewer VAFb,

We sincerely appreciate your valuable comments, which were extremely helpful in improving our draft. Below, we address each comment in detail.

[W1-1,1-2] Comparison with existing studies

VIP [1] and XIRL [2] share similarities with REDS in learning visual representations from demonstration videos and generating progressive reward signals. However, they require access to goal images to compute rewards, which is often impractical in real-world settings. In contrast, REDS autonomously generates reward signals without relying on additional environmental information, enabling fully autonomous online RL training.

Sermanet et al. [3] propose an unsupervised approach to derive reward signals by discovering intermediate steps from visual features. This method has a limitation in that it depends entirely on expert demonstrations and pre-trained visual features, making it susceptible to reward misattribution in visually ambiguous scenarios (e.g., confusing leg insertion with alignment in the "One Leg" task). REDS addresses these limitations by leveraging subtask-segmented demonstrations to train additional models on top of pre-trained representations, offering precise, context-aware guidance in long-horizon and complex tasks. Furthermore, REDS effectively incorporates suboptimal demonstrations, mitigating reward misspecification and enhancing robustness.

[W1-3] Using large-scale models to alleviate annotation burdens

Thank you for pointing this out. Following your suggestions, we revised the draft to remove this mention and added an in-depth discussion on limitations and future works (please refer to page 11 of the revised draft).

[W2] Clarification of motivation

The proposed method's motivation stems from the limitations of prior approaches like VIPER, as shown in Figure 11. VIPER assigns lower rewards to later subtasks (beyond point 4), making agents stagnate in earlier phases. In contrast, our method generates context-aware reward signals that maintain consistency across subtasks, enabling effective learning for complex, long-horizon tasks. We have revised the manuscript to elaborate on this motivation and its empirical support.

[W3-1] Clarification on rules for subtask segmentation

To ensure consistent subtask definitions, we assume tasks comprise a sequence of object-centric subtasks, where each subtask involves manipulating a single object, following prior work [4,5]. For instance, Door Open (Figure 2a of the draft) can be divided into (i) reaching the door handle (which involves motion relative to a handle.) and (ii) pulling the door to the goal position (which involves motion relative to a green sphere). For new tasks, subtasks can be similarly defined by identifying changes in the target object, and this approach aligns with human intuition and is repeatable. We have not tried different rules and suspect that excessive task horizons without proper subtask decomposition can degrade RL performance. In the revised manuscript, we have expanded the discussion on subtask definition rules.

[W3-2] Clarification on rules for threshold $T_{U}$

For each subtask $U_i$, we compute similarity scores between the visual observations within the subtask from expert demonstrations and their corresponding instructions. The threshold {T_{U_i} is set to the 75th percentile of these scores to account for demonstration variability while capturing the most relevant matches. We include this clarification in the revised manuscript.

[W3-3] Extension to other domains

Thank you for your valuable feedback. We agree that discussing the adaptability of REDS to other domains is important. REDS can be extended with minimal modifications, as its framework is inherently domain-agnostic. The primary considerations involve defining subtasks and decomposing expert demonstrations. To ensure robustness, the task horizon should be constrained to avoid excessively large state distributions, which could increase susceptibility to reward hacking. For example, in web control tasks, the problem can be decomposed into subtasks aligned with changes in target components (e.g., screens, icons, or scrollbars). Moreover, leveraging the reasoning capabilities of (Multimodal) Large Language Models (MLLMs) offers a scalable approach to subtask definition and demonstration decomposition. This has been demonstrated in domains such as navigation [6, 7], agent control [8, 9], and mobile manipulation [10, 11]. We will further clarify these points in the final manuscript, including a discussion of potential failure scenarios and domain-specific considerations, to strengthen the generalizability and applicability of REDS.

Dear Reviewer Vi9Z,

We sincerely appreciate your valuable comments, which were extremely helpful in improving our draft. Below, we address each comment in detail.

[W1] Reliance on domain-based task decomposition

Thank you for pointing this out. In defining subtasks, we designed a rule that avoids manual, task-specific adjustments and broadly applies to robotic manipulation tasks. To this end, we assumed that tasks can be decomposed into a sequence of object-centric subtasks, where each subtask involves manipulation relative to a single object, following prior work [1, 2]. For instance, in Door Open (Figure 2a of the draft), the subtasks are (i) reaching the door handle (motion relative to the handle) and (ii) pulling the door to the goal position (motion relative to a green sphere). This approach is intuitive because humans naturally perceive tasks as sequences of discrete object interactions, each with a clear and measurable goal. For new tasks, subtasks can be defined by identifying the manipulated object and its goal state, following the same objective criteria. This ensures that the method is consistent, repeatable, and easily generalizable across different tasks.

To further address concerns about scalability and reliance on predefined instructions, we plan to integrate ideas from concurrent works, generating high-level plans with instructions for solving each subtask and specifying current progress using the reasoning capability of Multimodal Large Language models (MLLM) [3] or decomposing long-horizon robotics demonstrations into subtasks using pre-trained vision language models [4, 5], in the future work. We have added an in-depth discussion of these extensions in the revised manuscript.

[W2] Clarification of subtask identification function $\psi$

Thank you for your detailed question. $\psi$ maps observations at each timestep to the index of the ongoing subtask, incrementing its output as subtasks are completed. This results in the form of a step function, where each step corresponds to the completion of a subtask. Please refer to the graph in the center of Figure 1 in the revised draft for visual examples.

[W3] Rationale and comparisons for EPIC loss selection in reward learning

As demonstrated in prior work [6], a low EPIC distance between a learned reward $\hat{R}$ and the true reward $R$ predicts low regret, which implies that policies $\pi_{\hat{R}}$​ and $\pi_R$, optimized for $\hat{R}$ and $R$, respectively, achieve similar returns under the ground-truth reward $R$. Leveraging this insight, we adopt EPIC distance as an optimization objective to train a dense reward function that induces the same set of optimal policies as $\psi$, ensuring robust and consistent performance across tasks.

For empirical comparison, we provide qualitative analyses of different reward functions in Figure 10. Discriminator-based methods are observed to suffer from mode collapse, leading to suboptimal reward signals. Similarly, video prediction-based approaches tend to assign lower rewards to later subtasks, causing agents to stagnate in earlier phases of the task. In contrast, trained with EPIC loss, REDS produces consistent and task-progressive reward signals, allowing agents to transition effectively through subtasks. These results demonstrate our approach's robustness in addressing prior methods' limitations.

[W4] Dealing with subtasks sharing similar text descriptions or video contents

Thank you for your constructive feedback. Section 5.4 demonstrates REDS's generalization capabilities in unseen environments requiring similar subtasks with unseen objects (Window Close), where the text instructions differ only by the target object. Figure 5 shows that REDS generates effective reward signals and achieves comparable or superior RL performance, even under these conditions. While the alignment mechanism currently focuses on minimizing distances between relevant embeddings, we acknowledge that incorporating a loss term to maximize distances between irrelevant embeddings could further improve robustness in tasks with similar subtasks. We will explore this enhancement in future work and have added a discussion of this limitation in the revised manuscript.

Dear Reviewer Vi9Z,

Thank you again for your time and efforts in reviewing our paper.

As the discussion period draws close, we kindly remind you that two days remain for further comments or questions. We would appreciate the opportunity to address any additional concerns you may have before the discussion phase ends.

Thank you very much!

Many thanks,
Authors

[Training Phase]
We begin by decomposing the task into a sequence of object-centric subtasks, where each subtask involves manipulation relative to a single object (as detailed in Section 4.1 of the revised draft). Expert demonstrations are segmented to label ongoing subtasks from the sequence of object-centric subtasks at each timestep. These segmentations are used to train REDS with the objectives described in Section 4, including a contrastive learning objective that aligns video embeddings with their corresponding subtask embeddings. Notably, no additional segmentation is required once training begins.

[Inference Phase]
Although our approach aligns with the overarching goal of [1], REDS implements a distinct mechanism for task decomposition. Instead of training a stage-recognition network, REDS infers subtasks during online interaction by comparing video and text embeddings. Specifically, a video encoder processes the observation history, while task instructions for all subtasks are pre-encoded using a text encoder. The ongoing subtask $i$ is identified by selecting the subtask embedding $\mathbf{e}_{i}$ with the highest cosine similarity to the video embedding. This process is made robust by training video representations to align closely with subtask embeddings via the contrastive learning objective (Equation 5 in the revised draft). This alignment significantly enhances subtask inference and improves RL performance, as demonstrated in Figure 9a.

Please let us know if further clarification would be helpful.

References
[1] Learning multi-stage tasks with one demonstration via self-replay, CoRL 2021.

Thank you for your response. We are happy to hear that we have addressed most of your concerns. Following your suggestion, we will further clarify the advantages of our method in the final manuscript. Again, Thank you for the valuable suggestion and your positive assessment of our work.

Thank you very much!
Authors

Dear Reviewer j9z4,

We sincerely appreciate your valuable comments, which were extremely helpful in improving our draft. Below, we address each comment in detail.

[W1, Q2] Explanation of subtask definition

Thank you for pointing this out. In defining subtasks, we designed a rule that avoids manual, task-specific adjustments and broadly applies to robotic manipulation tasks. To this end, we assumed that tasks can be decomposed into a sequence of object-centric subtasks, where each subtask involves manipulation relative to a single object, following prior work [1, 2]. For instance, in Door Open (Figure 2a of the draft), the subtasks are (i) reaching the door handle (motion relative to the handle) and (ii) pulling the door to the goal position (motion relative to a green sphere). This approach is intuitive because humans naturally perceive tasks as sequences of discrete object interactions, each with a clear and measurable goal. Furthermore, this assumption can be generally applied to different manipulation skills (e.g., pick-and-place, inserting, assembling) with diverse objects.

For new tasks, subtasks can be defined by identifying the manipulated object and its goal state, following the same objective criteria. This ensures that the method is consistent, repeatable, and easily generalizable across different tasks. In the revised manuscript, we have elaborated on the rules for subtask definition.

[Q1] Using imitation-from-observation methods as a baseline

We would like to clarify that several of our baselines, including ORIL, DrS, and R2R, are variants of discriminator-based reward learning methods. These approaches train a discriminator to differentiate between expert demonstrations and agent trajectories, using the discriminator's output as a reward signal based on the trajectory's alignment with expert behavior. REDS consistently and significantly outperforms these baselines in our experiments across all tasks.

Reference
[1] MimicGen: A Data Generation System for Scalable Robot Learning using Human Demonstrations, CoRL 2023.
[2] Learning multi-stage tasks with one demonstration via self-replay, CoRL 2022.

Dear Reviewer j9z4,

Thank you again for your time and efforts in reviewing our paper.

As the discussion period draws close, we kindly remind you that two days remain for further comments or questions. We would appreciate the opportunity to address any additional concerns you may have before the discussion phase ends.

Thank you very much!

Many thanks,
Authors

Thanks for the author's reply. After reading the reply, my concern still holds that the subtask definition heavily relies on human annotations. As a result, I decide to keep the score unchanged.

Dear Reviewer mmyJ,

We sincerely appreciate your valuable comments, which were extremely helpful in improving our draft. Below, we address each comment in detail.

[W1, Q1] Relation between subtask identification and RL performance

We agree that the reward model trained only with expert demonstrations might produce unreliable reward signals with incorrect subtask identification in the exploration phase. To solve this issue, as mentioned in Section 4.4, We fine-tuned our reward model using additional suboptimal demonstrations obtained from RL agents initially trained with expert demonstrations for preventing reward misspecification.

To quantify the effect of fine-tuning, we measure the accuracy of the reward model's subtask identification using unseen datasets before and after fine-tuning with additional suboptimal demonstrations in the revised draft. Table 7 of the revised draft demonstrates that fine-tuning notably enhances precision on suboptimal demonstrations, thereby improving RL performance.

[W2, Q2] Effect of hyperparameter $\epsilon$

For choosing the hyper-parameter $\epsilon$, which enforces progressive reward signals, we conducted a grid search in the range [0.0, 1.0] and selected 0.5 for all experiments. We included new RL performance results across different $\epsilon$ values in Figure 12 of the revised draft. We observe that smaller values failed to provide sufficient progressive signals, while larger values reduced the accuracy of subtask inference, degrading the reward function's overall effectiveness.

[W3] Convergence of baselines on some tasks

Thank you for your pointer! To address this concern, we extended the training for RL agents up to 5 million environmental interactions for the Push and Coffee Pull tasks. We have updated the results (see Figure 1) in the revised draft. Despite the increased training interactions, we observed that the overall performance trends remained consistent with our initial findings.

[W4, Q4] Comparison with prior work [1,2]

Both VLM-RM [1] and LIV [2] share similarities with REDS in terms of producing dense rewards from visual observations and text instructions. However, REDS distinguishes itself by its ability to infer ongoing subtasks and generate dynamic, context-aware reward signals tailored to specific text instructions during online interactions. In contrast, VLM-RM and LIV require static text instructions describing the overall task or depend on external mechanisms for subtask identification. We have included a comparison in Figure 12 of the revised draft, where RL agents trained using REDS significantly outperform those using VLM-RM and LIV, demonstrating the effectiveness of our approach.

References
[1] Vision-Language Models are Zero-Shot Reward Models for Reinforcement Learning, ICLR 2024.
[2] LIV: Language-Image Representations and Rewards for Robotic Control, ICML 2023.

Dear Reviewer mmyJ,

Thank you again for your time and efforts in reviewing our paper.

As the discussion period draws close, we kindly remind you that two days remain for further comments or questions. We would appreciate the opportunity to address any additional concerns you may have before the discussion phase ends.

Thank you very much!

Many thanks,
Authors