DrS: Learning Reusable Dense Rewards for Multi-Stage Tasks
Tongzhou Mu Minghua Liu Hao Su



[
    Leonie Vollmar^1,2,*, Rick Bebon^3,*, Julia Schimpf^1,2, Bastian Flietel^4, Sirin Celiksoy^4, Carsten Sönnichsen^4, Aljaž Godec^3 and Thorsten Hugel^1,5
    April 29, 2024
============================================================================================================================================================





  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 (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 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 https://sites.google.com/view/iclr24drs project page for more details.
  
  



§ INTRODUCTION




The success of many reinforcement learning (RL) techniques heavily relies on dense reward functions <cit.>, which are often tricky to design by humans due to heavy domain expertise requirements and tedious trials and errors. 
In contrast, sparse rewards, such as a binary task completion signal, are significantly easier to obtain (often directly from the environment).
For instance, in pick-and-place tasks, the sparse reward could simply be defined as the object being placed at the goal location. Nonetheless, sparse rewards also introduce challenges (e.g., exploration) for RL algorithms <cit.>. Therefore, a crucial question arises: can we learn dense reward functions in a data-driven manner?







Ideally, the learned reward will be reused to 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 morphologies. For clarity, we refer to each variant as a task and the set of all possible pick-and-place tasks as a 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 to directly transfer a policy operating a two-finger gripper for pick-and-place to a three-finger gripper due to action space misalignment, but a reward inducing the approach-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 a reward function from demonstrations is known as inverse RL in the literature <cit.>. More recently, adversarial imitation learning (AIL) approaches have been proposed <cit.> and gained popularity. Following the paradigm of GANs <cit.>, AIL approaches employ a policy network to generate trajectories and train a discriminator to distinguish between agent trajectories from demonstration ones. 
By using the discriminator score as rewards, <cit.> shows that a policy can be trained to imitate the demonstrations. 
Unfortunately, such rewards are not reusable across tasks – at convergence, the discriminator outputs 1/2 for both the agent trajectories and the demonstrations, as discussed in <cit.>, making it unable to learn useful information for solving new tasks.



In contrast to AIL, we propose a 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, we train a discriminator to classify success trajectories and failure trajectories based on the binary sparse reward. Please refer to Fig. <ref> (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 is 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 is a relatively weak requirement as it is often an inherent component of the task definition.


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. <ref>, there are three stages: 1) 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.
[Stage indicators are only required during RL training, but not required when deploying policy to real world.]
By utilizing these stage indicators, we can learn a dense reward for each stage and combine them into a 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 <cit.>, which offers a large number of task variants within each task family.We evaluate our approach on three task families: Pick-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 a 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-engineered reward functions.

Moreover, our approach drastically reduces the human effort needed for reward engineering. For instance, while the human-engineered 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 <ref> for a detailed example illustrating how our method reduces the required human effort.



Our contributions can be summarized as follows: 
*[itemize]labelindent=,itemindent=0pt,leftmargin=20pt

    
  * We propose DrS (Dense reward learning from Stages), a 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 effectiveness of our approach in generating high-quality and reusable dense rewards.









§ RELATED WORKS





Learning Reward from Demo (Offline)
Designing rewards is challenging due to domain knowledge requirements, so approaches to learning rewards from data have gained attention. 
Some methods adopt classification-based rewards, i.e., training a reward by classifying goals <cit.> or demonstration trajectories <cit.>. Other methods <cit.> 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 a wrong reward signal, as studied in <cit.>.

Learning Reward from Demo (Online)
The above issue can 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 a reward function given expert demonstrations. 
Traditional IRL methods <cit.> often require multiple iterations of Markov Decision Process solvers <cit.>, resulting in poor sample efficiency.
In recent years, adversarial imitation learning (AIL) approaches are proposed
<cit.>. 
They operate similarly to generative adversarial networks (GANs) <cit.>, in which a 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 <cit.> infer the reward through human preference queries on trajectories or explicitly asking for trajectory rankings <cit.>. Another line of works <cit.> 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-horizon tasks. In 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 <cit.> study the forms of shaped rewards that induce the same optimal policy as the ground-truth reward. 
Recently, some works <cit.> have shaped the rewards as the distance to the goal, similar to some offline reward learning methods mentioned above. 
Another idea <cit.> involves shaping delayed reward by ranking trajectories based on a fine-grained preference oracle. 
In contrast to these reward shaping approaches, our method leverages demonstrations, which are available in many real-world problems <cit.>. 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 <cit.> break down policies into sub-policies to solve specific sub-tasks. Skill chaining methods <cit.> focus on solving 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 <cit.>. In contrast to these approaches that utilize stage structures in policy space, our work explores an orthogonal direction by designing rewards with stage structures. 




§ PROBLEM SETUP





In this work, we adopt the Markov Decision Process (MDP) ℳ:=⟨ S,A,T,R,γ⟩ as the theoretical framework, where R is a reward function that defines the goal or purpose of a task. Specifically, we focus on tasks with sparse rewards. In this context, "sparse reward" denotes a binary reward function that gives a value of 1 upon successful task completion and 0 otherwise:




    R_sparse(s) = 
         1    task is completed by reaching one of the success states  s 
     0    otherwise




Our objective is to learn a 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 task family. A task family refers to a 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. 
A stage indicator is a binary function that indicates whether the current state corresponds to a specific stage of the task. An example of stage indicators is in Fig. <ref>. 
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 as a semi-sparse reward, and it serves as a strong baseline:

    R_semi-sparse(s) = k , when state s is at stage k


We aim to design an approach that learns a 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 is also possible to obtain them in the real world by various techniques (robot proprioception, tactile sensors <cit.>, visual detection/tracking <cit.>, large vision-language models <cit.>, etc.).





§ DRS: DENSE REWARD LEARNING FROM STAGES



Dense rewards are often tricky to design by humans (see an example in appendix <ref>), so
we aim to learn a 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. <ref> (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 a regular RL training process, we only discuss the reward learning phase in this section. 
We first explain how our approach learns a dense reward in one-stage tasks (Sec. <ref>). Then, we extend this approach to multi-stage tasks (Sec. <ref>). 



 §.§ Reward Learning on One-Stage Tasks



In line with previous work <cit.>, we employ a classification-based 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 (AIL) methods 
<cit.> 
used discriminators as classifiers/rewards 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/2, as observed in GANs <cit.> (also noted by <cit.>). This makes it unable to learn useful info for solving new tasks. 





Our approach introduces a simple modification to existing AIL methods to ensure that the discriminator continues to learn meaningful information even at 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 RL agent. Intuitively, a 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. <ref>(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-policy manner. While the original GAIL is on-policy, recent AIL methods <cit.> 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.





r0.33
    

    
    < g r a p h i c s >

    

    An illustration of our learned reward, which fills the gaps in semi-sparse rewards, resulting in a smooth reward curve.
    
    






 §.§ Reward Learning on Multi-Stage Tasks




In multi-stage tasks, it is desirable for the reward of a 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. <ref>) 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 as 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, a 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. <ref> gives an intuitive illustration of our learned reward, which fills the gaps in semi-sparse rewards, resulting in a 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(τ: (s_0, s_1, ...))=max_i  StageIndex(s_i),

where τ is a trajectory and s_i are the states in τ. For the discriminator associated with stage k, positive data consists of trajectories that progress beyond stage k (StageIndex >k), and negative data consists of trajectories that reach up to stage k (StageIndex ≤ k).


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. <ref>) 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(s') = k + α·tanh(Discriminator_k(s'))

where k is the stage index of s' and α is a 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 α<1/2 ensures that the reward of a state in stage k+1 is always higher than that of stage k. In practice, we use α=1/3 and it works well. 













R0.6
0.6













 §.§ Implementation


From the implementation perspective, our approach is similar to GAIL, but with a different training process for discriminators. While the original GAIL is combined with TRPO <cit.>, <cit.> found that using state-of-the-art off-policy RL algorithms (like SAC <cit.> or TD3 <cit.>) can greatly improve the sample efficiency of GAIL. Therefore, we also combine our approach with SAC, and the full algorithm is summarized in Algo. <ref>.



In addition to the regular replay buffer used in SAC, our approach maintains N different stage buffers to store trajectories corresponding to different stages(defined by Eq. <ref>). Each trajectory 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  <cit.>. However, our approach is also compatible with alternative forms of input, such as (s,a) or (s, a, s'). 





§ EXPERIMENTS





 §.§ Setup and Task Descriptions





We evaluated our approach on three challenging physical manipulation task families from the ManiSkill <cit.>: Pick-and-Place, Turn Faucet, and Open Cabinet Door. 
Each task family includes a set of different objects to 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.<ref>.
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 we 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, it is crucial to have a diverse set of tasks that exhibit similar structures and goals but possess variations in other aspects. However, most existing benchmarks lack an adequate number of task variations within the same task family. As a result, we primarily conducted our evaluation on the ManiSkill benchmark, which offers a range of object variations within each task family. 
This allowed us to thoroughly evaluate our learned rewards in a 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 is completed if the object is 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 is 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,600 EGAD objects.

Turn Faucet: A robot arm is tasked to turn on a faucet by rotating its handle. The task is completed if the handle reaches a target angle. The stage indicators include: (a) the target handle starts moving, (b) the handle reaches a 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 a designated target door on a cabinet. The task is completed if the target door is opened to a 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 4 cabinet doors and reuse rewards on 6 cabinet doors. Note that we remove all single-door cabinets in this task family, as they can 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 <ref> for a detailed description of the object sets, action space, state space, and demonstration trajectories.





 §.§ Baselines





Human-Engineered  The original human-written dense rewards in the benchmark, which require a 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, as discussed in Eq. <ref>. The agent receives a reward of k when it is in stage k. This baseline extends the binary sparse reward.


VICE-RAQ <cit.> An improved version of VICE <cit.>. It learns a 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 <cit.> A representative offline reward learning method, where the agent does not interact with the environments but purely learns from the demonstrations.
It learns a classifier (reward) to distinguish between the states from success trajectories and random trajectories.




 §.§ 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. <ref>, 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 <cit.>. 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.
Unsurprisingly, 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 to overfit the provided dataset.
When using such rewards in RL, the flaws in the learned rewards are easily exploited by the RL agents.








 §.§ Ablation Study



We examined various design choices within our approach on the Pick-and-Place task family.



  §.§.§ Robustness to Stage Configurations


Though many tasks present a natural structure of stages, there are still different ways to divide a 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 1-stage case. 
Our results, presented in Fig. <ref>, indicate that the learned rewards with 2 stages can still effectively train RL agents in test tasks, albeit with lower sample efficiency than those with 3 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.



  
Definition of Stages 
The stage indicator "object is placed" is initially defined as 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 5cm and 10cm, respectively. The results, as depicted in Fig. <ref>, demonstrate that changing the distance threshold within a 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 as it yields the highest reward according to Eq. <ref>. 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.





  §.§.§ 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, our approach also produces a policy as a byproduct in the reward learning phase. This policy can also be fine-tuned using various rewards in new tasks, providing an alternative to training RL 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. <ref>, all policies improve rapidly at the beginning due to 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 a 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 RL agents, even with a very good initialization.




  §.§.§ Additional Ablation Studies

Additional ablation studies are provided in appendix <ref>, with key conclusions summarized as follows:



    
  * DrS is compatible with various modalities of reward input, including point cloud data. <ref> 
    
  * Reward learned by GAIL, even with stage indicators, is not reusable. <ref>
    
  * The way of combining the dense rewards from each stage matters. <ref>











§ CONCLUSION AND LIMITATIONS




To make RL a more widely applicable tool, we have developed a 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 high-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, not needed at policy deployment time), we 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 <cit.> (by our testing, they suggest stages highly aligned to the ones we adopt by intuition for all tasks in this work) or employing information-theoretic approaches. Further discussions regarding this point can be found in appendix <ref>.

Secondly, the reliance on stage indicators adds a level of inconvenience when directly training RL agents in 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 <cit.>. 
For example, the “object is grasped” indicator can be acquired by tactile sensors <cit.>, and the “object is placed” indicator can be obtained by forward kinematics, visual detection/tracking techniques <cit.>, or even large vision-language models <cit.>.














iclr2024_conference









§ TASK DESCRIPTIONS




For all tasks, we use consistent setups for state spaces, action spaces, and demonstrations. The state spaces adhere to a standardized template that includes proprioceptive robot state information, such as joint angles and velocities of the robot arm, and, if applicable, the mobile base. Additionally, task-specific goal information is included within the state. Please refer to the ManiSkill paper <cit.> for more details. Below, we present the key details pertaining to the tasks used in this paper.




 §.§ Pick-and-Place


    
  * Stage Indicators:
        
            
  * Object is grasped: Both of the robot fingers contact the object, and the impulse (force) at the contact points is non-zero.
            
  * Object is placed: The distance between the object and the goal position is less than 2.5 cm. This is given by the success signal of the original task, not designed by us.
            
  * Robot and object are stationary: The joint velocities of all robot joints are less than 0.2 rad/s. The object velocity is less than 3 cm/s. This is given by the success signal of the original task, not designed by us.
        
    
  * Object Set: The objects in training tasks are from the YCB dataset, including 74 objects. And the objects in test tasks are from the EGAD dataset, including around 1600 objects.
    
  * Action Space: Delta position of the end-effector and the joint positions of the gripper. 
    
  * Demonstrations: We use 100 demonstration trajectories in total for this task family (around 1.4 trajectories per task). The demonstrations are from a trained RL agent.





 §.§ Turn Faucet


    
  * Stage Indicators:
        
            
  * Handle is moving: The joint velocity of the target joint is greater than 0.01 rad/s.
            
  * Handle reached the target angle: The joint angle is greater than 90% of the limit. This is given by the success signal of the original task, not designed by us.
        
    
  * Object Set: The objects in training and test tasks are both from the PartNet-Mobility dataset. The training tasks include 10 faucets, and the test tasks include 50 faucets.
    
  * Action Space: Delta pose of the end-effector and joint positions of the gripper. 
    
  * Demonstrations: We use 100 demonstration trajectories in total for this task family (around 10 trajectories per task). The demonstrations are from a trained RL agent.




 §.§ Open Cabinet Door


    
  * Stage Indicators:
        
            
  * Handle is grasped: Both of the robot fingers contact the handle, and the impulse (force) at the contact points is non-zero.
            
  * Door is open enough: The joint angle is greater than 90% of the limit. This is given by the success signal of the original task, not designed by us.
            
  * Door is stationary: The velocity of the door is less than 0.1 m/s, and the angular velocity is less than 1 rad/s. This is given by the success signal of the original task, not designed by us.
        
    
  * Object Set: The objects in training and test tasks are both from the PartNet-Mobility dataset. The training tasks include 4 cabinet doors, and the test tasks include 6 cabinet doors. We remove all single-door cabinets in this task family, as they can be solved by kicking the side of the door and this behavior can be readily learned by sparse rewards.
    
  * Action Space: Joint velocities of the robot arm joints and mobile robot base, and joint positions of the gripper.
    
  * Demonstrations: We use 200 demonstration trajectories in total for this task family (around 50 trajectories per task). The demonstrations are from a trained RL agent.




§ COMPARISON OF HUMAN EFFORT: STAGE INDICATORS VS. HUMAN-ENGINEERED REWARDS 



This section explains why designing stage indicators is much easier than designing a full dense reward.

The key challenges in reward engineering lies in designing reward candidate terms and tuning associated hyperparameters. To illustrate, let us use the “Open Cabinet Door” task familly as an example. The code of human engineered reward is in Listing <ref>, and the code of our stage indicators is in Listing <ref>.

The human-engineered reward involves the following reward candidate terms:

    
  * Distance between the robot gripper quaternion, and a set of manually designed grasp quaternions
    
  * Distance between robot hand and door handle
    
  * Signed-distance between tool center point (center of two fingertips) and door handle
    
  * Robot joint velocity
    
  * Door handle velocity
    
  * Door handle angular velocity
    
  * Door joint velocity
    
  * Door joint position
    
  * Multiple boolean functions to determine task stages


Each reward candidate term needs 1∼4 hyperparameters (e.g., normalization function, clip upper bound, clip lower bound, scaling coefficient). In total, this reward function involves more than 20 hyperparameters to tune. The major effort of reward engineering is thus spent iterating over these candidate terms and tuning the hyperparameters by trail and error. This process is laborious but critical for the success of human-engineered rewards. According to the authors of ManiSkill, they spend over one month crafting the dense reward for the “Open Cabinet Door” tasks. 


In contrast, our stage indicators for “Open Cabinet Door” tasks only requires to design two boolean functions: whether the robot has grasped the handle and whether the door is open enough. 
The third stage indicator is given by the tasks success signal so we do not need to design it.
This trims the number of hyperparameters down from 20+ to just  1 (the first boolean function requires one hyperparameter, and the second boolean function is directly taken from the task’s success condition so no hyperparamters), and reduces the lines of code from 100+ to 7 (with a utility function to check grasping, which is from the original codebase).


Therefore, our approach significantly reduces the human effort required for reward engineering.



[label=cabinet_human_reward, language=Python, caption=Human-engineered rewards for Open Cabinet Door tasks. The code is from the ManiSkill2 github repo (commit id: ). ]
    def _compute_grasp_poses(self, mesh: trimesh.Trimesh, pose: sapien.Pose):
        # NOTE(jigu): only for axis-aligned horizontal and vertical cases
        mesh2: trimesh.Trimesh = mesh.copy()
        # Assume the cabinet is axis-aligned canonically
        mesh2.apply_transform(pose.to_transformation_matrix())

        extents = mesh2.extents
        if extents[1] > extents[2]:  # horizontal handle
            closing = np.array([0, 0, 1])
        else:  # vertical handle
            closing = np.array([0, 1, 0])

        # Only rotation of grasp poses are used. Thus, center is dummy.
        approaching = [1, 0, 0]
        grasp_poses = [
            self.agent.build_grasp_pose(approaching, closing, [0, 0, 0]),
            self.agent.build_grasp_pose(approaching, -closing, [0, 0, 0]),
        ]

        pose_inv = pose.inv()
        grasp_poses = [pose_inv * x for x in grasp_poses]

        return grasp_poses

    def _compute_handles_grasp_poses(self):
        self.target_handles_grasp_poses = []
        for i in range(len(self.target_handles)):
            link = self.target_links[i]
            mesh = self.target_handles_mesh[i]
            grasp_poses = self._compute_grasp_poses(mesh, link.pose)
            self.target_handles_grasp_poses.append(grasp_poses)

    def compute_dense_reward(self, *args, info: dict, **kwargs):
        reward = 0.0

        # —————————————————– #
        # The end-effector should be close to the target pose
        # —————————————————– #
        handle_pose = self.target_link.pose
        ee_pose = self.agent.hand.pose

        # Position
        ee_coords = self.agent.get_ee_coords_sample()  # [2, 10, 3]
        handle_pcd = transform_points(
            handle_pose.to_transformation_matrix(), self.target_handle_pcd
        )
        # trimesh.PointCloud(handle_pcd).show()
        disp_ee_to_handle = sdist.cdist(ee_coords.reshape(-1, 3), handle_pcd)
        dist_ee_to_handle = disp_ee_to_handle.reshape(2, -1).min(-1)  # [2]
        reward_ee_to_handle = -dist_ee_to_handle.mean() * 2
        reward += reward_ee_to_handle

        # Encourage grasping the handle
        ee_center_at_world = ee_coords.mean(0)  # [10, 3]
        ee_center_at_handle = transform_points(
            handle_pose.inv().to_transformation_matrix(), ee_center_at_world
        )
        # self.ee_center_at_handle = ee_center_at_handle
        dist_ee_center_to_handle = self.target_handle_sdf.signed_distance(
            ee_center_at_handle
        )
        # print("SDF", dist_ee_center_to_handle)
        dist_ee_center_to_handle = dist_ee_center_to_handle.max()
        reward_ee_center_to_handle = (
            clip_and_normalize(dist_ee_center_to_handle, -0.01, 4e-3) - 1
        )
        reward += reward_ee_center_to_handle

        # pointer = trimesh.creation.icosphere(radius=0.02, color=(1, 0, 0))
        # trimesh.Scene([self.target_handle_mesh, trimesh.PointCloud(ee_center_at_handle)]).show()

        # Rotation
        target_grasp_poses = self.target_handles_grasp_poses[self.target_link_idx]
        target_grasp_poses = [handle_pose * x for x in target_grasp_poses]
        angles_ee_to_grasp_poses = [
            angle_distance(ee_pose, x) for x in target_grasp_poses
        ]
        ee_rot_reward = -min(angles_ee_to_grasp_poses) / np.pi * 3
        reward += ee_rot_reward

        # ————————————————- #
        # Stage reward
        # ————————————————- #
        coeff_qvel = 1.5  # joint velocity
        coeff_qpos = 0.5  # joint position distance
        stage_reward = -5 - (coeff_qvel + coeff_qpos)
        # Legacy version also abstract coeff_qvel + coeff_qpos.

        link_qpos = info["link_qpos"]
        link_qvel = self.link_qvel
        link_vel_norm = info["link_vel_norm"]
        link_ang_vel_norm = info["link_ang_vel_norm"]

        ee_close_to_handle = (
            dist_ee_to_handle.max() <= 0.01 and dist_ee_center_to_handle > 0
        )
        if ee_close_to_handle:
            stage_reward += 0.5

            # Distance between current and target joint positions
            # TODO(jigu): the lower bound 0 is problematic? should we use lower bound of joint limits?
            reward_qpos = (
                clip_and_normalize(link_qpos, 0, self.target_qpos) * coeff_qpos
            )
            reward += reward_qpos

            if not info["open_enough"]:
                # Encourage positive joint velocity to increase joint position
                reward_qvel = clip_and_normalize(link_qvel, -0.1, 0.5) * coeff_qvel
                reward += reward_qvel
            else:
                # Add coeff_qvel for smooth transition of stagess
                stage_reward += 2 + coeff_qvel
                reward_static = -(link_vel_norm + link_ang_vel_norm * 0.5)
                reward += reward_static

                # Legacy version uses static from info, which is incompatible with MPC.
                # if info["cabinet_static"]:
                if link_vel_norm <= 0.1 and link_ang_vel_norm <= 1:
                    stage_reward += 1

        # Update info
        info.update(ee_close_to_handle=ee_close_to_handle, stage_reward=stage_reward)

        reward += stage_reward
        return reward



[label=cabinet_stage_indicators, language=Python, caption=Our stage indiactors for Open Cabinet Door tasks  which is way more easier to design than the human-engineered rewards. ]
def compute_stage_indicators(self):
    stage_indicators = [
        self.agent.check_grasp(self.target_link), # this utility function is given by the original ManiSkill2 codebase. It requires one hyperparameter `max_angle` but we just use the default value
        self.link_qpos >= self.target_qpos, # door is open enough
        # the 3rd stage indicator is just the task success signal, so we don't need to include it here
    ]
    for i in range(1, len(stage_indicators)):
        stage_indicators[i-1] |= stage_indicators[i]
    return stage_indicators





§ COMPARISON WITH TEXT2REWARD



Text2Reward <cit.> is a concurrent work with our paper. We offer a comparison in this section to help readers understand the differences between our paper and Text2Reward.

 

While both <cit.> and our paper share the common goal of generating rewards for new tasks, they employ fundamentally distinct setups and methodologies. In short, the primary distinction lies in the fact that our approach learns rewards from training tasks and success signals (or stage indicators), while <cit.> generates rewards based on exemplar reward codes and the knowledge embedded in Large Language Models (LLMs).

To elaborate, the following disparities exist in respective setups and assumptions:


    
  * Both <cit.> and our methods need to interact with environments. However, we emphasize more on evaluating the learned rewards on unseen test tasks.
    
  * <cit.> assumes access to a pool of instruction-reward code pairs, while our method requires training on relevant training tasks instead.
    
  * <cit.> assumes access to the source code of the tasks, allowing them to provide LLMs with a Pythonic environment abstraction and various utility functions. In contrast, our method solely relies on success signals (or stage indicators) and does not require the code of the tasks.
















§ IMPLEMENTATION DETAILS





 §.§ Reward Learning Phase



  §.§.§ Network Architectures


    
  * Actor Network: 4-layer MLP, hidden units (256, 256, 256)
    
  * Critic Networks: 4-layer MLP, hidden units (256, 256, 256)
    
  * Discriminator Networks (Reward): 2-layer MLP, hidden units (32)




  §.§.§ Hyperparameters

We use SAC <cit.> as the backbone RL algorithm in the reward learning phase of DrS. The related hyperparameters are listed in Table <ref>.





 §.§ Reward Reuse Phase



  §.§.§ Network Architectures


    
  * Actor Network: 4-layer MLP, hidden units (256, 256, 256)
    
  * Critic Networks: 4-layer MLP, hidden units (256, 256, 256)




  §.§.§ Hyperparameters

During the reward reuse phase, we use different rewards to train agents by SAC <cit.>. The related hyperparameters are listed in Table <ref>.










§ ADDITIONAL ABLATION STUDY



In this section, we present more ablation studies that are not included in the main paper due to the space limit. These experiments are conducted on the Pick-and-Place task family. 












 §.§ Modality of the Inputs to the Rewards







Our approach is able to accommodate various input modalities for the reward functions, including both low-dimensional state vectors and high-dimensional visual inputs. To demonstrate this compatibility, we conducted an additional experiment using point cloud inputs. In this experiment, the reward function (discriminator) not only considers the low-dimensional state but also takes a point cloud as input, with the point cloud being processed by a PointNet. The results of this experiment are depicted in Fig. <ref>.

We can see that the reward with point cloud input performs similarly to the one with state input, which shows that our approach is perfectly compatible with high-dimensional visual inputs. However, the techniques about visual inputs are a bit orthogonal to our focus (reward learning), and learning with visual inputs takes significantly more time, so we still keep most of our experiments on state inputs. 

The results reveal that the reward function utilizing point cloud input performs comparably to the one utilizing state input, demonstrating the seamless integration of our approach with high-dimensional visual inputs. However, it is worth noting that the techniques about visual inputs, while compatible with our framework, are a little bit orthogonal to our focus (reward learning). Moreover, learning with visual inputs typically takes a significantly longer training time. Consequently, the majority of our experiments primarily use state inputs, allowing us to concentrate on the core aspects of reward learning.







 §.§ Discriminator Modification and Stage Indicators








In contrast to GAIL <cit.>, our approach incorporates two critical modifications in the training of discriminators to facilitate the learning of reusable dense rewards. These modifications entail: (a) replacing the agent-demonstration discriminator with the success-failure discriminator, and (b) employing stage indicators by utilizing a separate discriminator for each stage. To ascertain the significance of these modifications, we devised two ablation baselines:


    
  * GAIL w/ Stage Indicators: 
    This baseline serves as an equivalent representation of our method without the incorporation of the success-failure discriminator. In GAIL, the discriminator solely distinguishes between agent and expert trajectories, making it incapable of learning separate rewards for each stage. To incorporate the stage indicators within the GAIL framework, we first train the original GAIL on the training tasks. During the reward reuse phase, we linearly combine the GAIL reward with the semi-sparse reward, thus leveraging the stage information. Through experimentation, we explored different weightings to strike an optimal balance between these two reward components.
    
  * Ours w/o Stage Indicators: 
    In this baseline, we exclude the stage indicators and solely rely on the task completion signal to train the discriminator. This approach is equivalent to the one-stage reward learning discussed in Sec. <ref>.



Fig. <ref> illustrates the comparison between the two ablation baselines and our method during both the reward learning phase and reward reuse phase. While both "GAIL w/ Stage Indicators" and "Ours w/o Stage Indicators" demonstrate similar success rates as our method at the conclusion of the reward learning phase, it is crucial to emphasize that the learned rewards from both ablation baselines fail to be reused to the test tasks. In contrast, our method achieves the acquisition of high-quality reward functions capable of effectively training new RL agents in the test tasks. This outcome substantiates the indispensability of the two proposed components in facilitating the acquisition of reusable dense rewards.




 §.§ Reward Formulation






In our approach, we leverage the stage indicators and define the reward function as the sum of the semi-sparse reward and the discriminator's bounded prediction for the current stage, as expressed in Eq. <ref>. This formulation ensures that the reward strictly increases across stages. To evaluate the effectiveness of this formulation, we compare it with a straightforward variant, denoted as ∑_k tanh(Discriminator_k(s')), which sums up the discriminator predictions for all stages. As depicted in Fig. <ref>, the simple variant exhibits significantly poorer performance, underscoring the importance of focusing on the dense reward specific to the current stage.










§ AUTOMATICALLY GENERATING STAGE INDICATORS



This section discusses a few promising solutions to automatically generate stage indicators, drawing inspiration from some recent publications.
Though this topic is a little bit beyond the scope of our paper, we believe this is a valuable discussion for the readers.



 §.§ Employ LLMs for Code Generation of Stage Indicators


Beyond task decomposition, LLMs demonstrate the capability to directly write code <cit.> for robotic tasks. A recent study <cit.> exemplified how LLMs, when prompted with the appropriate APIs, can generate success conditions (code snippets) for each subtask. Given the swift advancements in the domain of large models, it is entirely feasible to generate both stage structures and stage indicators using them.



 §.§ Infer Stages via Keyframe Discovery

The boundaries between stages can be viewed as keyframes in the trajectories. A recent approach introduced by <cit.> suggests the automated extraction of such keyframes from trajectories, leveraging reconstruction errors. Given these keyframes, one intuitive solution is to develop a keyframe classifier that can act as a stage indicator. However, this requires a certain degree of consistency across keyframes, and we believe it is an interesting direction to explore.






§ ADDITIONAL EXPERIMENTS ON OTHER DOMAINS




 §.§ Navigation




  §.§.§ Introduction


In this section, we incorporated experiments on navigation tasks, which were conducted during the initial stages of our project. We do not include these results in the main paper, as we found these simple navigation tasks to be less interesting compared to the robot manipulation tasks.



  §.§.§ Setup






  
Task Description We have developed a 2D navigation task conceptually similar to MiniGrid, as visually represented in Fig. <ref>. The maps are 17x17, where the agent is randomly placed in the bottom room and needs to navigate to the star, randomly located in the top room.



  
Observation Observations provided to the agent include its xy coordinates, the xy coordinates of the goal, and a 3x3 patch around itself.



  
Action The agent has a choice of 5 actions: moving up, down, left, right, or remaining stationary.



  
Training and Test Set The reward is learned on the map shown in <ref> and then reused on the map in <ref>. The difference between these two maps lies in the positions of two gates.






  §.§.§ Results




Our method is also effective in learning reusable rewards for navigation tasks. Given the relative simplicity of this specific navigation task, our approach's one-stage version suffices, eliminating the need for additional stage information. The results for this experiment are shown in Fig. <ref>.



The results clearly demonstrate that the learned reward from our approach successfully guides the RL agent to complete the task perfectly. In contrast, RL agents with sparse rewards show poor performance. Note that the map used in the test task differs from the training one, so directly transferring policy would not work. We also visualize the learned reward in Fig. <ref>. See the caption for a detailed analysis.







 §.§ Locomotion




  §.§.§ Introduction



While it can be tricky to divide locomotion tasks into stages, our method (specifically, the one-stage version) is capable of effectively handling such tasks, if they have a short horizon. In this section, we demonstrate that our approach can learn reusable rewards for Half Cheetah, a representative locomotion task in MuJoCo.

For tasks that are long-horizon and hard to specify stages, such as the https://robotics.farama.org/envs/maze/ant_maze/Ant Maze, crafting rewards is very challenging even for experienced human experts. Therefore, we leave these tasks for future work.




  §.§.§ Setup




  
Task Description Our experiment uses  from Gymnasium. The  task has a predefined reward threshold of 4800, as specified in https://github.com/Farama-Foundation/Gymnasium/blob/main/gymnasium/envs/__init__.py#L277their code, which is used to gauge task completion according to https://gymnasium.farama.org/api/registry/their documentation. Thus, we define the sparse reward (success signal) for this task as achieving an accumulative dense reward greater than 4800.




  
Training and Test Set In the reward learning phase, we use the standard  task. In the reward reuse phase, we modify the task by increasing the damping of the front leg joints (thigh, shin, and foot joints) by 1.5 times. This increased damping makes it more challenging for the cheetah to achieve high speeds.




  §.§.§ Results




Our method has successfully demonstrated its capability to learn reusable rewards in the Half Cheetah task. The results are illustrated in Fig. <ref>. Notably, the performance achieved using the learned reward is comparable to that of the human-engineered reward, while the sparse reward proved ineffective in training an RL agent. Given that many locomotion tasks emphasize low-level control and are typically of a shorter horizon, our approach's one-stage version proves to be highly effective. Additionally, this version does not require any stage information, further underscoring its efficiency and adaptability in handling such tasks.









§ DISCUSSION ON THE DESIRED PROPERTIES OF DENSE REWARDS







 §.§ Overview

Our paper primarily focuses on learning a dense reward, so one important question we want to discuss is: What kind of dense reward do we aspire to learn?

It is somewhat challenging to strictly distinguish dense rewards from sparse rewards, due to the lack of strict definitions of dense rewards in the existing literature (to the best of our knowledge). However, this does not preclude a meaningful discussion about the desired properties of dense rewards. Unlike sparse rewards, which typically only provide reward signals when the task is solved, dense rewards offer more frequent and immediate feedback regarding the agent's actions. 

We posit that the fundamental property of an effective dense reward is its capacity to enhance the sample efficiency of RL algorithms. The rationale behind this property is straightforward: a well-structured dense reward should reduce the need for extensive exploration during RL training. By providing direct guidance and immediate feedback, the agent can quickly discover optimal actions, thereby accelerating the learning process.

In line with this philosophy, an ideal dense reward should allow the derivation of optimal policies with minimal effort. By analyzing a simple tabular case, we find that our learned reward exhibits this great property. To be more specific, in the example below, we can obtain the optimal policy by greedily following the path of maximum reward at each step.



 §.§ Analysis on a Simple Tabular Case


Under certain assumptions, we can obtain the optimal policy by greedily following the path of maximum reward at each step, i.e.,

    π^*(s)=_a R^†(s,a)

, where π^* is the optimal policy and R^† is the learned reward.



  §.§.§ Setup and Assumptions

In this analysis, we consider a MDP with the following assumptions:

    
  * Deterministic transitions: s'=P(s,a)
    
  * Discrete and finite state/action space: S={s_0,s_1,...}, A={a_0,a_1,...} 
    
  * Given sparse reward: R(s,a,s')=1 if s'=s_goal, otherwise 0
    
  * Discount factor: γ<1


Other assumptions about our approach:

    
  * Only one stage, so the one-stage version of our approach is applied.
    
  * The buffers for success trajectories and failure trajectories are large enough, but not infinite.
    
  * After training for a sufficiently long time, policy converges to the optimal policy π^*. (This is a strong assumption, but it is possible in theory.)




  §.§.§ Notations


    
  * Learned reward: R^†(s,a)=tanh(Discriminator(s,a)), so R^†(s,a) ∈ (-1,1)
    
  * Buffer for success trajectories ℬ^+, buffer for failure trajectories ℬ^-
    
  * Optimal policy: π^*(a|s), which represents the probability of choosing action a at state s. Here we overload the notation π^* to capture the potential multi-modal output of the policy.




  §.§.§ Connection between Optimal Policy and Learned Reward


Here, we want to demonstrate that the learned reward of an optimal action is always higher than that of any non-optimal action in each state. If this holds, it then becomes feasible to straightforwardly identify the optimal action at each state by adopting a greedy strategy that selects the action yielding the highest reward.

When γ<1, π^* will go to s_goal by the shortest paths, so π^*(a|s)=1/k_s or 0, where k_s is the number of optimal actions at s.

∀ s, there are two kinds of actions a^+ and a^-


    
  * π^*(a^+|s)>0, which means a^+ is one of the optimal actions. 
   Then (s,a^+) must be in ℬ^+, possibly be in ℬ^-.
   Therefore, R^†(s,a^+)> -1 + ϵ, when the discriminator converges. This is because the buffers are finite-size, (s,a^+) will be sampled into positive training data of the discriminator with a probability larger than 0.
   
  * π^*(a^-|s)=0, which means a^- is NOT one of the optimal actions.
   Then (s,a^-) will only be in ℬ^-, and will NOT be in ℬ^+
   Therefore, R^†(s,a^-)→ -1, when the discriminator converges. This is because (s,a^-) will only show in the negative training data of the discriminator.


Therefore, we have R^†(s,a^+) > R^†(s,a^-) for all states s. By employing a greedy strategy that selects _a R^†(s,a), we can reach the goal states in the same way as how the optimal policy π^* reaches the goal.




 §.§ Further Discussions


This subsection is dedicated to addressing additional questions the readers may raise after reading the above analysis.



  §.§.§ Does the above conclusion generalize to more complicated cases?


Although our analysis highlights a desirable property of the learned reward in a simple tabular case, this finding should not be hastily generalized to more complex cases, such as the robotic manipulation tasks used in our paper. This caution is due to two primary reasons:


    
  * In environments where the state and action spaces are continuous, the ability of the neural network to interpolate plays a significant role in shaping the final learned reward.
    
  * Practically, achieving convergence for both the policy and the discriminator can be a very time-consuming process.





  §.§.§ The Necessity of Learned Reward Despite Its Similarity to Policy



The learned reward might appear redundant at first glance, as it seems to convey the same information as the learned policy. This observation raises a potential question: why is there a need for a learned reward if we already have a learned policy? Couldn't we just utilize the learned policy directly?

The answer lies in the distinct advantages that the learned reward offers, particularly when adapting to new tasks. When the environment dynamics change, a new policy can be effectively retrained using the learned reward in conjunction with the new environmental dynamics. Directly transferring the policy, or fine-tuning it with a sparse reward, can be less efficient in certain situations. 
For a practical illustration of this concept, refer to Fig. <ref> and Sec, <ref>. These sections provide a compelling example where the transfer of rewards demonstrates success, in contrast to the less effective transfer of policies.





