Summary
The paper proposes Diffusion-Reward Adversarial Imitation Learning (DRAIL) an Adversarial Imitation Learning method where the discriminator is parameterized by the loss of a conditional diffusion model.
The diffusion model takes in a state-action pair $(s, a)$ and a binary label $c$ indicating whether the state-action pair comes from the discriminator (label $c^-$) or from the training set (label $c^+$). Its loss on these inputs is denoted $\mathcal{L}\_{\textrm{diff}}(s, a, c)$. The discriminator is then defined as $D(s, a) = \sigma(\mathcal{L}\_{\textrm{diff}}(s, a, c^+) - \mathcal{L}\_{\textrm{diff}}(s, a, c^-))$, where $\sigma$ is the sigmoid function.
The method resembles DiffAIL (Wang et al. 2023), which defines the discriminator as $\exp(-\mathcal{L}\_{\textrm{diff}}(s, a))$ (unconditional diffusion model), as well as Diffusion Classifier (Li et al. 2023).
They benchmark DRAIL on six settings, including long-horizon planning (e.g. Maze2d), high-dimensional control (e.g. AntReach) and dexterous manipulation (e.g. HandRotate). They compare against methods including Behavior Cloning, the original GAIL, and variants such as Wasserstein AIL and DiffAIL. They also compare against Diffusion Policy, which uses a diffusion model to parameterize a policy rather than a reward function.
In terms of raw performance, DRAIL obtains:
- visibly better results than the baselines in 4/6 environments;
- a result very close to that of DiffAIL in 1/6 environment;
- a result worse than DiffAIL in 1/6 environment.
Further experiments show:
- DRAIL is more robust to noise on initial states and goal positions than the baselines on a block-pushing task (FetchPush);
- DRAIL is more data-efficient than the baselines on the Walker environment (continuous high-dimensional control) and FetchPush;
- DRAIL learns more generalizable rewards than GAIL on a synthetic 2d point dataset.
Strengths
1. Clarity: the paper is well-written and the proposed method is explained clearly.
1. Extensive comparison with DiffAIL, albeit in the appendix.
1. Methodological novelty: the authors propose using a conditional diffusion model to parameterize the discriminator for AIL, which removes limitations from DiffAIL (such as an arbitrary decision boundary at $\ln 2$; see Appendix A).
1. Significance and quality of experimental results: the paper demonstrates that defining the discriminator using a conditional diffusion model improves performance, robustness and data-efficiency, compared to DiffAIL (which uses an unconditional diffusion model in the discriminator) and to other baselines. Hence, one can argue it constitutes a step forward in Adversarial Imitation Learning.
1. Referencing strongly related prior work: the authors point to DiffAIL and Diffusion Classifier, which strongly resemble different parts of DRAIL (namely using the loss of a diffusion model to parameterize a discriminator, and how to construct such a discriminator). I personally was not aware of these two papers before, and commend the authors for bringing these to the reader’s attention.
Weaknesses
## Better delimiting the author’s contributions
Despite DiffAIL and Diffusion Classifier being mentioned and their relationship to DRAIL being discussed, my impression is that more credit should be materially given to these two prior works.
For example, in line 143, the authors state “we propose to leverage the training procedure of DDPM and develop a mechanism to provide learning signals to a policy using a single denoising step”. However, it seems to me that this much was already present in DiffAIL.
Similarly, line 145 states “Our key insight is to leverage the derivations developed by Kingma et al. [22], Song et al. [47], which suggest that the diffusion loss, i.e., the difference between the predicted noise and the injected noise, indicates how well the data fits the target distribution since the diffusion loss is the upper bound of the negative log-likelihood of data in the target distribution”. However, it seems that a similar insight had already been used by Li et al. (2023), as mentioned in line 177.
Hence, my understanding is that, in the context of these two prior works, DRAIL can be seen as a combination of DiffAIL and Diffusion Classifier (DC). Please feel free to correct me if this understanding is incorrect.
This still constitutes novelty, although less than if DiffAIL and DC hadn’t preceded this work. In fact, in Appendix A the authors make what I consider to be a compelling case for why the discriminator should be “symmetric”, rather than “one-sided” like in DiffAIL.
As such, I believe this paper’s contributions could be presented in a way that better separates its original contribution from that present in prior work. One way of doing this could be:
- Add a section on DiffAIL and Diffusion Classifier to Section 3 (preliminaries).
- In Section 4 (Approach), discuss the limitations of DiffAIL, e.g. as done in Appendix A. Argue that using a conditional diffusion model with a symmetric discriminator (via Diffusion Classifier) would address these limitations.
- Hence present the final methodology as the combination of DiffAIL and DC.
Again, I am open to changing my mind about the above assessment, and will raise my score accordingly in that case, or in case a refactoring of the presentation of the method is implemented by the authors.
## Baselines against planning methods that use diffusion models
The authors baseline DRAIL against 3 AIL methods and Diffusion Policy. The motivation for including the latter is to “compare learning a diffusion model as a policy (diffusion policy) or reward function (ours)”. As per a recent survey of Zhu et al. (2023), diffusion models have been used in sequential decision-making not only for parameterizing policies, but also as planners, as is the case in the seminal work of Janner et al. (2022).
As such, for a more holistic comparison against methods outside of the AIL framework, the authors could include diffusion-based planning baselines such as Diffuser (Janner et al. 2022).
Recent work by Nuti et al. (2023) has also shown reward functions can be recovered from pairs of diffusion models in the setting of Janner et al. (2022), so that it would also be possible to compare the rewards learned via DRAIL against those obtained from Diffuser.
### References
Zhu, Zhengbang, et al. "Diffusion models for reinforcement learning: A survey." arXiv preprint arXiv:2311.01223 (2023).
Nuti, Felipe, Tim Franzmeyer, and João F. Henriques. "Extracting reward functions from diffusion models." Advances in Neural Information Processing Systems 36 (2024).
Limitations
The authors consider the limitations of their work in Appendix G and societal impacts in Appendix I.
My main concern with this work is the extent to which the authors claim originality over aspects of their method that are arguably already present in prior work. If the authors can clarify that this concern is misplaced, or if they modify the exposition of their contributions as outlined in the Weaknesses sections, I am open to raising my score.