Diffusion-Reward Adversarial Imitation Learning

Authors: chun-mao  lai
hsiang-chun  wang
ping-chun  hsieh
yu-chiang frank wang
min-hung  chen
shao-hua  sun
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-03 14:32:46

Abstract: Imitation learning aims to learn a policy from observing expert demonstrations without access to reward signals from environments. Generative adversarial imitation learning (GAIL) formulates imitation learning as adversarial learning, employing a generator policy learning to imitate expert behaviors and discriminator learning to distinguish the expert demonstrations from agent trajectories. Despite its encouraging results, GAIL training is often brittle and unstable. Inspired by the recent dominance of diffusion models in generative modeling, we propose Diffusion-Reward Adversarial Imitation Learning (DRAIL), which integrates a diffusion model into GAIL, aiming to yield more robust and smoother rewards for policy learning. Specifically, we propose a diffusion discriminative classifier to construct an enhanced discriminator, and design diffusion rewards based on the classifier’s output for policy learning. Extensive experiments are conducted in navigation, manipulation, and locomotion, verifying DRAIL’s effectiveness compared to prior imitation learning methods. Moreover, additional experimental results demonstrate the generalizability and data efficiency of DRAIL. Visualized learned reward functions of GAIL and DRAIL suggest that DRAIL can produce more robust and smoother rewards. Project page: https://nturobotlearninglab.github.io/DRAIL/

Main paper: 
 introduction: 
 Imitation learning, i.e., learning from demonstration [24, 41, 49] , aims to acquire an agent policy by observing and mimicking the behavior demonstrated in expert demonstrations. Various imitation learning methods [53, 60] have enabled deploying reliable and robust learned policies in a variety of tasks involving sequential decision-making, especially in the scenarios where devising a reward function is intricate or uncertain [7, 32, 34] , or when learning in a trial-and-error manner is expensive or unsafe [14, 17] .Among various methods in imitation learning, generative adversarial imitation learning (GAIL) [21] has been widely adopted due to its effectiveness and data efficiency. GAIL learns a generator policy to imitate expert behaviors through reinforcement learning and a discriminator to differentiate between the expert and the generator's state-action pair distributions, resembling the idea of generative adversarial networks (GANs) [16] . Despite its established theoretical guarantee, GAIL training is notoriously brittle and unstable. To alleviate this issue, significant efforts have been put into improving GAIL's sample efficiency, scalability, robustness, and generalizability by modifying loss functions [12] , developing improved policy learning algorithms [30] , and exploring various similarity measures of distributions [2, 8, 12] .Inspired by the recent dominance of diffusion models in generative modeling [22] , this work explores incorporating diffusion models into GAIL to provide more robust and smoother reward functions for policy learning as well as stabilize adversarial training. Specifically, we propose a diffusion discriminative classifier, which learns to classify a state-action pair into expert demonstrations or agent trajectories with merely two reverse diffusion steps. Then, we leverage the proposed diffusion discriminative classifier to devise diffusion rewards, which reward agent behaviors that closely align with expert demonstrations. Putting them together, we present Diffusion-Reward Adversarial Imitation Learning (DRAIL), a novel adversarial imitation learning framework that can efficiently and effectively produce reliable policies replicating the behaviors of experts.We extensively compare our proposed framework DRAIL with behavioral cloning [44] , Diffusion Policy [6, 42] , and AIL methods, e.g., GAIL [21] , WAIL [2] , and DiffAIL [59] , in diverse continuous control domains, including navigation, robot arm manipulation, locomotion, and games. This collection of tasks includes environments with high-dimensional continuous state and action spaces, as well as covers both vectorized and image-based states. The experimental results show that our proposed framework consistently outperforms the prior methods or achieves competitive performance. Moreover, DRAIL exhibits superior performance in generalizing to states or goals unseen from the expert's demonstrations. When varying the amounts of available expert data, DRAIL demonstrates the best data efficiency. At last, the visualized learned reward functions show that DRAIL captures more robust and smoother rewards compared to GAIL.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: Our introduction and abstract accurately outline the contributions of the paper, which include proposing a novel framework called Diffusion-Reward Adversarial Imitation Learning (DRAIL). We address the limitations of existing methods, particularly the brittleness and instability of generative adversarial imitation learning (GAIL), by integrating diffusion models into the learning process. Our proposedDRAIL framework aims to provide more robust and smoother rewards for policy learning while enhancing stability in adversarial training. The abstract and introduction clearly articulate these contributions, aligning with the experimental results and analyses presented in the paper.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
related work: 
 Imitation learning enables agents to learn from expert demonstrations to acquire complex behaviors without explicit reward functions. Its application spans various domains, including robotics [49, 60] , autonomous driving [36] , and game AI [18] .Behavioral Cloning (BC). BC [44, 55] imitates an expert policy through supervised learning without interaction with environments and is widely used for its simplicity and effectiveness across various domains. Despite its benefits, BC struggles to generalize to states not covered in expert demonstrations because of compounding error [10, 48] . Recent methods have explored learning diffusion models as policies [6, 42] , allowing for modeling multimodal expert behaviors, or using diffusion models to provide learning signals to enhance the generalizability of BC [5] . In contrast, this work aims to leverage a diffusion model to provide learning signals for policy learning in online imitation learning.Inverse Reinforcement Learning (IRL). IRL methods [37] aim at inferring a reward function that could best explain the demonstrated behavior and subsequently learn a policy using the inferred reward function. Nevertheless, inferring reward functions is an ill-posed problem since different reward functions could induce the same demonstrated behavior. Hence, IRL methods often impose constraints on reward functions or policies to ensure optimality and uniqueness [1, 37, 54, 61 ]. Yet, these constraints could potentially restrict the generalizability of learned policies.Adversarial Imitation Learning (AIL). AIL methods aim to directly match the state-action distributions of an agent and an expert through adversarial training. Generative adversarial imitation learning (GAIL) [21] and its extensions [23, 25, 31, 56, 59, 62] train a generator policy to imitate expert behaviors and a discriminator to differentiate between the expert and the generator's state-action pair distributions, which resembles the idea of generative adversarial networks (GANs) [16] . Thanks to its simplicity and effectiveness, GAIL has been widely applied to various domains [3, 28, 45] . Over the past years, researchers have proposed numerous improvements to enhance GAIL's sample efficiency, scalability, and robustness [40] , including modifications to discriminator's loss function [12] , extensions to off-policy RL algorithms [30] , addressing reward bias [31] , and exploration of various similarity measures [2, 8, 11, 12] . Another line of work avoids adversarial training, such as IQ-Learn [15] , which learns a Q-function that implicitly represents the reward function and policy. In this work, we propose to use the diffusion model as a discriminator in GAIL.Diffusion Model-Based Approaches in Reinforcement Learning. Diffuser [57] and Nuti et al. [39] apply diffusion models to reinforcement learning (RL) and reward learning, their settings differ significantly from ours. Diffuser [57] is a model-based RL method that requires trajectory-level reward information, which differs from our setting, i.e., imitation learning, where obtaining rewards is not possible. Nuti et al. [39] focus on learning a reward function, unlike imitation learning, whose goal is to obtain a policy. Hence, Nuti et al. [39] neither present policy learning results in the main paper nor compare their method to imitation learning methods. Moreover, they focus on learning from a fixed suboptimal dataset, AIL approaches and our method are designed to learn from agent data that continually change as the agents learn.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We explain several limitations of our proposed DRAIL framework in Appendix H. These include its restriction to learning from demonstration (LfD) settings, the requirement for state and action sequences, the assumption of optimal expert demonstrations, and its inability to leverage environmental rewards.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully. (c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
preliminaries: 
 We propose a novel adversarial imitation learning framework that integrates a diffusion model into generative adversarial imitation learning. Hence, this section presents background on the two topics.GAIL [21] establishes a connection between generative adversarial networks (GANs) [16] and imitation learning. GAIL employs a generator, G θ , that acts as a policy π θ , mapping a state to an action. The generator aims to produce a state-action distribution (ρ π θ ) which closely resembles the expert state-action distribution ρ π E ; discriminator D ω functions as a binary classifier, attempting to differentiate the state-action distribution of the generator (ρ π θ ) from the expert's (ρ π E ). The optimization equation of GAIL can be formulated using the Jensen-Shannon divergence, which is equivalent to the minimax equation of GAN. The optimization of GAIL can be derived as follows:EQUATIONwhere ρ π θ and ρ π E are the state-action distribution from an agent π θ and expert policy π E respectively. The loss function for the discriminator is stated as-(E x∼ρπ θ [log D ω (x)]+E x∼ρπ E [log(1-D ω (x))]).For a given state, the generator tries to take expert-like action; the discriminator takes state-action pairs as input and computes the probability of the input originating from an expert. Then the generator uses a reward function-E x∼ρπ θ [log D ω (x)] or -E x∼ρπ θ [log D ω (x)] + λH(π θ )to optimize its network parameters, where the entropy term H is a policy regularizer controlled by λ ≥ 0.x TReverse diffusion process q(x t | x t-1 )p ϕ (x t-1 | x t )Forward diffusion processx 0 x t x t-1Learn to denoise Adding Gaussian noise Diffusion models have demonstrated stateof-the-art performance on various tasks [9, 29, 38, 51] . This work builds upon denoising diffusion probabilistic models (DDPMs) [22] that employ forward and reverse diffusion processes, as illustrated in Figure 1 . The forward diffusion process injects noise into data points following a variance schedule until achieving an isotropic Gaussian distribution. The reverse diffusion process trains a diffusion model ϕ to predict the injected noise by optimizing the objective:EQUATIONwhere T represents the set of discrete time steps in the diffusion process, ϵ is the noise applied by the forward process, ϵ ϕ is the noise predicted by the diffusion model, and ᾱt is the scheduled noise level applied on the data samples.Beyond generative tasks, diffusion models have also been successfully applied in other areas, including image classification and imitation learning. Diffusion Classifier [35] demonstrates that conditional diffusion models can estimate class-conditional densities for zero-shot classification. In imitation learning, diffusion models have been used to improve Behavioral Cloning (DBC) [5] by using diffusion model to model expert state-action pairs. Similarly, DiffAIL [58] extends GAIL [21] by employing diffusion models to represent the expert's behavior and incorporating the diffusion loss into the discriminator's learning process. However, DiffAIL's use of an unconditional diffusion model limits its ability to distinguish between expert and agent state-action pairs. We provide a detailed explanation of its limitation in Section 4.3 and Section A.  
approach: 
 We propose a novel adversarial imitation learning framework incorporating diffusion models into the generative adversarial imitation learning (GAIL) framework, illustrated in Figure 2 . Specifically, we employ a diffusion model to construct an enhanced discriminator to provide more robust and smoother rewards for policy learning. We initiate our discussion by describing a naive integration of the diffusion model, which directly predicts rewards from Gaussian noises conditioned on stateaction pairs, and the inherent issues of this method in Section 4.1. Subsequently, in Section 4.2, we introduce our proposed method that employs a conditional diffusion model to construct a diffusion discriminative classifier, which can provide diffusion rewards for policy learning. Finally, the overall algorithm of our method is outlined in Section 4.3.Conditional diffusion models are widely adopted in various domains, e.g., generating an image x from a label y. Intuitively, one can incorporate a conditional diffusion model as a GAIL discriminator by training it to produce a real or fake label conditioned on expert or agent state-action pairs. Specifically, given a denoising time step t and a state-action pair (s, a) ∈ (S × A), where S, A stand for state and action spaces, respectively, as a condition, the diffusion model p ϕ (r t-1 |r t , s, a) learns to denoise a reward label r 0 ∈ {0, 1}, i.e., 1 for expert (real) state-action pairs and 0 for agent (fake) state-action pairs through a reverse diffusion process.To train a policy, we can use the diffusion model to produce a reward r given a state-action pair (s, a) from the policy through a generation process by iteratively denoising a sampled Gaussian noise, i.e., noisy reward, conditioned on the state-action pair. Then, the policy learns to optimize the rewards predicted by the diffusion model. Nevertheless, the reward generation process is extremely time-consuming since predicting a reward for each state-action pair from the policy requires running T (often a large number) denoising steps, and policy learning often takes tens of millions of samples, resulting in a billion-level overall training scale. Consequently, it is impractical to integrate a diffusion model into the GAIL framework by using it to predict "realness" rewards for policy learning from state-action pairs.Our goal is to yield a diffusion model reward given an agent state-action pair without going through the entire diffusion generation process. Inspired by previous work [5, 59] , we extract the learning signal from a portion of the diffusion denoising steps, rather than using the entire process. Building on these insights, we adapt the training procedure of DDPM to develop a mechanism that provides a binary classification signal using just one denoising step.Our key insight is to leverage the derivations developed by Kingma et al. [26] , Song et al. [52] , 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. In this work, we propose calculating "realness" rewards based on the diffusion loss computed by denoising the state-action pairs from the policy, which indicates how well the state-action pairs fit the expert behavior distributions. We formulate the diffusion loss L diff as follows:EQUATIONwhere c ∈ {c + , c -}, and the real label c + corresponds to the condition for fitting expert data while the fake label c -corresponds to agent data. We implement c + as 1 and c -as 0.To approximate the expectation in Eq. 3, we use random sampling, allowing us to achieve the result with just a single denoising step. Subsequently, given a state-action pair (s, a), L diff (s, a, c + ) measures how well (s, a) fits the expert distribution and L diff (s, a, c -) measures how well (s, a) fits the agent distribution 1 . That said, given state-action pairs sampled from expert demonstration, L + diff should be close to 0, and L - diff should be a large value; on the contrary, given agent state-action pairs, L + diff should be a large value and L - diff should close to 0. While L + diff and L - diff can indicate the "realness" or the "fakeness" of a state-action pair to some extent, optimizing a policy using rewards with this wide value range [0, ∞) can be difficult [20] . To address this issue, we propose transforming this diffusion model into a binary classifier that provides "realness" in a bounded range of [0, 1]. Specifically, given the diffusion model's output L +,- diff , we construct a diffusion discriminative classifier D ϕ : S × A → R:EQUATIONwhere σ(x) = 1/(1 + e -x ) denotes the sigmoid function. The classifier integrates L + diff and L - diff to compute the "realness" of a state-action pair within a bounded range of [0, 1], as illustrated in Figure 2 . Since the design of our diffusion discriminative classifier aligns with the GAIL discriminator [21] , learning a policy with the classifier enjoys the same theoretical guarantee, i.e., optimizing this objective can bring a policy's occupancy measure closer to the expert's. Consequently, we can optimize our proposed diffusion discriminative classifier D ϕ with the loss function:L D = E (s,a)∈τ E [-log(D ϕ (s, a))] L expert BCE +E (s,a)∈τi [-log(1 -D ϕ (s, a))] L agent BCE (5)where L D sums the expert binary cross-entropy loss L expert BCE and the agent binary cross-entropy loss L agent BCE , and τ E and τ i represent a sampled expert trajectory and a collected agent trajectory by the policy π at training step i. We then update the diffusion discriminative classifier parameters ϕ based on the gradient of L D to improve its ability to distinguish expert data from agent data.Intuitively, the discriminator D ϕ is trained to predict a value closer to 1 when the input state-action pairs are sampled from expert demonstration (i.e., trained to minimize L + diff and maximize L - diff ), and 0 if the input state-action pairs are obtained from the agent online interaction (i.e., trained to minimize L - diff and maximize L + diff ). Note that our idea of transforming the diffusion model into a classifier is closely related to Li et al. [35] , which shows that minimizing the diffusion loss is equivalent to maximizing the evidence lower bound (ELBO) of the log-likelihood [4] , allowing for turning a conditional text-to-image diffusion model into an image classifier by using the ELBO as an approximate class-conditional log-likelihood log p(x|c). By contrast, we employ a diffusion model for imitation learning. Moreover, we take a step further -instead of optimizing the diffusion loss L diff , we directly optimize the binary cross entropy losses calculated based on the denoising results to train the diffusion model as a binary classifier.Algorithm 1 Diffusion-Reward Adversarial Imitation Learning (DRAIL) 1: Input: Expert trajectories τ E , initial policy parameters θ 0 , initial diffusion discriminator parameters ϕ 0 , and discriminator learning rate η ϕ 2: for i = 0, 1, 2, . . . Compute the output of diffusion discriminative classifier D ϕ (Eq. 4) and the loss function L D (Eq. 5)5:Update the diffusion modelϕ i+1 ← ϕ i -η ϕ ∇L D 6:Compute the diffusion reward r ϕ (s, a) with Eq. 6 7:Update the policy θ i+1 ← θ i with any RL algorithm w.r.t. reward r ϕ 8: end for Our proposed method adheres to the fundamental AIL framework, where a discriminator and a policy are updated alternately. In the discriminator step, we update the diffusion discriminative classifier with the gradient of L D following Eq. 5. In the policy step, we adopt the adversarial inverse reinforcement learning objective proposed by Fu et al. [12] as our diffusion reward signal to train the policy:r ϕ (s, a) = log(D ϕ (s, a)) -log(1 -D ϕ (s, a)). (6)The policy parameters θ can be updated using any RL algorithm to maximize the diffusion rewards provided by the diffusion discriminative classifier, bringing the policy closer to the expert policy. In our implementation, we utilize PPO as our policy update algorithm. The algorithm is presented in Algorithm 1, and the overall framework is illustrated in Figure 2 .Among the related works, DiffAIL [59] is the closest to ours, as it also uses a diffusion model for adversarial imitation learning. DiffAIL employs an unconditional diffusion model to denoise state-action pairs from both experts and agents. However, this approach only implicitly reflects the likelihood of state-action pairs belonging to the expert class through diffusion loss, making it challenging to explicitly distinguish between expert and agent behaviors.In contrast, our method, DRAIL, uses a conditional diffusion model that directly conditions real (c + ) and fake (c -) labels. This allows our model to explicitly calculate and compare the probabilities of state-action pairs belonging to either the expert or agent class. This clearer and more robust signal for binary classification aligns more closely with the objectives of the GAIL framework, leading to more stable and effective learning. For further details and the mathematical formulation, please refer to Section A. 
experiments: 
 We extensively evaluate our proposed framework DRAIL in diverse continuous control domains, including navigation, robot arm manipulation, and locomotion. We also examine the generalizability and data efficiency of DRAIL in Section 5.4 and Section 5.5. The reward function learned by DRAIL is presented in Section 5.6.This section describes the environments, tasks, and expert demonstrations used for evaluation.MAZE. We evaluate our approach in the point mass MAZE navigation environment, introduced in Fu et al. [13] (maze2d-medium-v2), as depicted in Figure 3a . In this task, a point-mass agent is trained to navigate from a randomly determined start location to the goal. The agent accomplishes the task by iteratively predicting its acceleration in the vertical and horizontal directions. We use the expert dataset provided by Lee et al. [33] , which includes 100 demonstrations, comprising 18,525 transitions.FETCHPUSH. We evaluate our approach in a 7-DoF Fetch task, FETCHPUSH, depicted in Figure 3b , where the Fetch is required to push a black block to a designated location marked by a red sphere. We use the demonstrations from Lee et al. [33] , consisting of 20,311 transitions (664 trajectories).HANDROTATE. We further evaluate our approach in a challenging environment HANDROTATE with a high-dimensional continuous action space introduced by Plappert et al. [43] . Here, a 24-DoF Shadow Dexterous Hand is tasked with learning to in-hand rotate a block to a target orientation, as depicted in Figure 3c . This environment features a high-dimensional state space (68D) and action space (20D). We use the demonstrations collected by Lee et al. [33] , which contain 515 trajectories (10k transitions). ANTREACH. The goal of ANTREACH, a location and navigation task, is for a quadruped ant to reach a goal randomly positioned along the perimeter of a half-circle with a radius of 5 meters, as depicted in Figure 3d . The 132D high-dimensional continuous state space encodes joint angles, velocities, contact forces, and the goal position relative to the agent. We use the demonstrations provided by Lee et al. [33] , which contain 1,000 demonstrations (25k transitions).WALKER. The objective of WALKER is to let a bipedal agent move at the highest speed possible while preserving its balance, as illustrated in Figure 3e . We trained a PPO expert policy with environment rewards and collected 5 successful trajectories, each containing 1000 transitions, as an expert dataset.CARRACING. We evaluate our method in a racing game, CARRACING, illustrated in Figure 3f , requiring driving a car to navigate a track. This task features a 96 × 96 RGB image-based state space and a 3-dimensional action space (steering, braking, and accelerating). We trained a PPO expert policy on CARRACING environment and collected 671 transitions as expert demonstrations.Further details of the tasks can be found in Section B.We compare our method DRAIL with the following baselines of our approach.• Behavioral Cloning (BC) trains a policy to mimic the actions of an expert by supervised learning a mapping from observed states to corresponding expert actions [44, 55] . • Diffusion Policy represents a policy as a conditional diffusion model [6, 42, 46] , which predicts an action conditioning on a state and a randomly sampled noise. We include this method to compare learning a diffusion model as a policy (diffusion policy) or reward function (ours). • Generative Adversarial Imitation Learning (GAIL) [21] learns a policy from expert demonstrations by training a discriminator to distinguish between trajectories generated by the learned generator policy and those from expert demonstrations. • Generative Adversarial Imitation Learning with Gradient Penalty (GAIL-GP) is an extension of GAIL that introduces a gradient penalty to achieve smoother rewards and stabilize the discriminator. • Wasserstein Adversarial Imitation Learning (WAIL) [2] extends GAIL by employing Wasserstein distance, aiming to capture smoother reward functions. • Diffusion Adversarial Imitation Learning (DiffAIL) [58] integrates a diffusion model into AIL by using the diffusion model loss to provide reward e -Ldiff .We present the success rates (MAZE, FETCHPUSH, HANDROTATE, ANTREACH) and return (WALKER, CARRACING) of all the methods with regards to environment steps in Figure 4 . Each task was trained using five different random seeds. Note that BC and Diffusion Policy are offline imitation learning algorithms, meaning they cannot interact with the environment, so their performances are represented as horizontal lines. Detailed information on model architectures, training, and evaluation can be found in Section F and Section G.Overall, our method DRAIL consistently outperforms prior methods or achieves competitive performance compared to the best-performing methods across all the environments, verifying the effectiveness of integrating our proposed diffusion discriminative classifier into the AIL framework.DRAIL vs. DiffAIL. Both DRAIL and DiffAIL integrate the diffusion model into the AIL framework. In 5 out of 6 tasks, our DRAIL outperforms DiffAIL, demonstrating that our proposed discriminator provides a more effective learning signal by closely resembling binary classification within the GAIL framework.DRAIL vs. BC. AIL methods generally surpass BC in most tasks due to their ability to learn from interactions with the environment and thus handle unseen states better. However, BC outperforms all other baselines in the locomotion task (WALKER). We hypothesize that WALKER is a monotonic task requiring less generalizability to unseen states, allowing BC to excel with sufficient expert data. Additionally, our experiments with varying amounts of expert data, detailed in Section 5.5, suggest that DRAIL surpasses BC when less expert data is available.We empirically found that our proposed DRAIL is robust to hyperparameters, especially compared to GAIL and WAIL, as shown in the hyperparameter sensitivity experiment inSection D.To examine the generalizability to states or goals that are unseen from the expert demonstrations of different methods, we extend the FETCHPUSH tasks following the setting proposed by Lee et al. [33] . Specifically, we evaluate policies learned by different methods by varying the noise injected into initiate states (e.g., position and velocity of the robot arm) and goals (e.g., target block positions in FETCHPUSH). We experiment with different noise levels, including 1×, 1.25×, 1.5×, 1.75×, and 2.0×, compared to the expert environment. That said, 1.5× means the policy is evaluated in an environment with noises 1.5× larger than those injected into expert data collection. Performing well in a high noise level setup requires the policy to generalize to unseen states. To investigate the data efficiency of DRAIL, we vary the amount of expert data used for learning in WALKER and FETCHPUSH. Specifically, for WALKER, we use 5, 3, 2, and 1 expert trajectories, each containing 1000 transitions; for FETCHPUSH, we use 20311, 10000, 5000, and 2000 state-action pairs. The results reported in Figure 6 demonstrate that our DRAIL learns faster compared to the other baselines, indicating superior data efficiency in terms of environment interaction. In WALKER, our DRAIL maintains a return value of over 4500 even when trained with a single trajectory. In contrast, BC's performance is unstable and fluctuating, while the other baselines experience a dramatic drop.In FETCHPUSH, our DRAIL maintains a success rate of over 80% even when the data size is reduced by 90%, whereas the other AIL baselines' performance drops below 50%. To visualize and analyze the learned reward functions, we design a SINE environment with onedimensional state and action spaces, where the expert state-action pairs form a discontinuous sine wave a = sin (20sπ) + N (0, 0.05 2 ), as shown in Figure 7a . We train GAIL and our DRAIL to learn from this expert state-action distribution and visualize the discriminator output values D ϕ to examine the learned reward function, as presented in Figure 7 .Figure 7b reveals that the GAIL discriminator exhibits excessive overfitting to the expert demonstration, resulting in its failure to provide appropriate reward values when encountering unseen states. In contrast, Figure 7c shows that our proposed DRAIL generalizes better to the broader state-action distribution, yielding a more robust reward value, thereby enhancing the generalizability of learned policies. Furthermore, the predicted reward value of DRAIL gradually decreases as the state-action pairs deviate farther from the expert demonstration. This reward smoothness can guide the policy even when it deviates from the expert policy. In contrast, the reward distribution from GAIL is relatively narrow outside the expert demonstration, making it challenging to properly guide the policy if the predicted action does not align with the expert.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [No] Justification: We plan to release the codes, models, and expert datasets upon acceptance.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusion: 
 This A Relation to DiffAIL Among all the related works, DiffAIL [58] , which also employs a diffusion model for adversarial imitation learning, is the closest to our work. This section describes the differences of our work and DiffAIL.Unconditional diffusion models (DiffAIL) vs. conditional diffusion models (DRAIL). Wang et al. [58] proposed to learn an unconditional diffusion model D ϕ (s, a) to denoise state-action pairs from an expert and an agent. The diffusion model should denoise well when the state-action pairs are sampled from the expert demonstrations, while denoise poorly given the state-action pairs sampled from the agent policy. On the other hand, our proposed framework employs a conditional diffusion model D ϕ (s, a, c) that conditions on the real label c + or the fake label c -. In the following, we discuss how using a conditional diffusion model leads to better learning signals for policy learning.DiffAIL set up the discriminator output as D ϕ (s, a) = e -Ldiff(s,a) ∈ [0, 1], indicating how likely (s, a) is sampled from the expert distribution. That said, a e -Ldiff(s,a) close to 1 implies that the sample is likely from expert and a small value of e -Ldiff(s,a) means it is less likely that this sample comes from the expert. However, e -Ldiff(s,a) does not explicitly represent the probability of the "negative class" (agent), which makes it difficult to provide stable rewards for policy learning.We start by revisiting the GAIL formulation. The binary classifier in GAIL, i.e., the discriminator, is trained to minimize the binary cross-entropy loss, which leads to maximizing the likelihood of predicting the correct class probabilities. Subsequently, the discriminator outputs the predicted probability of the "positive class" (expert) relatively to the "negative class" (agent). In DiffAIL, the predicted probability of the "positive class" (expert) is e -Ldiff and the predicted probability of the "negative class" (agent) can be given by 1 -e -Ldiff . Following this definition, the discriminator thinks the sample if from an expert when e -Ldiff > 1 -e -Ldiff . We also know that L diff > 0 since it is the diffusion model denoising loss. This gives us 0 ≤ L diff ≤ ln 2 for the expert class, and ln 2 < L diff for the agent class. This can be problematic since the fixed boundary ln 2 my not reflect the real boundary between the real (expert) and the fake (agent) distributions.On the other hand, in our method (DRAIL), the predicted probability of the positive (expert) class of a given state-action pair (s, a) is defined as:EQUATIONand the predicted probability of the negative (agent) class of the same state-action pair (s, a) is explicitly defined as:EQUATIONThis binary classification formulation aligns with GAIL. Hence, the discriminator would think a given (s, a) comes from expert data only when e -Ldiff(s,a,c + ) e -Ldiff(s,a,c + ) + e -Ldiff(s,a,c -) > e -Ldiff(s,a,c -) e -Ldiff(s,a,c + ) + e -Ldiff(s,a,c -) ,which means e -Ldiff(s,a,c + ) > e -Ldiff(s,a,c -) .The resulting relative boundary can provide better learning signals for policy learning, especially when the behaviors of the agent policy become similar to those of expert, which can be observed in the tasks where the agent policies can closely follow the experts, such as FETCHPICK, FETCHPUSH, and ANTREACH. Also, we hypothesize this leads to the superior performance of our method compared to DiffAIL in most of the generalization experiments.Experimental setup. Wang et al. [58] evaluated DiffAIL and prior methods in locomotion tasks; in contrast, our work extensively compares our proposed framework DRAIL with various existing methods in various domains, including navigation (MAZE and ANTREACH), locomotion (WALKER and ANTREACH), robot arm manipulation (FETCHPUSH and FETCHPICK), robot arm dexterous (HANDROTATE), and games (CARRACING). Additionally, we present experimental results on generalization to unseen states and goals on the goal-oriented tasks, and on varying amounts of expert data.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: We have provided the experimental details including task details in Section B, model architecture and the policy update algorithm in Section F, and training details and their corresponding parameters in Section G.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes]Justification: We have considered the potential implications of our work on biases and discriminatory behaviors, as highlighted in our impact statement (Appendix J).Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: This paper does not pose risks associated with the release of high-risk data or models. We use a self-trained gymnasium dataset, which does not involve pre-trained language models, image generators, or scraped datasets that could be misused. As such, specific safeguards for data or model release are not necessary in this context.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.• We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: We gave credit to the diffusion model used in our image-based experiments and the expert datasets for some tasks and cited these papers in Appendix F.2 and Section 5.1. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] Justification: This paper does not release new assets. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Description. In a 2D maze environment, a point-maze agent learns to navigate from a starting location to a goal location. The agent achieves this by iteratively predicting its x and y velocity. The initial and final positions of the agent are randomly selected. The state space includes position, velocity, and goal position. The maximum episode length for this task is set at 400, and the episode terminates if the goal is reached earlier.Expert dataset. The expert dataset comprises 100 demonstrations, which includes 18, 525 transitions provided by Lee et al. [33] .Description. In the FETCHPUSH task, the agent is required to push an object to a specified target location. On the other hand, in the FETCHPICK task, the objective is to pick up an object from a table and move it to a target location.According to the environment setups stated in Lee et al. [33] , the 16-dimensional state representation includes the angles of the robot joints, and the initial three dimensions of the action vector represent the intended relative position for the next time step. The first three dimensions of the action vector denote the intended relative position in the subsequent time step. In the case of FETCHPICK, an extra action dimension is incorporated to specify the distance between the two fingers of the gripper. The maximum episode length for this task is set at 60 for FETCHPUSH and 50 for FETCHPICK, and the episode terminates if the agent reaches the goal earlier.Expert dataset. The expert dataset for FETCHPUSH comprises 664 trajectories, amounting to 20, 311 transitions, and the expert dataset for FETCHPICK comprises 303 trajectories, amounting to 10, 000 transitions provided by Lee et al. [33] .Description. In the task HANDROTATE proposed by Plappert et al. [43] , a 24-DoF Shadow Dexterous Hand is designed to rotate a block in-hand to a specified target orientation. The 68D state representation includes the agent's joint angles, hand velocities, object poses, and target rotation. The 20D action vector represents the joint torque control of the 20 joints. Notably, HANDROTATE is challenging due to its high-dimensional state and action spaces. We follow the experimental setup outlined in Plappert et al. [43] and Lee et al. [33] , where rotation is constrained to the z-axis, and allowable initial and target z rotations are within [-π 12 , π 12 ] and [ π 3 , 2π 3 ], respectively. The maximum episode length for this task is set at 50, and the episode terminates if the hand reaches the goal earlier.Expert dataset. We use the demonstrations collected by Lee et al. [33] , which contain 515 trajectories (10k transitions).Description. The ANTREACH task features a four-leg ant robot reaching a randomly assigned target position located within a range of half-circle with a radius of 5 meters. The task's state is represented by a 132-dimension vector, including joint angles, velocities, and the relative position of the ant towards the goal. Expert data collection for this task is devoid of any added noise, while random noise is introduced during the training and inference phases. Consequently, the policy is required to learn to generalize to states not present in the expert demonstrations. The maximum episode length for this task is set at 50, and the episode terminates if the ant reaches the goal earlier.Expert dataset. The expert dataset comprises 10000 state-action pairs provided by Lee et al. [33] .Description. WALKER task involves guiding an agent to move towards the x-coordinate as fast as possible while maintaining balance. An episode terminates either when the agent experiences predefined unhealthy conditions in the environment or when the maximum episode length (1000) is reached. The agent's performance is evaluated over 100 episodes with three different random seeds. The return of an episode is the cumulative result of all time steps within that episode. The 17D state includes joint angles, angular velocities of joints, and velocities of the x and z-coordinates of the top. The 6D action defines the torques that need to be applied to each joint of the walker avatar.Expert dataset. We trained a PPO expert policy with environment rewards and collected 5 successful trajectories, each containing 1000 transitions, as an expert dataset.Description. In the CARRACING task, the agent must navigate a track by controlling a rear-wheel drive car. The state space of CARRACING is represented by a top-down 96 × 96 RGB image capturing the track, the car, and various status indicators such as true speed, four ABS sensors, steering wheel position, and gyroscope readings. The agent controls the car using three continuous action values: steering, acceleration, and braking. Episodes have a maximum length of 1000 steps, and termination occurs if the car completes the track before reaching the maximum episode length.In our experiment settings, we preprocess the state image by converting it to grayscale and resizing it to 64 × 64 pixels. We then concatenate two consecutive frames to form a single state data, providing temporal context for the agent's decision-making process.Expert dataset. We trained a PPO expert policy on CARRACING environment and collected 671 transitions as expert demonstrations.For MAZE, FETCHPUSH, FETCHPICK, HANDROTATE, and ANTREACH, we collected only the successful trajectories, resulting in a success rate of 100% for experts on these tasks. The expert performance for WALKER and CARRACING is 5637 ± 55 and 933 ± 0.9, respectively.To show our approach's better generalization capabilities, we extend the environment scenarios following the setting stated in Lee et al. [33] : (1) In MAZE main experiment, the initial and goal states of the expert dataset only constitute 50% of the potential initial and goal states. In the generalization experiment, we gather expert demonstrations from some lower and higher coverage: 25%, 75%, and 100%. (2) In FETCHPICK, FETCHPUSH, and HANDROTATE main experiments, the demonstrations are collected in a lower noise level setting, 1×. Yet, the agent is trained within an environment incorporating 1.5× noise, which is 1.5 times larger noise than the collected expert demonstration, applied to the starting and target block positions. In the generalization experiment, we train agents in different noise levels: 1×, 1.25×, 1.75×, 2.0×. (3) In ANTREACH main experiment, no random noise is added to the initial pose during policy learning. In ANTREACH generalization experiment, we train agents in different noise levels: 0 (default), 0.01, 0.03, 0.05.These generalization experiments simulate real-world conditions. For example, because of the expenses of demonstration collection, the demonstrations may inadequately cover the entire state space, as seen in setup (1) . Similarly, in setups ( 2) and (3), demonstrations may be acquired under controlled conditions with minimal noise, whereas the agent operating in a real environment would face more significant noise variations not reflected in the demonstrations, resulting in a broader distribution of initial states.Our DRAIL outperforms baselines or performs competitively against DiffAIL across all demonstration coverages, as shown in Figure 8 . Particularly, BC, WAIL, and GAIL's performance decline rapidly in the low coverage case. In contrast, diffusion model-based AIL algorithms demonstrate sustained performance, as shown in Figure 4a . This suggests that our method exhibits robust generalization, whereas BC and GAIL struggle with unseen states under limited demonstration coverage.FETCHPICK and FETCHPUSH. In FETCHPICK, our method outperforms all baselines in most noise levels, as shown in Figure 8 . In the 2.00× noise level, our method DRAIL achieves a success rate of 87.22, surpassing the best-performing baseline Diffusion Policy, which achieves only around 76.64%. GAIL, on the other hand, experiences failure in 3 out of the 5 seeds, resulting in a high standard deviation (mean: 39.17, standard deviation: 47.98) despite our thorough exploration of various settings for its configuration. In FETCHPUSH, our method DRAIL exhibits more robust results, in generalizing to unseen states compared to other baselines. This showcases that the diffusion reward guidance could provide better generalizability for the AIL framework. Moreover, our DRAIL is quite sample-efficient regarding interaction with the environment during training compared to other baselines in FETCHPICK FETCHPUSH environment.HANDROTATE. DiffAIL and Our DRAIL show robustness to different levels of noise in HANDRO-TATE, as illustrated in Figure 8 . Specifically, DiffAIL and our DRAIL achieve a success rate of higher than 90% at a noise level of 2.0, while GAIL and WAIL only reach approximately 42.82% and 8.80%, respectively.We empirically found that our proposed method, DRAIL, is robust to hyperparameters and easy to tune, especially compared to GAIL and WAIL. In this section, we present additional ablation experiments to examine how hyperparameter tuning affects the performance of DRAIL. Like most AIL methods, the key hyperparameters of DRAIL are the learning rates of the policy and discriminator. the key hyperparameters for DRAIL are the learning rates of the policy and discriminator. We experimented with various learning rate values, including 5x, 2x, 1x, 0.5x, and 0.2x of the value used in the main results for the MAZE environment. The results, presented in Figure 9 , demonstrate that our method is robust to variations in hyperparameters.This section reports the quantitative results of the converged performance across all experiments, including the main results in Section 5.3, the generalization experiments in Appendix C.2, and the data efficiency experiments in Section 5.5. The results are presented in Table 1 .This section presents the model architecture of all the experiments. Appendix F.1 describe the model architecture of all methods used in Section 5.3.In Section 5.3, we conducted a comparative analysis between our proposed DRAIL, along with several baseline approaches (BC, Diffusion Policy, GAIL, GAIL-GP, WAIL, and DiffAIL) across six diverse environments. We applied Multilayer Perceptron (MLP) for the policy of BC, the conditional diffusion model in Diffusion Policy, as well as the policy and the discriminator of GAIL, GAIL-GP, and WAIL. For DiffAIL and our proposed DRAIL, MLPs were employed in the policy and diffusion model of the diffusion discriminative classifier. The activation functions used for the MLPs in the diffusion model were ReLU, while hyperbolic tangent was employed for the others. The total timestep T for all diffusion models in this paper is set to 1000 and the scheduler used for diffusion models is cosine scheduler [38] . Further details regarding the parameters for the model architecture can be found in Table 2 .BC. We maintained a concise model for the policy of BC to prevent excessive overfitting to expert demonstrations. This precaution is taken to mitigate the potential adverse effects on performance when confronted with environments exhibiting higher levels of noise.Diffusion Policy. Based on empirical results and Chen et al. [5] , the Diffusion Policy performs better when implemented with a deeper architecture. Consequently, we have chosen to set the policy's number of layers to 5. GAIL & GAIL-GP. The detailed model architecture for GAIL and GAIL-GP is presented in Table 2 . For GAIL-GP, the gradient penalty is set to 1 across all environments.WAIL. We set the ground transport cost and the type of regularization of WAIL as Euclidean distance and L 2 -regularization. The regularization value ϵ is provided in Table 2 .DiffAIL. In DiffAIL, the conditional diffusion model is not utilized as it only needs to consider the numerator of Equation ( 4). Consequently, the diffusion model takes only the noisy state-action pairs as input and outputs the predicted noise value.DRAIL. The conditional diffusion model of the diffusion discriminative classifier in our DRAIL is constructed by concatenating either the real label c + or the fake label c -to the noisy state-action pairs as the input. The model then outputs the predicted noise applied to the state-action pairs. The dimensions of both c + and c -are reported in Table 2 .For the CARRACING task, we redesigned the model architecture for our DRAIL and all baseline methods to handle image-based input effectively.In CARRACING, the policy for all baselines utilizes a convolutional neural network (CNN) for feature extraction followed by a multi-layer perceptron (MLP) for action prediction. The CNN consists of three downsampling blocks with 32, 64, and 64 channels respectively. The kernel sizes for these blocks are 8, 4, and 3, with strides of 4, 2, and 1, respectively. After feature extraction, the output is flattened and passed through a linear layer to form a 512-dimensional feature vector representing the state data. This state feature vector is subsequently processed by an MLP with three layers, each [47] structure, comprising 3 down-sampling blocks and 3 up-sampling blocks. Each block consists of 2 ResNet [19] layers, with group normalization applied using 4 groups. The channel sizes for each pair of down-sampling and up-sampling blocks are 4, 8, and 16, respectively. The condition label is incorporated through class embedding, with the number of classes set to 2, representing the real label c + and the fake label c -. Finally, we apply running normalization at the output to ensure stable training and accurate discrimination.To accommodate both image-based state and vector-based action data within the diffusion model, we flatten the action data into an image with a channel size equivalent to the action dimension. Subsequently, we concatenate the state and transformed action data as input to the U-Net. In DRAIL, we use 1 to represent c + and 0 to represent c -as condition data. For DiffAIL, since condition labels are not required, we simply assign a constant value of 0 as the condition label.The hyperparameters employed for all methods across various tasks are outlined in Table 3 . The Adam optimizer [27] is utilized for all methods, with the exception of the discriminator in WAIL, for which RMSProp is employed. Linear learning rate decay is applied to all policy models.Due to the potential impact of changing noise levels on the quality of agent data input for the discriminator, the delicate balance between the discriminator and the AIL method's policy may be disrupted. Therefore, we slightly adjusted the learning rate for the policy and the discriminator for different noise levels on each task. The reported parameters in Table 3 correspond to the noise levels presented in Figure 4 .As explained in Section 4.3, we adopt the optimization objective proposed by [12] as diffusion reward signal for the policy learning in our DRAIL. To maintain fairness in comparisons, we apply the same reward function to DiffAIL and GAIL. In CARRACING, we observe that adapting GAIL's optimization objective could lead to better performance; hence, we use it for DRAIL, DiffAIL, and GAIL. For WAIL, we adhere to the approach outlined in the original paper, wherein the output of the discriminator directly serves as the reward function.In our experiments, we employ Proximal Policy Optimization (PPO) [50] , a widely used policy optimization method, to optimize policies for all the AIL methods. We maintain all hyperparameters of PPO constant across methods for a given task, except the learning rate, which is adjusted for each method. The PPO hyperparameters for each task are presented in Table 4 .This work presents an adversarial imitation learning framework DRAIL by employing a diffusion model as a discriminator. While DRAIL achieves encouraging results in various domains, including robot arm manipulation, robot hand dexterous manipulation, locomotion, and games, our proposed framework is fundamentally limited to the learning from demonstration (LfD) setting. That said, DRAIL requires both state and action sequences and, therefore, cannot learn from videos or state-only sequences, i.e., learning from observation (LfO). Moreover, DRAIL assumes expert demonstrations to be optimal, and its performance may not be satisfactory if the demonstrations contain a certain level of noise or the demonstrators are suboptimal. Finally, DRAIL with its imitation learning nature, is not designed to learn from environmental rewards; therefore, even when environments can provide rewards, there is no apparent mechanism to utilize them with the current formulation of DRAIL.For our experiments, we used the following three workstations:• In this work, we propose a novel adversarial imitation learning framework, diffusion rewards guided adversarial imitation learning (DRAIL), which integrates a diffusion model into GAIL. The proposed framework can potentially reinforce the biases captured by expert demonstrations, which can lead to sub-optimal, unsafe, or even discriminatory behaviors. To address this issue, we encourage future works to focus on alleviating these issues in imitation learning, e.g., fairness in machine learning, and responsible AI.NeurIPS paper checklistFor simplicity, we will use the notations L + diff and L - diff to represent Ldiff(s, a, c + ) and Ldiff(s, a, c -), respectively, in the rest of the paper. Additionally, L +,- diff denotes L + diff and L - diff given a state-action pair. 


Decision: This paper proposes a diffusion-based adversarial imitation learning method to mitigate the instability of the prior work, generative adversarial imitation learning (GAIL). All of the reviewers found that the empirical result across a variety of benchmarks is strong enough to show that the proposed method is robust and sample-efficient. Although there were questions and concerns regarding the position of the work relative to the previous work and the lack of certain baselines, the authors clarified them and added additional baselines during the rebuttal period. As a result, all of the reviewers acknowledged that their major concerns were addressed. Assuming that the authors will make the contribution clearer in the camera-ready version, I recommend to accept this paper.

Reviews: review: paper_summary: paper_summary: Imitation learning (IL) is a research area that focuses on learning a policy from expert demonstrations. One of the most popular imitation learning techniques, General Adversarial Imitation Learning (GAIL), has been proposed to mitigate some of the issues that naive IL algorithms suffer from. Although GAIL has seen a lot of success, it is known to be very unstable and difficult to tune. The authors propose Diffusion Reward Adversarial Imitation Learning (DRAIL) to mitigate some of the difficulties of training GAIL. In particular, the authors use a diffusion model to train the reward function. With this change in reward function design, the authors showed in practice that policies trained with DRAIL performed better than past approaches that attempted to mitigate GAIL issues. Additionally, the authors demonstrated that they could learn a reward function with a very low number of expert trajectories compared to other approaches.
strength_weakness: strength_weakness: - The paper needs more novelty. DiffAIL proposed using diffusion as a reward, and DiffAIL also proposed using the training error as a reward signal.
- The paper needs to include comparisons to the two missing cites [1] and [2].
- The authors do not provide any insight into why the proposed reward function is better than DiffAIl, especially since DIffAIL sometimes matches or performs better than the proposed algorithms.
questions: questions: - What denoise step values did you try? And how did they perform?
- Could you explain why the proposed reward model performs better than DiffAIL?
- If BC performs better than all the baselines on Walker, then that means the task should be easy for all the baselines. - Did you try increasing or decreasing the number of expert samples used for learning? Are you using D4RL? If so, what dataset split are you using?
- Did you try training various levels of optimal demonstration, such as walker2d-expert-v0, walker2d-medium-v0, and walker2d-random-v0, to see how well the algorithms perform when noise is introduced?
- Are you taking the argmax or sampling from the learned imitation learning policies?
- Did you try other divergences besides Jensen-Shannon?
Missing citations [1], [2],


[1] A Coupled Flow Approach to Imitation Learning by Freund 2023

[2] Diffusion Model-Augmented Behavioral Cloning by Chen et al. 2024
limitations: limitations: Yes
paper_summary: paper_summary: The paper proposes a novel imitation learning framework that integrates a diffusion model into Generative Adversarial Imitation Learning (GAIL). The primary aim is to address the instability and brittleness associated with GAIL by introducing more robust and smoother reward functions for policy learning. The authors develop a diffusion discriminative classifier to enhance the discriminator's performance and generate diffusion rewards. 

Extensive experiments in various domains such as navigation, manipulation, and locomotion demonstrate that DRAIL outperforms or is competitive with existing imitation learning methods. The paper highlights DRAIL’s effectiveness, generalizability, and data efficiency, offering a significant contribution to the field of imitation learning.
strength_weakness: strength_weakness: 1. **Complexity**: The introduction of diffusion models adds complexity to the framework, which may pose challenges in implementation and require significant computational resources.
2. **Limited Scope in Real-World Applications**: While the experiments are diverse, the applicability in highly dynamic and unpredictable real-world environments is not thoroughly explored.
3. **Comparative Analysis Depth**: Although the paper compares DRAIL with several baselines, a deeper analysis of why certain methods perform better in specific tasks could enhance understanding.
4. **Scalability**: The scalability of DRAIL in larger and more complex environments is not extensively evaluated, which could be a limitation for broader adoption.
5. **Hyperparameter Sensitivity**: The sensitivity of DRAIL to various hyperparameters is not discussed, which could impact its robustness across different settings.

###
questions: questions: 1. How does DRAIL handle environments with rapidly changing dynamics, and what are the limitations in such scenarios?
2. What are the computational overheads associated with incorporating diffusion models, and how does it compare with other state-of-the-art methods in terms of efficiency?
3. How sensitive is DRAIL to the choice of hyperparameters, and what guidelines can be provided for tuning these parameters in different environments?
4. Can the diffusion discriminative classifier be extended or modified to further improve the reward robustness and policy performance in more complex tasks?
5. What specific strategies were employed to ensure the stability of the diffusion model training, and how do these strategies affect the overall performance of DRAIL?
limitations: limitations: Shown in weakness
paper_summary: paper_summary: This paper aims to address the training instability problem in generative adversarial imitation learning. The authors propose a diffusion discriminative classifier that helps achieve a more stable policy learning process by enhancing the smoothness and robustness of the reward model. Additionally, the author experimentally validates the effectiveness of DRIL across various control tasks, including manipulation, locomotion, and navigation benchmarks.
strength_weakness: strength_weakness: - The advantages of the diffusion discriminative classifier compared to diffusion reward [1] are not well explained. The authors claim that their proposed diffusion discriminative classifier only requires two reverse steps to obtain the corresponding reward, thereby reducing the computational resource consumption needed for sampling in diffusion models. However, I do not fully agree with this. According to equations (4) and (6), DRAIL needs to compute $L_{\text{diff}}(s,a,c^+)$ and $L_{\text{diff}}(s,a,c^-)$ when obtaining rewards. From equation (3), we know both terms require taking an expectation over the diffusion steps $T$. Therefore, to obtain an accurate reward, a number of inference steps equivalent to the diffusion steps are still needed.
- This paper lacks ablation experiments, including a sensitivity analysis of the hyperparameters.
- Some important AIL baselines are missing, such as DAC [2] and IQLearn [3].

[1] DiffAIL: Diffusion Adversarial Imitation Learning. AAAI, 2024.

[2] Discriminator-Actor-Critic: Addressing Sample Inefficiency and Reward Bias in Adversarial Imitation Learning. ICLR, 2019.

[3] IQ-Learn: Inverse soft-Q Learning for Imitation. NeurIPS, 2021.
questions: questions: - Why does DRAIL only require two reverse steps to obtain an accurate reward? Can you provide a more detailed explanation or analysis?
- I believe that the requirement for stable algorithm training includes robustness to hyperparameters. A major issue with AIL methods is the need for extensive hyperparameter tuning, and well-tuned AIL can achieve good performance. What are the key hyperparameters for DRAIL? Can you conduct additional ablation experiments on these hyperparameters?
- Can you compare DAC and IQLearn in your experiments? Both DAC and IQLearn use the gradient penalty for reward smoothing, and I am very interested in seeing a comparison between the diffusion discriminative classifier and the gradient penalty.

I would like to raise my score if my concerns are well addressed.
limitations: limitations: The authors have discussed the limitations of this paper. DRAIL is unable to learn from state-only trajectories and suboptimal trajectories.
paper_summary: paper_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.
strength_weakness: strength_weakness: ## 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).
questions: questions: See Weaknesses.
limitations: 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.

Comments: title: Reminder: The reviewer-author discussion period ends in 20 hours
Reminder: The reviewer-author discussion period ends in 20 hours
Reminder: The reviewer-author discussion period ends in three days
Reminder: The reviewer-author discussion period ends in three days
Reminder: The reviewer-author discussion period ends in three days
Re: Official Comment by Reviewer PZWW
comment: We would like to express our sincere gratitude to the reviewer for the thorough and constructive feedback. We are confident that our responses and additional experimental results adequately address the concerns raised by the reviewer. Please consider our rebuttal and kindly let us know if the reviewer has any additional concerns.
We would like to express our sincere gratitude to the reviewer for the thorough and constructive feedback. We are confident that our responses and additional experimental results adequately address the concerns raised by the reviewer. Please consider our rebuttal and kindly let us know if the reviewer has any additional concerns.
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 description of computational resources
- A discussion of the broad applicability of our proposed method
- An additional hyperparameter sensitivity experiment
- A discussion of further improving the robustness and stability of our proposed method

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.
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 detailed description of our contributions compared to DiffAIL
- An intuitive explanation of why our method works
- A discussion of Freund et al., 2023 (CFIL) and Chen et al., 2024 (DBC)
- A description of data efficiency experiments, i.e., learning from different amounts of expert data, provided in the main paper
- An explanation of training and sampling from learned policies using PPO
- A discussion of using different divergence measures  

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.
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 our contributions and a detailed plan to reorganize our paper based on the reviewer's suggestion
- Discussions of Janner et al., 2022 (Diffuser) and Nuti et al., 2023 (Extracting Reward Functions from Diffusion Models)

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.
We sincerely thank the reviewer for acknowledging our rebuttal and for helping us to improve our submission.
Thank you for your detailed answer to my questions. The response addressed my main concerns, therefore, I increased my score to 5.