Subtask-Aware Visual Reward Learning from Segmented Demonstrations

TL;DR

REDS tackles the challenge of reward engineering in reinforcement learning for robotics by learning dense, subtask-aware rewards from action-free videos segmented into subtasks. The approach combines a subtask-conditioned reward model with EPIC-based objective, a progressive reward regularizer, and a contrastive video-text alignment to infer ongoing subtasks online. Empirical results on Meta-World and FurnitureBench show improved sample efficiency, strong generalization to unseen tasks and embodiments, and robust performance under visual distractions, with competitive or superior performance to dense, hand-designed rewards. This work promises scalable, minimal-supervision reward learning for long-horizon robotic manipulation in diverse real-world settings.

Abstract

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.

Figures (17)

• Figure 1: Illustration of REDS. Our main idea is to leverage expert demonstrations annotated with the ongoing subtask as the source of implicit reward signals (left). We train a reward model conditioned on video segments and corresponding subtasks with 1) contrastive loss to attract the video segments and corresponding subtask embeddings and 2) EPIC EPIC loss to generate reward equivalent to subtask segmentations (middle). In online RL, REDS infers the ongoing subtask using only video segments at each timestep and computes the reward with that (right).
• Figure 1: Online fine-tuning results of IQL agents in One Leg from FurnitureBench. We report the initial performance after offline RL (left) and the performance after 150 episodes of online RL (right).
• Figure 2: Examples of visual observations used in our experiments. We consider a variety of robotic manipulation tasks from Meta-world yu2020meta and FurnitureBench furniturebench.
• Figure 2: EPIC EPIC distance (lower is better) between learned reward functions and hand-engineered reward functions (Meta-world) / subtask segmentations (FurnitureBench) in unseen data.
• Figure 3: Learning curves of DreamerV3 hafner2023mastering agents trained with different reward functions for solving eight robotic manipulation tasks from Meta-world yu2020meta, measured by success rate (%). The solid line and shaded regions represent the mean and stratified bootstrap interval across 4 runs.