Augmented Lagrangian Multiplier Network for State-wise Safety in Reinforcement Learning

2026-05-01

AI Summary

This work addresses training instability in safe reinforcement learning under state-dependent safety constraints, which arises from dual gradient ascent-induced oscillations. To mitigate this issue, the paper proposes the ALaM framework, which introduces, for the first time, a stable training mechanism for state-dependent Lagrange multiplier networks. The approach incorporates an augmented Lagrangian quadratic penalty term to alleviate update delays and enhance local convexity, while employing a supervised regression objective to optimize the multiplier network and improve convergence. Theoretical analysis guarantees convergence of the multipliers and recovery of the optimal policy. Implemented on top of SAC as SAC-ALaM, the method significantly outperforms existing safe RL algorithms across multiple benchmarks, achieving higher safety rates and superior returns while enabling stable training and learning well-calibrated, risk-sensitive multipliers.

Abstract

Safety is a primary challenge in real-world reinforcement learning (RL). Formulating safety requirements as state-wise constraints has become a prominent paradigm. Handling state-wise constraints with the Lagrangian method requires a distinct multiplier for every state, necessitating neural networks to approximate them as a multiplier network. However, applying standard dual gradient ascent to multiplier networks induces severe training oscillations. This is because the inherent instability of dual ascent is exacerbated by network generalization -- local overshoots and delayed updates propagate to adjacent states, further amplifying policy fluctuations. Existing stabilization techniques are designed for scalar multipliers, which are inadequate for state-dependent multiplier networks. To address this challenge, we propose an augmented Lagrangian multiplier network (ALaM) framework for stable learning of state-wise multipliers. ALaM consists of two key components. First, a quadratic penalty is introduced into the augmented Lagrangian to compensate for delayed multiplier updates and establish the local convexity near the optimum, thereby mitigating policy oscillations. Second, the multiplier network is trained via supervised regression toward a dual target, which stabilizes training and promotes convergence. Theoretically, we show that ALaM guarantees multiplier convergence and thus recovers the optimal policy of the constrained problem. Building on this framework, we integrate soft actor-critic (SAC) with ALaM to develop the SAC-ALaM algorithm. Experiments demonstrate that SAC-ALaM outperforms state-of-the-art safe RL baselines in both safety and return, while also stabilizing training dynamics and learning well-calibrated multipliers for risk identification.

Problem

• state-wise constraints
• Lagrangian multiplier network
• training oscillations
• safe reinforcement learning
• multiplier instability

Innovation

• Augmented Lagrangian
• Multiplier Network
• State-wise Constraints
• Safe Reinforcement Learning
• Stable Training