| """DAgger for Snake (Ross et al., 2011). |
| |
| Behaviour cloning only sees states the oracle visits; once the student makes a |
| mistake it lands in states it never trained on and errors compound. DAgger lets |
| the student play, labels *its* states with the oracle, aggregates them with the |
| original data and fine-tunes. Cheap (no reward engineering, the oracle is |
| free) and directly targets the student's own failure modes. |
| """ |
| from __future__ import annotations |
|
|
| import copy |
| import json |
| import shutil |
| from pathlib import Path |
|
|
| from minijev.datasets.decision_data import DecisionData |
| from minijev.inference.runtime import TorchDecider |
| from minijev.training.supervised import train |
|
|
| from .dataset import QUESTION, collect_with_policy |
|
|
|
|
| def run_dagger(cfg: dict, checkpoint: str, iterations: int = 2, samples: int = 100_000, |
| run_dir: str | None = None) -> list[dict]: |
| import torch |
|
|
| |
| |
| ck_cfg = torch.load(checkpoint, map_location="cpu", weights_only=False)["config"] |
| cfg = {**copy.deepcopy(ck_cfg), **{k: cfg[k] for k in ("dagger", "env", "eval") if k in cfg}} |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| root = Path(run_dir or f"runs/{ck_cfg.get('name', 'snake')}-dagger") |
| root.mkdir(parents=True, exist_ok=True) |
| base_dir = Path(cfg["data"]["dir"]) |
| agg = DecisionData.load(base_dir / "train.npz") |
| dcfg = cfg.get("dagger", {}) |
| board_sizes = cfg["data"].get("generate", {}).get("board_sizes", [cfg.get("env", {}).get("board_size", 10)]) |
| history = json.loads((root / "dagger.json").read_text()) if (root / "dagger.json").exists() else [] |
| for rec in history: |
| checkpoint = rec["checkpoint"] |
| if history: |
| agg = DecisionData.load(root / f"data-it{history[-1]['iteration']}" / "train.npz") |
| print(f"[dagger] resuming after iteration {history[-1]['iteration']}") |
| for it in range(len(history) + 1, iterations + 1): |
| dec = TorchDecider(checkpoint, threads=4, device=device) |
| |
| policy = lambda states: dec.probs(states, QUESTION, list(dec.task.default_choices)) |
| print(f"[dagger] iteration {it}: collecting {samples:,} student states ...") |
| new = collect_with_policy(policy, samples, seed=1000 + it, board_sizes=board_sizes, |
| keep_prob=dcfg.get("keep_prob", 0.2), explore_eps=dcfg.get("explore_eps", 0.02), |
| env_cfg={k: v for k, v in cfg.get("env", {}).items() if k != "board_size"}, |
| game_id_base=10**9 * it, parallel=256) |
| agg = DecisionData.concat([agg, new]) |
| data_dir = root / f"data-it{it}" |
| data_dir.mkdir(exist_ok=True) |
| agg.save(data_dir / "train.npz") |
| for split in ("val", "test"): |
| shutil.copy(base_dir / f"{split}.npz", data_dir / f"{split}.npz") |
| c = copy.deepcopy(cfg) |
| c["name"] = f"{cfg.get('name', 'snake')}-dagger{it}" |
| c["data"] = {"dir": str(data_dir)} |
| c["training"]["init_checkpoint"] = checkpoint |
| c["training"]["epochs"] = dcfg.get("epochs", 1) |
| c["training"]["lr"] = dcfg.get("lr", cfg["training"].get("lr", 1e-3) * 0.5) |
| c["training"]["warmup_steps"] = 100 |
| it_dir = root / f"it{it}" |
| summary = train(c, run_dir=it_dir, resume=(it_dir / "checkpoints" / "last.pt").exists()) |
| checkpoint = summary["final_checkpoint"] |
| from .evaluate import model_policy, play_games |
| games = play_games(model_policy(TorchDecider(checkpoint, threads=4, device=device)), |
| cfg.get("eval", {}).get("games", 200), cfg.get("env", {}), parallel=200) |
| (root / f"it{it}" / "games.json").write_text(json.dumps(games, indent=2)) |
| rec = {"iteration": it, "rows": len(agg), "new_rows": len(new), "checkpoint": checkpoint, |
| "test_acc": summary["test_calibrated"]["accuracy"], "test_ece": summary["test_calibrated"]["ece"], |
| "avg_score": games["avg_score"], "fatal_move_rate": games["fatal_move_rate"]} |
| print(f"[dagger] it {it}: avg score {games['avg_score']:.1f} | test acc {rec['test_acc']:.4f} " |
| f"| fatal-move {games['fatal_move_rate']:.4%}") |
| history.append(rec) |
| (root / "dagger.json").write_text(json.dumps(history, indent=2)) |
| return history |
|
|
