# DiffAIL: Diffusion Adversarial Imitation Learning

Bingzheng Wang1, Guoqiang Wu1*, Teng Pang1, Yan Zhang1, Yilong Yin1*

###### Abstract

Imitation learning aims to solve the problem of defining reward functions in
real-world decision-making tasks. The current popular approach is the
Adversarial Imitation Learning (AIL) framework, which matches expert state-
action occupancy measures to obtain a surrogate reward for forward
reinforcement learning. However, the traditional discriminator is a simple
binary classifier and doesn’t learn an accurate distribution, which may result
in failing to identify expert-level state-action pairs induced by the policy
interacting with the environment. To address this issue, we propose a method
named diffusion adversarial imitation learning (DiffAIL), which introduces the
diffusion model into the AIL framework. Specifically, DiffAIL models the
state-action pairs as unconditional diffusion models and uses diffusion loss
as part of the discriminator’s learning objective, which enables the
discriminator to capture better expert demonstrations and improve
generalization. Experimentally, the results show that our method achieves
state-of-the-art performance and significantly surpasses expert demonstration
on two benchmark tasks, including the standard state-action setting and state-
only settings. Our code can be available at the link https://github.com/ML-
Group-SDU/DiffAIL.

## Introduction

Deep Reinforcement Learning has achieved significant success in many decision-
making tasks, including AlphaGo (Silver et al. 2016), Atari games (Mnih et al.
2013), MuJoCo environment tasks (Todorov, Erez, and Tassa 2012), and robot
control tasks (Mnih et al. 2015), where these tasks are typically defined with
clear reward functions to guide the agent for decision-making. However, in
real-world scenarios, the reward function is regularly challenging to obtain
or define, such as in the case of autonomous driving, where it is complicated
to delineate what behavior is beneficial for the agent. The emergence of
_Imitation Learning_ provides a practical solution to the problem of
inaccessible rewards. In imitation learning, the agent does not rely on actual
rewards but instead utilizes expert demonstration data to learn a similar
expert policy.

A relatively old method in imitation learning is clone learning (Pomerleau
1991), which uses supervised learning to learn from expert data. Although such
methods are straightforward to implement, they are prone to serious
extrapolation errors when visiting out-of-distribution data while practically
interacting with environments. To alleviate the aforementioned errors, Dagger
(Ross, Gordon, and Bagnell 2011) proposes access to expert policies, where the
agent continuously interacts with the environment online and asks for expert-
level actions from the expert policy to expand the dataset. AdapMen (Liu et
al. 2023) proposes an active imitation learning framework based on teacher-
student interaction, and theoretical analysis shows that it can avoid
compounding errors under mild conditions.

To effectively tackle the aforementioned extrapolation error, Adversarial
imitation learning (AIL) (Ho and Ermon 2016; Fu, Luo, and Levine 2017;
Kostrikov et al. 2018; Ghasemipour, Zemel, and Gu 2020; Zhang et al. 2020;
Garg et al. 2021; Zhang et al. 2022) has become the most popular approach in
imitation learning. Rather than minimizing the divergence between the expert
policy and the agent policy, AIL employs online interactive learning to focus
on minimizing the divergence between the joint state-action distribution
induced by the expert policy and the learned policy. However, typical AIL
methods often use a simplistic discriminator that does not learn a
distribution, which may not accurately classify specific expert-level actions
generated by the policy during interactions with the environment. This can
result in the naive discriminator failing to distinguish expert-level
behaviors generated by the policy, potentially hindering agent learning.
Therefore, a powerful discriminator that can accurately capture the
distribution of expert data is crucial for AIL.

Currently, diffusion model (Ho, Jain, and Abbeel 2020) processes a powerful
distribution matching capability and has achieved remarkable success in image
generation, surpassing other popular generative models such as GAN (Goodfellow
et al. 2014) and VAE (Kingma and Welling 2013) in producing higher-quality and
more diverse samples.

Based on diffusion model’s powerful ability to capture data distribution, we
propose a new method called Diffusion Adversarial Imitation Learning (DiffAIL)
in this work. DiffAIL employs the same framework as traditional AIL but
incorporates the diffusion model to model joint state-action distribution.
Therefore, DiffAIL enables the discriminator to accurately capture expert
demonstration and improve generalization, facilitating the successful
identification of expert-level behaviors generated by the policy during
interactions with the environment. Furthermore, we conducted experiments on
representatively continuous control tasks in Mujoco, including the standard
state-action setting and state-only setting. Surprisingly in both settings, we
share the hyperparameters instead of readjusting them. The experiment results
show that our method can achieve state-of-the-art (SOTA) performance and
significantly surpass expert demonstration on these two benchmark tasks.

Overall, our contributions can be summarized as follows:

  * •

We propose the DiffAIL method, which combines the diffusion model into AIL to
improve the discriminator ability of distribution capturing.

  * •

The experimental results show that DiffAIL can achieve SOTA performance,
including standard state-action and state-only settings.

## Preliminaries

### Problem Setting

A sequential decision-making problem in reinforcement learning can be modeled
as a Markov Decision Process (MDP) defined as a tuple
$M=(S,A,\rho_{0},P,r,\gamma)$, where $S$ denotes the state space, $A$ denotes
the action space, $\rho_{0}$ denotes the initial state distribution,
$P:S\times A\times S\rightarrow[0,1]$ describes the environment’s dynamic
model by specifying the state transition function, $r:S\times
A\rightarrow\mathbb{R}$ denotes the reward function, and ${\gamma\in[0,1]}$
denotes the discount factor used to weigh the importance of future rewards. We
represent the agent policy as $\pi:S\rightarrow A$. Policy $\pi(a|s)$
interacts with the environment to generate transitions
$(s_{t},a_{t},r_{t},s_{t+1})$,where $t$ denotes the timestep,
$s_{0}\sim\rho_{0},a_{t}\sim\pi(\cdot|s_{t}),r_{t}\sim
r(s_{t},a_{t}),s_{t+1}\sim p(\cdot|s_{t},a_{t})$. The goal of reinforcement
learning is to learn a policy $\pi(a|s)$ maximizing cumulative discount
rewards.

In imitation learning, the environment reward signal is not accessible.
Rather, we can access expert demonstrations
$D=\\{(s_{t},a_{t})\\}_{t=1}^{k}\sim\pi_{e}$ given by the unknown expert
policy $\pi_{e}$. Imitation learning aims to learn a policy $\pi$ that can
recover $\pi_{e}$ based on expert demonstrations without relying on reward
signals. This setting can be purely offline to learn with only expert
demonstrations or additionally online to interact with the environment (reward
unavailable) by behavior policy

### Adversarial Imitation Learning

For AIL, it aims to find an optimal policy that minimizes the state-action
distribution divergence induced by the agent policy and the expert policy, as
measured by the occupancy measure:

$\pi^{*}=\arg\min_{\pi}{D_{f}(d_{\pi}(s,a)||d_{\pi_{e}}(s,a))},$ (1)

where
$d_{\pi}(s,a)=(1-\gamma)\sum\limits_{t=0}^{\infty}\gamma^{t}p(s_{t}=s,a_{t}=a)$
denotes state-action distribution of $\pi_{\theta}$, $(1-\gamma)$ term denotes
normalization factor. Similarly, $d_{\pi_{e}}(s,a)$ denotes expert
demonstration distribution with a similar form. $D_{f}$ can be an arbitrary
distance formulation with the corresponding min-max target through the
corresponding dual representation, which is also applied in f-gan (Nowozin,
Cseke, and Tomioka 2016). Gail (Ho and Ermon 2016) adopts the Jensen-Shannon
(JS) divergence as the chosen distance metric $D_{f}$. By leveraging the dual
representation of JS divergence, we can obtain the minimax optimization
objective of Gail:

$\displaystyle\min_{\pi_{\theta}}\max_{D_{\phi}}\
\mathbb{E}_{(s,a)\sim\pi_{e}}$ $\displaystyle[\log(D_{\phi}(s,a))]$ (2)
$\displaystyle+\mathbb{E}_{(s,a)\sim\pi_{\theta}}[\log(1-D_{\phi}(s,a))].$

The discriminator $D_{\phi}$ is used to identify between samples produced by
the learning policy and expert demonstrations. Typically, the discriminator is
a binary classifier whose output tends to be $1$ for expert data and $0$ for
data generated by the policy. The optimal discriminator satisfies the
following:

$\log{D^{*}(s,a)}-\log{(1-D^{*}(s,a))}=\log{\frac{d_{\pi_{e}}(s,a)}{d_{\pi_{\theta}}(s,a)}},$
(3)

where $D^{*}=\frac{d_{\pi_{e}}(s,a)}{d_{\pi_{e}}(s,a)+d_{\pi_{\theta}}(s,a)}$.
The log-density ratio can be directly used as a surrogate reward function for
forward reinforcement learning in CFIL (Freund, Sarafian, and Kraus 2023).
Other popularly used surrogate reward function settings include
$R(s,a)=\log{D_{\phi}(s,a)}$ or $R(s,a)=-\log{(1-D_{\phi}(s,a))}$ used by Gail
(Ho and Ermon 2016) to minimize the Jensen-Shannon (JS) divergence and
$R(s,a)=\log{D_{\phi}(s,a)}-\log{(1-D_{\phi}(s,a))}$ used by AIRL (Fu, Luo,
and Levine 2017) to minimize the reverse Kullback-Leibler (KL) divergence.
Although these surrogate reward functions have some prior bias in the absorbed
state (Kostrikov et al. 2018), they still achieve good results in practical
use.

All of these surrogate reward functions contain the discriminator term, which
means that an advanced discriminator can provide better guidance for policy
learning.

With these surrogate reward functions, Gail and other imitation methods can be
combined with any forward reinforcement learning algorithm for policy
optimization, such as PPO (Schulman et al. 2017) or TRPO (Schulman et al.
2015) for on-policy algorithms and TD3 (Fujimoto, Hoof, and Meger 2018) or SAC
(Haarnoja et al. 2018) for off-policy algorithms.

### Diffusion Model

Diffusion model (Ho, Jain, and Abbeel 2020) is a latent variable model mapped
to potential space using a Markov chain. In the forward diffusion process,
noise is gradually added to the data $x_{t}$ at each time step $t$ using a
pre-defined variance schedule $\beta_{t}$. As $t$ increases, $x_{t}$ gets
closer to pure noise, and when $T\rightarrow\infty$, $x_{T}$ becomes the
standard Gaussian noise.

$q(x_{1:T}|x_{0}):=\prod_{t=1}^{T}{q(x_{t}|x_{t-1})},$ (4)
$q(x_{t}|x_{t-1}):=\mathcal{N}(x_{t};\sqrt{1-\beta_{t}}~{}x_{t-1},\beta_{t}I).$
(5)

By defining $\alpha_{t}=1-\beta_{t}$ and
$\overline{\alpha}_{t}=\prod_{s=1}^{t}(1-\beta_{s})$, we can use the
reparameterization trick to directly sample $x_{t}$ at any time step $t$ from
$x_{0}$:

$x_{t}=\sqrt{\overline{\alpha}_{t}}\ x_{0}+\sqrt{1-\overline{\alpha}_{t}}\
\epsilon,$ (6)

where $\epsilon\sim\mathcal{N}(0,1)$. The reverse diffusion process involves
gradual sampling from the pure Gaussian noise to recover $x_{0}$. This process
is modeled by latent variable models of the form $p_{\phi}(x_{0}):=\int
p_{\phi}(x_{0:T})dx_{1:T}$.

$p_{\phi}(x_{0:T}):=p(x_{T})\prod_{t=1}^{T}p_{\phi}(x_{t-1}|x_{t}),$ (7)
$p_{\phi}(x_{t-1}|x_{t}):=\mathcal{N}(x_{t-1};~{}\mu_{\phi}(x_{t},~{}t),{\Sigma}_{\phi}(x_{t},~{}t)),$
(8)

where
$\mu_{\phi}(x_{t},t)=\frac{1}{\sqrt{1-\beta_{t}}}(x_{t}-\frac{\beta_{t}}{\sqrt{1-\overline{\alpha}_{t}}}\epsilon_{\phi}(x_{t},~{}t))$
and $\epsilon_{\phi}$ is used to predict $\epsilon$ from $x_{t}$ in the
reverse process by parameter $\phi$. The covariance matrix
${\Sigma}_{\phi}(x_{t},t)$ can be fixed for non-trainable or parameterized for
learning.

The diffusion model maximizes the log-likelihood of the predicted distribution
of the model, which is given by
$\mathbb{L}=\mathbb{E}_{q(x_{0})}[\log{p_{\phi}(x_{0})}]$. Subsequently,
similar to VAE (Kingma and Welling 2013), the diffusion model maximizes the
variational lower bound (VLB) defined as
$\mathbb{E}_{q(x)}{[\ln{\frac{p_{\phi}(x_{0:T})}{q(x_{1:T}|x_{0})}}]}$. The
continuous derivation of VLB results in a simple loss function for the
diffusion model, which can be expressed as follows:

$\displaystyle L_{simple}(\phi)=\mathbb{E}_{\ t\sim\mu(1,T),\
\epsilon\sim\mathcal{N}(0,I),\ x_{0}\sim{q(x_{0})}}[||$
$\displaystyle\epsilon-$
$\displaystyle\epsilon_{\phi}(\sqrt{\overline{\alpha}_{t}}\
{x}_{0}+\sqrt{1-\overline{\alpha}_{t}}\ \epsilon$
$\displaystyle,~{}t)||^{2}],$ (9)

where $\overline{\alpha}_{t}=\prod_{s=1}^{t}{(1-\beta_{s})}$, $\mu$ is a
uniform distribution from between 1 and $T$ and $x_{0}$ sampled from real
data.

## Method

In this section, we present our proposed method - DiffAIL. First, we describe
how we model the unconditional diffusion process on the joint distribution of
state-action pairs. Second, we incorporate the diffusion model into the
discriminator of an AIL framework to improve the ability that the
discriminator to capture data distribution.

### Diffusion over state-action pairs

To simplify the notation, we abuse notation and define
$x_{t}=(s_{i},a_{i})_{t}$, where $t$ denotes the time step in the diffusion
process and $i$ denotes the time at which a particular state is visited or an
action in the trajectory. $x_{t}$ represents the state-action pairs available
in the data. Since the forward diffusion process is parameter-free, we denote
our forward diffusion process as follows:

$q(x_{t}|x_{t-1})=\mathcal{N}(x_{t};~{}\sqrt{1-\beta_{t}}\
x_{t-1},\beta_{t}I),$ (10)

where $\beta_{t}$ denotes the variance schedule at different time steps $t$.
We then define the reverse process of the parameterized diffusion model as
follows:

$\displaystyle
p_{\phi}(x_{0:T}|x_{t})=\mathcal{N}(x_{t-1};~{}\mu_{\phi}(x_{t},~{}t),{\Sigma}_{\phi}(x_{t},~{}t)).$
(11)

In our work, the covariance matrix is fixed as
${\Sigma}_{\phi}(x_{t},~{}t)=\sigma_{t}^{2}I=\beta_{t}I$ to be non-trainable
and represented using a pre-defined time schedule $\beta_{t}$, following the
same as in DDPM (Ho, Jain, and Abbeel 2020). According to the forward process,
$x_{t}$ can be derived from $x_{0}$ at any time step $t$; the mean can be
expressed as a noisy contained function:

$\displaystyle\mu_{\phi}(x_{t},~{}t)=\frac{1}{\sqrt{1-\beta_{t}}}(x_{t}-\frac{\beta_{t}}{\sqrt{1-\overline{\alpha}_{t}}}\
\epsilon_{\phi}(x_{t},~{}t)).$ (12)

Diffuser over state-action pairs is ultimately modelled as predicting the
noise at each time step in the reverse diffusion process:

$\displaystyle L(\phi)=\mathbb{E}_{x_{0}\sim
D,\epsilon\sim\mathcal{N}(0,I),t\sim\mu(1,T)}[Diff_{\phi}(x_{0},\epsilon,t)],$
(13) $\displaystyle
Diff_{\phi}(x_{0},\epsilon,t)=||\epsilon-\epsilon_{\phi}(\sqrt{\overline{\alpha}_{t}}x_{0}+\sqrt{1-\overline{\alpha}_{t}}\
\epsilon,~{}t||^{2},$ (14)

where $\mu$ is the uniform distribution, $\epsilon$ follows the standard
Gaussian noise distribution $\mathcal{N}(0,I)$ and
$\alpha_{t}=\prod_{s=1}^{T}(1-\beta_{s})$. In this work, we do not employ the
diffusion model to infer for generating samples but use its loss as a
distribution matching technique to improve the discriminator’s ability to
identify the expert demonstrations and policy data. Therefore, practically,
$Diff_{\phi}$ can refer to either expert demonstrations $(s_{i},a_{i})\sim
d_{\pi_{e}}(s,a)$ or data generated from interactions with the environment by
policy $(s_{i},a_{i})\sim d_{\pi}(s,a)$.

Although many studies (Song, Meng, and Ermon 2020; Lu et al. 2022a, b; Zheng
et al. 2023) have focused on designing accelerated sampling algorithms for
diffusion models to reduce the time cost of generating samples while
maintaining high sample quality, we found that only a few tens of diffusion
steps were sufficient to achieve excellent performance when modeling various
environment tasks in MuJoCo, unlike image generation tasks that require
thousands of sampling steps. Therefore, in our work, the primitive form of the
diffusion model (Ho, Jain, and Abbeel 2020) is adequate and efficient.

### Diffusion Adversarial Imitation Learning

Task | Hopper | HalfCheetah | Ant | Walker2d
---|---|---|---|---
Expert | 3402 | 4463 | 4228 | 6717
BC | 2376.22 ± 754.80 | 2781.85 ± 1143.11 | 2942.24 ± 344.62 | 1091.46 ± 187.60
Gail | 2915.67 ± 328.12 | 4383.44 ± 62.08 | 4530.43 ± 133.30 | 1942.81 ± 957.68
Valuedice | 2455.25 ± 658.43 | 4734.59 ± 229.67 | 3781.56 ± 571.15 | 4606.50 ± 927.80
CFIL | 3234.61 ± 266.62 | 4976.73 ± 60.41 | 4222.07 ± 201.32 | 5039.77 ± 307.17
DiffAIL | 3382.03 ± 142.86 | 5362.25 ± 96.92 | 5142.60 ± 90.05 | 6292.28 ± 97.65
Table 1: Learned policy best performance during training for different sota
imitation learning algorithms with 1 trajectory over 5 seeds in the standard
state-action setting.

Next, we describe how diffusion loss is integrated into the AIL framework to
improve discriminator generalization.

We introduce the diffusion model into adversarial imitation learning, which
combines the simple loss function of the diffusion model with the traditional
discriminator structure of AIL. For convenience, we abbreviate $x_{0}$ as $x$
in the following text. The objective function used in our method is the same
as Gail (Ho and Ermon 2016):

$\displaystyle\min_{{\pi_{\theta}}}\max_{D_{\phi}}\mathbb{E}_{\epsilon\sim\mathcal{N}(0,I),t\sim\mu(1,T)}[$
$\displaystyle\mathbb{E}_{x\sim{\pi_{e}}}[\log(D_{\phi}(x,\epsilon,t))]$
$\displaystyle+~{}\mathbb{E}_{x\sim\pi_{\theta}}[$
$\displaystyle\log(1-D_{\phi}(x,\epsilon,t))]],$ (15)

where the $D_{\phi}(x,\epsilon,t)$ is a discrimination composed of diffusion
noise loss by parameter $\phi$, and the policy network’s parameters are
represented as $\theta$. At the inner level maximization discriminator
provides higher values for expert data demonstrations and penalizes all other
data areas. The specific formulation of $D_{\phi}(x,\epsilon,t)$ is defined as
follows:

$\displaystyle D_{\phi}(x,\epsilon,t)$
$\displaystyle=\exp{(-Diff_{\phi}(x,\epsilon,t))}$ (16)
$\displaystyle=\exp{(-||\epsilon-\epsilon_{\phi}(\sqrt{\overline{\alpha}_{i}}\
x+\sqrt{1-\overline{\alpha}_{t}}\ \epsilon,~{}t)||^{2})}.$

We apply the exponential operation on the negative diffusion loss, which
allows DiffAIL to restrict the discriminator output strictly within $(0,1)$
and continue satisfying the optimization objective of minimizing the JS
divergence. This is consistent with the logit value range after the sigmoid
function in the original Gail and coincides with the monotonicity requirement.

By combining Eq. (Diffusion Adversarial Imitation Learning) and Eq. (16), we
can obtain our final optimization objective, which utilizes the diffusion loss
as the part of discriminator in AIL:

$\displaystyle\min_{{\pi_{\theta}}}\max_{D_{\phi}}\mathbb{E}_{\epsilon\sim\mathcal{N},t\sim\mu}[\mathbb{E}$
${}_{x\sim{\pi_{e}}}[\log(\exp(-Diff_{\phi}(x,\epsilon,t)))]$
$\displaystyle+\mathbb{E}_{x\sim\pi_{\theta}}[\log($ $\displaystyle
1-\exp(-Diff_{\phi}(x,\epsilon,t)))]].$ (17)

In this adversarial framework, training the diffusion model loss is an
adversarial process. $Diff_{\phi}(x,\epsilon,t)$ will minimize the diffusion
error for $x$ from the expert $\pi_{e}$ and maximize for $x$ generated samples
from the policy $\pi_{\theta}$. Therefore, $Diff_{\phi}(x,\epsilon,t)$ will
learn to distinguish between expert demonstrations and policy yield samples.
$\pi_{\theta}$ is a generator to generate state-action pairs similar to expert
data, confusing the discriminator.

For outer minimization of Eq.(Diffusion Adversarial Imitation Learning), it is
usually to optimize $\pi_{\theta}$ by an arbitrary reinforcement learning
algorithm with reward signal. So next, we use a diffusion-based discriminator
to obtain our surrogate reward function like Gail (Ho and Ermon 2016):

$R_{\phi}(x,\epsilon)=-\frac{1}{T}\sum_{t=1}^{T}\log(1-\exp{(-Diff_{\phi}(x,\epsilon,t)})).$
(18)

During the training process for the diffusion loss, we perform uniform
sampling from the interval $[1,T]$. When using diffusion loss as part of the
surrogate reward function, we only need to compute the average loss of the
diffusion model over $N$ timesteps instead of uniform sampling. We give high
rewards to the $x$ with low diffusion errors and vice versa. $\pi_{\theta}$
can conduct policy improvement from any forward reinforcement learning
algorithms with the surrogate reward function. We summarize the overall
process of the method in Algorithm 1.

Figure 1: Top: Learning curve for different sota imitation learning algorithms
with 1 trajectory 5 five seeds in the standard state-action setting. Bottom:
Learning curve for different sota imitation learning algorithms with one
trajectory over five seeds in the state-only setting. The x-axis denotes
timesteps, and the y-axis denotes the average return. The shadow areas
represent the standard deviation. Algorithm 1 Diffusion Adversarial Imitation
Learning

Require: initial policy network with parameter $\theta$, $Q$ network with
parameter $\omega$ and diffusion discriminator network with parameter $\phi$,
expert data $\mathcal{R}_{e}$, empty replay buffer $\mathcal{R}$, batch size
$k$, total timesteps $N$, diffusion timesteps $T$, learning rate
$\eta_{\phi},\eta_{\theta},\eta_{\omega}$ for parameters $\phi,\theta,\omega$
respectively.

1: let $n=0$.

2: while $n<N$ do

3: Collect ($s,a,s^{\prime}$) according to policy $\pi_{\theta}$ during
interaction and add samples into $\mathcal{R}$.

4: Sample
$\\{(s_{i}^{\pi_{\theta}},a_{i}^{\pi_{\theta}})\\}_{i=1}^{k}\sim\mathcal{R}$,
$\\{(s_{i}^{\pi_{e}},a_{i}^{\pi_{e}})\\}_{i=1}^{k}\sim\mathcal{R}_{e}$.

5: Sample
$\\{(t_{i})\\}_{i=1}^{k}\sim\mu(1,T),\\{(\epsilon_{i})\\}_{i=1}^{k}\sim\mathcal{N}(0,I)$
.

6: Compute the gradient of diffusion discriminator $\phi$:
$\displaystyle~{}~{}\nabla\mathcal{L}$
$\displaystyle=~{}\nabla_{\phi}\log(\exp(-Diff_{\phi}(s_{i}^{\pi_{e}},a_{i}^{\pi_{e}},\epsilon_{i},t_{i})))$
$\displaystyle+\nabla_{\phi}\log(1-\exp(-Diff_{\phi}(s_{i}^{\pi_{\theta}},a_{i}^{\pi_{\theta}},\epsilon_{i},t_{i}))).$

7: Update diffusion discriminator $\phi$ by gradient ascent:
$\phi\leftarrow\phi+\eta_{\phi}\nabla\mathcal{L}$

8: Sample
$\\{(s_{i}^{\pi_{\theta}},a_{i}^{\pi_{\theta}})\\}_{i=1}^{k}\sim\mathcal{R},\\{(\epsilon_{i})\\}_{i=1}^{k}\sim\mathcal{N}(0,I)$.

9: Compute the surrogate reward
$R_{\phi}(s_{i}^{\pi_{\theta}},a_{i}^{\pi_{\theta}},\epsilon_{i})$ by Eq.
(18).

10: Update SAC parameters $\pi_{\theta}$ and $Q_{\omega}$ by $\eta_{\phi}$ and
$\eta_{\omega}$ with the surrogate reward.

11: end while

## Related Work

### Adversarial Imitation Learning

AIL can be viewed as a unified perspective to minimize f-divergences,
including Gail (Ho, Jain, and Abbeel 2020), AIRL (Fu, Luo, and Levine 2017),
FAIRL (Ghasemipour, Zemel, and Gu 2020), f-Gail (Zhang et al. 2020), etc.
These methods aim to minimize divergence by matching the occupancy measure of
expert demonstrations and policy data. The policy improvement process and the
discriminator optimization process are similar to the GAN (Goodfellow et al.
2014). However, these methods typically rely on on-policy sampling and often
require millions or even tens of millions of interactions. To address this
issue, DAC (Kostrikov et al. 2018) extends Gail to the off-policy algorithm,
significantly improving the learning efficiency of AIL methods and reducing
the necessary online interaction to hundreds of thousands of steps. AEAIL
(Zhang et al. 2022) proposes to use the reconstruction loss of Autoencoder as
the reward function for AIL.

Additionally, a novel category of AIL methods, such as Valuedice (Kostrikov,
Nachum, and Tompson 2019) and CFIL (Freund, Sarafian, and Kraus 2023), employ
distribution correction estimation (Nachum et al. 2019) to derive the
objective function. They recover the log-density ratio by solving the optimal
point of the Donsker-Varadhan representation of the KL divergence. Valuedice
builds a fully off-policy optimization objective and outperforms BC in a
strictly offline setting. However, recent work (Li et al. 2022) proved that BC
is optimal in the strictly offline setting, which raises doubts about the
performance of Valuedice. CFIL argues that the log-density ratio cannot be
accurately estimated without appropriate modelling tools and proposes using a
coupled normalization flow to model the log-likelihood ratio explicitly.

### Diffusion Model with Reinforcement Learning

In RL, many methods leverage the powerful generative capabilities of diffusion
models to generate state-action pairs or trajectories. Diffusion-BC (Pearce et
al. 2023) addresses the limitation of insufficient expressiveness in
traditional behavior cloning by modeling diffusion model as a policy. Diffuser
(Janner et al. 2022) uses an unconditional diffusion model to generate
trajectories consisting of states and actions but requires an additional
reward function to guide the reverse denoising process towards high-return
trajectories. Diffuser Decision (Ajay et al. 2022) employs a conditional
diffusion model with classifier-free guidance to model trajectories containing
only states. However, since the modelling process does not involve actions, it
is necessary to additionally learn an additional inverse dynamics model to
select actions for any two adjacent states along the planned trajectory. BESO
(Reuss et al. 2023) applies the conditional diffusion model to goal-
conditional imitation learning by directly modelling state and action as
conditional distributions and using future states as conditions. Diffusion Q
(Wang, Hunt, and Zhou 2022) represents the diffusion model as a policy,
explicitly regularizes it, and adds the maximum action value function term to
the training objective to adjust the diffusion policy to select actions with
the highest Q value.

Our method also utilizes the diffusion model, which differs from the
abovementioned diffusion-based methods. Specifically, our approach remains
within the AIL framework, unlike BESO, Diffuser Decision, and Diffuser, which
directly employ the diffusion model as an expression policy rather than in
standard reinforcement learning or imitation learning. Moreover, DiffAIL
models an unconditional joint distribution that consists of state-action pairs
instead of the conditional distribution.

Figure 2: Ablation study on the diffusion timestamps of the diffusion
discriminator using N grids [2, 5, 10, 20] with 4 trajectories over 5 seeds.
Our findings suggest that a diffusion step of 10 yields good results for the
diffusion discriminator. Note that at timestep 2, Walker2d did not show
results for some random seeds due to training collapse.

## Experiments

In our experimental section, we explore the performance of the diffusion
discriminator on various tasks. Our experiments focus on four key aspects:

  * •

With the popular Mujoco environment, can our method achieve superior results
with a small amount of demonstration data in state-action and state-only
settings?

  * •

On those expert-level unseen state-action pairs, can DiffAIL improve the
generalization of successful identification for the discriminator compared to
a naïve discriminator?

  * •

Compared to naïve Gail, can DiffAIL provide a more linearly correlational
surrogate reward?

  * •

Since the number of diffusion steps significantly impacts model capability,
will this phenomenon also occur in our methods?

Benchmark & Dataset. In our experiments, we compare DiffAIL with the SOTA
algorithm in the Mujoco environments. We choose four representative tasks,
namely Hopper, HalfCheetah, Walker2d and Ant. We have $40$ trajectories on
each task dataset, each with $1000$ steps. We randomly subsample $n=[1,4,16]$
trajectories from the pool of $40$ trajectories. Please refer to Appendix A
for detailed dataset descriptions.

Baselines. We benchmark our method against state-of-the-art imitation learning
algorithms, including (1) Behavior cloning (Pomerleau 1991), the typical
supervised learning approach; (2) Gail (Ho and Ermon 2016), the classical AIL
method; (3) Valuedice (Kostrikov, Nachum, and Tompson 2019) which utilize
distribution correction estimation; (4) CFIL (Freund, Sarafian, and Kraus
2023), which employs a coupled flow to evaluate the log-density ratio both
used in state-action and state-only settings; (5) OPOLO (Zhu et al. 2020)
adjusts policy updates by using a reverse dynamic model to accelerate learning
for the state-only setting.

Implementation details. Our denoising model uses a $3$-layer MLP with Mish
activation function and $128$ hidden neurons in each layer rather than a
standard U-Net architecture. The training uses the Adam optimizer(Kingma and
Ba 2014) with a learning rate of $0.00003$ and the gradient penalty(Gulrajani
et al. 2017) for stability. Besides bc and Valuedice, other methods all use
SAC(Haarnoja et al. 2018) for forward reinforcement learning. For more
implementation details, please refer to Appendix B.

For a fair comparison, we implemented an off-policy version of Gail to improve
its sample efficiency, but it did not involve special processing of the bias
of the absorbed state reward function in DAC (Kostrikov et al. 2018). All
algorithms are evaluated under the same experimental setting. Experiments use
five random seeds $(0,1,2,3,4)$, run for an equal number of steps per task,
with $10$ thousand steps per epoch. Model performance is evaluated every $10$
thousand steps by averaging rewards over ten episodes. We then smooth the
results using a sliding window of $4$ before plotting means and standard
deviations across the $5$ seeds.

### Comparison to SOTA methods

We conduct experiments on multiple tasks from the popular MuJoCo benchmark,
using $1$, $4$, and $16$ expert trajectories to evaluate performance across
varying trajectory numbers. For each algorithm, we assess the real average
episodic return the agent achieves when interacting in the environment.

With only $1$ expert trajectory, the top of Figure 1 shows our method
outperforms other imitation learning baselines. Our approach achieves near
expert-level asymptotic performance on all tasks and even significantly
surpasses the expert on HalfCheetah and Ant. This expert-surpassing result is
not performed by other methods. Especially table 1 clearly shows that our
method outperforms all other methods significantly on best performance during
training.

Subsequently, we conducted our experiment under a state-only setting and found
that the advantages brought by the powerful discriminator are very
significant, not just in modeling state-action distribution. Under the state-
only setting, DiffAIL models the state and the next state distribution.
Importantly, we don’t need to modify any hyperparameters, and then we directly
validate the performance of our method. We compare with OPPLO and CFIL. From
the bottom of Figure 1, it can be seen that our method can reach expert-level
performance with 1 trajectory and outperform other methods, especially on
HalfCheetah and Ant.

All of these significant advantages come from the diffusion discriminator
being able to accurately capture the data distribution of experts, whether it
is state action pairs or state next state pairs.

While our method achieves expert-level performance on multiple tasks, it is
less time-efficient than Valuedice’s optimized implementation. We acknowledge
Valuedice’s impressive efficiency gains, though algorithmic speed is not our
primary focus. However, in one trajectory situation, we find it can’t maintain
optimality in later training stages stably. Still, our approach performs
similarly or even more efficiently compared to other methods, except for
Valuedice.

The experimental results of training additional $4$ and $16$ trajectories can
be seen in Appendix E.

### Discriminate on unused expert trajectories

Env | Gail | DiffAIL
---|---|---
Hopper | 83.3% ± 1.2% | 92.8% ± 2.5%
HalfCheetah | 83.8% ± 0.9% | 94.5% ± 0.6%
Walker2d | 90.5% ± 0.8% | 96.1% ± 0.6%
Ant | 81.7% ± 1.6% | 98.8% ± 0.1%
Table 2: A comparison of the ability to distinguish expert demonstrations out
of data for Gail and DiffAIL with four trajectories over five seeds. It
suggests that DiffAIL has a significant improvement over Gail’s discriminator.

We conduct an additional experiment to verify that our method improves the
generalization of unseen expert-level state-action pairs for discrimination.
We take the discriminator which has completed training on just $4$
trajectories, and test it on $36$ held-out expert trajectories from the
remaining dataset. If the discriminator can still successfully identify these
new trajectories as expert-like, it indicates improved generalization.

However, selecting an evaluation metric is challenging. Theoretically, the
discriminator will eventually converge to $0.5$(Goodfellow et al. 2014), where
it cannot distinguish between expert demonstrations and policy data. But in
practice, training ends before the discriminator and generator reach the ideal
optimum, and the discriminator often overfits the expert demonstrations.
Therefore, we adopted a simple evaluation criterion: if the discriminator
output is greater than or equal to $0.5$, the expert data is considered
successfully distinguished; otherwise, the discrimination fails.

As shown in Table 2, it represents the percentage of state-action pairs
successfully distinguished by the discriminator from $36$ trajectories, i.e.,
$36000$ state-action pairs. From the benchmark of HalfCheetah and Ant, DiffAIL
can significantly improve the discrimination generalization of unseen expert
demonstrations. In other tasks, DiffAIL also exhibits stronger discriminative
abilities than Gail, which sufficiently exhibits the powerful advantage of the
diffusion discriminator in capturing the expert distribution. Notably, all
trajectories of the same task are yielded by interactions between the same
trained policy. Therefore, substantial similarities may exist across state-
action pairs from different trajectories, which also explains why Gail can
achieve decent discrimination as well.

Figure 3: Correlation between average discriminator return and average actual
return with Gail and DiffAIL in 3 seeds. The returns have been normalized for
display convenience.

### Return correlations

As Eq. (3) shows, a better discriminator can provide better surrogate rewards
to guide forward reinforcement learning. We verify this by experimentally
analyzing the Pearson coefficient(PC) between the surrogate discriminator
return and the actual return. We compare ours with the vanilla discriminator
in GAIL, as shown in Figure 3: The PC for DiffAIL are
$[97.1\%,98.8\%,99.1\%,99.6\%]$, while for GAIL are
$[97.0\%,54.1\%,-89.0\%,92.1\%]$. This indicates that our diffusion surrogate
reward function is more linearly correlated with actual rewards than GAIL,
which can better guide policy learning.

### Ablation study

We conduct ablation studies to investigate how various diffusion steps impact
the distribution-matching ability of the diffusion discriminator on state-
action pairs.

Figure 2 shows that we set the diffusion steps to $t=[2,5,10,20]$,
respectively. The results show that increasing diffusion steps leads to more
stable performance and faster convergence across all tasks, aligning with
results in diffusion models for image generation. However, we also observe
that more diffusion steps increase training time, as exhibited in Appendix
(Table 4). Therefore, we finally select $t=10$ as the diffusion step to
balance discriminator quality and training cost.

## Conclusion

In this work, we propose DiffAIL, which utilizes diffusion models for AIL.
Unlike prior methods that directly apply diffusion models as an expressive
policy, our approach leverages the diffusion model to enhance the distribution
matching ability of the discriminator, which enables the discriminator to
improve generalization for expert-level state-action pairs. In both the state-
action setting and state-only setting, We model the joint distribution of
state-action pairs and state and next state pairs, respectively, with an
unconditional diffusion model. Surprisingly, in both settings with a single
expert trajectory, the result shows our method achieves SOTA performance
across all tasks. Additionally, we discuss the limitations of this work and
some promising future directions in Appendix D.

## Acknowledgments

This work was supported by the NSFC (Nos. 62206159, 62176139), the Natural
Science Foundation of Shandong Province (Nos. ZR2022QF117, ZR2021ZD15), the
Fundamental Research Funds of Shandong University, the Fundamental Research
Funds for the Central Universities. G. Wu was sponsored by the TaiShan
Scholars Program.

## References

  * Ajay et al. (2022) Ajay, A.; Du, Y.; Gupta, A.; Tenenbaum, J.; Jaakkola, T.; and Agrawal, P. 2022.  Is conditional generative modeling all you need for decision-making?  _arXiv preprint arXiv:2211.15657_.
  * Bao et al. (2022) Bao, F.; Li, C.; Zhu, J.; and Zhang, B. 2022.  Analytic-dpm: an analytic estimate of the optimal reverse variance in diffusion probabilistic models.  _arXiv preprint arXiv:2201.06503_.
  * Freund, Sarafian, and Kraus (2023) Freund, G.; Sarafian, E.; and Kraus, S. 2023.  A Coupled Flow Approach to Imitation Learning.  _arXiv preprint arXiv:2305.00303_.
  * Fu, Luo, and Levine (2017) Fu, J.; Luo, K.; and Levine, S. 2017.  Learning robust rewards with adversarial inverse reinforcement learning.  _arXiv preprint arXiv:1710.11248_.
  * Fujimoto, Hoof, and Meger (2018) Fujimoto, S.; Hoof, H.; and Meger, D. 2018.  Addressing function approximation error in actor-critic methods.  In _International conference on machine learning_ , 1587–1596. PMLR.
  * Garg et al. (2021) Garg, D.; Chakraborty, S.; Cundy, C.; Song, J.; and Ermon, S. 2021.  Iq-learn: Inverse soft-q learning for imitation.  _Advances in Neural Information Processing Systems_ , 34: 4028–4039.
  * Ghasemipour, Zemel, and Gu (2020) Ghasemipour, S. K. S.; Zemel, R.; and Gu, S. 2020.  A divergence minimization perspective on imitation learning methods.  In _Conference on Robot Learning_ , 1259–1277. PMLR.
  * Goodfellow et al. (2014) Goodfellow, I.; Pouget-Abadie, J.; Mirza, M.; Xu, B.; Warde-Farley, D.; Ozair, S.; Courville, A.; and Bengio, Y. 2014.  Generative adversarial nets.  _Advances in neural information processing systems_ , 27.
  * Gulrajani et al. (2017) Gulrajani, I.; Ahmed, F.; Arjovsky, M.; Dumoulin, V.; and Courville, A. C. 2017.  Improved training of wasserstein gans.  _Advances in neural information processing systems_ , 30.
  * Haarnoja et al. (2018) Haarnoja, T.; Zhou, A.; Abbeel, P.; and Levine, S. 2018.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  In _International conference on machine learning_ , 1861–1870. PMLR.
  * Ho and Ermon (2016) Ho, J.; and Ermon, S. 2016.  Generative adversarial imitation learning.  _Advances in neural information processing systems_ , 29.
  * Ho, Jain, and Abbeel (2020) Ho, J.; Jain, A.; and Abbeel, P. 2020.  Denoising diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 33: 6840–6851.
  * Janner et al. (2022) Janner, M.; Du, Y.; Tenenbaum, J. B.; and Levine, S. 2022.  Planning with diffusion for flexible behavior synthesis.  _arXiv preprint arXiv:2205.09991_.
  * Kingma and Ba (2014) Kingma, D. P.; and Ba, J. 2014.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_.
  * Kingma and Welling (2013) Kingma, D. P.; and Welling, M. 2013.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_.
  * Kostrikov et al. (2018) Kostrikov, I.; Agrawal, K. K.; Dwibedi, D.; Levine, S.; and Tompson, J. 2018.  Discriminator-actor-critic: Addressing sample inefficiency and reward bias in adversarial imitation learning.  _arXiv preprint arXiv:1809.02925_.
  * Kostrikov, Nachum, and Tompson (2019) Kostrikov, I.; Nachum, O.; and Tompson, J. 2019.  Imitation learning via off-policy distribution matching.  _arXiv preprint arXiv:1912.05032_.
  * Li et al. (2022) Li, Z.; Xu, T.; Yu, Y.; and Luo, Z.-Q. 2022.  Rethinking ValueDice: Does It Really Improve Performance?  _arXiv preprint arXiv:2202.02468_.
  * Liu et al. (2022) Liu, L.; Ren, Y.; Lin, Z.; and Zhao, Z. 2022.  Pseudo numerical methods for diffusion models on manifolds.  _arXiv preprint arXiv:2202.09778_.
  * Liu, Gong, and Liu (2022) Liu, X.; Gong, C.; and Liu, Q. 2022.  Flow straight and fast: Learning to generate and transfer data with rectified flow.  _arXiv preprint arXiv:2209.03003_.
  * Liu et al. (2023) Liu, X.-H.; Xu, F.; Zhang, X.; Liu, T.; Jiang, S.; Chen, R.; Zhang, Z.; and Yu, Y. 2023.  How To Guide Your Learner: Imitation Learning with Active Adaptive Expert Involvement.  _arXiv preprint arXiv:2303.02073_.
  * Lu et al. (2022a) Lu, C.; Zhou, Y.; Bao, F.; Chen, J.; Li, C.; and Zhu, J. 2022a.  Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps.  _Advances in Neural Information Processing Systems_ , 35: 5775–5787.
  * Lu et al. (2022b) Lu, C.; Zhou, Y.; Bao, F.; Chen, J.; Li, C.; and Zhu, J. 2022b.  Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models.  _arXiv preprint arXiv:2211.01095_.
  * Mnih et al. (2013) Mnih, V.; Kavukcuoglu, K.; Silver, D.; Graves, A.; Antonoglou, I.; Wierstra, D.; and Riedmiller, M. 2013.  Playing atari with deep reinforcement learning.  _arXiv preprint arXiv:1312.5602_.
  * Mnih et al. (2015) Mnih, V.; Kavukcuoglu, K.; Silver, D.; Rusu, A. A.; Veness, J.; Bellemare, M. G.; Graves, A.; Riedmiller, M.; Fidjeland, A. K.; Ostrovski, G.; et al. 2015.  Human-level control through deep reinforcement learning.  _nature_ , 518(7540): 529–533.
  * Nachum et al. (2019) Nachum, O.; Chow, Y.; Dai, B.; and Li, L. 2019.  Dualdice: Behavior-agnostic estimation of discounted stationary distribution corrections.  _Advances in Neural Information Processing Systems_ , 32.
  * Nowozin, Cseke, and Tomioka (2016) Nowozin, S.; Cseke, B.; and Tomioka, R. 2016.  f-gan: Training generative neural samplers using variational divergence minimization.  _Advances in neural information processing systems_ , 29.
  * Pearce et al. (2023) Pearce, T.; Rashid, T.; Kanervisto, A.; Bignell, D.; Sun, M.; Georgescu, R.; Macua, S. V.; Tan, S. Z.; Momennejad, I.; Hofmann, K.; et al. 2023.  Imitating human behaviour with diffusion models.  _arXiv preprint arXiv:2301.10677_.
  * Pomerleau (1991) Pomerleau, D. A. 1991.  Efficient training of artificial neural networks for autonomous navigation.  _Neural computation_ , 3(1): 88–97.
  * Reuss et al. (2023) Reuss, M.; Li, M.; Jia, X.; and Lioutikov, R. 2023.  Goal-conditioned imitation learning using score-based diffusion policies.  _arXiv preprint arXiv:2304.02532_.
  * Ross, Gordon, and Bagnell (2011) Ross, S.; Gordon, G.; and Bagnell, D. 2011.  A reduction of imitation learning and structured prediction to no-regret online learning.  In _Proceedings of the fourteenth international conference on artificial intelligence and statistics_ , 627–635. JMLR Workshop and Conference Proceedings.
  * Schulman et al. (2015) Schulman, J.; Levine, S.; Abbeel, P.; Jordan, M.; and Moritz, P. 2015.  Trust region policy optimization.  In _International conference on machine learning_ , 1889–1897. PMLR.
  * Schulman et al. (2017) Schulman, J.; Wolski, F.; Dhariwal, P.; Radford, A.; and Klimov, O. 2017.  Proximal policy optimization algorithms.  _arXiv preprint arXiv:1707.06347_.
  * Silver et al. (2016) Silver, D.; Huang, A.; Maddison, C. J.; Guez, A.; Sifre, L.; Van Den Driessche, G.; Schrittwieser, J.; Antonoglou, I.; Panneershelvam, V.; Lanctot, M.; et al. 2016.  Mastering the game of Go with deep neural networks and tree search.  _nature_ , 529(7587): 484–489.
  * Song, Meng, and Ermon (2020) Song, J.; Meng, C.; and Ermon, S. 2020.  Denoising diffusion implicit models.  _arXiv preprint arXiv:2010.02502_.
  * Song et al. (2023) Song, Y.; Dhariwal, P.; Chen, M.; and Sutskever, I. 2023.  Consistency Models.  _arXiv preprint arXiv:2303.01469_.
  * Todorov, Erez, and Tassa (2012) Todorov, E.; Erez, T.; and Tassa, Y. 2012.  Mujoco: A physics engine for model-based control.  In _2012 IEEE/RSJ international conference on intelligent robots and systems_ , 5026–5033. IEEE.
  * Wang, Hunt, and Zhou (2022) Wang, Z.; Hunt, J. J.; and Zhou, M. 2022.  Diffusion policies as an expressive policy class for offline reinforcement learning.  _arXiv preprint arXiv:2208.06193_.
  * Zhang et al. (2022) Zhang, K.; Zhao, R.; Zhang, Z.; and Gao, Y. 2022.  Auto-Encoding Adversarial Imitation Learning.
  * Zhang et al. (2020) Zhang, X.; Li, Y.; Zhang, Z.; and Zhang, Z.-L. 2020.  f-gail: Learning f-divergence for generative adversarial imitation learning.  _Advances in neural information processing systems_ , 33: 12805–12815.
  * Zheng et al. (2023) Zheng, K.; Lu, C.; Chen, J.; and Zhu, J. 2023.  DPM-Solver-v3: Improved Diffusion ODE Solver with Empirical Model Statistics.  _arXiv preprint arXiv:2310.13268_.
  * Zhu et al. (2020) Zhu, Z.; Lin, K.; Dai, B.; and Zhou, J. 2020.  Off-policy imitation learning from observations.  _Advances in Neural Information Processing Systems_ , 33: 12402–12413.

## Appendices

Tasks | HalfCheetah | Hopper | Walker2d | Ant
---|---|---|---|---
Optimizer | Adam | Adam | Adam | Adam
Batch Size | $256$ | $256$ | $256$ | $256$
diffusion hidden layer | $128$ | $128$ | $128$ | $128$
diffusion timesteps | $10$ | $10$ | $10$ | $10$
noise schedule | linear | linear | linear | linear
Diffusion learning rate | $3\times 10^{-4}$ | $3\times 10^{-4}$ | $3\times 10^{-4}$ | $3\times 10^{-4}$
use gradient penalty | True | False | True | True
gradient penalty weight | 0.1 | NAN | 0.001 | 0.1
Discount factor $\gamma$ | $0.99$ | $0.99$ | $0.99$ | $0.99$
$Q$ learning rate | $3\times 10^{-3}$ | $3\times 10^{-3}$ | $3\times 10^{-3}$ | $3\times 10^{-3}$
$\pi$ learning rate | $3\times 10^{-3}$ | $3\times 10^{-3}$ | $3\times 10^{-3}$ | $3\times 10^{-3}$
Table 3: Optimal hyperparameters of DiffAIL in both state-action and state-only settings. DiffAIL parameters change only on the gradient penalty term, indicating that it is very applicable.  diffusion steps | Gail | DiffAIL$(2)$ | DiffAIL$(5)$ | DiffAIL$(10)$ | DiffAIL$(20)$
---|---|---|---|---|---
HalfCheetah $(timesteps=50)$ | 3.1h | 5.5h | 5.9h | 6.2h | 8.0h
Ant $(timesteps=50)$ | 3.2h | 5.7h | 6.1 | 6.4 | 8.2h
Hopper $(timesteps=75)$ | 4.5h | 6.5h | 6.9h | 7.5h | 9.3h
Walker2d $(timesteps=75)$ | 4.8h | 8.4h | 8.9h | 9.5h | 11.4h
Table 4: A comparison of running time for different diffusion steps with
DiffAIL on $4$ Mujoco Tasks, with a unit equal to $1$ hour. The first line
denotes the algorithms, where the numbers in parentheses represent the
diffusion steps. The first column represents the Mujoco task, where the number
in parentheses represents the training steps in units of $10,000$. Figure 4:
Top: Learning curve for different sota imitation learning algorithms with $4$
trajectories over $5$ seeds in the standard state-action setting. Bottom:
Learning curve for different sota imitation learning algorithms with $4$
trajectories over $5$ seeds in the state-only setting. The x-axis denotes
timesteps, and the y-axis denotes the average return. The shadow areas
represent the standard deviation. Figure 5: Top: Learning curve for different
sota imitation learning algorithms with $16$ trajectories over $5$ seeds in
the standard state-action setting. Bottom: Learning curve for different sota
imitation learning algorithms with $16$ trajectories over $5$ seeds in the
state-only setting. The x-axis denotes timesteps, and the y-axis denotes the
average return. The shadow areas represent the standard deviation.

### A Dataset discription

We utilize benchmark datasets from two sources: The Halfcheetah, Walker2d and
Ant datasets are selected from the open-source valuedice(Kostrikov, Nachum,
and Tompson 2019) repository. For the Hopper, we use the SAC algorithm to
interact with the environment for $1$ million steps to obtain an expert
policy. We then collect $40$ expert trajectories by executing this policy in
the environment. Since adversarial imitation learning requires a small amount
of expert data, we subsample n trajectories from $40$ trajectories. We test
$n=[1,4,16]$ to evaluate performance with varying trajectory set sizes.

### B Implementation details

We build Gail(Ho and Ermon 2016) and DiffAIL on top of the implementation from
the FAIRL(Ghasemipour, Zemel, and Gu
2020)111https://github.com/KamyarGh/rl_swiss. For the diffusion model in
DiffAIL, we modify the code of Diffusion Q(Wang, Hunt, and Zhou
2022)222https://github.com/zhendong-wang/diffusion-policies-for-offline-rl to
create an unconditional variant. Our denoising model uses a $3$-layer MLP with
Mish activations and $128$ hidden units rather than a U-Net architecture. We
embed the timestep with a linear function. The training uses the Adam
optimizer(Kingma and Ba 2014) with a learning rate of $0.00003$ and the
gradient penalty(Gulrajani et al. 2017) for stability. Besides bc and
Valuedice, other methods use sac for forward reinforcement learning. The
specific optimal parameters of DiffAIL can be seen in Table 3. For the
baselines, we leverage open-source implementations of BC,
Valuedice333https://github.com/google-research/google-
research/tree/master/value_dice, and
CFIL444https://github.com/gfreund123/cfil,
OPOLO555https://github.com/illidanlab/opolo-code, which allow rapid
benchmarking against prior state-of-the-art methods. We implemented all our
experiments on a server with $8$ Tesla P100 graphics cards.

### C Optimal hyperparameters

For the DiffAIL experiments, we summarize the optimal hyperparameters in Table
3. We mainly considered two hyperparameters - the learning rate of the
diffusion model and the gradient penalty for the discriminator. The learning
rate of the diffusion model was determined via grid search over
$[0.0003,0.00003,0.000003]$, and the gradient penalty was determined via grid
search over $[0.001,0.01,0.1]$. For forward reinforcement learning, we employ
the SAC algorithm with fixed parameters across all tasks.

### D Limitaion and future work

In this work, we introduced diffusion models into the adversarial imitation
learning framework and proposed DiffAIL. Experiments show that this approach
could significantly enhance adversarial imitation learning performance, owing
to the powerful distribution-capturing capability of the diffusion
discriminator. Although sampling inference is unnecessary during this
adversarial training, computing the noise mse loss $Diff_{\phi}$ still
requires multi-step noise sampling, which undoubtedly increases the
computational cost of DiffAIL. However, our ablation studies on diffusion
steps showed that more steps lead to better discriminator training, posing a
trade-off between the two. Meanwhile, we also noticed that much research has
focused on accelerating sampling(Song, Meng, and Ermon 2020; Bao et al. 2022;
Lu et al. 2022a; Liu, Gong, and Liu 2022; Liu et al. 2022; Song et al. 2023)
in diffusion models - these methods can maintain generation quality while
drastically reducing sampling steps, even enabling one-step generation(Liu,
Gong, and Liu 2022; Song et al. 2023). Therefore, future Diffusion AIL could
leverage advanced samplers to accelerate training and reduce computational
costs.

Additionally, we note that this framework of utilizing diffusion models in
adversarial training is generalizable. This raises the question of whether
introducing this framework into GANs could improve the quality and diversity
of samples generated by adversarial networks.

### E Compared with sota on more trajectories

We compared our method with the currently state-of-the-art baselines on more
trajectories, as shown at the top of Figures 4 and 5. It can be observed that
the performance of all methods improves as the number of trajectories
increases, with faster and more stable convergence. Valuedice utilizes
distribution correction estimation for off-policy and even offline policy
learning. Its performance on 16 trajectories surprised us, especially the
extremely fast and stable convergence on the walker2d task. However, Valuedice
requires full trajectory demonstrations, which is a drawback. Our method also
exhibited strong performance, achieving expert-level scores across all tasks
and significantly surpassing the expert demonstrations substantially on the
HalfCheetah and Ant tasks.

Additionally, we conducted multi-trajectory experiments under a state-only
setting, with results shown at the bottom of Figures 4 and 5. Our method
demonstrates solid performance, especially on the Ant and HalfCheetah tasks,
substantially surpassing the expert demonstration performance. CFIL and OPOLO
do not achieve this level of performance.