Skip to content

Shielding Architecture

This note documents the canonical conservative, budget-augmented TransitionShield implementation.

Why the old VI was unsound

The previous src/shield/sound_value_iteration.py initialized the upper bound to 1.0 everywhere and then iterated only that upper bound downward. In finite avoidance problems that is not enough: states that are safe almost surely remain stuck at the trivial fixed point unless the algorithm is seeded with known zero states or run as an interval method.

In practice this meant toy cases like:

  • one action reaches unsafe with probability 1
  • another action stays in a safe absorbing component

could still produce q(s) = 1.0 for the safe state. That is the core bug that drove this overhaul.

Current VI design

sound_value_iteration_min_reach_upper(...) now computes a lower and upper interval for the minimal probability of eventually reaching unsafe:

  • unsafe states are fixed to 1
  • the almost-sure safe set is computed first and fixed to 0
  • lower bounds are iterated upward
  • upper bounds are iterated downward
  • missing transition mass is conservatively sent to an unsafe sink

The returned SoundVIResult keeps:

  • q: upper bound on minimal unsafe reachability
  • lower_q: lower bound
  • upper and lower action risks
  • convergence gap / iteration metadata
  • missing-mass and state-action coverage diagnostics

The shield runtime uses the upper bound only, but the lower bound and gap are useful for debugging and regression tests.

Per-level shield bundle

The canonical runtime wrapper is TransitionShield.

Its intended construction path is:

  1. load env_transition_graph.pkl
  2. build one induced focal MDP per opponent-model level with build_iop_induced_graph(...)
  3. run sound VI on each level
  4. freeze the results into a TransitionShieldBundle

The runtime wrapper can still accept precomputed vi_levels, but the preferred backend path is the immutable per-level bundle because it keeps graph induction, unsafe labelling, and VI metadata tied together.

Budgeted runtime rule

In the default runtime mode the shield:

  • reads argmax(iop.mix_alpha)
  • raises a monotone level floor that never decreases
  • considers all accessible levels from that floor upward
  • exposes the focal observation as flatten(obs) || current_safety_bound || level_floor
  • keeps the focal action space equal to the base environment Discrete action
  • exposes the robust budget-admissible actions as a pre-sampling mask
  • replaces an unsafe external proposal with the highest-reward admissible pure action, using lower risk as the tie-break
  • updates the remaining budget from the realized successor state's projected budget on every step

This keeps the project’s conservative multi-level rule:

  • once the opponent model appears to have progressed to a higher level, we stay there
  • higher remaining levels must all be tolerated
  • the successor lower bound is beta(s') = V_floor(s'), where V_floor is the robust floor-wise eventual-unsafe certificate over all remaining levels
  • missing coverage contributes budget demand 1.0
  • a floor raise leaves the scalar budget unchanged and only changes the future admissible-action test

The base environment state contract is unchanged, but the wrapper adds budget and floor scalars outside the base env so the learner sees the shielded training state.

shield_opponent_mode can instead select the posterior-credible monotone suffix, fixed all_levels, or posterior-predictive bayesian_mixture ablations. The Bayesian mode appends the posterior weights to the wrapper observation/state. The aggregation and guarantee differences are maintained in opponent-uncertainty-modes.md.

Level-raise patience

TransitionShield now supports patience, with a default of 64.

  • patience counts successful calls to iop.update_mixture_with_observation(...), so in the common runtime path it is one observed opponent joint-action sample per timestep.
  • A raise is allowed only after at least patience new samples have supported the same next floor candidate. Elapsed samples alone are not enough; a transient higher-level argmax(mix_alpha) blip resets unless it persists.
  • When the posterior jumps multiple levels at once, the shield still raises by only one level per patience window. For example, if argmax(mix_alpha) jumps from level 0 to level 2 and stays there, the floor moves 0 -> 1 after the first full patience window and 1 -> 2 only after another full patience window.
  • patience=0 disables this gate and restores the old eager behavior, where the monotone floor immediately matches the current argmax(mix_alpha).

The candidate counter is attached to the runtime shield instance rather than to episodes, so sustained evidence carries across environment resets within the same training run.

For monotone_floor, the estimate is the posterior MAP level. For credible_floor, it is the largest suffix floor retaining the scheduled posterior mass. Patience only delays either raise, so runtime decisions use the patience-filtered floor rather than independently recomputing it.

On-policy action consistency

Shielded IPPO must not store a proposed action when the wrapper executes a different primitive action. PPO's old log probability and likelihood ratio would then describe the wrong behavior distribution.

The wrapper therefore exposes action_mask(agent_id) from the current public state, safety budget, and OM floor. IPPO applies the mask to categorical logits before sampling, stores the mask in the rollout, and applies the same mask again during every PPO epoch. It also checks last_executed_actions after the step and raises if any wrapper still rewrote an action. This is the replacement-shield analogue of MASA's augmented action: in both cases the action optimized by PPO is exactly the action contract consumed by the environment.

Telemetry modes

TransitionShield(debug_mode=...) supports:

  • off: no shield-specific info fields
  • minimal: override flag, proposed/final action, eventual and immediate risk, predicted reward, reward delta, current/realized-next budget, admissible set, level floor, levels considered, missing coverage, and decision reason
  • verbose: adds per-action aggregated risks, which level supplied the winning risk, per-action immediate risks and robust rewards, state key, convergence gaps, current mix_alpha, projected successor budgets, and an in-memory debug trace

Verbose traces are stored in a bounded deque and can be read with get_debug_trace().

Episode summaries

Shield episode summaries are recorded through LabelledEnv rather than a second logging system.

Per-episode EpisodeLog entries now include:

  • safety outcome: whether any violation occurred, violation rate, and first violation step
  • override, unsafe-proposal, missing-coverage, and infeasible-budget counts/rates
  • mean/max eventual proposed and final risk plus final safety margins
  • current and realized-next budget mean/min
  • admissible-action count/availability and OM floor summaries
  • proposed/final immediate risk and one-step Brier score
  • first-step reachability risk and a clearly named finite-episode Brier diagnostic (not an infinite-horizon guarantee estimate)
  • predicted proposed/final reward and cumulative replacement reward delta
  • decision-reason, proposed/final action, and replacement-transition histograms

Because the shield sits outside the labelled wrapper, record_shield_step(...) also patches the most recent EpisodeLog when the final transition of an episode already triggered logging.