Rethinking Inverse Reinforcement Learning: from Data Alignment to Task Alignment

Many imitation learning (IL) algorithms use inverse reinforcement learning (IRL) to infer a reward function that aligns with the demonstrations. However, the inferred reward function often fails to capture the underlying task objective. In this paper, we propose a novel framework for IRL-based IL that prioritizes task alignment over conventional data alignment. Our framework is a semi-supervised approach that leverages expert demonstrations as weak supervision signals to derive a set of candidate reward functions that align with the task rather than only with the data. It adopts an adversarial mechanism to train a policy with this set of reward functions to gain a collective validation of the policy's ability to accomplish the task. We provide theoretical insights into this framework's ability to mitigate task-reward misalignment and present a practical implementation. Our experimental results show that our framework outperforms conventional IL baselines in complex and transfer learning scenarios. The complete code are available at https://github.com/zwc662/PAGAR. achieves high --but not necessarily optimal --performance. The rationale is that achieving high performance under a task-aligned reward function is often adequate for real-world applications.

Building on this premise, we leverage IRL to derive the set of candidate task-aligned reward functions and propose Protagonist Antagonist Guided Adversarial Reward (PAGAR), a semi-supervised framework designed to mitigate task-reward misalignment by training a policy with this candidate reward set. PAGAR adopts an adversarial training mechanism between a protagonist policy and an adversarial reward searcher, iteratively improving the policy learner to attain high performance across the candidate reward set. This method moves beyond relying on deriving a single reward function from data, enabling a collective validation of the policy's similarity to expert demonstrations in terms of effectiveness in accomplishing tasks. Experimental results show that our algorithm outperforms baselines on complex IL tasks with limited demonstrations and in challenging transfer environments. We summarize our contributions below.