"""Oracle-labelled Snake datasets. Games are played by a behaviour policy (oracle + a little exploration), every visited state is labelled with the oracle's move, and a random subset of states (`keep_prob`) is kept so a dataset of N rows spans many independent games. Rows carry their game id (`group`) for leakage-free splits and an exact state hash to drop train/test duplicates (e.g. near-identical opening positions). """ from __future__ import annotations import json import multiprocessing as mp import os import sys from concurrent.futures import ProcessPoolExecutor from pathlib import Path import numpy as np from minijev.datasets.decision_data import DecisionData, split_by_group from .agents import make_agent, safe_moves from .env import ACTIONS, SnakeConfig, SnakeEnv, SnakeState from .task import SnakeTask QUESTION = SnakeTask.default_question def state_hash(s: SnakeState) -> int: return hash((s.height, s.width, tuple(s.body), s.food)) def play_and_label(n_samples: int, seed: int, board_sizes=(10,), keep_prob: float = 0.1, explore_eps: float = 0.05, env_cfg: dict | None = None, behaviour=None, labeler: str = "oracle", game_id_base: int = 0): """Collect `n_samples` labelled states. `behaviour(state, rng) -> action` (default: the labeler itself with epsilon-random safe moves). Returns (states, labels, groups).""" rng = np.random.default_rng(seed) oracle = make_agent(labeler) states, labels, groups, optimal = [], [], [], [] game = 0 while len(states) < n_samples: size = int(board_sizes[game % len(board_sizes)]) cfg = SnakeConfig.from_dict({**(env_cfg or {}), "height": size, "width": size}) gid = game_id_base + game env = SnakeEnv(cfg, seed=int(rng.integers(2**31))) while not env.done and len(states) < n_samples: s = env.state() label, good, _ = oracle.analyze(s) if rng.random() < keep_prob: states.append(s) labels.append(label) groups.append(gid) optimal.append(good) if behaviour is not None: action = behaviour(s, rng) elif rng.random() < explore_eps: safe = safe_moves(s) action = safe[int(rng.integers(len(safe)))] if safe else label else: action = label env.step(action) game += 1 return states, labels, groups, optimal def _worker(args): n, seed, sizes, keep, eps, env_cfg, gbase = args return play_and_label(n, seed, sizes, keep, eps, env_cfg, game_id_base=gbase) def to_decision_data(states, labels, groups, optimal=None, meta: dict | None = None) -> DecisionData: """`soft` = uniform distribution over the oracle's equally-optimal moves: used for the "optimal-move accuracy" metric and available as a soft-label training target.""" task = SnakeTask() lab = {a: i for i, a in enumerate(ACTIONS)} n = len(states) soft = None if optimal is not None: soft = np.zeros((n, 4), np.float32) for i, good in enumerate(optimal): soft[i, [lab[a] for a in good]] = 1.0 / len(good) return DecisionData( task="snake", state=task.states_to_arrays(states), questions=[QUESTION], question_idx=np.zeros(n, np.int32), choices=list(ACTIONS), choice_idx=np.tile(np.arange(4, dtype=np.int32), (n, 1)), target=np.array([lab[a] for a in labels], np.int64), group=np.array(groups, np.int64), hash=np.array([state_hash(s) for s in states], np.int64), soft=soft, meta=meta or {}, ) def generate(n_samples: int, seed: int = 0, board_sizes=(10,), keep_prob: float = 0.1, explore_eps: float = 0.05, env_cfg: dict | None = None, workers: int | None = None) -> DecisionData: workers = workers or max(1, (os.cpu_count() or 2) - 1) workers = min(workers, max(1, n_samples // 2000)) per = [n_samples // workers + (1 if i < n_samples % workers else 0) for i in range(workers)] jobs = [(per[i], seed * 1000 + i, tuple(board_sizes), keep_prob, explore_eps, env_cfg, i * 10_000_000) for i in range(workers)] if workers == 1: results = [_worker(jobs[0])] else: ctx = mp.get_context("fork") if sys.platform == "linux" else None with ProcessPoolExecutor(workers, mp_context=ctx) as ex: results = list(ex.map(_worker, jobs)) states = [s for r in results for s in r[0]] labels = [a for r in results for a in r[1]] groups = [g for r in results for g in r[2]] optimal = [o for r in results for o in r[3]] meta = {"generator": "oracle", "board_sizes": list(board_sizes), "keep_prob": keep_prob, "explore_eps": explore_eps, "seed": seed, "env": env_cfg or {}} return to_decision_data(states, labels, groups, optimal, meta) def generate_splits(out_dir: str | Path, n_samples: int, seed: int = 0, fractions=(0.9, 0.05, 0.05), preview: int = 50, **kw) -> dict: data = generate(n_samples, seed=seed, **kw) splits = split_by_group(data, fractions, seed=seed) out = Path(out_dir) out.mkdir(parents=True, exist_ok=True) summary = {} for name, d in splits.items(): d.save(out / f"{name}.npz") summary[name] = {"rows": len(d), "games": int(len(np.unique(d.group))), "dropped_duplicates": d.meta.get("dropped_duplicates", 0), "label_distribution": {a: int((d.target == i).sum()) for i, a in enumerate(ACTIONS)}} (out / "meta.json").write_text(json.dumps({"task": "snake", **data.meta, "splits": summary}, indent=2)) if preview: export_jsonl(splits["train"], out / "preview.jsonl", limit=preview) return summary def export_jsonl(data: DecisionData, path: str | Path, limit: int | None = None) -> None: """Human-readable export in the canonical {state, question, choices, target} format.""" from minijev.tasks import ragged_take n = len(data) if limit is None else min(limit, len(data)) with open(path, "w") as f: for i in range(n): w = int(data.state["width"][i]) cells, _ = ragged_take(data.state["body"], data.state["body_offsets"], np.array([i])) food = int(data.state["food"][i]) state = {"height": int(data.state["height"][i]), "width": w, "body": [list(divmod(int(c), w)) for c in cells], "food": list(divmod(food, w)) if food >= 0 else None, "direction": ACTIONS[int(data.state["direction"][i])]} choices = data.row_choices(i) f.write(json.dumps({"state": state, "question": data.questions[data.question_idx[i]], "choices": choices, "target": choices[data.target[i]]}) + "\n") def collect_with_policy(policy_probs, n_samples: int, seed: int = 0, board_sizes=(10,), keep_prob: float = 0.2, explore_eps: float = 0.0, env_cfg: dict | None = None, parallel: int = 64, game_id_base: int = 0) -> DecisionData: """DAgger collection: the *student* drives (`policy_probs(states) -> (B, 4) probs`), the oracle labels every kept state. Games run `parallel` at a time for batched inference.""" rng = np.random.default_rng(seed) oracle = make_agent("oracle") states, labels, groups, optimal = [], [], [], [] next_game = 0 def spawn(): nonlocal next_game size = int(board_sizes[next_game % len(board_sizes)]) cfg = SnakeConfig.from_dict({**(env_cfg or {}), "height": size, "width": size}) env = SnakeEnv(cfg, seed=int(rng.integers(2**31))) env.gid = game_id_base + next_game next_game += 1 return env envs = [spawn() for _ in range(parallel)] while len(states) < n_samples: cur = [e.state() for e in envs] probs = policy_probs(cur) for i, (env, s) in enumerate(zip(envs, cur)): if rng.random() < keep_prob and len(states) < n_samples: label, good, _ = oracle.analyze(s) states.append(s) labels.append(label) groups.append(env.gid) optimal.append(good) if rng.random() < explore_eps: safe = safe_moves(s) or list(ACTIONS) action = safe[int(rng.integers(len(safe)))] else: action = ACTIONS[int(probs[i].argmax())] env.step(action) if env.done: envs[i] = spawn() return to_decision_data(states, labels, groups, optimal, {"generator": "dagger", "keep_prob": keep_prob, "seed": seed})