Multi Task Inverse Reinforcement Learning for Common Sense Reward

## Abstract

Abstract One of the challenges in applying reinforcement learning in a complex real-world environment lies in providing the agent with a sufficiently detailed reward function. Any misalignment between the reward and the desired behavior can result in unwanted outcomes. This may lead to issues like “reward hacking” where the agent maximizes rewards by unintended behavior. In this work, we propose to disentangle the reward into two distinct parts. A simple task-specific reward, outlining the particulars of the task at hand, and an unknown common-sense reward, indicating the expected behavior of the agent within the environment. We then explore how this common-sense reward can be learned from expert demonstrations. We first show that inverse reinforcement learning, even when it succeeds in training an agent, does not learn a useful reward function. That is, training a new agent with the learned reward does not impair the desired behaviors. We then demonstrate that this problem can be solved by training simultaneously on multiple tasks. That is, multi-task inverse reinforcement learning can be applied to learn a useful reward function.

## 1 Introduction

Reinforcement Learning (RL) is a machine learning paradigm where an agent learns to make decisions by interacting with an environment (Sutton & Barto, ). The agent seeks to maximize a cumulative reward signal received in response to its actions. By maximizing this objective, the agent learns a policy that dictates its behavior within the environment. However, in many real-world applications, one needs to design the reward function to precisely define the target behavior. In these scenarios, designing a suitable reward function becomes a key challenge practitioners must address in order to train an agent effectively. Since the agent aims to maximize the reward, any misalignment between the reward and the agent’s intended actions can lead to undesirable behaviors. In the extreme case, this misalignment can lead to “reward hacking” (Skalse et al., ) where the agent successfully maximizes the reward without achieving the desired goal.
For instance, in (Clark & Amodei, ), the authors described a scenario where an agent, trained on the CoastRunners boat racing game learned unwanted behavior. The agent repeatedly knocked out targets in an infinite loop to maximize rewards without ever completing the race.
Hence, a crucial element in RL is the design of an effective reward function. This is essential to ensure that agents, trained to maximize this reward, learn the desired behavior, especially when designing agents for operation within complex real-world environments.

To address the reward design problem we first argue that it is natural to split the reward function into two components. One is a simple task-specific reward that defines the goal the agent aims to accomplish. The second is a task-agnostic reward that describes how the agent should behave in the environment while achieving this goal. We refer to this task-agnostic reward as the *common-sense* reward, cs-reward for short, as it represents “the basic level of practical knowledge and judgment that we all need to help us live in a reasonable and safe way”, using the Cambridge dictionary definition of common sense.

Furthermore, we argue that the tasks-specific reward should be relatively simple to design. For example, consider a scenario where a household robot is assigned chores like throwing garbage and mopping floors. Crafting a reward using computer vision models to verify goal completion is not a significant challenge. However, the cs-reward should be much more complex as it needs to account for a wide variety of cases the agent might encounter. Beyond completing specific tasks, the robot must safely navigate spaces, handle delicate objects with care, conserve electricity, and perform other considerations, such as closing cupboard doors.

Based on this observation, we assume that the reward function can be separated into the task-specific reward, which we assume to be known, and the common-sense reward, which is unknown. We then try to learn the shared common-sense reward from expert demonstrations. A natural approach is to use Inverse Reinforcement Learning (IRL, (Arora & Doshi, ) where an agent is trained to imitate the behavior of an expert by simultaneously learning a reward and an agent that tries to maximize said reward. Unfortunately, we show empirically that even when the IRL produces an agent with the desired behavior, it does not learn a meaningful reward. In other words, when we try to train a new agent from scratch using the learned reward, we don’t achieve the desired behavior.

Figure: Figure 1: MT-CSIRL architecture overview
Refer to caption: extracted/5411392/figures/mtirl.png

The challenge in deriving a useful reward function, even when IRL successfully imitates the expert, highlights an important distinction between this work and prior works on IRL. Typically,
in IRL, the reward serves as a means to an end, i.e., to imitate the expert, and has close connections to the discriminator in Generative Adversarial Networks (GANs) (Ho & Ermon, ; Finn et al., ). In contrast, this work focuses on the reward itself. Our main question is “will a *new* agent trained from scratch with our cs-reward gain the designed behavior?”

We aim to address this challenge by leveraging our previous assumptions about the reward structure. Specifically, we utilize multi-task IRL to learn a task-independent shared cs-reward. We term our approach MT-CSIRL. Intuitively, this allows us to combine information from multiple different experts to avoid learning task-specific behavior and spurious correlations. This directs the learned reward to emphasize the underlying shared common-sense reward.

We experiment on the Meta-world benchmark (Yu et al., ) using two simple common-sense rewards we designed. We show that even in this simple scenario IRL fails to learn a useful reward and show the importance of multi-task learning over a variety of tasks to learn a useful and transferable reward.

To conclude, one important contribution is the disentanglement of the reward into the task-specific and task-independent components, formulating the common-sense reward. This formulation has many important advantages. First, we can easily combine information from different tasks. Second, it allows us to inject prior knowledge on what each expert is trying to achieve, instead of inferring this from observations as well. Finally, the learned cs-reward can easily be transferred to new unseen tasks. Another contribution is designing a simple MTL approach to learn the cs-reward from expert demonstrations, and evaluating it empirically. Finally, our experiments indicate that training on a wide variety of different tasks is a key component in learning the cs-reward.

Figure: Algorithm 1 MT-CSIRL

Figure: Figure 2: Visualization of train and test tasks for CS-Reward Transfer in MTL CS-Reward transfer experiment
Refer to caption: extracted/5411392/vis_exp.png

## 2 Related Work

#### Multi-task RL and IRL.

In multi-task learning, a model is trained to perform multiple tasks simultaneously while leveraging shared representations across tasks to improve overall performance and efficiency (Ruder, ; Caruana, ). Extensive research has been conducted on multi-task RL and IRL, focusing on utilizing shared properties and structures among tasks (Arora et al., ; Yang et al., ; ; Sodhani et al., ; Chen et al., ; Zhang et al., ), overcoming negative transfer (Yu et al., ; Liu et al., ; Navon et al., ) and performing rapid adaptation to novel tasks (Yu et al., ). MTL was also employed in the IRL setting by learning a policy that imitates a mixture of expert demonstrations. MT-IRL generally focuses on the meta-learning setup.
Seyed Ghasemipour et al. (); Yu et al. () propose learning a policy conditioned on a trained task context vector. Finn et al. (); Yu et al. () propose a MAML-based (Finn et al., ) approach that can adapt a trained policy to a new task with only a few gradient steps. (Xu et al., ) proposed to learn a prior over reward functions. In (; Yu et al., ) the authors learn a family of rewards by using probabilistic context variables. Differently from these approaches, this paper proposes a method to transfer desired behavior that is not task-specific, without additional adaptation.

Inverse Constrained Learning
One line of research that has strong similarities to this work is inverse constrained learning (Malik et al., ), and more specifically, multi-task inverse constrained learning (Lindner et al., ; Kim et al., ). In inverse constrained learning the goal is to learn a set of safety constraints from expert demonstrations. While these safety constraints are conceptually similar to our common-sense reward, our common-sense reward is more general. One can consider a hard constraint as a reward that is $-\infty$ for the forbidden set and zero otherwise. This cannot take into account more subtle effects such as a small negative reward for using up a resource, e.g. energy, that we want to discourage the agent from needlessly wasting, but we do not want to stop it from doing so entirely.

## 3 Background

#### A Markov Decision Process (MDP).

An MDP is a mathematical framework for modeling sequential decision-making in stochastic environments, central to reinforcement learning.
An MDP consists of a tuple $(\mathcal{S},\mathcal{A},\mathcal{T},r,\gamma)$, where $\mathcal{S}$ represents the set of states, $\mathcal{A}$ is the set of actions, $\mathcal{T}$ denotes the state transition probabilities, $r$ is a reward function, and $\gamma$ is the discount factor. Agents in an MDP interact with an environment via a policy $\pi$ that selects actions from $\mathcal{A}$ in a given state $s$ from the distribution $\pi(a|s)$. After action $a$ is selected, the agent transitions to a new state $s^{\prime}$ and receives a reward $r(s,a,s^{\prime})$. The state transition probabilities are defined as $\mathcal{T}(s^{\prime}|s,a)$, representing the probability of transitioning to state $s^{\prime}$ given that action $a$ is taken in state $s$. Importantly, the transition and rewards are Markovian, depending solely on the current state and action. The main goal in reinforcement learning is to train an agent to maximize the total future discounted rewards $\mathbb{E}_{\pi}[\sum_{t=0}^{\infty}\gamma^{t}r(s_{t},a_{t},s_{t+1})]$ in an MDP by repeated interactions with the environment.

Figure: Figure 3: Single task cs-reward: Velocity and action norm during IRL, and during RL training with the learned cs-reward. The red horizontal line represents the target value in the ground truth reward, see Eq. 8 and 9.
Refer to caption: extracted/5411392/Screen_Shot.png

### 3.1 Inverse Reinforcement Learning

An important RL scenario is imitation learning where the goal is to train an agent on an MDP without having access to the reward, but with expert demonstrations. Commonly, we assume the expert is optimal or near-optimal with regard to the unknown reward. One approach to imitation learning is Inverse Reinforcement Learning (IRL) (Arora & Doshi, ), where we imitate the expert by learning a reward and a policy maximizing it to match the expert demonstrations.

In many early IRL approaches, the policy was fully trained at every stage until convergence by using a current reward (Ng & Russell, ; Abbeel & Ng, ). However, this approach is computationally expensive as it requires training an RL agent from scratch for each reward update. Hence, this approach does not scale well to recent deep reinforcement learning, where the training process can be much more expensive. Current IRL approaches are centered more around training both the reward and the agent simultaneously, updating each one in turn. An important implication of this approach is that it makes the meaning of the learned reward much more uncertain, as there are no guarantees that an agent trained from scratch on the final learned reward will train properly. A visualization of such phenomena is shown in Figure 3. It presents a comparison between training an agent with the learned reward, during IRL process, versus training an agent from scratch using the final learned reward.

Figure: (a) Action norm CS reward
Refer to caption: extracted/5411392/cs_an.png

#### Adversarial Inverse Reinforcement Learning.

The pioneering works (Ho & Ermon, ; Finn et al., ) were the first to explore the connection of generative adversarial networks (GANs) to inverse reinforcement learning (IRL), specifically the connection between IRL reward and the GAN discriminator. Their adversarial approach has proven to be the leading method that is been used to train IRL systems (Fu et al., ; Jeon et al., ; Han et al., ). For simplicity, we will base our work on the Adversarial Inverse Reinforcement Learning (AIRL) (Fu et al., ) framework, which we found to work well in our experiments.

In AIRL, we train a generator and discriminator simultaneously where the generator is the stochastic policy which we train to fool the discriminator, while the discriminator is trained to distinguish between expert trajectories and generated trajectories. The discriminator in AIRL is formulated as:

$$ $D(s,a,s^{\prime})=\frac{\exp{f_{\theta}(s,a,s^{\prime})}}{(\exp{f_{\theta}(s,a ,s^{\prime})})+\pi(a\mid s)},$ (1) $$

where $f_{\theta}(s,a,s^{\prime})$ encapsulates the learned reward structure, and $\pi(a\mid s)$ is the policy’s action probability. This formulation leads to a reward function $r_{\theta}(s,a)$ derived as:

$$ $\displaystyle r_{\theta}(s,a,s^{\prime})$ $\displaystyle=\log D(s,a,s^{\prime})-\log(1-D(s,a,s^{\prime}))$ (2) $\displaystyle=f_{\theta}(s,a,s^{\prime})-\log\pi(a\mid s).$ (3) $$

Intuitively, we get a high reward when the discriminator is confident that $(s,a,s^{\prime})$ belongs to and expert, and a low reward when it is confident $(s,a,s^{\prime})$ belongs to the agent.

The adversarial loss in AIRL aims to optimize the discriminator to accurately distinguish between expert and agent trajectories. It is defined as:

$$ $\mathcal{L}(\theta)=-\mathbb{E}_{\mathcal{D}}[\log D_{\theta}(s,a,s^{\prime})] -\mathbb{E}_{\pi}[\log(1-D_{\theta}(s,a,s^{\prime}))],$ (4) $$

where $\mathcal{D}$ represents expert demonstrations,
and $\pi$ is the policy of the agent.

## 4 Our Approach

Our goal is to learn the task-agnostic common-sense reward from expert demonstrations by employing Multi-Task Inverse Reinforcement Learning. We assume each task loss $r_{t}(s,a,s^{\prime})$ has the following structure

$$ $r_{t}(s,a,s^{\prime})=\bar{r}_{t}(s,a,s^{\prime})+r_{CS}(s,a,s^{\prime}),$ (5) $$

and that $\bar{r}_{t}$ is known. This allows for a simple MT-IRL approach where each task policy is trained with a separate weight vector $\pi_{w_{t}}(a|s)$ and the discriminator for each task share weights, capturing the common-sense reward, via

$$ $D^{t}_{\theta}(s,a,s^{\prime})=\frac{\exp{(f_{\theta}(s,a,s^{\prime})+\bar{r}_ {t}(s,a,s^{\prime}))}}{\exp{(f_{\theta}(s,a,s^{\prime})+\bar{r}_{t}(s,a,s^{ \prime}))}+\pi_{w_{t}}(a|s)}.$ (6) $$

At each iteration, we pick a task $t$, and update the policy based on the following reward:

$$ $r_{t}(s,a,s^{\prime})=\bar{r}_{t}(s,a,s^{\prime})+f_{\theta}(s,a,s^{\prime}),$ (7) $$

where $f_{\theta}(\cdot)$ is our learned shared cs-reward. We then update the task-specific discriminator $D^{t}_{\theta}$ with the adversarial loss. Full description of our MT-CSIRL method is provided in Alg. 1.

The key aspect of our approach is the disentanglement between task-specific and task-agnostic rewards. This allows us to train on multiple environments using shared weights for all the different discriminators (see Figure 1). Intuitively, training the shared weights of the different discriminators to simultaneously distinguish expert demonstrations across various environments, avoids learning task-specific behaviors or spurious correlations. We will show empirically in Sec. 5 that training on a variety of tasks is important for learning a transferable common-sense reward function.

### 4.1 Curriculum Learning

We found in our experiments, that training directly with the discriminator in Eq. (6) returns sub-optimal performance on the main task. We hypothesis that, despite having the exact task-specific reward, training with it poses challenges due to the noise introduced by the common-sense reward. This challenge is particularly prominent in the early stages of optimization, given its random initialization.

To overcome this issue and avoid any serious degradation to the task-specific performance, we devised a simple curriculum learning approach (Bengio et al., ). Instead of updating the policy with the combined reward $r_{t}(s,a,s^{\prime})=\bar{r}_{t}(s,a,s^{\prime})+f_{\theta}(s,a,s^{\prime})$, we use $r_{t}(s,a,s^{\prime})=\bar{r}_{t}(s,a,s^{\prime})+\alpha(\bar{r}_{t}(s,a,s^{
\prime}))f_{\theta}(s,a,s^{\prime})$ with the adaptive weighting $\alpha(\bar{r}_{t}(s,a,s^{\prime}))=\frac{\bar{r}_{t}(s,a,s^{\prime})}{R_{MAX}}$ assuming $\bar{r}_{t}(s,a,s^{\prime})\in[0,R_{MAX}]$.

This allows us to gradually increase the weighting of our learned reward as the policy improves on the main task. Additionally, the common-sense reward does not have a significant effect at the beginning of training, and thus only impacts the policy after it had several update rounds. We note that this does not affect the discriminator update, which still uses Eq. 6. Also, when training a new agent with our learned reward we train from scratch using the standard aggregation $r_{t}(s,a,s^{\prime})=\bar{r}_{t}(s,a,s^{\prime})+f_{\theta}(s,a,s^{\prime})$.

**Table 1: Velocity cs-reward. Each row represents a different task for MT-CSIRL training, each column represents a task for RL training and evaluation. The Numbers represent the ratio between the agents’ velocity and its target value.**
|  | reach-wall | button-press-<br>topdown-wall | drawer-close | push-back | coffee-button |
| --- | --- | --- | --- | --- | --- |
| reach-wall | $\mathbf{0.93}$ | $0.40$ | $0.43$ | $0.40$ | $0.25$ |
| button-press-topdown-wall | $0.36$ | $\mathbf{0.82}$ | $0.50$ | $0.31$ | $0.21$ |
| drawer-close | $0.34$ | $0.48$ | $\mathbf{0.89}$ | $0.35$ | $0.24$ |
| push-back | $0.25$ | $0.25$ | $0.20$ | $\mathbf{0.88}$ | $0.22$ |
| coffee-button | $0.28$ | $0.32$ | $0.35$ | $0.38$ | $\mathbf{0.91}$ |

## 5 Experiments

For our experiments, we use the Meta-world benchmark (Yu et al., ). It provides a diverse set of robotic manipulation tasks, which share the same robot, action space, and observation space. Thus, this benchmark is suitable for evaluating the effectiveness and transferability of our inverse reinforcement learning method across various scenarios. We trained all policies using the Soft Actor-Critic (SAC) algorithm (Haarnoja et al., ).
In order to emphasize our primary goal of learning transferable rewards, we limit our experimental setup to Meta-world tasks on which SAC performed well, according to the Meta-world’s paper. We used the exact training process and SAC hyperparameters as in (Yu et al., ). We repeat the experiments five times using different random initializations and report the mean and standard deviation of the performance.

Using the Meta-world terminology, there are several distinct tasks, e.g., reach-wall and drawer-close. For each task, there are several variations with distinct targets. For example, in the drawer-close task, a different target would be a different location of the drawer. We also note that Meta-world has two versions for each task. All experiments in this paper use the second version, i.e., reach-wall is reach-wall-V2, and we omit the V2 for brevity.

The goal of our experiments is to demonstrate that IRL fails to learn a useful cs-reward and that training the cs-reward on multiple tasks is the key to learning such a reward. To show that we perform several experiments, each time learning the cs-reward (which we will describe shortly) from a more diverse set of tasks. See Fig. 2
for a visual summery of our experiments.

### 5.1 Common-Sense Rewards

As Meta-world only contains task-specific rewards, we designed two simple common-sense rewards for our experiments.
Our first cs-reward directs the agent to move one of the key points along the robotic arm (whose 3D location is part of the observation vector) with a target velocity $v_{target}$. The reward function we use is

$$ $r_{CS}(s,a,s^{\prime})=-C_{v}\cdot|\|\ell(s^{\prime})-\ell(s))\|_{2}-v_{target}|$ (8) $$

where $\ell(s)$ is the 3D location of the selected key point given the observation vector. The second reward directs the $L_{1}$ norm of the action towards a target value $n_{target}$. The reward is defined as follows,

$$ $r_{CS}(s,a,s^{\prime})=-C_{n}\cdot|\|a\|_{1}-n_{target}|$ (9) $$

We name these the velocity and action norm cs-rewards, respectively. See Appendix C.1 for further details.

One important property of these rewards is that they do not depend on the task or environment, and as such should be easily transferable. We specifically designed these simple common-sense behaviors to show how IRL struggles to learn a transferable reward.

Figure: Figure 5: Scatter plot of learned cs-reward and ground-truth cs-reward in the single-task transfer experiment, with velocity as common sense.
Refer to caption: extracted/5411392/figures/corr_Single_Task_velocity_Task_Rewards_Correlation.png

**Table 2: Action Norm cs-reward. Each row represents a different task for MT-CSIRL training, each column represents a task for RL training and evaluation. The Numbers represent the ratio between the agents’ action norm and its target value.**
|  | reach-wall | button-press-topdown-wall | drawer-close | push-back | coffee-button |
| --- | --- | --- | --- | --- | --- |
| reach-wall | $\mathbf{0.91}$ | $0.43$ | $0.60$ | $0.54$ | $0.58$ |
| button-press-topdown-wall | $0.60$ | $\mathbf{0.91}$ | $0.58$ | $0.60$ | $0.59$ |
| drawer-close | $0.59$ | $0.53$ | $\mathbf{0.90}$ | $0.63$ | $0.57$ |
| push-back | $0.56$ | $0.56$ | $0.51$ | $\mathbf{0.90}$ | $0.55$ |
| coffee-button | $0.58$ | $0.57$ | $0.54$ | $0.58$ | $\mathbf{0.92}$ |

### 5.2 CS-Reward Trained on a Single Task

Here, we empirically demonstrate that, even for simple common-sense rewards for which the IRL process successfully trains an agent, the final reward cannot be used to train a new agent from scratch. This limitation persists even when applying the reward on the exact same task and target.

To that end, we train an expert with the task-specific (i.e., reach-wall) and ground truth common-sense reward for each of our common-sense rewards. In both cases, we successfully train experts who maximize both the task rewards and the common-sense rewards. Then, for each expert, we sample a set of 60K trajectories for the IRL process. The cs-reward is learned using Alg. 1 with a single task (i.e., $T=1$). Finally, we train a new agent on the exact same task and target. In all of the following experiments in this section, all trained agents achieved 100% success rate on the main task. Therefore, we only report the results on the common-sense rewards.

We present our results in Figure 3.
During the IRL process, the trained agent effectively acquired the desired common-sense behavior, marked in the horizontal red line. However, when we train a new RL agent with our learned reward (and the task-specific reward) does not achieve the same desired behavior, even when training on the exact same target. Similar results with different tasks are shown in Appendix A. We note that this phenomenon has been also previously observed in other bi-level optimization scenarios (Navon et al., ; Vicol et al., ).

In Figure 5, we also present a scatter plot of the learned cs-reward versus the ground-truth cs-reward. This shows a very small correlation between the learned reward and the ground-truth reward (correlation coefficient of 0.18), which further demonstrates that the IRL fails to capture the desired reward.

Figure: Figure 6: Scatter plot of ground-truth VS learned cs-reward for MT-CSIRL on velocity experiment. The rewards are calculated on an unseen task.
Refer to caption: extracted/5411392/figures/corr_MTE_velocity_Task_Rewards_Correlation.png

### 5.3 Learning CS-Reward with Multiple Targets

In the next experiment we show that when we train our IRL on expert demonstration from different targets, we manage to learn a reward that can be used to train new agents on unseen targets. However, it does not transfer to novel tasks. Again, all agents trained in this section achieve 100% success rate so we only show results for the cs-reward.

To do that, we train our MT-CSIRL method with experts trained on different targets for the same task, for each of our cs-rewards. We then test how these learned rewards can be used to train on new targets in this environment and how they transfer to other tasks. We execute this experiment on five different Meta-world tasks: reach-wall, button-press-topdown-wall, drawer-close, push-back and coffee-button. For each task, we did the following: (i) Train an expert for each of the 5 different targets. (ii) Sampled 10K trajectories from each expert. (iii) Train cs-reward using MT-CSIRL. (iv) Train with the learned cs-reward and evaluate using both new targets and new tasks.

The results for both velocity and action norm are presented in Table 1 and 2, respectively. The numbers represent the ratio between the agents’ velocity or action norm and its target value, i.e., the closer to one, the better. On the diagonal, we see the results on different targets for the original task, and on the off-diagonal, we see the results when testing on new tasks.

When looking at the diagonal elements of Tables 1 and 2, we see that our reward transfers well to new targets for the seen task. While our agents do not reach the target values, we still get relatively close to them. However, the drop in performance on the off-diagonal elements indicates that the learned reward does not transfer well to novel tasks. This is especially interesting when we consider our simple cs-rewards, which do not depend on the specifics of the environment (and in the action norm case, is entirely independent of the state).

**Table 3: MT-CSIRL experiment results on unseen task with action norm cs-reward.**
| Agent | Avg Action Norm | Success-rate |
| --- | --- | --- |
| Expert | $2.95\pm 0.01$ | $100\pm 0.00$ |
| SAC | $1.57\pm 0.48$ | $100\pm 0.00$ |
| MT-AIRL | $1.30\pm 0.47$ | $63.5\pm 1.20$ |
| MT-CSIRL | $2.75\pm 0.18$ | $99.9\pm 0.01$ |

**Table 4: MT-CSIRL experiment results on unseen task with velocity cs-reward.**
| Agent | Avg velocity | Success-rate |
| --- | --- | --- |
| Expert | $0.81\pm 0.02$ | $100\pm 0.00$ |
| SAC | $0.32\pm 0.16$ | $100\pm 0.00$ |
| MT-AIRL | $0.11\pm 0.01$ | $67.5\pm 5.90$ |
| MT-CSIRL | $0.69\pm 0.07$ | $98.9\pm 0.02$ |

### 5.4 Learning CS-Reward with Multiple Tasks

In our final experiment, we show that when we train the IRL process on experts’ trajectories from multiple tasks, we learn a useful cs-reward that can be transferred to new unseen tasks. In order to do that, we perform our MT-CSIRL method again, but this time we train each expert on a different meta-world task, as shown in Figure 2.

We evaluate the following methods: (i) Expert: RL trained with the task and ground-truth cs-reward (ii) SAC: An RL trained with only the task reward (iii) MT-AIRL: For this baseline, we employ a multi-task AIRL for learning the entire combined reward per task. This baseline also receives expert trajectories from the test task. (iv) MT-CSIRL: Our proposed method. An RL agent trained with the task and learned cs-reward on an unseen task.

The results for velocity as cs-reward, and action norm as cs-reward are presented in table 3 and 4, respectively.
As can be easily observed, the MT-CSIRL results show that our learned cs-reward manages to transfer the desired behavior even to unseen tasks, albeit not as strongly as the ground-truth reward. We included the MT-AIRL, which unsurprisingly does not work well, to show the importance of our split between task-specific and task-independent rewards. Without this distinction combining different tasks is non-trivial and can easily harm performance as we see here. We also note that our MT-AIRL baseline shares similarities with (Gleave & Habryka, ), however, they train with a meta-learning approach similar to MAML.

To further illustrate our results we show in Fig. 4 the ground truth cs-reward (both for velocity and action norm) on the unseen test task. The figures show the positive impact of our cs-reward, making the agent’s behavior significantly more aligned with the expert. Finally, we show in Fig. 6 a scatter plot of the ground-truth cs-reward versus our learned reward on a new unseen task. As one can see, these rewards are highly correlated (correlation coefficient of 0.88) which again shows that our agent managed to learn the desired reward function. The results in Fig. 6 are for the velocity experiment, the action norm scatter plot is in the appendix.

## 6 Conclusions

In this work, we addressed the important question of how can we get rewards capable of guiding reinforcement learning agents to exhibit desired behavior within our environment. An important step in achieving this goal is to disentangle the task-independent reward, which we name common-sense reward from the task reward. We believe this simple observation can have an important impact as it allows us to
focus on learning *how* the agent should behave and not on *what* the agent should be doing.

Furthermore, our experiments show that not only does our new framework allow us to easily combine expert observations from different tasks, but that learning from a variety of different tasks is a key to learning meaningful reward functions.

Finally, we observe that current IRL methods, while successful in imitation learning, do not learn a proper reward function. We advocate for further work in this direction, as the reward function can be of great importance and not just an auxiliary for imitation learning.

## Appendix A Additional experimental results for Single Task Experiment

In section 5.2 we explained about figure 3, showing that during the IRL process the trained agent successfully learned the desired common-sense behavior. However, a new agent training on the learned reward does not achieve the same desired behavior. Here we will show those results on two more tasks.
In the paper, figure 3 shows graphs for the ”reach-wall-v2” task.
Results fot ”drawer-close-v2”, and ”button-press-topdown-wall-v2” are shown in figures 7 and 8, respectively.

Figure: Figure 7: Single task cs-reward: Velocity and action norm during IRL, and during RL training with the learned cs-reward, on ”drawer-close” task.
Refer to caption: extracted/5411392/appen_vel_first_experiment.png

Figure: Figure 8: Single task cs-reward: Velocity and action norm during IRL, and during RL training with the learned cs-reward, on ”button-press-topdown-wall” task.
Refer to caption: extracted/5411392/appen_vel_first_experiment_2.png

## Appendix B Additional experimental results for Multi Task Experiment

In section 5.4, we showed in a scatter plot the correlation between the learned cs reward and the ground truth reward. We showed it for the velocity case, and here in figure 9 we will show it also for the action norm case. The Correlation coefficient between ground truth reward and cs reward is 0.81. In the action norm case, similarly to the velocity case, we can see that the cs reward is highly correlated with ground truth reward.

Figure: Figure 9: Scatter plot of ground-truth VS learned cs-reward for MT-CSIRL on velocity experiment. The reward is an agent trained on an unseen task.
Refer to caption: extracted/5411392/corr_Single_Task_Action_Norm_Task_Rewards_Correlation.png

## Appendix C More Experiments’ Details

In Alg. 1, we explain the MT-CSIRL training process. The inputs of the algorithm are:

- •
Expert demonstrations $\mathcal{D}_{t_{1}}...\mathcal{D}_{t_{T}}$, which in our experiments set to 10K demonstrations for each task.
- •
Discriminator updates per iteration $D_{\text{updates}}$ is set to 10
- •
Number of tasks $T$: in the multi target experiment is set to 5, and in the multi task experiment is set to 5 targets * 6 tasks.
- •
For the RL generator, we used Soft Actor Critic, with HP as in (Yu et al., )
- •
SAC was updates per iteration $G_{\text{updates}}$,is set to 20.

### C.1 Common Sense Reward with Sawyer Robot

In our experimental framework, we utilize the Meta-World benchmark, which incorporates the Sawyer robot, a collaborative robotic arm known for its high precision and 7-degree-of-freedom. This setup enables us to simulate and evaluate complex tasks, and to define a common sense, which we showed in our paper how we learned ans transfer it between tasks.
Our velocity ground truth reward, was calculated on the ”end effector” part of the sawyer robot shown in figure 10.
The end effector located at the end of the arm, there is an end effector, which can be equipped with different tools or grippers, depending on the task. In simulation environments like MetaWorld, the end effector is often a gripper used for tasks like picking and placing objects.

The location of the end effector is given us by both Mujoco and Meta-worlds observation space. In each step we have an access to the end effector of $s$ ans $s^{\prime}$, and we use those loations to calculate the velocity, as explained in 5.1

We introduce two common sense behaviors, as shown in section 5.1. In the velocity case, we calculate the velocity of the Sawyer’s end effector. Our reward function is maximized when the end effector velocity get closer to $C_{v}$.
In our experiments, we set $C_{v}$ to be 0.8. As for the action norm, our ground truth reward is maximized when the $l_{1}$ of the action space, gets closer to $C_{n}$.In our experiments, we set $C_{n}$ to be 3. It is important to mention that action space is the same for all meta worlds tasks.

Figure: Figure 10: Visualization of Sawyer Robot: The end effector which is the part we calculated the velocity on, marked with an arrow.
Refer to caption: extracted/5411392/sawyer.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 1
questions: 1. IQ-Learn [2] showed that they could recover the reward function and should be added as a baseline for comparison. 
2. When the task reward functions are known, then using Eq 6,7 with multiple tasks should recover the common reward function. However, for unknown task rewards (Section 4.4), Eq 8 will try to recover the task reward function. Given that Eq 8 uses the expert demonstration, will it recover the overall reward function? 
3. If point 2 holds, a reward function trained with Eq 9 might not learn anything informative. This might be true because the learned task reward should learn about the sum of rewards from the expert trajectories. Here, I assume that the demonstration is obtained by experts trained with task rewards (Eq 10 /11). 
3. At line 380, how are the trajectories sampled?
4. In Section 5.2, how will the performance be impacted if the agents cannot recover 100% success rate?
5. In Section 5.3, the method MT-CSIRL-LT uses a similar mechanism to IRL for learning the reward function and the policy. Since IRL does not recover good reward function, how well does the task reward correlate with the original task reward?
6. Overall, I feel the experiment section is hard to understand and needs to be explained better about the training and evaluation setups. The results presented in Table 1 are not clear. 
7. As the episodes can terminate as soon as the task is solved, is it possible that the bias in rewards described in [3] leads to agents not following the constraints, and thereby the baselines in Table 2 do not learn desired behaviors?

#### References
[1] Swamy et al., Of Moments and Matching: A Game-Theoretic Framework for Closing the Imitation Gap, ICML'21\
[2] Garg et al., IQ-Learn: Inverse soft-Q Learning for Imitation, NeurIPS '21\
[3] Jena et al., Addressing reward bias in Adversarial Imitation Learning with neutral reward functions. Deep RL workshop at NeurIPS'2020.
soundness: 2
strengths: - The problem of recovering reward functions using demonstration is interesting as this will facilitate better transfer and generalization across tasks. 
- The experiments on Meta-World show that the method can recover the reward functions that are common across tasks.
summary: This work proposes a method to recover a “common-sense” reward using demonstrations across multiple tasks. This reward represents the part of the reward function that may be common across tasks. The paper discusses why existing IRL methods like AIRL, GAIL fail to recover the true reward function upon convergence. For stable learning, a mechanism to scale the rewards using historical averages is proposed. Experiments on a few tasks in MetaWorld show that the proposed method can extract this reward and with the learned common-sense reward, the agent can learn new tasks with known (standard RL) or unknown task (IRL for task) rewards.
weaknesses: - The motivation is not clear. The point that an IRL algorithms (AIRL, GAIL) can fail to recover the true reward function is highlighted multiple times. However, it depends on the formalism as described in [1], where methods like GAIL is a primal approach and the learned reward at convergence might not align with the true reward, whereas solving the dual form of IRL where the reward function is learned with no-regret approach, the method should recover a meaningful reward function. So, this might explain that IRL methods when trained in Dual form can recover the true reward function and AIRL is not able to learn it.
- The paper is not well written and hard to follow, especially the experiments section. 
- According to the paper, the IRL methods do not learn good reward functions. It is not clear why the proposed method still uses this formalism for the scenario with unknown task rewards (Section 4.4).

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: I listed the questions in the above section together with weaknesses.
soundness: 2
strengths: - **[Motivation]**: This paper addresses a compelling and fundamental technical challenge in reinforcement learning: the need to learn common-sense rewards for improved generalization.

- **[Technical Soundness]**: The method is generally sound, with no major technical flaws. Its simplicity facilitates easy implementation and potential extension to other RL domains.

- **[Clarity of Presentation]**: The overall presentation is clear and easy to follow.
summary: This paper presents a framework that learns disentangled task-specific and task-shared common-sense rewards from multiple tasks using adversarial inverse reinforcement learning (IRL). The approach is simple and straightforward, based on the hypothesis that the reward has an additive structure. The method involves decomposing the reward term of traditional adversarial IRL methods. Results from robotic control experiments (Meta-World) indicate that learning from diverse tasks is crucial for capturing common-sense rewards, and the proposed method outperforms approaches that do not incorporate this common-sense reward learning.

Overall, this paper addresses an intriguing and significant challenge in reinforcement learning: how to learn useful common-sense rewards from demonstration data to facilitate generalization and adaptation. However, I think that certain technical aspects of the work, particularly the hypothesis regarding the reward structure, could be explored further to enhance its applicability to more general tasks and strengthen the framework. Therefore, I would assign a borderline rating (slightly negative), but I will consider adjusting this based on the authors' rebuttal and discussions with other reviewers.
weaknesses: - **[Reward structure]**: The authors assume an additive reward structure, which may be overly simplistic. It would be beneficial to consider more complex functional forms that are representative of real-world scenarios, such as general non-linear functions. Additionally, incorporating noise factors in both common-sense and task-specific rewards—like additive noise, linear non-Gaussian models, or post-nonlinear models with noise—could enhance the framework. A discussion and analysis of these more general function forms, both theoretically (ensuring that disentanglement factors remain identifiable) and empirically (demonstrating that they can be learned within the framework), would be valuable.

- **[Common-sense reward concepts]**: The definition of common-sense rewards in this work is quite basic. Can the framework be scaled to learn more generalizable common-sense rewards, such as a foundational understanding of physical world models or other relevant norms applicable to more complex tasks, like those involving embodied agents? Incorporating language might also be a relevant avenue to explore [1]. This point is more of an open question, but any discussion around it would be appreciated.

- **[Relation to unsupervised RL]** In unsupervised RL, especially mutual-information-based skill discovery, methods can also learn a set of common skills from tasks. It might be beneficial to discuss that line of research and the major benefit of using this framework to learn disentangled common-sense rewards. 

- **[Evaluation]**: Is there any theoretical or formal analysis addressing why common-sense reward learning performs poorly in single-task or single-target transfer scenarios compared to multi-task learning?

[1] Zhao, Zirui, Wee Sun Lee, and David Hsu. "Large language models as commonsense knowledge for large-scale task planning." Advances in Neural Information Processing Systems 36 (2024).

---

confidence: 3
contribution: 1
initial_rating: 1
presentation: 3
questions: None
soundness: 1
strengths: 1. The paper addresses an important problem of learning a policy with learned environmental constraints using maxentIRL.

2.  The paper is easy to read and communicates the method clearly.
summary: The paper proposes to use multi-task experts to learn environmental constraints via a common sense reward. Then the authors demonstrate that by representing any task's reward function as common-sense reward+task-specific reward, agents can learn quickly and safely.
weaknesses: 1. Motivation : The paper's motivation is based on unsubstantiated claims regarding capabilities and limitations of IRL. In line 76-82 the paper claims that "IRL fails to learn correct reward functions due to its connection to GANs". IRL is not equivalent to GAIL; rather GAIL is a way to solve distribution matching and obtained from a primal-dual perspective. While, methods like GAIL are not aiming to learn reward functions there exists method like f-IRL[2] and MaxEntIRL[3] that explicitly try to learn reward functions. f-IRL has looked at learning "stationary reward functions" - reward functions on which the policy can be retrained.
2. Literature review: I found the literature reviews to be somewhat shallow. For example, IRL was attributed to Arora and Doshi (2021); IRL existed long before and its origins can be traced back to [1] and possibly earlier. Authors are encouraged to use correct attributions and perform an in-depth literature review.
3. Unprincipled method: The reward function for any task is assumed to decompose as common sense reward+task-specific reward. Under such a decomposition there is no guarantee that optimizing task-specific reward results in same optimal policy as the combination with common sense reward.  Abeel and Ng[4] have proposed a class of reward functions that result in same optimal policy.
4: Assuming known structure of common sense rewards: In experiments it was presumed that common sense reward had a fix structure. I think that is a major limitation as the structure cannot be determined without domain knowledge. 
5. Section 4.3 is titled curriculum learning, but I do not see the connection to curriculum learning from the method proposed. 
6. Experiments: Standard IRL methods are missing from comparisons like MaxentIRL and f-IRL.


[1]: Abbeel, Pieter, and Andrew Y. Ng. "Apprenticeship learning via inverse reinforcement learning." Proceedings of the twenty-first international conference on Machine learning. 2004.
[2]: Ni, Tianwei, et al. "f-irl: Inverse reinforcement learning via state marginal matching." Conference on Robot Learning. PMLR, 2021.
[3]: Ziebart, Brian D., et al. "Maximum entropy inverse reinforcement learning." Aaai. Vol. 8. 2008.
[4]: Ng, Andrew Y., Daishi Harada, and Stuart Russell. "Policy invariance under reward transformations: Theory and application to reward shaping." Icml. Vol. 99. 1999.

---

confidence: 4
contribution: 2
initial_rating: 1
presentation: 3
questions: * For MT-CSIRL transferred onto an unseen task, do you  have access to that task’s task-specific reward and demonstrations?  Do you do any finetuning for the cs-reward?
* How are the MT-AIRL methods evaluated in Table 2?  Are they trained over all tasks’ demonstrations including the unseen task’s?  Do you use task-specific reward for MT-AIRL as well?
soundness: 1
strengths: * Many of the ideas presented are interesting and worth exploring: disentangled reward functions, learning common sense rewards from multi-task data
summary: The authors propose to do multi-task inverse reinforcement learning by decomposing the reward function into the sum of a task-specific reward, which they assume is given, and a common sense reward, which is learned and shared between the tasks.  Their proposed method uses AIRL to learn a discriminative common sense reward over multiple tasks, then transfers this common sense reward to a new task.  The authors also propose an extension of their method that does not require pre-specified task-specific reward by learning the task specific and common sense reward at the same time with two AIRL discriminators.  They conduct experiments on a subset of Meta-World tasks to demonstrate that their method requires multi-task training in order to learn an effective common sense reward and compares with other multi task IRL methods.
weaknesses: * Their main method assumes task-specific rewards are given.  This is an unrealistic assumption and changes their problem setting from IRL to demonstration-guided RL.  In practice their method looks more like learning an auxiliary reward from multi-task demonstrations to supplement the task reward given by the environment.
* The unknown task reward extension of their method naively learns a task specific and common sense reward jointly without any effort at disentangling the two.  This will not automatically separate the reward functions and it’s unclear why this method would be better than learning a combined reward function.

There are also some serious issues with the experimental setup, which makes the paper’s main claims unconvincing to me.
* The authors only experiment on selected Meta-World tasks, which are all relatively similar tabletop manipulation tasks.  I am not convinced these results are generalizable to other benchmarks and types of tasks.
* Comparison methods.  While their method has access to both task specific rewards and expert demonstrations, the baselines methods seem to only have access to one or the other (reward for SAC, demos for IRL methods), so this is an unfair comparison.  Furthermore, they are missing other methods that decompose reward functions including Reward Network Distillation which they cite in related work.
* The definition of the target velocity and action norm common sense rewards seem contrived since they are not critical or even necessarily helpful to learning the tasks (since almost all experiments achieve 100% task success).  The common sense reward effectively learns a style preference from the demonstrations that’s not present in the task reward function.  The setup in section 5.4 where the algorithm extracts both task specific and common sense reward from demonstrations makes much more sense.
* Section 5.4 is missing any baseline methods, especially doing AIRL on each task individually, so I am not convinced that MT-CSIRL+LT is effectively decomposing the reward functions and transferring them.