🤖 machine learning

Augmented Lagrangian Method for Last-Iterate Convergence for Constrained MDPs

This paper proposes a general framework based on the inexact augmented Lagrangian method that achieves provable global last-iterate convergence for constrained Markov decision processes across tabular, log-linear, and complex non-linear policy settings, addressing the practical limitations of existing mixture-policy approaches.

Original authors: Michael Lu, Max Qiushi Lin, Mo Chen, Sharan Vaswani

Published 2026-05-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Michael Lu, Max Qiushi Lin, Mo Chen, Sharan Vaswani

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are training a robot to play a video game. The goal is simple: get the highest score possible. But there's a catch. The robot has a strict rule: it cannot run out of battery before the game ends.

This is the core problem the paper tackles, known in the tech world as a Constrained Markov Decision Process (CMDP). The robot (the "agent") needs to maximize its reward (score) while staying under a budget (battery life).

The Problem with Current Methods: The "Mix-and-Match" Mess

Most existing AI methods for this problem work like a chef trying to make a perfect soup. They try many different recipes (policies) one by one. At the end, instead of serving you the single best recipe they found, they tell you: "Here is a bowl of soup made by randomly mixing a little bit of every recipe we tried."

While this "mixed soup" works well on paper (mathematically, it satisfies the rules), it's a nightmare in the real world:

Memory Heavy: You have to remember every single recipe you ever tried to make the mix.
Unpredictable: If you actually serve a single spoonful from that mix, it might be terrible. One random spoon might be pure salt (violating the battery rule), even if the average bowl is fine.
Oscillation: The robot's behavior often swings wildly back and forth, never settling down.

The paper argues that in real life (like self-driving cars or medical devices), we can't rely on a "random mix." We need one single, final robot that is safe and effective right out of the box. This is called "Last-Iterate Convergence."

The Solution: The "Augmented Lagrangian" (The Strict Coach)

The authors propose a new way to train the robot using a classic technique from math called the Augmented Lagrangian (AL) method.

Think of the AL method as a strict coach who doesn't just yell "Go faster!" (maximize reward) but also carries a heavy penalty weight on the robot's back if it breaks the rules.

Here is how the coach works:

The Penalty Weight: If the robot gets too close to running out of battery, the coach adds a heavy, quadratic penalty (like a heavy backpack) to the robot's goal. The more it violates the rule, the heavier the backpack gets, making it harder to move forward.
The Adjustment: The coach doesn't just leave the weight there. They constantly adjust how heavy the backpack is based on how well the robot is doing.
- If the robot is safe, the coach lightens the load slightly.
- If the robot is risky, the coach makes the load heavier immediately.
The Result: Instead of the robot swinging wildly between "too fast" and "too slow," the AL method guides it smoothly toward a single, stable path where it gets a high score and stays safe.

The "Magic" Ingredient: Projected Q-Ascent (PQA)

The paper's biggest breakthrough is figuring out how to make this "Strict Coach" work efficiently, even when the robot is learning complex skills (like walking or flying).

They use a specific training technique called Projected Q-Ascent (PQA).

The Analogy: Imagine the robot is trying to climb a hill to find the highest peak (the best score). But the hill has a "No-Go Zone" (the safety constraint).
The Old Way: The robot might try to climb, realize it's in the No-Go Zone, and then jump back and forth, never settling.
The PQA Way: The robot takes a step up the hill. If that step would put it in the No-Go Zone, PQA acts like a magnetic wall. It gently but firmly pushes the robot back to the edge of the safe zone, but keeps it moving in the best possible direction. It "projects" the robot's movement onto the safe path.

What Did They Prove?

The authors didn't just build a cool robot; they proved mathematically that this approach works:

It Converges: The robot will eventually stop swinging and settle on a single, final policy.
It's Safe: That final policy will satisfy the safety rules (battery limit) with high certainty, not just on average.
It's Efficient: They showed this works for simple grids (tabular) and complex, real-world tasks (like continuous control in video games) without needing to store thousands of past versions of the robot.

The Results in the Real World

The team tested their method (which they call PPQA-ALM or SPMA-ALM) on standard safety benchmarks (like a robot navigating a maze without hitting walls).

Comparison: They compared it against other popular methods (like PPO-Lag and CPO).
Outcome: Their method was just as good at getting high scores, but it was much more stable. It didn't oscillate. It found a single, reliable solution that respected the safety constraints, whereas the other methods sometimes struggled to settle down or required complex "mixing" tricks to work.

Summary

In short, this paper introduces a smarter way to train AI agents with safety rules. Instead of relying on a confusing "average" of many failed attempts, they use a Strict Coach with a Penalty Backpack and a Magnetic Wall to guide the AI to a single, perfect, and safe final behavior. This makes the technology ready for real-world applications where safety is non-negotiable.

Technical Summary: Augmented Lagrangian Method for Last-Iterate Convergence for Constrained MDPs

Problem Formulation

The paper addresses the problem of policy optimization for infinite-horizon, discounted Constrained Markov Decision Processes (CMDPs). In this setting, an agent must maximize a cumulative reward $V^\pi_r(\rho)$ while ensuring that cumulative constraint rewards $V^\pi_{c_i}(\rho)$ satisfy prescribed thresholds $b_i$ for $m$ constraints. Unlike standard MDPs where optimal policies are deterministic, optimal policies for CMDPs are generally stochastic.

A significant practical challenge arises from the limitations of existing Primal-Dual (PD) methods. While theoretical guarantees for PD methods exist, they typically apply only to a mixture policy (a uniform average of all intermediate policies generated during training). Deploying such a mixture is computationally and memory-intensive, as it requires storing all intermediate policies. Furthermore, in safety-critical applications, relying on a mixture is risky because a randomly sampled policy from the mixture may violate constraints significantly, even if the mixture satisfies them in expectation. Consequently, practitioners often deploy the last-iterate policy (the final policy generated), but standard PD methods lack theoretical guarantees for this specific policy, often exhibiting oscillatory behavior and failing to converge to a feasible solution.

Recent works attempting to prove last-iterate convergence have largely been limited to tabular settings or rely on algorithmic variants (e.g., heavy regularization) that are difficult to implement or impractical for large-scale problems.

Methodology

The authors propose a general framework based on the Augmented Lagrangian (AL) method, a classic approach in constrained optimization, to achieve provable last-iterate convergence for CMDPs.

1. The Augmented Lagrangian Framework

Instead of the standard Lagrangian used in PD methods, the framework optimizes the Augmented Lagrangian function $L_\beta(\pi, \lambda)$ :
$L_\beta(\pi, \lambda) := V^\pi_r(\rho) + \frac{\beta}{2} \sum_{i=1}^m \left( -\min\left\{ V^\pi_{c_i}(\rho) - b_i - \frac{\lambda_i}{\beta}, 0 \right\}^2 + \frac{\lambda_i^2}{\beta^2} \right)$
where $\beta > 0$ is a penalty parameter and $\lambda$ are Lagrange multipliers. The quadratic penalty term discourages constraint violations. The algorithm iteratively updates the policy (primal variable) to approximately maximize this function and updates the dual variables (Lagrange multipliers) based on constraint violations.

2. Solving the AL Subproblem

The core challenge is that $L_\beta(\pi, \lambda)$ is non-concave with respect to the policy $\pi$ , making the subproblem difficult to solve. The authors decompose the solution into two settings:

Tabular Setting: The authors propose solving the AL subproblem using Projected Q-Ascent (PQA). They derive the policy gradient for the AL subproblem, which involves a non-stationary pseudo-reward $\Gamma(\pi)$ that depends on the current policy's constraint satisfaction level. By using PQA, they establish global last-iterate convergence guarantees for the subproblem.
Function Approximation (Log-Linear Policies): To handle large state-action spaces, the framework is extended to parametric policies (specifically log-linear policies). The authors derive an AL policy gradient that decomposes linearly into the original reward and constraint components, allowing standard estimation tools (like those used in PPO or SPMA) to be employed. They introduce a Projected PQA (PPQA) update, which involves:
1. Computing an idealized policy update (PQA step).
2. Projecting this update onto the feasible policy class using forward KL-divergence, formulated as a surrogate optimization problem (multi-class logistic regression with soft labels).

3. Convergence Analysis

The authors provide a modular analysis:

They first prove that the inexact AL method converges to an $\epsilon$ -optimal solution for the CMDP if the subproblem solver achieves a specific accuracy ( $O(1/\epsilon^2)$ iterations for the outer loop).
They then analyze the specific solvers (PQA and PPQA) to show they can achieve the required accuracy for the subproblem.
Combining these results, they prove that the overall algorithm achieves global last-iterate convergence with a sample complexity of $O(1/\epsilon^6)$ gradient evaluations for both tabular and log-linear settings.

Key Contributions

Generic Inexact AL Framework: The paper introduces a modular framework for solving CMDPs using the inexact AL method. It separates the analysis of the AL iterations from the choice of the subproblem solver, proving that $O(1/\epsilon^2)$ outer iterations suffice to find an $\epsilon$ -optimal policy, provided the subproblem is solved with sufficient accuracy.
Tabular Last-Iterate Convergence: By instantiating the framework with Projected Q-Ascent (PQA) for the tabular setting, the authors establish global last-iterate convergence with a complexity of $O(1/\epsilon^6)$ . This addresses the non-stationarity of the pseudo-reward in the AL subproblem.
Extension to Function Approximation: The framework is generalized to log-linear policies. The authors derive the AL policy gradient and propose a projected variant (PPQA) that achieves comparable $O(1/\epsilon^6)$ convergence guarantees. Unlike prior work, this method does not require non-standard algorithmic modifications (like projecting onto restricted simplices) or restrictive structural assumptions (like compatible function approximation) and supports multiple constraints.
Empirical Validation: The framework is instantiated using PPO and SPMA (a recent policy optimization algorithm) and evaluated on large-scale continuous control tasks from Safety-Gymnasium. The results demonstrate that the proposed approach performs comparably to heavily engineered baselines like PPO-Lag and CPO, while offering theoretical guarantees for the final policy.

Results

Theoretical: The paper proves that the proposed algorithm returns a stochastic policy that is an $\epsilon$ -approximate solution to the CMDP problem after $O(1/\epsilon^6)$ gradient evaluations. This holds for both tabular and log-linear policy classes.
Empirical (Tabular): In tabular environments (Constrained Cliff World, Constrained Deep Sea Treasure), the proposed PQA-ALM and PPQA-ALM methods converge smoothly to the optimal policy with minimal oscillation, outperforming or matching other last-iterate methods like RPG-PD and ReLOAD.
Empirical (Continuous Control): On Safety-Gymnasium tasks (SafetyAntVelocity, SafetyHalfCheetahVelocity, SafetyHopperVelocity), the SPMA-ALM and PPO-ALM variants achieve returns and constraint violations comparable to PPO-Lag and CPO. Notably, the AL-based methods exhibit less oscillatory behavior and more consistent constraint satisfaction (lower mean and tail violations) compared to baselines.

Significance and Claims

The paper claims to bridge the gap between theoretical guarantees and practical deployment in constrained reinforcement learning.

Practicality: Unlike previous last-iterate convergence works that are either limited to tabular settings or rely on impractical regularizations, this framework is designed to be efficient and scalable to large problems using standard function approximation.
Safety: By guaranteeing convergence for the last-iterate policy rather than a mixture, the method ensures that the deployed policy itself satisfies constraints, which is critical for safety-critical applications where random sampling from a mixture is unacceptable.
Principled Approach: The method relies on the classic Augmented Lagrangian method, providing a clear theoretical foundation without the need for ad-hoc heuristics often found in practical constrained RL implementations (e.g., scaling objectives by multipliers).

The authors acknowledge limitations, specifically the strong dependence on the minimum state visitation probability ( $\rho_{min}$ ) in the tabular setting and the bias term ( $\epsilon_{bias}$ ) for log-linear policies, suggesting these as directions for future work. They also note that future directions include incorporating actor-critic or off-policy methods into the theoretical framework.