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 2024 poster (2024)
Keywords: Imitation Learning
Adversarial Imitation Learning
Diffusion Model
Primary area: reinforcement_learning
TL;DR: This work proposes a novel adversarial imitation learning framework that integrates a diffusion model into generative adversarial imitation learning.

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/

Diffusion-Reward Adversarial Imitation Learning
Chun-Mao Lail*
Hsiang-Cuie Chun Wang Ping-Chun Hsieh
Yu-Chiang Frank Wang" 1.3
Min-Hung Chen
Shao-Hua Sun'
Nattonal Taiwan University
National Yang Ming Chiao Tung University
NNVIDIA
Abstract
Imitation learning aims to learn policy from observing expert demonstrations without access to reward signals from environments. Generative adversarial imi- tation learning (GAIL) formulates imitation learning as adversarial learning, em- ploying generator policy learning to imitate expert behaviors and discriminator learning to distinguish the expert demonstrations from agent trajectories. De- spite its encouraging results, GAIL training is often brittle and unstable. In- spired by the recent dominance of diffusion models in generative modeling, we propose Reward Adversarial Imitation Learning (DRAIL), which inte- grates diffusion model into GAIL, aiming to yield more robust and smoother rewards for policy learning. Specifically, tive classifier to construct an enhanced discriminator, and design diffusion re- wards based the classifier's output for policy learning. Extensive experi- ments are conducted in navigation, manipulation, and locomotion, verifying DRAIL's effectiveness compared prior imitation learning methods. More- over, 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::/nturrbotteeeeee
propose diffusion discrimina-
github. 1o/DRAIL/
1 Introduction
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 reward function is intricate or uncertain 7 32 or when learning in trial-and-error manner is expensive or unsafe [7].
Imitation learning, i.e., learning from demonstration 24
Among various methods in imitation learning, generative adversarial imitation learning (GAIL) has been widely adopted due its effectiveness and data efficiency. GAIL learns generator policy to imitate expert behaviors through reinforcement learning and a discriminator to differentiate between the expert and the generator's 's state-action pair distributions, resembling the idea of generative adversarial networks (GANs) [16]. Despite its established theoretical guarantee, GAIL training 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 (3), and exploring various similarity
measures of distributions [2] [2].
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
"Equal contribution.
Correspondence to: Shao-Hur Hua Sun cushacchase@tu..duu
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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-Rewii Imitation Learning (DRAIL), a novel adversarial imitation learning framework that can efficiently and effectively produce reliable policies replicating the behaviors of experts.
Adversarial
We extensively compare our proposed framework DRAIL with behavioral cloning 44 Diffusion Policy 6 42]. and AIL methods, e.g., GAIL (2I) 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 '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.
2 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 (6), and game Al [1]
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 [T0 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.
Behavioral Cloning (BC). BC
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 distribu- tions of an agent and an expert through adversarial training. Generative adversarial imitation learning (GAIL) 21 and its extensions 123 25 31, [56] 59 62] train a generator policy to imitate expert behaviors and diseriminator 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 various domains [ 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 is loss func- tion 12), extensions to off-policy RL algorithms [30]. addressing reward bias [31], and exploration of various similarity measures 12 1112) 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 model-based RL method that requires trajectory- reward information, which differs from our setting, iee,, imitation learning, where obtaining rewards is not possible. Nuti et al. [3 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. Morcover, they focus on learning

---

from fixed suboptimal dataset, AIL approaches and our method are designed to learn from agent data that continually change as the agents learn.
3
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.
3.1 Generative Adversarial Imitation Learning (GAIL)
GAIL [2 establishes a connection between generative adversarial networks (GANs) [T6] and imitation learning. GAIL employs generator, Gg. that acts policy иa, mapping state to an action. The generator aims to produce a state-action distribution (Pxg) which closely resembles the expert state-action distribution Prgi discriminator D. functions as a binary classifier, attempting to differentiate the state-action distribution of the generator (Px,) from the expert's (Prg). The optimization equation of GAIL can be formulated using the equivalent to the minimax equation of GAN. The optimization of GAIL can be derived as follows:
divergence, which is
min max Exmprn [log (x)) + Ex~peg [log(1 1 - (x)))
(1)
where Pmo and Prg are the state-action distribution from an agent my and expert policy म. respectively, (og(11D(xx)) For given state, the generator tries to take expert-like action; the discriminator takes state-action pairs input and computes the probability of the input originating from an expert. Then the generator uses a reward function - -x [log Dw (x)] or -Ex~p. [log (x)) + (TT) to optimize its network parameters, where the entropy term H is policy regularizer controlled by 2 0.
The loss function for the discriminator is stated as -Exxpp [log D.. (x)]+ (x))
3.2
Diffusion models
Diffusion models have demonstrated state- of-the-art performance on various tasks [
Reverse diffusion process Learn to denoise
PalX,-1 x
29
38 51.. This work builds upon
denoising diffusion probabilistic models (DDPMs) [22) that employ forward and reverse diffusion processes, as illustrated in Figure D The forward diffusion pro- cess injects noise into data points follow- ing a variance schedule until achieving an isotropic Gaussian distribution. The re- verse diffusion process trains a diffusion model é to predict the injected noise by optimizing the objective:
X
Adding Gaussian noise Forward diffusion process
Figure 1: Denoising diffusion probabilistic model. Latent variables I1, ..., EN are produced from the data point via forward diffusion process, i.e. gradually adding noises latent variables. diffusion model learns to reverse diffusion process by denoising the noisy data to reconstruct the original data point To-
CDM = EENT.ENN [|looVaz + VI - EzE, t) -El
(2)
where T' represents the set of discrete time steps in the diffusion process, e is the noise applied by the forward process, eş the noise predicted by the diffusion model, and a; is the scheduled noise level applied on the data samples.
Beyond generative tasks, diffusion models have also been successfully applied in other areas, in- cluding image classification and imitation lcarning. Diffusion Classifier [35] demonstrates that conditional diffusion models 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 model expert state-action pairs. Similarly, DiffAIL (58 extends GAIL [21 by employing diffusion models to represent the expert's 's behavior and incorporating the diffusion loss into the discriminator's learning process. However, DiffAIL's 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
3

---

Diffusion Discriminative Clessifier D
Learnsble mapping
Frozen meeping
Learning objective
Lar
"Real"
Expert Data fgg
D 1218p
C BCB
Diffusion Discrimintivee Classifier Da
Lar
Lan
Diftusion
d
Policy
Diffusion Model @b
D2(5.a)
"Fake"
Agent
Lia
Data
D2(5_-
BCF
4
Diffusion Reward rg(s.a)
logtD log(l Dy(s.ai)
(8) Learning Diffusion Discriminative Classifier
(b) Leamiing Policy with Diffusion Rewards
Figure 2: Diffusion-Reward Adversarial Imitation Learning, Our proposed framework DRAIL incorporates a diffusion model into GAIL. (a) Our proposed diffusion discriminative classifier D Iearns distinguish expert data (Sg, as) ~ TE from agent data (8+, a-) Ti using diffusion model. De trained to predict value closer to when the input state-action pairs are sampled from expert demonstration and predict value closer to 0 otherwise. (b) The policy Ty learns to maximize the diffusion reward r computed based on the output of De that takes the state-action pairs from the policy as input. The closer the policy resembles expert behaviors, the higher the rewards can obtain.
4 Approach
We propose novel adversarial imitation learning framework incorporating diffusion models into the generative adversarial imitation learning (GAIL) framework, illustrated in Figure Specifically, we employ diffusion model construct an enhanced discriminator 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 state- action 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 diffuston model construct diffusion discriminative classifier, which can provide diffusion rewards for policy learning. Finally, the overall algorithm of our method is outlined Section 4.3
4.1 Reward prediction with a conditional diffusion model
Conditional diffusion models are widely adopted various domains, e.g., generating an image from label y. Intuitively, one can incorporate a conditional diffusion model as GAIL. discriminator by training it to produce real or fake label conditioned on expert or agent state-action pairs. Specifically, given a denoising time step & and state-action pair (s, a) (S x A), where S, A stand for state and action spaces, respectively, as a condition, the diffusion model (T [Tt, s, a) learns to denoise a reward label E {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 policy, we can use the diffusion model to produce reward r given state-action pair (s, a) from the policy through a generation process by iteratively denoising 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 extremely time-consuming since predicting reward for each state-action pair from the policy requires running (often large number) denoising steps, and policy learning often takes tens of millions of samples, resulting billion-level overall training scale. Consequently, it is impractical to integrate diffusion model into the GAIL framework by using to predict "realness" rewards for policy learning from state-action pairs.
4.2 Diffusion discriminative classifier
Our goal is to yield diffusion model reward given an agent state-action pair without going through the entire diffusion generation process. Inspired by previous work [5] [5]], we extract the learning signal from 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. 226. Song et al. 52 which suggest that the diffusion loss, Lee, 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 Lai as follows:
Can . a. c) = EsT (ss a, c, -cl"]
(3)
where с {ct " and the real label c corresponds the condition for fitting expert data while the fake label c corresponds to agent data. We implement cx as 1 and c as O.
To approximate the expectation in Eq. result with just a single denoising step. Subsequently, given a state-action pair (s. a), Lai (s, a, c+) measures how well (s. a) fits the expert distribution and Lair(s, (s, a, c-) measures how well (s, a) fits the agent distribution That said, given state-action pairs sampled from expert demonstration, Lait should be close to 0, and far should be a large value; on the contrary, given agent state-action pairs, fai should be large value and aai should close to 0.
use random sampling, allowing us to achieve the
While Cai and dif can indicate the "realness" the "fakeness" of state-action pair to some extent, optimizing a policy using rewards with this wide value range [o, oo) can be difficult 220). To address this issue, we propose transforming this diffusion model into binary classifier that provides "realness" in bounded range of [0, IJ. Specifically, given the diffusion model's output C we construct a diffusion discriminative classifier D& S + R:
E C_n(ss,,et)
Dø(s, a)
C. (., Caai(sa..)-) o (Lairss, a, Lair(s, a, ct)),
where (x) = =/ denotes the sigmoid function. The classifier integrates C and L 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. ss Consequently, we can optimize our proposed diffusion discriminative classifier D& with the loss function:
LD ((S.))TT log( (Do (s, a + (s,a))) [- log(1 D& a))]
(5)
RCR
CRC
where LD sums the expert binary cross-entropy loss LBCE and the agent binary cross-entropy loss LCC. and TE and t; represent sampled expert trajectory and collected agent trajectory by the policy т at training step i. We then update the diffusion discriminative classifier parameters @ based on the gradient of LD to improve its ability distinguish expert data from agent data.
Intuitively, the discriminator D& trained to predict a value closer to when the input state-action pairs are sampled from expert demonstration (i.e., trained to minimize 4 and maximize LTr), and 0 if the input state-action pairs are obtained from the agent online interaction (i.e., trained to minimize fai and maximize L
Note that our idea of transforming the diffusion model into a classifier is closely related to Li et al. B35. which shows that minimizing the diffusion loss equivalent maximizing the evidence lower bound (ELBO) of the log-likelihood [4] allowing for turning conditional text-to-image diffusion model into an image classifier by using the ELBO as an approximate class-conditional log likelihood log p(z|c). By contrast, we employ diffusion model for imitation learning. Moreover, we take a step further instead of optimizing the diffusion loss Las, we directly optimize the binary cross entropy losses calculated based on the denoising results to train the diffusion model binary classifier.
'For simplicity, we will use the notations L and fun represent Lan(s, a, ct) and Cam (s. a, ), respectively, in the rest of the paper. Additionally. L.t denotes Lig and Lūn given state-action pair.
5

---

4.3 SelfusiionRewaary Adversarial Imitation Learning
Our proposed method adheres to the fundamental AlL framework, where discriminator and policy up- dated alternately. In the discriminator step, we update the diffusion discriminative classifier with the gradient of LD following Eq. [ 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:
Algorithm 1 Diffusion-Rewaan Adversarial Imitation Learning (DRAIL)
1: Input: Expert trajectories TE, initial policy parameters db, initial diffusion discriminator parameters Φo. and dis- criminator learning rate Ns
2: for i = 0, 1, do
3: Sample agent transitions Tı wo;
No (s. a) = log (s) - log(1 - (s(s, (6)
4: Compute the output of diffusion dis- criminative classifier D& (Eq. and the loss function LD (Eq. 55)
The policy parameters 0 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 implementa- tion, utilize PPO as our policy update algorithm. The algorithm is presented in Algorithm and the overall framework is illustrated in Figure 2
5:
Update the diffusion model Φ++1
O VLD
6:
Compute the diffusion reward
r a) with Eq. 6
7: Update the policy 0s+1 + 8 with
any RL algorithm w.r.t. reward ra
Among the related works, DiffAIL. 59 the closest to ours, as it also uses diffusion model for adversarial
8: end for
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 challenging to explicitly distinguish between expert and agent behaviors.
In contrast, our method, DRAIL, uses a conditional diffusion model that directly conditions real (ct) 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
5 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 Section5.44and 5.4 Section The reward function learned by DRAIL is presented in Section 5.6
5.1 Experimental setup
This section describes the environments, tasks, and expert demonstrations used for evaluation.
MAZE. We evaluate approach in the point mass MAZE navigation environment, introduced in Fu et al. [3] (maze2d-medium-
as depicted in Figure 3a In this task, a point-mass agent is trained navigate from 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 approach in 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 al. 33 consisting of 20,31 transitions (664 trajectories).
HANDROTATE. We further evaluate our approach in a challenging environment HANDROTATE with 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 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 al. 33 which contain 515 trajectories (10k transitions).

---

GAS
(a) MAZE
(b) FETCHPUSH
(c) HANDROTATE
(d) ANTREACH
(C) WALKER
(f) CARRACING
Figure 3: Environments & tasks. (a) MAZE: A point-mass agent (green) within 2D maze is trained to move from its initial position reach the goal (red). (b) FETCHPUSH: The manipulation task implemented with a 7-DoF Fetch robotics arm. FETCHPUSH requires picking up or pushing an object to target location (red). (c) HANDROTATE: For this dexterous manipulation task, a Shadow Dexterous Hand is employed to in-hand rotate block to achieve a target orientation. (d) ANTREACH: This task trains a quadruped ant to reach a goal randomly positioned along the perimeter of half-cirele with radius of 5 m. (e) WALKER: This locomotion task requires training bipedal walker policy achieve the highest possible walking speed while maintaining balance, (f) CARRACING This image-basee racing game task requires driving a car to navigate track as quickly as possible.
ANTREACH. The goal of ANTREACH, a location and navigation task, is for quadruped ant to reach a goal randomly positioned along the perimeter of a half-circle with a radius 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 al. 33 which contain 1,000 demonstrations (25k transitions).
The objective of WALKER let a bipedal agent move the highest speed possible while preserving its balance, as illustrated in Figure 3e We trained a PPO expert policy with environment rewards and collected successful trajectories, each containing 1000 transitions, as an expert dataset.
WALKER.
CARRACING. We evaluate our method in a racing game, CARRACING, illustrated in Figure 3f requiring driving car 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 PPO expert policy on CARRACING environment and collected 671 transitions expert demonstrations.
Further details of the tasks can be found in Section
5.2 Baselines
We compare our method DRAIL with the following baselines of our approach.
Behavioral Cloning (BC) trains 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 policy as conditional diffusion model 6 42
46 which
predicts action conditioning on state and 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) (1 learns a policy from expert demonstrations by training a discriminator 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 gradient penalty 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 (DiffAlL.) 58
integrates a diffusion model into AlL by using the diffusion model loss to provide reward c Can
5.3 Experimental results
We present the success rates (MAZE, FETCHPUSH, HANDROTATE, ANTREACH) and return (WALKER, CARRACINg) of all the methods with regards to environment steps in Figure Each task
7

---

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 and Section
Overall, our method DRAIL consistently outperforms prior methods or achieves competitive per- formance compared 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 AlL framework. In out tasks, our DRAIL outperforms DiffAlL, demonstrating that our proposed discriminator provides 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 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 WAll., as shown in the hyperparameter sensitivity experiment inSection D
5.4 Generalizability
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. [3.. Specifically, 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 1x, 1.25x, 1.5x, 1.75x, and 2.0x, compared to the expert environment. That said, 1.5x means the policy is evaluated in an environment with noises 1.5x larger than those injected into expert data collection. Performing well in high noise level setup requires the policy to generalize to unseen states.
M RC
Dithiiiiin Polic - GAI
GAII GP
WAI
DRAI
DRAI
(Ous) ..... Pxpert
1.0
1.0
1.0
Success Rale
0.8
Success Rate
0.6
0.4
0.4
0.4
0.2
0.2
0.2
0.0
10
20
0.0
0
0.0
Step (M)
Step (M)
Step (M)
(a) MAZE
(b) FETCHPUSH
(c) HANDROTATE
1.0
BWW
1000
800
0.8
Return
4000
Return
600
Success
2000
200
0.2
0
0.0
10.0
0
20
0.0
1.0
Step (M)
Step (M)
Step
(d) ANTREACH
(e) WALKER
(f) CARRACING
Figure 4: Learning efficiency. We report success rates (MAZE, FETCHPUSH, HANDROTATE, ANTREACH) and return (WALKER. CCINAC evaluated over five random seeds. Our method DRAIL learns more stably. faster, and achieves higher competitive performance compared to the best- best-performing bascline in all the tasks.

---

Pally
OPNL
Export
Rale
Success Ralts
Sucosss Rate
0.0
Successs
300
00
0.0
Step IM)
Step
Step
Step (M)
(a) FETCHPUSH
(b) FETCHPUSH 1.25x
FETCHPUSH 1.75x
(d) FETCHPUSH 2.0
Figure 5: Generalization experiments in FETCHPUSH. We present the performance of our proposed DRAIL and haselines in the FETCHPUSH task, under varying levels of noise initial states and goal locations. The evaluation spans three random seeds, and the training curve illustrates the success rate dynamics.
The results FETCHPUSH under 1x, 1.25x, 1.75x, and 2.0x noise level are presented in Figure 5 Across all noise levels, approaches utilizing the diffusion model generally exhibit better performance. Notably, our proposed DRAIL demonstrates the highest robustness towards noisy environments. Even at the highest noise level of 2.00x, DRAIL maintains success rate of over 95%, surpassing the best-pent pest-perfrrmig baseline, Diffusion Policy, which achieves only around 79.20% success rate. In contrast, DiffAIL experiences failures (mean: 40.90, standard deviation: 47 45), despite our extensive efforts on experimenting with various configurations and a wide range of hyperparameter values.
out of the 5 seeds, resulting high standard deviation
The extended generalization experiments results in MAZE, FETCHPUSH, HANDROTATE. and the new task FETCHPICK presented in Section C Overall, method outperforms performs competitively against the best-performing baseline, demonstrating its superior generalization ability.
Mi BC . Dww.--ittionnoooe Policy
GAL
GAIL GP
WAIL
DRAL (Ous)
...... Export
0000
Walber20
6000
Whlee20
6000
Wa
6000
Walter2D
4000
4000
Raluun
Return
Return
2000
Step (M)
Step (M) (b) 3 trajectories
Step (M)
Step (M)
(a) 5 trajectories
(c) 2 trajectorics
(d) 1 trajectory
Fetcch
Stro00ss Rata 36
Sucoss
0.0
Suugosss 12
Sucpess
300
Step (M)
Step
Step
(e) 20311 transitions
(f) 10000 transitions
(g) 5000 transitions
(h) 2000 transitions
Figure Data efficiency. We experiment learning with varying amounts of expert data in WALKER and FETCHPUSH. The results show that our proposed method DRAIL more data efficient, iee. can learn with less expert data, compared to other methods.
5.5 Data efficiency
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, and expert trajectories, each containing 1000 transitions; for FETCHPUSH, we use 20311, 1, 10000, 5000, and 2000 state-action pairs. The results reported in Figure (nguure(ddemoottral that our DRAIL learns faster compared 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. contrast, BC's performance is unstable and fluctuating, while the other baselines experience dramatic drop. In FETCHPUSH, our DRAIL maintains a success rate of over 80% even when the data size reduced by 90%, whereas the other AIL baselines' performance drops below 50%.

---

expent didtiuction
Expert Distribution
(b) GAIL Reward
(c) DRAIL Reward
Figure 7: Reward function visualization. We present visualizations of the learned reward values by the discriminative classifier GAIL the diffusion discriminative classifier of our DRAIL. target expert demonstration for imitation is depicted in (a), which is a discontinuous sine function. The reward distributions of GAIL and our DRAIL. are illustrated in (b) and (c), respectively.
5.6 Reward function visualization
To visualize and analyze the learned reward functions, design SINE environment with one- dimensional state and action spaces, where the expert state-action pairs form discontinuous sine wave sin (20sT) + (0, 0.05²), as shown in Figure We train GAIL and DRAIL to learn from this expert state-action distribution and visualize the discriminator output values D& to examine the learned reward function, presented in Figure 7
Figure [b]reveals that the GAIL discriminator exhibits excessive overfitting to the expert demonstra- tion, resulting in its failure provide appropriate reward values when encountering unseen states. contrast, Figure 7c shows that our proposed DRAIL generalizes better to the broader state-action distribution, yielding 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 guide the policy even when deviates from the expert policy. In contrast, the reward distribution from GAIL is relatively narrow outside the expert demonstration, making challenging properly guide the policy the predicted action does not align with the expert.
Conclusion
This work proposes novel adversarial imitation learning framework that integrates diffusion model into generative adversarial imitation learning (GAIL). Specifically, we propose a diffusion diseriminative classifier that employs a diffusion model to construct an enhanced discriminator, yielding more robust and smoother rewards. Then, we design diffusion rewards based the classifier's output for policy learning. Extensive experiments navigation, manipulation, locomotion, and game justify our proposed framework's 's effectiveness, generalizability, and data efficiency. Future work could apply DRAIL to image-based robotic tasks in real-world or simulated environments and explore its potential in various domains outside robotics, such as autonomous driving, to assess its generalizability and adaptability. Additionally, exploring other divergences and distance metrics, such the Wasserstein distance f-divergences, could potentially further improve training stability.
Acknowledgement
This work was supported by NVIDIA Taiwan Research & Development Center (TRDC) under the funding code 22HT911007, Shao-Hua Sun supported the Yushan Fellow Program by the Ministry of Education, Taiwan. Ping-Chun Hsieh is supported in part by the National Science and Technology Council (NSTC), Taiwan under Contract No. NSTC A-49-2222 We thank Tim Pearce and Ching-An Cheng for the fruitful discussions.
10

---

References
[1] Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In International Conference on Machine Learning, 2004.
[2] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. International Conference on Machine Learning. 2017.
[3] Yusuf Aytar, Tobias Pfaff, David Budden, Thomas Paine, Ziyu Wang, and Nando De Freitas. Playing hard exploration games by watching youtube. In Neural Information Processing Systems, 2018.
[4] David M Blei, Alp Kucukelbir, and Jon D McAuliffe, Variational inference: A review for statisticians. Journal of the American statistical Association, 2017.
[5] Shang-Fu Chen, Hsiang-Chun Wang, Ming-Hao Hsu, Chun-Mao Lai, and Shao-Hua Sun. Diffusion model-augmentted behavioral cloning. In International Conference on Machine Learning. 2024.
[6] Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. In Robotics: Science and Systems, 2023.
[7] Paul Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. In Neural Information Processing Systems, 2017.
[8] Robert Dadashi, Léonard Hussenot, Matthieu Geist, and Olivier Pietquin. Primal wasserstein imitation learning. arXiv preprint urXiv:2006.04677 2020.
[9] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. In Neural Information Processing Systems, 2021.
[10] Pete Florence, Corey Lynch, Andy Zeng, Oscar A Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning. In Conference on Robotic Learning, 2022.
[11] Gideon Joseph Freund, Elad Sarafian, and Sarit Kraus. A coupled flow approach imitation learning. In International Conference on Machine Learning, pages 10357-10372. PMLR, 2023.
[12] Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adverserial inverse reinforcement learning. In International Conference on Learning Representations, 2018.
[13] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.072111 2020.
[14] Javier Garcia and Fernando Fernández. A comprehensive survey on safe reinforcement learning. Journal of Machine Leurning Research, 2015.
[15] Divyansh Garg, Shuvam Chakraborty, Chris Cundy, Jiaming Song. and Stefano Ermon. Iq-learn: Inverse soft-q learning for imitation. Advances in Neural Information Processing Systems, 34: 4028 4039, 2021.
[16] Ian Goodfellow, Jean Pouget-Abadic Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Neural Information Processing Systems, 2014.
[17] Shangding Gu, Long Yang, Yali Du, Guang Chen, Florian Walter, Jun Wang, Yaodong Yang, and Alois Knoll. review of safe reinforcement learning: Methods, theory and applications. arXiv preprint arXiv:2205. 10330, 2022.
[18] Jack Harmer, Linus Gisslén, Jorge del Val, Henrik Holst, Joakim Bergdahl, Tom Olsson, Kristoffer Sjöö, and Magnus Nordin. Imitation learning with concurrent actions in 3d games. In IEEE Conference Computational Intelligence and Games, 2018.
11

---

[19] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEEE Conference on Computer Vision and Pattern Recognition, 2016.
[20] Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. In AAAI Conference on Artificial Intelligence, 2018.
[21] Jonathan Ho and Stefano Ermon, Generative adversarial imitation learning. In Neural Informa- fion Processing Systems, 2016.
[22] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Neural Information Processing Systems, 2020.
[23] Bo-Ruei Huang, Chun-Kai Yang, Chun-Mao Lai., and Shao-Hua Sun. Diffusion imitation from observation. In Neural Information Processing Systems, 2024.
[24] Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne. Imitation learning: survey of learning methods. ACM Computing Surveys, 2017.
[25] Rohit Jena, Changliu Liu, and Katia Sycara. Augmenting gail with be for sample efficient imitation learning. In Conference on Robot Learning. 2021.
[26] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. In Neural Information Processing Systems, 2021.
[27] Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In Interna- tional Conference on Learning Representations, 2015.
[28] B Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A Al Sallab, Senthil Yogamani, and Patrick Pérez. Deep reinforcement learning for autonomous driving: A survey, IEEE Transactions on Intelligent Transportation Systems, 2021.
[29] Po- Po-Chen Ko, Jiayuan Mao, Yilun Du, Shao-Hus Sun, and Joshua B. Tenenbaum. Learning act from actionless videos through dense correspondences. In International Conference on Learning Representations, 2024.
[30] Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tompson. Discriminator-actov aat.ttorrrrtooo-rttiiic Addressing sample inefficiency and reward bias in adversarial imitation learning. In International Conference on Machine Leurning, 2019.
[31] Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tompson. Discriminator-actor actortoorrrttoo-titiiici Addressing sample inefficiency and reward bias in adversarial imitation learning. In International Conference on Learning Representations, 2019.
[32] Youngwoon Lee, Shao-Hua Sun, Sriram Somasundaram, Edward S. Hu, and Joseph J. Lim. Composing complex skills by learning transition policies. In International Conference on Learning Representations, 2019.
[33] Youngwoon Lee, Andrew Szot, Shao-Hua Sun, and Joseph J. Lim. Generalizable imitation learning from observation via inferring goal proximity. Neural Information Processing Systems, 2021.
[34] Jan Leike, David Krueger, Tom Everitt, Miljan Martic, Vishal Maini, and Shane Legg. Scalable agent alignment via reward modeling: a research direction. arXiv preprint arXiv: 181 1.0787,, 2018.
[35] Alexander C Li, Mihir Prabhudesai, Shivam Duggal, Ellis Brown, and Deepak Pathak. Your diffusion model is secretly a zero-shot classifier. In IEEE International Conference on Computer Vision, 2023.
[36] Abdoulaye O Ly and Moulay Akhloufi. Learning to drive by imitation: An overview of deep behavior cloning methods. IEEE Transactions on Intelligent Vehicles, 2020.
[37] Andrew Ng and Stuart Russell. Algorithms for inverse reinforcement learning. In International Conference on Machine Learning, 2000.
12

---

[38] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, 2021.
[39] Felipe Nuti, Tim Franzmeyer, and João F Henriques. Extracting reward functions from diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
[40] Manu Orsini, Anton Raichuk, Léonard Hussenot, Damien Vincent, Robert Dadashi, Sertan Girgin, Matthieu Geist, Olivier Bachem, Olivier Pietquin, and Marcin Andrychowicz. What matters for adversarial imitation learning? In Neural Information Processing Systems, 2021.
[41] Takayuki Osa, Joni Pajarinen, Gerhard Neumann, J Andrew Bagnell, Pieter Abbeel, Jan Peters, et al. An algorithmic perspective on imitation learning. Foundations and Trends@ in Robotics, 2018.
[42] Tim Pearce, Tabish Rashid, Anssi Kanervisto, David Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcareel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, and Sam Devlin. Imitating human behaviour with diffusion models. In International Conference on Learning Representations, 2023.
[43] Matthias Plappert, Marcin Andrychowicz, Alex Ray, Bob McGrew, Bowen Baker, Glenn Powell, Jonas Schneider, Josh Tobin, Maciek Chociej, Peter Welinder, et al. Multi-goal reinforcement learning: Challenging robotics environments and request for research. arXiv preprint arXiv:1802.004664 2018.
[44] Dean A Pomerleau. Alvinn: An autonomous land vehicle in a neural network. In Neural Information Processing Systems, 1989.
[45] Harish Ravichandar, Athanasios Polydoros, Sonia Chernova, and Aude Billard. Recent advances in robot learning from demonstration. Annual review of control, robolics, and autonomous systems, 2020.
[46] Moritz Reuss, Maximilian Li, Xiaogang Jia, and Rudolf Lioutikov. Goal-conditioned imitation learning using score-based diffusion policies. arXiv preprint arXiv:2304.02532, 2023.
[47] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-assisted intervention Intervention-MICCAI 2015: 18th international conference, Munich, Germany, October 5.9, 2015, proceedings. part III 18, 2015.
[48] Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. structured prediction to no-regret online learning. In International Conference on Artificial Intelligence and Statistics, 2011.
A reduction of imitation learning and
[49] Stefan Schaal. Learning from demonstration. In Neural Information Processing Systems, 1997.
[50] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 2017.
[51] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, 2015.
[52] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Machine Learning, 2021.
[53] Gokul Swamy, David Wu, Sanjiban Choudhury, Drew Bagnell, and Steven Wu. Inverse reinforcement learning without reinforcement learning. In International Conference on Machine Learning, 2023.
[54] Umar Syed, Michael Bowling, and Robert E Schapire. Apprenticeship learning using linear programming. In Proceedings of the 25th international conference on Machine learning, 2008
[55] Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. In International Joint Conference Artificial Intelligence, 2018.
13

---

[56] Faraz Torabi, Garrett Warnell, and Peter Stone. Generative adversarial imitation from observa- tion. In International Conference on Machine Learning. 2019.
[57] Patrick von Platen, Suraj Patil, Anton Lozhkov, Pedro Cuenca, Nathan Lambert, Kashif Rasul, Mishig Davaadorj. Dhruv Nair, Sayak Paul, William Berman, Yiyi Xu, Steven Liu, and Thomas Wolf. Diffusers: State-of-the-art diffusion models. Sttts://qithun com/ huggingface diffusers 2022.
[58] Bingzheng Wang, Yan Zhang, Teng Pang, Guoqiang Wu, and Yilong Yin. Diffail: Diffusion adversarial imitation learning. arXiv preprint urXiv:2312.06344 2023.
[59] Bingzheng Wang, Guoqiang Wu, Teng Pang, Yan Zhang, and Yilong Yin. Diffail: Diffusion adversarial imitation learning. AAAI Conference Artificial Intelligence, 2024.
[60] Tony Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn. Learning fine-grained bimanual manipulation with low-cost hardware. In Robotics: Science and Systems, 2023.
[61] Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In AAAI Conference on Artificial Intelligence, 2008.
[62] Konrad Zolna, Scott Reed, Alexander Novikov, Sergio Gomez Colmenarejo, David Budden, Serkan Cabi, Misha Denil, Nando de Freitas, and Ziyu Wang. Task-relevant adversarial imitation learning. In Conference on Robotic Learning, 2021.
14

---

Appendix
Table of Contents
Relation to DiffAII.
B Environment task details
B.1 MAZE
B.2 FETcHPUSH & FETCHPICK
17
B.3 HANDROTATE
17
B.4 ANTREACH
17
B.5 WALKER
17
B.6 CARRACING
18
B.7 Expert performance
18
C Extended results of generalization experiments
18
C.I. Experiment settings
18
C.2 Experiment results
18
D Hyperparameter Sensitivity Experiment
19
E Converged performance
20
F Model architecture
20
F.1 Model architecture of DRAIL, DiffALL, and the baselines
20
F2 Image-based model architecture DRAIL, DiffAlL, and the baselines
21
Training details
23
Training hyperparamters
23
G.2 Reward function details
23
H Limitations
24
1 Computational resources and time
24
1.1 Computational resources
24
1.2 Computational time
24
J Impact statements
25
A Relation to DiffAIL
Among all the related works, DiffAIL (58 which also employs a diffusion model for adversarial imitation learning, is the closest our work. This section describes the differences of work and DiffAIL.
Unconditional diffusion models (DiffAIL) conditional diffusion models (DRAIL), Wang et al. 58) proposed to leam an unconditional diffusion model De (s, a) to denoise state-action pairs from expert and an agent. The diffusion model should denoise well when the ittt-acction pairs 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 conditional diffusion model (s,,, a, c) that conditions on the real label e+ or the fake label c In the following, we discuss how using conditional diffusion model leads to better learning signals for policy learning.
DiffALL set up the discriminator output (s, a) = Cami(s,a) E [0, 1]. indicating how likely (s, a) is sampled from the expert distribution. That said, Canr(sa.)) close to 1 implies that the sample is likely from expert and small value of e (Lai(8.a) means it is less likely that this sample
15

---

comes from the expert. However, e Can((ss does not explicitly represent the probability of the "negative class" (agent), which makes difficult 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 the "negative class" (agent). In DiftAIL, the predicted probability of the "positive class" (expert) is e EL and the predicted probability of the Lam "negative class" (agent) can be given by 1-ee----- Lam Following this definition, the discriminator thinks the sample if from an expert when E Fair > - e Can, We also know that Lair > since it is the diffusion model denoising loss. This gives us < Cair I 2 for the expert class, and In < Cait for the agent class. This can be problematic since the fixed boundary in 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 given state-action pair (s, a) is defined as:
(-ma(aaa,et)) (ø(s, C-Lan(B,a,ct) + C--iis.... = (Lan(ss,a,c a, (s, a, a,
(7)
and the predicted probability of the negative (agent) class of the same state-action pair (s, a) is explicitly defined
e (am
1 - (s,,a)))) = c-Lan(s,a,c)) - Lan(s,a.c = OC a, ct) Laitt a, as,
(8)
This binary classification formulation aligns with GAIL. Hence, the discriminator would think given (s, a) comes from expert data only when
_ann(s,a.c")
-arr(s,a.e"
(9)
e (an(x, + - Car(za.e-)
(ar(s,,, + Can(s,a,e-)
which means
e c_aar(s,,.o))
- -Lan(s,,a
(10)
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 ANTREACII), locomotion (WALKER and ANTREACH), robot arm manipulation (FEtCHPUSH and FETCHPICK), robot arm dexterous (HANDROTATE), and games (CARRACING). Additionally, we present experimental results on generalization unseen states and goals on the goal-oriented tasks, and on varying amounts of expert data.
B Environment & task details
B.1 MAZE
Description. In 2D maze environment, a point-maze agent learns to navigate from a starting location to goal location. The agent achieves this by iteratively predicting its and y velocity. The initial and final positions of the agent are randomly selected. The state space includes position,
16

---

velocity, and goal position. The maximum episode length for this task is set at 400, and the episode terminates if the goal reached earlier.
Expert dataset. The expert dataset comprises 100 demonstrations, which includes 18, 525 transitions provided by Lee al. 33
B.2 FETCHPUSH & FETCHPICK
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 object from table and move it to 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 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. 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
B.3 HANDROTATE
Description. In the task HANDROTATE proposed by Plappert et al. [4], 24-DoF Shadow Dex- terous 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 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 within [- I fel and [ 2 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).
B.4 ANTREACH
Description. The ANTREACH task features four-leg ant robot reaching randomly assigned target position located within a range of half-cirole with a radius of 5 meters. The task's 's state is represented by 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 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
B.5 WALKER
Description, WALKER task involves guiding agent move towards the x-coordinate 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 episode is the cumulative result of all time steps within that episode. The 1 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 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 expert dataset.
17

---

B.6 CARRACING
Description. In the CARRACING task, the agent must navigate track by controlling a rear- wear-wheel drive car. The state space of CARRACING is represented by top-down 96 x 96 RGB image capturing the track, the car, and various status indicators such 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 to 64 x 64 pixels. We then concatenate two consecutive frames to form a single state data, providing temporal context for the agent's decision-making process.
grayscale and resizing
Expert dataset. We trained a PPO expert policy on CARRACING environment and collected 671 transitions as expert demonstrations.
B.7 Expert performance
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.
C Extended results of generalization experiments
C.1 Experiment settings
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. 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 collected in lower noise level setting, 1x. Yet, the agent is trained within environment incorporating 1.5x 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, train agents in different noise levels: 1 xx, x.25x, 1.75x, 2.0x. (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 demonstration collection, the demonstrations may inadequately cover the entire state space, 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 broader distribution of initial states.
C.2 Experiment results
MAZE. Our DRAIL outperforms baselines or performs competitively against DiffAIL across all demonstration coverages, as shown Figure Particularly, BC, WAIL, and GAIL's performance decline rapidly in the low coverage case. In contrast, diffusion model-based AlL algorithms demon- strate sustained performance, as shown in Figure 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 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 out of the seeds, resulting in a high standard deviation (mean: 39.17, standard deviation: 47.98) despite our thorough exploration of various settings for configuration. In FETCHPUSH, our method DRAIL. exhibits more robust results, in generalizing to unseen states compared to other baselines. This showcases that the diffusion
18

---

-R.OOR Patky
Dxpor
Su00es Raate
Sucoess Rate
Success State
04
0.0
00g
0%
Step
(a) MAZE 100%
(b) MAZE 75%
(c) MAZE 50%
(d) MAZE 25%
Su0225 Ralo
Success Rate
10
Stop (M)
(M)
Stop (M)
Stop
(e) FETCHPICK 1.00x
(1) FETCHPICK 1.25x
(g) FETCHPICK 1.75x
(h) FETCHPICK 2.0
1.0
10
Su00es5 Rate
Success Rate
Success Rate
90
0.0g
Step (M)
Step (M)
Step (M)
Stop (M)
FETCHPUSH 1x
FETCHPUSH 1.25x
(k) FETCHPUSH 1.75x
(1) FETCHPUSH 2.0
Suggess Fale
Pale
Success Rote
8
Rabe
06
Sucress
02
003
0.0
00g
0.0g
Step
Step (M)
Step (M)
Step (M)
(m) HANDROTATE X (n) HANDROTATE 1.25 (o) HANDROTATE 1.75x (p) HANDROTATE 2.0x
Figure 8: Extended results of generalization experiments. MAZE is evaluated with different coverages of state and goal locations in the expert demonstrations, while FETCHPICK, FETCHPUSH, and HANDROTATE environments evaluated in environments of different noise levels. The number indicates the amount of additional noise agent learning compared that in the expert demonstrations, with more noise requiring harder generalization. The noise level rises from left to right.
reward guidance could provide better generalizability for the AIL framework. Moreover, our DRAIL quite sample-efficient regarding interaction with the environment during training compared to other baselines in FETCUPICK and FETCHIPUSII environment.
HANDROTATE. DiffAIL and Our DRAIL show robustness to different levels of noise in HANDRO- TATE, as illustrated in Figure g Specifically, DiffAIL and our DRAIL achieve a success rate of higher than 90% at noise level of 2.0, while GAIL and WAIL only reach approximately 42.82% and 8.80% respectively.
D Hyperparameter Sensitivity Experiment
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.
19

---

Maze
Maze
1.0
1.0
Success Rate
0.8
Success Rate
0.8
0.6
0.6
0.5x
0.5x
0.4
0.4
0.2x
5x
5x
0.2
0.2
1x (Ours)
1x (Ours)
0.0
0.0
10
20
10 Step (M)
20
Step (M)
(a) Discriminator Learning Rate
(b) Policy Learning Rate
Figure 9: Hyperparameter Sensitivity of DRAIL in the MAZE Environment. The results show the perfor-
mance of DRAIL under varying learning rates for the discriminator @ and policy (b). Different scaling factors (5x, 2x. 1x, 0.Sx, 0.2x) of the baseline learning rate tested. The results demonstrate that DRAIL remains robust across these variations, maintaining stable performance the MAZE environment.
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.
E
Converged performance
This section reports the quantitative results 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
Model architecture
This section presents the model architecture of all the experiments. Appendix architecture of all methods used in Section 5.3
F. describe the model
F1
Model architecture of DRAIL, DifTAIL, and the baselines
In Section 5.3 we conducted a comparative analysis between our proposed DRAIL, along with several baseline approaches (BC, Diffusion Policy, GAIL,, GAIL--P, WAlL,, and DiffAlL.) across six diverse environments. We applied Multilayer Perceptron (MLP) for the policy of BC, the conditional diffusion model in Diffusion Policy, well the policy and the discriminator of GAIL,, GAIL--P, 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 for all diffusion models in this paper
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 2
BC. We maintained a concise model for the policy BC to prevent excessive overfitting to expert demonstrations. This precaution is taken to mitigate the potential adverse effects 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 the policy's number of layers to 5.
20

---

Table I: Converged performance. We report the quantitative results of the converged performance across all experiments.
Envinnments
Settings
Diflusion Palicy
DiSAII
DRAIL
Main Results
MAZE
61. 61.66/666 3.58%
68.01% = 1..22%
68 08802%
= 8..0%
68.80%
= -1.77%
94.00%% + 1.67%
96.88%
0..4%%
FITCHPUSH
82-00% + 21 13%
69.405% + +-2008%
73.77%%
+.37 .55%
91.6455
0.005%
95.66%
S.1555
95.78%
0..8%%
HANDROTAVE
91.31% 2.61%
90.24%
61.615
43.825%
29.255
97.49%
0.48%
93.85%
3..05%
ANTREACH
53.00/5 4.34%
33,60%F
1.96%
21.08%
5.00%%
33.84%%
2.655
29 3093
1.05%
39.02%
1.40%
WALKER
6431.19 267.40
4749.03
2230,95
2080.11
677..0
2321.00 1021000
4639.
G381
115.25
CARRACING
536.92
45.02
689.62
7/60.35
Generalization Expertments
25%
122.67% d: 6.1186
52.101% d: 0.8225
38.115% =: 91.30%
722 53.555%%
= 6..80%
89.67% = 4.335%
91.00%
3.34%
MAZE
50%
75%
00/5
6058 3.0885
0.00%6
68 08.80%%-
56,005F
2. 4555
1. 77%%
68.67%
68.62%
8..3%%
5. 5.79%
65. 8595
72.47%
5.77%
1.05%
96.58%
D4 06.0055 + 1.67%%
0.515
96.8355
96 83%
0.55%
0..49%
100%
0.825
98.03%
06.477
0..0%%
69.
98.48%
0.32%
08.20%
0.04%
1,00%
47 40%
10%
92.40%
36.74%
0.00%
0.00%
50%
16.92%
0.007
PETCHPIC
L.00X
50.00%
7-13%
19.01296
1005
1L.0%%
86.
L.75x
76.06%
7.48%
19.155%
2317%
1L.0%%
0.005
2.14%
6.28%
8075
40.735
2.00%
28 87% 6816
76.6436 7.635%
39.177%
47.08%
0.00%
0.00%
0.00%
0.0005
87.2255
1.99955
1.00%
1055
98 60%
0.48%
54.18%
4.87%
99.68%
0.30%
98. 1255 0.75%
00.00% 0.5895
1.25%
83. STAL 17.09
97. 20%
0.983
94.34%
1.62%
10.52%
89.999
37 80.8
95. 845
1.18%
FErcHPPUNI
1.50X
42.00% + 21 18%
80. 40%
06%
83.78%
+- 66.%
81 1446695-
= 9.08%
+ 8.1596
US. 759%
0.8K05
1.75%
71 00% + 25 13%
64.80%
32%
43.71%
+ 51%
0.00%% 0.000%
55. 62%
45 5655
96. 155% 2: 2:95
2,00x
61, 42% 1 71/%
79.20%% 701%
16,997
25.008
0.00% 0,000
50.50%
47.50%
99.09% 0.47%
L
94. 00% a: 2,5956
90.06% 2 101
72.78%
35.57%
51.95%
88.88%
95. 60%
0.10%
996 31%
0. 3.39%
1,25x
00/66 2.9455
92.71%
27%
05.00%
+-111.2 21%
26. 90%
38 88%
98. 98.03%
0.615
97.17%
1.5095
HANDROTATE
1,50x
91.31 31.31 2.61%
90.24%
615
63.825%
315%
65%
97.497 40%
0.48%
3.25%%
1.76x
1.7056
82% 92.725
992.00%
611.11%%
.10%
28.12%
20%
0.25%
91.7.8%
8.18%%
8.90%%
9.77%
0.80%
0.77%
96.33%
76, 60%
75
355
0.65%%
ANTREACH
0.03x
64.67%
5.25%
63.33%
33.0075
21 059
13.75%
0596
8355
1.47%
1.05%
02.
1.011
0.05%
17.30% 2.8756
21.33%6
0.94%
50.433
1.088%
21.9005
2.15%
13 5855
+ 0.8155
20.9255
9.7455
Duta Elliciency
20311
82 00%,
13%
59,40% 1. 2.00%
33.733% 1.37 665%
01.64%%
0,08%
95.66 05.66% = 5.15%
95.78%
0.80%
100000
N0 20%
71%%
90.20%
6..64%
24.18% 21.11.77 75.%
62.10%
05. 45.625
95.31%
0..4%
FETCHPUSH
5000
56.60%
5955 48.61%
95.70%
0.75%
2000
80.20%
14.70%
84.40% 1.83%
30.28% 24.075
41.21%
47.47%
37.37%
86.44% 13.44%
taas
5431.
267.40
4749.00
223.96
a0s1.18
574.20
2521.07 =. M63
4639.47 4. 282
5381.13
trajs
73.67
3478.225-1
31 33001
2837.70
1028.000 0G
3200.00
4584.24
200.69
6266641
+ 67.93
WALKER
2 trajs
4957 94
658 54
1872.29
328.88
2523.08
888.00
2067.75
409.00
$250, 25
1810.07
5083.32
119.98
mm)
3066.27
à: 1834.77
1165.22
d: 147.79
960066 =: 731.99
1017.87 07.21
3067.36
1530.88
49600.10
164.57
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 L2 re-gguulaiizttio. The regularization value is provided in Table 2
DiffAIL. In DiffALL, 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.
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 cm 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 t and e are reported in Table 2
DRAIL.
F.2 Image-based model architecture of DRAIL, DHITAIL, and the baselines
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 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 51 1i-dimmesio feature vector representing the state data. This state feature vector is subsequently processed by an MLP with three layers, each
21

---

Table 2: Model architectures of policies and discriminators. We report the architectures used for all the methods all the tasks. Note that denotes the neural network policy, represents multilayer perceptron discriminator used GAIL, GAIL-GP, and WAlL, De represents diffusion model discriminator used DiffALL and our method DRAIL.
Method
Models
Component
MAZE
PETCHPIKK
FETCHPPSHH
HANDROTATE
ANTREACI
WALKER
Layers
3
4
Input Dim.
16
68
132
17
Hidden Dim.
256
256
256
512
256
256
Output Dim.
20
Layers
Diffusion Policy
Input Dim. 8
20
19
88
140
23
Hidden Dim.
256
1200
1200
2100
1200
1024
Output Dim.
20
ayers
3
4
Input Dim. 8
20
19
88
140
23
D
Hidden
64
GAIL
Output Dim.
GAIL-GP
Layers
3
3
3
3
3
Input Dim.
68
Hidden Dim.
64
256
64
256
256
Ouiput Dim.
2
4
20
8
Layers
3
4
5
5
Input Dim.
20
88
D
Hidden Dim.
64
64
64
128
64
64
Ourput Dim.
1
WAIL
Reg. Value
0.01
Layers
3
Input Dim.
6
16
16
68
132
17
Hidden
64
Output Dim.
2
20
Layers
5
4
3
5
5
D3
Input Dim.
88
140
Hidden Dim.
128
128
1024
128
1024
1024
Output Dim.
8
20
19
88
140
23
DiffAIL
Layers
3
3
Input Dim.
68
Hidden Dim.
64
64
256
64
256
256
Ourput Dim.
2
4
20
8
Layers
3
Input Dim.
18
98
150
De
Hidden Dim.
128
128
1024
128
1024
1024
Ourput
88
DRAIL
Label Dim. lel
10
10
10
10
# Layers
3
3
3
3
3
Input Dim.
68
132
17
Hidden Dim.
64
64
256
64
256
256
Output Dim.
2
4
20
8
having hidden dimension of 256, to predict the appropriate action. In Diffusion Policy, we only use the downsampling part to extract features.
Diffusion Policy. Diffusion Policy represents policy as a conditional diffusion model, which state and randomly sampled noise. Our condition diffusion model is implemented using the diffusers package by von Platen et al. [57]. The state CARRACING image of size 64 x 64, so first use convolutional neural network (CNN) extract the feature. The CNN based on U-Net [47] structure, comprising 3 down- nown-sampling blocks. Each block consists of ResNet [19 layers, with group normalization applied using 4 groups. The channel sizes for each pair of down-sampling blocks are 4, 8, and 16, respectively.
predicts an action conditioning
Discriminator of GAIL, GAIL-GP, & WAIL. The discriminators of GAIL, GAIL-GP, and WAIL are similar to the policy model: the only difference is that the last linear layer outputs a 1-dimensional value that indicates the probability of given state-action pair being from the expert demonstrations.
Discriminator DiffAIL & DRAIL. Our diffusion model implemented using the diffusers package by von Platen et al. [57] The architecture of both DiffAIL and DRAIL is based on U-Net
22

---

Table 3: Hyperparameters. This table provides overview of the hyperparameters used for all methods across various tasks. Th denotes the learning rate of the discriminator, while 7/. denotes the learning of the policy.
Method
Hyperparameler
MAZE
FETCHPICK
FETCHPUSH
HANDROTATE
ANTREACH
WALKER
CARRACING
Leaming Rase
0.0008
0.0001
BC
Batch Size
128
128
128
128
128
128
128
# Epochs
2000
1000
1000
5000
1000
1000
25001
Leaming Rate
0.00001
0.0001
0.00001
0.0001
Dillosion Policy
Batch Size # Epochs
128
128
128
128
128
128
128
20000
20000
10000
2000
10000
5000
100000
GAIl.
NV
0.001
0,00001
0U000OX
0.0001
0.0001
00000005
0
Vh Bov. Steps
0,00008
0.9001
0.0001
GAIL-GP
25000000
25000000
5000000
5000000
100XXXX
25000000
2000000
0.00001
0.0001
0,00008
0.0001
0,00001
0,000008
0.0001
Rhe Bav. Steps
0.0005
0.0001
0.0001
25000000
25000000
5000000
5000000
1000000
25000000
2000000
0.0011
0.9001
0.0001
0.0001
0.0001
0,0001
0,00001
DiTAIL
0.0001
0.0001
0,000055
0.3001
0.0011
0
Steps
25000000
25000000
5000000
5000000
10000X0X
25000000
0.001
0.0001
0.001
0.0001
0.001
0.0002
0.0001
DRAIL (Ours)
2/ Stops
0.0001
0.0001
0.0001
0.0001
0.0001
0.0001
0.0001
25000000
25000000
5000000
25000000
20000000
Table 4: PPO training parameters. This table reports the PPO training hyperparameters used for each task.
Hyperparameter
MAZE
FETCHPiCK
FETCHPUSH
HANDROTATE
ANTREACH
WALKER
CARRACING
Clipping Range
0.2
0.2
0.2
0.2
0.2
0.2
Discount Factor
0.99
0.99
0.99
0.99
GAE Parameter
0.95
0.95
0.95
0.95
0.95
0.95
0.95
Value Function Coefficient
0.5
Entropy Coefficiet
0.0001
0,0001
0,001
0.0001
0,001
0,001
[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 et and the fake label c Finally, 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 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 represent e+ and to represent c as condition data. For DiffAll,, since condition labels are not required, we simply assign a constant value of O as the condition label.
G Training details
G.1 Training hyperparamters
The hyperparameters employed for 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 WALL, 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 cooreespodd to the noise levels presented in Figure 4
G.2 Reward function details
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 DiffAlL. and GAIL. CARRACING, we observe that adapting GAIL's optimization objective could lead to better performance; hence, we use it for DRAIL, DiffAlL, and
23

---

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, optimize policies for all the AlL methods. We maintain all hyperparameters of PPO constant across methods for given task, except the learning rate, which is adjusted for each method. The PPO hyperparameters for each task are presented in Table
H
Limitations
This work presents an adversarial imitation learning framework DRAIL by employing a diffusion model discriminator. While DRAIL achieves encouraging results various domains, including robot arm manipulation, robot hand dexterous manipulation, locomotion, and games, our proposed framework is fundamentally limited 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 the demonstrators are suboptimal. Finally, DRAIL with imitation learning nature, is not designed to learn from environmental rewards; therefore, even when environments can provide rewards, there no apparent mechanism to utilize them with the current formulation of DRAIL.
I Computational resources and time
I.1 Computational resources
For our experiments, we used the following three workstations:
Machine 1 & Machine 2: ASUS WS880T workstation
-CPU: an Intel Xeon W-2255 (10C/20T, 19.25M, 4.5GHz) 48-Lane CPU
-GPUs: an NVIDIA RTX 3080 Ti GPU and an NVIDIA RTX 3090 GPU
Memory: 128GB memory
Machine 3: ASUS WS880T workstation
CPPU: an Intel Xeon W-2255 (10C/20T, 19.25M, 4.5GHz) 48-Lane CPU
GPUs: two NVIDIA RTX 3080 GPUs
Memory: 128GB memory
1.2
Computational time
In the following, we report the total approximate training GPU hours for all algorithms across all environments, with each algorithm trained on 5 random seeds.
Main Experiments: 1945 GPU hours
Generalization Experiments: 7300 GPU hours
Data Efficiency Experiments: 2920 GPU hours
Reward Function Visualization Experiments: 8 GPU hours
We conducted all the experiments on the following three workstations:
M1: ASUS WS880T workstation with an Intel Xeon W-2255 (10C/20T, 19.25M, 4.5GHz) 48-Lane CPU, 64GB memory, an NVIDIA RTX 3080 GPU, and an NVIDIA RTX 3090 Ti GPU
M2: ASUS WS880T workstation with an Intel Xeon W-2255 (10C/20T, 19.25M, 4.5GHz) 48-Laane CPU, 64GB memory, NVIDIA RTX 3080 Ti GPU, and NVIDIA RTX 3090
Ti GPU
M3: ASUS WS880T workstation with an Intel Xeon W-2255 (10C/20T, 19.25M, 4.5GHz) 48-Lane CPU, 64GB memory, and two NVIDIA RTX 3080 Ti GPUs

---

Impact statements
In this work, we propose 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.
25

---

NeurIPS paper checklist
1. Claims
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-Rewart 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 the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made 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.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
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 Ap- pendix 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 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 few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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 seale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
26

---

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification: This paper does not include theoretical results. We elaborate that the design of our diffusion discriminative classifier aligns with the GAIL discriminator (21), learning a policy with the classifier enjoys the same theoretical guarantee as stated in GAIL.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Thcorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We described every step of our method Section 4.2] and provided the model architecture and details of how our experiment employed the conditional diffusion model in Appendix Appendix Additionally, we provided the origin of our expert datasets in Appendix B
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
27

---

(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome 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., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions 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.
Please the NeurlPS code and data submission guidelines (https://nips... cc/
public/ guid/s/ Sewuddee/odeeuunnneeee Code swriommisssooooooo...mm ssionPolicy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
httppsscccpbblc//ggcg/ Cttl://gggidees/oeeuuuubmb Cote/ asi iinoo icy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, 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 Section model architecture and the policy update algorithm in Section 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 the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
28

---

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We ensure statistical significance by including multiple runs with different random seeds in each experiment. Error bars are included in every figure in the paper, and we provide numerical results in [ Additionally, we consider confidence intervals to demonstrate the significance of observed differences. Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report I-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: We provide detailed information about the computational resources used for the experiments, including the specific workstations, CPUs, GPUs, and memory confi sonfigurations, in Section [.T] and provide the approximate total GPU hours in Section [.2
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https::/neurrssii httpp/ccccppcccctcccccccp csGuide nes?
Answer: [Yes]
Justification: Our research conducted this paper aligns with the principles outlined the NeurlPS Code of Ethics. We have thoroughly reviewed the guidelines and ensured that our research respects the rights and dignity of individuals, promotes fairness, and prioritizes socictal well-being.
29

---

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 cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification: We have considered the potential implications of our work biases and discriminatory behaviors, as highlighted in our impact statement (Appendix
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 societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious 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 particular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out. For example, it is legitimate point that 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 train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and funetioning correctly, harms that could arise when the technology is being used 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper does not pose risks associated with the release of high-risk data models. We use self-trained gymnasium dataset, which does 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 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.
30

---

We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
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?
Question:
Answer: [Yes]
Justification: We gave credit to the diffusion model used in our image-based experiments and the expert datasets for 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 used and, 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, paper saperswil wwwrrsswtthte thooce com/ dat astasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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 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 obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
31

---

According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based on the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
32


Reviews:
[{"review_id": "vcSF892zbB", "reviewer": "Reviewer_uNxe", "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.", "questions": "- What denoise step values did you try? And how did they perform?\n- Could you explain why the proposed reward model performs better than DiffAIL?\n- 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?\n- 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?\n- Are you taking the argmax or sampling from the learned imitation learning policies?\n- Did you try other divergences besides Jensen-Shannon?\nMissing citations [1], [2],\n\n\n[1] A Coupled Flow Approach to Imitation Learning by Freund 2023\n\n[2] Diffusion Model-Augmented Behavioral Cloning by Chen et al. 2024", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper was well written, and the author's contribution was easy to follow.\n- The authors performed a thorough empirical investigation to show how their proposed approach compared against baselines.\n- The authors propose an algorithm that takes full advantage of recent advancements in generative modeling.", "weaknesses": "- The paper needs more novelty. DiffAIL proposed using diffusion as a reward, and DiffAIL also proposed using the training error as a reward signal.\n- The paper needs to include comparisons to the two missing cites [1] and [2].\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the thorough and constructive comments. Please find the response to your questions below.\n\n> The paper needs more novelty. DiffAIL proposed using diffusion as a reward, and DiffAIL also proposed using the training error as a reward signal.\n\n> ​​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.\n\n> Could you explain why the proposed reward model performs better than DiffAIL?\n\nWe explicitly discuss how our proposed method differs from DiffAIL and why our proposed method produces better rewards in Section A. We will reorganize the paper to make this clear from the main paper.\n\nDiffAIL uses how well a diffusion model can denoise a state-action pair to indicate the “realness” of the pair for adversarial imitation learning. To this end, DiffAIL learns an unconditional diffusion model to denoise expert state-action pairs well, while denoise agent state-action pairs poorly. In contrast, we propose to directly train an expert/agent binary classifier by optimizing the binary cross-entropy (BCE) loss according to the denoising performance of a conditional diffusion model conditioning on expert/agent label, and our novelty lies in such a design. Our method significantly differs from DiffAIL since we formulate distinguishing expert and agent state-action pairs as a binary classification task instead of a denoising task, and this design aligns better with the GAIL formulation compared to DiffAIL. We extensively discuss this in Section A in the submission. We will revise the paper to highlight this discussion in the main paper to avoid confusion.\n\nMoreover, DiffAIL’s evaluation only considers locomotion tasks, while our work extensively compares our method 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 and on varying amounts of expert data. \n\n> The paper needs to include comparisons to the two missing cites [1] and [2].\n\nWe thank the reviewer for providing these references. We will revise the paper to include the following discussions and cite these works.\n\nFreund et al. (2023) [1] introduce CFIL, which employs normalizing flows for state and state-action distribution modeling in imitation learning. Chen et al. (2024) [2] augment behavioral cloning using a diffusion model focusing on offline imitation learning. In contrast, AIL and our method aim to leverage online environment interactions. Also, we would like to note that Chen et al. (2024) [2] was published at ICML 2024 (July 2024), which is later than the NeurIPS 2024 submission deadline (May 2024).\n\n> What denoise step values did you try? And how did they perform?\n\nAs mentioned in Section E.1, we set the total timesteps $T$ to 1000 in all experiments. During the training phase, we uniformly randomly sample the denoising steps from the range 0 to T. During the inference phase, we followed the same procedure as in the training phase, i.e., uniformly randomly sampling from [0, T]. \n\nAs requested by the reviewer, we additionally experimented with setting the denoising step to constant values \\{250, 500, 750\\}. The result is presented in Figure R.1 in the PDF attached to the overall response. The result shows that following the same denoising step sampling procedure as in the training phase, as adopted in our method, achieved the best performance. Also, consistently sampling with a large denoise step (750) hurts the policy learning performance.\n\n\n> Did you try increasing or decreasing the number of expert samples used for learning in Walker? Are you using D4RL?\n\nAs mentioned in Section 5.1 (line 228), the expert Walker dataset was collected from a PPO expert policy instead of D4RL datasets. We evaluate and discuss the data efficiency, i.e., the amount of expert data required, of all the methods in Section 5.5. Figure 6 shows that our proposed method can learn with fewer demonstrations compared to the baselines in Walker and FetchPush.\n\n> 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?\n\nWe didn’t try training with various levels of optimal demonstration in Walker. We evaluate the ability to generalize to varying levels of noise in Section 5.4 in FetchPush by randomizing initial states and goal locations. The results in Figure 5 show that our proposed method demonstrates the highest robustness towards noisy environments.\n\n> Are you taking the argmax or sampling from the learned imitation learning policies?\n\nWe use PPO as our RL algorithm to learn the agent policy. Hence, the policy is stochastic during training and deterministic during inference.\n> Did you try other divergences besides Jensen-Shannon?\n\nWe use JS divergence following the GAIL’s original setting. Exploring other divergences and distance metrics, such as the Wasserstein distance or f-divergences, is a promising future direction. We thank the reviewer for the suggestion and will revise the paper to include this discussion.\n\n**References**\n\n[1] Freund et al. “A Coupled Flow Approach to Imitation Learning.” In PMLR, 2023.\n\n[2] Chen et al. “Diffusion Model-Augmented Behavioral Cloning.” In ICML, 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BsZ7op02Dt", "reviewer": "Reviewer_3pMZ", "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. \n\nExtensive 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.", "questions": "1. How does DRAIL handle environments with rapidly changing dynamics, and what are the limitations in such scenarios?\n2. 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?\n3. How sensitive is DRAIL to the choice of hyperparameters, and what guidelines can be provided for tuning these parameters in different environments?\n4. Can the diffusion discriminative classifier be extended or modified to further improve the reward robustness and policy performance in more complex tasks?\n5. 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": "Shown in weakness", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. **Innovative Integration**: The integration of diffusion models into GAIL is innovative and addresses the common issue of instability in adversarial imitation learning.\n2. **Comprehensive Experiments**: The authors conduct extensive experiments across diverse domains, including navigation, manipulation, and locomotion, which provide robust evidence of DRAIL's effectiveness.\n3. **Generalizability and Data Efficiency**: The paper demonstrates superior generalizability to unseen states and goals, as well as high data efficiency, making the approach practical for real-world applications.\n4. **Robust Reward Mechanism**: The diffusion discriminative classifier enhances the robustness of the reward functions, contributing to more stable and reliable policy learning.\n5. **Visualizations**: The visualized comparisons of learned reward functions between GAIL and DRAIL effectively illustrate the advantages of the proposed approach.\n\n###", "weaknesses": "1. **Complexity**: The introduction of diffusion models adds complexity to the framework, which may pose challenges in implementation and require significant computational resources.\n2. **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.\n3. **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.\n4. **Scalability**: The scalability of DRAIL in larger and more complex environments is not extensively evaluated, which could be a limitation for broader adoption.\n5. **Hyperparameter Sensitivity**: The sensitivity of DRAIL to various hyperparameters is not discussed, which could impact its robustness across different settings.\n\n###", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the thorough and constructive comments. Please find the response to your questions below.\n\n> Complexity: The introduction of diffusion models adds complexity to the framework, which may pose challenges in implementation and require significant computational resources. 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?\n\nWe would like to clarify that our proposed method does not require significant computational resources, as it only requires two feedforward passes to obtain a reward for a state-action pair. This is because our design does not require going through the diffusion model generation process. Our method exhibits better efficiency compared to the most widely adopted method, GAIL, performing significantly better than GAIL with the same number of feedforward passes computed.\n\n> Limited Scope, Scalability, and rapidly changing dynamics\n\nThis paper aims to propose a fundamental algorithm that improves the AIL framework. While previous AIL papers, such as GAIL, WAIL, and DiffAIL, have only tested their algorithms in locomotion environments, we evaluate our proposed method in various domains, including locomotion (Walker and AntReach), navigation (Maze and AntReach), robot arm manipulation (FetchPush and FetchPick), robot arm dexterity (HandRotate), and games (CarRacing). We believe this sufficiently demonstrates the broad applicability of our proposed method.\n\n>  A deeper analysis of why certain methods perform better in specific tasks could enhance understanding\n\nWe thank the reviewer for the suggestion. We will provide deeper analyses of the performance of different methods on all the tasks in the revised paper.\n\n> How sensitive is DRAIL to the choice of hyperparameters, and what guidelines can be provided for tuning these parameters in different environments?\n\nWe empirically found that our proposed method, DRAIL, is robust to hyperparameters and easy to tune, especially compared to GAIL and WAIL.\n\nLike most AIL methods, the key hyperparameters of DRAIL are the learning rates of the policy and discriminator. We additionally experimented with various values of the learning rates of the policy and discriminator and reported the results in Figure R.2 in the PDF attached to the rebuttal summary. The results show that our method is robust to hyperparameter variations, including 5x, 2x, 1x, and 0.5x.\n\nWe thank the reviewer for inspiring us to conduct this hyperparameter sensitivity experiment. We will revise the paper to include it.\n\n> Can the diffusion discriminative classifier be extended or modified to further improve the reward robustness and policy performance in more complex tasks?\n\nWe empirically observe that the rewards produced by our diffusion model classifier are robust, resulting in stable policy learning curves, as demonstrated in the experiments. It is potentially beneficial to incorporate the recent advancements in developing diffusion models into DRAIL, which is left for future work.\n\n> 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?\n\nWe empirically observe that the training of diffusion models is very stable with smoothly decreased losses. We believe this justifies the design of our BCE loss.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5pCtUBYGdL", "reviewer": "Reviewer_PZWW", "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.", "questions": "- Why does DRAIL only require two reverse steps to obtain an accurate reward? Can you provide a more detailed explanation or analysis?\n- 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?\n- 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.\n\nI would like to raise my score if my concerns are well addressed.", "limitations": "The authors have discussed the limitations of this paper. DRAIL is unable to learn from state-only trajectories and suboptimal trajectories.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The instability of AIL training is a key problem to its application, making research on this issue valuable.\n- DRAIL performs well across multiple benchmarks, improving sample efficiency and training stability.", "weaknesses": "- 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.\n- This paper lacks ablation experiments, including a sensitivity analysis of the hyperparameters.\n- Some important AIL baselines are missing, such as DAC [2] and IQLearn [3].\n\n[1] DiffAIL: Diffusion Adversarial Imitation Learning. AAAI, 2024.\n\n[2] Discriminator-Actor-Critic: Addressing Sample Inefficiency and Reward Bias in Adversarial Imitation Learning. ICLR, 2019.\n\n[3] IQ-Learn: Inverse soft-Q Learning for Imitation. NeurIPS, 2021.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the thorough and constructive comments. Please find the response to your questions below.\n\n> The advantages of the diffusion discriminative classifier compared to diffusion reward [1] are not well explained.\n\nWe extensively discuss how our proposed method differs from DiffAIL and why our proposed method produces better rewards in Section A. We will reorganize the paper to make this clear from the main paper.\n\nDiffAIL uses how well a diffusion model can denoise a state-action pair to indicate the “realness” of the pair for adversarial imitation learning. To this end, DiffAIL learns an unconditional diffusion model to denoise expert state-action pairs well, while denoise agent state-action pairs poorly. In contrast, we propose to directly train an expert/agent binary classifier by optimizing the binary cross-entropy (BCE) loss according to the denoising performance of a conditional diffusion model conditioning on expert/agent label, and our novelty lies in such a design. Our method significantly differs from DiffAIL since we formulate distinguishing expert and agent state-action pairs as a binary classification task instead of a denoising task, and this design aligns better with the GAIL formulation compared to DiffAIL. \n\nMoreover, DiffAIL’s evaluation only considers locomotion tasks, while our work extensively compares our method 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 and on varying amounts of expert data. \n\n> The authors claim that their proposed diffusion discriminative classifier only requires two reverse steps to obtain the corresponding reward. However, I do not fully agree with this. According to equations (4) and (6), DRAIL needs to compute $\\mathcal{L}\\_{diff}(𝑠,𝑎,𝑐^+)$ and $\\mathcal{L}\\_{diff}(𝑠,𝑎,𝑐^-)$ when obtaining rewards. From equation (3), we know both terms require taking an expectation over the diffusion steps 𝑇.\n\nWe would like to clarify that we use sampling to approximate the expectation in Equation 3 instead of sampling all timesteps $T$ and averaging them. In all the experiments, we sample one single denoising time step $t$ to compute the reward, and we empirically found that sampling multiple denoising time steps does not consistently improve performance. \n\n> Why does DRAIL only require two reverse steps to obtain an accurate reward? Can you provide a more detailed explanation or analysis?\n\nOur method redesigns a conditional diffusion model to learn to perform an expert/agent binary classification task. Therefore, it can naturally provide a “realness” reward according to how indistinguishable agent state-action pairs are from the expert state-action pairs. To obtain the “realness” reward in Equation 4 given an agent state-action pair, we have to compute two $\\mathcal{L}_{\\text{diff}}$ in Equation 3 with the positive/expert condition ($c^+$) and the negative/agent condition ($c^-$), and therefore it requires two reverse steps.\n\n>  What are the key hyperparameters for DRAIL? Can you conduct additional ablation experiments on these hyperparameters?\n\nWe empirically found that our proposed method, DRAIL, is robust to hyperparameters and easy to tune, especially compared to GAIL and WAIL.\n\nLike most AIL methods, the key hyperparameters of DRAIL are the learning rates of the policy and discriminator. We additionally experimented with various values of the learning rates and reported the results in Figure R.2 in the PDF attached to the rebuttal summary. The results show that our method is robust to hyperparameter variations, including 5x, 2x, 1x, and 0.5x.\n\nWe thank the reviewer for inspiring us to conduct this hyperparameter sensitivity experiment. We will revise the paper to include it.\n\n> Can you compare DAC and IQLearn in your experiments? I am very interested in seeing a comparison between the diffusion discriminative classifier and the gradient penalty.\n\nWe have conducted additional experiments to address it.\n\n- **Comparison to gradient penalty**: We would like to note that one of our baselines, WAIL, has already implemented the gradient penalty and gradient clipping. As requested by the reviewer, we additionally implemented and evaluated GAIL with the gradient penalty (GAIL+GP) in CarRacing. The results are shown in Figure R.3 in the PDF attached to the rebuttal summary. GAIL with gradient penalty (GAIL+GP) initially shows a faster improvement but is unstable, and its overall performance is inferior to our method, DRAIL.\n- **Comparison to IQ-Learn**: As requested by the reviewer, we additionally implemented and evaluated IQ-Learn in CarRacing. The results in Figure R.3 show that IQ-Learn struggles at this task despite our substantial effort in experimenting with different hyperparameters, including actor learning rate ($10^{-4}$, $5 \\times 10^{-5}$, $3 \\times 10^{-5}$), critic learning rate ($10^{-3}$, $5 \\times 10^{-4}$, $3 \\times 10^{-4}$), and the entropy coefficient ($10^{-1}$, $5 \\times 10^{-2}$, $2 \\times 10^{-2}$, $10^{-2}$, $5 \\times 10^{-3}$, $10^{-3}$), as well as trying various setups, including $\\chi^2$-divergence and gradient penalty.\n- **Comparison to DAC**: The main contribution of DAC is to assign proper rewards to the “absorbing states” that an agent enters after the end of episodes. DAC requires additional annotations to determine if a state is an absorbing state or not. However, all the methods evaluated in our work do not have access to such information, making comparing them to DAC unfair. We believe the contributions of DAC and our work are orthogonal and could be combined.\n\nWe will revise the paper to include the new results and the discussion.\n\n**References**\n\n[1] Wang et al. “DiffAIL: Diffusion Adversarial Imitation Learning.” In AAAI, 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0zlC4G9vMq", "reviewer": "Reviewer_5b3J", "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. \n\nThe 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.\n\nThe 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).\n\nThey 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.\n\nIn terms of raw performance, DRAIL obtains:\n\n- visibly better results than the baselines in 4/6 environments; \n\n- a result very close to that of DiffAIL in 1/6  environment;\n\n- a result worse than DiffAIL in 1/6 environment.\n\nFurther experiments show:\n\n- DRAIL is more robust to noise on initial states and goal positions than the baselines on a block-pushing task (FetchPush);\n\n- DRAIL is more data-efficient than the baselines on the Walker environment (continuous high-dimensional control) and FetchPush;\n\n- DRAIL learns more generalizable rewards than GAIL on a synthetic 2d point dataset.", "questions": "See Weaknesses.", "limitations": "The authors consider the limitations of their work in Appendix G and societal impacts in Appendix I.\n\nMy 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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Clarity: the paper is well-written and the proposed method is explained clearly.\n\n1. Extensive comparison with DiffAIL, albeit in the appendix.\n\n1. Methodological novelty: the authors propose using a conditional diffusion model to parameterize the discriminator for AIL, which removes limitations from DiffAIL (such as an arbitrary decision boundary at $\\ln 2$; see Appendix A).\n\n1. Significance and quality of experimental results: the paper demonstrates that defining the discriminator using a conditional diffusion model improves performance, robustness and data-efficiency, compared to DiffAIL (which uses an unconditional diffusion model in the discriminator) and to other baselines. Hence, one can argue it constitutes a step forward in Adversarial Imitation Learning.\n\n1. Referencing strongly related prior work: the authors point to DiffAIL and Diffusion Classifier, which strongly resemble different parts of DRAIL (namely using the loss of a diffusion model to parameterize a discriminator, and how to construct such a discriminator). I personally was not aware of these two papers before, and commend the authors for bringing these to the reader’s attention.", "weaknesses": "## Better delimiting the author’s contributions\n\nDespite 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. \n\nFor 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. \n\nSimilarly, 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.\n\nHence, 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.\n\nThis 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.\n\nAs 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:\n- Add a section on DiffAIL and Diffusion Classifier to Section 3 (preliminaries).\n- 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.\n- Hence present the final methodology as the combination of DiffAIL and DC.\n\nAgain, 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.\n\n## Baselines against planning methods that use diffusion models\n\nThe 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).\n\nAs 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). \n\nRecent 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.\n\n\n### References\n\nZhu, Zhengbang, et al. \"Diffusion models for reinforcement learning: A survey.\" arXiv preprint arXiv:2311.01223 (2023).\n\nNuti, Felipe, Tim Franzmeyer, and João F. Henriques. \"Extracting reward functions from diffusion models.\" Advances in Neural Information Processing Systems 36 (2024).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the thorough and constructive comments. Please find the response to your questions below.\n\n> 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.\n\n> This still constitutes novelty, ... 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.\n\nWe thank the reviewer for the detailed suggestions for better organizing the paper and discussing the novelty of our work. We would like to clarify that DRAIL is not a combination of DiffAIL and Diffusion Classifier (DC).\n\nDC turns a text-to-image diffusion model, which optimizes the denoising MSE loss, into an image classifier without further training. DiffAIL follows the intuition of DC and uses how well a diffusion model can denoise a state-action pair to indicate the “realness” of the pair for adversarial imitation learning. To this end, DiffAIL learns an unconditional diffusion model to denoise expert state-action pairs well, while denoise agent state-action pairs poorly.\n\nIn contrast, we propose to directly train an expert/agent binary classifier by optimizing the binary cross-entropy (BCE) loss according to the denoising performance of a conditional diffusion model conditioning on expert/agent label, and our novelty lies in such a design. In other words, our method significantly differs from DC and DiffAIL since we formulate distinguishing expert and agent state-action pairs as a binary classification task instead of a denoising task, and this design aligns with the GAIL formulation. Hence, our method is not a combination of DC and DiffAIL.\n\nMoreover, DiffAIL’s evaluation only considers locomotion tasks, while our work extensively compares our method 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 and on varying amounts of expert data. \n\nWe completely agree with and highly appreciate the reviewer’s detailed suggestion for reorganizing the paper. We will revise the paper by presenting DC and DiffAIL in Section 3 and discussing the limitations of DiffAIL in Section 4 to motivate our method. We believe this will make our contributions more clear.\n\n> 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),\n\nWe thank the reviewer for providing these references. We will revise the paper to include the following discussions.\n\n- Diffuser [1] 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. Therefore, it is not trivial to directly compare our method to the diffuser.\n- Nuti et al. [2] focus on learning a reward function, unlike imitation learning, whose goal is to obtain a policy. Hence, Nuti et al. [2] 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, while adversarial imitation learning (AIL) approaches and our method are designed to learn from agent data that continually change as the agents learn.\n\n**References**\n\n[1] Janner et al. “Planning with diffusion for flexible behavior synthesis.” In ICML, 2022.\n\n[2] Nuti et al. “Extracting Reward Functions from Diffusion Models.” In NeurIPS, 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
The attached PDF file contains the following content:
- **[Reviewer uNxe] The Effect of Denoising Time Step**: We experimented with computing rewards using different constant denoising time steps and reported the result in Figure R.1. The result shows that following the same denoising step sampling procedure as in the training phase, i.e., uniformly randomly sample from [0, T], as adopted in our method, achieves the best performance, justifying our design choice.
- **[Reviewer 3pMZ, Reviewer PZWW] Hyperparameter Ablation Study**:  We experimented with varying hyperparameters, including the discriminator learning rate ($\eta_\phi$) and policy learning rates ($\eta_\pi$). The results are shown in Figure R.2, demonstrating that DRAIL maintains robust performance with varying hyperparameters.
- **[Reviewer PZWW] Comparison to Gradient Penalty and IQ-Learn**: We implemented and evaluated two additional baselines, GAIL with Gradient Penalty and IQ-Learn, as suggested by the reviewer. The results shown in Figure R.3 highlight the superior performance of our proposed method.

Meta-review: 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.
Decision: Accept (poster)