How Does the Lagrangian Guide Safe Reinforcement Learning through Diffusion Models?
This work addresses safety in online reinforcement learning for domains requiring stable and safe decision-making, representing an incremental improvement over existing diffusion-based methods.
The paper tackles the instability of diffusion-based safe reinforcement learning in online settings by proposing ALGD, which uses an augmented Lagrangian to convexify the energy landscape, resulting in stable policy generation and training without changing the optimal policy distribution.
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.