Skip to content

Shielded Notebook Integration

The marimo experiment notebooks now run the per-level TransitionShield end to end instead of stopping at the old deferred shield placeholder.

Artifact flow

The shielded experiment stage depends on the same artifact chain used elsewhere in the project:

  1. wm/env_transition_graph.pkl
  2. om/iop_stack.pt
  3. cached shield bundle at exports/.../shield/transition_shield_bundle.pkl
  4. shielded experiment histories and figures under exports/.../experiments/

notebooks/utils/shielding.py is the shared orchestration layer for this path.

  • load_or_build_transition_shield_bundle(...) loads the world-model graph and IOP checkpoint, builds the immutable per-level shield bundle when needed, and caches it under the environment export directory.
  • make_shielded_env_factory(...) reuses the cached bundle but reloads a fresh IOP checkpoint for each training run so posterior state does not leak between runs.
  • The shared factory now also probes reset() once during setup so obviously infeasible shield configurations fail in the notebook setup cell instead of surfacing later inside PPO rollout code.
  • summarize_transition_shield_bundle(...) gives the notebook a short markdown summary with cache status, artifact paths, per-level state/edge counts, and VI max-gap diagnostics.
  • shield_experiment_spec(...) gives each uncertainty/reward condition stable run, timing, history, checkpoint, diagnostic, and plot names. The default monotone-floor/Bayesian-reward condition uses the canonical ippo_shielded_* paths. Non-default modes, credible-set parameters, robust reward aggregation, and non-v0 shield versions receive deterministic suffixes so ablation runs cannot silently overwrite one another.

run_jax_stage(...) resolves that identity for every learned-shield stage. Mode-aware notebook cells also use the resolved specification explicitly for reload paths and plot labels.

Cache invalidation

The bundle cache is intentionally conservative. A cached bundle is reused only when all of these match:

  • graph path, file size, and mtime_ns
  • IOP checkpoint path, file size, and mtime_ns
  • current labelled-env signature:
  • public state() width
  • focal action dimension
  • labelled env type
  • hashed safe(...) implementation source
  • shield build parameters: shield_semantics, agent_id, focal_idx, vi_tol, vi_max_iter, vi_round_outward

Anything else forces a rebuild and overwrites the cache sidecar.

Cache validation, rebuilding, and publication are serialized by a process lock next to the bundle. This matters when Monotone and Credible jobs become eligible together: only the first process constructs the common per-level bundle, while later processes wait and then receive a cache hit. Publication uses a temporary file plus atomic replacement, so an interrupted pickle write cannot expose a partial cache; a legacy truncated pickle is treated as a miss and rebuilt while holding the lock.

Learned-versus-exact diagnostic scope

The exact true-shield graphs certify transition dynamics and safety but store zero reward placeholders. Consequently, wm/transition_tv_report.json reports graph-exhaustive transition TV and coverage but marks exact reward MAE as unavailable. Likewise, shield/learned_vs_oracle_report.json reports risk and safe-set agreement plus whether the learned reward-selected action is oracle-safe, but it does not claim oracle reward regret or reward-based action agreement. Reward prediction quality comes from the held-out environment data in wm/wm_report.txt instead.

Runtime learning semantics

TransitionShield is now the budget-augmented replacement shield used directly by training, and shielded notebook runs still update the opponent-level posterior online.

  • The focal observation is the flattened base observation with appended current_safety_bound and level_floor scalars. The Bayesian-mixture ablation also appends the opponent-level posterior.
  • The focal action space is the original base Discrete action space.
  • By default, the current action decision uses the pre-step monotone floor and robust floor-wise certificate V_floor(s') shared across all remaining levels. shield_opponent_mode selects the credible-floor, all-level, and Bayesian-mixture ablations documented in opponent-uncertainty-modes.md.
  • After the environment transition, the shield updates the remaining budget from the realized successor state's projected budget on every step.
  • After the environment step, the wrapper encodes the realized joint opponent action from the action dict and calls iop.update_mixture_with_observation(...) with the pre-step binarized state.
  • The floor can rise after that posterior update, but it never decreases within the same training run, and the scalar budget is not remapped when it rises.
  • Each new notebook training run starts from a fresh iop_stack.pt load, so posterior drift does not carry across runs.

The runtime risk semantics are eventual unsafe reachability, not one-step unsafe probability. That distinction matters in repeated matrix games: a configuration can look plausible at the stage level and still be infeasible at reset once the shield asks for a small upper bound on eventual unsafe reachability.

That repeated-game failure mode still matters conceptually, but the current matrix envs are no longer in the old broken state:

  • Bertrand now treats only price_war as unsafe, so its learned and exact shielded baselines can run again at max_risk=0.2.
  • Congestion now exposes a safe Detour fallback, so exact shielding can run again at max_risk=0.2.
  • Congestion learned shielding still depends on regenerated WM/OM artifacts, because the detour action changed the public action/state contract.

Visual outputs

The existing cumulative and episodic reward/safety figures remain unchanged except that they continue to overlay shield level raises from episode histories.

A dedicated shield_telemetry.png figure is now also saved by each experiment notebook when shielded histories are available. The figure is a fixed 2x2 summary:

  • cumulative overrides
  • cumulative unsafe proposals
  • cumulative missing coverage
  • per-episode proposed/final risk means and maxima

Verbose shield telemetry also includes budget-specific fields such as the current budget, projected next budget, proposed/final focal actions, and projected successor budgets.

Credible-floor runs additionally record their total and per-decision tail allowance, excluded/retained posterior mass, whether the irreversible suffix still meets its posterior-tail constraint, observation count, and the nominal summable additive risk quantity. Episode histories aggregate posterior-tail constraint misses and posterior masses for reporting; these fields do not establish transition coverage or contribute a measured delta_cov.