Adversarial Diffusion for Robust Reinforcement Learning

We thank the reviewer for their detailed and thoughtful evaluation of our work. We are pleased that the reviewer recognized the novel integration of adversarial guidance into diffusion-based planning for robust reinforcement learning and appreciated our method’s ability to improve both robustness and performance across multiple tasks. We are also grateful for the reviewer’s recognition of the theoretical grounding of our approach and the broader potential of diffusion models in robust sequential decision-making. At the same time, we acknowledge and value the reviewer’s detailed comments, and we now provide responses to each point below.

---

Q1. The proposed AD-RRL method is much slower than the PPO and TRPO baselines. How can it be improved using a more efficient diffusion model?

Thank you for raising this important practical point. Our method is approximately $1.5\times$ more computationally expensive than PolyGRAD, which does not include adversarial guidance (note that AD-RRL requires an additional model).

• As a model-based method, higher computational cost compared to model-free baselines is expected. However, we believe this overhead is justified by the significant gains in robustness and nominal performance, as demonstrated in Figures 1–3 and Table 1. AD-RRL consistently outperforms all baselines — including model-free and model-based approaches — in both nominal and perturbed environments, illustrating a favorable trade-off between computational resources and robustness.
• Importantly, the proposed adversarial conditioning method is architecture-agnostic and can be applied on top of more efficient diffusion models. This means that AD-RRL can directly benefit from advances in efficient generative modeling, such as latent diffusion or EDM (as mentioned by the reviewer).

We agree that reducing computational overhead is an important avenue for future research and consider integrating such efficient architectures a promising direction.

---

Q2. How does AD-RRL compare with CVaR-PPO in terms of robustness and computational efficiency?

We thank the reviewer for suggesting this additional baseline. For a fair comparison, we trained CVaR-PPO (CPPO) for 1.5 M steps — consistent with AD-RRL and the other baselines — and present point-wise evaluations below.

Table 1 — Performance across mass-relative changes for HalfCheetah (best results in bold).

Table 2 — Performance across mass-relative changes for Hopper (best results in bold).

Table 3 — Performance across mass-relative changes for Walker (best results in bold).

In terms of efficiency, one CPPO run takes roughly 30 minutes, compared to slightly under 4 days for AD-RRL. This is expected, since CPPO is a model-free method that does not rely on learned dynamics or trajectory guidance. However, as the tables show, AD-RRL vastly outperforms CPPO and is significantly more robust to variations in environment parameters. Its performance slowly degrades as the mass parameter shifts, whereas CPPO’s performance declines sharply. We believe this greatly increased robustness justifies the additional overhead, and again highlight the future potential of incorporating more efficient architectures into AD-RRL.

---

Q3. The proposed adversarial diffusion guidance is based on a smooth approximation of the classifier. How tight is this approximation? What are the limitations of this approach?

Thank you for this interesting question.

Our approach builds on prior work (e.g., [1]), which casts control as probabilistic inference. There, the probability of optimality at each time step is approximated as $P(O_t = 1 \mid s_t, a_t) = \exp(r(s_t, a_t)),$ a smooth relaxation of the hard indicator $P(O_t = 1 \mid s_t, a_t) = \mathbf{1}_{(s_t, a_t) \in C}$.

• In our setting, a diffusion model generates entire trajectories rather than individual actions. We therefore model the optimality of an entire trajectory $\tau$:

  • Let $O(\tau)$ be a binary variable indicating whether a trajectory is optimal. We approximate $P \bigl(O(\tau)=1 \mid \tau\bigr) = \exp\bigl(-c_0 \sum_t \gamma^t r_t\bigr),$ where $c_0$ controls the sharpness of the approximation.
  • This smooth form allows gradient-based guidance to steer the diffusion model toward low-return (i.e., adversarial) trajectories.
  • Without such an approximation, adversarial guidance would be intractable in diffusion-based planning.

• The tightness of this approximation is governed by the scalar $c_0$:

  • Larger values of $c_0$ make the approximation closer to a hard threshold; smaller values produce smoother weighting.
  • Similarly, at each diffusion step $i$, a corresponding $c_i$ governs the sharpness of the step-specific approximation.

• In practice, we observe that $c_i$ increases as $i$ decreases (i.e., as the model approaches the final, denoised trajectory $\tau_0$), making the approximation sharper at later steps. This effect arises from Eq. (19) and from the cosine noise schedule of the diffusion model [2], which reduces variance as $i$ decreases — thus increasing $c_i$.

• A limitation is that this smooth approximation may not sharply distinguish trajectories that lie near the CVaR threshold. Nevertheless, it provides a practical and differentiable surrogate for otherwise intractable objectives.

We acknowledge that exploring tighter or alternative approximations, and empirically validating their effects, would be a valuable direction for future work and will mention this in the final version.

---

Q4. AD-RRL outperforms PolyGRAD, PPO, and other baselines across several MuJoCo tasks, even in the training (non-adversarial) environments.

Indeed, while AD-RRL is designed for robustness, it also performs strongly in the nominal setting. We believe this stems from two main factors:

• First, the CVaR objective lower-bounds the expected return. Maximizing CVaR thus implicitly promotes strong average performance.

• Second, adversarial guidance encourages the policy to learn from failure modes that might be underexplored in other approaches. This leads to more informative training signals and improved generalization, even under nominal dynamics.

--

We hope that, given the detailed comparison with CVaR-PPO and our in-depth answers to the reviewer's questions, they might consider raising the final score for our paper.

--

[1] Levine, Sergey. "Reinforcement learning and control as probabilistic inference: Tutorial and review." arXiv preprint arXiv:1805.00909 (2018).

[2] Nichol, Alexander Quinn, and Prafulla Dhariwal. "Improved denoising diffusion probabilistic models." International conference on machine learning. PMLR, 2021.

We thank the reviewer for the helpful comments and valuable questions. We’re pleased to hear that the paper was found to be well-motivated, with an approach that is both intuitive and reasonable. Below, we address each comment in turn.

---

Q1. Could the work be compared and contrasted with [1] that generates trajectories with worst kernel with theoretical guarantees?

Thank you for pointing to this relevant work. The key idea of EWoK presented in [1] is to bias transitions toward next states with lower estimated value, thereby encouraging the policy to learn from pessimistic outcomes.

However, there are a number of distinctions to highlight:

1. EWoK assumes access to a generative model, allowing multiple samples to be drawn from any $(s,a)$ pair. AD-RRL does not.

   • This assumption is often impractical in real-world environments where the agent cannot reset and resample transitions on demand.
   • In contrast, AD-RRL learns a dynamics model and uses it to generate entire trajectories into the future, enabling adversarial sampling without requiring additional environment access during training.

2. While EWoK can offer strong theoretical guarantees due to its assumption of a known transition function, AD-RRL does not make such assumptions. We believe this leads to a more realistic framework for robust RL in environments with unknown or partially observable dynamics.

We agree that comparing AD-RRL with EWoK would be valuable for future work, and we will strive to include such a comparison in the final version of the manuscript.

---

Q2. The paper is well motivated and the approach looks good. Can we say anything about theoretical performance of the approach?

Thank you for the positive feedback and for this question.

• In the paper, we explore how the dual CVaR formulation connects to robustness. In particular, Proposition 4.3 formally describes how to select the scalar $c_i$ such that the adversarial guidance at each diffusion step produces trajectories within the CVaR risk envelope.
• Choosing $c_i$ in this way ensures that the perturbations introduced by the adversarial objective remain valid from a risk-sensitive optimization perspective, and provides a principled foundation for the guided sampling process. Note that $c_i$ controls the sharpness of the approximation to the indicator function $p_\theta(\bar O(\tau_i) = 1 \mid \tau_i) = 1_{\tau_i \in C_\alpha}$ used in the CVaR formulation (see, e.g., Eq. 12 and 13).
  • Larger values of $c_i$ make the function sharper (closer to a hard threshold), while smaller values smooth the weighting more broadly.
• In practice, the values of $c_i$ increase as $i$ decreases. This means that as the model approaches the final denoised trajectory $\tau_0$, the guidance becomes more selective, making the approximation sharper and closer to a hard threshold.
  • This trend results from the definition of $c_i$ (see Eq. 19) and the diffusion model’s cosine schedule.
  • We use the cosine schedule proposed by [3], which decreases the variance of the diffusion model as $i$ decreases. This reduces the denominator in Eq. 19 and thus increases $c_i$ as $i$ decreases.

While we acknowledge that exploring other theoretical properties—such as regret bounds for AD-RRL—is an exciting direction, we consider this beyond the scope of the current paper.

---

Q3. Could you please summarize the key challenges of the paper?

Thank you for the question. The key challenges we address in this work are:

• Learning robust policies in uncertain environments: Robust RL aims to prepare agents for rare or worst-case scenarios, but identifying and training on such situations is non-trivial—especially when they are not encountered during standard exploration.
• Generating informative, low-return trajectories: Our challenge is to guide a diffusion model to generate low-return (adversarial) trajectories from the tail of the return distribution (i.e., within a CVaR risk envelope), while ensuring they remain realistic and actionable.
• Providing theoretical guarantees on adversarial guidance: We derive and prove conditions (Proposition 4.3) under which the gradient guidance coefficients $c_i$ ensure that the perturbed samples remain within the CVaR risk envelope. This offers a principled foundation for adversarial sampling throughout the diffusion process.
• Maintaining model fidelity while increasing robustness: Training on adversarial samples can cause a mismatch with the data distribution. A key challenge is to ensure that the generated trajectories remain likely under the current policy, which we address by leveraging PolyGRAD’s policy-conditioned diffusion model.

---

Q4. What are the limitations of the work?

As also discussed in Appendix G of our paper, we acknowledge the following key limitations:

• Computational overhead: While AD-RRL significantly improves robustness and performance, it does so at the cost of increased computational overhead, due to both the learned diffusion model and the adversarial guidance mechanism. As discussed in our response to Reviewer TZ4u, we believe this cost is justified by the observed robustness and nominal performance gains, and could be mitigated in future work by incorporating more efficient generative architectures.
• Smoothness assumptions: Our derivation employs a Gaussian approximation and requires computing gradients $\nabla_{\tau}Z(\tau)$, which presupposes reasonably smooth rewards and state transitions. Tasks with non-smooth dynamics may violate this assumption, leading to less accurate guidance. Extending adversarial diffusion to such domains is an important direction for future work.
• Benchmark: Finally, our experiments are currently limited to continuous control tasks in MuJoCo. While these are standard benchmarks, evaluating AD-RRL in higher-dimensional settings—such as vision-based environments (e.g., Atari)—would further clarify the method’s scalability.

---

Q5. There are many works on uncertainty sets bounded by $L_p$ balls, revealing many structural properties such as “adversary is rank-one perturbation of the nominal kernel” [2], etc. Can these properties be used to design better diffusion-based algorithms for robust RL?

Thank you for the insightful question. Indeed, many works in robust MDPs (e.g., [2]) define uncertainty sets using $L_p$-norm balls around nominal transition kernels, which leads to useful structural properties such as rectangularity, convexity, and tractable dual forms.

Our method generates samples at the trajectory level using a learned diffusion model, which makes it difficult to directly apply kernel-level uncertainty set formulations.

That said, these structural insights could certainly inspire future improvements. For instance, incorporating constraints or priors on the learned dynamics model (e.g., Lipschitz continuity or rectangularity) may lead to better-calibrated adversarial sampling and stronger theoretical guarantees.

---

Q6. Clarity of presentation

Thank you for suggesting a concrete way to improve the clarity of our paper. In the final version, we will include an additional illustration to provide a simplified overview of the AD-RRL pipeline.

--

We hope that our detailed comparison with EWoK, along with the comprehensive responses to the reviewer’s questions, helps clarify the theoretical contributions, the challenges (both present and future), and the strengths of our work. We the reviewer might consider raising the final score when assessing their final evaluation.

--

[1] Gadot, Uri, et al. "Bring your own (non-robust) algorithm to solve robust MDPs by estimating the worst kernel." Forty-first International Conference on Machine Learning. 2024.

[2] Kumar, Navdeep, et al. "Policy gradient for rectangular robust markov decision processes." Advances in Neural Information Processing Systems 36 (2023): 59477-59501.

[3] Nichol, Alexander Quinn, and Prafulla Dhariwal. "Improved denoising diffusion probabilistic models." International conference on machine learning. PMLR, 2021.

We thank the reviewer for the interesting questions and for the time dedicated to reading and evaluating our paper.

• We are pleased to see that the reviewer found the paper well-structured, clearly motivated, and enjoyable to read, and appreciated the significance and theoretical soundness of our contribution.
• We are especially grateful for the positive assessment of both the clarity of our mathematical formulation and the strength of our experimental results, as well as the overall recommendation that the paper exceeds the acceptance threshold for NeurIPS.

Below, we address their questions in detail.

---

Q1. By design, diffusion models aim to replicate the historical data distribution within the dataset. Given this, how is it possible for the model to generate “trajectories that are either rare in the current environment or originate from unexplored regions of the domain,” as described in line 51 of the paper?

We thank the reviewer for highlighting this point. While diffusion models are indeed trained to replicate the distribution of historical trajectories, our approach modifies the sampling process via CVaR-based adversarial guidance, which biases the model toward trajectories associated with low return (i.e., areas where the agent tends to fail). Among low-return trajectories, we expect some to come from under-explored regions of the environment, since a low cumulative return may indicate that the agent has not learned how to behave optimally in these regions.

We agree that our original phrasing may have been ambiguous and will revise the sentence in line 51 to make this distinction clearer in the final version.

---

Q2. I can understand the approximation that $p_\theta(\tau_0 \in C_\alpha \mid \tau_0) \approx \exp(-c_0 \sum \gamma^t r_t)$. Could the authors provide more intuition for why $p_\theta(\tau_i \in C_\alpha \mid \tau_i) \approx \exp(-c_i \sum \gamma^t r_t^{(i)})$, and how to obtain $r_t^{(i)}$ for each $i$?

Thank you for the question; we are happy to clarify.

The reason we need this approximation at each step is as follows:

• The goal is to be able to sample from $\bar p_{ \theta}(\tau_0)$. From eq. (16), and equation below line 201, we have that $\bar p_{ \theta_0}(\tau_0)\propto p_{\theta}(\tau_0)p_\theta(\tau_0\in C_\alpha|\tau_0)$.

  • However, multiplying the distribution by the second term (a conditional probability) is usually costly and difficult. For diffusion models, the second term can be either treated as a small perturbation in each step in the diffusion process, or (often) obtained adding a multiplicative factor to the distribution in each diffusion step.
  • Our approach consists in adding a multiplicative factor in each diffusion step (see also proof of eq. 16). Therefore, we use the same kind of approximation $p_\theta(\tau_i\in C_\alpha|\tau_i)\approx exp(-c_i\sum\gamma^t r_t^{(i)})$, so that we can exploit its smoothness for guiding the model toward adversarial rollouts.
  • The intuition is that this form enables us to reweight trajectories at every diffusion step based on their cumulative return.

• How we get $r_t^{(i)}$: note that for every generated trajectory $\tau_i$ we predict the cumulative reward $\sum_t \gamma^tr_t^{(i)}$ using a separate network, that takes as input $\tau_i$ and the diffusion step $i$. The implementation of the guidance procedure is described in section 5 of our paper (see, e.g., lines 270/271 onwards). More details on the implementation are presented in Appendix E.

---

Q3. In Algorithm 2, $c_i$ is computed with Eq. (18), but Eq. (18) is an inequality. As stated in line 243, should $c_i$ instead be calculated using Eq. (19) and set as an equality?

Thank you for pointing this out. You are correct — in Algorithm 2, $c_i$ should be computed using Eq. (19), where the value is set using an equality. However, we would like to note that using Eq. (18) would not change the computed value of $c_i$ in this setting, as the first term in the $\min$ is always greater than the second term, as shown in Proposition 4.3.

We will fix this in the final version by using Eq. (19) with equality to compute $c_i$ in Algorithm 2.

---

Q4. Minor issues and notation fixes

We thank the reviewer for carefully reading the paper and noting these points. We will correct the notation and writing issues in the final version.

---

We hope that, given our detailed answers to the reviewer's questions, they might consider raising the final score for our paper.

We thank the reviewer for their feedback and for the time spent evaluating our paper. We are glad to see that the reviewer found the paper well-written and easy to follow, and appreciated the integration of diffusion-based trajectory modeling with adversarial robustness as a compelling and timely contribution to the field. We are also pleased that the reviewer highlighted the strength of our empirical results, noting that our method matches or outperforms strong baselines, particularly in noisy or perturbed test environments, thereby demonstrating the robustness benefits of our approach.

Below, we provide detailed responses to their comments and questions.

---

Q1. The proposed method may tend to sample trajectories that are rare or insufficiently explored during training. This raises concerns about increased modeling errors, as the diffusion model may not have been adequately trained on such trajectories.

We appreciate the concern raised by the reviewer, and are happy to clarify this point.

Note that our CVaR-based sampling process simply biases the diffusion process toward trajectories with low return, which do not need to be rare in general.

• While diffusion models are typically trained to replicate the distribution of historical trajectories, our approach modifies the sampling process through CVaR-based adversarial guidance, biasing the model toward low-return trajectories, i.e., regions where the agent tends to fail. Among these low-return trajectories, we expect some to originate from under-explored areas of the environment, as poor cumulative return may suggest the agent has not yet learned optimal behavior in those regions. However, low returns can also arise for other reasons, such as stochastic dynamics or a sub-optimal current policy.
• Regarding model accuracy, the training curves in the appendix show that the agent’s performance on the nominal environment continuously improves throughout training, suggesting small modeling errors.
  • Importantly, our architecture builds on PolyGRAD, conditioning the diffusion process on the policy gradient. Thus, the sampled trajectories are not arbitrary outliers but are likely under the current policy, ensuring they remain realistic.
  • The combination of adversarial guidance (AD-RRL) and policy-informed conditioning (PolyGRAD) steers sampling toward challenging yet plausible trajectories, balancing robustness and model fidelity.

We will clarify these points in the final version of the manuscript.

---

Q2. The experiments are limited to relatively small-scale locomotion tasks. Evaluating the method’s scalability to more complex, high-dimensional environments would significantly strengthen the paper.

We understand the reviewer’s concern. Our current experiments focus on well-established continuous-control benchmarks that allow clear comparison to prior work (see for example [1,2,3,4], a small selection of papers using these same benchmarks), highlighting the advantage of our adversarial guidance mechanism. We agree that evaluating scalability to more complex, high-dimensional environments (e.g., vision-based settings) is a valuable future direction.

Furthermore, our approach is modular: the adversarial guidance does not depend on domain-specific structure, suggesting that it should generalize well to more complex tasks. We aim to explore such extensions in future work.

---

Q3. Training guided diffusion models and computing return gradients at every diffusion step can be computationally expensive.

This is an important practical point, but note that our methods is approximately $1.5\times$ more computationally expensive than PolyGRAD, which does not include adversarial guidance (note that compared to PolyGrad our method requires an additional model).

• Being a model-based approach, it is natural to expect it to be more expensive than model-free approaches. However, we believe this overhead is justified by the significant gains in robustness and nominal performance, as demonstrated in Figures 1–3 and Table 1.
• AD-RRL consistently outperforms all baselines, including model-free and model-based alternatives, on both nominal environments and unobserved test-time conditions. This highlights a positive tradeoff between computational resources and robustness.

We agree that reducing this computational overhead is an important direction for future work.

---

Q4. The choice of the approximation formula seems a bit arbitrary. Can the authors clarify why that specific formula was used? For example, what happens if you remove the exponential term?

We would like to highlight that we are not the first to employ such ideas. In fact, such approximation is also used in other relevant prior works (e.g., [5,6]) focusing on RL for continuous tasks.

The main idea borrows from [5], which casts the control problem as a probabilistic inference problem. Practically speaking, they introduce a binary random variable $O_t$ indicating whether timestep $t$ is optimal, and parameterize the probability of $O_t$ as $P(O_t=1|s_t,a_t)=\exp(r(s_t,a_t))$. This is a smooth approximation of the indicator function $P(O_t=1|s_t,a_t)={\bf 1}_{(s_t,a_t)\in C}$, where $C$ is the set of optimal state-action pairs.

• While such approximation may seem arbitrary, in fact it leads to a very natural posterior distribution $P(\tau|O_1,\dots,O_T)\propto P(\tau)\exp(\sum_t r(s_t,a_t))$. That is, the probability of observing a trajectory (given its optimality) is given by the product between its probability to occur and the exponential of the total reward along that trajectory.
• In AD-RRL we follow this approach, but need to adapt it to our setting: a diffusion process does not generate one state-action, but an entire trajectory at the time. For this reason, we do not model the optimality of a time-step, but rather the optimality of an entire trajectory $\tau$:
  • We introduce a binary random variable $O(\tau)$ indicating wether $\tau$ is optimal, and parametrize such random variable by $P(O(\tau)=1|\tau)=\exp(-c_0\sum_t \gamma^t r_t)$ (we consider the discounted return since the optimality of a trajectory is based on its total discounted return).
  • Since this is a smooth parametrization, it enables gradient-based guidance in the diffusion process that steers sampling toward low-return (adversarial) trajectories.
  • Without such an approximation, the adversarial objective would be intractable for diffusion-based planning.

---

Q5. To determine if guiding the diffusion model to generate adversarial trajectories is necessary, it would be helpful to compare the proposed method with simply introducing random noise to the diffusion model (similar to DR-PPO).

This is a great point, we thank the reviewer for the suggestion. We believe that just adding random noise to the diffusion model would not yield improvements in performance or robustness, since diffusion models inherently inject Gaussian noise at every denoising step.

• As an alternative, for this rebuttal, we implemented a Domain Randomization variant of PolyGRAD (which we refer to as DR-PolyGRAD), where the agent is trained to maximize expected return over a uniform distribution of randomized environment dynamics, in line with the DR-PPO formulation.
• Note that Domain Randomization methods assume access to the environment simulator to randomize physical parameters. This assumption is often impractical in many real-world settings where the agent does not have access to a simulator. In contrast, AD-RRL learns a dynamics model and uses it to generate entire trajectories into the future, enabling adversarial sampling without requiring additional environment access during training.
• Due to time and resource constraints, we were able to compare the algorithms only on one environment (Hopper). Since we are not allowed to share plots over anonymized links, we provide point-wise evaluations of the two methods in the following table. We can see that AD-RRL still outperforms DR-PolyGRAD in robustness to unobserved test conditions. AD-RRL also achieves better performance than DR-PolyGRAD when the mass is close to its nominal value.

Table 1 — Performance across mass-relative changes for Hopper (best results in bold).

---

We hope that these answers address the reviewer’s concerns and that, given the additional comparison with DR-PolyGRAD, they will consider raising their final score for our paper.

---

[1] Pinto, Lerrel, et al. "Robust adversarial reinforcement learning." International conference on machine learning. PMLR, 2017.

[2] Ying, Chengyang, et al. "Towards safe reinforcement learning via constraining conditional value-at-risk." arXiv preprint arXiv:2206.04436 (2022).

[3] Rigter, Marc, Jun Yamada, and Ingmar Posner. "World models via policy-guided trajectory diffusion." arXiv preprint arXiv:2312.08533 (2023).

[4] Zeng, Siliang, et al. "Maximum-likelihood inverse reinforcement learning with finite-time guarantees." Advances in Neural Information Processing Systems 35 (2022): 10122-10135.

[5] Levine, Sergey. "Reinforcement learning and control as probabilistic inference: Tutorial and review." arXiv preprint arXiv:1805.00909 (2018).

[6] Janner, Michael, et al. "Planning with diffusion for flexible behavior synthesis." arXiv preprint arXiv:2205.09991 (2022).

We thank the reviewer for their positive feedback and insightful questions. In our rebuttal we have:

1. Clarified the trajectory-generation mechanism, showing how CVaR guidance steers sampling toward low-return yet policy-plausible trajectories, thereby limiting modelling error.
2. Explained our focus on standard continuous-control benchmarks, which enable direct comparison with prior work. Because the guidance mechanism is modular, we expect it to generalize readily to higher-dimensional tasks.
3. Discussed the trade-off between computational cost and performance: although AD-RRL is more expensive than some baselines, it consistently outperforms them in both nominal and perturbed conditions, which we believe to be a worth trade-off.
4. Clarified the exponential approximation and why a smooth, differentiable form is essential for gradient-based guidance.
5. Implemented the DR-PolyGRAD baseline you suggested; AD-RRL still outperforms it across mass variations in Hopper.

If any additional clarification or experiment would be helpful before the deadline, we are happy to provide it. Otherwise, we hope these additions address your concerns and will be reflected in your final assessment.

Thank you for your time and consideration.