Safe reinforcement learning: constrained MDPs, Lagrangian methods and shielding

Safe reinforcement learning is reinforcement learning in which the agent must respect safety requirements as well as maximise reward, both in the policy it ends up with and, ideally, while it is still learning. The standard formalism is the constrained Markov decision process: maximise expected return subject to an upper limit on expected cost, usually solved with Lagrangian primal-dual methods. Where constraints must hold on every step, methods such as shielding and safe exploration restrict which actions the agent may take at all.

This page builds the subject from a small example. We define a constrained MDP, solve one exactly, train a primal-dual agent on it in numpy, and show why the Lagrangian method oscillates around the constraint and how a PI-controlled multiplier fixes it. Then we cover constraint satisfaction during training, shielding, safe exploration, and how the same machinery is used to train language models. The interactive figure reruns the training in your browser as you move the cost limit.

A worked example: the shortcut through the hazard#

A robot starts at the bottom of a 6 by 6 grid and earns a reward of 1 for reaching a goal at the top, with discount factor 0.95 so that shorter routes are worth more. Between them lies a hazard band two rows deep that covers the middle four columns; every step that lands in it costs 1. Moves slip in a random other direction 10% of the time. Two natural policies:

• Straight through the hazard: 5 steps, return 0.772, discounted cost 2.06.
• Around the side: 9 steps, return 0.621, discounted cost 0.069 (the cost is not zero because slips near the edge sometimes land in the hazard).

Now require the expected discounted cost to be at most $d=0.5$. Neither policy is the answer: the shortcut violates the limit and the detour leaves most of the budget unused. The best policy under the constraint, worth 0.655, takes the shortcut about a fifth of the time and the detour otherwise. That randomisation is not a quirk; optimal policies for constrained MDPs can need it, which is one of the first differences from the unconstrained theory in Markov decision processes, where a deterministic optimal policy always exists [1].

What is a constrained MDP?#

A constrained MDP adds cost functions to an ordinary MDP. With one cost $c$ and limit $d$, the problem is

$J_r$ is the expected discounted return, $J_c$ the expected discounted cost, $\gamma$ the discount factor and $d$ the cost budget. Costs are kept separate from reward deliberately: folding a hazard penalty into the reward forces you to choose an exchange rate between reward and safety in advance, while a constraint lets you state the safety requirement directly and lets the optimiser find the exchange rate. Altman's monograph is the standard reference; it shows that the problem can be written as a linear program over discounted state-action occupancy measures, which is why optimal solutions are mixtures of at most a few deterministic policies [1].

The Lagrangian and the frontier#

Attach a multiplier $\lambda \ge 0$ to the constraint:

For a fixed $\lambda$, maximising $\mathcal{L}$ over $\pi$ is an ordinary MDP with reward $r-\lambda c$, solvable by value iteration. Each $\lambda$ gives one point (cost, return) on the frontier of achievable trade-offs, and the constrained optimum for any $d$ lies on the upper concave hull of those points. We computed the frontier for the gridworld from 161 values of $\lambda$: at $\lambda =0$ the best response is the shortcut; at $\lambda =1$ it is the detour; at $\lambda =100$ it avoids risk so strongly that it barely reaches the goal (return 0.0006). Despite the non-convexity of policy space, constrained RL has zero duality gap under mild conditions, so solving the dual recovers the constrained optimum [8].

The marginal price of safety is the slope of the frontier, and the optimal multiplier equals that slope. Between costs 0.1 and 1.93 the frontier is a straight line with slope 0.078 return per unit of cost, so every unit of risk you allow buys about 0.08 of return here.

Lagrangian primal-dual training#

In practice neither the model nor the frontier is known, and the policy is learned. Primal-dual methods alternate a policy step that increases $\mathcal{L}$ with a multiplier step that raises $\lambda$ when the constraint is violated and lowers it otherwise:

$\theta$ are the policy parameters, $\eta$ and ${\eta }_{\lambda }$ the primal and dual step sizes. The policy step can be any policy-gradient method (see the policy gradient theorem and TRPO and PPO) applied to the combined advantage $A_r-\lambda A_c$. Reward Constrained Policy Optimisation runs the multiplier on a slower timescale than the policy and proves convergence to a constraint-satisfying solution under that separation [5].

The gridworld, trained#

We trained a tabular softmax policy on the gridworld from a uniform random start with the exact natural-policy-gradient step for this policy class, ${\theta }_{s,a}\leftarrow {\theta }_{s,a}+\frac{\eta }{1-\gamma }(A_r(s,a)-\lambda A_c(s,a))$ with $\eta =0.05$, and exact policy evaluation by solving the Bellman equations as a linear system (see the Bellman equation). This removes sampling noise, so what remains is the dynamics of the primal-dual game itself. 600 iterations per run, cost limit 0.5.

Results (the computation is deterministic, so there is no seed variation to report):

The pattern matches what Stooke and colleagues reported for deep RL: Lagrangian learning dynamics oscillate and overshoot, producing constraint violations during training, because the multiplier is an integral controller that only reacts to accumulated violation; adding proportional and derivative terms on the constraint error damps the oscillation [6]. Here the plain Lagrangian's last iterate never settles: its cost swings between 0.31 and 0.66 for the final hundred iterations at dual step 0.2. Its average over iterates comes closer (return 0.652 at cost 0.468, against the optimum of 0.655 at 0.500), which is the usual guarantee for primal-dual methods: the average converges, the last iterate need not. A larger dual step makes the agent conservative on average (cost 0.169) rather than accurate. The PI version reached return 0.655 at cost 0.500, the constrained optimum to three decimals. Its final multiplier, 0.078, equals the slope of the frontier at the optimum, as the theory says it should.

Constraint satisfaction during training#

Every run above began with a policy whose cost was 3.47, seven times the limit, and the Lagrangian runs went back over the limit many times as the multiplier oscillated. A CMDP constraint is a statement about the final policy's expected cost; it says nothing about what happens while learning. When violations during training are themselves unacceptable (a real robot, a production system), you need one or more of:

• Start from a safe policy. Initialise from a conservative demonstrator or a policy known to satisfy the constraint, and keep updates small. Constrained Policy Optimisation does this with a trust region and offers guarantees of near-constraint satisfaction at each iteration [4].
• Per-step action correction. A safety layer learns a linear model of how each action changes the constraint value and projects the policy's action onto the safe set at every step, in closed form [11].
• Shielding. Block unsafe actions using a model of the safety-relevant dynamics (next section).
• Safe exploration with uncertainty. Only visit states whose safety is certified under a model's uncertainty, and expand that set as evidence accumulates (below).
• Tighter or different constraints. Expected-cost limits allow rare catastrophic episodes as long as the average is low. Chance constraints and conditional value-at-risk limits constrain the tail instead [7].

What is shielding in safe RL?#

A shield is a component, synthesised from a formal safety specification and a model of the environment, that sits next to the learning agent and prevents it from taking actions that would violate the specification. Alshiekh and colleagues described two placements [3]. A preemptive shield acts before the agent decides: at each step it hands the agent the list of actions that are safe in the current state, so the agent only ever chooses among them. A post-posed shield acts after: it monitors the agent's chosen action and replaces it only if it would violate the specification. Either way the specification holds during learning as well as after, provided the model the shield is built from is correct.

That proviso is the main limitation. Shields need a model of the dynamics that matter for safety, at least an abstraction of them, and a specification written in a formal language such as temporal logic. In the gridworld, a shield that removed every action whose intended move enters the hazard would leave only the cost caused by slips but would also make the constrained optimum (which deliberately uses the hazard a fifth of the time) unreachable: shields enforce hard constraints, not budgets. In language-model systems, input and output filters play a role similar to a post-posed shield; our guide to AI guardrails covers how to calibrate them, and guardrail monitoring how to watch them in production.

Safe exploration#

Exploration is how an RL agent learns, and it is the part most likely to be unsafe. Three representative ideas from the literature:

• Stay recoverable. Restrict exploration to policies that keep the agent able to return to where it started, a notion of ergodicity with respect to a safe home state [12].
• Expand a certified safe set. Model an unknown safety feature with a Gaussian process and only visit states that are safe with high probability and from which the agent can return; the safe set grows as observations reduce uncertainty [13].
• Use control theory. With a Lyapunov function and a statistical model of the dynamics, update the policy only within a region of attraction that is certified stable with high probability, and enlarge that region as data accumulates [14].

These methods give guarantees that primal-dual methods do not, at the price of assumptions (smoothness, known structure, low dimensions) that are hard to meet in large systems. Benchmarks such as Safety-Gymnasium make it possible to compare safe RL algorithms on both return and cost, including cost during training [9]; when you evaluate, report cost with confidence intervals across seeds as you would return, as described in how to evaluate RL agents. Surveys by García and Fernández [2] and by Gu and colleagues [10] organise the much larger literature.

Safe RL for language models#

The CMDP view transfers directly to language-model alignment. Safe RLHF trains separate reward and cost models, one for helpfulness and one for harmlessness, from separately collected preference data, and optimises the policy with a Lagrangian so that expected cost stays below a threshold while helpfulness is maximised [15]. Constrained RLHF uses the same machinery for a different purpose: it treats each component reward model as useful only up to a threshold beyond which it stops tracking human ratings, and constrains the policy not to exceed it, which counters over-optimisation [16]. The KL penalty used in most RLHF is itself a fixed Lagrange multiplier on a divergence constraint; KL-regularised RL makes that connection. The oscillation problem above appears in this setting too, so multiplier dynamics deserve as much monitoring as the losses. For the broader context, see what is AI safety, and for what happens when the cost model itself is exploited, reward hacking.

Which safe RL method should you use?#

Code: a CMDP primal-dual solver in numpy#

The full gridworld run. It runs in well under a second and reproduces the PI row of the table.

import numpy as np

N, GAMMA, SLIP, START, GOAL = 6, 0.95, 0.1, (5, 2), (0, 2)
HAZ = {(r, c) for r in (2, 3) for c in (1, 2, 3, 4)}
MOVES = [(-1, 0), (0, 1), (1, 0), (0, -1)]
S, A = N * N, 4

P = np.zeros((S, A, S)); Rw = np.zeros((S, A)); Cw = np.zeros((S, A))
for r in range(N):
    for c in range(N):
        s = r * N + c
        if (r, c) == GOAL:
            P[s, :, s] = 1.0
            continue
        for a in range(A):
            for b in range(A):                      # intended move with prob 0.9, else slip
                p = 1 - SLIP if b == a else SLIP / 3
                r2 = min(N - 1, max(0, r + MOVES[b][0])); c2 = min(N - 1, max(0, c + MOVES[b][1]))
                P[s, a, r2 * N + c2] += p
                Rw[s, a] += p * ((r2, c2) == GOAL)
                Cw[s, a] += p * ((r2, c2) in HAZ)
s0 = START[0] * N + START[1]

def evaluate(pi):
    M = np.eye(S) - GAMMA * np.einsum("sa,sat->st", pi, P)
    return np.linalg.solve(M, (pi * Rw).sum(1)), np.linalg.solve(M, (pi * Cw).sum(1))

def primal_dual(d=0.5, eta=0.05, ki=0.2, kp=1.0, iters=600):
    theta, lam, integral = np.zeros((S, A)), 0.0, 0.0
    for _ in range(iters):
        pi = np.exp(theta - theta.max(1, keepdims=True)); pi /= pi.sum(1, keepdims=True)
        Vr, Vc = evaluate(pi)
        Ar = Rw + GAMMA * P @ Vr - Vr[:, None]          # advantages for reward and cost
        Ac = Cw + GAMMA * P @ Vc - Vc[:, None]
        theta += eta / (1 - GAMMA) * (Ar - lam * Ac)    # natural policy gradient, softmax policy
        err = Vc[s0] - d
        integral = max(0.0, integral + err)
        lam = max(0.0, kp * err + ki * integral)        # kp = 0 gives the plain Lagrangian update
    return Vr[s0], Vc[s0], lam

print(primal_dual())   # return about 0.655, cost about 0.500

Failure modes#

Reporting the last iterate of an oscillating run. Depending on where the oscillation stopped, the final policy can be well inside or well outside the constraint. Plot cost over training and report its range, not only the final value.

Treating a CMDP constraint as a training-time guarantee. It constrains the converged policy's expected cost. If violations during learning matter, use a method designed for that.

A cost signal that can be gamed. A learned cost model is a proxy like any reward model; a policy under a cost constraint is optimising to find where the cost model under-reports. Audit cost with an independent measure.

Expected cost hiding the tail. A budget of 0.5 can be met by never entering the hazard or by entering it disastrously once in a while. If the tail matters, constrain the tail.

Shields built on the wrong model. A shield is exactly as safe as the abstraction it was synthesised from. Test it on states outside that abstraction.

Frequently asked questions#

Sources#

1. Eitan Altman. Constrained Markov Decision Processes, 1999. www-sop.inria.fr
2. Javier García, Fernando Fernández. A Comprehensive Survey on Safe Reinforcement Learning, 2015. jmlr.org
3. Mohammed Alshiekh, Roderick Bloem, Ruediger Ehlers, et al.. Safe Reinforcement Learning via Shielding, 2017. arxiv.org
4. Joshua Achiam, David Held, Aviv Tamar, Pieter Abbeel. Constrained Policy Optimization, 2017. arxiv.org
5. Chen Tessler, Daniel J. Mankowitz, Shie Mannor. Reward Constrained Policy Optimization, 2018. arxiv.org
6. Adam Stooke, Joshua Achiam, Pieter Abbeel. Responsive Safety in Reinforcement Learning by PID Lagrangian Methods, 2020. arxiv.org
7. Yinlam Chow, Mohammad Ghavamzadeh, Lucas Janson, Marco Pavone. Risk-Constrained Reinforcement Learning with Percentile Risk Criteria, 2015. arxiv.org
8. Santiago Paternain, Luiz F. O. Chamon, Miguel Calvo-Fullana, Alejandro Ribeiro. Constrained Reinforcement Learning Has Zero Duality Gap, 2019. arxiv.org
9. Jiaming Ji, Borong Zhang, Jiayi Zhou, et al.. Safety-Gymnasium: A Unified Safe Reinforcement Learning Benchmark, 2023. arxiv.org
10. Shangding Gu, Long Yang, Yali Du, et al.. A Review of Safe Reinforcement Learning: Methods, Theory and Applications, 2022. arxiv.org
11. Gal Dalal, Krishnamurthy Dvijotham, Matej Vecerik, et al.. Safe Exploration in Continuous Action Spaces, 2018. arxiv.org
12. Teodor Mihai Moldovan, Pieter Abbeel. Safe Exploration in Markov Decision Processes, 2012. arxiv.org
13. Matteo Turchetta, Felix Berkenkamp, Andreas Krause. Safe Exploration in Finite Markov Decision Processes with Gaussian Processes, 2016. arxiv.org
14. Felix Berkenkamp, Matteo Turchetta, Angela P. Schoellig, Andreas Krause. Safe Model-based Reinforcement Learning with Stability Guarantees, 2017. arxiv.org
15. Josef Dai, Xuehai Pan, Ruiyang Sun, et al.. Safe RLHF: Safe Reinforcement Learning from Human Feedback, 2023. arxiv.org
16. Ted Moskovitz, Aaditya K. Singh, DJ Strouse, et al.. Confronting Reward Model Overoptimization with Constrained RLHF, 2023. arxiv.org