Diffusion Imitation from Observation

We sincerely thank the reviewer for the thorough and constructive comments. Please find the response to your questions below.

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] [2] It would be convincing to see comparisons with this line of work.

We thank the reviewer for providing these references. We will revise the paper to discuss these works. As requested by the reviewer, we additionally implemented Optimal Transport (OT) [1]. OT uses proxy rewards derived from Sinkhorn distances rather than directly obtaining rewards from the raw logits of the discriminator, which enhances stability compared to AIL methods.

We report the result in Figure R.1 in the PDF file attached to the rebuttal summary, which shows that our method consistently outperforms OT across all evaluated tasks. We hypothesize that this is because OT computes distances at the trajectory level rather than the transition level, which requires monotonic trajectories. Consequently, OT performs well in environments like Walker and AdroitDoor, where trajectory variety is limited. However, it struggles in tasks with diverse trajectories, such as navigation, where the initial and goal states vary significantly. In contrast, our method generates rewards at the transition level, allowing us to identify transition similarities even when facing substantial trajectory variability. This flexibility enables our method to succeed in more complex environments with diverse trajectories.

We will revise the paper to include the results and the discussion.

It would be great to see more evaluations on image-based environments.

As requested by the reviewer, we additionally conducted experiments on the drawer close task introduced in Meta-World [2] using image-based states. This table-top manipulation task requires the agent to control a Sawyer robotic arm to close a drawer. An illustration of this task is shown in Figure R.5(a) in the PDF file attached to the rebuttal summary.

Figure R.5(b) presents the learning efficiency of our proposed method and the baselines. Our proposed method achieves an 80% success rate with only 70k environment interactions, outperforming BC, BCO, GAIfO, WAILfO, AIRLfO, IQ-Learn, and DePO.

We will revise the paper to include the results of this new image-based environment.

It seems unclear to me why behavioral cloning, with access to ground truth action labels, is doing worse than LfO methods in Figure 3.

While BC has access to ground truth action, it may suffer from compounding errors, i.e., the accumulation of errors from small initial deviations, caused by covariate shifts [3, 4, 5]. Because BC relies solely on learning from the observed expert dataset, unlike the LfO methods that utilize online interaction with environments, BC is susceptible to accumulating errors from states with small deviations, ultimately reaching an irrecoverable point. Therefore, BC is well-known for requiring a substantial amount of expert data to achieve coverage of the dataset and reduce unseen states. However, under our experiment setting with limited expert data, BC doesn’t have sufficient data to generalize to unseen states, leading to poor performance.

Does DIFO work with negative samples generated offline, e.g. adding Gaussian noise to the previous state, instead of policy online rollouts?

We thank the reviewer for the insightful idea. This idea of generating and utilizing negative samples offline for imitation learning resembles the idea of Implicit Behavior Cloning (Implicit BC) [6]. Implicit BC learns the joint probability of expert state-action pairs $p(s, a)$ via contrastive learning using expert data and offline generated negative data without policy online rollouts.

We believe it is possible to generate negative data offline by adding Gaussian noise to the previous state, as suggested by the reviewer, and learning our diffusion classifier to classify expert state transitions and the generated state transitions. However, unlike Implicit BC, our work focuses on learning from observation, where we do not have access to action labels in expert demonstrations. Hence, without online interactions, it would be impossible to know the action space, let alone learn a policy. Thus, we believe that online interactions are essential for policy learning.

References

[1] Papagiannis et al. "Imitation learning with Sinkhorn Distances." In ECML PKDD, 2022.

[2] Yu et al. "Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning." In CoRL, 2019.

[3] Ross et al. “Efficient Reductions for Imitation Learning.” In AISTATS, 2010.

[4] Ross et al. “A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning.” In AISTATS, 2011.

[5] Laskey et al. “SHIV: Reducing supervisor burden in DAgger using support vectors for efficient learning from demonstrations in high dimensional state spaces.” In ICRA, 2016.

[6] Florence et al. “Implicit Behavioral Cloning.” In CoRL, 2022.

We sincerely thank the reviewer for acknowledging our rebuttal and helping us improve our submission.

We sincerely thank the reviewer for the thorough and constructive comments. Please find the response to your questions below.

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.

Diffusion models have shown their ability to become classifiers with diffusion loss [1]. DIFO-NA demonstrates using diffusion loss from an offline pre-trained model provides valid rewards, as shown in Figure 3(a), which indicates diffusion loss could be a reasonable metric for the discriminator. However, as the agent continues to learn, it will generate transitions that increasingly resemble those of the expert, causing the offline model to lose its ability to provide precise rewards eventually. Hence, we integrate the AIL framework to improve the discriminator simultaneously with the policy.

It appears more like using diffusion loss as regularization for discriminator training. Do other regularization techniques lead to similar improvements?

WAILfO builds on GAIfO's work by incorporating gradient penalty as a regularization to improve performance. Additionally, our implementations of GAIfO, AIRLfO, and WAILfO already employ L2 regularization. The experimental results show that our proposed method outperforms these methods with various regularizations. We would like to highlight that these regularizations could also be applied to our method.

Why are diffusion models more suitable in this setting (an intuitive explanation)?

We hypothesize that there are two reasons why our diffusion model classifier performs better than the GAIL discriminator. First, the instability of GAIL arises from the tendency of discriminators to overfit, resulting in the policy's inability to learn from discriminator rewards. Unlike GAIL's MLP binary classifier, which maps a high-dimensional input to a one-dimensional logit, our diffusion model learns to predict high-dimensional noises, which is inherently more difficult to overfit. Second, diffusion models excel at modeling multimodal distributions and thus can outperform GAIL's discriminator when expert demonstrations exhibit significant variability. We will revise the paper to include these intuitions.

More recent LfO baselines should be compared and discussed [1, 2].

We thank the reviewer for these references. We will revise the paper to discuss these works. As requested by the reviewer, we additionally included the results of two recent methods: Decoupled Policy Optimization (DePO) [2] and Optimal Transport (OT) [3] provided by Reviewer P9FE, which is considered a more general version compared to ADS [4] provided by Reviewer Dab5. We report the results in Figure R.1 in the PDF file attached to the rebuttal summary. The results show that our proposed method outperforms both DePO and OT.

DePO decouples the policy into a high-level state planner and an inverse dynamics model (IDM), utilizing embedded decoupled policy gradient and generative adversarial training. The decouple feature makes the planner and IDM transferable across different domains. However, it still suffers from the same challenges as the ones in GAIL, such as overfitting and the inability to model multimodality. Moreover, it also requires training with an IDM, making it more difficult to tune.

OT uses proxy rewards derived from Sinkhorn distances rather than directly obtaining rewards from the raw logits of the discriminator, which enhances stability compared to AIL methods. However, OT computes distances at the trajectory level rather than the transition level, which requires monotonic trajectories. Consequently, OT performs well in environments like Walker and AdroitDoor, where trajectory variety is limited. However, it struggles in tasks with diverse trajectories, such as navigation, where the initial and goal states vary significantly. In contrast, our method generates rewards at the transition level, allowing us to identify transition similarities even when facing substantial trajectory variability. This flexibility enables our method to succeed in more complex environments with diverse trajectories.

We will include the results and discussion of these two methods in the revised paper.

The diffusion discriminator is different from prior methods in that we can sample from the distribution 𝐷(𝑠,𝑠′). Can the resulting discriminator itself be used as a policy when combined with an inverse dynamics model 𝑝(𝑎|𝑠,𝑠′)?

As suggested by the reviewer, we believe that combining an ideal inverse dynamics model (IDM) with our diffusion discriminator could indeed result in an effective policy. This potential is illustrated by our successful trajectory generation in PointMaze experiments (Section 5.5 and Figure 5), where our diffusion discriminator shows the capability to accurately predict the next state based on the current state.

However, obtaining a high-quality IDM is challenging, as discussed in lines 27-29. Training IDMs requires data that is well-aligned with the expert's data distribution, while, collecting such data relies on having an effective policy, creating a deadlock. Even if this problem is solved, the planner may also produce invalid transitions for the IDM, posing challenges for both planners and IDMs.

Given these difficulties, our method directly learns a policy with the rewards produced by the diffusion discriminator instead of addressing the difficulties of obtaining an ideal IDM. We will revise the paper to make this clear.

References

[1] Li et al. "Your diffusion model is secretly a zero-shot classifier." In ICCV, 2023.

[2] Liu et al. "Plan your target and learn your skills: Transferable state-only imitation learning via decoupled policy optimization." In ICML, 2022.

[3] Papagiannis et al. "Imitation learning with Sinkhorn Distances." In ECML PKDD, 2022.

[4] Liu et al. “Imitation Learning from Observation with Automatic Discount Scheduling.” In ICLR, 2024.

We sincerely thank the reviewer for the thorough and constructive comments. Please find the response to your questions below.

There are other methods [1, 2, 3] that use metrics, i.e. likelihood, ELBO and entropy, from generative model for ILfO. I suggest discussing them in the related works.

We thank the reviewer for providing the references. We will revise the paper to discuss these methods using metrics from generative models for ILfO.

How many samples of and are needed to compute the discriminator score in eq. 4? Is the value stable across different samples?

We sample a single denoising timestep (n=1) to compute the reward. As suggested by the reviewer, to investigate the robustness of our rewards with a varying number of timestep samples, we conducted experiments in PointMaze by averaging different numbers of timestep samples (n=1, 2, 5, 10) as rewards. The result presented in Figure R.3 in the PDF file attached to the rebuttal summary shows that our method can learn smoothly with a single sample.

To verify the stability of the rewards, we present the standard deviation to mean ratio of the rewards from 500 timesteps across different numbers of timestep samples in the table below. The values are averaged from a batch of 4096 transitions. The result below suggests that the reward is stable.

Although you are motivated to improve the robustness of AILfO, I don't see a direct connection that using a diffusion model instead of an FFN should make things more stable

We hypothesize that the instability of most existing AIL frameworks arises from the tendency of FFN discriminators to overfit, resulting in the policy's inability to learn from discriminator rewards. Unlike FFN binary classifiers, which map a high-dimensional input to a one-dimensional logit, our diffusion model learns to predict high-dimensional noises, which is inherently more difficult to overfit.

Furthermore, expert demonstrations typically exhibit significant variability, which could be difficult for FFNs to model effectively [1]. Diffusion models, however, are adept at handling such multimodality, providing a more robust approach for capturing the diverse patterns present in expert behavior. We will revise the paper to include these insights.

While DIFO demonstrates sample efficiency in terms of expert demonstrations, it lacks efficiency in environment interactions.

We respectfully disagree with the reviewer. Figure 3 in the main paper shows that our method, DIFO, is very efficient in environmental interactions, i.e., sample efficiency. In PointMaze, DIFO only requires about 175k environment steps to achieve a 60% success rate, while the best-performing baseline AIRLfO needs over 400k steps. Similarly, in AdroitDoor, DIFO achieves a 75% success rate with only 4M steps, yet the best-performing baseline WAILfO needs around 10M steps. That said, DIFO is very sample-efficient.

The table below presents the number of environment steps required for each method to achieve 50% expert performance. The result shows that our method, DIFO, can achieve the same performance with significantly fewer environmental steps (environment interactions). That said, DIFO is very sample-efficient.

The experimental scope is limited, with tests conducted on only six environments and comparisons made against relatively older baselines.

We conducted extensive experiments to evaluate our method across various aspects.

• The six test environments we selected cover a wide range of domains, including navigation (PointMaze and AntMaze), manipulation (FetchPush and AdroitDoor), locomotion (Walker), and games (CarRacing). In contrast, most existing works [2, 3, 4, 5] focus only on locomotion tasks.
• The tasks included in our work cover both vectorized-state-based environments (PointMaze, AntMazmethod'sPush, AdroitDoor, and Walker) and image-based environments (CarRacing and DrawerClose), while most existing works [2, 3, 4, 5, 6, 7, 8, 9] only consider vectorized-state-based environments.
• Our work analyzes the data efficiency by varying the amount of available expert demonstrations and evaluating the learning performance of all methods in Section 5.4.
• We designed and conducted a toy environment to visualize the reward functions learned by our proposed method and GAIfO.

More baselines

Please see the response to Reviewer Dab5.

References

[1] Li et al. "Infogail: Interpretable imitation learning from visual demonstrations." In NIPS, 2017.

[2] Liu et al. "Plan your target and learn your skills: Transferable state-only imitation learning via decoupled policy optimization." ICML, 2022.

[3] Papagiannis et al. "Imitation learning with sinkhorn distances." In ECML PKDD, 2022.

[4] Ni et al. "f-irl: Inverse reinforcement learning via state marginal matching." In CoRL, 2021.

[5] Liu et al. “CEIL: Generalized Contextual Imitation Learning.”. In NeurIPS, 2023.

[6] Garg et al. "Iq-learn: Inverse soft-q learning for imitation." In NeurIPS, 2021.

[7] Fu et al. "Learning robust rewards with adversarial inverse reinforcement learning." In ICLR, 2018.

[8] Ho et al. "Generative adversarial imitation learning." In NIPS, 2016.

[9] Xiao et al. "Wasserstein adversarial imitation learning." arXiv, 2019.

We sincerely thank the reviewer for the thorough and constructive comments. Please find the response to your questions below.

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.

We thank the reviewer for the question. The objective of RL is to maximize the overall return. The RL algorithm should converge to the optimal policy as long as the reward function produces the highest reward at the optimal transition. A smoother contour helps guide the RL algorithm without misguiding the final policy. Our method's reward peaks align with expert distribution (see Figure 7(a) in the main paper), ensuring accurate guidance.

Moreover, under the AIL framework, the reward function evolves alongside the policy. The reward function becomes sparser as the policy improves and eventually converges to a relatively sparse condition. Detailed theoretical guarantees can be found in the GAIL paper [1]. Besides, many previous imitation learning and inverse RL papers visualize the reward function with similar plots like ours [2, 3, 4, 5]. Hence, we believe that this approach is widely accepted for evaluating reward functions.

Lastly, we would like to highlight that we evaluate DIFO in diverse domains, including navigation, manipulation, locomotion, and control, in both state- and image-based environments. Our evaluation results demonstrate that DIFO is broadly applicable and not restricted to specific environments.

Why do we only learn the MSE on the expert data but not also the agent data?

Please see the overall response.

Is the expert stochastic? Does the expert policy exhibit diverse behaviors?

All of our expert data incorporates stochasticity to enhance diversity in trajectories. In the PointMaze environment, our expert trajectory data is generated using a breadth-first search planner with added stochasticity to create different path choices, resulting in various trajectories for the same start and end positions. In other environments, we add a small amount of noise to the experts' actions. These methods provide stochasticity and multimodality in expert behaviors. Our model effectively learns from such expert demonstrations, demonstrating its capability to handle multimodal distributions.

Diffusion models are usually used for modeling multimodal distribution.

As discussed in Section 5.5 and Section D of the main paper, we use a trained diffusion model to generate maze trajectories. As previously mentioned, our expert data exhibits a certain degree of multimodality, meaning that there are different paths in the dataset for the same start and goal locations. The generated results support our observations, as our model effectively produces multimodal possible paths. This demonstrates the benefit of modeling multimodal distributions using diffusion models.

Is it more convincing to have stochastic environments rather than purely deterministic environments?

Please see the overall response.

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.

Our method builds upon the established theories of Diffusion Classifier [6] and GAIL [1], as described in Section 4. The Diffusion Classifier paper shows that the diffusion loss (Equations 2 and 3) approximates ELBO, effectively transforming a diffusion model into a classifier. GAIL proves that IRL is a dual problem of an occupancy measure matching problem. Consequently, running RL with the discriminator IRL reward, which recovers the primal optimum from the dual optimum, ensures that the resulting policy is the primal optimum. As a result, by adversarially training the discriminator and the policy, we can match the distribution of the agent's policy to that of the expert. By incorporating the diffusion model as the discriminator into the AIL framework, our method leverages the strengths of the diffusion model while maintaining optimality guarantees. We will revise the paper to include these details.

We believe our experiments are comprehensive. Please see the response to Reviewer AzJp.

Tuning the lambda seem difficult as they can vary by ~20% success rate according to figure 7.

Figure 7 aims to investigate the effect of $\lambda$, and therefore, we experimented with a very wide range of values. The performance only drops by 20% when the lambda is increased or decreased by two orders of magnitude from the optimal value we identify. In Figure 7(a), the performance of $\lambda=0.1$ and $0.01$ is comparable, with overlapping standard deviations, indicating no statistically significant difference between them. Similarly, in Figure 7(b), although the performance appears to vary slightly, the standard deviations overlap. Therefore, we believe DIFO is not sensitive to $\lambda$. Any $\lambda$s in a reasonable range [0.1~0.001] can yield good performance.

References

[1] Ho et al. "Generative adversarial imitation learning." In NIPS, 2016.

[2] Garg et al. "Iq-learn: Inverse soft-q learning for imitation." In NeurIPS, 2021.

[3] Fu et al. "Learning robust rewards with adversarial inverse reinforcement learning." In ICLR, 2018.

[4] Ni et al. "f-irl: Inverse reinforcement learning via state marginal matching." In CoRL, 2021.

[5] Liu et al. "Energy-based imitation learning." In AAMAS, 2021.

[6] Li et al. "Your diffusion model is secretly a zero-shot classifier." In ICCV, 2023.