Diffusion Imitation from Observation

Summary

This paper introduces Diffusion Imitation from Observation, an adversarial imitation learning method using a conditional diffusion model as the discriminator for policy learning. DIFO learns a diffusion model on expert and agent state transitions, with an auxiliary binary classification objective to discriminate between the expert and the agent. Experiments show that DIFO can learn an imitation policy online, given only expert demonstration observations, outperforming relevant LfO baselines. Ablations and variants show that both loss terms are important for the diffusion model training, and justify the choice of diffusion model conditioning input.

Strengths

- Using conditional diffusion models for adversarial imitation learning from observations is a novel approach. - The loss objectives in the diffusion model are well-motivated and shown to be necessary for downstream policy performance. - The trained diffusion model captures the expert distribution, and exhibits some generalizability in the generated trajectories. - Experiment results show that DIFO works across multiple state-based and image-based environments, outperforming previous methods for learning from observation.

Weaknesses

- This method requires online interactions to train the diffusion discriminator and the policy. One popular alternative approach in this setting (online + expert demonstrations) is optimal transport-based RL, which also works with only access to expert demo observations (optimal transport over expert and agent state trajectories as the reward). [\[1\]](https://arxiv.org/pdf/2008.09167) [\[2\]](https://arxiv.org/pdf/2206.15469) It would be convincing to see comparisons with this line of work. - It would be great to see more evaluations on image-based environments.

Questions

- It seems unclear to me why behavioral cloning, with access to ground truth action labels, is doing worse than LfO methods in Figure 3. - Does DIFO work with negative samples generated offline, e.g. adding Gaussian noise to the previous state, instead of policy online rollouts?

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

- Most experiments are on state-based environments. It would be convincing to see more evaluations on image-based environments. - This method requires online rollouts for collecting negative transitions for discriminator training. It would be great if this method can be shown to work completely offline.

Summary

The paper introduces a novel method named Diffusion Imitation from Observation (DIFO), which integrates diffusion models into the adversarial imitation learning from observation (LfO) framework. Traditional adversarial imitation learning methods often struggle with hyperparameter sensitivity and training stability. DIFO leverages the generative capabilities of diffusion models to improve the imitation process. Specifically, the diffusion model generates the next state based on the current state, and its learning objective is reformulated as a binary classification task to distinguish between expert and agent transitions. This model then provides "realness" rewards to guide the policy learning process. The paper demonstrates that DIFO achieves superior performance across various continuous control tasks compared to existing LfO methods.

Strengths

1. The paper introduces a novel approach by integrating diffusion models into the adversarial imitation learning from observation (LfO) framework, enhancing both stability and performance. 2. The proposed DIFO method consistently outperforms existing LfO methods across various continuous control tasks, demonstrating improved data efficiency.

Weaknesses

1. Although the authors name their paper "Diffusion Imitation from Observation," the performance gain does not seem to be due to the use of the diffusion model. As shown in Section 5.7, using the diffusion loss alone demonstrates very poor results. The major contributing factor is the BCE loss and the discriminator as a whole. It appears more like using diffusion loss as regularization for discriminator training. Do other regularization techniques lead to similar improvements, and why are diffusion models more suitable in this setting (an intuitive explanation)? 2. The LfO baselines in the experiments are rather outdated. More recent LfO baselines should be compared and discussed [1, 2]. 3. Minor problems: 1. In Line 111, there seems to be a missing symbol following the first tilde, and the $\phi$ should be the "parameters" of the diffusion model. 2. In the Preliminary section, the symbol $t$ is used to denote the environment step, while later $t$ is again used as the diffusion step, which might cause confusion. It is common practice in diffusion-based RL papers to denote the diffusion step and environment step using different symbols, respectively written as superscripts and subscripts. 3. The "Expert" lines are missing in Figure 3 and Figure 4. [1] Liu, M., Zhu, Z., Zhuang, Y., Zhang, W., Hao, J., Yu, Y., & Wang, J. (2022, June). Plan Your Target and Learn Your Skills: Transferable State-Only Imitation Learning via Decoupled Policy Optimization. In *International Conference on Machine Learning* (pp. 14173-14196). PMLR. [2] Liu, Y., Dong, W., Hu, Y., Wen, C., Yin, Z. H., Zhang, C., & Gao, Y. Imitation Learning from Observation with Automatic Discount Scheduling. In *The Twelfth International Conference on Learning Representations*.

Questions

1. The diffusion discriminator is different from prior methods in that we can sample from the distribution $D(s, s')$. Can the resulting discriminator itself be used as a policy when combined with an inverse dynamics model $p(a|s, s')$, and how does it perform during evaluation?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Authors listed the limitation of their method in the appendix.

Summary

This paper leverages a diffusion model to learn a expert-state transition and additionally a discriminator that can differentiate expert and agent states. The paper conducts experiments on standard RL environments and demonstrate better performance and data efficiency.

Strengths

- The idea is simple and its experimental results demonstrate strong performance and data efficiency - The paper is well written in general

Weaknesses

- Under section 4.3, why do we only learn the MSE on the expert data but not also the agent data? My understanding is that both data will impact the discriminator output based on section 4.2. - Diffusion models are usually used for modelling multimodal distribution. I have two questions: - Is the expert stochastic? Does the expert policy exhibit diverse behaviours? - Is it more convincing to have stochastic environments rather than purely deterministic environments? - The experimentation seems very limited for the amount of theory it provides---in particular there is no theory regarding whether the proposed objective is truly sound. - Tuning the $\lambda$'s seem difficult as they can vary by ~20% success rate according to figure 7. How should one address this?

Questions

**Questions** - On page 8, line 298, the paper indicates that a "smoother contour" allows for bringing a learning agent closer to the expert. I was wondering if this property is only useful for specific environments but not in general. Dense reward may be helpful but often can misguide policy learning. **Possible typos** - Page 3, line 90: "can't" should be "cannot"

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

The approach leverages only expert data, it remains to extend this to setting with suboptimal demonstrations.

Looks good, I will increase my score

I would like to thank for the author's rebuttal. The additional experiments and clarifications do make the results more convencing. In the light of this, I will increase my score from 5 to 6.

Reminder: The reviewer-author discussion period ends in three days

We would like to express our sincere gratitude to the reviewer for the thorough and constructive feedback. We are confident that our responses adequately address the concerns raised by the reviewer, including the following points. - A clarification of the importance of the BCE and denoising (MSE) losses - Additional results of employing AIL regularization techniques - An intuitive explanation of why diffusion models are suitable for AIL - Additional results of more recent LfO baselines: Decoupled Policy Optimization and Optimal Transport - A discussion of combining our diffusion model and an inverse dynamics model Please kindly let us know if the reviewer has any additional concerns or if further experimental results are required. We are fully committed to resolving any potential issues, should time permit. Again, we thank the reviewer for all the detailed review and the time the reviewer put into helping us to improve our submission.

Please respond to author response if you haven't

Dear Reviewers, Before the author-reviewer discussion period ends, please make sure that you have read the author responses, acknowledge that your review reflects them, and engage with the authors if necessary. Thank you!