DrS: Learning Reusable Dense Rewards for Multi-Stage Tasks
Authors: Tongzhou Mu
Minghua Liu
Hao Su
Venue: ICLR 2024 poster (2024)
Keywords: Reward Learning
Multi-stage Task
Primary area: reinforcement learning
TL;DR: We propose DrS, a novel reward learning approach that learns reusable dense rewards for multi-stage tasks.

Abstract: The success of many RL techniques heavily relies on human-engineered dense rewards, which typically demands substantial domain expertise and extensive trial and error. In our work, we propose **DrS** (**D**ense **r**eward learning from **S**tages), 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 *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.

Published as a conference paper at ICLR 2024
DRS: LEARNING REUSABLE DENSE REWARDS FOR MULTI-STAGE TASKS
Tongzhou Mu, Minghua Liu, Hao Su
UC San Diego (t.3mu, mi m1700,, haosu) Bucsd.eds
ABSTRACT
The success of many RL techniques heavily relies on human-en -ngiineerd dense rewards, which typically demands 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 reused in unseen tasks, thus reducing the human effort for reward engineering. Extensive experiments 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 RI algorithms. The learned rewards even achieve comparable performance to human- -enianee rewards on some tasks, See our project page for more details.
three physical robot manipulation
INTRODUCTION
The success of many reinforcement learning (RL) techniques heavily relies on dense reward functions (Hwangbo et al. 2019; Peng et al.] 2018), which are often tricky to design by humans due to heavy domain expertise requirements and tedious trials and errors. In contrast, sparse rewards, such binary task completion signal, are significantly easier to obtain (often directly from the environment). For instance, in mick-and-place tasks, the sparse reward could simply be defined the object being placed at the goal location. Nonetheless, sparse rewards also introduce challenges (e.g., exploration) for RI. algorithms (Pathak et al.) 2017 Burda et al.) 2018. Ecoffet et al. 2019). Therefore, a crucial question arises: can we learn dense reward functions in data-driven manner?
Ideally, the learned reward will be reused efficiently solve new tasks that share similar success conditions with the task used to learn the reward. For example, in pick-and-place tasks, different objects may need to be manipulated with varying dynamics, action spaces, and even robot morpholo- gies. For clarity, refer to each variant task and the set of all possible teck-and-place tasks as task family. Importantly, the reward function, which captures approaching, grasping, and moving the object toward the goal position, can potentially be transferred within this task family. This observation motivates us to explore the concept of reusable rewards, which can be learned as a function from some tasks and reused in unseen tasks. While existing literature in RL primarily focuses on the reusability (generalizability) of policies, we argue that rewards can pose greater flexibility for reuse across tasks. For example, it is nearly impossible directly transfer policy operating two-finger gripper for pick-and-pp -nd--lace to three-finger gripper due to action space misalignment, but a reward inducing the approach- saroouch-grasp--nn grasp-move workflow may apply for both types of grippers.
However, many existing works on reward learning do not emphasize reward reuse for new tasks. The field of learning reward function from demonstrations known as inverse RL in the literature (Ng
et al. 2000, Abbeel & Ng. 2004 Ziebart et al. 2008). More recently, adversarial imitation learning
(AIL) approaches have been proposed (Ho & Ermon 2016 Kostrikov et al. 2018 Fu et al. 2017 Ghasemipour et al. 2020) and gained popularity. Following the paradigm of GANs (Goddellow et al.
2020), AlL approaches employ a policy network to generate trajectories and train discriminator to distinguish between agent trajectories from demonstration ones. By using the discriminator score as rewards, (Ho & Ermon, 2016) shows that a policy can trained to imitate the demonstrations. Unfortunately, such rewards are reusable across tasks convergence, the discriminator outputs

---

Published as a conference paper at ICLR 2024
for both the agent trajectories and the demonstrations, as discussed in GGoodfello et al. 2020 Fu
et al. 2017). making it unable to learn useful information for solving new tasks.
In contrast to AIL, we propose novel approach for learning reusable rewards. Our approach involves incorporating sparse rewards as a supervision signal in lieu of the original signal used for classifying demonstration and agent trajectories. Specifically, success trajectories and failure trajectories based on the binary sparse reward. Please refer to Fig. 2 (a)(b) for an illustrative depiction. Our formulation assigns higher rewards to transitions in success trajectories and lower rewards to transitions within failure trajectories, which is consistent throughout the entire training process. As a result, the reward will be reusable once the training completed. Expert demonstrations can be included as success trajectories in our approach, though they are not mandatory. We only require the availability of a sparse reward, which a relatively weak requirement as it is often an inherent component of the task definition.
train a discriminator classify
Our approach can be extended to leverage the inherent structure of multi-stage tasks and derive stronger dense rewards. Many tasks naturally exhibit multi-stage structures, and it is relatively easy to assign a binary indicator on whether the agent has entered a stage. For example, in the "Open Cabinet Door" task depicted in Fig. | there are three stages: I) approach the door handle, 2) grasp the handle and pull the door, and 3) release the handle and keeping it steady. If the agent is grasping the handle of the door but the door has not been opened enough. then we can simply use a corresponding binary indicator asserting that the agent is in the 2nd stage. | By utilizing these stage indicators, we can learn dense reward for each stage and combine them into more structured reward. Since the horizon for each stage is shorter than that of the entire task, learning a high-quality dense reward becomes more feasible. Furthermore, this approach provides flexibility in incorporating extra information beyond the final success signal. We dub our approach as DrS (Dense reward learning from Stages).
Our approach exhibits strong performance on challenging tasks. To assess the reusability of the rewards learned by our approach, we employ the ManiSkill benchmark (Mu et al. [021]] Gu et al. 2023), which offers a large number of task variants within each task family. We evaluate our approach on three task families: Bock-and-Place Open Cabinet Door, and Turn Faucet, including 1000+ task variants. Each task variant involves manipulating a different object and requires precise low-level physical control, thereby highlighting the need for good dense reward. Our results demonstrate that the learned rewards can be reused across tasks, leading to improved performance and sample efficiency of RL algorithms compared to using sparse rewards. In certain tasks, the learned rewards even achieve performance comparable to those attained by human engineee reward functions.
Moreover, our approach drastically reduces the human effort needed for reward engincering. For instance, while the human-enginedd
reward for "Open Cabinet Door" involves over 100 lines of code, 10 candidate terms, and tons of "magic" parameters, our approach only requires two boolean functions as stage indicators: if the robot has grasped the handle and if the door is open enough. See appendix B for detailed example illustrating how our method reduces the required human effort.
Our contributions can be summarized as follows:
We propose DrS (Dense reward learning from Stages), novel approach for learning reusable dense rewards for multi-stage tasks, effectively reducing human efforts in reward engineering.
Extensive experiments on 1,000+ task variants from three task families showcase the effective- ness of our approach in generating high-quality and reusable dense rewards.
RELATED WORKS
Learning Reward from Demo (Offtine) Designing rewards is challenging due domain knowledge requirements, so approaches to learning rewards from data have gained attention. Some methods adopt classification-bust based rewards. i.c., training a reward by classifying goals (Smith et al. 2019
Kalashnikov et al.] 2021; Du et al., (023) or demonstration trajectories Zolna et al.) 2020). Other methods (Zkka et al. 2022 Aytar et al. 2018) use the distance to goal as a reward function, where the distance is usually computed in a learned embedding space, but these methods usually require
that the goal never changes in a task. These rewards are only trained on offline datasets, hence they can easily be exploited by an RL agent, i.e., an RL can enter a state that is not in the dataset and get wrong reward signal, as studied in (Vecerik et al. 2019 Xu & Denil 2021).
Stage indicators are only required during RI. training, but not required when deploying policy real world.
2

---

Published as a conference paper at ICLR 2024
I Stage Indicator: Grasp the Handle
x Siage Indicator: Door is Open Enough
x Stage Indicator: Door Static
Stage 1: Approach the Handle
Stage Pull the Handle
Stage Leave the Door Still
Success
Figure 1:
An illustration of stage indicators in an OpenCabinetDoor task, which 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.
Learning Reward from Demo (Online) The above issue be addressed by allowing agents to verify the reward in the environment, and inverse reinforcement learning (IRL) is the prominent paradigm. IRL aims to recover reward function given expert demonstrations. Traditional IRL methods (Ng et al. 2000; Abbeel & Ng. 2004 Ziebart et al., 2008 Ratliff et al., (206) often require
multiple iterations of Markov Decision Process solvers (Puterman 2014), resulting poor sample efficiency. In recent years, adversarial imitation learning (AIL) approaches are proposed (Ho &
Ermon 2016: Kostrikov et al. 2018: Fu al. 2017 Ghasemipour al. 2020) Liu al. 2019). They
operate similarly to generative ad vdeesariial networks (GANs) (Goodfe et al. 2020), in which generator (the policy) is trained to maximize the confusion of a discriminator, and the discriminator (serves the role of rewards) is trained to classify the agent trajectories and demonstrations, However, such rewards are not reusable as we discussed in the introduction classifying agent trajectories and demonstrations is impossible at convergence. In contrast, our approach gets rid of this issue by classifying the success/failure trajectories instead of expert/agent trajectories.
Learning Reward from Human Feedback Recent studies (Christiano et al.) 2017 Ibarz et al.) 2018 Jain et al. 2013) infer the reward through human preference queries on trajectories or explicitly asking for trajectory rankings (Brown al. 2019). Another line of works (Fu al. 2018 Singh et al., 2019) involves humans specifying desired outcomes or goals to learn rewards. However, in these methods, the rewards only distinguish goal from non-goal states, offering relatively weak incentives to agents at the beginning of an episode, especially in long- long-horizon tasks. contrast, our approach classifies all the states in the trajectories, providing strong guidance throughout the entire episode.
Reward Shaping Reward shaping methods aim to densify sparse rewards. Earlier works (Ng et al. 1999) study the forms of shaped rewards that induce the same optimal policy the ground-truth reward. Recently, some works (Toott et al. 2019 Wu et al. 202T) have shaped the rewards as the distance the goal, similar some offline reward learning methods mentioned above. Another idea (Memarian et al. (02T) involves shaping delayed reward by ranking trajectories based on fine- fine-grained preference oracle. contrast to these reward shaping approaches, our method leverages demonstrations, which are available in many real-world problems (Sun et al. 2020; Dasari et al.) 2019). This not only boosts the reward learning process but also reduces the additional domain knowledge required by these methods.
Task Decomposition The decomposition of tasks into stages/sub-tasks has been explored in various domains. Hierarchical RL approaches (Frans et al. 2018 Nachum et al., 2018 Levy et al.) 2018) break down policies into sub-policies to solve specific sub-tasks. Skill chaining methods (Lee et al. Gu al. 2022 Lee al. 2019 focus on solving long- long-horizon tasks by combining multiple short-horizon policies or skills. Recently, language models have also been utilized to break the whole task into sub-tasks (hhn et al. (2022). In contrast to these approaches that utilize stage structures in policy space, our work explores an orthogonal direction by designing rewards with stage structures.
2021
PROBLEM SETUP
In this work, we adopt the Markov Decision Process (MDP) M: AS, R,7) as the theoretical framework, where R is reward function that defines the goal purpose of task. Specifically, focus on tasks with sparse rewards. In this context, "sparse reward" denotes a binary reward function that gives value of upon successful task completion and 0 otherwise:
3

---

Published as a conference paper at ICLR 2024
doppiibanoo
discribiator
diectiii nator D
diaccteen nationt
zraining
pos
pos
agont
buffer
bufferi
buffer
Saaakia
EragActoris
taska
Innjecticine
Irajeclettee
and strape
trajectoriee
Irjectories bee
Phaze
Phosee roward rouse
roward leaming
(5) (oll-golioy)
(9) Reward Learning (One Stage)
(0) Reward Learning
(d) Pipeline
(Muli Stage)
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 stati-stage tasks, our approach train a separate discriminator for each stage. The discriminator for stage l 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.
task is completed by reaching one of the success states $ {
R Rsparse (6)
otherwise
Our objective is to learn dense reward function from a set of training tasks, with the intention of reusing it for unseen test tasks. Specifically, we aim to successfully train RL agents from scratch on the test tasks using the learned rewards. The desired outcome is to enhance the efficiency of RL training, surpassing the performance achieved by sparse rewards.
We assume that both the training and test tasks are in the same lask family. A task family refers to set of task variants that share the same success criteria, but may differ in terms of assets, initial states, transition functions, and other factors. For instance, the task family of object grasping includes tasks such as "Alice robot grasps an apple" and "Bob robot grasps a pen." The key point is that tasks within the same task family share a common underlying sparse reward.
Additionally, we posit that the task can be segmented into multiple stages, and the agent has access to several stage indicators obtained from the environment. stage indicator a binary function that indicates whether the current state corresponds to specific stage of the task. An example of stage indicators is Fig. T This assumption is quite general as many long-term tasks have multi-stage structures, and determining the current stage of the task is not hard in many cases. By utilizing these stage indicators, it becomes possible to construct a reward that is slightly denser than the binary sparse reward, which we refer to semi-sparse reward, and it serves strong baseline:
Rsemi-sparael
k, when state is at stage
We aim to design an approach that learns dense reward based on the stage indicators. When expert demonstration trajectories are available, they can also be incorporated to boost the learning process.
Note that the stage indicators are only required during RL training, but not required when deploying the policy to the real world. Training RL agents directly in the real world is often impractical due to cost and safety issues. Instead, a more common practice is to train the agent in simulators and then transfer/deploy it to the real world. While obtaining the stage indicators in simulators is fairly easy, it also possible to obtain them the real world by various techniques (robot proprioception, tactile
sensors Lin et al. (2022); Melnik et al. 12221) visual detection/tracking Kalashnikov et al. (2018)
2021
large vision- language models Du et al. 2023). etc.).
4 DRS: DENSE REWARD LEARNING FROM STAGES
Dense rewards are often tricky to design by humans (see an example in appendix B), so we aim to learn reusable dense reward function from stage indicators in multi-stage tasks and demonstrations when available. Overall, our approach has two phases, as shown in Fig. (d):
Reward Learning Phase: learn the dense reward function using training tasks.
Reward Reuse Phase: reuse the learned dense reward to train new RL agents in test tasks.
Since the reward reuse phase is just regular RL training process, only discuss the reward learning phase in this section. We first explain how our approach learns a dense reward in one-stage tasks (Sec. Then, extend this approach to multi-stage tasks (Sec. 4.2).
4

---

Published as a conference paper at ICLR 2024
4.1 REWARD LEARNING ON ONE-STAGE TASKS
In line with previous work (Vecerik et al. 2019 Fü et al. 2018), we employ a classification--ass dense reward. We train a classifier to distinguish between good and bad trajectories, utilizing the learned classifier as dense reward. Essentially, states resembling those in good trajectories receive higher rewards, while states resembling bad trajectories receive lower rewards. While previous Adversarial Imitation Learning (AII.) methods (Ho & Ermon 2016, Kostrikov et al., 2018) used discriminators as classifiersrrewards to distinguish between agent and demonstration trajectories, these discriminators cannot be directly reused as rewards to train new RL agents. As the policy improves, the agent trajectories (negative data) and the demonstrations (positive data) can become nearly identical. Therefore, at convergence, the discriminator output for both agent trajectories and demonstrations tends to approach 1. as observed in GANs (Goodfellow et al. 2020) (also noted by (Fu et al. 2017 Xu & Denil 202T)). This makes it unable to learn useful into for solving new tasks.
Our approach introduces a simple modification to existing AIL methods to ensure that the discrim- inator continues to learn meaningful information even convergence. The key issue previously mentioned arises from the diminishing gap between agent and demonstration trajectories over time, making it challenging to differentiate between positive and negative data. To address this, we propose training the discriminator to distinguish between success and failure trajectories instead of agent and demonstration trajectories. By defining success and failure trajectories based on the sparse reward signal from the environment, the gap between them remains intact and does not shrink. Consequently, the discriminator effectively emulates the sparse reward signal, providing dense reward signals to the RI. agent. Intuitively, state that is closer to the success states in terms of task progress (rather than Euclidean distance) receives a higher reward, as it is more likely to occur in success trajectories. Fig. 2 a) and (b) illustrate the distinction between our approach and traditional AIL methods.
To ensure that the training data consistently includes both success and failure trajectories, we use replay buffers to store historical experiences, and train the discriminator in an off off-policy manner. While the original GAIL is on- on-policy, recent AIL methods (Kostrikov et al. 2018) Orsini et al. 202T) have adopted off-policy training for better sample efficiency. Note that although our approach shares similarities with AIL methods, it is not adversarial in nature. In particular, our policy does not aim to deceive the discriminator, and the discriminator does not seek to penalize the agent's trajectories.
4.2 REWARD LEARNING ON MULTI-STAGE TASKS
Illstrationn of the Learned Reward
In multi-stage tasks, it is desirable for the reward of state in stage k +++ 1 to be strictly higher than that of stage k to incentivize the agent to progress towards later stages. The semi-sparse reward (Eq. 2 aligns with this intuition, but it is still a bit too sparse. If each stage of the task is viewed as an individual task, the semi-sparse reward acts a sparse reward for each stage. In the case of a one-stage task, a discriminator can be employed to provide a dense reward. Similarly, for multi-stage tasks, separate discriminator can be trained for each stage to serve as a dense reward for that particular stage. By training stage-specific discriminators, we can effectively address the sparse reward issue and guide the agent's progress through the different stages of the task. Fig. gives an intuitive illustration of our learned reward, which fills the gaps in semi-sparse rewards, resulting in a smooth reward curve.
Reward
Semi sparse
Task Progress (fo a S-stage tank)
Figure 3: An illustration of our learned reward, which fills the gaps in semi-sparse rewards, re- sulting in smooth reward curve.
To train the discriminators for different stages, we need to establish the positive and negative data for each discriminator. In one-stage tasks, positive data comprises success trajectories and negative data encompasses failure trajectories. In multi-stage tasks, we adopt a similar approach with a slight modification. Specifically, we assign a stage index to each trajectory, which is determined as the highest stage index among all states within the trajectory:
StageIndex(r (so, s1.)..) = max StageIndex(ss).
where is trajectory and s; are the states in For the discriminator associated with stage k. positive data consists of trajectories that progress beyond stage (StageIndex > k),, and negative data consists of trajectories that reach up stage l (StageIndex & k).
5

---

Published as a conference paper at ICLR 2024
Once the positive and negative data for each discriminator have been established, the next step is to combine these discriminators to create a reward function. While the semi-sparse reward (Eq. 2 lacks incentives for the agent at stage k until it reaches stage k + 1, we can fill in the gaps in the semi-sparse reward by the stage-specific discriminators. We define our learned reward function for a multi-stage task as follows:
R(5') = + a- - taulk(ssrrimiiattrs((]]
where k is the stage index of s' and a is hyperparameter. Basically, the formula incorporates a dense reward term into the semi-sparse reward. The tanh function is used to bound the output of the discriminators. As the range of the tanh function is (-1, 1), any < 2 ensures that the reward of a state in stage k + 1 is always higher than that of stage k. In practice, we use
and it works well.
4.3 IMPLEMENTATION
Algorithm 1 DrS (Dense reward learning from Stages
Require: Task MDP .M, Number of stages in task N, Demonstra-
From the implementation perspec-
tion dataset D [-" _²,, (optional)
tive, our approach similar to GAIL,
1: Initialize policy и., critic Q, replay buffer BR
but with different training process for discriminators. While the orig- inal GAIL is combined with TRPO
2: Initialize discriminators fur fI, ..., fN-1. stage buffers
Ba, B2..., Bxx
3: Fill demo D into Bx: By By UD
(Schulman et al.) 2015). (Orsini et al.)
4: for each iteration do
2021) found that using state-of-the-art
5: Collect trajectories {To, - ......) by executing x M
off-policy RL algorithms (like SAC
6: Add trajectories to replay buffer: Bn 4- Bn U (T2, T.....
(aarnoja et al.] 2018) TD3 (Fu-
7:
for each trajectory = in (f2: do
jimoto et al. 2018)) can greatly im- prove the sample efficiency of GAIL.
8:
j StageIndex(rz) according to Eq.
9:
B, < B, (T*)
Therefore, we also combine our ap- proach with SAC, and the full algo-
10:
for each gradient step for discriminators do for each discriminator fe do
11:
rithm is summarized in Algo. μ
12:
Sample negative data from U B.
13:
Sample positive data from U B Update fx using BCE loss
In addition to the regular replay buffer
14:
used in SAC, our approach maintains
15:
for each gradient step for the policy do
N different stage buffers to store tra-
16:
Sample from B2
jectories corresponding to different
17:
Compute rewards according to Eq. Update т and Q by SAC (Harnojj el al. 2018
stages(defined by Eq. 3). Each tra-
18:
jectory is assigned to only one stage
buffer based on its stage index. During the training of the discriminators, we sample data from the union of multiple buffers. In practice, we early stop the discriminator training of k once its success rate is sufficiently high, as we find it reduces the computational cost and makes the learned reward more robust. Note that our approach uses the next state s' as the input to the reward. which aligns with common practices in human reward engineering (Gu et al.) [023] Zhu et al. (020). However, our approach is also compatible with alternative forms of input, such as (s, a) Or (8,0,8)
EXPERIMENTS
5.1 SETUP AND TASK DESCRIPTIONS
We evaluated our approach on three challenging physical manipulation task families from the ManiSkill (Mu et al. 2021 Gu et al. 2023): Pick-and-Placee Turn Faucet, and Open Cabinet Door. Each task family includes set of different objects be manipulated. To assess the reusability of the learned rewards, we divided the objects within each task family into non-overlapping training and test sets, as depicted in Fig 4 During the reward learning phase, we learned the rewards by training an agent for each task family to manipulate all training objects. In the subsequent reward reuse phase, the learned reward rewards are reused to train an agent to manipulate all test objects for each task family. And compare with other baseline rewards in this reward reuse phase. It is important to note that our learned rewards are agnostic to the specific RL algorithm employed. However, we utilized the Soft actor-Critic (SAC) algorithm to evaluate the quality of the different rewards.
To assess the reusability of the learned rewards, is crucial to have a diverse set of tasks that exhibit similar structures and goals but possess variations in other aspects. However, most existing

---

Published as a conference paper at ICLR 2024
Training Tasks (-100 tasks)
Test Tasks (more than 1,000 tasks)
Pick And Place
Turn Faucet
Open Cabinet Door
Figure 4: We evaluated approach DrS more than 1,000 task variants from three task families in ManiSkill (Mu et al. 2021 Gu et al., (023). Each task variant associated with a different object. All tasks require level physical control. The objects in training and test tasks are non-overlapped.
benchmarks lack an adequate number of task variations within the same task family. As a result, primarily conducted our evaluation the ManiSkill benchmark, which offers range of object variations within each task family. This allowed us to thoroughly evaluate our learned rewards in realistic and comprehensive manner.
Pick-and-Place A robot arm is tasked with picking up an object and relocating it to a random goal position in mid-air. The task completed the object in close proximity to the goal position, and both the robot arm and the object remain stationary. The stage indicators include: (a) the gripper grasps the object, (b) the object close the goal position, and (c) both the robot and the object are stationary. We learn rewards on 74 YCB objects and reuse rewards on 1 1,600 EGAD objects.
Turn Faucet: robot arm tasked turn on faucet by rotating its handle. The task is completed if the handle reaches target angle. The stage indicators include: (a) the target handle starts moving, (b) the handle reaches target angle. We learn rewards on 10 faucets and reuse rewards on 50 faucets.
Open Cabinet Door: A single-arm mobile robot is required to open designated target door on a cabinet. The task is completed if the target door opened to sufficient degree and remains stationary. The stage indicators include: (a) the robot grasps the door handle, (b) the door is open enough, and (c) the door is stationary. We learn rewards on cabinet doors and reuse rewards on cabinet doors. Note that remove all single-door cabinets in this task family, they be solved by kicking the side of the door and this behavior can be readily learned by sparse rewards.
We employed low-level physical control for all task families. Please refer to the appendix A for detailed description of the object sets, action space, state space, and demonstration trajectories.
5.2 BASELINES
Human- Engineered The original human--ritten dense rewards in the benchmark, which require significant amount of domain knowledge, thus can be considered as an upper bound of performance.
Semi-Sparse The rewards constructed based on the stage indicators, discussed Eq The agent receives a reward of & when it is in stage A: This baseline extends the binary sparse reward.
VICE-RAQ (Singh et al. 2019) An improved version of VICE (Fu et al. 2018). It learns classifier, where the positive samples are successful states annotated by querying humans, and the negative samples are all other states collected by the agent. Since our experiments do not involve human feedback, we let VICE-RAQ query the oracle success condition infinitely for a fair comparison.
ORIL (Zolna et al. 2020) representative offline reward learning method, where the agent does not interact with the environments but purely leams from the demonstrations. It learns a classifier (reward) distinguish between the states from success trajectories and random trajectories.
5.3 COMPARISON WITH BASELINE REWARDS
We trained RL agents using various rewards and assessed the reward quality based on both the sample efficiency and final performance of the agents. The experimental results, depicted in Fig.
7

---

Published as a conference paper at ICLR 2024
PickAndPlace
TurnFaucet
OpenCabinetDoor
80
80
80
40
40
40
0.0
Environment Steps
Environment Steps
Environment Steps
1e6
Human-Engiannn Engineered
Ours
Sami-Sparse
VICE-RAD
ORIL
Figure 5: Evaluation results of reusing learned rewards. All curves use SAC train, but with different rewards. VICE-RAQ and ORIL get no success. 5 random seeds, the shaded region is std.
demonstrate that our learned reward surpasses semi-sparse rewards and all other reward learning methods across all three task families. This outcome suggests that our approach successfully acquires high-quality rewards that significantly enhance RL training. Remarkably, our learned rewards even achieve performance comparable to human-engineered rewards in Pick-and-Place and Turn Faucet.
Semi-sparse rewards yielded limited success within the allocated training budget, suggesting that RL agents face exploration challenges when confronted with sparse reward signals. VICE-RAQ failed in all tasks. Notably, it actually failed during the reward learning phase on the training tasks, rendering the learned rewards inadequate for supporting RL training on the test tasks. This failure aligns with observations made by (Wu et al. 202Tb. We hypothesize that by only classifying the success states from other states, it cannot provide sufficient guidance during the early stages of training. where most states are distant from the success states and receive low rewards. Unsurprisingy. ORIL does not get any success on all tasks either. Without interacting with the environments to gather more data, the learned reward functions easily tend overfit the provided dataset. When using such rewards in RL, the flaws in the learned rewards are easily exploited by the RL agents.
5.4 ABLATION STUDY
We examined various design choices within our approach on the Pick-and- aeccaend task family.
5.4.1 ROBUSTNESS TO STAGE CONFIGURATIONS
Though many tasks present natural structure of stages, there are still different ways divide task into stages. To assess the robustness of our approach in handling different task structures, we experiment with different numbers of stages and different ways to define stage indicators.
Number of Stages The Pick-and-Place task family originally consisted of three stages: (a) approach the object, (b) move the object to the goal, and (c) make everything stationary. We explored two ways of reducing the number of stages to two, namely merging stages (a) and (b) or merging stages (b) and (c), as well as the I-stage case. Our results, presented Fig. 5 indicate that the learned rewards with 2 stages efficiency than those with stages. Specifically, the reward that preserves stage (c) "make everything stationary" performs slightly better than the reward that preserves stage (a) "approach the object". This suggests that it may be more challenging for a robot to learn to stop abruptly without a dedicated stage. However, when reducing the number of stages to 1, the learned reward failed to train RL agents in test tasks, demonstrating the benefit of using more stages in our approach.
still effectively train RL agents in test tasks, albeit with lower sample
Definition of Stages The stage indicator "object is placed" is initially defined if the distance between the object and the goal is less than 2.5 cm. We create two variants of it, where the distance thresholds are 5em and 10cm, respectively. The results, as depicted Fig. 7 demonstrate that changing the distance threshold within reasonable range does not significantly affect the efficiency of RL training. Note that the task success condition is unchanged, and our rewards consistently encourage the agents to reach the success state it yields the highest reward according to Eq. Al The stage definitions solely affect the efficiency of RL training during the reward reuse phase.
Overall, the above results highlight the robustness of our approach to different stage configurations, indicating that it is not heavily reliant on intricate stage designs. This robustness contributes to a significant reduction in the burden of human reward engineering.
8

---

Published as a conference paper at ICLR 2024
Different Numbers of Stages
Different Stage Definitions
Fine-tune Policy
100
100
Succeee Ratte %
Buscess Rate
Ours (3-ttgel
(öst c2.5cll
Gurs (Pastage, x-bc]
Ours (dit x Scm)
wi Taman-nngeeere
Ours 12 stags, abx-)
Dura (daa - 10cm))
w/ Eeeri Sparse
Ours it stage)
Semi Sparse
Wr Ouu
0.0
0.5
1.0
1.5
200
00
0.5
1.0
1.5
2.0
0.25
0.50
1.00
Emvironmentt Stops
145
Endrronment Stops
Eovironment Steps
Figure 6: Ablation study on the number of stages, see here
Figure Ablation study on the stage definitions, see here
Figure 8: Fine-tune the policy from reward learning, see here
5.4.2 FINE-TUNING POLICY
In our previous experiments, we assessed the quality of the learned reward by reusing it in training RL agents from scratch since it is the most common and natural way to use a reward. However, approach also produces policy
byproduct in the reward learning phase. This policy can also be fine-tuned using various rewards in new tasks, providing an alternative to training RI. agents from scratch. We compare the fine-tuning of the byproduct policy using human-engineered rewards, semi-sparse rewards, and our learned rewards.
As shown in Fig. 8 all policies improve rapidly at the beginning due the good initialization of the policies. However, fine-tuning with our learned reward yields the best performance (even slightly better than the human-engineered reward), indicating the advantages of utilizing our learned dense reward even with good initialization. Furthermore, the significant variance observed when fine-tuning the policy with semi-sparse rewards highlights the limitations of sparse reward signals in effectively training RI. agents, even with very good initialization.
5.4.3 ADDITIONAL ABLATION STUDIES
Additional ablation studies are provided in appendix E with key conclusions summarized as follows:
DrS is compatible with various modalities of reward input, including point cloud data. E.
Reward learned by GAIL, even with stage indicators, is not reusable.
The way of combining the dense rewards from each stage matters.
CONCLUSION AND LIMITATIONS
To make RL a more widely applicable tool, have developed data-driven approach for learning dense reward functions that can be reused in new tasks from sparse rewards. We have evaluated the effectiveness of our approach on robotic manipulation tasks, which have nigh-dimensional action spaces and require dense rewards. Our results indicate that the learned dense rewards are effective in transferring across tasks with significant variation in object geometry. By simplifying the reward design process, our approach paves the way for scaling up RL in diverse scenarios.
We would like to discuss two main limitations when using the multi-stage version of our approach.
Firstly, though our experiments show the substantial benefits of knowing the multi-stage structure of tasks (at training time, needed at policy deployment time), did not specifically investigate how this knowledge can be acquired. Much future work on be done here, by leveraging large language models such as ChatGPT (OpenAl 2023) (by our testing, they suggest stages highly aligned the ones we adopt by intuition for all tasks in this work) or employing iefommation-thoorttii approaches. Further discussions regarding this point can be found in appendix
Secondly, the reliance on stage indicators adds a level of inconvenience when directly training RL agents the real world. While it is infrequent to directly train RL agents in the real world due to cost and safety issues, when necessary, stage information can still be obtained using existing techniques, similar to (Kalashnikov et al.] 2018 2021). For example, the "object is grasped" indicator can be acquired by tactile sensors (Lin et al. 2022, Melnik et al.) (021), and the "object is placed" indicator can be obtained by forward kinematics, visual detection/trackin techniques (Kalashnikov et al. 2018 202T), or even large vision-language models (Du al. 2023)

---

Published as a conference paper at ICLR 2024
REFERENCES
Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In Proceedings of the twenty-first international conference on Machine learning. pp. 2004.
Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, et al. Do as can, not say: Grounding language in robotic affordances. arXiv preprint arXiv: 22X1::2204.1691,, 2022.
Yusuf Aytar, Tobias Pfaff, David Budden, Thomas Paine, Ziyu Wang. and Nando De Freitas. Playing hard exploration games by watching youtube. Advances in neural information processing systems, 31, 2018.
Daniel Brown, Wonjoon Goo, Prabhat Nagarajan, and Scott Nickum. Extrapolating beyond sub- optimal demonstrations via inverse reinforcement learning from observations. In International conference on machine learning, pp. 783-792. PMLR, 2019.
Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. arXiv preprint arXiv: 18/0. /2894, 2018.
Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. Advances in neural information processing systems, 30, 2017.
Sudeep Dasari, Frederik Ebert, Stephen Tian, Suraj Nair, Bernadette Bucher, Karl Schmeckpeper, Siddharth Singh, Sergey Levine, and Chelsea Finn. Robonet: Large-scale multi-robot learning. arXiv preprint arXiv:1910.11211 2019.
Yuqing Du, Ksenia Konyushkova, Misha Denil, Akhil Raju, Jessica Landon, Felix Hill, Nando de Freitas, and Serkan Cabi. Vision-language models as success detectors. arXiv preprint arXiv:2303.07280, 2023.
Adrien Ecoffet, Joost Huizinga, Joel Lehman, Kenneth O Stanley, and Jeff Clune. Go-explore: new approach for hard-exploration problems. arXiv preprint 19X1. 2019.
Kevin Frans, Jonathan Ho. Xi Chen, Pieter Abbeel, and John Schulman. Meta learning shared hierarchies. In International Conference on Learning Representations, 2018
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse reinforce- ment learning. arXiv preprint arXiv:1710.112: 2017.
Justin Fu, Avi Singh, Dibya Ghosh, Larry Yang, and Sergey Levine. Variational inverse control with events: A general framework for data-driven reward definition. Advances in neural information processing systems, 31, 2018.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International conference on machine learning, pp. 1587-1596. PMLR, 2018.
Seyed Kamyar Seyed Ghasemipour, Richard Zemel, and Shixiang Gu. A divergence minimization perspective on imitation learning methods. In Conference on Robot Learning, pp. 1259-1277 PMLR, 2020.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11):139-144, 2020.
Jiayuan Gu, Devendra Singh Chaplot, Hao Su, and Jitendra Malik. Multi-skill mobile manipulation for object rearrangement. arXiv preprint ArXiv:2209.022778 2022.
Jiayuan Gu, Fanbo Xiang, Xuanlin Li, Zhan Ling, Xiqiaing Liu. Tongzhou Mu, Yihe Tang, Stone Tao, Xinyue Wei, Yunchao Yao, Xiaodi Yuan, Pengwei Xie, Zhiao Huang, Rui Chen, and Hao Su. Maniskil12: A unified benchmark for generalizable manipulation skills. In International Conference Learning Representations, 2023.
10

---

Published as a conference paper at ICLR 2024
Huy Ha, Pete Florence, and Shuran Song. Scaling up and distilling down: Language-guinn guided robot skill acquisition. arXiv preprint arXiv:2307.14535 2023.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy
maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning. pp. 1861-1870. PMLR, 2018.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.
Jemin Hwangho, Joonho Lee, Alexey Dosovitskiy, Dario Bellicoso, Vassilios Tsounis, Vladlen Koltun, and Marco Hutter. Learning agile and dynamic motor skills for legged robots. Science Robotics, a(26))2a 2019.
Borja Ibarz, Jan Leike, Tobias Pohlen, Geoffrey Irving, Shane Legg, and Dario Amodei. Reward learning from human preferences and demonstrations in atari. Advances in neural information processing systems, 31, 2018.
Ashesh Jain, Brian Wojcik, Thorsten Joachims, and Ashutosh Saxena. Learning trajectory preferences for manipulators via iterative improvement. Advances in neural information processing systems, 26, 2013.
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, et al. Scalable deep reinforcement learning for vision-based robotic manipulation. Conference on Robot Learning, pp. 651-673.
PMLR, 2018.
Dmitry Kalashnikov, Jacob Varley, Yevgen Chebotar, Benjamin Swanson, Rico Jonschkowski, Chelsea Finn, Sergey Levine, and Karol Hausman. Mt-opt: Continuous multi-task robotic rein- forcement learning at scale. arXiv preprint arX1::2104.0821,, 2021.
Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tompson. Discriminator-acto--iiiin Addressing sample inefficiency and reward bias in adversarial imitation learning. arXiv preprint
089900299000.....00 2018.
Youngwoon Lee, Shao-Hua Sun, Sriram Somasundaram, Edward Hu, and Joseph J Lim. Com- posing complex skills by learning transition policies. In International Conference on Learning Representations, 2019.
Youngwoon Lee, Joseph J Lim, Anima Anandkumar, and Yuke Zhu. Adversarial skill chaining for long-horizon robot manipulation via terminal state regularization. arXiv preprint arXiv:2111.0799 2021.
Andrew Levy, George Konidaris, Robert Platt, and Kate Saenko. Learning multi-level hierarchies with hindsight. In International Conference Learning Representations, 2018.
Jacky Liang, Wenlong Huang, Fei Xia, Peng Xu, Karol Hausman, Brian Ichter, Pete Florence, and Andy Zeng. Code as policies: Language model programs for embodied control. 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 9493 IEEE, 2023.
Yijiong Lin, John Lloyd, Alex Church, and Nathan F Lepora. Tactile gym 2.0: Sim-to-real deep reinforcement learning for comparing low-cost high-resolution robot touch. IEEE Robotics and Automation Letters, 7(4):107441 2022.
Fangchen Liu, Zhan Ling, Tongzhou Mu, and Hao Su. State alignment- based imitation learning. arXiv preprint arXiv:1911.10% 2019.
Andrew Melnik, Luca Lach, Matthias Plappert, Timo Korthals, Robert Haschke, and Helge Ritter. Using tactile sensing to improve the sample efficiency and performance of deep deterministic policy gradients for simulated in-hand manipulation tasks. Frontiers in Robotics and AI, 8:538773, 2021.
11

---

Published as a conference paper at ICLR 2024
Farzan Memarian, Wonjoon Goo, Rudolf Lioutikov, Scott Nickum, and Ufuk Topcu. Self Self-supervised online reward shaping in sparse-reward environments. In 2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 1369-2375 IEEE, 2021.
Tongzhou Mu, Zhan Ling, Fanbo Xiang, Derek Cathera Yang, Xuanlin Li, Stone Tao, Zhiao Huang. Zhiwei Jia, and Hao Su. Maniskill: Generalizable manipulation skill benchmark with large-scale demonstrations. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
Ofrr Nachum, Shixiang Shane Gu, Honglak Lee, and Sergey Levine. Data-efficient hierarchical reinforcement learning. Advances in Neural Information Processing Systems, 31 31:3303-3313, 2018.
Andrew Y Ng. Daishi Harada, and Stuart Russell. Policy invariance under reward transformations:
Theory and application reward shaping. In leml, volume 99, pp. 278-287, 1999.
Andrew Ng, Stuart Russell, et al. Algorithms for inverse reinforcement learning. In leml, volume 1, pp. 2000.
OpenAl. Gpt-4 technical report, 2023.
Manu Orsini, Anton Raichuk, Léonard Hussenot, Damien Vincent, Robert Dadashi, Sertan Girgin, Matthieu Geist, Olivier Bachem, Olivier Pietquin, and Marcin Andrychowicz. What matters for adversarial imitation learning? Advances in Neural Information Processing Systems, 34: 14656-14668, 2021.
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International conference on machine learning, 2778-2787 PMLR, 2017.
Xue Bin Peng, Pieter Abbeel, Sergey Levine, and Michiel Van de Panne. Deepmimic: Example- guided deep reinforcement learning of physics- character skills. ACM Transactions On Graphics (TOG), 37(4):1-14, 2018.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John
Wiley & Sons, 2014.
Nathan D Ratliff Andrew Bagnell, and Martin A Zinkevich. Maximum margin planning. In Proceedings of the 23rd international conference on Machine learning, pp. 729-736, 2006.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pp. 1889-1897. PMLR, 2015.
Lucy Xiaoyang Shi, Archit Sharma, Tony Z Zhao, and Chelsea Finn. Waypoint-based imitation learning for robotic manipulation. arXiv preprint arXiv:2307. 2023.
Avi Singh, Larry Yang, Kristian Hartikainen, Chelsea Finn, and Sergey Levine. End-to-end robotic 2019.
reinforcement learning without reward engineering. arXiv preprint
Ishika Singh, Valts Blukis, Arsalan Mousavian, Ankit Goyal, Danfei Xu, Jonathan Tremblay, Dieter Fox, Jesse Thomason, and Animesh Garg. Progprompt: Generating situated robot task plans using large language models. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 11523-11530 IEEE, 2023.
Laura Smith, Nikita Dhawan, Marvin Zhang, Pieter Abbeel, and Sergey Levine. Avid: Learning multi-stage tasks via pixel-level translation of human videos. arXiv preprint 1911v:111..04443, 2019.
Pei Sun, Henrik Kretzschmar, Xerxes Dotiwalla, Aurelien Chouard, Vijaysai Patnaik, Paul Tsui, James Guo, Yin Zhou, Yuning Chai, Benjamin Caine, et al. Scalability in perception for autonomous driving: Waymo open dataset. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 2446-2454, 2020.
12

---

Published as a conference paper at ICLR 2024
Alexander Trott, Stephan Zheng, Caiming Xiong, and Richard Socher. Keeping your distance: Solving sparse reward tasks using hafl-balnniin shaped rewards. Advances in Neural Information Processing Systems, 32, 2019.
Mel Vecerik, Oleg Sushkov, David Barker, Thomas Rothörl, Todd Hester, and Jon Scholz. A practical approach to insertion with variable socket position using deep reinforcement learning. In 2019 international conference on robotics and automation (ICRA), pp. 754-760. IEEE, 2019.
Zheng Wu, Wenzhao Lian, Vaibhav Unhelkar, Masayoshi Tomizuka, and Stefan Schaal. Learning dense rewards for contact-rich manipulation tasks. In 2021 IEEE International Conference Robotics and Automation (ICRA), pp. 6214-6221, IEEE, 2021.
Tianbao Xie, Siheng Zhao, Chen Henry Wu, Yitao Liu, Qian Luo, Victor Zhong, Yanchao Yang. and Tao Yu. Text2reward: Automated dense reward function generation for reinforcement learning. arXiv preprint arXiv:2309. /1489, 2023.
Danfei Xu and Misha Denil. Positive-umlabeled reward learning. In Conference on Robot Learning, pp. 205-201220 219. PMLR, 2021.
Wenhao Yu, Nimrod Gileadi, Chuyuan Fu, Sean Kirmani, Kuang-Huet Lee, Montse Gonzalez Arenas, Hao-Tien Lewis Chiang, Tom Erez, Leonard Hasenclever, Jan Humplik, et al. Language rewards for robotic skill synthesis. arXiv preprint arXiv:2306. 086477, 2023.
Kevin Zakka, Andy Zeng, Pete Florence, Jonathan Tompson, Jeannette Bohg, and Debidatta Dwibedi. Xirl: Cross-embodiment inverse reinforcement learning. In Conference on Robot Learning, pp- 537-546. PMLR, 2022.
Yuke Zhu, Josiah Wong, Ajay Mandlekar, and Roberto
robosuite: A modular simulation framework and benchmark for robot learning. arXiv preprint (tti:::000.2299990 2020.
Brian Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Aaai, volume 8, pp. 1433-1438. Chicago, IL, USA, 2008
Konrad Zolna, Alexander Novikov, Ksenia Konyushkova, Caglar Gulcehre, Ziyu Wang. Yusuf Aytar, Misha Denil, Nando de Freitas, and Scott Reed. Offline learning from demonstrations and unlabeled experience, arXiv preprint arXiv:201 13885, 2020.
13


Reviews:
[{"review_id": "dyO1LDtFi9", "reviewer": "Reviewer_ck68", "summary": "The paper presents a method, Dense reward learning from Stages (DrS), for learning reusable dense reward functions using only sparse task reward signals and task-family-specific “stage indicators”. DrS learns to classify the stage reached by each trajectory and adds the discriminator score to the index of the stage reached, leading to a reward function that increases both across stages and within progress along a stage. Prior adversarial imitation learning approaches to reward learning do not lead to reward functions that can be reused because the discriminator cannot distinguish policy trajectories from demonstrations at convergence. Experiments show that DrS-learned rewards can be reused for test tasks and can even compete with human-engineered rewards in some cases, despite requiring much less manual engineering (just the stage indicator function).", "questions": "Did you try running DrS without providing demonstrations? Or varying the number of demonstrations?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "* As far as I am aware, the overall design of the algorithm (exploiting stage indicators) as well as the form of the learned reward function are novel.\n* The method significantly reduces the amount of engineering required to learn reward functions when the task can be broken down into identifiable stages.\n* DrS handily outperforms several reasonable baselines and competes with the hand-engineered reward in some cases.\n* In addition to the ablation study in the body of the paper, the appendix includes extensive additional experiments, such as varying the number of training tasks, input modality, dynamics, action space, robot morphology, and form of the reward function. The authors also verify the claim that GAIL-generated rewards are not reusable.\n* The paper is clear, easy to read, and well-motivated.", "weaknesses": "* The paper states that demonstrations are optional, but it sounds like they were used in all experiments. I imagine that sample efficiency would deteriorate substantially if no demonstrations are provided and the first stage cannot be solved easily by random exploration, or more generally if any stage cannot be easily solved by a noisy version of a policy that solves the previous stage.\n* It is not clear that all sparse-reward tasks can be broken up into stages, and as shown in the ablation study, the method struggles when there is only one stage. So DrS is not universally applicable to sparse-reward tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NLofNO2i2k", "reviewer": "Reviewer_oxbF", "summary": "This paper proposes a method for learning reusable dense rewards for multi-stage tasks using binary stage indicators. The idea is to train a discriminator that distinguishes successful trajectories from unsuccessful ones for each stage, and then use the discriminator output as dense reward. The authors demonstrate that the learned reward functions generalize to thousands of tasks across three Maniskill domains, resulting in policies that perform nearly as well as those trained from human-designed rewards. This is a step towards automating reward design, which is a long-standing problem in RL.", "questions": "- Does the method work with pixel observations?\n- If the method only works with states, does generalization come as a result of state abstraction or the method itself?\n- Do the learned dense reward functions resemble hand-designed rewards?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "- This paper makes a contribution towards automating reward design, which is of paramount importance in the field of RL. Having access to dense rewards takes the burden off of exploration, which in turn reduces the number of samples required to solve a task.\n- The method is a niche application of contrastive discriminator learning, which is well-established in the literature.", "weaknesses": "- The method requires success and failure trajectories for each stage in the training data, which can be expensive to collect. \n- The scope of the method is limited to a family of tasks that can be divided into stages. This prevents it from being applied to other tasks such as locomotion. It also means the method is less general compared to LLM-based rewards with external knowledge [1, 2].\n- Similarly, the need for stage indicators prevents the method from scaling to real-world problems, which would arguably benefit more from automated reward design than simulated domains.\n- The experiments do not demonstrate new skills that are unachievable by human-designed rewards.\n\nReferences:\n1. Tianbao Xie, Siheng Zhao, Chen Henry Wu, Yitao Liu, Qian Luo, Victor Zhong, Yanchao Yang, and Tao Yu. Text2reward: Automated dense reward function generation for reinforcement learning. arXiv preprint, 2023.\n2. Yecheng Jason Ma, William Liang, Guanzhi Wang, De-An Huang, Osbert Bastani, Dinesh Jayaraman, Yuke Zhu, Linxi “Jim” Fan, Anima Anandkumar. Eureka: Human-Level Reward Design via Coding Large Language Models. arXiv preprint, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qnRCiTsXHn", "reviewer": "Reviewer_GMmw", "summary": "This paper proposes `DrS` (Dense reward learning from Stages), an approach for learning reusable dense rewards for multi-stage tasks, effectively reducing human efforts in reward engineering. By breaking down the dask into stages, this method learns dense reward from sparse ones using human demonstrations, where the focus is to learn reusable representations that can potentially by used on unseen tasks with similar structure. An example of structure is illustrated using the `Open Cabinet Door` which can be naturally divided into stages, such as approaching the handle, grasping and pulling, and then releasing it. Recognizing which stage an agent is in can be done using simple binary indicators. By applying these indicators, this work cultivates a dense reward for every stage.\n\nExperiments are `PickNPlace`, `TurnFaucet`, `OpenCabinetDoor`.\n\nBaselines are `dense reward`,  `VICE-RAQ`, `ORIL`", "questions": "- I did not understand what Figure 3 is presenting and there is no reference in the paper. Does this plot mean that `DrS` is learning a sigmoid liked function of the sparse reward?\n- In the limitation section, the authors talk about the usage of language models such as ChatGPT and I think some discussion can be built around methods such as say-can[1]. I think the goal of both frameworks are aligned.\n\n\n\n\n\n[1] Do As I Can, Not As I Say: Grounding Language in Robotic Affordances -  https://say-can.github.io/", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The goal of this work is deriving a dense reward function from an array of training tasks to be repurposed for new, unseen tasks.\n- The notion of capturing representations for a `task family` is important for enabling RL agents to learn multi-purpose policies.\n-  Operating on the understanding that tasks can be broken down into several segments is also logical.\n- I like this paper because I think it's important to move away from engineered dense rewards, to more tangible methodologies for learning rewards, specially in stage-drived manner, using demonstrations.", "weaknesses": "- Some ablation study of the robustness of this method against bad demonstrations (i.e. suboptimal, noisy etc) could be nice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "da3IIougUh", "reviewer": "Reviewer_35ZH", "summary": "This paper considers learning reusable dense rewards for multi-stage tasks and proposes a method named as DrS. DrS classifies each collected trajectory by its stage and relabels them into failure and success ones. It then defines the learned reward function as the stage index plus the output of the corresponding discriminator. Extensive experiments from three task families showcase the effectiveness of DrS.", "questions": "1. The authors asked \"can we learn dense reward functions in a data-driven manner?\" But what kind of dense reward do we really need? Say, there are three reward function $r_1$, $r_2$ and $r_3$, where $r_1=1$ if reaching the goal state, otherwise $r_1=0$; $r_2=0$ if reaching the goal state, otherwise $r_2=-1$; $r_3=2$ if reaching the goal state, otherwise $r_3=1$. Are $r_2$ and $r_3$ your considered dense reward functions and better than $r_1$? Why?\n2. How is the training time? DrS has to classify each collected trajectory based on the maximal stage index of all the transitions' stage indexes, which seems to be quite time consuming.\n3. Why is GAIL unable to learn reusable rewards? I understand that at convergence, the discriminator outputs 1/2 for every expert state-action transition, but correspondingly, other state-action pairs will be assigned lower values. It seems to be a reasonable dense reward function. In Figure 12, the authors compare DrS with \"GAIL w/ stage indicators\", but what if GAIL? \n4. In line 17 of Algorithm 1's pseudocode, we need to calculate the reward for each transition. But there are multiple discriminators, so which one should we select? Based on the stage index of the trajectory? If so, the same next state may get different rewards because of being in different trajectories. Will it cause the training process unstable?\n\nI am willing to raise my scores if you could sovle my concerns.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. This paper is clearly written and easy to follow.\n2. Authors conducted extensive experiments over 1000 task variants from threes task families to empirically evaluate their proposed method.\n3. Learning reusable rewards is an interesting and important topic for RL.", "weaknesses": "1. The experimental results are all from the ManiSkill benchmark. Considering other domains such as navigation (say, the Minigrid benchmark) would make the results more convincing.\n2.  Lack of discussion on some related work, such as [1, 2].\n3. The proposed method heavily relies on the stage structures of the task (see Figure 6, 1-stage fails), but there may exist tasks hard to specify stages (e.g., locomotion tasks of mujoco).\n4. Classifying the collected trajectories into success and failure ones and learning a corresponding reward function may not incentive the agent to finish the task optimally since no matter how sub-optimal the trajectory is, as long as it reaches the highest stage, it will be regarded as the most successful one.\n\n[1] Yoo, Se-Wook, and Seung-Woo Seo. \"Learning Multi-Task Transferable Rewards via Variational Inverse Reinforcement Learning.\" 2022 International Conference on Robotics and Automation (ICRA). IEEE, 2022.\n\n[2] Luo, Fan-Ming, Xingchen Cao, and Yang Yu. \"Transferable Reward Learning by Dynamics-Agnostic Discriminator Ensemble.\" arXiv preprint arXiv:2206.00238 (2022).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)