1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
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()
|