Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Xiaoyuan Cheng, Wenxuan Yuan, Boyang Li, Yuanchao Xu, Yiming Yang, Hao Liang, Bei Peng, Robert Loftin, Zhuo Sun, Yukun Hu

TL;DR

This work tackles safe online reinforcement learning with expressive diffusion policies. It reframes safety as an energy-based constraint via a Lagrangian and introduces an augmented Lagrangian $\mathcal{L}_A$ to locally convexify the energy landscape, preserving the optimal Boltzmann policy $\pi^*(a|s) \propto \exp(-\mathcal{L}(s,a,\lambda^*)/\beta)$. The proposed ALGD framework guides the reverse diffusion with an energy-aware score, employing an ensemble of cost critics and Monte Carlo score estimation to achieve stable training and reduced constraint violations. Empirical results on Safety-Gym and velocity-constrained MuJoCo demonstrate competitive rewards with improved safety, and ablations validate the roles of Monte Carlo sampling, critic ensembles, and the convexification strength. Overall, ALGD provides a principled, scalable approach to online safe RL with diffusion-based policies, enabling safer, multimodal action exploration in robotics and autonomous systems.

Abstract

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.

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

TL;DR

This work tackles safe online reinforcement learning with expressive diffusion policies. It reframes safety as an energy-based constraint via a Lagrangian and introduces an augmented Lagrangian to locally convexify the energy landscape, preserving the optimal Boltzmann policy . The proposed ALGD framework guides the reverse diffusion with an energy-aware score, employing an ensemble of cost critics and Monte Carlo score estimation to achieve stable training and reduced constraint violations. Empirical results on Safety-Gym and velocity-constrained MuJoCo demonstrate competitive rewards with improved safety, and ablations validate the roles of Monte Carlo sampling, critic ensembles, and the convexification strength. Overall, ALGD provides a principled, scalable approach to online safe RL with diffusion-based policies, enabling safer, multimodal action exploration in robotics and autonomous systems.

Abstract

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.
Paper Structure (36 sections, 6 theorems, 81 equations, 10 figures, 6 tables, 1 algorithm)

This paper contains 36 sections, 6 theorems, 81 equations, 10 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Problem Formulation.
  4. Method
  5. When Diffusion Meets Safe RL
  6. Locally Convexified Energy Landscape
  7. Practical Algorithm and Discrepancy Analysis
  8. Numerical Experiments
  9. Conclusion
  10. More Literature Review
  11. Limitation of Our Work and Future Direction
  12. Review Limitations of Original Lagrangian Method
  13. Iterative Dual Updates.
  14. Why Directly Using the Lagrangian in Diffusion Models Is Problematic.
  15. Theoretical Analysis
  16. ...and 21 more sections

Key Result

Proposition 3.1

Consider the VE SDE with the reverse-time formulation in eq:langevin_dynamics, the Lagrangian $\mathcal{L}(s,a,\lambda)$ is globally Lipschitz continuous over the entire state-action spaces $\mathcal{S}$ and $\mathcal{A}$. By fixing the Lagrange multiplier $\lambda$ and sampling trajectories $(s_t, with the posterior distribution $a^{0|\tau} \sim \mathcal{N}(a^{\tau}, \sigma^2(\tau) I)$ and the f

Figures (10)

  • Figure 1: Visualization of energy landscapes for a differential-drive mobile robot contreras2017controllers after 100 training episodes based on our methods (see the two algorithms implementation in Appendix \ref{['append:algorithm']}). Top: The landscape of standard Lagrangian induces a highly irregular and non-convex energy surface with sharp curvature, reflecting unstable denoising dynamics. Bottom: The landscape of augmented Lagrangian yields a smoother and locally convexified energy landscape, effectively regularizing the score field. In both contour plots, the white circle indicates the safe policy region, within which policies satisfy the safety constraints.
  • Figure 2: Comparative analysis of training stability between the standard Lagrangian (yellow) and the augmented Lagrangian (blue). (Top) Evaluation rewards; (Middle) Dual variable $\lambda$ updates; (Bottom) Average constraint violation (calculated as $\max(0,c(s)-h)$). The augmented formulation directly addresses the oscillation of dual variables (L2) and the instability of the induced Boltzmann distribution (L3) by regularizing the local curvature of the energy landscape. This results in more stable denoising dynamics and dual updating, allowing the policy to reach the safety threshold significantly faster and with zero dual variables (see full result in Figure \ref{['fig:compare_full_auglag_lag']} in Appendix \ref{['append:full_comparative_analysis']}).
  • Figure 3: Comparison of performance across Safety-Gym and MuJoCo benchmarks. For each benchmark, the first row reports the evolution of test reward versus environment steps, the second row shows the corresponding test safety cost (with the dashed line indicating the cost budget), and the third row presents box plots of the training cost distribution over four equal step intervals. Overall, ALGD achieves competitive rewards while exhibiting improved stability and safety compared to baselines.
  • Figure 4: Ablation Studies of Monte Carlo sample size $N$. Due to the space limitation, see full results in Figure \ref{['fig:ablation_full_1']} in Appendix \ref{['append:more_ablation']}.
  • Figure 5: Task environments used in our experiments. Top row: Safe-Gym manipulation and navigation tasks, including Point Button, Car Button, and Point Push, which require the agent to accomplish goal-oriented behaviors while satisfying safety constraints such as obstacle avoidance and region constraints. Bottom row: Velocity-constrained MuJoCo locomotion tasks, including HalfCheetah, Hopper, Ant, and Humanoid, where agents must learn stable and efficient locomotion policies under explicit velocity limits and safety-related constraints.
  • ...and 5 more figures

Theorems & Definitions (7)

  • Proposition 3.1: Lagrangian-guided score function
  • Theorem 3.2
  • Theorem 3.3
  • Lemma 4.1: Monte Carlo Score Estimation Error
  • Lemma 4.2: Pathwise KL Divergence under Drift Perturbation via Girsanov
  • Corollary 4.3: Stability of score under Lagrangian approximation
  • Remark 4.4