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

TL;DR

This paper introduces ALGD, a novel off-policy safe RL algorithm using diffusion models and an augmented Lagrangian to stabilize policy training and ensure safety.

Contribution

It proposes ALGD, which stabilizes diffusion-based safe RL by convexifying the Lagrangian energy landscape, improving training stability and safety guarantees.

Findings

ALGD achieves stable and strong performance across diverse environments.

The augmented Lagrangian approach stabilizes policy generation and training.

Theoretical analysis confirms the effectiveness of the convexification method.

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.