DPAIL: Training Diffusion Policy for Adversarial Imitation Learning without Policy Optimization

(Weakness 1) Inference latency and reward-world applicability

As noted in the section on limitations, DPAIL incurs longer inference time due to the use of diffusion models as policies. This is a trade-off for their high expressiveness and ability to capture multi-modal behaviors, which are essential for our goal of modeling diverse expert demonstrations. There has been promising progress in accelerating diffusion model inference (e.g., through distillation or fast samplers), and integrating such techniques into our framework is a promising direction for improving scalability and real-time applicability, which we leave as future work.

(Weakness 2) Sensitivity of the additional hyperparameter

In addition to the hyperparameter planning horizon $H$ and sampling step $N$ discussed in the experiments section, we conducted additional sensitivity analysis on the number of generated samples per iteration in AntGoal-v3 environment. The results are shown below:

Our findings indicate that when the number of generated samples is too small, performance degrades due to insufficient learning from the generated samples. However, once the sample size exceeds a certain threshold, performance improves and then saturates, exhibiting only marginal gains with further increases.

(Weakness 3) Differences from DRAIL and DiffAIL

While DPAIL shares similarities with prior works such as DRAIL and DiffAIL, the underlying mechanisms are fundamentally different. Both DRAIL and DiffAIL use explicit diffusion-based discriminators, and then used to provide rewards for training a separate, uni-modal Gaussian policy, often via RL. In contrast, DPAIL does not use an explicit discriminator at all. Instead, the discriminator is implicitly defined by the diffusion policy itself through the ASAF framework. This implicit formulation enables DPAIL to capture and reproduce diverse behaviors from multi-modal demonstrations, whereas DRAIL and DiffAIL are limited in this regard since their diffusion models are used solely as a reward signal for training unimodal policies.

(Question 1 & 2) Hyperparameter Tuning

We made efforts to ensure a fair and meaningful comparison across all baseline algorithms. For PPO-based baselines (e.g., GAIL, DiffAIL, InfoGAIL), we performed hyperparameter tuning on several factors that impact performance. For example, we searched over the # of rollout-length per iteration, using values [5e3, 1e4, 5e4, 1e5]. In our experiments on MuJoCo tasks, rollout lengths of $\geq 5 \times 10^4$ generally led to stable and successful training, so we selected those as default settings for the PPO-based baselines.

Another hyperparameter is K, the # of iterations. Since different algorithms require different numbers of iterations $K$ to reach their converged performance, we determined a sufficiently large value of $K$ for each method and task during hyperparameter tuning. This ensures that performance is evaluated after sufficiently training rather than at a fixed iteration cutoff. While our method requires fewer environment interactions to converge, we allowed other baselines to consume more interactions if needed to reach their converged performance. This ensures a fair comparion based on fully trained policies, rather than under a fixed interaction budget. All experiments were run using the same set of random seeds (0-4). For details about the tuned hyperparameters, please refer to Appendix D (Implementation Details).

(Question 3) Extension to latent variables and Conditional Diffusion Policies

Yes, we have preliminary ideas on integrating latent variables $z$ for mode conditioning, and such an extension is compatible with the current DPAIL lower bound. Below, we briefly outline the idea. We define a conditional diffusion model with a latent variable $z$, where the joint probability $p\_\theta(\tau^{0:N}, z)$ factorizes as:

$p\_\theta(\tau^{0:N}, z) = q(\tau^N)q(z) \prod\_{i=1}^{N} p\_\theta(\tau^{i-1} | \tau^i, z)$.

where $q(z)$ is the prior of $z$, $q(\tau^N)$ is a standard normal distribution.

To evaluate the log-ratio $\log \frac{p\_\theta(\tau^0)}{p\_{\theta^\text{old}}(\tau^0)}$, we rewrite it with the posterior $q(z|\tau^0)$ as:

$\log \frac{p\_\theta(\tau^0)}{p\_{\theta^\text{old}}(\tau^0)} = \log \frac{\prod\_{i=1}^N q(\tau^i|\tau^{i-1})}{\prod\_{i=1}^N q(\tau^i | \tau^{i-1})} \frac{q(z|\tau^0)}{q(z|\tau^0)} \frac{p\_\theta(\tau^0)}{p\_{\theta^\text{old}}(\tau^0)} =\log \frac{\prod\_{i=1}^N p\_\theta(\tau^{i-1}|\tau^{i}, z)}{\prod\_{i=1}^N p\_{\theta^{\text{old}}}(\tau^{i-1} | \tau^{i}, z)} \frac{q(z)}{q(z)} \frac{q(\tau^N)}{q(\tau^N)}$,

where $q(\tau^N)$ and $q(z)$ cancel out. Taking expectation of this log-ratio over $\tau^{1:N} \sim q(\cdot | \tau^0), z~\sim q(z|\tau^0)$ does not change the value since it is still a ratio of marginal. Thus, the objective can be rewritten as:

$\max\_\theta \mathbb{E}\_{\tau^0 \sim p_E} \left[ \log \sigma \left ( \mathbb{E}\_{\substack{z \sim q\_(z | \tau^0),\\\\ \tau^{1:N} \sim q(\tau^{1:N} | \tau^0)}} \log \frac{\prod\_{i=1}^N p\_\theta(\tau^{i-1}|\tau^{i}, z)} {\prod\_{i=1}^N p\_{\theta_{\text{old}}}(\tau^{i-1}|\tau^{i}, z)} \right) \right].$

Given the expert sample $\tau^0\sim p_E$, the approximate variational posterior $q_\phi(z|\tau^0)$, the reverse process $p_\theta(\tau^{i-1}| \tau^i, z)$, this leads to the following lowerbound:

$$ &\mathbb{E}_{\tau^{0} \sim p_{E}} \left [ \log \sigma \left( \mathbb{E}_{\substack{z \sim q_\phi(z|\tau^0) \\ {\tau^{1:N} \sim q(\tau^{1:N}| \tau^0)}}} \log \frac{\prod_{i=1}^{N} p_\theta(\tau^{i-1}|\tau^{i},z)} {\prod_{i=1}^{N} p_{\theta^{\text{old}}}(\tau^{i-1}|\tau^{i},z)} \right) \right] \\ &= \mathbb{E}_{\tau^{0} \sim p_{E}} \left [ \log \sigma \left( \mathbb{E}_{\substack{z\sim q_\phi(z|\tau^0) \\ {\tau^{1:N}\sim q(\tau^{1:N}| \tau^0)}}} \sum_{i=1}^{N} \log \frac{p_\theta(\tau^{i-1}| \tau^i, z)}{p_{\theta^{\text{old}}}(\tau^{i-1}| \tau^{i}, z)} \right) \right] \\ &= \mathbb{E}_{\tau^{0} \sim p_{E}} \left [ \log \sigma \left( \sum_{i=1}^{N} \mathbb{E}_{\substack{z\sim q_\phi(z|\tau^0) \\ {\tau^{1:N} \sim q(\tau^{1:N}| \tau^0)}}} \log \frac{p_\theta(\tau^{i-1}| \tau^i, z)}{p_{\theta^{\text{old}}}(\tau^{i-1}| \tau^{i},z)} \right) \right] \\ &= \mathbb{E}_{\tau^{0} \sim p_{E}} \left[ \log \sigma \left( \sum_{i=1}^{N} \mathbb{E}_{\substack{z\sim q_\phi(z|\tau^0) \\ \tau^{i-1}, \tau^{i} \sim q(\tau^i | \tau^0) q(\tau^{i-1}|\tau^i, \tau^0)}} \log \frac{p_\theta(\tau^{i-1}| \tau^i, z)}{p_{\theta^{\text{old}}}(\tau^{i-1}| \tau^{i},z)} \right) \right] \\ &\geq \mathbb{E}_{\tau^{0} \sim p_{E}} \left [ \mathbb{E}_{\substack{i \sim \mathcal{U}(1, N),\\ z\sim q_\phi(z|\tau^0), \tau^{i} \sim q(\tau^i|\tau^0)}} \log \sigma \left( {N} \mathbb{E}_{\tau^{i-1} \sim q(\tau^{i-1}|\tau^i, \tau^0)} \log \frac{p_\theta(\tau^{i-1}| \tau^i, z)}{p_{\theta^{\text{old}}}(\tau^{i-1}| \tau^{i}, z)} \right) \right] (\text{concavity}) \\ &= \mathbb{E}_{\tau^{0} \sim p_{E}}\left [ \mathbb{E}_{\substack{i\sim \mathcal{U}(1, N),\\ z\sim q_\phi(z|\tau^0), \tau^{i} \sim q(\tau^i|\tau^0)}} \log \sigma \left( N \mathbb{E}_{\tau^{i-1} \sim q(\tau^{i-1}|\tau^i, \tau^0)} \log \frac {q(\tau^{i-1}|\tau^i, \tau^0)}{p_{\theta^{\text{old}}}(\tau^{i-1}| \tau^i,z)} -\log \frac {q(\tau^{i-1}|\tau^i, \tau^0)}{p_\theta(\tau^{i-1}| \tau^{i},z)}\right) \right]

$$

$\therefore \mathcal{L}\_{\text{DPAIL}}^{(1)}(\theta,\theta^{\text{old}},\tau^0, z) = \mathbb{E}\_{\tau^0, z, i, \epsilon} \\Big[ \log \sigma \\Big (N \\left (\\| \epsilon - \epsilon\_{\theta^{\text{old}}}(\tau^i,z,i)\\|^2 \- \\| \epsilon - \epsilon\_\theta(\tau^i,z,i)\\|^2 \\right )\\Big )\\Big],$

Here, $z \sim q\_\phi(z|\tau^0)$, $\epsilon \sim \mathcal{N}(0, I)$, and $i \sim \mathcal{U}(1, N)$.

Similarly, for given generative samples \bar{\tau}\sim p\_{\theta^\{\text{old}}}, we can also get $\mathcal{L}\_{\text{DPAIL}}^{(2)}(\theta,\theta^{\text{old}},\bar{\tau}^0, z)$.

The overall objective becomes:

$\therefore \mathbb{E}\_{\substack{\tau^0 \sim p\_E, \\\\ z \sim q\_\phi(z|\tau^0)}}\\left[ \mathcal{L}^{(1)}\_{\text{DPAIL}} (\theta, \theta^{\text{old}}, \tau^0,z) \\right] \+ \mathbb{E}\_{\substack{\bar{\tau}^0 \sim p\_{\theta^{\text{old}}}, \\\\ z \sim q\_\phi(z|\bar{\tau}^0)}} \\left[ \mathcal{L}^{(2)}\_{\text{DPAIL}}(\theta, \theta^{\text{old}}, \bar{\tau}^0, z) \\right] + \text{MI}\_\phi(\bar{\tau}^0, z)$

Where the last term encourages learning meaningful unsupervised representation $z \sim q\_\phi(z|\bar{\tau})$ by maximizing mutual information.

(Weakness 1) Clarifying the phrase “without policy optimization”

The phrase “without policy optimization” in our paper specifically refers to the absence of reward-driven policy gradient updates, which are commonly used in adversarial imitation learning frameworks. Our proposed method, DPAIL, does not perform RL optimization steps, such as policy gradients of value-based updates that maximize expected rewards.

Instead, we train the diffusion policy using a lower bound objective derived from the discriminator, which does not incorporate any reward signal. As such, the optimization does not involve computing gradients of a reward-based objective with respect to the policy parameters. We will clarify this distinction in the paper.

(Weakness 2 & Question 1) Clarifying the discriminator formulation in Eq. (5) and Eq. (6)

In Eq. (5), we define the structured discriminator $D\_{\pi, {\pi'}}$ as $D\_{\pi, \pi'}(\tau) = {{p_{\pi'} (\tau) }\over{ p\_{\pi'} (\tau) + p\_\pi (\tau)}}$, which matches the form of the optimal discriminator in the inner maximization of eq. (4):

$\max\_D \mathcal{L}(D, p_\pi) = \max_D \mathbb{E}\_{\tau \sim p\_E}[\log D(\tau)] + \mathbb{E}\_{\tau \sim p_\pi} [\log (1 - D(\tau))]$.

At optimality, the derivative of $\mathcal{L}(D, p\_\pi)$ with respect to $D$ should be zero for $\tau$:

$\nabla\_D \mathcal{L}(D, p\_\pi) = \frac{p\_E(\tau) }{ D(\tau) } - \frac{p\_\pi(\tau)}{1-D(\tau)} = 0$.

Rewriting this gives: $\frac{1-D(\tau)}{ D(\tau) }= \frac{p\_\pi(\tau) }{p\_E(\tau)}$.

Therefore, the optimal solution is : $D^{*}\_{\pi, \pi\_E} = \frac{ p\_{\pi\_E(\tau)} }{ p\_{\pi\_E}(\tau) + p\_\pi (\tau) } = \frac{1}{1+ \frac{p\_\pi (\tau)}{p\_{\pi\_E(\tau)}}}=\frac{1}{1+ \exp(-\log\frac{p\_{\pi\_E}(\tau)}{p\_{\pi}(\tau)})} = \sigma \left(\log\frac{p\_{\pi\_E}(\tau)}{p\_{\pi}(\tau)} \right)$,

where $\sigma(\cdot)$ is the sigmoid function.

Thus, our definition in eq. (5) generalizes this form by replacing the expert policy $\pi\_E$ with a variable $\pi’$. We will clarify this in the paper.

(Weakness 3) How to construct multi-modal demonstration

To construct a multi-modal demonstration dataset, we use K sufficiently trained expert policies, each using a distinct and different reward function. For example, in AntGoal-v3, each reward function corresponds to the negative distance to a different target goal position from the current position. We collect demonstrations from all K experts and aggregate them into a single dataset. We consider all collected trajectories as perfect demonstrations. For detailed reward functions, please refer to Appendix B (Environment Details).

(Weakness 1) Motivation

In real-world settings, expert demonstrations are often limited in quantity and exhibit multi-modality. For example, there may be multiple distinct but valid strategies to solve the same task. Diffusion models are particularly well-suited to capture such multi-modal behavior due to their strong generative modeling capabilities. However, when only a few expert demonstrations are available, purely supervised learning on these data can lead to overfitting even for diffusion models. To address this, we bridge diffusion modeling with adversarial learning by leveraging diffusion-generated samples as negative examples. This enables the model to contrast expert data against generated ones, thereby promoting better alignment with expert distribution[1,2] and enhancing robustness[3] and generalization[4,5] of the learned policy.

[1] Self-Play Fine-Tuning Converts Weak Language Models to Strong Language Models, Chen et al.,ICML 2024.

[2] Adversarial Soft Advantage Fitting: Imitation Learning without Policy Optimization, Barde et al., NeurIPS 2020.

[3] On the Convergence and Robustness of Adversarial Training, Wang et al., ICML 2019.

[4] Adversarial Training Helps Transfer Learning via Better Representations, Deng et al., NeurIPS 2021.

[5] Improved OOD Generalization via Adversarial Training and Pre-training, Yi et al., ICML 2021.

(Weakness 2) Necessity of an alternative approach for diffusion policies in AIL

AIL methods typically involve training a discriminator and a generator in a minimax objective, which often leads to training instability. In particular, when the discriminator becomes too accurate, it can cause vanishing gradients for the generator, hindering the learning process. When applying such adversarial training to diffusion policies, these issues are further exacerbated. Specifically, computing policy gradient based on reward signals requires backpropagation through the entire diffusion steps, resulting in substantial computational cost and additional instability [6,7]. To address these challenges, several recent works have been proposed. For instance, [6] introduces a technique that gradually increases the number of trainable diffusion steps, while [7] proposes a reverse sampling method that estimates the mapping from noise to clean data to stabilize training. However, these techniques still require careful scheduling and additional sampling, making the training process complex. In contrast, our proposed method, DPAIL, builds upon the ASAF framework, which allows us to sidestep these difficulties. By eliminating both the explicit discriminator and policy gradient-based optimization, DPAIL enables stable and efficient training of diffusion models, allowing them to be effectively applied to multi-modal generation.

[6] Towards Better Alignment: Training Diffusion Models with Reinforcement Learning Against Sparse Rewards, Hu et al., CVPR, 2025.

[7] Efficient Online Reinforcement Learning for Diffusion Policy, Ma et al., ICML, 2025.

(Weakness 3) Evaluation on challenging manipulation task

We conducted additional experiments on two challenging environments: FrankaKitchen and adroit, which go beyond standard benchmarks and better reflect the complexity of real-world tasks. The FrankaKitchen environment involves goal completion tasks that require the agent to perform household activities, such as opening a microwave or turning on a light switch. For experiments, we collected a mix of 5 trajectories per task for 4 tasks: ‘microwave’, ‘kettle’, ‘light switch’, and ‘kettle’. In addition, we evaluate DPAIL on the adroit pen manipulation task, which requires a dexterous robot hand to control and orient a pen. This task is particularly challenging due to its high-dimensional and contact-rich nature. We used 25 human demonstrations collected from [8]. The normalized returns are shown below:

These results demonstrate that DPAIL consistently outperforms baseline methods, validating its effectiveness in realistic scenarios.

[8] D4RL: Datasets for Deep Data-Driven Reinforcement Learning, Fu et al., arxiv.

(Weakness 1) High computation cost and longer inference time

As noted in the section on limitations, DPAIL incurs longer inference time due to the use of diffusion models as policies. This is a trade-off for their high expressiveness and ability to capture multi-modal behaviors, which are essential for our goal for modeling diverse expert demonstrations. There has been promising progress in accelerating diffusion model inference (e.g., through distillation or fast samplers), and integrating such techniques into our framework is a promising direction for improving its scalability and real-time applicability, which we leave as future work.

(Weakness 2 & Question 1) Scaling factor $N$ inside the sigmoid function

In our experiments, we set $N=50$, and we did not observe any numerical issues or optimization instability resulting from scaling factor $N$. Although $N$ increases the magnitude of the input to the sigmoid function in Eq. (9-10), this did not push it into the flat region of sigmoid in our experiments. This is because DPAIL resets the parameter via $\theta^{\text{old}} \leftarrow \theta$ at every iteration $k$, which causes the difference in noise prediction errors, the input of the sigmoid, to be zero at the start of each iteration. As a result, the input to the sigmoid remains within a manageable range. Specifically, for N = 50, the value typically lies within [0.0~8.0], where the log-sigmoid still provides meaningful gradients.

For larger scaling factors (e.g N $\geq 100$), we investigated whether clipping the input to the log-simoid function would be necessary for the optimization. To assess this, we conducted additional experiments with N=100, both with and without clipping the input (e.g to a maximum value of 9).

In practice, we observed no significant difference in performance between them, likely because the input to the log-sigmoid resets to near-zero at the start of each iteration. However, the clipping may still be beneficial when using very large values of N for the effective optimization.