Title: Diffusion-Reward Adversarial Imitation Learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: Imitation Learning

Adversarial Imitation Learning

Diffusion Model
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

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/

1 Introduction

Imitation learning, i . e ., learning from demonstration [24, 41, 49], aims to acquire an agent policy by observing and mimicking the behavior demonstrated in expert demonstrations. Various imitation learning methods [53, 60] have enabled deploying reliable and robust learned policies in a variety of tasks involving sequential decision-making, especially in the scenarios where devising a reward function is intricate or uncertain [7, 32, 34], or when learning in a trial-and-error manner is expensive or unsafe [14, 17].

Among various methods in imitation learning, generative adversarial imitation learning (GAIL) [21] has been widely adopted due to its effectiveness and data efficiency. GAIL learns a generator policy to imitate expert behaviors through reinforcement learning and a discriminator to differentiate between the expert and the generator's state-action pair distributions, resembling the idea of generative adversarial networks (GANs) [16]. Despite its established theoretical guarantee, GAIL training is notoriously brittle and unstable. To alleviate this issue, significant efforts have been put into improving GAIL's sample efficiency, scalability, robustness, and generalizability by modifying loss functions [12], developing improved policy learning algorithms [30], and exploring various similarity measures of distributions [2, 8, 12].

Inspired by the recent dominance of diffusion models in generative modeling [22], this work explores incorporating diffusion models into GAIL to provide more robust and smoother reward functions

for policy learning as well as stabilize adversarial training. Specifically, we propose a diffusion discriminative classifier, which learns to classify a state-action pair into expert demonstrations or agent trajectories with merely two reverse diffusion steps. Then, we leverage the proposed diffusion discriminative classifier to devise diffusion rewards, which reward agent behaviors that closely align with expert demonstrations. Putting them together, we present Diffusion-Reward Adversarial Imitation Learning (DRAIL), a novel adversarial imitation learning framework that can efficiently and effectively produce reliable policies replicating the behaviors of experts.

We extensively compare our proposed framework DRAIL with behavioral cloning [44], Diffusion Policy [6, 42], and AIL methods, e . g ., GAIL [21], WAIL [2], and DiffAIL[59], in diverse continuous control domains, including navigation, robot arm manipulation, locomotion, and games. This collection of tasks includes environments with high-dimensional continuous state and action spaces, as well as covers both vectorized and image-based states. The experimental results show that our proposed framework consistently outperforms the prior methods or achieves competitive performance. Moreover, DRAIL exhibits superior performance in generalizing to states or goals unseen from the expert's demonstrations. When varying the amounts of available expert data, DRAIL demonstrates the best data efficiency. At last, the visualized learned reward functions show that DRAIL captures more robust and smoother rewards compared to GAIL.

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 [36], and game AI [18].

Behavioral Cloning (BC). BC [44, 55] imitates an expert policy through supervised learning without interaction with environments and is widely used for its simplicity and effectiveness across various domains. Despite its benefits, BC struggles to generalize to states not covered in expert demonstrations because of compounding error [10, 48]. Recent methods have explored learning diffusion models as policies [6, 42], allowing for modeling multimodal expert behaviors, or using diffusion models to provide learning signals to enhance the generalizability of BC [5]. In contrast, this work aims to leverage a diffusion model to provide learning signals for policy learning in online imitation learning.

Inverse Reinforcement Learning (IRL). IRL methods [37] aim at inferring a reward function that could best explain the demonstrated behavior and subsequently learn a policy using the inferred reward function. Nevertheless, inferring reward functions is an ill-posed problem since different reward functions could induce the same demonstrated behavior. Hence, IRL methods often impose constraints on reward functions or policies to ensure optimality and uniqueness [1, 37, 54, 61]. Yet, these constraints could potentially restrict the generalizability of learned policies.

Adversarial Imitation Learning (AIL). AIL methods aim to directly match the state-action distributions of an agent and an expert through adversarial training. Generative adversarial imitation learning (GAIL) [21] and its extensions [23, 25, 31, 56, 59, 62] train a generator policy to imitate expert behaviors and a discriminator to differentiate between the expert and the generator's state-action pair distributions, which resembles the idea of generative adversarial networks (GANs) [16]. Thanks to its simplicity and effectiveness, GAIL has been widely applied to various domains [3, 28, 45]. Over the past years, researchers have proposed numerous improvements to enhance GAIL's sample efficiency, scalability, and robustness [40], including modifications to discriminator's loss function [12], extensions to off-policy RL algorithms [30], addressing reward bias [31], and exploration of various similarity measures [2, 8, 11, 12]. Another line of work avoids adversarial training, such as IQ-Learn [15], which learns a Q-function that implicitly represents the reward function and policy. In this work, we propose to use the diffusion model as a discriminator in GAIL.

Diffusion Model-Based Approaches in Reinforcement Learning. Diffuser [57] and Nuti et al. [39] apply diffusion models to reinforcement learning (RL) and reward learning, their settings differ significantly from ours. Diffuser [57] is a model-based RL method that requires trajectory-level reward information, which differs from our setting, i . e ., imitation learning, where obtaining rewards is not possible. Nuti et al. [39] focus on learning a reward function, unlike imitation learning, whose goal is to obtain a policy. Hence, Nuti et al. [39] neither present policy learning results in the main paper nor compare their method to imitation learning methods. Moreover, they focus on learning

from a fixed suboptimal dataset, AIL approaches and our method are designed to learn from agent data that continually change as the agents learn.

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 [21] establishes a connection between generative adversarial networks (GANs) [16] and imitation learning. GAIL employs a generator, G θ , that acts as a policy π θ , mapping a state to an action. The generator aims to produce a state-action distribution ( ρ π θ ) which closely resembles the expert state-action distribution ρ π E ; discriminator D ω functions as a binary classifier, attempting to differentiate the state-action distribution of the generator ( ρ π θ ) from the expert's ( ρ π E ). The optimization equation of GAIL can be formulated using the Jensen-Shannon divergence, which is equivalent to the minimax equation of GAN. The optimization of GAIL can be derived as follows:

where ρ π θ and ρ π E are the state-action distribution from an agent π θ and expert policy π E respectively. The loss function for the discriminator is stated as -( E x ∼ ρ π θ [log D ω ( x )]+ E x ∼ ρ π E [log(1 -D ω ( x ))]) . For a given state, the generator tries to take expert-like action; the discriminator takes state-action pairs as input and computes the probability of the input originating from an expert. Then the generator uses a reward function -E x ∼ ρ π θ [log D ω ( x )] or -E x ∼ ρ π θ [log D ω ( x )] + λH ( π θ ) to optimize its network parameters, where the entropy term H is a policy regularizer controlled by λ ≥ 0 .

3.2 Diffusion models

Diffusion models have demonstrated stateof-the-art performance on various tasks [9, 29, 38, 51]. This work builds upon denoising diffusion probabilistic models (DDPMs) [22] that employ forward and reverse diffusion processes, as illustrated in Figure 1. The forward diffusion process injects noise into data points following a variance schedule until achieving an isotropic Gaussian distribution. The reverse diffusion process trains a diffusion model ϕ to predict the injected noise by optimizing the objective:

x

T

Reverse di

ff

usion process

q

(

x

t

|

x

t

-1

)

p

ϕ

(

x

t

-1

|

x

t

)

Forward di

ff

usion process

x

0

x

t

x

t

-1

Learn to denoise

where T represents the set of discrete time steps in the diffusion process, ϵ is the noise applied by the forward process, ϵ ϕ is the noise predicted by the diffusion model, and ¯ α t is the scheduled noise level applied on the data samples.

Beyond generative tasks, diffusion models have also been successfully applied in other areas, including image classification and imitation learning. Diffusion Classifier [35] demonstrates that conditional diffusion models can estimate class-conditional densities for zero-shot classification. In imitation learning, diffusion models have been used to improve Behavioral Cloning (DBC) [5] by using diffusion model to model expert state-action pairs. Similarly, DiffAIL [58] extends GAIL [21] by employing diffusion models to represent the expert's behavior and incorporating the diffusion loss into the discriminator's learning process. However, DiffAIL's use of an unconditional diffusion model limits its ability to distinguish between expert and agent state-action pairs. We provide a detailed explanation of its limitation in Section 4.3 and Section A.

Dϕ

(

s

E

,

a

E

)

(a) Learning Diffusion Discriminative Classifier

0

1

/uni03F5

a

E

s

E

+

+

/uni03F5

a

π

s

π

+

+

Di

ff

usion

Model

ϕ

Di

ff

usion Discriminative Classifier

Dϕ

c

+

c

-

c

+

c

-

Expert

Data

τ

E

Agent

Data

τ

i

'Real'

'Fake'

Dϕ

(

s

π

,

a

π

)

Learning objective

Learnable mapping

Frozen mapping

(b) Learning Policy with Diffusion Rewards

/uni03F5

a

s

+

+

Di

ff

usion

Model

ϕ

c

+

c

-

Policy

π

θ

s

Di

ff

usion Reward

r

ϕ

(

s

,

a

)

log(

Dϕ

(

s

,

a

)) - log(1 -

Dϕ

(

s

,

a

))

Di

ff

usion Discriminative Classifier

Dϕ

Dϕ

(

s

,

a

)

4 Approach

We propose a novel adversarial imitation learning framework incorporating diffusion models into the generative adversarial imitation learning (GAIL) framework, illustrated in Figure 2. Specifically, we employ a diffusion model to construct an enhanced discriminator to provide more robust and smoother rewards for policy learning. We initiate our discussion by describing a naive integration of the diffusion model, which directly predicts rewards from Gaussian noises conditioned on stateaction pairs, and the inherent issues of this method in Section 4.1. Subsequently, in Section 4.2, we introduce our proposed method that employs a conditional diffusion model to construct a diffusion discriminative classifier, which can provide diffusion rewards for policy learning. Finally, the overall algorithm of our method is outlined in Section 4.3.

4.1 Reward prediction with a conditional diffusion model

Conditional diffusion models are widely adopted in various domains, e . g ., generating an image x from a label y . Intuitively, one can incorporate a conditional diffusion model as a GAIL discriminator by training it to produce a real or fake label conditioned on expert or agent state-action pairs. Specifically, given a denoising time step t and a state-action pair ( s , a ) ∈ ( S × A ) , where S , A stand for state and action spaces, respectively, as a condition, the diffusion model p ϕ ( r t -1 | r t , s , a ) learns to denoise a reward label r 0 ∈ { 0 , 1 } , i . e ., 1 for expert (real) state-action pairs and 0 for agent (fake) state-action pairs through a reverse diffusion process.

To train a policy, we can use the diffusion model to produce a reward r given a state-action pair ( s , a ) from the policy through a generation process by iteratively denoising a sampled Gaussian noise, i . e ., noisy reward, conditioned on the state-action pair. Then, the policy learns to optimize the rewards predicted by the diffusion model. Nevertheless, the reward generation process is extremely time-consuming since predicting a reward for each state-action pair from the policy requires running T (often a large number) denoising steps, and policy learning often takes tens of millions of samples, resulting in a billion-level overall training scale. Consequently, it is impractical to integrate a diffusion model into the GAIL framework by using it to predict 'realness' rewards for policy learning from state-action pairs.

4.2 Diffusion discriminative classifier

Our goal is to yield a diffusion model reward given an agent state-action pair without going through the entire diffusion generation process. Inspired by previous work [5, 59], we extract the learning signal from a portion of the diffusion denoising steps, rather than using the entire process. Building

on these insights, we adapt the training procedure of DDPM to develop a mechanism that provides a binary classification signal using just one denoising step.

Our key insight is to leverage the derivations developed by Kingma et al. [26], Song et al. [52], which suggest that the diffusion loss, i . e ., the difference between the predicted noise and the injected noise, indicates how well the data fits the target distribution since the diffusion loss is the upper bound of the negative log-likelihood of data in the target distribution. In this work, we propose calculating 'realness' rewards based on the diffusion loss computed by denoising the state-action pairs from the policy, which indicates how well the state-action pairs fit the expert behavior distributions. We formulate the diffusion loss L diff as follows:

where c ∈ { c + , c -} , and the real label c + corresponds to the condition for fitting expert data while the fake label c -corresponds to agent data. We implement c + as 1 and c -as 0 .

To approximate the expectation in Eq. 3, we use random sampling, allowing us to achieve the result with just a single denoising step. Subsequently, given a state-action pair ( s , a ) , L diff ( s , a , c + ) measures how well ( s , a ) fits the expert distribution and L diff ( s , a , c -) measures how well ( s , a ) fits the agent distribution 1 . That said, given state-action pairs sampled from expert demonstration, L + diff should be close to 0 , and L -diff should be a large value; on the contrary, given agent state-action pairs, L + diff should be a large value and L -diff should close to 0 .

While L + diff and L -diff can indicate the 'realness' or the 'fakeness' of a state-action pair to some extent, optimizing a policy using rewards with this wide value range [0 , ∞ ) can be difficult [20]. To address this issue, we propose transforming this diffusion model into a binary classifier that provides 'realness' in a bounded range of [0 , 1] . Specifically, given the diffusion model's output L + , -diff , we construct a diffusion discriminative classifier D ϕ : S × A → R :

where σ ( x ) = 1 / (1 + e -x ) denotes the sigmoid function. The classifier integrates L + diff and L -diff to compute the 'realness' of a state-action pair within a bounded range of [0 , 1] , as illustrated in Figure 2. Since the design of our diffusion discriminative classifier aligns with the GAIL discriminator [21], learning a policy with the classifier enjoys the same theoretical guarantee, i . e ., optimizing this objective can bring a policy's occupancy measure closer to the expert's. Consequently, we can optimize our proposed diffusion discriminative classifier D ϕ with the loss function:

where L D sums the expert binary cross-entropy loss L expert BCE and the agent binary cross-entropy loss L agent BCE , and τ E and τ i represent a sampled expert trajectory and a collected agent trajectory by the policy π at training step i . We then update the diffusion discriminative classifier parameters ϕ based on the gradient of L D to improve its ability to distinguish expert data from agent data.

Intuitively, the discriminator D ϕ is trained to predict a value closer to 1 when the input state-action pairs are sampled from expert demonstration ( i . e ., trained to minimize L + diff and maximize L -diff ), and 0 if the input state-action pairs are obtained from the agent online interaction ( i . e ., trained to minimize L -diff and maximize L + diff ).

Note that our idea of transforming the diffusion model into a classifier is closely related to Li et al. [35], which shows that minimizing the diffusion loss is equivalent to maximizing the evidence lower bound (ELBO) of the log-likelihood [4], allowing for turning a conditional text-to-image diffusion model into an image classifier by using the ELBO as an approximate class-conditional log-likelihood log p ( x | c ) . By contrast, we employ a diffusion model for imitation learning. Moreover, we take a step further - instead of optimizing the diffusion loss L diff, we directly optimize the binary cross entropy losses calculated based on the denoising results to train the diffusion model as a binary classifier.

4.3 Diffusion-Reward Adversarial Imitation Learning

Our proposed method adheres to the fundamental AIL framework, where a discriminator and a policy are updated alternately. In the discriminator step, we update the diffusion discriminative classifier with the gradient of L D following Eq. 5. In the policy step, we adopt the adversarial inverse reinforcement learning objective proposed by Fu et al. [12] as our diffusion reward signal to train the policy:

The policy parameters θ can be updated using any RL algorithm to maximize the diffusion rewards provided by the diffusion discriminative classifier, bringing the policy closer to the expert policy. In our implementation, we utilize PPO as our policy update algorithm. The algorithm is presented in Algorithm 1, and the overall framework is illustrated in Figure 2.

Among the related works, DiffAIL [59] is the closest to ours, as it also uses a diffusion model for adversarial

Algorithm 1 Diffusion-Reward Adversarial Imitation Learning (DRAIL)

1: Input: Expert trajectories τ E , initial policy parameters θ 0 , initial diffusion discriminator parameters ϕ 0 , and discriminator learning rate η ϕ

2: for i = 0 , 1 , 2 , . . . do

3: Sample agent transitions τ i ∼ π θ i

4: Compute the output of diffusion discriminative classifier D ϕ (Eq. 4) and the loss function L D (Eq. 5)

5: Update the diffusion model ϕ i +1 ← ϕ i -η ϕ ∇L D

6: Compute the diffusion reward r ϕ ( s , a ) with Eq. 6

7: Update the policy θ i +1 ← θ i with any RL algorithm w.r.t. reward r ϕ

8: end for

imitation learning. DiffAIL employs an unconditional diffusion model to denoise state-action pairs from both experts and agents. However, this approach only implicitly reflects the likelihood of state-action pairs belonging to the expert class through diffusion loss, making it challenging to explicitly distinguish between expert and agent behaviors.

In contrast, our method, DRAIL, uses a conditional diffusion model that directly conditions real ( c + ) and fake ( c -) labels. This allows our model to explicitly calculate and compare the probabilities of state-action pairs belonging to either the expert or agent class. This clearer and more robust signal for binary classification aligns more closely with the objectives of the GAIL framework, leading to more stable and effective learning. For further details and the mathematical formulation, please refer to Section A.

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 Section 5.4 and Section 5.5. 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 our approach in the point mass MAZE navigation environment, introduced in Fu et al. [13] (maze2d-medium-v2), as depicted in Figure 3a. In this task, a point-mass agent is trained to navigate from a randomly determined start location to the goal. The agent accomplishes the task by iteratively predicting its acceleration in the vertical and horizontal directions. We use the expert dataset provided by Lee et al. [33], which includes 100 demonstrations, comprising 18,525 transitions.

FETCHPUSH. We evaluate our approach in a 7-DoF Fetch task, FETCHPUSH, depicted in Figure 3b, where the Fetch is required to push a black block to a designated location marked by a red sphere. We use the demonstrations from Lee et al. [33], consisting of 20,311 transitions (664 trajectories).

HANDROTATE. We further evaluate our approach in a challenging environment HANDROTATE with a high-dimensional continuous action space introduced by Plappert et al. [43]. Here, a 24-DoF Shadow Dexterous Hand is tasked with learning to in-hand rotate a block to a target orientation, as depicted in Figure 3c. This environment features a high-dimensional state space (68D) and action space (20D). We use the demonstrations collected by Lee et al. [33], which contain 515 trajectories (10k transitions).

(a) MAZE

(b) FETCHPUSH

(c) HANDROTATE

1

0001

ANTREACH. The goal of ANTREACH, a location and navigation task, is for a quadruped ant to reach a goal randomly positioned along the perimeter of a half-circle with a radius of 5 meters, as depicted in Figure 3d. The 132D high-dimensional continuous state space encodes joint angles, velocities, contact forces, and the goal position relative to the agent. We use the demonstrations provided by Lee et al. [33], which contain 1,000 demonstrations (25k transitions).

WALKER. The objective of WALKER is to let a bipedal agent move at the highest speed possible while preserving its balance, as illustrated in Figure 3e. We trained a PPO expert policy with environment rewards and collected 5 successful trajectories, each containing 1000 transitions, as an expert dataset.

CARRACING. We evaluate our method in a racing game, CARRACING, illustrated in Figure 3f, requiring driving a car to navigate a track. This task features a 96 × 96 RGB image-based state space and a 3-dimensional action space (steering, braking, and accelerating). We trained a PPO expert policy on CARRACING environment and collected 671 transitions as expert demonstrations.

Further details of the tasks can be found in Section B.

5.2 Baselines

We compare our method DRAIL with the following baselines of our approach.

Behavioral Cloning (BC) trains a policy to mimic the actions of an expert by supervised learning a mapping from observed states to corresponding expert actions [44, 55].

Diffusion Policy represents a policy as a conditional diffusion model [6, 42, 46], which predicts an action conditioning on a state and a randomly sampled noise. We include this method to compare learning a diffusion model as a policy (diffusion policy) or reward function (ours).

Generative Adversarial Imitation Learning (GAIL) [21] learns a policy from expert demonstrations by training a discriminator to distinguish between trajectories generated by the learned generator policy and those from expert demonstrations.

Generative Adversarial Imitation Learning with Gradient Penalty (GAIL-GP) is an extension of GAIL that introduces a gradient penalty to achieve smoother rewards and stabilize the discriminator.

Wasserstein Adversarial Imitation Learning (WAIL) [2] extends GAIL by employing Wasserstein distance, aiming to capture smoother reward functions.

Diffusion Adversarial Imitation Learning (DiffAIL) [58] integrates a diffusion model into AIL by using the diffusion model loss to provide reward e -L diff .

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 4. Each task

was trained using five different random seeds. Note that BC and Diffusion Policy are offline imitation learning algorithms, meaning they cannot interact with the environment, so their performances are represented as horizontal lines. Detailed information on model architectures, training, and evaluation can be found in Section F and Section G.

Overall, our method DRAIL consistently outperforms prior methods or achieves competitive performance compared to the best-performing methods across all the environments, verifying the effectiveness of integrating our proposed diffusion discriminative classifier into the AIL framework.

DRAIL vs. DiffAIL. Both DRAIL and DiffAIL integrate the diffusion model into the AIL framework. In 5 out of 6 tasks, our DRAIL outperforms DiffAIL, demonstrating that our proposed discriminator provides a more effective learning signal by closely resembling binary classification within the GAIL framework.

DRAIL vs. BC. AIL methods generally surpass BC in most tasks due to their ability to learn from interactions with the environment and thus handle unseen states better. However, BC outperforms all other baselines in the locomotion task (WALKER). We hypothesize that WALKER is a monotonic task requiring less generalizability to unseen states, allowing BC to excel with sufficient expert data. Additionally, our experiments with varying amounts of expert data, detailed in Section 5.5, suggest that DRAIL surpasses BC when less expert data is available.

We empirically found that our proposed DRAIL is robust to hyperparameters, especially compared to GAIL and WAIL, as shown in the hyperparameter sensitivity experiment inSection D.

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. [33]. Specifically, we evaluate policies learned by different methods by varying the noise injected into initiate states ( e . g ., position and velocity of the robot arm) and goals ( e . g ., target block positions in FETCHPUSH). We experiment with different noise levels, including 1 × , 1 . 25 × , 1 . 5 × , 1 . 75 × , and 2 . 0 × , compared to the expert environment. That said, 1 . 5 × means the policy is evaluated in an environment with noises 1 . 5 × larger than those injected into expert data collection. Performing well in a high noise level setup requires the policy to generalize to unseen states.

Expert

DRAIL (Ours)

DiffAIL

WAIL

Diffusion Policy

BC

GAIL-GP

(a) MAZE

(b) FETCHPUSH

(c) HANDROTATE

(d) ANTREACH

(e) WALKER

(f) CARRACING

1.0

1.0

1.0

Success Rate

0.8

Success Rate

0.8

Success Rate

0.8

0.6

0.6

0.6

0.4

0.4

0.2

0.2

0.2

0.0

0

10

20

0.0

0

2

4

0.0

0

2.

4

Step (M)

Step (M)

Step (M)

1.0

1UUU

800

Success Rate

0.8

4000

0.6

Return

Return

009

0.4

400

2000

200

0.2

0.80

0

0

2.5

5.0

7.5

10.0

0

10

20

0.0

0.5

1.0

1.5

Step (M)

Step (M)

Step (k)

Expert

DRAIL (Ours)

DiffAIL

WAIL

Diffusion Policy

BC

GAIL-GP

(a) FETCHPUSH

1

×

(b) FETCHPUSH

1

.

25

×

(c) FETCHPUSH

1

.

75

×

(d) FETCHPUSH

2

.

0

×

1.0

1.0

1.0

1.0

Rate

0.8

Rate

0.8

Rate

0.8

Rate

0.8

Success

0.6

Success

0.6

Success

0.6

Success

0.6

0.4

0.4

0.4

0.4

0.2

0.2

0.2

0.2

0.0

0

Step (M)

2

4

0.0

0

Step (M)

2

4

0.0

Q

Step (M)

2

4

0.0

0

Step (M)

2

4

The results of FETCHPUSH under 1 × , 1 . 25 × , 1 . 75 × , and 2 . 0 × 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 . 00 × , DRAIL maintains a success rate of over 95% , surpassing the best-performing baseline, Diffusion Policy, which achieves only around a 79 . 20% success rate. In contrast, DiffAIL experiences failures in 2 out of the 5 seeds, resulting in a high standard deviation (mean: 40 . 90 , standard deviation: 47 . 59 ), despite our extensive efforts on experimenting with various configurations and a wide range of hyperparameter values.

The extended generalization experiments results in MAZE, FETCHPUSH, HANDROTATE , and the new task FETCHPICK are presented in Section C. Overall, our method outperforms or performs competitively against the best-performing baseline, demonstrating its superior generalization ability.

Expert

DRAIL (Ours)

DiffAIL

WAIL

Diffusion Policy

BC

GAIL-GP

(a)

5

trajectories

(b)

3

trajectories

(c)

2

trajectories

(d)

1

trajectory

(e)

20311

transitions

(f)

10000

transitions

(g)

5000

transitions

(h)

2000

transitions

6000

6000

6000

6000

Return

4000

Return

4000

Return

4000

Return

4000

2000

2000

2000

2000

0

0

Step (M)

10

20

0

0

Step (M)

10

20

10

0

Step (M)

10

20

0

0

Step (M)

10

20

FetchPush

FetchPush

FetchPush

FetchPush

1.0

1.0

1.0

1.0

Success Rate

Success Rate

0.8

Success Rate

0.8

Success Rate

0.8

0.6

0.6

0.4

0.4

0.2

0.2

0.2

0.2

0.0

0

Step (M)

2

4

0.0

0

Step (M)

2

4

0.0

0

Step (M)

2

4

0.0

0

Step (M)

2

4

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, 2, and 1 expert trajectories, each containing 1000 transitions; for FETCHPUSH, we use 20311, 10000, 5000, and 2000 state-action pairs. The results reported in Figure 6 demonstrate that our DRAIL learns faster compared to the other baselines, indicating superior data efficiency in terms of environment interaction. In WALKER, our DRAIL maintains a return value of over 4500 even when trained with a single trajectory. In contrast, BC's performance is unstable and fluctuating, while the other baselines experience a dramatic drop. In FETCHPUSH, our DRAIL maintains a success rate of over 80% even when the data size is reduced by 90% , whereas the other AIL baselines' performance drops below 50% .

(a) Expert Distribution

(b) GAIL Reward

(c) DRAIL Reward

Sine Expert Distribution

2.0

1.5

:expert distribution

: sampled (s, a)

0'2

1.5

2.5

0°2

1.5

0.7

0.8

1.0

1.0

0.0

1.0

0.6

Action

0.5

0.5

2.5

0.5

0.0

0.0

5.0

00

0.4

0.5

3.5

7.5

3.5

1.0

1.0

10.0

1.0

0.2

1.5

1.5

-12.5

1.5

0.1

2.0

0

4

State

6

8

10

2.0

0

10

2.0

0

B

10

5.6 Reward function visualization

To visualize and analyze the learned reward functions, we design a SINE environment with onedimensional state and action spaces, where the expert state-action pairs form a discontinuous sine wave a = sin (20 s π ) + N (0 , 0 . 05 2 ) , as shown in Figure 7a. We train GAIL and our DRAIL to learn from this expert state-action distribution and visualize the discriminator output values D ϕ to examine the learned reward function, as presented in Figure 7.

Figure 7b reveals that the GAIL discriminator exhibits excessive overfitting to the expert demonstration, resulting in its failure to provide appropriate reward values when encountering unseen states. In contrast, Figure 7c shows that our proposed DRAIL generalizes better to the broader state-action distribution, yielding a more robust reward value, thereby enhancing the generalizability of learned policies. Furthermore, the predicted reward value of DRAIL gradually decreases as the state-action pairs deviate farther from the expert demonstration. This reward smoothness can guide the policy even when it deviates from the expert policy. In contrast, the reward distribution from GAIL is relatively narrow outside the expert demonstration, making it challenging to properly guide the policy if the predicted action does not align with the expert.

6 Conclusion

This work proposes a novel adversarial imitation learning framework that integrates a diffusion model into generative adversarial imitation learning (GAIL). Specifically, we propose a diffusion discriminative classifier that employs a diffusion model to construct an enhanced discriminator, yielding more robust and smoother rewards. Then, we design diffusion rewards based on the classifier's output for policy learning. Extensive experiments in navigation, manipulation, locomotion, and game justify our proposed framework'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 as the Wasserstein distance or f-divergences, could potentially further improve training stability.

References

Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In International Conference on Machine Learning , 2004.

Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning , 2017.

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.

David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association , 2017.

Shang-Fu Chen, Hsiang-Chun Wang, Ming-Hao Hsu, Chun-Mao Lai, and Shao-Hua Sun. Diffusion model-augmented behavioral cloning. In International Conference on Machine Learning , 2024.

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.

Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. In Neural Information Processing Systems , 2017.

Robert Dadashi, Léonard Hussenot, Matthieu Geist, and Olivier Pietquin. Primal wasserstein imitation learning. arXiv preprint arXiv:2006.04678 , 2020.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. In Neural Information Processing Systems , 2021.

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.

Gideon Joseph Freund, Elad Sarafian, and Sarit Kraus. A coupled flow approach to imitation learning. In International Conference on Machine Learning , pages 10357-10372. PMLR, 2023.

Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adverserial inverse reinforcement learning. In International Conference on Learning Representations , 2018.

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 , 2020.

Javier Garcıa and Fernando Fernández. A comprehensive survey on safe reinforcement learning. Journal of Machine Learning Research , 2015.

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.

Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Neural Information Processing Systems , 2014.

Shangding Gu, Long Yang, Yali Du, Guang Chen, Florian Walter, Jun Wang, Yaodong Yang, and Alois Knoll. A review of safe reinforcement learning: Methods, theory and applications. arXiv preprint arXiv:2205.10330 , 2022.

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 on Computational Intelligence and Games , 2018.

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.

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.

Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Neural Information Processing Systems , 2016.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Neural Information Processing Systems , 2020.

Bo-Ruei Huang, Chun-Kai Yang, Chun-Mao Lai, , and Shao-Hua Sun. Diffusion imitation from observation. In Neural Information Processing Systems , 2024.

Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne. Imitation learning: A survey of learning methods. ACM Computing Surveys , 2017.

Rohit Jena, Changliu Liu, and Katia Sycara. Augmenting gail with bc for sample efficient imitation learning. In Conference on Robot Learning , 2021.

Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. In Neural Information Processing Systems , 2021.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations , 2015.

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.

Po-Chen Ko, Jiayuan Mao, Yilun Du, Shao-Hua Sun, and Joshua B. Tenenbaum. Learning to act from actionless videos through dense correspondences. In International Conference on Learning Representations , 2024.

Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tompson. Discriminator-actor-critic: Addressing sample inefficiency and reward bias in adversarial imitation learning. In International Conference on Machine Learning , 2019.

Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tompson. Discriminator-actor-critic: Addressing sample inefficiency and reward bias in adversarial imitation learning. In International Conference on Learning Representations , 2019.

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.

Youngwoon Lee, Andrew Szot, Shao-Hua Sun, and Joseph J. Lim. Generalizable imitation learning from observation via inferring goal proximity. In Neural Information Processing Systems , 2021.

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:1811.07871 , 2018.

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.

Abdoulaye O Ly and Moulay Akhloufi. Learning to drive by imitation: An overview of deep behavior cloning methods. IEEE Transactions on Intelligent Vehicles , 2020.

Andrew Y. Ng and Stuart J. Russell. Algorithms for inverse reinforcement learning. In International Conference on Machine Learning , 2000.

Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning , 2021.

Felipe Nuti, Tim Franzmeyer, and João F Henriques. Extracting reward functions from diffusion models. Advances in Neural Information Processing Systems , 36, 2024.

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.

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.

Tim Pearce, Tabish Rashid, Anssi Kanervisto, David Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, and Sam Devlin. Imitating human behaviour with diffusion models. In International Conference on Learning Representations , 2023.

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.09464 , 2018.

Dean A Pomerleau. Alvinn: An autonomous land vehicle in a neural network. In Neural Information Processing Systems , 1989.

Harish Ravichandar, Athanasios S Polydoros, Sonia Chernova, and Aude Billard. Recent advances in robot learning from demonstration. Annual review of control, robotics, and autonomous systems , 2020.

Moritz Reuss, Maximilian Li, Xiaogang Jia, and Rudolf Lioutikov. Goal-conditioned imitation learning using score-based diffusion policies. arXiv preprint arXiv:2304.02532 , 2023.

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-assisted intervention-MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18 , 2015.

Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In International Conference on Artificial Intelligence and Statistics , 2011.

Stefan Schaal. Learning from demonstration. In Neural Information Processing Systems , 1997.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , 2015.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Machine Learning , 2021.

Gokul Swamy, David Wu, Sanjiban Choudhury, Drew Bagnell, and Steven Wu. Inverse reinforcement learning without reinforcement learning. In International Conference on Machine Learning , 2023.

Umar Syed, Michael Bowling, and Robert E Schapire. Apprenticeship learning using linear programming. In Proceedings of the 25th international conference on Machine learning , 2008.

Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. In International Joint Conference on Artificial Intelligence , 2018.

Faraz Torabi, Garrett Warnell, and Peter Stone. Generative adversarial imitation from observation. In International Conference on Machine Learning , 2019.

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. https://github.com/ huggingface/diffusers , 2022.

Bingzheng Wang, Yan Zhang, Teng Pang, Guoqiang Wu, and Yilong Yin. Diffail: Diffusion adversarial imitation learning. arXiv preprint arXiv:2312.06348 , 2023.

Bingzheng Wang, Guoqiang Wu, Teng Pang, Yan Zhang, and Yilong Yin. Diffail: Diffusion adversarial imitation learning. In AAAI Conference on Artificial Intelligence , 2024.

Tony Z Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn. Learning fine-grained bimanual manipulation with low-cost hardware. In Robotics: Science and Systems , 2023.

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.

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.

Appendix

Table of Contents

A Relation to DiffAIL

Among all the related works, DiffAIL [58], which also employs a diffusion model for adversarial imitation learning, is the closest to our work. This section describes the differences of our work and DiffAIL.

Unconditional diffusion models (DiffAIL) vs. conditional diffusion models (DRAIL). Wang et al. [58] proposed to learn an unconditional diffusion model D ϕ ( s , a ) to denoise state-action pairs from an expert and an agent. The diffusion model should denoise well when the state-action pairs are sampled from the expert demonstrations, while denoise poorly given the state-action pairs sampled from the agent policy. On the other hand, our proposed framework employs a conditional diffusion model D ϕ ( s , a , c ) that conditions on the real label c + or the fake label c -. In the following, we discuss how using a conditional diffusion model leads to better learning signals for policy learning.

DiffAIL set up the discriminator output as D ϕ ( s , a ) = e -L diff ( s , a ) ∈ [0 , 1] , indicating how likely ( s , a ) is sampled from the expert distribution. That said, a e -L diff ( s , a ) close to 1 implies that the sample is likely from expert and a small value of e -L diff ( s , a ) means it is less likely that this sample

comes from the expert. However, e -L diff ( s , a ) does not explicitly represent the probability of the 'negative class' (agent), which makes it difficult to provide stable rewards for policy learning.

We start by revisiting the GAIL formulation. The binary classifier in GAIL, i . e ., the discriminator, is trained to minimize the binary cross-entropy loss, which leads to maximizing the likelihood of predicting the correct class probabilities. Subsequently, the discriminator outputs the predicted probability of the 'positive class' (expert) relatively to the 'negative class' (agent). In DiffAIL, the predicted probability of the 'positive class' (expert) is e -L diff and the predicted probability of the 'negative class' (agent) can be given by 1 -e -L diff . Following this definition, the discriminator thinks the sample if from an expert when e -L diff > 1 -e -L diff . We also know that L diff > 0 since it is the diffusion model denoising loss. This gives us 0 ≤ L diff ≤ ln 2 for the expert class, and ln 2 < L diff for the agent class. This can be problematic since the fixed boundary ln 2 my not reflect the real boundary between the real (expert) and the fake (agent) distributions.

On the other hand, in our method (DRAIL), the predicted probability of the positive (expert) class of a given state-action pair ( s , a ) is defined as:

and the predicted probability of the negative (agent) class of the same state-action pair ( s , a ) is explicitly defined as:

This binary classification formulation aligns with GAIL. Hence, the discriminator would think a given ( s , a ) comes from expert data only when

which means

The resulting relative boundary can provide better learning signals for policy learning, especially when the behaviors of the agent policy become similar to those of expert, which can be observed in the tasks where the agent policies can closely follow the experts, such as FETCHPICK, FETCHPUSH, and ANTREACH. Also, we hypothesize this leads to the superior performance of our method compared to DiffAIL in most of the generalization experiments.

Experimental setup. Wang et al. [58] evaluated DiffAIL and prior methods in locomotion tasks; in contrast, our work extensively compares our proposed framework DRAIL with various existing methods in various domains, including navigation (MAZE and ANTREACH), locomotion (WALKER and ANTREACH), robot arm manipulation (FETCHPUSH and FETCHPICK), robot arm dexterous (HANDROTATE), and games (CARRACING). Additionally, we present experimental results on generalization to unseen states and goals on the goal-oriented tasks, and on varying amounts of expert data.

B Environment & task details

B.1 MAZE

Description. In a 2D maze environment, a point-maze agent learns to navigate from a starting location to a goal location. The agent achieves this by iteratively predicting its x and y velocity. The initial and final positions of the agent are randomly selected. The state space includes position,

velocity, and goal position. The maximum episode length for this task is set at 400 , and the episode terminates if the goal is reached earlier.

Expert dataset. The expert dataset comprises 100 demonstrations, which includes 18 , 525 transitions provided by Lee et al. [33].

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 an object from a table and move it to a target location.

According to the environment setups stated in Lee et al. [33], the 16-dimensional state representation includes the angles of the robot joints, and the initial three dimensions of the action vector represent the intended relative position for the next time step. The first three dimensions of the action vector denote the intended relative position in the subsequent time step. In the case of FETCHPICK, an extra action dimension is incorporated to specify the distance between the two fingers of the gripper. The maximum episode length for this task is set at 60 for FETCHPUSH and 50 for FETCHPICK, and the episode terminates if the agent reaches the goal earlier.

Expert dataset. The expert dataset for FETCHPUSH comprises 664 trajectories, amounting to 20 , 311 transitions, and the expert dataset for FETCHPICK comprises 303 trajectories, amounting to 10 , 000 transitions provided by Lee et al. [33].

B.3 HANDROTATE

Description. In the task HANDROTATE proposed by Plappert et al. [43], a 24-DoF Shadow Dexterous Hand is designed to rotate a block in-hand to a specified target orientation. The 68D state representation includes the agent's joint angles, hand velocities, object poses, and target rotation. The 20D action vector represents the joint torque control of the 20 joints. Notably, HANDROTATE is challenging due to its high-dimensional state and action spaces. We follow the experimental setup outlined in Plappert et al. [43] and Lee et al. [33], where rotation is constrained to the z-axis, and allowable initial and target z rotations are within [ -π 12 , π 12 ] and [ π 3 , 2 π 3 ] , respectively. The maximum episode length for this task is set at 50 , and the episode terminates if the hand reaches the goal earlier.

Expert dataset. Weuse the demonstrations collected by Lee et al. [33], which contain 515 trajectories (10k transitions).

B.4 ANTREACH

Description. The ANTREACH task features a four-leg ant robot reaching a randomly assigned target position located within a range of half-circle with a radius of 5 meters. The task's state is represented by a 132-dimension vector, including joint angles, velocities, and the relative position of the ant towards the goal. Expert data collection for this task is devoid of any added noise, while random noise is introduced during the training and inference phases. Consequently, the policy is required to learn to generalize to states not present in the expert demonstrations. The maximum episode length for this task is set at 50 , and the episode terminates if the ant reaches the goal earlier.

Expert dataset. The expert dataset comprises 10000 state-action pairs provided by Lee et al. [33].

B.5 WALKER

Description. WALKER task involves guiding an agent to move towards the x-coordinate as fast as possible while maintaining balance. An episode terminates either when the agent experiences predefined unhealthy conditions in the environment or when the maximum episode length ( 1000 ) is reached. The agent's performance is evaluated over 100 episodes with three different random seeds. The return of an episode is the cumulative result of all time steps within that episode. The 17D state includes joint angles, angular velocities of joints, and velocities of the x and z-coordinates of the top. The 6D action defines the torques that need to be applied to each joint of the walker avatar.

Expert dataset. We trained a PPO expert policy with environment rewards and collected 5 successful trajectories, each containing 1000 transitions, as an expert dataset.

B.6 CARRACING

Description. In the CARRACING task, the agent must navigate a track by controlling a rear-wheel drive car. The state space of CARRACING is represented by a top-down 96 × 96 RGBimage capturing the track, the car, and various status indicators such as true speed, four ABS sensors, steering wheel position, and gyroscope readings. The agent controls the car using three continuous action values: steering, acceleration, and braking. Episodes have a maximum length of 1000 steps, and termination occurs if the car completes the track before reaching the maximum episode length.

In our experiment settings, we preprocess the state image by converting it to grayscale and resizing it to 64 × 64 pixels. We then concatenate two consecutive frames to form a single state data, providing temporal context for the agent's decision-making process.

Expert dataset. We trained a PPO expert policy on CARRACING environment and collected 671 transitions as expert demonstrations.

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. In the generalization experiment, we gather expert demonstrations from some lower and higher coverage: 25% , 75% , and 100% . (2) In FETCHPICK, FETCHPUSH, and HANDROTATE main experiments, the demonstrations are collected in a lower noise level setting, 1 × . Yet, the agent is trained within an environment incorporating 1 . 5 × noise, which is 1.5 times larger noise than the collected expert demonstration, applied to the starting and target block positions. In the generalization experiment, we train agents in different noise levels: 1 × , 1 . 25 × , 1 . 75 × , 2 . 0 × . (3) In ANTREACH main experiment, no random noise is added to the initial pose during policy learning. In ANTREACH generalization experiment, we train agents in different noise levels: 0 (default), 0 . 01 , 0 . 03 , 0 . 05 .

These generalization experiments simulate real-world conditions. For example, because of the expenses of demonstration collection, the demonstrations may inadequately cover the entire state space, as seen in setup (1). Similarly, in setups (2) and (3), demonstrations may be acquired under controlled conditions with minimal noise, whereas the agent operating in a real environment would face more significant noise variations not reflected in the demonstrations, resulting in a broader distribution of initial states.

C.2 Experiment results

MAZE. Our DRAIL outperforms baselines or performs competitively against DiffAIL across all demonstration coverages, as shown in Figure 8. Particularly, BC, WAIL, and GAIL's performance decline rapidly in the low coverage case. In contrast, diffusion model-based AIL algorithms demonstrate sustained performance, as shown in Figure 4a. This suggests that our method exhibits robust generalization, whereas BC and GAIL struggle with unseen states under limited demonstration coverage.

FETCHPICK and FETCHPUSH. In FETCHPICK, our method outperforms all baselines in most noise levels, as shown in Figure 8. In the 2 . 00 × noise level, our method DRAIL achieves a success rate of 87 . 22 , surpassing the best-performing baseline Diffusion Policy, which achieves only around 76 . 64% . GAIL, on the other hand, experiences failure in 3 out of the 5 seeds, resulting in a high standard deviation (mean: 39 . 17 , standard deviation: 47 . 98 ) despite our thorough exploration of various settings for its configuration. In FETCHPUSH, our method DRAIL exhibits more robust results, in generalizing to unseen states compared to other baselines. This showcases that the diffusion

Expert

DRAIL (Ours)

DiffAIL

WAIL

Diffusion Policy

BC

GAIL-GP

(a) MAZE

100%

(b) MAZE

75%

(c) MAZE

50%

(d) MAZE

25%

(e) FETCHPICK

1

.

00

×

(f) FETCHPICK

1

.

25

×

(g) FETCHPICK

1

.

75

×

(h) FETCHPICK

2

.

0

×

(i) FETCHPUSH

1

×

(j) FETCHPUSH

1

.

25

×

(k) FETCHPUSH

1

.

75

×

(l) FETCHPUSH

2

.

0

×

1.0

1.0

1.0

1.0

Rate

0.8

Success Rate

0.8

Rate

0.8

Success Rate

0.8

Success

0.6

0.6

Success

0.6

0.6

0.4

0.4

0.4

0.4

0.2

0.2

0.2

0.2

0.0

0

10

20

0.0

0

10

20

0.0

Q

10

20

0.0

10

20

Step (M)

Step (M)

Step (M)

Step (M)

1.0

1.0

1.0

1.0

Success Rate

0.8

Success Rate

0.8

Success Rate

0.8

Success Rate

0.8

0.6

0.6

0.6

0.6

0.4

0.4

0.4

0.4

0.2

0.2

0.2

0.2

00

0

10

20

0.0

0

10

20

0.0

0

10

20

0.0

0

10

20

Step (M)

Step (M)

Step (M)

Step (M)

1.0

1.0

1.0

1.0

SuccessRate

0.8

SuccessRate

0.8

SuccessRate

0.8

Success Rate

0.8

.6

0.6

0.6

30.4

0.4

FO.41

0.2

0.2

0.2

0.2

0.0

0

Step (M)

2

4

0.0

0

Step (M)

2

4

0.0

0

Step (M)

2

4

0.0

0

Step (M)

2

1.0

1.0

1.0

1.0

Rate

0.8

Rate

0.8

Rate

0.8

Rate

0.8

Success

0.6

Success

0.6

Success

0.6

Success

0.6

0.4

0.4

0.4

0.4

0.2

0.2

0.2

0.2

0.0

0

Step (M)

2

0.0

0

Step (M)

2

0.0

Step (M)

2

0.0

0

Step (M)

2

reward guidance could provide better generalizability for the AIL framework. Moreover, our DRAIL is quite sample-efficient regarding interaction with the environment during training compared to other baselines in FETCHPICK and FETCHPUSH environment.

HANDROTATE. DiffAIL and Our DRAIL show robustness to different levels of noise in HANDROTATE, as illustrated in Figure 8. Specifically, DiffAIL and our DRAIL achieve a success rate of higher than 90% at a noise level of 2 . 0 , while GAIL and WAIL only reach approximately 42 . 82% and 8 . 80% , respectively.

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.

Maze

Maze

1.0

1.0

Rate

0.8

Rate

0.8

Success

0.6

0.5x

Success

0.6

0.5x

0.2x

0.2x

0.4

0.4

5x

5x

0.2

2x

0.2

2x

1x (Ours)

1x (Ours)

0.0

0.0

0

10

20

0

10

20

Step (M)

Step (M)

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 of the converged performance across all experiments, including the main results in Section 5.3, the generalization experiments in Appendix C.2, and the data efficiency experiments in Section 5.5. The results are presented in Table 1.

F Model architecture

This section presents the model architecture of all the experiments. Appendix F.1 describe the model architecture of all methods used in Section 5.3.

F.1 Model architecture of DRAIL, DiffAIL, 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-GP, WAIL, and DiffAIL) across six diverse environments. We applied Multilayer Perceptron (MLP) for the policy of BC, the conditional diffusion model in Diffusion Policy, as well as the policy and the discriminator of GAIL, GAIL-GP, and WAIL. For DiffAIL and our proposed DRAIL, MLPs were employed in the policy and diffusion model of the diffusion discriminative classifier. The activation functions used for the MLPs in the diffusion model were ReLU, while hyperbolic tangent was employed for the others. The total timestep T for all diffusion models in this paper is set to 1000 and the scheduler used for diffusion models is cosine scheduler [38]. Further details regarding the parameters for the model architecture can be found in Table 2.

BC. We maintained a concise model for the policy of BC to prevent excessive overfitting to expert demonstrations. This precaution is taken to mitigate the potential adverse effects on performance when confronted with environments exhibiting higher levels of noise.

Diffusion Policy. Based on empirical results and Chen et al. [5], the Diffusion Policy performs better when implemented with a deeper architecture. Consequently, we have chosen to set the policy's number of layers to 5.

GAIL & GAIL-GP. The detailed model architecture for GAIL and GAIL-GP is presented in Table 2. For GAIL-GP, the gradient penalty is set to 1 across all environments.

WAIL. We set the ground transport cost and the type of regularization of WAIL as Euclidean distance and L 2 -regularization. The regularization value ϵ is provided in Table 2.

DiffAIL. In DiffAIL, the conditional diffusion model is not utilized as it only needs to consider the numerator of Equation (4). Consequently, the diffusion model takes only the noisy state-action pairs as input and outputs the predicted noise value.

DRAIL. The conditional diffusion model of the diffusion discriminative classifier in our DRAIL is constructed by concatenating either the real label c + or the fake label c -to the noisy state-action pairs as the input. The model then outputs the predicted noise applied to the state-action pairs. The dimensions of both c + and c -are reported in Table 2.

F.2 Image-based model architecture of DRAIL, DiffAIL, 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 a multi-layer perceptron (MLP) for action prediction. The CNN consists of three downsampling blocks with 32, 64, and 64 channels respectively. The kernel sizes for these blocks are 8, 4, and 3, with strides of 4, 2, and 1, respectively. After feature extraction, the output is flattened and passed through a linear layer to form a 512-dimensional feature vector representing the state data. This state feature vector is subsequently processed by an MLP with three layers, each

having a 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 a policy as a conditional diffusion model, which predicts an action conditioning on a state and a randomly sampled noise. Our condition diffusion model is implemented using the diffusers package by von Platen et al. [57]. The state in CARRACING image of size 64 × 64 , so we first use a convolutional neural network (CNN) to extract the feature. The CNN is based on a U-Net [47] structure, comprising 3 down-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 blocks are 4, 8, and 16, respectively.

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 a given state-action pair being from the expert demonstrations.

Discriminator of DiffAIL & DRAIL. Our diffusion model is implemented using the diffusers package by von Platen et al. [57]. The architecture of both DiffAIL and DRAIL is based on a U-Net

[47] structure, comprising 3 down-sampling blocks and 3 up-sampling blocks. Each block consists of 2 ResNet [19] layers, with group normalization applied using 4 groups. The channel sizes for each pair of down-sampling and up-sampling blocks are 4, 8, and 16, respectively. The condition label is incorporated through class embedding, with the number of classes set to 2, representing the real label c + and the fake label c -. Finally, we apply running normalization at the output to ensure stable training and accurate discrimination.

To accommodate both image-based state and vector-based action data within the diffusion model, we flatten the action data into an image with a channel size equivalent to the action dimension. Subsequently, we concatenate the state and transformed action data as input to the U-Net. In DRAIL, we use 1 to represent c + and 0 to represent c -as condition data. For DiffAIL, since condition labels are not required, we simply assign a constant value of 0 as the condition label.

G Training details

G.1 Training hyperparamters

The hyperparameters employed for all methods across various tasks are outlined in Table 3. The Adam optimizer [27] is utilized for all methods, with the exception of the discriminator in WAIL, for which RMSProp is employed. Linear learning rate decay is applied to all policy models.

Due to the potential impact of changing noise levels on the quality of agent data input for the discriminator, the delicate balance between the discriminator and the AIL method's policy may be disrupted. Therefore, we slightly adjusted the learning rate for the policy and the discriminator for different noise levels on each task. The reported parameters in Table 3 correspond to the noise levels presented in Figure 4.

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 DiffAIL and GAIL. In CARRACING, we observe that adapting GAIL's optimization objective could lead to better performance; hence, we use it for DRAIL, DiffAIL, and

GAIL. For WAIL, we adhere to the approach outlined in the original paper, wherein the output of the discriminator directly serves as the reward function.

In our experiments, we employ Proximal Policy Optimization (PPO) [50], a widely used policy optimization method, to optimize policies for all the AIL methods. We maintain all hyperparameters of PPO constant across methods for a given task, except the learning rate, which is adjusted for each method. The PPO hyperparameters for each task are presented in Table 4.

H Limitations

This work presents an adversarial imitation learning framework DRAIL by employing a diffusion model as a discriminator. While DRAIL achieves encouraging results in various domains, including robot arm manipulation, robot hand dexterous manipulation, locomotion, and games, our proposed framework is fundamentally limited to the learning from demonstration (LfD) setting. That said, DRAIL requires both state and action sequences and, therefore, cannot learn from videos or state-only sequences, i . e ., learning from observation (LfO). Moreover, DRAIL assumes expert demonstrations to be optimal, and its performance may not be satisfactory if the demonstrations contain a certain level of noise or the demonstrators are suboptimal. Finally, DRAIL with its imitation learning nature, is not designed to learn from environmental rewards; therefore, even when environments can provide rewards, there is no apparent mechanism to utilize them with the current formulation of DRAIL.

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

-CPU: an Intel Xeon W-2255 (10C/20T, 19.25M, 4.5GHz) 48-Lane CPU

-GPUs: two NVIDIA RTX 3080 Ti GPUs

-Memory: 128GB memory

I.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 Ti 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-Lane CPU, 64GB memory, an NVIDIA RTX 3080 Ti GPU, and an 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

J Impact statements

In this work, we propose a novel adversarial imitation learning framework, diffusion rewards guided adversarial imitation learning (DRAIL), which integrates a diffusion model into GAIL. The proposed framework can potentially reinforce the biases captured by expert demonstrations, which can lead to sub-optimal, unsafe, or even discriminatory behaviors. To address this issue, we encourage future works to focus on alleviating these issues in imitation learning, e . g ., fairness in machine learning, and responsible AI.

NeurIPS paper 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-Reward Adversarial Imitation Learning (DRAIL). We address the limitations of existing methods, particularly the brittleness and instability of generative adversarial imitation learning (GAIL), by integrating diffusion models into the learning process. Our proposedDRAIL framework aims to provide more robust and smoother rewards for policy learning while enhancing stability in adversarial training. The abstract and introduction clearly articulate these contributions, aligning with the experimental results and analyses presented in the paper.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

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 Appendix H. These include its restriction to learning from demonstration (LfD) settings, the requirement for state and action sequences, the assumption of optimal expert demonstrations, and its inability to leverage environmental rewards.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a 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 important 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], so 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 crossreferenced.

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 to provide a 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.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We described every step of our method in Section 4.2, and provided the model architecture and details of how our experiment employed the conditional diffusion model in Appendix Appendix F. 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 be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a 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 a 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 to 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 to replicate the results, access to a hosted model (e.g., in the case of a 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 all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: We plan to release the codes, models, and expert datasets upon acceptance.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on 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. If only a 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) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We have provided the experimental details including task details in Section B, model architecture and the policy update algorithm in Section F, and training details and their corresponding parameters in Section G.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

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 Table 1. 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 error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run 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.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in 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 computer resources (type of compute workers, memory, time of execution) needed to 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 configurations, in Section I.1, and we provide the approximate total GPU hours in Section I.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 (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Our research conducted in this paper aligns with the principles outlined in the NeurIPS 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 societal well-being.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We have considered the potential implications of our work on biases and discriminatory behaviors, as highlighted in our impact statement (Appendix J).

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper does not pose risks associated with the release of high-risk data or models. We use a self-trained gymnasium dataset, which does not involve pre-trained language models, image generators, or scraped datasets that could be misused. As such, specific safeguards for data or model release are not necessary in this context.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We gave credit to the diffusion model used in our image-based experiments and the expert datasets for some tasks and cited these papers in Appendix F.2 and Section 5.1. Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: This paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an 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 to 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 crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS 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 an equivalent approval/review based on the requirements of your country or 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 to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n## 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).\n\nQUESTIONS:\nSee Weaknesses.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThe authors have discussed the limitations of this paper. DRAIL is unable to learn from state-only trajectories and suboptimal trajectories.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. **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###\n\nWEAKNESSES:\n1. **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###\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nShown in weakness", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nImitation 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\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- 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.\n\nQUESTIONS:\n- 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\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-22"}