How Does the Lagrangian Guide Safe Reinforcement Learning through Diffusion Models?

Diffusion policy sampling enables reinforcement learning (RL) to represent multimodal action distributions beyond suboptimal unimodal Gaussian policies. However, existing diffusion-based RL methods primarily focus on offline settings for reward maximization, with limited consideration of safety in online settings. To address this gap, we propose Augmented Lagrangian-Guided Diffusion (ALGD), a novel algorithm for off-policy safe RL. By revisiting optimization theory and energy-based model, we show that the instability of primal-dual methods arises from the non-convex Lagrangian landscape. In diffusion-based safe RL, the Lagrangian can be interpreted as an energy function guiding the denoising dynamics. Counterintuitively, direct usage destabilizes both policy generation and training. ALGD resolves this issue by introducing an augmented Lagrangian that locally convexifies the energy landscape, yielding a stabilized policy generation and training process without altering the distribution of the optimal policy. Theoretical analysis and extensive experiments demonstrate that ALGD is both theoretically grounded and empirically effective, achieving strong and stable performance across diverse environments.

💡 Research Summary

The paper tackles a fundamental limitation of safe reinforcement learning (RL) – the instability of primal‑dual methods when enforcing safety constraints – by marrying safe RL with diffusion‑based policy generation. The authors first reinterpret the Lagrangian of the constrained RL problem as an energy function that governs the reverse‑time diffusion process. In this view, the score function used in diffusion models (the gradient of the log‑density) is directly proportional to the gradient of the Lagrangian with respect to actions. While this connection is elegant, it reveals a critical problem: the raw Lagrangian is typically highly non‑convex and its gradient is noisy because it depends on learned Q‑functions for reward and cost. Moreover, the Lagrange multiplier λ is updated online and can fluctuate dramatically when cost estimates are inaccurate. Consequently, the energy landscape that drives the diffusion becomes unstable, leading to oscillating dual variables, poor policy updates, and frequent safety violations – exactly the issues observed in traditional primal‑dual algorithms.

To resolve this, the authors introduce an augmented Lagrangian formulation. By adding a quadratic penalty term with coefficient ρ>0, the augmented Lagrangian L_A(s,a,λ)=−Q(s,a)+(