# Multi-Stage Manipulation with Demonstration-Augmented Reward, Policy, and World Model Learning

arXiv: https://arxiv.org/abs/2503.01837 · HTML: https://mlsanity.com/paper/2503.01837/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Placed in Embodied and robot world models, one of 11 directions in World models.
- Builds on 39 of its 58 references that are in the map.
- First posted on arXiv 3 Mar 2025.

## Placed in

- World models / [Embodied and robot world models](https://mlsanity.com/wm/embodied-and-robot-world-models/)
- Deep reinforcement learning methods / [RL for robot locomotion and manipulation](https://mlsanity.com/deeprl/rl-for-robot-locomotion-and-manipulation/)

## Builds on (most cited first, 30 of 39)

- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Learning Complex Dexterous Manipulation with Deep Reinforcement Learning and Demonstrations](https://mlsanity.com/paper/1709.10087/) (2017)
- [MOReL : Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2005.05951/) (2020)
- [Eureka: Human-Level Reward Design via Coding Large Language Models](https://mlsanity.com/paper/2310.12931/) (2023)
- [robosuite: A Modular Simulation Framework and Benchmark for Robot Learning](https://mlsanity.com/paper/2009.12293/) (2020)
- [Learning Robust Rewards with Adversarial Inverse Reinforcement Learning](https://mlsanity.com/paper/1710.11248/) (2017)
- [Temporal Difference Learning for Model Predictive Control](https://mlsanity.com/paper/2203.04955/) (2022)
- [Efficient Online Reinforcement Learning with Offline Data](https://mlsanity.com/paper/2302.02948/) (2023)
- [Overcoming Exploration in Reinforcement Learning with Demonstrations](https://mlsanity.com/paper/1709.10089/) (2017)
- [Reward Design with Language Models](https://mlsanity.com/paper/2303.00001/) (2023)
- [One-Shot Imitation Learning](https://mlsanity.com/paper/1703.07326/) (2017)
- [MT-Opt: Continuous Multi-Task Robotic Reinforcement Learning at Scale](https://mlsanity.com/paper/2104.08212/) (2021)
- [Vision-Language Models are Zero-Shot Reward Models for Reinforcement Learning](https://mlsanity.com/paper/2310.12921/) (2023)
- [Text2Reward: Reward Shaping with Language Models for Reinforcement Learning](https://mlsanity.com/paper/2309.11489/) (2023)
- [The Unsurprising Effectiveness of Pre-Trained Vision Models for Control](https://mlsanity.com/paper/2203.03580/) (2022)
- [Vision-Language Models as Success Detectors](https://mlsanity.com/paper/2303.07280/) (2023)
- [Model Predictive Path Integral Control using Covariance Variable Importance Sampling](https://mlsanity.com/paper/1509.01149/) (2015)
- [End-to-End Robotic Reinforcement Learning without Reward Engineering](https://mlsanity.com/paper/1904.07854/) (2019)
- [XIRL: Cross-embodiment Inverse Reinforcement Learning](https://mlsanity.com/paper/2106.03911/) (2021)
- [AVID: Learning Multi-Stage Tasks via Pixel-Level Translation of Human Videos](https://mlsanity.com/paper/1912.04443/) (2019)
- [Playing hard exploration games by watching YouTube](https://mlsanity.com/paper/1805.11592/) (2018)
- [Adversarial Motion Priors Make Good Substitutes for Complex Reward Functions](https://mlsanity.com/paper/2203.15103/) (2022)
- [Keeping Your Distance: Solving Sparse Reward Tasks Using Self-Balancing Shaped Rewards](https://mlsanity.com/paper/1911.01417/) (2019)
- [Offline Learning from Demonstrations and Unlabeled Experience](https://mlsanity.com/paper/2011.13885/) (2020)
- [MoDem: Accelerating Visual Model-Based Reinforcement Learning with Demonstrations](https://mlsanity.com/paper/2212.05698/) (2022)
- [Graph Inverse Reinforcement Learning from Diverse Videos](https://mlsanity.com/paper/2207.14299/) (2022)

## Built on by (most cited first, 11 of 11)

- [STARE-VLA: Progressive Stage-Aware Reinforcement for Fine-Tuning Vision-Language-Action Models](https://mlsanity.com/paper/2512.05107/) (2025)
- [Continual Learning for Generative AI: From LLMs to MLLMs and Beyond](https://mlsanity.com/paper/2506.13045/) (2025)
- [From Instruction to Event: Sound-Triggered Mobile Manipulation](https://mlsanity.com/paper/2601.21667/) (2026)
- [SLOPE: Optimistic Potential Landscape Shaping for Model-based Reinforcement Learning](https://mlsanity.com/paper/2602.03201/) (2026)
- [Not All Actions Are Equal: Rethinking Conditioning for Dexterous World Model](https://mlsanity.com/paper/2606.27325/) (2026)
- [Stage-Transition Dense Reward Modeling for Reinforcement Learning](https://mlsanity.com/paper/2606.31377/) (2026)
- [SUN: Agentic Robot Policy Learning with Persistent Task Programs](https://mlsanity.com/paper/2608.31167/) (2026)
- [A GPU-Parallel Framework for Heterogeneous Multi-Task Reinforcement Learning](https://mlsanity.com/paper/2606.03335/) (2026)
- [Grasp-Then-Plan with Failure Attribution: A Closed Two-Stage Framework for Precise and Generalizable Robotic Manipulation](https://mlsanity.com/paper/2606.03385/) (2026)
- [Where Success Breaks: Failure-Boundary Learning for Robust Vision-Language-Action Models](https://mlsanity.com/paper/2609.06114/) (2026)
- [Demonstration-Free Success-Probability Reward Learning for Generalist Robot Policies](https://mlsanity.com/paper/2609.33653/) (2026)

## Abstract (arXiv)

Long-horizon tasks in robotic manipulation present significant challenges in reinforcement learning (RL) due to the difficulty of designing dense reward functions and effectively exploring the expansive state-action space. However, despite a lack of dense rewards, these tasks often have a multi-stage structure, which can be leveraged to decompose the overall objective into manageable subgoals. In this work, we propose DEMO3, a framework that exploits this structure for efficient learning from visual inputs. Specifically, our approach incorporates multi-stage dense reward learning, a bi-phasic training scheme, and world model learning into a carefully designed demonstration-augmented RL framework that strongly mitigates the challenge of exploration in long-horizon tasks. Our evaluations demonstrate that our method improves data-efficiency by an average of 40% and by 70% on particularly difficult tasks compared to state-of-the-art approaches. We validate this across 16 sparse-reward tasks spanning four domains, including challenging humanoid visual control tasks using as few as five demonstrations.