| """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}) |
|
|
