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 setting 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 modeling, 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; counter-intuitively, 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, without altering the distribution of 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.

Modern reinforcement learning methods allow an artificial agent to learn how to make decisions through interaction with its environment. In many real-world applications, however, an agent should not only achieve high rewards, but also avoid unsafe or undesirable behavior during learning. This is challenging because the agent often needs to explore different possible actions, and some of them may violate safety requirements. This paper proposes a new method called Augmented Lagrangian-Guided Diffusion for safer reinforcement learning. The key idea is to use a flexible action-generation model that can represent multiple possible good actions, rather than forcing the agent to choose actions from an overly simple distribution. At the same time, our method incorporates safety constraints directly into the action-generation process. A main difficulty in safe reinforcement learning is that common ways of balancing reward and safety can make training unstable. We show that this instability comes from the shape of the objective being used to guide learning. Our method modifies this guidance so that learning becomes more stable, while preserving the same desired optimal behavior. Through theoretical analysis and experiments on a range of tasks, we show that the proposed method learns policies that are both effective and safer than existing alternatives. These results suggest that the method can be a useful step toward reinforcement learning systems that perform well while respecting safety requirements.