Diffusing States and Matching Scores: A New Framework for Imitation Learning
Runzhe Wu, Yiding Chen, Gokul Swamy, Kianté Brantley, Wen Sun
Abstract
Adversarial Imitation Learning is traditionally framed as a two-player zero-sum game between a learner and an adversarially chosen cost function, and can therefore be thought of as the sequential generalization of a Generative Adversarial Network (GAN). However, in recent years, diffusion models have emerged as a non-adversarial alternative to GANs that merely require training a score function via regression, yet produce generations of higher quality. In response, we investigate how to lift insights from diffusion modeling to the sequential setting. We propose diffusing states and performing score-matching along diffused states to measure the discrepancy between the expert's and learner's states. Thus, our approach only requires training score functions to predict noises via standard regression, making it significantly easier and more stable to train than adversarial methods. Theoretically, we prove first-and second-order instance-dependent bounds with linear scaling in the horizon, proving that our approach avoids the compounding errors that stymie offline approaches to imitation learning. Empirically, we show our approach outperforms both GAN-style imitation learning baselines and discriminator-free imitation learning baselines across various continuous control problems, including complex tasks like controlling humanoids to walk, sit, crawl, and navigate through obstacles.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers6
- Scaling Offline RL via Efficient and Expressive Shortcut ModelsNicolas A. Espinosa Dice, Yiyi Zhang, Yiding Chen, Bradley Guo et al.NeurIPS 2025 · 28 citations
- A Smooth Sea Never Made a Skilled SAILOR: Robust Imitation via Learning to SearchArnav Kumar Jain, Vibhakar Mohta, Subin Kim, Atiksh Bhardwaj et al.NeurIPS 2025 · 27 citations
- Fourier Features Let Agents Learn High Precision Policies with Imitation LearningBalázs Gyenes, Emiliyan Gospodinov, Jan Frieling, Enrico Krohmer et al.ICML 2026 · 3 citations
- Distributional Inverse Reinforcement LearningFeiyang Wu, Ye Zhao, Anqi WuICML 2026 · 1 citation
- Matching Multiple Experts: On the Exploitability of Multi-Agent Imitation LearningAntoine Bergerault, Volkan Cevher, Negar MehrICLR 2026
Builds on38
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- High-Resolution Image Synthesis with Latent Diffusion ModelsRobin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser et al.CVPR 2022 · 13,123 citations
- DiffWave: A Versatile Diffusion Model for Audio SynthesisZhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao et al.ICLR 2021 · 1,902 citations
- Score-Based Generative Modeling through Stochastic Differential EquationsYang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar et al.ICLR 2021 · 1,270 citations
- Planning with Diffusion for Flexible Behavior SynthesisMichael Janner, Yilun Du, Joshua B. Tenenbaum, Sergey LevineICML 2022 · 1,115 citations
Related papers
- Diffusion Imitation from ObservationBo-Ruei Huang, Chun-Kai Yang, Chun-Mao Lai, Dai-Jie Wu et al.NeurIPS 2024 · 15 citations
- Diffusion-Reward Adversarial Imitation LearningChun-Mao Lai, Hsiang-Chun Wang, Ping-Chun Hsieh, Yu-Chiang Frank Wang et al.NeurIPS 2024 · 28 citations
- DiffAIL: Diffusion Adversarial Imitation LearningBingzheng Wang, Guoqiang Wu, Teng Pang, Yan Zhang et al.AAAI 2024 · 24 citations
- Time-series Generation by Contrastive ImitationDaniel Jarrett, Ioana Bica, Mihaela van der SchaarNeurIPS 2021 · 29 citations
- Diffusion Based Representation LearningSarthak Mittal, Korbinian Abstreiter, Stefan Bauer, Bernhard Schölkopf et al.ICML 2023 · 71 citations
