Diffusion Model-Augmented Behavioral Cloning

Imitation learning addresses the challenge of learning by observing an expert's demonstrations without access to reward signals from environments. Most existing imitation learning methods that do not require interacting with environments either model the expert distribution as the conditional probability p(a|s) (e.g., behavioral cloning, BC) or the joint probability p(s, a). Despite the simplicity of modeling the conditional probability with BC, it usually struggles with generalization. While modeling the joint probability can improve generalization performance, the inference procedure is often time-consuming, and the model can suffer from manifold overfitting. This work proposes an imitation learning framework that benefits from modeling both the conditional and joint probability of the expert distribution. Our proposed Diffusion Model-Augmented Behavioral Cloning (DBC) employs a diffusion model trained to model expert behaviors and learns a policy to optimize both the BC loss (conditional) and our proposed diffusion model loss (joint). DBC outperforms baselines in various continuous control tasks in navigation, robot arm manipulation, dexterous manipulation, and locomotion. We design additional experiments to verify the limitations of modeling either the conditional probability or the joint probability of the expert distribution, as well as compare different generative models. Ablation studies justify the effectiveness of our design choices.

Paper

Similar papers

Reviewer Q6yg7/10 · confidence 4/52023-06-12

Summary

The paper proposes a new algorithm for behavior cloning (BC) where the BC learning objective is modified with a diffusion modeling loss that models the joint state-action distribution of the expert data. The paper demonstrates the benefits of modeling both the conditional probability and joint probability of the expert distribution. The results of the proposed method have been compared with baselines modeling both conditional probability and joint probability distributions and have been supported with appropriate ablation studies.

Strengths

- The paper is well-written and easy to understand. - The problem has been well motivated. The authors claim that while solely modeling the conditional probability p(a|s) of the expert data struggle to generalize to unseen states during training, solely modeling the joint probability can improve generalization to unseen states at the cost of high inference time. Further, though generative models provide encouraging results in modeling stochastic and multimodal behaviors, they struggle with manifold overfitting (shown in Fig. 4(c)). Accordingly, the paper proposes a loss function that combines modeling the conditional and joint probabilities of the expert data while using a state-of-the-art generative model (diffusion model) to benefit from their superior behavior modeling capabilities (as shown in [1] which has also been cited in the paper). - An interesting feature of the proposed method is that since it only uses a diffusion model to guide the learning of conditional probability, it is able to circumvent the issue of high inference time associated with the noising and denoising procedure in diffusion models (mentioned in Sec. 5.3). - The primary results have been compared with baselines modeling both conditional probability and joint probability distributions. Table 1 shows that DBC outperforms BC (conditional probability), Implicit BC (joint probability), and Diffusion policy (a state-of-the-art generative model) on all tasks except the Walker. However, since DBC performs comparably to the best performing baseline (BC) (the authors provide some justification for this in Sec. 5.3 Locomotion). - The paper provides interesting insights on modeling high dimensional action spaces, inference efficiency, generalization to unseen goals, and manifold overfitting. Additional ablation studies to justify design choices such as the choice of the generative model, loss coefficient value, and effect of the normalization term have also been provided. [1] Chi, Cheng, et al. "Diffusion policy: Visuomotor policy learning via action diffusion." arXiv preprint arXiv:2303.04137 (2023).

Weaknesses

- The limitations of the proposed approach can only be found in the limitations. Though I believe that is fine given space limitations, it would be nice to mention it briefly in the main paper for completeness. - Since the paper focuses on imitation learning from expert data, it would be interesting to also add *Action Chunking with Transformers* (ACT) [1] as a baseline. ACT also uses a generative model (a conditional variational autoencoder) as a policy and given the reported results, it would be interesting to see how DBC compares with ACT. I understand that the paper does compare with a VAE. However, ACT with its transformer encoder and decoder would be an interesting baseline to compare with. - I am curious about how the performance of the proposed method scales with the dataset size. Specifically, since diffusion models seem to be “data hungry” (based on results from the computer vision and NLP community), it would be interesting how this method scales with dataset size. Since this method solely uses a diffusion model to guide the learning of policy, I believe a study of differences in the performance of DBC and the baselines with different amounts of training data might give interesting insights. [1] Zhao, Tony Z., et al. "Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware." arXiv preprint arXiv:2304.13705 (2023).

Questions

It would be great if the authors could address the points mentioned in the “Weaknesses” section.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes. However, the limitations and societal impact of the work can only be found in the appendix. It would be nice to either have a brief mention or a reference to the appendix section in the main paper for completeness.

Reviewer Pa843/10 · confidence 4/52023-07-05

Summary

The submission proposes an imitation learning method optimizing a loss that is a weighted sum of a behavioral cloning loss and a loss based on a diffusion model. The diffusion model loss penalizes the policy for generating actions that are unlikely under the diffusion model, which is pre-trained to maximize the likelihood of joint state-action pairs sampled from the data.

Strengths

Originality: I am not aware of any previous work combining a behavioral cloning loss with a diffusion model loss. Quality: I think there is one interesting "key trick" that seems to be required for the method to work, which is using the diffusion-model-loss incurred by the expert demonstration as a baseline. I can see how if the diffusion model is wrong somewhere, it could inappropriately penalize the policy for doing something reasonable. Only penalizing the loss in excess of the expert's loss mitigates the effect of errors in the diffusion model. The paper's strongest point is probably the experimental results in section 5.3. A reasonably diverse set of tasks were evaluated, and the proposed method performs significantly better than the baselines in most cases. I also appreciated that the experiments were targeted towards testing specific hypotheses—e.g., the results in section 5.5, which validate the hypotheses that the BC loss is necessary, and that the diffusion model is the best choice of EBM for this purpose. Significance: The relative simplicity of the approach, along with positive experimental results in non-trivial tasks, could encourage more research into combining BC with EBM-type losses.

Weaknesses

Originality: Although the specific combination of BC and DDPMs is novel, there is some precedent for combining BC losses with other losses—for example [A], which combines a BC loss with the GAIL loss. [A] Jena et al. Augmenting GAIL with BC for sample efficient imitation learning. https://arxiv.org/pdf/2001.07798.pdf Clarity: Some key aspects of the submission's presentation could be improved. In particular, the mathematical notation is lacking in key places. Section 4.2.2 is missing notation to describe where we are taking various expectations, and with respect to which distributions. For example, should eq. 5 read $E_{(s,a) \sim D, a \sim \pi(s)} \max(..., 0)$ or $\max (E_{(s,a) \sim D, a \sim \pi(s)} \dots, 0)$? Are any expectations taken with respect to $s$ sampled from the DM? I believe the answer is no, but I had to check the algorithm description in the appendix to be sure. The answers to these questions should be obvious from the math, but this is unfortunately not the case. Another point of confusion is in section 4.3, which mentions "jointly optimizing the proposed diffusion model loss." I generally interpret this to mean simultaneously optimizing the parameters of all the models, but this is not the case. I again had to check the algorithm description in the appendix to verify that the method is strictly a two-phase method, where the DM is pretrained and fixed before optimizing eq. 6. Quality: One of the paper's main weaknesses is confusion as to the motivation for the method and why the method seems to work. Some of the explanations offered by the paper include: 1. Diffusion models / EBMs generalize better because they model the joint probability p(s,a) as opposed to the conditional probability p(a|s) (lines 52-54). 2. Combining BC with an EBM helps because EBMs alone suffer from "manifold overfitting." (lines 127-128) Both of these explanations seem a bit tenuous. As for the first statement (EBMs generalize better because they model the joint probability), it is simultaneously vague, debatable, and it is unclear how that property would benefit the method, even if it were true. It is stated (line 113) that "These methods demonstrate superior generalization performance in diverse domains," but this statement seems too vague to be useful. Even assuming that modeling the joint probability produces better generalization in some sense, how does that translate to better results in the proposed method? No plausible mechanism is given for this. For example, I could imagine that if one were to sample joint state-action configurations from a good model, and then train behavioral cloning using those state-action pairs, then that could conceivably perform better than training BC on raw data, because sampling from generated states may serve as a form of data augmentation. However, this is not what the method does, according to the algorithm description in the appendix—states are sampled strictly from the data distribution. As for the statement about manifold overfitting, it is true that this is theoretically an issue with EBMs—however, there is a trivial solution for this that works fairly well, which is to perturb the training examples with a small amount of random noise to break any manifold constraints in the data. So, it is not necessary to add a BC loss to EBMs in order to solve the manifold overfitting problem. That said, I do believe it is plausible that combining the BC and DM losses helps, but not for any of the reasons above. A more plausible explanation for why the method works is that the loss optimizes (ignoring eq. 4) an approximation to "forward" KL divergence plus "reverse" KL divergence—since the DM loss is a lower bound on the log-likelihood of the "data," where the "data" in this case consists of samples from the model. The forward KL divergence encourages mode coverage, at the expense of occasionally producing bad samples, whereas the reverse KL has the opposite properties—it produces "good" samples, at the expense of losing modes. Therefore, these losses have complementary characteristics that benefit each other when combined. However, one has to be careful not to inadvertently penalize good samples in the reverse KL term, which is probably why adding the human baseline in equation 5, helps. Notably, imitation learning based on reverse KL divergence and more general f-divergences has been previously explored, e.g., in [B]. This would make one of the methods suggested in [B] a fair comparison in experiments. I noted a few other issues with the experiments. As stated earlier, I am suspicious of the idea that it is essential that the DM model the joint probability as opposed to the conditional probability—it seems plausible to me that the method would work just as well if the DM were trained as a conditional model (since the method doesn't use states sampled from the DM), so it would be interesting to test this experimentally. As also mentioned previously, I believe the manifold problem can be trivially solved by adding a small amount of random noise to the training examples. Doing so may significantly affect the results in Table 2—i.e., adding noise would probably boost the performance of the "without BC" column. Similar comments apply to the "manifold overfitting" results in section 5.4. I think this is an important experiment to run, because it may weaken one of the core claims, which is that adding BC to EBMs helps because it helps address the manifold overfitting problem. [B] Ke, Liyiming, et al. "Imitation learning as f-divergence minimization." Algorithmic Foundations of Robotics XIV: Proceedings of the Fourteenth Workshop on the Algorithmic Foundations of Robotics 14. Springer International Publishing, 2021.

Questions

How do the results change (esp. the manifold overfitting experiments) when the data is perturbed with a small amount of random noise? Have you tried comparing the method to BC with some sort of simple data augmentation? Have you considered comparing to another method that optimizes another type of f-divergence (as exemplified by [B])?

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

1 poor

Presentation

2 fair

Contribution

2 fair

Limitations

I didn't see much discussion in the way of limitations. What about the efficiency and stability of training?

Reviewer 2m353/10 · confidence 4/52023-07-06

Summary

The paper proposes a method to augment a behavior cloning (BC) agent with additional diffusion loss. The goal is to leverage the conditional probability learned by the BC loss and the joint probability learned by the diffusion loss. The diffusion loss is calculated using the prediction error of a pre-trained diffusion model on both the agent's and expert's state-action pairs. Experimental results demonstrate that this method outperforms baseline methods, achieving improved performance.

Strengths

1. Effective Experimental Results: The paper provides experimental results that showcase the superior performance of DBC compared to several baselines, including BC, Implicit BC, and Diffusion Policy. 2. Leveraging the Advantage of Offline Learning: The proposed method operates purely in an offline setting, eliminating the need for interaction with the environment. This reduction in training complexity is an advantage.

Weaknesses

1. Sensitivity to the $\lambda$ Parameter: The paper introduces a $\lambda$ parameter for BC loss and diffusion model loss balancing, but the analysis is limited to the maze environment. It is crucial to investigate and analyze the performance of the algorithm across various environments when the $\lambda$ parameter is modified. 2. Lack of Theoretical Analysis: The paper lacks a theoretical analysis of the proposed diffusion model loss. Providing a theoretical foundation for the diffusion loss would enhance the paper's scientific rigor.

Questions

1. Mechanism of Knowledge Transfer: The paper does not sufficiently explain how the diffusion model loss leads to agent's improved generalizability to unseen states. While the modeling the joint distribution of state-action pair can extrapolate/interpolate actions better than original BC on unseen states [1], it remains unclear how the knowledge generated by modeling the joint distribution is transferred to the agent. Since DBC trains the agent with the diffusion model loss on **expert states only**, it is essential to provide a theoretical proof or analysis to address this question. The authors should clarify the mechanism behind the transfer of knowledge from the diffusion model to the agent and explain how this transfer leads to improved performance on **unseen states**. 2. Redundancy of Diffusion Model Loss: The paper describes the proposed diffusion model loss as a measure of how well a state-action pair $(s^e,\hat{a})$ fits the expert state-action pair distribution, where $s^e$ is the state of $(s^e,a^e)$ sampled from the expert's state-action distribution and $\hat{a}\sim \pi(s^e)$ is the action predicted by agent. Additionally, the BC loss $||\hat{a}-a^e||^2$ is used to ensure that the agent's prediction $\hat{a}$ is as close to $a^e$ as possible. Theoretically, the state-action pair $(s^e,a^e)$ can fully minimize the diffusion model loss since it already exists in the expert state-action pair distribution. Consequently, minimizing the BC loss should also minimize the diffusion model loss on all training data (which is the expert data), potentially rendering the diffusion model loss redundant. The authors should address this observation and clarify the necessity of the diffusion model loss despite its potential redundancy. 3. Maximizing Joint Probability vs. Conditional Probability: The paper claims that the agent benefits from the diffusion model loss because it models the joint probability $p(s,a)$. However, maximizing the joint probability reduces to maximizing the conditional probability when the state $s$ is fixed, since the states in training data cannot be modified by the agent. The authors should clarify this discrepancy and provide further insights. 4. Performance in Other Locomotion Tasks: The paper's locomotion task evaluation focuses solely on the Walker2d environment. To establish the generalizability of the proposed method, it is crucial to conduct experiments on other well-established benchmarks, including Hopper, HalfCheetah, Ant, and Humanoid. By evaluating the performance across these mainstream locomotion tasks, the authors can provide a more comprehensive assessment and demonstrate the effectiveness of the proposed method in diverse and challenging environments.
 5. Minor Concern: Conducting experiments with only three random seeds may not provide sufficient statistical significance. Increasing the number of random seeds would strengthen the experimental evaluation. [1] Florence, Pete, et al. "Implicit behavioral cloning." Conference on Robot Learning. PMLR, 2022.

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

1. Limited Task Scope: The experimental evaluation is confined to specific tasks, raising questions about DBC's performance in other types of tasks or more complex environments. 2. Offline Learning Restriction: As an offline imitation learning algorithm, DBC lacks the capability to acquire knowledge through online interaction with the environment.

Senior Area Chairs2023-08-18

Please engage with author rebuttal

Please engage with the author rebuttal as soon as possible. It is critical for the review process to have engagement from all reviewers, as we cannot easily judge the validity of the original review, nor fairly calibrate across papers, without commentary on the rebuttal. Thank you.

Reviewer 2m352023-08-19

Thank you for your rebuttal. I appreciate the effort made in addressing my concerns; however, I still have reservations regarding the points raised. Though the diffusion model might have a better interpolating and extrapolating ability on “unseen states”. But in this work, the agent is not acting with the diffusion model. The agent is trained with the loss calculated by this diffusion model instead. And the agent is only minimizing the diffusion model loss on “expert states”. So it is still unclear how the agent can achieve better performance on “unseen states” after training with the diffusion loss. If training the agent with the BC loss solely and using the diffusion model loss as a performance measurement, will the diffusion model loss reach its optimal value when the BC loss converges? When training the agent with the diffusion model loss $||\hat{\epsilon}(s,\hat{a},n)-\epsilon||^2$, the state $s$ is fixed, the agent is just optimizing the conditional probability $\pi(a|s)$ on the expert’s state, and given that optimizing the BC loss solely can optimize the conditional probability $\pi(a|s)$. It is unclear what’s the difference between training the agent with these two loss in theoretical perspective.

Authorsrebuttal2023-08-21

> Though the diffusion model might have a better interpolating and extrapolating ability on “unseen states”. But in this work, the agent is not acting with the diffusion model. The agent is trained with the loss calculated by this diffusion model instead. And the agent is only minimizing the diffusion model loss on “expert states”. So it is still unclear how the agent can achieve better performance on “unseen states” after training with the diffusion loss. We thank the reviewer for initiating further discussion on the mechanism of knowledge transfer. We propose to transfer the generalization ability acquired by the diffusion model to the policy by optimizing our proposed diffusion model loss, which captures the joint distribution of state-action pairs $p(s, a)$. As noted in our rebuttal, the joint probability $p(s, a)$ can be represented as the product of the marginal state probability and the conditional action probability using the Bayes Rules, i.e., $p(s, a) = p(s)p(a|s)$. We observe that the policy can achieve better performance on unseen states with the distribution of expert states $p(s)$ taken into account, despite only observing the expert states during training. Extensive experiment results support the above observations. In Section 5.4, we show that the policy $\pi\_{\text{DM}}$ modeling the joint distribution $p(s, a)$ can generalize to unseen test regions while the policy $\pi\_{\text{BC}}$ modeling the conditional distribution $p(a|s)$ fails the task; In Section E, we show that our proposed DBC generalizes to unseen initial and target locations better than baselines that only model conditional distributions. Therefore, we believe that with the guidance of the diffusion model, which models the joint distribution $p(s, a)$ and considers the distribution of expert states $p(s)$, the generalization ability of the agent can be improved. We will revise the paper to make this clear. > If training the agent with the BC loss solely and using the diffusion model loss as a performance measurement, will the diffusion model loss reach its optimal value when the BC loss converges? As requested by the reviewer, we have conducted additional experiments investigating the trend of $\mathcal{L}\_{\text{BC}}$ and $\mathcal{L}\_{\text{DM}}$ when selectively optimizing either each one of them or both. Specifically, BC only optimizes $\mathcal{L}\_{\text{BC}}$, DM-only soley optimizes $\mathcal{L}\_{\text{DM}}$, and our proposed DBC optimizes $\mathcal{L}\_{\text{BC}} + \lambda \mathcal{L}\_{\text{DM}}$. In the following tables, we report the loss values of these three methods during training at 500, 1000, 1500, and 2000 epochs, on the Maze task. **BC** | #Epoch | 500 | 1000 | 1500 | 2000 | | --- | --- | --- | --- | --- | | $\mathcal{L}\_{\text{BC}}$ | 0.2403 (0.0200) | 0.1953 (0.0142) | 0.1829 (0.0129) | 0.1726 (0.0121) | | $\mathcal{L}\_{\text{DM}}$ | 0.0333 (0.0307) | 0.0285 (0.0262) | 0.0252 (0.0233) | 0.0228 (0.0217) | **DM-only** | #Epoch | 500 | 1000 | 1500 | 2000 | | --- | --- | --- | --- | --- | | $\mathcal{L}\_{\text{BC}}$ | 0.6067 (0.0349) | 0.5934 (0.0327) | 0.5297 (0.0266) | 0.5671 (0.0452) | | $\mathcal{L}\_{\text{DM}}$ | 0.0394(0.0214) | 0.0273 (0.0186) | 0.0196 (0.0190) | 0.0158 (0.0139) | **DBC (Ours)** | #Epoch | 500 | 1000 | 1500 | 2000 | | --- | --- | --- | --- | --- | | $\mathcal{L}\_{\text{BC}}$ | 0.3413 (0.0266) | 0.3127 (0.0268) | 0.2899 (0.0029) | 0.2698 (0.0174) | | $\mathcal{L}\_{\text{DM}}$ | 0.0401 (0.0342) | 0.0310(0.0247) | 0.0213 (0.0196) | 0.0162 (0.0256) | The results show that even BC only optimizes $\mathcal{L}\_{\text{BC}}$, $\mathcal{L}\_{\text{DM}}$ also reduces. However, $\mathcal{L}\_{\text{DM}}$ of BC converges to a higher value (0.0228), compared to only optimizing $\mathcal{L}\_{\text{DM}}$, where DM-only achieves a $\mathcal{L}\_{\text{DM}}$ value of 0.0158. On the other hand, our proposed DBC can effectively optimize both $\mathcal{L}\_{\text{BC}}$ and $\mathcal{L}\_{\text{DM}}$, demonstrating the compatibility of the two losses, which justifies the proposed combination of the two losses.

Authorsrebuttal2023-08-21

> When training the agent with the diffusion model loss, the state is fixed, the agent is just optimizing the conditional probability on the expert’s state, and given that optimizing the BC loss solely can optimize the conditional probability. It is unclear what’s the difference between training the agent with these two loss in theoretical perspective. Given **a single state $s$**, maximizing $\mathcal{L}_{\text{DM}}$ reduces to maximizing the conditional probability $p(a|s)$. However, we consider learning **an expert demonstration dataset $D$**, where the distribution of states $p(s)$ is unknown, so modeling $p(s)$ is still required and cannot be seen as a fixed constant. From a theoretical perspective, the joint probability $p(s, a)$ can be represented as the product of the marginal state probability and the conditional action probability using the Bayes Rules, i.e., $p(s, a) = p(s)p(a|s)$. In short, our $\mathcal{L}\_{\text{DM}}$ takes $p(s)$ into account to model the joint distribution while $\mathcal{L}\_{\text{BC}}$ optimizes $p(a|s)$ directly. To be more specific, MSE is a “non-learnable” measure, dubbed $\mathcal{L}\_{\text{BC}}$ in our paper. In contrast, the proposed diffusion model loss $\mathcal{L}\_{\text{DM}}$ is a learned distance, which measures the distance between an expert action $a$ and a predicted action $\hat{a}$ given a state $s$. As discussed in the paper and the rebuttal, while $\mathcal{L}\_{\text{BC}}$ and $\mathcal{L}\_{\text{DM}}$ both aim to bring the learner policy $\pi$ to the expert policy $\pi\_{\text{expert}}$, they are not identical. For example, $\mathcal{L}\_{\text{BC}}$ brings $\hat{a}$ closer to $a$ without considering the given state $s$; yet, $\mathcal{L}\_{\text{DM}}$ is calculated based on $(s, \hat{a})$ and $(s, a)$, which take $s$ into account. Note that despite their difference, when $\pi$ converges to $\pi\_{\text{expert}}$, both $\mathcal{L}\_{\text{BC}}$ and $\mathcal{L}\_{\text{DM}}$ converge to 0, indicating that these two losses are not conflicting. We thank the reviewer for the question. We will revise the paper to clarify this.

Reviewer uqof3/10 · confidence 4/52023-07-07

Summary

This paper presents a method for guiding behavior cloning via state-action joint distribution learning. They train a diffusion model to maximize the log-likelihood of state and action pairs in conjunction with an imitation learning model that learns to mimic expert actions given state observations. They combine the two objectives to obtain a policy that predicts actions that fit the expert joint probability captured by a diffusion model. The authors demonstrate experiments on four continuous control domains, and show ablations to demonstrate the effects of their design choices.

Strengths

1) The paper is well-structured and easy to follow. The experimental setup is clearly described with relevant details. The diffusion model for learning a joint distribution over state-action pairs is novel and is presented concisely. 2) The considered tasks for this method are challenging and results on 4 out of 5 continuous control tasks show improvements over the considered baselines.

Weaknesses

1) Intro: The authors state that implicit behavioral cloning (IBC) learns a joint distribution of state and action p(s,a). This is incorrect. IBC learns the joint “energy” E(s,a) but the learned distributions are still conditionals p(a|s). This is evident from the fact that IBC is trained to maximize the log-likelihood of actions in the dataset, and minimize those of sampled negative actions, given state inputs, and never trained to generate states or minimize the energy of negatively-sampled states. This makes a major claim of this paper incorrect. 2) Section 3.1: Talking about reinforcement learning as a preliminary seems very absurd and misleading for a subsection on imitation learning, and also when the paper has nothing to do with learning from rewards. 3) The core approach of this paper does not seem technically sound to me. Learning two distributions over the same random variables (state and action) seems to bring inconsistency to the probabilistic model where at least one of the two distributions is bound to be approximate.

Questions

1) According to the authors, what are some design choices that are critical to learning the joint over two variables of different modalities? 2) Table 2 shows the performance in Maze when using different generative models for guiding policy learning. Can the authors comment on the ranking of all methods in this table, why GANs performed the worst and Diffusion Models the best?

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

1 poor

Presentation

2 fair

Contribution

2 fair

Limitations

The authors have addressed the limitations of their method in the supplementary submission.

Senior Area Chairs2023-08-18

Please engage with author rebuttal

Please engage with the author rebuttal as soon as possible. It is critical for the review process to have engagement from all reviewers, as we cannot easily judge the validity of the original review, nor fairly calibrate across papers, without commentary on the rebuttal. Thank you.

Reviewer uqof2023-08-18

Response to rebuttal

I thank the authors for their rebuttal. Please find my response below. > Implicit BC The authors correctly pointed out the formulation of implicit BC by Florence et al., however, I would again like to make the same observation - while the energies are joint, the distribution on the action is a conditional p(a|s). As I pointed out in my review, IBC is trained to maximize the log-likelihood of actions in the dataset, and minimize those of sampled negative actions, given state inputs, and never trained to generate states or minimize the energy of negatively-sampled states. While the authors could decide to learn a joint distribution of actions and observations, using the findings of IBC as a major standing ground for their research is problematic. > Learning two distributions over the same variables (s, a) I note that Reviewer 2m35 raised a similar concern where they pointed out that both losses were theoretically redundant. Having read both the response to that question as well as mine, I absolutely do not align with the authors’ position. In the rebuttal, the authors mention that “BC loss brings the predicted actions and the expert actions closer in terms of a heuristic MSE distance” - MSE is not a heuristic distance, rather the negative log-likelihood of the data in parameterized normal distribution. Diffusion models also minimize the same distance but with a different parameterization. While I believe that minimizing both objectives separately can give us different results, the authors are unable to provide a technically sound argument as to why we expect any benefits at all using both objectives together. Simply pointing out that one is an MSE while the other a learned distance does not put much weight into the argument. Having pointed out theoretical irregularities between this paper and the existing literature, as well as missing technically sound arguments to defend the research, my assessment remains the same.

Authorsrebuttal2023-08-19

> Implicit BC We thank the reviewer for initiating further discussion on Implicit BC. We agree with the reviewer that Implicit BC learns from sampled negative actions given state inputs, and does not learn to generate states nor minimize the energy of negatively-sampled states. We will revise the paper to make this clear to the readers. Specifically, in our revision, we will motivate the proposed method based on the improved generalization performance of modeling state-action joint distribution shown in Section 5.4, where Figure 4 shows that $\pi\_\{\text{DM}}$ exhibits better interpolate and extrapolate ability. That said, we will include and discuss Implicit BC as a baseline, instead of mentioning it as an example of modeling joint distribution or using the findings of IBC as a standing ground for our proposed method. We sincerely thank the reviewer for this detailed feedback, which helps us clarify our work. > Learning two distributions over the same variables (s, a) We thank the reviewer for raising this point. We would like to apologize for misusing the term “heuristic”. To clarify, our intention was to emphasize that MSE is a “non-learnable” measure, dubbed $\mathcal{L}\_{\text{BC}}$ in our paper. In contrast, the proposed diffusion model loss $\mathcal{L}\_{\text{DM}}$ is a learned distance, which measures the distance between an expert action $a$ and a predicted action $\hat{a}$ given a state $s$. As discussed in the paper and the rebuttal, while $\mathcal{L}\_{\text{BC}}$ and $\mathcal{L}\_{\text{DM}}$ both aim to bring the learner policy $\pi$ to the expert policy $\pi\_{\text{expert}}$, they are not identical. For example, $\mathcal{L}\_{\text{BC}}$ brings $\hat{a}$ closer to $a$ without considering the given state $s$; yet, $\mathcal{L}\_{\text{DM}}$ is calculated based on $(s, \hat{a})$ and $(s, a)$, which takes $s$ into account. Note that despite their difference, when $\pi$ converges to $\pi\_{\text{expert}}$, both $\mathcal{L}\_{\text{BC}}$ and $\mathcal{L}\_{\text{DM}}$ converge to 0, indicating that these two losses are not conflicting. Our extensive experiments in the paper show that optimizing a combination of these two losses results in the best performance, compared to solely optimizing each of them. Table 1 shows that DBC ($\mathcal{L}\_{\text{BC}} + \lambda \mathcal{L}\_{\text{DM}}$) outperforms BC ($\mathcal{L}\_{\text{BC}}$). Table 2 shows that optimizing $\mathcal{L}\_{\text{BC}} + \lambda \mathcal{L}\_{\text{DM}}$ outperforms solely optimizing $\mathcal{L}\_{\text{DM}}$. Furthermore, Table 2 demonstrates that this is not specific to diffusion models; instead, combining a learned loss from an energy-based model (EBM) or a variational autoencoder (VAE) with a BC loss also achieves better performance than learning from an EBM or a VAE alone, or optimizing only $\mathcal{L}\_{\text{BC}}$. While we agree with the reviewer that our paper does not provide a theoretical ground for combining the two losses, we still firmly believe that our contributions are solid and novel, given extensive experimental results and analyses.

Reviewer Bk1n6/10 · confidence 3/52023-07-09

Summary

This paper presents a novel approach in the field of imitation learning. The authors address the challenge of learning from expert demonstrations without access to reward signals from the environment. They propose a framework called Diffusion Model-Augmented Behavioral Cloning (DBC) that combines the benefits of modeling both the conditional and joint probability of the expert distribution. The authors demonstrate the effectiveness of DBC through extensive experiments in various continuous control tasks, including navigation, robot arm manipulation, dexterous manipulation, and locomotion. However, there are certain strengths, weaknesses, limitations, and questions that need to be addressed regarding this work.

Strengths

1) Extensive experiments: The article presents a wide range of experiments conducted on diverse tasks, including navigation, robot arm manipulation, dexterous manipulation, and locomotion. This comprehensive evaluation demonstrates the effectiveness of DBC in various scenarios. 2) Novel approach: The proposed DBC framework offers a promising approach to imitation learning. By combining behavioral cloning with a diffusion model, the authors achieve more stable training compared to methods that combine behavioral cloning with GANs, such as GAIL. 3) Clear comparison: The authors compare DBC with existing baselines, providing a clear understanding of its advantages over other methods in terms of performance and generalization.

Weaknesses

1) Coefficient selection: The combination of the behavioral cloning loss and diffusion model loss in DBC relies on a simple addition of coefficients. However, the sensitivity of these coefficients to different environments should be further investigated. A more elegant approach and thorough ablation experiments on diverse tasks are needed to validate the coefficient selection process. Currently, this paper only provide ablation on this coefficient in a single environment (Maze). 2) Lack of comprehensive comparison: Although DBC is compared with baselines, the article does not provide a comprehensive comparison with other state-of-the-art methods in the field of imitation learning. It would be beneficial to evaluate DBC against a wider range of approaches, including GAIL and other recent advancements, to gain a more comprehensive understanding of its relative performance. This would provide a clearer perspective on the strengths and weaknesses of DBC in comparison to alternative methods.

Questions

1) How does the proposed DBC framework handle scenarios with complex and high-dimensional state spaces? Are there any specific limitations or challenges encountered in such cases? 2) Are there any plans to extend the evaluation of the diffusion model to other environments beyond the MAZE setting? How does the diffusion model compare to other generative models in terms of performance and efficiency?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

While the article highlights improved generalization performance, there is limited analysis or discussion on the factors contributing to this improvement. A deeper analysis of the generalization capabilities and limitations of DBC would enhance the understanding of its strengths and weaknesses.

Reviewer Q6yg2023-08-11

Thank you for the rebuttal

I thank the authors for the rebuttal. All of my concerns have been addressed.

Reviewer Pa842023-08-17

Rebuttal response

Thank you for trying the experiment to add noise to address the manifold overfitting issue. However, I don't find the results that convincing—I think either one or both of the following hypotheses could explain the fact that adding noise only degraded performance: 1. The reason for the poor performance is not manifold overfitting, so adding noise does not help. 2. The magnitude of the added noise was too great. Ideally, one should try a variety of noise levels and select the one that yields the best performance. Looking at figures 1b and 2b in the rebuttal PDF, I would expect that the results of $\pi_{DM}$ would vary continuously as a function of the noise level. The fact that there is a huge discrepancy between the two results would seem to indicate that the noise level is too high. If manifold overfitting is the problem, I would be fairly surprised if adding noise would not help. Manifold overfitting arises because the learner receives unbounded benefit in log-likelihood from squeezing the model PDF onto the manifold, regardless of whether it fits the data well otherwise. Adding noise places an upper limit on the amount of benefit we can receive from solely squeezing the PDF onto the manifold. So, I don't see why it shouldn't solve the problem.

Authorsrebuttal2023-08-19

Re: Rebuttal response

### Addressing Manifold Overfitting by Adding Noise to Expert Data **Modeling Expert Distribution** As suggested by the reviewer, we have conducted additional experiments to extensively evaluate diffusion models trained with various levels of noise added to the expert actions from the spiral dataset. Since providing result figures is not allowed during this phase, we calculate the average MSE distance between expert actions and the reconstruction of the trained diffusion models, which indicates how well diffusion models capture the expert distribution. We report the result below. | Noise level | 0 | 0.002 | 0.005 | 0.01 | 0.02 | 0.05 | 0.1 | | --- | --- | --- | --- | --- | --- | --- | --- | | MSE distance | 0.0162 | 0.0152 | **0.0148** | 0.0164 | 0.0188 | 0.0300 | 0.0491| We observe that applying a noise level of 0.005 results in the lowest MSE distance (0.0148) and indicates that the diffusion model best fits the expert data with this magnitude of injected noises, which is slightly better than directly modeling the original distribution (i.e., with a noise level of 0). This finding aligns with the reviewer’s statement (i.e., adding noise to expert distribution can alleviate the manifold overfitting issue). However, we would like to point out that adding noise still does not entirely address the issue, and there is still a discrepancy between the learned and expert distributions. **Guiding Policy Learning** Then, we investigate the performance of using the learned diffusion models to guide policy learning. Specifically, we train policies to optimize the diffusion model loss $\mathcal{L}\_{\text{DM}}$ provided by either the diffusion model learning from a noise level of 0 and the diffusion model learning from a noise level of 0.005, dubbed $\pi\_{\text{DM-0}}$ and $\pi\_{\text{DM-0.005}}$, respectively. We evaluate the performance of the policies by rolling out each policy and calculating the distance between the end location of the policy and the expert end location. A policy rollout is considered successful if the distance is not greater than 0.1. We report the average success rate of each policy over 100 episodes below. | Method | $\pi\_{\text{DM-0}}$ | $\pi\_{\text{DM-0.005}}$ | $\pi\_{\text{BC}}$ | | --- | --- | --- | --- | | Success Rate | 27% | 44% | 93% | The result suggests that the diffusion model learning from expert distribution added with a preferable magnitude noise can also better guide policy learning, achieving a success rate of 44%, outperforming the original diffusion model that suffers more from the manifold overfitting with a success rate of 27%. Yet, directly learning to model the conditional probability (i.e., $\pi\_{\text{BC}}$) achieves a much higher success rate of 93%. This result verifies the advantage of modeling the conditional probability on this task, which motivates us to incorporate $\mathcal{L}\_{\text{BC}}$ in our proposed learning objective, instead of solely optimizing $\mathcal{L}\_{\text{DM}}$. We sincerely thank the reviewer for pointing out the importance of selecting a preferable magnitude of noises and helping us better investigate the mechanism of our proposed method. We will revise our paper to incorporate these findings. ### Other Concerns Finally, since the reviewer’s response only discusses the manifold overfitting experiment, we would like to know if other concerns in the original reviews were addressed by our rebuttal. If not, we would love to have the chance to discuss further with the reviewers. Thank you very much.

Reviewer Bk1n2023-08-18

Thanks for the the authors' thoughtful response

Thanks for the the authors' thoughtful response. My queries have been clarified, and I believe the response was well-articulated.

Program Chairsdecision2023-09-21

Decision

Reject

© 2026 NYSGPT2525 LLC