Shielding API¶

Shielding APIs cover the learned transition shield, conservative value iteration, runtime inspection, and exact true-shield baselines.

Learned Transition Shield¶

src.shield.transition_shield.TransitionShield ¶

TransitionShield(
    env: ParallelEnv,
    *,
    agent_id: str = "player_0",
    init_safety_bound: float,
    bundle: TransitionShieldBundle | None = None,
    transition_graph_path: str | None = None,
    iop=None,
    vi_levels: Sequence[SoundVIResult] | None = None,
    focal_idx: Optional[int] = None,
    binarize_threshold: float = 0.5,
    debug_mode: str = "minimal",
    debug_trace_limit: int = 256,
    verbose_data: bool = False,
    data_trace_stride: int = 100,
    max_trace_rows: int = 100000,
    missing_coverage_log_limit: int = 5,
    vi_tol: float = 1e-08,
    vi_max_iter: int = 1000000,
    vi_round_outward: float = 1e-12,
    budget_granularity: int = 20,
    patience: int = 64,
    opponent_mode: str = MONOTONE_FLOOR_MODE,
    reward_mode: str = BAYESIAN_REWARD_MODE,
    credible_delta: float = 0.05,
    credible_schedule: str = SUMMABLE_CREDIBLE_SCHEDULE,
)

Bases: BaseParallelWrapper

Budget-augmented opponent-model transition shield.

Runtime behavior

use the monotone OM-level floor by default
optionally run credible-floor, fixed all-level, or Bayesian-mixture ablations
expose [flat_obs, budget, floor], plus the posterior in Bayesian mode
keep the focal action space identical to the base environment
keep a monotone floor over opponent-model levels using argmax(iop.mix_alpha) with optional sample-based raise patience
expose admissible actions to on-policy trainers before they sample
replace only unsafe external proposals with the highest-reward robust-safe action
update the remaining budget from the realized successor state's projected budget

last_executed_actions `property` ¶

last_executed_actions: dict[str, Any] | None

Actions applied to the wrapped environment by the most recent step.

action_mask ¶

action_mask(agent: str) -> np.ndarray

Expose state-dependent admissibility without changing action spaces.

On-policy trainers use this before sampling, making the sampled, stored, and executed primitive actions identical. Non-focal agents remain unrestricted.

get_evaluation_state ¶

get_evaluation_state() -> dict[str, Any]

Return the persistent controller state needed for IID policy replay.

Episode-local budget and decision state are deliberately omitted: a live evaluation starts every replicate with :meth:reset, which reinstates init_safety_bound. The posterior, monotone floor, and patience counters persist across ordinary resets and therefore must be snapshotted explicitly before evaluating independent replicates.

set_evaluation_state ¶

set_evaluation_state(state: Mapping[str, Any]) -> None

Restore a state produced by :meth:get_evaluation_state.

src.shield.transition_shield.TransitionShieldBundle `dataclass` ¶

TransitionShieldBundle(
    agent_id: str,
    focal_idx: int,
    n_actions: int,
    levels: tuple[ShieldLevelBuildResult, ...],
    transition_graph_path: str | None = None,
    semantics_version: str = SHIELD_SEMANTICS_VERSION,
    floor_budget_lower_bounds: tuple[dict[str, float], ...]
    | None = None,
)

src.shield.transition_shield.ShieldLevelBuildResult `dataclass` ¶

ShieldLevelBuildResult(
    level: int,
    vi: SoundVIResult,
    n_states: int,
    n_edges: int,
    successor_kernel: dict[
        str, tuple[SuccessorDistribution, ...]
    ],
    budget_lower_bounds: dict[str, float],
    action_rewards: dict[str, tuple[float, ...]] = dict(),
)

src.shield.transition_shield.SuccessorDistribution `dataclass` ¶

SuccessorDistribution(
    successor_keys: tuple[str, ...],
    probabilities: tuple[float, ...],
)

src.shield.transition_shield.build_transition_shield_bundle ¶

build_transition_shield_bundle(
    *,
    env: ParallelEnv,
    transition_graph_path: str,
    iop,
    agent_id: str = "player_0",
    focal_idx: Optional[int] = None,
    vi_tol: float = 1e-08,
    vi_max_iter: int = 1000000,
    vi_round_outward: float = 1e-12,
) -> TransitionShieldBundle

Parameterized Learned Shield¶

src.shield.parameterized_transition_shield.ParameterizedTransitionShield ¶

ParameterizedTransitionShield(
    env,
    *,
    action_mode: str,
    projection_eps: float = 1e-08,
    **kwargs,
)

Bases: TransitionShield

OmSh wrapper whose learned action includes a safe mixture and/or budgets.

The focal policy emits a Dict action with categorical and raw-Gaussian components. This wrapper deterministically maps that augmented action to a robustly feasible primitive distribution and successor certificate, then samples the primitive environment action. The projection is therefore part of the augmented environment rather than an unrecorded PPO rewrite.

augmented_action_masks ¶

augmented_action_masks(
    agent: str,
) -> tuple[np.ndarray, np.ndarray]

Return padded categorical and continuous masks for exact PPO scoring.

src.shield.augmented_action.project_edge_mixture ¶

project_edge_mixture(
    *,
    beta: Sequence[float] | ndarray,
    per_level_action_probabilities: Sequence[ndarray]
    | ndarray,
    budget: float,
    action_i: int,
    action_j: int,
    proposed_mix: float,
    fallback_action: int | None = None,
    eps: float = DEFAULT_PROJECTION_EPS,
) -> EdgeMixtureProjection

Project a proposed point on an action edge into every level constraint.

The returned distribution is (1-x) e_i + x e_j. Unlike masking the two endpoints separately, intersecting the level-wise intervals can retain a mixture whose endpoints violate different robust constraints.

src.shield.augmented_action.project_successor_budget_logits ¶

project_successor_budget_logits(
    *,
    beta: Sequence[float] | ndarray,
    level_distributions: Sequence[ndarray] | ndarray,
    budget: float,
    raw_logits: Sequence[float] | ndarray,
    eps: float = DEFAULT_PROJECTION_EPS,
) -> SuccessorBudgetProjection

Safely project learned raw budget logits along their proposed direction.

src.shield.augmented_action.robust_egalitarian_fill ¶

robust_egalitarian_fill(
    *,
    beta: Sequence[float] | ndarray,
    level_distributions: Sequence[ndarray] | ndarray,
    budget: float,
    eps: float = 1e-06,
) -> np.ndarray

Max-min additive successor-budget fill under every level constraint.

Sound Value Iteration¶

src.shield.sound_value_iteration.SoundVIResult `dataclass` ¶

SoundVIResult(
    q: ndarray,
    lower_q: ndarray,
    action_risk: Dict[tuple[str, int], float],
    action_risk_lower: Dict[tuple[str, int], float],
    keys: List[str],
    key_to_idx: Dict[str, int],
    n_actions: int,
    unsafe_keys: frozenset[str],
    almost_sure_safe_keys: frozenset[str],
    iterations: int,
    converged: bool,
    max_gap: float,
    tol: float,
    missing_mass: ndarray,
    transition_mass: ndarray,
    successor_counts: ndarray,
    n_missing_state_actions: int,
    n_partial_state_actions: int,
    max_missing_mass: float,
)

src.shield.sound_value_iteration.sound_value_iteration_min_reach_upper ¶

sound_value_iteration_min_reach_upper(
    *,
    graph: Graph,
    unsafe: Iterable[str],
    n_actions: Optional[int] = None,
    tol: float = 1e-08,
    max_iter: int = 1000000,
    round_outward: float = 1e-12,
) -> SoundVIResult

Compute sound lower/upper bounds on the minimal probability of eventually reaching the unsafe set U.

Conservative modeling: any missing transition mass is sent to an absorbing unsafe sink (value 1), ensuring over-approximation for the upper bound.

The interval is initialized with

lower_q = 0 outside unsafe,
upper_q = 1 outside unsafe,
lower_q = upper_q = 1 on unsafe,
lower_q = upper_q = 0 on the almost-sure safe set.

Then both bounds are iterated monotonically until the sup-norm gap between them is within tol, or max_iter is reached.

src.shield.sound_value_iteration.unsafe_keys_from_env ¶

unsafe_keys_from_env(
    *,
    graph: Graph,
    labelled_env,
    agent_id: str | None = None,
) -> set[str]

Label graph nodes through a labelled environment's safety model.

src.shield.sound_value_iteration.unsafe_keys_from_abstraction ¶

unsafe_keys_from_abstraction(
    *,
    graph: Graph,
    safety_model: SafetyAbstraction[Any],
    agent_id: str,
) -> set[str]

Label graph nodes with an explicit environment safety abstraction.

src.shield.sound_value_iteration.almost_sure_safe_keys_from_graph ¶

almost_sure_safe_keys_from_graph(
    *,
    graph: Graph,
    unsafe: Iterable[str],
    n_actions: Optional[int] = None,
    mass_tol: float = 1e-12,
) -> set[str]

Largest safe set W such that every state in W has at least one action whose full support stays inside W and has no missing transition mass.

These are the states from which the controller can stay safe almost surely, so their minimal reachability to unsafe is exactly zero.

Runtime Inspection¶

src.shield.inspection.ShieldLevelInspection `dataclass` ¶

ShieldLevelInspection(
    level: int,
    key: str,
    has_key: bool,
    q: float,
    lower_q: float,
    gap: float,
    max_gap: float,
    unsafe: bool,
    almost_sure_safe: bool,
    action_risks: tuple[float, ...],
    action_risk_lower: tuple[float, ...],
    immediate_action_risks: tuple[float, ...],
    missing_mass: tuple[float, ...],
    transition_mass: tuple[float, ...],
    successor_counts: tuple[int, ...],
    successor_keys: tuple[tuple[str, ...], ...],
    successor_probabilities: tuple[tuple[float, ...], ...],
)

src.shield.inspection.ShieldStateInspection `dataclass` ¶

ShieldStateInspection(
    key: str,
    current_budget: float,
    level_floor: int,
    levels: tuple[int, ...],
    has_key_each: tuple[bool, ...],
    missing_coverage: bool,
    successor_keys: tuple[str, ...],
    beta: ndarray,
    per_level_action_probs: tuple[ndarray, ...],
    pure_action_costs: ndarray,
    pure_action_immediate_costs: ndarray,
    pure_action_rewards: ndarray,
    robust_action_costs: tuple[float, ...],
    robust_action_immediate_costs: tuple[float, ...],
    robust_action_rewards: tuple[float, ...],
    worst_action_levels: tuple[int, ...],
    admissible_actions: tuple[int, ...],
    safe_action: int,
    safe_action_feasible: bool,
    mix_alpha: tuple[float, ...],
    level_weights: tuple[float, ...],
    opponent_mode: str,
    reward_mode: str,
    credible_delta: float | None,
    credible_excluded_mass: float | None,
    credible_retained_mass: float | None,
    posterior_tail_constraint_held: bool | None,
    credible_observation_count: int | None,
    convergence_gaps: tuple[float, ...],
)

src.shield.inspection.ShieldDecisionInspection `dataclass` ¶

ShieldDecisionInspection(
    state: ShieldStateInspection,
    key: str,
    proposed_action: int,
    final_action: int,
    proposed_risk: float,
    final_risk: float,
    proposed_immediate_risk: float,
    final_immediate_risk: float,
    proposed_reward: float,
    final_reward: float,
    overrode: bool,
    levels_considered: tuple[int, ...],
    level_floor: int,
    missing_coverage: bool,
    level_has_key: tuple[bool, ...],
    decision_reason: str,
    aggregated_risks: tuple[float, ...],
    aggregated_immediate_risks: tuple[float, ...],
    aggregated_rewards: tuple[float, ...],
    admissible_actions: tuple[int, ...],
    winning_levels: tuple[int, ...],
    mix_alpha: tuple[float, ...],
    level_weights: tuple[float, ...],
    opponent_mode: str,
    reward_mode: str,
    convergence_gaps: tuple[float, ...],
    current_budget: float,
    next_budget: float,
    successor_keys: tuple[str, ...],
    successor_budgets: tuple[float, ...],
)

src.shield.inspection.inspect_shield_level ¶

inspect_shield_level(
    bundle, *, key: str, level: int
) -> ShieldLevelInspection

src.shield.inspection.inspect_shield_state ¶

inspect_shield_state(
    bundle,
    *,
    key: str,
    current_budget: float,
    mix_alpha: Sequence[float] | ndarray | None,
    current_floor: int = 0,
    manual_floor: int | None = None,
    opponent_mode: str = MONOTONE_FLOOR_MODE,
    reward_mode: str = BAYESIAN_REWARD_MODE,
    credible_delta: float = 0.05,
    credible_observation_count: int = 0,
    warn_missing_coverage=None,
) -> ShieldStateInspection

src.shield.inspection.inspect_shield_decision ¶

inspect_shield_decision(
    bundle,
    *,
    key: str,
    proposed_action: Any,
    current_budget: float,
    mix_alpha: Sequence[float] | ndarray | None,
    current_floor: int = 0,
    manual_floor: int | None = None,
    opponent_mode: str = MONOTONE_FLOOR_MODE,
    reward_mode: str = BAYESIAN_REWARD_MODE,
    credible_delta: float = 0.05,
    credible_observation_count: int = 0,
    warn_missing_coverage=None,
) -> ShieldDecisionInspection

src.shield.inspection.credible_delta_at_step ¶

credible_delta_at_step(
    credible_delta: float,
    *,
    observation_count: int,
    schedule: str = SUMMABLE_CREDIBLE_SCHEDULE,
) -> float

Return the pointwise tail-mass allowance for an OM observation count.

summable allocates a total error budget using 6 * delta / (pi^2 * (t + 1)^2). fixed retains the same pointwise allowance and therefore has no time-uniform infinite-horizon guarantee.

src.shield.inspection.resolve_credible_tail_floor ¶

resolve_credible_tail_floor(
    *,
    mix_alpha: Sequence[float] | ndarray | None,
    n_levels: int,
    credible_delta: float,
    current_floor: int = 0,
    manual_floor: int | None = None,
) -> int

Resolve the largest monotone floor retaining 1 - credible_delta mass.

True Shield¶

src.true_shield.wrapper.TruePolicyTransitionShield ¶

TruePolicyTransitionShield(
    env: ParallelEnv,
    *,
    transition_graph_path: str,
    transition_graph: Graph | None = None,
    max_risk: float,
    agent_id: str = "player_0",
    focal_idx: Optional[int] = None,
    refresh_interval_updates: int = 1,
    binarize_threshold: float = 0.5,
    debug_mode: str = "minimal",
    debug_trace_limit: int = 256,
    verbose_data: bool = False,
    data_trace_stride: int = 100,
    max_trace_rows: int = 100000,
    missing_coverage_log_limit: int = 5,
    vi_tol: float = 1e-08,
    vi_max_iter: int = 1000000,
    vi_round_outward: float = 1e-12,
    budget_granularity: int = 20,
    patience: int = 64,
    initial_bundle: TransitionShieldBundle | None = None,
)

Bases: TransitionShield

src.true_shield.bundle.build_true_policy_transition_shield_bundle ¶

build_true_policy_transition_shield_bundle(
    *,
    env: ParallelEnv,
    opponent_policies: Sequence[Any],
    transition_graph_path: str | None = None,
    transition_graph: Graph | None = None,
    agent_id: str = "player_0",
    focal_idx: Optional[int] = None,
    vi_tol: float = 1e-08,
    vi_max_iter: int = 1000000,
    vi_round_outward: float = 1e-12,
) -> TransitionShieldBundle

src.true_shield.exact_graph.build_true_transition_graph ¶

build_true_transition_graph(env, **kwargs) -> Graph

Exact Anytime Certificate¶

src.true_shield.adversarial.ExactAdversarialSafetyCertificate `dataclass` ¶

ExactAdversarialSafetyCertificate(
    version: str,
    validator_version: str,
    agent_id: str,
    focal_idx: int,
    n_actions: int,
    joint_action_sizes: tuple[int, ...],
    keys: tuple[str, ...],
    beta: ndarray,
    unsafe_mask: ndarray,
    robust_action_risk: ndarray,
    robust_action_risk_numerator: tuple[
        tuple[int, ...], ...
    ],
    robust_action_risk_denominator: tuple[
        tuple[int, ...], ...
    ],
    worst_opponent_action: tuple[
        tuple[tuple[int, ...] | None, ...], ...
    ],
    worst_row_id: ndarray,
    admissible_at_beta: ndarray,
    reset_keys: tuple[str, ...],
    reset_support_kind: str,
    reset_support_complete: bool | None,
    reset_support_sample_count: int | None,
    claim_condition: str,
    p_total: float,
    p_shield: float,
    delta_cov: float,
    solver_converged: bool,
    solver_iterations: int,
    solver_gap: float,
    validation: ExactAdversarialValidation,
    graph_hash: str,
    state_encoding_hash: str,
    unsafe_label_hash: str,
    environment_config_hash: str,
)

Checked per-episode arbitrary-opponent reachability certificate.

src.true_shield.adversarial.compile_exact_adversarial_outer_kernel ¶

compile_exact_adversarial_outer_kernel(
    *,
    graph: Graph,
    n_actions: int,
    focal_idx: int,
    joint_action_sizes: Sequence[int] | None = None,
    mass_tolerance: float = 1e-12,
) -> ExactAdversarialOuterKernel

Compile (s, a_i, a_-i) rows without marginalizing opponents.

src.true_shield.adversarial.build_exact_adversarial_safety_certificate ¶

build_exact_adversarial_safety_certificate(
    *,
    graph: Graph,
    unsafe_keys: Iterable[str],
    n_actions: int,
    p_total: float,
    agent_id: str = "player_0",
    focal_idx: int = 0,
    joint_action_sizes: Sequence[int] | None = None,
    environment_config_hash: str = "",
    reset_support_kind: str = "declared",
    reset_support_complete: bool | None = None,
    reset_support_sample_count: int | None = None,
    tolerance: float = 1e-10,
    max_iterations: int = 1000000,
    mass_tolerance: float = 1e-12,
    allow_nonconverged_diagnostic: bool = False,
) -> ExactAdversarialSafetyCertificate

Synthesize and independently validate an exact arbitrary-opponent certificate.

src.true_shield.adversarial.validate_exact_adversarial_certificate ¶

validate_exact_adversarial_certificate(
    *,
    graph: Graph,
    certificate: ExactAdversarialSafetyCertificate,
    unsafe_keys: Iterable[str],
    expected_environment_config_hash: str | None = None,
    mass_tolerance: float = 1e-12,
) -> ExactAdversarialValidation

Audit the artifact with a slow raw-edge implementation.

src.true_shield.adversarial_wrapper.ExactAdversarialTransitionShield ¶

ExactAdversarialTransitionShield(
    env,
    *,
    certificate: ExactAdversarialSafetyCertificate,
    graph: Graph,
    unsafe_keys: Sequence[str],
    validated_environment_config_hash: str,
    action_scorer: ActionScorer | None = None,
)