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
"""Apply the frozen BabyAI B0 clean-selector rule."""
import argparse
import json
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
DEFAULT_RESULTS = ROOT / "results" / "babyai_shared" / "b0"
DEFAULT_OUT = ROOT / "results" / "babyai_shared" / "b0_selector.json"
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()
rows = []
for depth in (2, 4):
for learning_rate in (0.01, 0.03):
records = {}
for condition in ("bp", "clean_kp"):
path = args.results / (
f"d{depth}_lr{learning_rate}_{condition}.json")
with open(path, encoding="utf-8") as handle:
records[condition] = json.load(handle)
bp = records["bp"]
kp = records["clean_kp"]
bp_success = float(bp["rollout"]["success"])
kp_success = float(kp["rollout"]["success"])
bp_lesion = float(bp["mission_lesion_rollout"]["success"])
checks = {
"both_finite": bool(bp["finite"] and kp["finite"]),
"bp_success_at_least_0p8": bp_success >= 0.8,
"clean_kp_success_at_least_0p8": kp_success >= 0.8,
"bp_mission_lesion_drop_at_least_0p2": (
bp_success - bp_lesion >= 0.2),
}
rows.append({
"hidden_layers": depth,
"learning_rate": learning_rate,
"bp_rollout_success": bp_success,
"clean_kp_rollout_success": kp_success,
"bp_mission_lesion_success": bp_lesion,
"bp_mission_lesion_drop": bp_success - bp_lesion,
"bp_action_accuracy": float(bp["validation"]["accuracy"]),
"clean_kp_action_accuracy": float(
kp["validation"]["accuracy"]),
"eligible": all(checks.values()),
"checks": checks,
"source_files": [str(
(args.results / f"d{depth}_lr{learning_rate}_{condition}.json")
.relative_to(ROOT)) for condition in ("bp", "clean_kp")],
})
eligible = [row for row in rows if row["eligible"]]
selected = max(eligible, key=lambda row: (
row["clean_kp_rollout_success"], row["clean_kp_action_accuracy"],
-row["hidden_layers"], -row["learning_rate"])) if eligible else None
report = {
"stage": "babyai_shared_b0_selector",
"gate": "pass" if selected is not None else "fail",
"selection_rule": (
"highest clean-KP rollout success, then action accuracy, then "
"fewer layers, then smaller learning rate, among eligible rows"),
"candidates": rows,
"selected": ({
"hidden_layers": selected["hidden_layers"],
"width": 256,
"learning_rate": selected["learning_rate"],
"context_gain": 1.0,
"b1_epochs": 40,
"b1_model_and_shuffle_seeds": [4101, 4102, 4103],
} if selected is not None else None),
"raw_or_sdil_results_read": False,
"test_split_generated_or_read": False,
}
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()
|
