Research map / Reinforcement learning / Visual and representation RL
Multi-Stage Manipulation with Demonstration-Augmented Reward, Policy, and World Model Learning
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 a demonstration-augmented model-based RL method (built on TD-MPC2 and MoDem) that learns stage-wise dense rewards online via stage discriminators, alongside policy and world model, for visual multi-stage manipulation with sparse rewards.
Problem
Long-horizon multi-stage robotic manipulation from visual and proprioceptive inputs with only sparse stage indicators is hard for RL due to reward design difficulty and exploration; demonstrations are limited.
Approach
Use a few demonstrations to learn stage discriminators in the world model's latent space that densify sparse stage rewards online (tanh-bounded within stage intervals), combined with a bi-phase scheme (BC pretraining then planning with demo oversampling) so the world model, policy and reward are learned jointly from demonstrations.
Key results
- Averages a 40% data-efficiency improvement over baselines across 16 sparse-reward visual tasks, and about 70% on the hardest tasks.
- On ManiSkill3, reaches an average 75% success rate after 500k steps, about 50% better than the second-best baseline.
- Ablations show reward learning matters most in long-horizon tasks such as Peg Insertion and Pick Place, and the method is the only one reaching meaningful performance on Stack Cube and Peg Insertion with only 5 demonstrations.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Ours instead of the proposed baselines | multi-stage manipulation tasks across four domains · 16 tasks across four domains | data-efficiency: 40\% | Helped |
| Use Ours instead of the second-best baseline | ManiSkill3 · ManiSkill3 | success rate: 50\% | Helped |
| Use Ours instead of current methods | the most difficult tasks · most difficult tasks | steps to convergence: 4$x | Helped |
| Add policy pretraining and demonstration oversampling (ablation) | TD-MPC2 backbone without learned reward · ManiSkill3 manipulation tasks | success rate | Helped |
| Add learned reward (ablation) | Ours · long-horizon tasks, including Peg Insertion and Pick Place | success rate | Helped |
| Use 2 stages instead of 1 stage (ablation) | Ours · Stack Cube and Peg Insertion | success rate | Helped |
| Use dense rewards instead of stage-based rewards (ablation) | Ours · Stack Cube and Peg Insertion | task completion speed | Helped |
| Use only 5 demonstrations (ablation) | Ours · Stack Cube and Peg Insertion | success rate | Helped |
| Add online reward learning to the demonstration-augmented method | Ours · Robosuite | hours per 100k interaction steps: 4.84 -> 5.19 | Hurt |
Surprising findings
- Two-stage reward granularity achieves near-equal performance to finer stage divisions, suggesting stage indicators plus few demonstrations suffice.
- Pretraining and demonstration oversampling alone (no learned reward) gave a large jump in performance over TD-MPC2.
Limitations
- Experiments are limited to simulation; no real-robot evaluation.
- Only one demonstration source (TD-MPC2-generated state-based demonstrations) was studied; human teleoperation was not explored.
- Requires task-specific stage definitions and demonstration counts tuned per task to ensure baselines solve tasks.
Experimental setup
Datasets: ManiSkill3 demonstrations Meta-World demonstrations Robosuite demonstrations ManiSkill Humanoids demonstrations
Benchmarks: ManiSkill3 Meta-World Robosuite ManiSkill Humanoids
Metrics: success rate data-efficiency (interaction steps) wall-time (hours per 100k steps)
Baselines: MoDem LaNE TD-MPC2
Builds on: TD-MPC2 MoDem DrS LaNE
Closest papers by meaning
- DrS: Learning Reusable Dense Rewards for Multi-Stage Tasks (ICLR 2024 · similarity 0.89)
- GravMAD: Grounded Spatial Value Maps Guided Action Diffusion for Generalized 3D Manipulation (ICLR 2025 · similarity 0.89)
- VICtoR: Learning Hierarchical Vision-Instruction Correlation Rewards for Long-horizon Manipulation (ICLR 2025 · similarity 0.89)
- Text2Reward: Reward Shaping with Language Models for Reinforcement Learning (ICLR 2024 · similarity 0.89)
- Subtask-Aware Visual Reward Learning from Segmented Demonstrations (ICLR 2025 · similarity 0.88)
- Eureka: Human-Level Reward Design via Coding Large Language Models (ICLR 2024 · similarity 0.88)
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.