Diffusing States and Matching Scores: A New Framework for Imitation Learning

Thank you for your valuable feedback! Your questions are insightful, and we would like to share our thoughts on them below.

• Time Complexity of Evaluating the Cost Function $c^k(s)$ in Eq.3:

Our thoughts on this are twofold:

(1) Estimating the cost function does not require 5000 steps to be executed sequentially (i.e., 5000 steps one by one). This process is fully parallelized. In practice, we use pytorch’s parallelization techniques to optimize this step efficiently during our experiments.

(2) In fact, we don’t need to sample 5000 points to estimate the cost. Concentration inequalities imply that the estimation error decreases rapidly with the number of samples. In our experiments, we found that approximately 100–500 samples are sufficient.

Moreover, even if our algorithm may consume more computing resource, this is a training time only cost--it does not mean the learned policy is slow in running. We can also draw an analogy to image generation: while GANs generate images in a single pass, diffusion models require multiple steps; despite being slower, diffusion models produce higher-quality images.

• Choice of Data Aggregation: why learning from the mixture of $\pi^{(1:k-1)}$ rather than $\pi^{k-1}$

Empirically, learning from mixture policy is more stable and allows us to use a larger dataset for training, which helps with finite sample concerns. Additionally, the key reason why DAC outperforms GAIL is that DAC learns from a mixture of past policies, while GAIL learns solely on the most recent policy. This point is empirically observable and is also argued in the DAC paper.

• The rate of $\epsilon_\text{RL}$

The rate of $\epsilon_\text{RL}$ is a core question in RL theory. Numerous studies have explored and established rates for $\epsilon_\text{RL}$ under various settings, including those with continuous space--to list a few, we have the following:

[1] show that, for linear MDP, the rate is upper bounded by $O(\sqrt{d^3 H/M})$ where $d$ is the feature dimension and $M$ is the number of rollout trajectories. For linear mixture MDP, we have rate of $O(d/\sqrt{M})$ from [2]. More generally, [3] establishes the rate for low-rank MDPs, while [4] extends this to low-nonnegative-rank MDPs. Recall that we defined $\epsilon_\text{RL}$ as the upper bound on returns averaged over all $h$ so the $H$ dependence above is decreased.

The question of the rate of $\epsilon_\text{RL}$ is orthogonal to this work. In fact, any improvement in the rate of $\epsilon_\text{RL}$ can be seamlessly integrated into the theoretical framework we present.

[1] Jin, Chi, et al. "Provably efficient reinforcement learning with linear function approximation." Mathematics of Operations Research 48.3 (2023): 1496-1521.

[2] Zhou, Dongruo, Quanquan Gu, and Csaba Szepesvari. "Nearly minimax optimal reinforcement learning for linear mixture markov decision processes." Conference on Learning Theory. PMLR, 2021.

[3] Agarwal, Alekh, et al. "Flambe: Structural complexity and representation learning of low rank mdps." Advances in neural information processing systems 33 (2020): 20095-20107.

[4] Modi, Aditya, et al. "Model-free representation learning and exploration in low-rank mdps." Journal of Machine Learning Research 25.6 (2024): 1-76.

• how the proposed framework might be adapted or fit into established theoretical results in discrete spaces

Thanks for bringing this interesting point! In this work, we focus on continuous robotics control which itself is an important problem. Regarding discrete one, we have the following thought:

We have treated the diffusion model as a black box in this work. All we require is its training loss (score matching loss) to compute reward signal. We are entirely agnostic to its internal structure. Therefore, for discrete spaces, any new score-matching technique can be directly integrated into our framework and we can apply our algorithm in the same way (by just changing the score matching component in the algo).

Hence, we believe this is highly feasible. In this sense, our algorithm could be extensible and adaptable to future developments like this.

Thank you for your valuable feedback! We understand your concerns regarding our paper and would like to provide detailed responses to each of them below.

• They apply previously established score matching using diffusion models to IL. There have been many methods that used diffusion models in imitation learning (Section 2) and thus I rate the novelty as low.

We would like to clarify that our work is fundamentally distinct from prior works on diffusion models in either the problem setting or the methodology. We take four representative prior works below as example. More details can be found in our related work section.

1. Wang et al., “DiffAIL: Diffusion Adversarial Imitation Learning”, AAAI 2024: They insert the score matching loss directly into the f-divergence (specifically JS-divergence) objective of the discriminator. In other words, they are still training some specific discriminator based on JS-divergence. Hence, their method still belongs to the GAIL framework that we are comparing against. In contrast, we do not train any discriminator at all; instead, we only use score matching (i.e., just regression).

We have incoprate their algorithm as a baseline in our experiments; please refer to Figures 2 and 4 in our draft where we show that our algorithm consistently outperform theirs.

2. Pearce et al., “Imitating Human Behaviour with Diffusion Models”, ICLR 2024: They focus on policy parameterization and directly use a diffusion model to model the policy, which is orthogonal to what we are doing. We do not focus on policy parameterization; instead, we are focused on using a diffusion model to model state(-action) distribution (rather than modeling policies in BC) and generate reward signal for imitation. Note that our policy can have any structure, such as an MLP, and we can even integrate their diffusion-based policy into our framework.

3. Chen et al., “Diffusion Model-Augmented Behavioral Cloning”, ICML 2024: They propose a Behavior Cloning algorithm with an auxiliary loss function computed from diffusion model. In contrast, we consider inverse RL where diffusion models are used to model state(-action) distributions. So ours differs from theirs in methodology (it is BC vs IRL). Also note that, since their algorithm is BC, it requires expert action while we do not; in addition, BC is likely to suffer compounding error while it is known that IL w/ additional online interaction addresses compounding error (see [1], also as we did in this paper).

[1] Swamy, Gokul, et al. "Of moments and matching: A game-theoretic framework for closing the imitation gap." International Conference on Machine Learning. PMLR, 2021.

4. Lai et al., "Diffusion-Reward Adversarial Imitation Learning", NeurIPS 2024: They use diffusion model as discriminator and their Diffusion Reward is from the standard JS-divergence objective. So it also belongs to the GAIL framework that we are comparing against just like Wang et al.

In summary, prior works either consider a totally different problem setting (e.g. policy parameterization using diffusion model) or using different methodology--for example, some use the GAN-style discriminator that we are comparing against in the paper; in contrast, we do not train a discriminator at all; instead, we only use score matching (i.e., just regression). This is the central contribution of our paper. Given that no prior work has introduced such an approach, we believe this is a novelty. We will also incorporate a more comprehensive review of related work in the next version to clarify it.

• Can the authors clarify if a reward function is inferred or clarify on how SMILING is indeed an IRL method?

Yes, we do infer a reward (cost) function during training. See Eq. (3): $c^k(s)$ is the cost function. Our algorithm is really minimizing this cost function (see line 4 of our algorithm).

• connection to SAC/DreamerV3 / Were SAC and DreamerV3 used in the experiments?

Our algorithm requires an RL subroutine solver (see line 4 of our algorithm) to minimize the cost function we infer, and SAC and DreamerV3 are exactly the RL solvers we are using. Specifically, in line 4 of our algorithm, SAC or DreamerV3 can be employed to perform this minimization. Notably, our algorithm does not rely on any specific RL subroutine and can work with other suitable alternatives as well.

• It's unclear how Section 6.3 helps to show that score functions are more expressive than discriminators.

Section 6.3 (now moved to Appendix D.3 to save some space since we added more experiments) is designed to support the key arguments we made in Section 5.1 (now moved to Appendix A): using score functions, we can employ a less expressive function class to represent a richer class of distributions than f-divergence- and IPM-based approach, making them more effective in imitation learning.

There we provided an example to illustrate this point. When the target distribution belongs to the exponential family (a common class of distributions, expressed as $p(s) = \exp(w^T \phi(s))/Z$), the score function can represent it easily, as its score function is linear. In contrast, methods based on f-divergence or IPM will require nonlinear discriminators. These methods require more expressive function classes (empirically, larger neural networks), which in turn increase the sample complexity needed for learning.

If you think this part is important in building our claim, we can definitely move this section to the main content.

I thank the authors for the clarifications and all of the answers made sense. I have raised my score.

Thank you for your valuable feedback! We appreciate your comments about our experiments and have conducted additional experiments based on your suggestions. In particular, please find our responses to each of your concerns below.

• The experiments are not sufficient.

We have added empirical comparison to a related algorithm DiffAIL [1] and a state-of-the-art algorithm IQ-Learn [4]. We ran them on all our tasks. The results are already added in Figures 2 and 4, which suggest our method consistently outperforms all of them.

Regarding [2,3], we would like to clarify that they are truly orthogonal to ours. They focus on policy parameterization by representing the policy as a diffusion model. In contrast, our work centers on imitation learning, focusing on how to match state(-action) occupancies and generate reward signals.

Due to this orthogonality, we could even replace our current MLP-based policy by their diffusion policy to take their advantage; but this is orthogonal to what we are focusing on in this paper.

• the authors should report the specific performance values of these policies as reference points

Thanks for this suggestion! We have included the "unnormalized" versions of experiment results in Appendix D.4 for reference!

• If the number of training steps for the expert policy is on a similar scale (e.g., 10^6), directly applying reinforcement learning to learn the expert policy could be a simpler solution.

We want to clarify that we focus on imitation learning, where we only have access to expert states and have no reward information in both the dataset and the online interaction, so unclear how to do RL here. Actually, we merely use RL to train expert and thus generate demonstrations. In practice, we could imagine them coming from real people as is common in robot learning. In that case, the demonstrations naturally come without reward signals so people usually use imitation learning.

[1] Wang, Bingzheng, et al. "DiffAIL: Diffusion Adversarial Imitation Learning." Proceedings of the AAAI Conference on Artificial Intelligence (AAAI). 2024

[2] Chi, Cheng, et al. "Diffusion policy: Visuomotor policy learning via action diffusion." Robotics: Science and Systems (RSS). 2023

[3] Pearce, Tim, et al. "Imitating human behaviour with diffusion models." International Conference on Learning Representations (ICLR). 2023

[4] Garg, Divyansh, et al. "Iq-learn: Inverse soft-q learning for imitation." Advances in Neural Information Processing Systems 34 (2021): 4028-4039.

We are delighted that our experiments have addressed your previous concerns.

In response to your valuable suggestions, we have included additional details on expert training and data collection in our manuscript (see Lines 1353-1363 on Page 26). We will release the expert checkpoints, expert demos, and code upon acceptance.

To summarize, we use RL solvers (SAC for DMC tasks and DreamerV3 for HumanoidBench tasks) to train the expert policies on benchmarks’ ground truth rewards. Once the expert training is complete, we generate expert demonstrations by running the learned expert policy for five episodes for each task. Since each episode has a fixed horizon of 1K steps, the resulting dataset contains $5 \times 1,000 = 5,000$ samples per task. This approach is consistent with prior works on imitation learning (including GAIL and DAC) where the expert demos are also generated by first training a policy using RL w/ the ground truth reward, and then collecting trajectories from the RL policy.

We hope this clarifies our setup and are happy to engage in further discussion.

Thank you for your valuable feedback! We appreciate your input and have provided our response to your questions below.

• Can you elaborate on how your is novel compared to the following...

We would like to elaborate on it below:

[1]: This work uses diffusion models to represent policies and aims to learn a relative reward function which, when employed as classifier guidance, converts the base policy into the expert policy. It is relevant to our work as it also applies diffusion models in the context of imitation learning. However, unlike our method, which uses diffusion models with a score-matching loss to model state(-action) distributions and generate reward signals, their focus lies in diffusion policies and learning through classifier guidance. Additionally, their method is offline and is thus likely to suffer compounding error, while imitation learning w/ online interaction is known to address this (see [4], as well as we did in our paper).

[2] also explore score matching for imitation learning but consider a different class of models--energy-based models, so it is different from our focus on diffusion models.

[3] consider the problem from the other way around-—it uses inverse RL, specifically maximum entropy IRL, to train diffusion models. In contrast, our approach leverages diffusion models and its standard regression-based training method to perform IRL. In other words, they use diffusion for policy, while we use diffusion for reward (for imitation).

We have make sure to include [1,2] into the related work section where we have highlighted in red.

• Could the authors elaborate on the sentence...Specifically, what does "mixture of the preceding policies" mean?

"mixture of the preceding policies" means the mixture of all previous generated policies prior to the current round: let's say we are at round $k$, then the mixture is on $\pi_1,\dots,\pi_{k-1}$. The "mixture policy" is executed by first choosing a policy $\pi^i$ uniformly at random from all these policies and then executing the chosen policy. It is trajectory-level mixture instead of action-level. (Precise definition provided in line 378-380)

That sentence means: We have two diffusion models. One of them learns the state distribution of the expert policy, and the other learns the state distribution of the mixture policy. We then use the score functions of these two diffusion models to define the reward function, which is used to train a new policy at the next round.

If this explanation sounds good for you, we will incorporate it into the corresponding section to improve clarity.

• definition of $c^{(k)}$

We have made the definition of $c^{(k)}$ more clear in Eq.(3) in line 295.

[1] Nuti, Felipe, Tim Franzmeyer, and João F. Henriques. "Extracting reward functions from diffusion models." Advances in Neural Information Processing Systems 36 (2024).

[2] Liu, Minghuan, et al. "Energy-Based Imitation Learning." Proceedings of the 20th International Conference on Autonomous Agents and MultiAgent Systems. 2021.

[3] Yoon, Sangwoong, et al. "Maximum Entropy Inverse Reinforcement Learning of Diffusion Models with Energy-Based Models." The Thirty-eighth Annual Conference on Neural Information Processing Systems.

[4] Swamy, Gokul, et al. "Of moments and matching: A game-theoretic framework for closing the imitation gap." International Conference on Machine Learning. PMLR, 2021.

Thank you for your valuable feedback! We appreciate your comments regarding our experiments and have conducted additional experiments as per your suggestions. Moreover, We have provided a more thorough comparison with related works. Please find our response to each of your concerns below.

• Previous studies [A, B, C, D] have already explored integrating diffusion models with imitation learning.

We would like to clarify that our work is fundamentally distinct from [A] and [D] in the problem setting and [B] and [C] in methodology. We elaborate on this below:

[A] & [D]: [A] insert the score matching loss directly into the f-divergence (specifically JS-divergence) objective of the discriminator. Similarly, [D] use diffusion model as discriminator and their Diffusion Reward is from the standard JS-divergence objective. Hence, both of them are still training some specific discriminator based on JS-divergence, and thus, their methods still belong to the GAIL framework that we are comparing against.

In other words, both [A] and [D] are training a discriminator using JS-divergence objective (so it is still GAIL); in contrast, we do not train a discriminator at all; instead, we only use score matching (i.e., just regression). This is the central contribution of our paper. Given that no prior work has introduced such an approach, we believe this is a novelty.

In addition, we have incorporated [A] as a baseline in our experiments; please refer to Figures 2 and 4 in our draft where we show that our algorithm consistently outperform theirs.

[B]: They focus on policy parameterization and directly use a diffusion model to model the policy, which is orthogonal to what we are doing. We do not focus on policy parameterization; instead, we are focused on using a diffusion model to model state(-action) distribution (rather than modeling policies in BC) and generate reward signal for imitation. Note that our policy can have any structure, such as an MLP, and we can even integrate their diffusion-based policy into our framework.

[C]: They propose a Behavior Cloning algorithm with an auxiliary loss function computed from diffusion model. In contrast, we consider inverse RL where diffusion models are used to model state(-action) distributions. So ours differs from theirs in methodology (it is BC vs IRL). Also note that, since their algorithm is BC, it requires expert action while we do not; in addition, BC is likely to suffer compounding error while it is known that IL w/ additional online interaction addresses compounding error (see [H], also as we did in this paper).

We have made sure to include these works in the related work section where we have highlighted them in red. We will also incorporate a more comprehensive review of related work in the next version.

• SMILING is compared only against two obsolete baselines: BC and DAC. It is recommended to include comparisons with...

Thanks for these suggestions! We have added empirical comparison to related algorithm DiffAIL and a recent state-of-the-art algorithm IQ-Learn. Also, prior works have shown that ValueDICE is at best comparable to IQ-Learn (e.g.,[E], [F]), so we think comparision to IQ-Learn is enough. For these algorithms, we use their official implementation while making sure the neural networks are of same size with ours.

We ran them on all our tasks. The results are added in Figures 2 and 4, which suggest our method consistently outperforms all of them.

• experiments quantifying the relationship between the number of expert trajectories and the converged normalized score are essential for substantiating these claims

Thanks for this suggestions and we found this very reasonable! Hence, we conducted an additional experiment to study the impact of varying the number of expert demonstrations on the performance of the learned policy. The results are presented in Appendix D.2 (specifically in Figures 5 and 6). We observe that SMILING can achieve near-expert-level performance with just 5K expert states (Figure 5), whereas DAC requires 25K to reach similar performance. This indicates that SMILING is more efficient in utilizing expert demonstrations than GAN-style discrimator methods, which support our theoretical results (Theorem 1). If you find this section helpful in supporting our theory, we should move it back in the main content to emphasize it.

• experimental outcomes in both stochastic and deterministic cost scenarios

We hope the additional experiments we provided in the previous two responses have demonstrated the better efficiency of our algorithm. In addition, we want to emphasize that the experiments we conducted are not entirely deterministic. Cost is not the only source of randomness (also noting that we are not assuming determinisitic cost either, and moreover, our policies are also random). However, as demonstrated by our previous experiments on the relationship between the number of expert demonstrations and return, our algorithm can outperform others with significantly better sample efficiency (Figure 5) despite it is not deterministic.

• Q1. Is the theoretical result specific to state-matching only? or can it also be applied to (state-action)-pair score matching?

The theoretical results can be extended to state-action pairs for sure, as we can imagine an "augmented state" by concatenating the original state and action together. The entire proof remains valid.

• Q2. LobsDICE [G] highlights a limitation of the state-matching objective...

Thank you for bringing this interesting point! In our setting, this issue does not pose any problem.

Our goal is to maximize the return (or specifically, to ensure that the total return of the learned policy matches that of the expert). Recall that we have defined the reward solely based on the state (this seems the best we can do given that we only have expert state, but the following argument will be extendable to learning from states and actions as well). Hence, to maximize return, it suffices to align their state distributions $d(s)$, because for any two policies $\pi_1$ and $\pi_2$, we have:

$V^{\pi^1} - V^{\pi^2} =\sum_s \big(d^{\pi^1}(s) - d^{\pi^2}(s) \big) r(s)$

This implies that if $d(s)$ matches, the total returns will naturally match as well. In the counterexample provided in [G], while the policies may visit states in different order, the state distribution is same. Therefore, matching state distribution is sufficient for matching value, matching trajectory distributions is unnecessary.

If you find this explanation reasonable, we would be happy to include their argument and this clarification in our draft.

[A] Wang et al., “DiffAIL: Diffusion Adversarial Imitation Learning”, AAAI 2024.

[B] Pearce et al., “Imitating Human Behaviour with Diffusion Models”, ICLR 2024.

[C] Chen et al., “Diffusion Model-Augmented Behavioral Cloning”, ICML 2024.

[D] Lat et al., "Diffusion-Reward Adversarial Imitation Learning", NeurIPS 2024.

[E] Chang et al. "Adversarial Imitation Learning via Boosting." ICLR 2024.

[F] Garg et al., "IQlearn: Inverse soft-Q Learning for Imitation", NeurIPS 2021.

[G] Kim et al., “LobsDICE: Offline Learning from Observation via Stationary Distribution Correction Estimation”, NeurIPS 2022.

[H] Swamy et al. "Of moments and matching: A game-theoretic framework for closing the imitation gap." ICML, 2021.