Add and Test an Action Shield During Reinforcement Learning Exploration

Your agent keeps taking an action you already know is wrong. You add a shield. Now you cannot tell whether it is helping or quietly starving the learner.

This tutorial builds the smallest transparent version of that filter: one predicate, one fallback, one log. Then we read the log like evidence instead of trusting the demo.

What You Are Actually Building

If you have read the conceptual version of this topic, you already know why shields exist. This article is about the contract.

A shield is a function that sits between the policy and the environment. It receives a proposed action and returns either that action or a replacement. That is the whole interface.

Two integration points matter:

• Pre-shielding restricts the candidate set before selection. The policy chooses from allowed actions only.
• Post-shielding inspects and corrects the chosen action after the fact.

We use post-shielding here because it is easier to log. You see the proposed action and the substituted action, which is exactly the comparison you need later.

One boundary before we write code: the filter below encodes a constraint you wrote by hand. It is not a learned safety critic, and it carries no formal guarantee. It is a rule you asserted and then observed.

Set Up the Toy Task and the Constraint

We need a small discrete-action task where at least one action is hazardous in specific states. A grid world works. So does Taxi-style navigation, where "pick up" is invalid unless you are standing on the passenger.

Assumptions for everything below:

• Python 3.10 or newer
• A Gymnasium-style API (reset() returns (obs, info), step() returns a five-tuple)
• A tabular Q-learning agent, no GPU, no external services

Here is the environment skeleton. The is_unsafe predicate is the constraint — plain, readable, testable, and easy to break on purpose.

import numpy as np

class GridWorld:
    """5x5 grid. Action 4 is 'sprint' and is hazardous near the pit at (2, 2)."""
    def __init__(self):
        self.n_states = 25
        self.n_actions = 5  # 0-3 move, 4 sprint
        self.pit = (2, 2)

    def reset(self):
        self.pos = (0, 0)
        return self.pos, {}

    def step(self, action):
        r, c = self.pos
        if action == 0: r = max(0, r - 1)
        elif action == 1: r = min(4, r + 1)
        elif action == 2: c = max(0, c - 1)
        elif action == 3: c = min(4, c + 1)
        # action 4: sprint — moves two cells, but only safe away from the pit
        elif action == 4:
            r = min(4, r + 2)

        self.pos = (r, c)
        done = self.pos == (4, 4)
        reward = 1.0 if done else -0.01
        if self.pos == self.pit:
            reward = -1.0
            done = True
        return self.pos, reward, done, False, {}

Now the constraint as a predicate over (state, action). The key detail: the predicate must evaluate the resulting state, not the current one. A sprint from (0, 2) lands on the pit at (2, 2), so the predicate has to simulate the transition to catch it.

def is_unsafe(state, action):
    r, c = state
    if action == 4:
        # sprint moves two rows; check where it lands
        next_r = min(4, r + 2)
        return (next_r, c) == (2, 2)
    return False

Run the unconstrained baseline first. You need a reference point before the comparison means anything.

Write the Filter and the Fallback

The filter returns the proposed action when it passes the predicate. Otherwise it selects from the allowed set.

def shield(state, proposed_action, fallback="safest"):
    if not is_unsafe(state, proposed_action):
        return proposed_action, False

    allowed = [a for a in range(5) if not is_unsafe(state, a)]
    if not allowed:
        # Nothing is permitted. We still must return something.
        # A no-op is the honest choice here; document it.
        return 0, True

    if fallback == "safest":
        # Prefer the action that moves away from the pit.
        r, c = state
        away = 0 if r > 2 else 1
        return (away if away in allowed else allowed[0]), True
    if fallback == "fixed":
        return 0, True
    if fallback == "random":
        return int(np.random.choice(allowed)), True
    raise ValueError(fallback)

Three things to notice.

The fallback is a real design decision. "Safest available action" keeps the agent moving toward the goal. A fixed default is simpler but can stall the agent in place. A random allowed action preserves exploration but adds noise. Each changes what the agent learns, because each changes the data the Q-table sees.

The empty-allowed-set case is explicit. If nothing is permitted, the shield must still return something. That choice deserves a comment, not a silent default.

Every branch returns the same (action, blocked) tuple. That contract is what makes the caller simple. A branch that returns a bare action will break the unpacking in the next section — a mistake worth catching with a test:

assert shield((2, 2), 4)[1] is True   # sprint near pit is blocked
assert shield((0, 0), 4)[1] is False  # sprint far from pit passes
assert shield((0, 2), 4)[1] is True   # sprint from (0,2) lands on the pit

That third assertion is the one that matters. It is the case where a naive predicate — one that checks the current row instead of the landing row — passes a hazardous action through.

Log Every Intervention

Without a log, you cannot distinguish a working shield from a broken policy. The log is the artifact that makes the experiment interpretable.

Record per step:

def run_episode(env, q, epsilon, use_shield, fallback="safest"):
    state, _ = env.reset()
    log = []
    done = False
    while not done:
        if np.random.rand() < epsilon:
            proposed = np.random.randint(env.n_actions)
        else:
            proposed = int(np.argmax(q[state[0] * 5 + state[1]]))

        if use_shield:
            action, blocked = shield(state, proposed, fallback)
        else:
            action, blocked = proposed, False

        next_state, reward, done, _, _ = env.step(action)
        log.append({
            "state": state,
            "proposed": proposed,
            "blocked": blocked,
            "executed": action,
            "reward": reward,
        })
        # Q-update: bootstrap from the best valid action in the next state
        s_idx = state[0] * 5 + state[1]
        ns_idx = next_state[0] * 5 + next_state[1]
        best_next = np.max(q[ns_idx])
        q[s_idx, action] += 0.1 * (reward + 0.99 * best_next - q[s_idx, action])
        state = next_state
    return log

Aggregate per episode: intervention count, intervention rate, and how often the fallback fired because no action was allowed.

Common mistake: Treating the logged proposal as a counterfactual. You recorded what the agent wanted to do, but you did not observe the reward that action would have produced. The environment never ran it. If you want that number, you need a separate controlled rollout or a model — not a field in this log.

Compare Shielded and Unshielded Runs

Hold seeds, episode count, learning rate, and environment fixed. Change only the presence of the shield.

def train(use_shield, episodes=300, fallback="safest"):
    env = GridWorld()
    q = np.zeros((env.n_states, env.n_actions))
    returns = []
    for ep in range(episodes):
        log = run_episode(env, q, epsilon=0.1, use_shield=use_shield, fallback=fallback)
        returns.append(sum(step["reward"] for step in log))
    return returns

for seed in range(5):
    np.random.seed(seed)
    unshielded = train(use_shield=False)
    np.random.seed(seed)
    shielded = train(use_shield=True)

Plot or tabulate return and intervention rate across episodes. Do not report a single final score — that hides the shape of learning.

Expect one of three honest outcomes:

The confound to watch: a shield that blocks a lot can look like a better policy while actually shrinking the action space the agent explores. If intervention rate is high and return is flat, you have constrained the problem, not solved it.

Failure Modes and Debugging Signals

The shield can lie to you in four ways.

Intervention rate near 100%. The constraint is too tight or the predicate is wrong. The agent is effectively learning a different task — one where its own choices rarely survive.

Intervention rate near 0%. The shield is dead code, or the unsafe action was never reachable in the sampled states. Check whether your exploration actually visits the states where the predicate fires.

Fallback dominates the action distribution. The agent's policy is being overwritten. Its Q-values are learning from data it did not choose. This is the subtlest failure: the return curve may look fine while the policy underneath is meaningless.

Reward improves while the unshielded baseline also improves. The shield may be irrelevant. The gain came from something else — a lucky seed, a learning-rate interaction, or the environment itself.

Tip: Print the intervention rate alongside the return curve. A shield that never fires and a shield that always fires are both suspicious for different reasons.

One Modification Worth Trying

Pick one variant and state its success criterion before you run it.

Swap the fallback strategy. Run safest, fixed, and random on the same seed. Compare intervention logs. Success criterion: the fallback that preserves forward progress produces the highest return without inflating the intervention rate.

Relax the constraint mid-training. After 500 episodes, widen the unsafe radius or remove it entirely. Watch whether behavior degrades, holds, or improves. Success criterion: the agent's policy survives the relaxation without a return collapse.

Move from post-shielding to pre-shielding. Restrict the candidate set before selection instead of correcting after. Note what changes in the code — you now need the allowed set at selection time — and what changes in the log — you no longer see the proposed action, only the chosen one.

Each variant tests a different assumption. The fallback swap tests whether your replacement policy matters. The relaxation tests whether the agent learned a general strategy or just a constrained one. The pre/post shift tests whether observability was worth the integration cost.

What This Does Not Prove

A hand-written predicate over a toy environment is not a safety guarantee. It is a constraint you asserted and then observed.

The shield only enforces what you encoded. Anything you failed to specify is still free for the agent to do. Formal shields are synthesized from a model of the system and come with guarantees this tutorial does not attempt — they reason about all reachable states, not the ones your agent happened to visit.

The transferable skill is the loop: encode a constraint, log interventions, compare against a baseline, and refuse to overclaim.

Before trusting any shield, read its intervention log and ask one question: is the agent still learning from its own choices?

If the answer is no — if the fallback is doing most of the work — then you have not made the agent safer. You have replaced it with a rule and given the rule credit for the agent's reward.

Run the same comparison with a different fallback. Then relax the constraint and see what the policy actually learned. The log, not the reward curve, is what tells you whether the filter is doing what you think.