File size: 8,815 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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 | """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})
|