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

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.

• The paper presents a novel augmented Lagrangian multiplier network that enforces state-wise safety constraints in reinforcement learning.
• It employs supervised regression for dual updates and adaptive penalty scaling to reduce training oscillations and improve convergence.
• Empirical results on continuous control tasks demonstrate that SAC-ALaM robustly balances safety and performance, offering interpretable risk calibration.

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

Introduction and Problem Motivation

Safe RL requires satisfaction of safety constraints throughout an agent's trajectory, especially in real-world applications such as autonomous driving and robotics. The classical CMDP approach introduces safety as a single global constraint on cumulative costs, which is inadequate for enforcing local state-wise safety; this leads to agents potentially violating constraints in individual states while still satisfying global requirements.

Recent advances propose state-wise constraints, formulating the problem as an infinite set of per-state safety constraints. The Lagrangian relaxation method, ubiquitous in constrained RL, addresses these constraints through dual variables (multipliers). However, a single global multiplier cannot enforce state-wise feasibility. The practical solution is to parameterize multipliers as a neural network (multiplier network), mapping states to penalties. Naive dual gradient ascent, however, induces severe training oscillations due to delayed dual response, network generalization artifacts, and distribution shift during policy updates, particularly in the context of nonlinear function approximation.

Methodology: Augmented Lagrangian Multiplier Network (ALaM)

Augmented Lagrangian Formulation

ALaM introduces a quadratic penalty into the Lagrangian to remedy the inherent instability of standard dual updates. Specifically, the augmented Lagrangian for policy $\pi$ and statewise multiplier $\lambda$ is

$$L(\pi, \lambda, \rho) = -J(\pi) + \left\langle \lambda, \max\left\{-\frac{\lambda}{\rho}, F^\pi \right\} \right\rangle + \frac{\rho}{2} \left\| \max\left\{-\frac{\lambda}{\rho}, F^\pi \right\} \right\|^2$$

where $F^\pi$ is a feasibility function capturing state-wise constraint violation, and $\rho$ is the (adaptive) penalty parameter. This quadratic penalty directly compensates for dual delays and establishes local convexity, thereby ameliorating the ill-conditioned landscape of the minimax optimization and limiting policy divergence.

Multiplier Network Training via Supervised Regression

Unlike prior approaches that apply parameter gradient ascent to multiplier networks (which do not correspond to functional dual ascent), ALaM computes an analytic dual target for each state:

$$\lambda_\text{target}(x) = \max\{\lambda(x) + \rho F^\pi(x), 0\}$$

and directly trains the multiplier network parameters via supervised regression to this target:

$$L_\lambda = \mathbb{E}_x \left[(\lambda(x) - \lambda_\text{target}(x))^2\right]$$

This mechanism ensures that the network output closely tracks the theoretically correct dual variable, leading to more consistent convergence and improved generalization in risk calibration.

(Figure 1)

Figure 1: Comparison of normalized performance across 8 tasks. The data points and error bars denote the mean and 95% confidence intervals for episode returns and costs.

Adaptive Penalty Scaling

The penalty parameter $\rho$ is updated on-the-fly based on aggregate constraint violation, ensuring immediate feedback on infeasibility and preventing the system from stagnating in infeasible regimes. The update is controlled by the magnitude of constraint violation and capped by a maximum value to prevent gradient explosion.

Policy Optimization

ALaM integrates with off-policy actor-critic methods; specifically, this work implements SAC-ALaM, combining ALaM with Soft Actor-Critic, exploiting entropy-regularized exploration and robust value estimation. Policy updates are performed alternately with multiplier and penalty updates, with delayed multiplier updates to ensure near-converged policies when dual regression occurs.

Theoretical Analysis

ALaM’s convergence is analyzed in the infinite-dimensional Hilbert space setting, assuming compactness of state and action spaces for the continuous control regime. All dual operators—subdifferential mappings of the Lagrangian—are maximal monotone, and the update sequence for the statewise multipliers corresponds to projections onto the nonnegative cone in $L^2(\mathcal{X})$, ensuring weak convergence to feasible dual solutions.

A key result is that any weak sequential cluster point of the policy iteration converges to a globally optimal feasible policy for the constrained problem, under assumptions of strong continuity and convexity for reward and feasibility mappings.

Empirical Evaluation

Experimental Suite

SAC-ALaM is benchmarked on the Safety-Gymnasium suite, spanning diverse continuous control agents (Point, Car, Swimmer, HalfCheetah, Ant) and tasks (Velocity, Goal, Circle, Button) with statewise constraints. Baselines included scalar-multiplier Lagrangian SAC, FEASIBLE Actor-Critic (SAC-LagNet), PID-Lagrangian, and Augmented Lagrangian variants, with a unified SAC reward backbone for fairness.

Performance and Stability

SAC-ALaM exhibits superior performance—maximizing returns while strictly satisfying statewise constraints—across all environments. Critically, it substantially reduces training oscillations compared to both vanilla and penalty-augmented Lagrangian networks, as seen in cost and return learning curves.

Figure 2: Learning curves of SAC-ALaM and baselines across 8 environments. The solid lines represent the average performance over 5 seeds, and the shaded areas indicate the 95% confidence intervals. The first row shows the average episode cost, and the second row shows the average episode return.

The ablation study confirms that supervised dual regression is essential: substituting gradient ascent for regression (SAC-ALaM-GA) deteriorates training stability and final constraint satisfaction.

Figure 3: Ablation of multiplier update. Solid lines and shaded regions denote the mean and 95% confidence intervals across 5 seeds. The first row shows the average episode cost, and the second row shows the average episode return.

ALaM's robustness manifests in the stability of both return and cost metrics throughout training, not just in asymptotics but also transient behavior.

Figure 4: Training stability comparison in the PointCircle2 and SwimmerVelocity tasks.

Generalization and Multiplier Calibration

SAC-ALaM's learned multipliers provide physically-consistent and interpretable statewise risk quantification, even in rarely-explored or out-of-distribution regions, as demonstrated by fine-grained heatmaps. Unlike SAC-LagNet, which over-penalizes safe regions or collapses to uniform penalties, SAC-ALaM differentiates hazards based on both location and agent velocity, generalizing from limited training data.

Figure 5: Heatmaps of the multiplier for SAC-ALaM and SAC-LagNet across varying agent velocities.

Implications and Future Directions

Theoretically, ALaM establishes a rigorous foundation for optimization with parametric multiplier networks, a critical requirement for scalable safe RL in high-dimensional or continuous domains. Practically, SAC-ALaM demonstrates that with proper algorithmic design (quadratic penalties, supervised dual regression, and adaptive penalty scaling), it is possible to robustly train RL agents that enforce complex statewise constraints while optimizing performance.

Potential future work includes: (1) extension to multi-agent constrained settings; (2) adaptation to partially-observed or visually-rich RL environments; (3) exploration of alternative constraint aggregators and feasibility functions (e.g., Hamilton-Jacobi reachability); (4) further analysis of safety generalization under domain shift and adversarial perturbations; and (5) combining with model-based RL for safer real-world transfer.

Conclusion

ALaM introduces a robust, convergent architecture for learning statewise multipliers in safe RL via augmented Lagrangian and supervised regression. SAC-ALaM, the instantiation with Soft Actor-Critic, achieves superior safety-performance tradeoffs and stable training dynamics, setting a new state-of-the-art for scalable safe RL under complex, statewise constraints.