Imitation learning aims to solve the problem of defining reward functions in real-world decision-making tasks. The current popular approach is the Adversarial Imitation Learning (AIL) framework, which matches expert state-action occupancy measures to obtain a surrogate reward for forward reinforcement learning. However, the traditional discriminator is a simple binary classifier and doesn't learn an accurate distribution, which may result in failing to identify expert-level state-action pairs induced by the policy interacting with the environment. To address this issue, we propose a method named diffusion adversarial imitation learning (DiffAIL), which introduces the diffusion model into the AIL framework. Specifically, DiffAIL models the state-action pairs as unconditional diffusion models and uses diffusion loss as part of the discriminator's learning objective, which enables the discriminator to capture better expert demonstrations and improve generalization. Experimentally, the results show that our method achieves state-of-the-art performance and significantly surpasses expert demonstration on two benchmark tasks, including the standard state-action setting and state-only settings. Our code can be available at an anonymous link https://github.com/ML-Group-SDU/DiffAIL.

翻译：模仿学习旨在解决现实决策任务中奖励函数的定义问题。当前主流方法是采用对抗模仿学习（AIL）框架，通过匹配专家状态-动作占用度量来获取前向强化学习的替代奖励。然而，传统判别器仅为简单的二分类器，无法学习精确的分布，这可能导致无法识别由策略与环境交互产生的专家级状态-动作对。针对此问题，我们提出名为扩散对抗模仿学习（DiffAIL）的方法，将扩散模型引入AIL框架。具体而言，DiffAIL将状态-动作对建模为无条件扩散模型，并采用扩散损失作为判别器学习目标的一部分，从而增强判别器对专家演示的捕捉能力并提升泛化性能。实验结果表明，我们的方法在标准状态-动作设置与仅状态设置两项基准任务中均达到最优性能，且显著超越专家演示。代码可通过匿名链接 https://github.com/ML-Group-SDU/DiffAIL 获取。