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

arXiv: https://arxiv.org/abs/2607.17760 · HTML: https://mlsanity.com/paper/2607.17760/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Inverse reinforcement learning, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 28 of its 37 references that are in the map.
- First posted on arXiv 20 Jul 2026.

## Placed in

- Offline RL and imitation learning from data / [Inverse reinforcement learning](https://mlsanity.com/offlinerl/inverse-reinforcement-learning/)

## Builds on (most cited first, 28 of 28)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Concrete Problems in AI Safety](https://mlsanity.com/paper/1606.06565/) (2016)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [What Matters in Learning from Offline Human Demonstrations for Robot Manipulation](https://mlsanity.com/paper/2108.03298/) (2021)
- [Minigrid & Miniworld: Modular & Customizable Reinforcement Learning Environments for Goal-Oriented Tasks](https://mlsanity.com/paper/2306.13831/) (2023)
- [One-Shot Imitation Learning](https://mlsanity.com/paper/1703.07326/) (2017)
- [One-Shot Visual Imitation Learning via Meta-Learning](https://mlsanity.com/paper/1709.04905/) (2017)
- [Accelerating Reinforcement Learning with Learned Skill Priors](https://mlsanity.com/paper/2010.11944/) (2020)
- [One-Shot Imitation from Observing Humans via Domain-Adaptive Meta-Learning](https://mlsanity.com/paper/1802.01557/) (2018)
- [Primal Wasserstein Imitation Learning](https://mlsanity.com/paper/2006.04678/) (2020)
- [Optimal Goal-Reaching Reinforcement Learning via Quasimetric Learning](https://mlsanity.com/paper/2304.01203/) (2023)
- [Learning Generalizable Robotic Reward Functions from "In-The-Wild" Human Videos](https://mlsanity.com/paper/2103.16817/) (2021)
- [Discriminator-Weighted Offline Imitation Learning from Suboptimal Demonstrations](https://mlsanity.com/paper/2207.10050/) (2022)
- [Decomposing the Generalization Gap in Imitation Learning for Visual Robotic Manipulation](https://mlsanity.com/paper/2307.03659/) (2023)
- [Watch and Match: Supercharging Imitation with Regularized Optimal Transport](https://mlsanity.com/paper/2206.15469/) (2022)
- [Meta-Inverse Reinforcement Learning with Probabilistic Context Variables](https://mlsanity.com/paper/1909.09314/) (2019)
- [Optimal Transport for Offline Imitation Learning](https://mlsanity.com/paper/2303.13971/) (2023)
- [Hierarchical Few-Shot Imitation with Skill Transition Models](https://mlsanity.com/paper/2107.08981/) (2021)
- [Learning a Prior over Intent via Meta-Inverse Reinforcement Learning](https://mlsanity.com/paper/1805.12573/) (2018)
- [Few-Shot Goal Inference for Visuomotor Learning and Planning](https://mlsanity.com/paper/1810.00482/) (2018)
- [Robust Imitation Learning against Variations in Environment Dynamics](https://mlsanity.com/paper/2206.09314/) (2022)
- [PsiPhi-Learning: Reinforcement Learning with Demonstrations using Successor Features and Inverse Temporal Difference Learning](https://mlsanity.com/paper/2102.12560/) (2021)
- [Multi-task Maximum Entropy Inverse Reinforcement Learning](https://mlsanity.com/paper/1805.08882/) (2018)
- [Multi-task Hierarchical Adversarial Inverse Reinforcement Learning](https://mlsanity.com/paper/2305.12633/) (2023)
- [Expert Proximity as Surrogate Rewards for Single Demonstration Imitation Learning](https://mlsanity.com/paper/2402.01057/) (2024)

## Abstract (arXiv)

Inverse reinforcement learning (IRL) provides a powerful framework for learning from demonstrations. However, real-world tasks often exhibit substantial natural variations (e.g., picking up mugs with varying shapes), making it impractical to collect demonstrations that fully specify a new task under every possible scenario. In practice, while demonstrations for the target task are limited, it is often easier to obtain datasets of heterogeneous but related behaviors. This motivates the problem of few-shot IRL with multi-task demonstrations (FM-IRL), where an agent must learn a new task with substantial variations from only a limited number of target-task demonstrations, together with sufficient demonstrations of related tasks and online agent experience. To do so, we must both recover the expert distribution of the new task and provide guidance when the agent deviates from it. We introduce Multitask discriminator Proximity-Guided IRL (MPG), which learns two complementary reward components: (1) a generalizable discriminator that transfers shared structure across related tasks to identify expert behavior in a new task, and (2) a proximity function that measures how far a state deviates from expert behavior and provides corrective guidance during exploration. We demonstrate the effectiveness of our method on multiple challenging navigation and manipulation tasks under significant variations (e.g., object configurations, table layouts, and initial robot poses), achieving an average success rate of 81.2%, outperforming the strongest per-task baseline by an average of 24.7 percentage points.