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

arXiv: https://arxiv.org/abs/2602.02924 · HTML: https://mlsanity.com/paper/2602.02924/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for robotics, one of 14 directions in Diffusion models.
- Builds on 43 of its 66 references that are in the map.
- First posted on arXiv 2 Feb 2026.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)
- Deep reinforcement learning methods / [RL with diffusion and flow models](https://mlsanity.com/deeprl/rl-with-diffusion-and-flow-models/)

## Builds on (most cited first, 30 of 43)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [Constrained Policy Optimization](https://mlsanity.com/paper/1705.10528/) (2017)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning](https://mlsanity.com/paper/2208.06193/) (2022)
- [Reward Constrained Policy Optimization](https://mlsanity.com/paper/1805.11074/) (2018)
- [Risk-Constrained Reinforcement Learning with Percentile Risk Criteria](https://mlsanity.com/paper/1512.01629/) (2015)
- [A Lyapunov-based Approach to Safe Reinforcement Learning](https://mlsanity.com/paper/1805.07708/) (2018)
- [Diffusion Policy Policy Optimization](https://mlsanity.com/paper/2409.00588/) (2024)
- [How to Train Your Robot with Deep Reinforcement Learning; Lessons We've Learned](https://mlsanity.com/paper/2102.02915/) (2021)
- [Flow Q-Learning](https://mlsanity.com/paper/2502.02538/) (2025)
- [Efficient Diffusion Policies for Offline Reinforcement Learning](https://mlsanity.com/paper/2305.20081/) (2023)
- [Safety-Gymnasium: A Unified Safe Reinforcement Learning Benchmark](https://mlsanity.com/paper/2310.12567/) (2023)
- [Learning a Diffusion Model Policy from Rewards via Q-Score Matching](https://mlsanity.com/paper/2312.11752/) (2023)
- [Iterated Denoising Energy Matching for Sampling from Boltzmann Densities](https://mlsanity.com/paper/2402.06121/) (2024)
- [Diffusion-based Reinforcement Learning via Q-weighted Variational Policy Optimization](https://mlsanity.com/paper/2405.16173/) (2024)
- [Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review](https://mlsanity.com/paper/2407.13734/) (2024)
- [Diffusion Actor-Critic with Entropy Regulator](https://mlsanity.com/paper/2405.15177/) (2024)
- [SafeDiffuser: Safe Planning with Diffusion Probabilistic Models](https://mlsanity.com/paper/2306.00148/) (2023)
- [Constrained Variational Policy Optimization for Safe Reinforcement Learning](https://mlsanity.com/paper/2201.11927/) (2022)
- [Efficient Online Reinforcement Learning for Diffusion Policy](https://mlsanity.com/paper/2502.00361/) (2025)
- [Safe Offline Reinforcement Learning with Feasibility-Guided Diffusion Model](https://mlsanity.com/paper/2401.10700/) (2024)
- [Reachability Constrained Reinforcement Learning](https://mlsanity.com/paper/2205.07536/) (2022)
- [Constrained Update Projection Approach to Safe Policy Optimization](https://mlsanity.com/paper/2209.07089/) (2022)
- [Safe Reinforcement Learning Using Advantage-Based Intervention](https://mlsanity.com/paper/2106.09110/) (2021)
- [Iterative Reachability Estimation for Safe Reinforcement Learning](https://mlsanity.com/paper/2309.13528/) (2023)

## Built on by (most cited first, 5 of 5)

- [Scaling World-Model Reinforcement Learning Through Diffusion Policy Optimization](https://mlsanity.com/paper/2605.26282/) (2026)
- [Constrained Flow Policy Updates: A Generalized Schr"odinger Bridge View](https://mlsanity.com/paper/2609.32952/) (2026)
- [Safe Score Matching: Diffusion Policies with Hamilton-Jacobi Reachability for Online Safe Reinforcement Learning](https://mlsanity.com/paper/2609.33337/) (2026)
- [Brain-SAD: A Brain-Inspired Safe Autonomous Driving Control Framework with Dynamic Fear-Oriented Constraint on Dual-Policy](https://mlsanity.com/paper/2609.38016/) (2026)
- [Reachability-Aware Diffusion Policy Optimization](https://mlsanity.com/paper/2610.05969/) (2026)

## Abstract (arXiv)

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.