Budget-Augmented Robust Replacement Shielding¶

TransitionShield now implements a budget-augmented replacement shield rather than the paper/MASA encoded-action shield.

Core semantics¶

The shield keeps the base environment state contract unchanged:

• WM / OM / exact graph builders still reason over the public base state bits
• graph node keys still hash only those public bits
• the safety budget lives only in the outer shield wrapper

The default base shielded training wrapper augments the focal agent with two extra scalars:

• focal observation: flatten(obs) || current_safety_bound || level_floor
• focal action space: the original base-env Discrete(A)
• wrapper state(): flatten(base_env.state()) || current_safety_bound || level_floor

The optional bayesian_mixture ablation additionally appends the normalized opponent-level posterior because its action mask depends on those weights. See opponent-uncertainty-modes.md.

Non-focal agents keep the base environment observation and action spaces. The optional hybrid-action variants replace the focal Discrete(A) space with a categorical/raw-Gaussian Dict action; see parameterized-actions.md.

Bundle contents¶

Each ShieldLevelBuildResult now stores more than just SoundVIResult:

• the VI result itself
• a compiled successor kernel for every (state_key, a_self)
• per-successor budget lower bounds derived from q(s')
• robust focal-agent expected reward for every (state_key, a_self), compiled from the induced graph's reward-bearing edges

The runtime uses those kernels directly instead of only reading per-action risk tables. Missing transition mass is compiled as an explicit unsafe sink with budget 1.0.

Those stored risk bounds are not stage costs. They are conservative upper bounds on the minimal probability of eventually reaching unsafe from each successor state under the current induced model.

Bundles also store floor_budget_lower_bounds: one budget-certificate map per OM floor. For floor f, this map solves the robust Bellman equation over all levels l >= f:

V_f(s) = min_a max_{l >= f} E_l[V_f(s')]

Runtime successor-budget projection uses V_f(s'), not max_l q_l(s'). This is important for shrinking-budget semantics: a state can be cheap under each individual level because each level prefers a different action, while still being expensive under the robust single-action choice enforced at runtime. Using the floor-wise certificate keeps the budget carried to the next state at least as large as the next state's own robust eventual-unsafe demand.

Runtime decision rule¶

At a pre-step base state s the shield:

1. applies the configured opponent uncertainty mode (the monotone floor by default)
2. considers every accessible level from that floor upward
3. builds the union of all successor states reachable from s
4. sets beta(s') = V_floor(s'), defaulting to 1.0 on missing coverage
5. checks the proposed pure action against the current robust budget
6. keeps the proposal if it is safe; otherwise chooses the highest predicted reward among budget-admissible actions, using Bayesian expected reward by default and breaking reward ties by lower risk and then action index
7. computes one shared successor-budget vector W >= beta for the executed action
8. after the environment transition, updates current_safety_bound = W(s')
9. updates the OM posterior and monotone level floor from the realized opponent action

The successor-budget fill still uses the robust egalitarian rule: raise all non-saturated successor budgets equally while every accessible level still satisfies its expected-budget constraint for the executed action.

Practical notes¶

• TransitionShield accepts init_safety_bound, the initial safety budget for each episode. Notebook configuration may still call the experiment-level quantity max_risk before passing it into that constructor.
• The runtime feasibility check is therefore about eventual unsafe reachability, not about one-step unsafe probability on the current transition.
• Safety remains a hard constraint. Reward never trades against the budget: it ranks only actions whose eventual-unsafe demand is within the current budget. Across considered OM levels, the default reward score is the posterior expectation; shield_reward_mode=robust retains the minimum-reward ablation.
• TransitionShield.action_mask(agent_id) exposes that admissible set before an on-policy learner samples. IPPO stores each mask with its transition and uses it again when evaluating PPO likelihood ratios. A post-step execution rewrite is an error for IPPO, since it would invalidate the on-policy log probability.
• External policies can still propose an unsafe action. The wrapper applies the reward-aware replacement rule and exposes the executed action through last_executed_actions; this remains valid for off-policy IQL replay.
• Some repeated games can be infeasible at reset under a low initial budget even when their immediate unsafe probability looks moderate. The old Bertrand and congestion labels behaved that way before their env-specific repairs; see repeated-matrix-feasibility.md.
• The level floor still never decreases within a training run.
• The optional credible_floor mode chooses that monotone floor from a posterior-credible suffix and remains robust within the suffix. Its summable and pointwise guarantees are documented in credible-tail-floor.md.
• A floor raise does not directly change the numeric budget scalar; it only changes which actions are considered safe at the next decision.
• Posterior updates still happen after the step from the realized opponent action, using the pre-step base public state bits.
• vi_levels= is still accepted for tests, but only through a synthetic fallback kernel. Real runs should build a full bundle from the transition graph path.
• This wrapper is intentionally not paper-faithful: it is a robust action-replacement shield over an augmented (state, budget, floor) training state, not the paper’s richer encoded-action shielded MDP.
• The bundle semantics version is reward_aware_action_consistent_v4; older cached bundles do not contain action rewards and must be rebuilt.