Lagrangian Guided Safe Reinforcement Learning through Diffusion Models
Representative research papers from the Dynamic Systems Lab
Ensuring safety in reinforcement learning is challenging, especially in online settings where exploration is inherently risky and constraint violations can have severe consequences. While primal–dual methods provide a principled way to enforce safety constraints, they often suffer from severe instability due to oscillating dual variables and inaccurate cost estimation. At the same time, diffusion-based policies offer expressive multi-modal action distributions, but existing approaches are largely restricted to offline settings and rarely address safety during online interaction.
To bridge this gap, we propose Augmented Lagrangian-Guided Diffusion (ALGD), a novel framework that unifies safe reinforcement learning with diffusion-based policy generation. By revisiting constrained optimization from an energy-based perspective, we interpret the Lagrangian as the energy function governing the reverse diffusion process. We show that directly using the standard Lagrangian induces a highly non-convex energy landscape, leading to unstable denoising dynamics and unreliable policy sampling. ALGD resolves this issue by introducing an Augmented Lagrangian, which locally convexifies the energy landscape and stabilizes both policy generation and primal–dual training, without altering the optimal policy distribution.
Building on this insight, ALGD enables stable off-policy and online safe reinforcement learning with expressive diffusion policies. Extensive experiments on Safety-Gym and MuJoCo benchmarks demonstrate that ALGD achieves competitive returns while consistently reducing constraint violations and improving training stability compared to existing primal–dual and hard-constrained baselines.
Find out more information about this publication here.
