Research map / Reinforcement learning / Imitation learning
Subtask-Aware Visual Reward Learning from Segmented Demonstrations
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes REDS, a visual reward learning framework that trains a subtask-conditioned dense reward model from subtask-segmented action-free demonstrations using EPIC loss, contrastive subtask alignment, and iterative fine-tuning, evaluated on Meta-World, FurnitureBench, and RLBench.
Problem
Reward engineering for long-horizon, multi-subtask robotic manipulation is laborious, and existing video-based reward learners give context-unaware signals and need explicit subtask information or suffer mode collapse.
Approach
Treat subtask segmentation of demonstrations as a ground-truth reward and train a causal-transformer reward model conditioned on subtask text embeddings by minimizing the EPIC distance to it, plus a contrastive loss so the subtask can be inferred from video alone at online time.
Key results
- On Meta-world, REDS consistently improves DreamerV3 sample efficiency over baselines across eight tasks and even surpasses human-engineered rewards on some tasks (e.g., Drawer Open, Push, Coffee Pull).
- On FurnitureBench One Leg, REDS achieves a score of 2.45 after online fine-tuning, versus 1.8 reported by FurnitureBench for IQL with 500 expert demonstrations, using 300 demonstrations.
- REDS trained on three Meta-world tasks transfers to unseen tasks (Door Close, Window Close) without fine-tuning the reward model, achieving comparable or better RL performance than task-specific training.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use REDS instead of IQL trained with 500 expert demonstrations | FurnitureBench One Leg · FurnitureBench One Leg | completed subtasks: 2.45 compared to 1.8 | Helped |
| Replace EPIC loss with regression to subtask segmentation (ablation) | REDS · Meta-world Door Open and Drawer Open | success rate | Hurt |
| Remove subtask embeddings and use only video representations (ablation) | REDS · Meta-world Door Open and Drawer Open | success rate | Hurt |
| Remove the progressive reward regularization loss (ablation) | REDS · Meta-world Door Open and Drawer Open | success rate | Hurt |
| Replace pre-trained visual representations with a CNN image encoder (ablation) | REDS · Meta-world Door Open and Drawer Open | success rate | Hurt |
| Remove the causal transformer and concatenate pre-trained visual representations (ablation) | REDS · Meta-world Door Open and Drawer Open | success rate | Hurt |
| Fine-tune with additional suboptimal demonstrations (ablation) | REDS · Meta-world Door Open and Drawer Open | success rate | Helped |
| Increase the number of expert demonstrations used for reward learning (ablation) | REDS · Meta-world Door Open and Drawer Open | success rate | Helped |
| Set progressive-reward regularization strength epsilon to 0.5 (ablation) | REDS · Meta-world | success rate | Helped |
| Use a lower-percentile threshold for subtask inference (ablation) | REDS · Meta-world | success rate | Hurt |
| Fine-tune the reward model with additional suboptimal demonstrations (ablation) | REDS · Suboptimal demonstrations | subtask identification precision | Helped |
| Use REDS rather than baselines that do not explicitly incorporate subtask information | RL agents · Meta-world | success rate | Helped |
Surprising findings
- Reward model trained on Panda arm demonstrations transfers to an unseen Sawyer arm demonstration in RLBench.
- Using a lower percentile threshold for subtask inference degrades RL performance, since more observations are misclassified as successful.
Limitations
- Assumes object-centric subtask definitions are known and manually provided, which limits automation to new tasks and domains.
- Struggles to distinguish subtle changes such as screwing the leg in One Leg, and relies on natural-image pretrained representations.
- May produce inappropriate rewards for semantically different subtasks with similar video content and text instructions; out-of-distribution background or camera changes are untested.
Experimental setup
Datasets: Meta-world demonstrations (scripted) FurnitureBench One Leg demonstrations RLBench Take Umbrella Out of Umbrella Stand demonstrations (Panda and Sawyer)
Benchmarks: Meta-world FurnitureBench RLBench
Metrics: success rate average number of completed subtasks EPIC distance
Baselines: Human-engineered reward ORIL Rank2Reward VIPER DrS CLIP LIV Diffusion Reward
Builds on: EPIC VIPER DrS DreamerV3 CLIP
Closest papers by meaning
- Multi-Stage Manipulation with Demonstration-Augmented Reward, Policy, and World Model Learning (ICML 2025 · similarity 0.88)
- ELEMENTAL: Interactive Learning from Demonstrations and Vision-Language Models for Reward Design in Robotics (ICML 2025 · similarity 0.88)
- SOLE-R1: Video-Language Reasoning as the Sole Reward for On-Robot Reinforcement Learning (NeurIPS 2026 · similarity 0.88)
- VICtoR: Learning Hierarchical Vision-Instruction Correlation Rewards for Long-horizon Manipulation (ICLR 2025 · similarity 0.87)
- Learning to Act from Actionless Videos through Dense Correspondences (ICLR 2024 · similarity 0.87)
- RoboAlign-R1: Distilled Multimodal Reward Alignment for Robot Video World Models (NeurIPS 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.