results: close BabyAI population predictor rescue

1 files changed, 86 insertions, 0 deletions

diff --git a/experiments/analyze_babyai_population_p3.py b/experiments/analyze_babyai_population_p3.py
new file mode 100644
index 0000000..554447c
--- /dev/null
+++ b/experiments/analyze_babyai_population_p3.py
@@ -0,0 +1,86 @@
+#!/usr/bin/env python3
+"""Aggregate the frozen BabyAI population-predictor capacity screen."""
+
+import argparse
+import json
+from pathlib import Path
+import statistics
+
+
+ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_RESULTS = ROOT / "results" / "babyai_shared" / "population_p3"
+DEFAULT_OUT = ROOT / "results" / "babyai_shared" / "population_p3_analysis.json"
+SEEDS = (4101, 4102, 4103)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--results", type=Path, default=DEFAULT_RESULTS)
+    parser.add_argument("--out", type=Path, default=DEFAULT_OUT)
+    args = parser.parse_args()
+    records = []
+    for seed in SEEDS:
+        with open(args.results / f"seed{seed}.json", encoding="utf-8") as handle:
+            records.append(json.load(handle))
+    rows = [
+        {
+            "seed": record["model_seed"],
+            "epoch": audit["epoch"],
+            "layer": layer["layer"],
+            **layer["holdout"],
+        }
+        for record in records for audit in record["audits"]
+        for layer in audit["layers"]
+    ]
+    checks = {
+        "all_seed_gates_pass": all(record["gate"] == "pass"
+                                    for record in records),
+        "r2_at_least_0p8_every_seed_epoch_layer": all(
+            row["mean_per_cell_r2"] >= 0.8 for row in rows),
+        "residual_ratio_at_most_0p25_every_seed_epoch_layer": all(
+            row["residual_context_rms_ratio"] <= 0.25 for row in rows),
+        "zero_action_or_teaching_observations": all(
+            layer["action_observations"] == 0
+            and layer["teaching_observations"] == 0
+            for record in records for audit in record["audits"]
+            for layer in audit["layers"]),
+    }
+    epoch_layer = []
+    for epoch in (0, 1, 5, 10, 20, 40):
+        for layer in range(4):
+            selected = [row for row in rows
+                        if row["epoch"] == epoch and row["layer"] == layer]
+            epoch_layer.append({
+                "epoch": epoch,
+                "layer": layer,
+                "mean_holdout_r2": statistics.mean(
+                    row["mean_per_cell_r2"] for row in selected),
+                "mean_holdout_residual_ratio": statistics.mean(
+                    row["residual_context_rms_ratio"] for row in selected),
+            })
+    worst_r2 = min(rows, key=lambda row: row["mean_per_cell_r2"])
+    worst_ratio = max(
+        rows, key=lambda row: row["residual_context_rms_ratio"])
+    report = {
+        "stage": "babyai_population_p3_analysis",
+        "gate": "pass" if all(checks.values()) else "fail",
+        "checks": checks,
+        "minimum_holdout_r2": worst_r2,
+        "maximum_holdout_residual_context_rms_ratio": worst_ratio,
+        "epoch_layer_means_across_seeds": epoch_layer,
+        "raw_sdil_rollout_or_test_outcomes_read": False,
+        "decision": (
+            "Close the population-linear predictor rescue. It is highly "
+            "predictive at initialization but fails the frozen early-layer "
+            "threshold after clean task learning; do not run another "
+            "PickupLoc downstream endpoint."),
+    }
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+    with open(args.out, "w", encoding="utf-8") as handle:
+        json.dump(report, handle, indent=2, sort_keys=True)
+        handle.write("\n")
+    print(json.dumps(report, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()