Technical Summary: Generalize and Guide: Decomposing Rewards for Few-Shot Inverse Reinforcement Learning

1. Problem Formulation: Few-Shot IRL with Multi-Task Demonstrations (FM-IRL)

The paper addresses a critical bottleneck in Inverse Reinforcement Learning (IRL): the inability of traditional methods to generalize when faced with substantial intra-task variations (e.g., different object shapes, placements, or initial robot poses) given only a limited number of demonstrations for a specific target task.

While prior work has focused on Imitation Learning (IL) or Meta-IRL, these approaches face significant limitations in practical settings:

• Imitation Learning (IL): Methods like Behavior Cloning (BC) struggle to generalize beyond the specific state distribution of the few available demonstrations and cannot improve through online policy optimization without a reward function.
• Meta-IRL: These methods typically require access to a multi-task training environment to meta-train a reward function or context variable, an assumption often impractical in real-world robotics where recreating diverse environments is difficult.

The authors introduce a novel problem setting: Few-Shot IRL with Multi-Task Demonstrations (FM-IRL). In this setting, an agent must learn a new target task with significant variations using:

1. A small number of target-task demonstrations ($D_{target}$).
2. A large dataset of heterogeneous but related task demonstrations ($D_{multi}$).
3. Online interaction with the target task environment ($D_{\pi}$).

Crucially, the agent does not have access to the multi-task training environments, only the demonstration data from them. The goal is to recover the expert distribution of the new task and provide corrective guidance when the agent deviates from it.

2. Methodology: Multitask discriminator Proximity-Guided IRL (MPG)

To solve FM-IRL, the authors propose MPG, a framework that decomposes the reward function into two complementary components: a generalizable discriminator for expert recognition and a proximity function for recovery guidance.

2.1 Multi-Task Discriminator (Generalization)

The first component is a demonstration-conditioned discriminator, $d_\phi(\tau, (s, a))$, designed to transfer shared structural knowledge across related tasks.

• Architecture: It utilizes two separate encoders to process a demonstration trajectory $\tau$ and a state-action pair $(s, a)$. The embeddings are concatenated and passed through fully connected layers to output the probability that $(s, a)$ is consistent with the expert behavior for the task described by $\tau$.
• Training Objective: The discriminator is trained using a binary classification loss on all available demonstrations ($D_{target} \cup D_{multi}$). Positive pairs consist of state-action tuples from the same task as the demonstration; negative pairs consist of tuples from different tasks. Additionally, online policy samples are treated as negative pairs (adversarial objective) to prevent overfitting to expert data.
• Function: This component learns to identify expert behavior in the target task by leveraging the shared structure found in the multi-task dataset, effectively generalizing the expert distribution despite variations.

2.2 Proximity Function (Guidance)

The second component, $p_\theta(s)$, provides dense, informative rewards in non-expert states to guide the agent back toward the expert distribution.

• Concept: It estimates the temporal distance (minimum steps) from a state $s$ to the expert state distribution.
• Learning Objective: Inspired by quasimetric learning, the function is optimized to maximize the proximity value of policy states (pushing them away from the expert distribution in terms of value, but effectively measuring distance) while enforcing temporal consistency.
• Constraints: The learning is constrained by the triangle inequality: $p_\theta(s_t) \leq p_\theta(s_{t+1}) + \zeta$, where $\zeta$ is the cost of a single timestep. Expert states are anchored to a proximity of 0.
• Integration: To enrich the target task's expert distribution, policy states are re-labeled as "expert" if the multi-task discriminator's confidence exceeds a threshold ($d(s, a) > c_{thresh}$).

2.3 Combined Reward Function

The final reward function $\tilde{R}$ combines these two signals:
$$\tilde{R}(s_t, a_t, s_{t+1}) = \underbrace{d(s_t, a_t)}_{\text{Expert Behavior Signal}} + \lambda \underbrace{[p(s_t) - p(s_{t+1})]}_{\text{Proximity Improvement}}$$
The discriminator provides a sparse signal for expert-consistent behavior, while the proximity term provides a dense gradient for reducing the distance to the expert distribution. This combined reward is used to optimize the policy via Reinforcement Learning (PPO).

3. Key Contributions

1. Problem Setting: The paper formalizes the FM-IRL setting, addressing the under-specification of tasks with natural variations where multi-task environments are unavailable but multi-task demonstrations are accessible.
2. Algorithm (MPG): The proposal of a reward decomposition framework that synergistically combines a multi-task discriminator (for generalization across task variations) and a proximity function (for corrective guidance in unseen states).
3. Empirical Validation: The method is evaluated on diverse navigation and manipulation tasks (Maze2D, Block Stacking, and FactorWorld) under significant variations.

4. Experimental Results

The authors evaluated MPG against nine baselines, including traditional IRL (GAIL, MT-AIRL, PEMIRL), Imitation Learning (BC, SQIL, Transformer-BC), and other discriminator/proximity-based methods (DVD, GoalPro).

• Performance: MPG achieved the highest mean success rate on 8 out of 9 tasks.
• Comparison: It outperformed the strongest per-task baseline by an average of 24.7 percentage points in success rate.
• Ablation Studies:
  • Removing either the discriminator or the proximity function individually degraded performance, confirming the necessity of both components.
  • Removing the multi-task dataset ("No Multi-task Data") resulted in performance comparable to the "Proximity Only" ablation, highlighting the discriminator's reliance on multi-task data for generalization.
  • The method remained robust across varying numbers of target demonstrations (5 to 25) and multi-task tasks (2 to 8), though performance saturated after a moderate level of diversity.
• Data Efficiency: MPG performed well even with abundant demonstrations (200), though the advantage over baselines was most pronounced in the few-shot regime.

5. Significance and Claims

The paper claims that MPG effectively bridges the gap between limited supervision and the need for robust policy learning in real-world scenarios characterized by natural variations. By decomposing the reward into a generalizable expert identifier and a proximity-based guide, the method enables agents to "observe, practice, and master" new tasks with minimal supervision, mirroring human learning strategies.

The authors emphasize that their approach does not require the impractical assumption of access to multi-task training environments, relying instead on readily available demonstration datasets. The results suggest that leveraging heterogeneous multi-task demonstrations is a viable and powerful strategy for overcoming the sample inefficiency and generalization failures of traditional few-shot IRL and IL methods.

Limitations noted by the authors:

• The method requires a structured multi-task demonstration dataset from the same domain and agent.
• While efficient in demonstrations, the online training phase remains sample-hungry.
• The reward function is trained online and cannot be directly reused to train a new policy from scratch without retraining.