DrS: Learning Reusable Dense Rewards for Multi-Stage Tasks

TL;DR

DrS tackles reward engineering in RL by learning reusable dense rewards from stage structures in multi-stage tasks. It builds per-stage discriminators that enforce progress via a two-phase process: reward learning from training tasks and reward reuse to unseen tasks, using SAC as the backbone. Empirical results on ManiSkill show significant gains in sample efficiency and, in some tasks, parity with human-engineered rewards, while reducing human effort relative to manual reward design. The approach demonstrates robust transfer across object variations and highlights the practical potential of stage-aware reward learning for scalable RL, albeit with limitations in automatically obtaining task stage structures.

Abstract

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 \textit{reused} 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.

Figures (16)

• Figure 1: An illustration of stage indicators in an OpenCabinetDoor task, which can be naturally divided into three stages plus a success state. A stage indicator is a binary function representing whether the current state is in a certain stage, and it can be simply defined by some boolean functions.
• Figure 2: a) GAIL's discriminator aims to distinguish agent trajectories from demonstrations. b) In single-stage tasks, the discriminator in our approach aims to distinguish success trajectories from failure ones. c) In multi-stage tasks, our approach train a separate discriminator for each stage. The discriminator for stage $k$ aims to distinguish trajectories that reach beyond stage $k$ from those that only reach up to stage $k$. d) Overall, our approach has 2 phases: reward learning and reward reuse.
• Figure 3: An illustration of our learned reward, which fills the gaps in semi-sparse rewards, resulting in a smooth reward curve.
• Figure 4: We evaluated our approach DrS on more than 1,000 task variants from three task families in ManiSkill mu2021maniskillgu2023maniskill2. Each task variant is associated with a different object. All tasks require low-level physical control. The objects in training and test tasks are non-overlapped.
• Figure 5: Evaluation results of reusing learned rewards. All curves use SAC to train, but with different rewards. VICE-RAQ and ORIL get no success. 5 random seeds, the shaded region is std.