Frozen final-policy safety evaluation¶

Purpose¶

Training-window V/U metrics average a sequence of changing policies. They do not estimate the safety of the policy at the final optimizer update. The final policy tooling adds two distinct post-training measurements:

• IID live-environment rollouts of the frozen joint policy conditional on one fixed post-training controller snapshot, reporting reward, violations per step, survival, and P(T_unsafe <= H) with Wilson intervals at declared horizons;
• finite- and infinite-horizon unsafe reachability of a stationary frozen policy induced on the exact environment transition graph.

Neither quantity should be described as CPO acquiring an anytime guarantee. The exact result is a property of one frozen policy and exact benchmark model. It says nothing about safety during training or robustness to later opponent drift.

Checkpoint schema¶

src.rl.checkpoints owns schema version 3. Each checkpoint records:

• algorithm name, trainer seed, and algorithm configuration;
• a restorable final inference policy for every agent;
• the corresponding optimizer/training state in host-portable arrays.
• the post-training state of a stateful controller, when present. For OmSh this includes the opponent posterior, monotone level floor, floor-patience counters, opponent-observation count, and credible-schedule state.

PyTorch ICPO/CPO checkpoints store the actor architecture and nonempty module state dictionaries. JAX IPPO, Lagrangian, and OmSh checkpoints store the Flax module kind, architecture, adapter contract, parameters, and optimizer state. Schema-2 actor checkpoints remain loadable for stationary baselines, but they cannot reproduce a deployed OmSh controller. Loading a legacy unversioned checkpoint raises an explicit replay instruction. In particular, the old 96-byte ICPO files contain empty params mappings and cannot be repaired after the fact.

The experiment runner can additionally evaluate a final policy immediately after training without adding the evaluation episodes to its training history:

final_policy_eval_episodes=1000
final_policy_eval_horizons=(1,10,50,100,200)
final_policy_eval_seed_offset=1000000
final_policy_eval_confidence=0.95

Per-run results are written below experiments/runs/<tag>/final_policy/<artifact>_run<N>.json. Evaluation time is recorded separately from training time.

JAX policy adapters retain the original bound Flax apply function for checkpoint architecture discovery and cache a JIT-compiled batch-one inference path for live evaluation. The observation shape is fixed by the adapter contract, so repeated episode steps reuse one compilation without changing logits or action sampling.

The live evaluator restores the saved controller snapshot before every seeded episode and restores it once more on exit. Ordinary OmSh environment resets do not reset the online opponent posterior or monotone floor, so resetting only the environment would create one adaptive sequence rather than IID evaluation replicates. Correct reports use live-evaluation schema 2, record episode_sampling=iid_controller_snapshot for OmSh (or iid_environment_reset for stationary baselines), and include a stable digest of the controller snapshot. Reports without this protocol must not be pooled with corrected results, and their Wilson intervals are not valid as IID replicate intervals.

Exact post-processing¶

After a tagged condition has produced corrected checkpoints, run:

uv run python scripts/evaluate_final_policy_checkpoints.py \
  congestion final_cpo_risk_v1 \
  --artifact-stem icpo_histories \
  --episodes 1000 \
  --horizons 1 10 50 100 200 \
  --reset-samples 10000

The post-processor recreates the notebook environment contract, loads or builds the cached exact transition graph, samples the declared reset distribution, restores every seed's final policy, and writes JSON plus Markdown beneath the tag's final_policy/ directory.

For a complete final tag, scripts/evaluate_final_campaign_baselines.py runs the same validation and exact post-processing sequentially for IPPO, IPPO-Lagrangian, and CPO. This is the preferred dependency-job entry point on the cluster because one job can wait for all three training conditions and does not load three copies of a large exact graph concurrently. Pass horizons=1,2,10,64 through py_cpu.sh for an environment with a 64-step cap; the script accepts comma-separated runner values as well as ordinary space-separated CLI values.

When the training run already wrote matching per-run live evaluations, the postprocessor reuses them and computes only the exact graph quantities. Pass --rerun-live to deliberately generate a second independent live sample. The reuse check includes the trainer seed, exact evaluation-seed block, evaluation protocol, and controller-state digest; a same-shaped report from a different controller or sample is not silently reused. For py_cpu.sh, use the equivalent named forms environment-name=congestion tag=final_cpo_risk_v1; that runner translates only key=value arguments.

Interrupted-publication recovery¶

If training wrote a versioned checkpoint but the notebook failed before its frozen-policy JSON or completion manifest was published, recover the declared live rollout without rebuilding an exact graph:

uv run python scripts/evaluate_final_policy_checkpoints.py \
  congestion screen_seed1601 \
  --artifact-stem ippo_shielded_credible_floor_summable_delta_0p05_histories \
  --episodes 1000 \
  --horizons 1 10 50 100 200 \
  --evaluation-seed-offset 1001000 \
  --live-only

The recovered per-run JSON records the algorithm and trainer seed from the checkpoint. It uses the requested evaluation-seed block, so an independently trained seed shard can be retagged as the corresponding serial run without reusing rollout randomness.

For an environment retired from the active OmSh tree, pass --legacy-project-root pointing to the source-pinned worktree that created the checkpoint. The CLI loads only that worktree's environment registry and extends the src package search path with its source tree; checkpoint loading, evaluation logic, and output validation remain the current implementation. The resolved worktree path is recorded in the aggregate recovery report.

scripts/merge_final_policy_shards.py --recover-missing-manifests can recover a missing shard manifest only when the versioned checkpoint, live report, trainer seed, and algorithm agree. Recovery also requires an explicit --risk-threshold; the generated manifest records that training diagnostics are unavailable because notebook publication was interrupted. The default merge remains strict and requires native completion manifests.

Cluster runners that translate only key=value arguments should use evaluation_mode=live instead of the flag-form --live-only, and manifest-recovery=missing instead of the flag-form --recover-missing-manifests.

An OmSh actor checkpoint cannot be evaluated in the bare labelled environment: its observation contract contains shield state, and executed actions require the carried budget, opponent posterior, and replacement rule. Use scripts/recover_omsh_final_policy_checkpoint.py for an interrupted OmSh publication. That entry point bootstraps the source-pinned worktree before any project imports, rebuilds the original transition-shield wrapper from its cached WM/OM bundle, verifies the checkpoint trainer seed and source commit, and then runs the declared IID controller-snapshot rollout. Recovery requires a schema-3 checkpoint with the post-training controller state; an older adaptive evaluation cannot be retrofitted into IID replicates. The script records the complete shield recovery configuration in the output JSON. Evaluating only the neural actor against the bare environment is intentionally rejected by observation shape checks and must not be reported as OmSh.

For source-compatibility audits, --output-path writes the replay beside the canonical report rather than replacing it. Compare the two JSON payloads after removing only the replay-specific omsh_recovery provenance block; a seeded replay is expected to match exactly.

The aggregate report links each per-run live JSON and omits its large episodes_detail array; those details remain in the linked per-run artifact. This avoids duplicating hundreds of thousands of reward entries while keeping the horizon summaries and exact results self-contained.

The exact solver keeps transition deficits explicit. It reports lower and upper risk values rather than silently renormalizing a partial row. Complete exact graphs should have zero gaps. Eventual reachability is solved with sparse monotone lower and upper iterates after identifying public states that cannot reach unsafe and closed safe recurrent classes. States with no reachable closed safe class have eventual unsafe probability one without requiring a long numerical limit. Sparse matrix-vector products keep the 200-step curves tractable on the large Gathering and Markov Stag Hunt graphs. For an incomplete graph, lower and upper iterates may reach distinct fixed points. The report then sets solver_fixed_point=true but solver_converged=false and preserves the nonzero uncertainty interval rather than iterating fruitlessly to the iteration cap.

Restored float32 policy distributions are accepted within 1e-6 of unit mass and then normalized in float64 before composing the joint policy. Material probability deficits still raise instead of being hidden as environment risk. After composing an exact policy row, only residuals at or below 1e-12 are treated as arithmetic roundoff and normalized; their maximum and state count are reported. Larger deficits remain explicit missing mass in the lower/upper risk interval.

Start-distribution aggregation can likewise produce values a few ulps outside [0, 1] even though every state probability is bounded. Aggregate corrections at or below 1e-12 are clipped and reported as max_aggregate_probability_roundoff_correction; a larger excursion raises instead of being hidden.

Policy-kernel construction preallocates compact NumPy coordinate arrays for the exact graph's edges, then lets CSR construction aggregate duplicate successors. It does not materialize one Python dictionary per state plus a second set of Python coordinate lists. This bounds the additional edge storage at roughly 16 bytes per edge with 32-bit state indices and float64 mass, which is important for the large Markov Stag Hunt graph.

OmSh boundary¶

A plain CPO actor is stationary on the public observation, so it induces a finite Markov chain directly. The executed OmSh policy generally does not: its observation and action replacement depend on the carried successor budget, opponent posterior, and monotone level floor. The exact public graph therefore does not contain enough state to reconstruct OmSh execution.

The evaluator checks observation dimensions and rejects this mismatch instead of evaluating the unshielded neural actor and labelling that result "OmSh". For the current paper, report these quantities separately:

1. frozen-policy live-environment CPO risk;
2. exact-model frozen-policy CPO reachability;
3. frozen OmSh empirical risk, when the complete shield state is preserved;
4. OmSh's anytime certificate inside the learned WM/OM-induced model.

An exact executed-OmSh chain is future work requiring a finite representation of budget, posterior, and floor in the graph state.