Subtask-Aware Visual Reward Learning from Segmented Demonstrations
This addresses the problem of scalable reward learning for robotic tasks, enabling generalization to unseen tasks and robot embodiments, though it appears incremental as it builds on existing reward learning methods with segmentation.
The paper tackles the problem of reinforcement learning agents requiring human-engineered reward functions by introducing REDS, a reward learning framework that uses action-free video demonstrations segmented into subtasks as ground-truth rewards. The result shows REDS significantly outperforms baseline methods on complex robotic manipulation tasks in Meta-World and furniture assembly in FurnitureBench, with minimal human intervention.
Reinforcement Learning (RL) agents have demonstrated their potential across various robotic tasks. However, they still heavily rely on human-engineered reward functions, requiring extensive trial-and-error and access to target behavior information, often unavailable in real-world settings. This paper introduces REDS: REward learning from Demonstration with Segmentations, a novel reward learning framework that leverages action-free videos with minimal supervision. Specifically, REDS employs video demonstrations segmented into subtasks from diverse sources and treats these segments as ground-truth rewards. We train a dense reward function conditioned on video segments and their corresponding subtasks to ensure alignment with ground-truth reward signals by minimizing the Equivalent-Policy Invariant Comparison distance. Additionally, we employ contrastive learning objectives to align video representations with subtasks, ensuring precise subtask inference during online interactions. Our experiments show that REDS significantly outperforms baseline methods on complex robotic manipulation tasks in Meta-World and more challenging real-world tasks, such as furniture assembly in FurnitureBench, with minimal human intervention. Moreover, REDS facilitates generalization to unseen tasks and robot embodiments, highlighting its potential for scalable deployment in diverse environments.