File size: 4,680 Bytes
026f63e | 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 | """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
# model/training settings come from the checkpoint (DAgger fine-tunes *that* model);
# the dagger config only contributes its `dagger`, `env` and `eval` sections
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: # resume: finished iterations are kept, their data is rebuilt below
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)
# probs are aligned with ACTIONS order whatever the labels ("UP" or "move up")
policy = lambda states: dec.probs(states, QUESTION, list(dec.task.default_choices)) # noqa: E731
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"): # evaluation splits stay the original, untouched ones
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
|