# DrS: Learning Reusable Dense Rewards for Multi-Stage Tasks

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 48 of its 61 references that are in the map.
- First posted on arXiv 25 Apr 2024.

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

- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Scalability in Perception for Autonomous Driving: Waymo Open Dataset](https://mlsanity.com/paper/1912.04838/) (2019)
- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [Do As I Can, Not As I Say: Grounding Language in Robotic Affordances](https://mlsanity.com/paper/2204.01691/) (2022)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [Code as Policies: Language Model Programs for Embodied Control](https://mlsanity.com/paper/2209.07753/) (2022)
- [Curiosity-driven Exploration by Self-supervised Prediction](https://mlsanity.com/paper/1705.05363/) (2017)
- [Exploration by Random Network Distillation](https://mlsanity.com/paper/1810.12894/) (2018)
- [QT-Opt: Scalable Deep Reinforcement Learning for Vision-Based Robotic Manipulation](https://mlsanity.com/paper/1806.10293/) (2018)
- [ProgPrompt: Generating Situated Robot Task Plans using Large Language Models](https://mlsanity.com/paper/2209.11302/) (2022)
- [robosuite: A Modular Simulation Framework and Benchmark for Robot Learning](https://mlsanity.com/paper/2009.12293/) (2020)
- [Learning agile and dynamic motor skills for legged robots](https://mlsanity.com/paper/1901.08652/) (2019)
- [Learning Robust Rewards with Adversarial Inverse Reinforcement Learning](https://mlsanity.com/paper/1710.11248/) (2017)
- [Data-Efficient Hierarchical Reinforcement Learning](https://mlsanity.com/paper/1805.08296/) (2018)
- [ManiSkill2: A Unified Benchmark for Generalizable Manipulation Skills](https://mlsanity.com/paper/2302.04659/) (2023)
- [Reward learning from human preferences and demonstrations in Atari](https://mlsanity.com/paper/1811.06521/) (2018)
- [Language to Rewards for Robotic Skill Synthesis](https://mlsanity.com/paper/2306.08647/) (2023)
- [RoboNet: Large-Scale Multi-Robot Learning](https://mlsanity.com/paper/1910.11215/) (2019)
- [Extrapolating Beyond Suboptimal Demonstrations via Inverse Reinforcement Learning from Observations](https://mlsanity.com/paper/1904.06387/) (2019)
- [MT-Opt: Continuous Multi-Task Robotic Reinforcement Learning at Scale](https://mlsanity.com/paper/2104.08212/) (2021)
- [Go-Explore: a New Approach for Hard-Exploration Problems](https://mlsanity.com/paper/1901.10995/) (2019)
- [Discriminator-Actor-Critic: Addressing Sample Inefficiency and Reward Bias in Adversarial Imitation Learning](https://mlsanity.com/paper/1809.02925/) (2018)
- [ManiSkill: Generalizable Manipulation Skill Benchmark with Large-Scale Demonstrations](https://mlsanity.com/paper/2107.14483/) (2021)
- [A Divergence Minimization Perspective on Imitation Learning Methods](https://mlsanity.com/paper/1911.02256/) (2019)
- [Scaling Up and Distilling Down: Language-Guided Robot Skill Acquisition](https://mlsanity.com/paper/2307.14535/) (2023)
- [Vision-Language Models as Success Detectors](https://mlsanity.com/paper/2303.07280/) (2023)
- [End-to-End Robotic Reinforcement Learning without Reward Engineering](https://mlsanity.com/paper/1904.07854/) (2019)

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

- [SRPO: Self-Referential Policy Optimization for Vision-Language-Action Models](https://mlsanity.com/paper/2511.15605/) (2025)
- [Multi-Stage Manipulation with Demonstration-Augmented Reward, Policy, and World Model Learning](https://mlsanity.com/paper/2503.01837/) (2025)
- [Subtask-Aware Visual Reward Learning from Segmented Demonstrations](https://mlsanity.com/paper/2502.20630/) (2025)
- [ELMUR: External Layer Memory with Update/Rewrite for Long-Horizon RL Problems](https://mlsanity.com/paper/2510.07151/) (2025)
- [Attention-Based Reward Shaping for Sparse and Delayed Rewards](https://mlsanity.com/paper/2505.10802/) (2025)
- [RDD: Retrieval-Based Demonstration Decomposer for Planner Alignment in Long-Horizon Tasks](https://mlsanity.com/paper/2510.14968/) (2025)
- [SARM2: Multi-Task Stage Aware Reward Modeling for Self Improving Robotic Manipulation](https://mlsanity.com/paper/2606.10305/) (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)
- [Demonstration-Free Success-Probability Reward Learning for Generalist Robot Policies](https://mlsanity.com/paper/2609.33653/) (2026)

## Abstract (arXiv)

The success of many RL techniques heavily relies on human-engineered dense rewards, which typically demand substantial domain expertise and extensive trial and error. In our work, we propose DrS (Dense reward learning from Stages), a novel approach for learning reusable dense rewards for multi-stage tasks in a data-driven manner. By leveraging the stage structures of the task, DrS learns a high-quality dense reward from sparse rewards and demonstrations if given. The learned rewards can be textitreused in unseen tasks, thus reducing the human effort for reward engineering. Extensive experiments on three physical robot manipulation task families with 1000+ task variants demonstrate that our learned rewards can be reused in unseen tasks, resulting in improved performance and sample efficiency of RL algorithms. The learned rewards even achieve comparable performance to human-engineered rewards on some tasks. See our project page (https://sites.google.com/view/iclr24drs) for more details.