"""Build decision-question training samples from self-play traces. Three question families, mirroring the NanoJev/TypeSafe primitives: Choice (main task) state + "Which action is most likely to make progress?" + candidates label = the candidate that produced the best immediate outcome: +2 levels_completed increased +1 frame changed without losing levels -1 GAME_OVER within the next K steps (credit propagated backwards) Boolean (action-effect pretraining) "Does ACTION k change the frame?" -> p_true from frame diff Score (state value shaping) "How promising is this state?" on 2..10 ordered levels, label derived from exploration coverage + level progress heuristics. Output: HuggingFace `datasets` directory data/dataset/{train,val} """ from __future__ import annotations import argparse import glob import json import sys import zlib from collections import defaultdict from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from qwenjev.candidates import build_candidates # noqa: E402 from qwenjev.state_codec import canonical_state, grid_to_text # noqa: E402 CHOICE_INSTRUCTION = "Which candidate action is most likely to make progress?" BOOLEAN_INSTRUCTION = "Answer whether the proposition about this state is true." SCORE_INSTRUCTION = "Rate the promise of this state on the given ordered levels." def load_traces(in_dir: Path) -> list[dict]: records = [] for path in sorted(glob.glob(str(in_dir / "**" / "ep*.jsonl"), recursive=True)): with open(path) as fh: step_rows = [] for line in fh: if not line.strip(): continue try: row = json.loads(line) except json.JSONDecodeError: continue # truncated line from a killed collector # Slim down: downstream only uses grids_before[0] + # frame_hash/levels — drop the rest to fit in RAM. row.pop("grids_after", None) gb = row.get("grids_before") if isinstance(gb, list) and len(gb) > 1: row["grids_before"] = gb[:1] step_rows.append(row) for i, row in enumerate(step_rows): row["_next"] = step_rows[i + 1] if i + 1 < len(step_rows) else None records.append(row) return records def state_text(row: dict) -> str: return row.get("state_text") or canonical_state(grid_to_text(row["grids_before"][0])) def make_choice_sample(row: dict) -> dict | None: nxt = row["_next"] if nxt is None: return None changed = nxt["frame_hash"] != row["frame_hash"] leveled = nxt["levels_completed_after"] > row["levels_completed_before"] cands = build_candidates( row["available_actions"] or ["ACTION1"], row["grids_before"][0], row["grids_before"][0], # prev==cur at collection time; cheap proxy allow_reset=False, ) keys = [c.key() for c in cands] if row["action"] not in keys: return None label = keys.index(row["action"]) return { "family": "choice", "state": state_text(row), # NOTE: the frame grid rides along EVERY sample - the System-1 # vision tower consumes it on both training and inference paths. "grid": row["grids_before"][0], "instruction": CHOICE_INSTRUCTION, "candidates": [c.text or c.key() for c in cands], "label": label, "reward": 2 if leveled else (1 if changed else 0), "game_id": row["game_id"], } def make_boolean_sample(row: dict) -> dict | None: nxt = row["_next"] if nxt is None: return None changed = nxt["frame_hash"] != row["frame_hash"] return { "family": "boolean", "state": state_text(row), "grid": row["grids_before"][0], "instruction": BOOLEAN_INSTRUCTION, "proposition": f"{row['action']} changes the frame", "label": int(changed), "game_id": row["game_id"], } LEVEL_COUNT = 10 def make_score_sample(row: dict, coverage: dict) -> dict: seen = coverage.get(row["game_id"], 1) leveled = max(0, row["levels_completed_before"]) score01 = min(1.0, 0.6 * (leveled / (leveled + 2)) + 0.4 * (row["step"] / max(seen, 1))) level = max(0, min(LEVEL_COUNT - 1, int(score01 * (LEVEL_COUNT - 1)))) return { "family": "score", "state": state_text(row), "grid": row["grids_before"][0], "instruction": SCORE_INSTRUCTION, "description": "promise of the current state", "label": level, "game_id": row["game_id"], } def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--in", dest="in_dir", default="data/traces") ap.add_argument("--out", dest="out_dir", default="data/dataset") ap.add_argument("--val-ratio", type=float, default=0.05) args = ap.parse_args() records = load_traces(Path(args.in_dir)) if not records: raise SystemExit(f"no traces found under {args.in_dir}; run collect_selfplay.py first") coverage = defaultdict(int) for r in records: coverage[r["game_id"]] = max(coverage[r["game_id"]], r["step"]) samples = [] for row in records: c = make_choice_sample(row) if c: samples.append(c) samples.append(make_boolean_sample(row)) samples.append(make_score_sample(row, coverage)) by_game = sorted({r["game_id"] for r in records}) # Game-heldout validation when there are enough games; otherwise (single # game / smoke runs) fall back to a row-hash split so train is non-empty. if len(by_game) >= 5: val_games = set(by_game[: max(1, int(len(by_game) * args.val_ratio))]) def is_val(s): return s["game_id"] in val_games else: val_games = set() thr = int(args.val_ratio * 0xFFFFFFFF) def is_val(s): return zlib.crc32(s["state"].encode()) < thr out = Path(args.out_dir) for split in ("train", "val"): (out / split).mkdir(parents=True, exist_ok=True) counts = {"choice": 0, "boolean": 0, "score": 0} with open(out / "train" / "samples.jsonl", "w") as ftr, \ open(out / "val" / "samples.jsonl", "w") as fva: for s in samples: if s is None: continue counts[s["family"]] += 1 writer = fva if is_val(s) else ftr writer.write(json.dumps(s) + "\n") (out / "stats.json").write_text(json.dumps( {"total": sum(counts.values()), **counts, "games": len(by_game), "val_games": sorted(val_games)}, indent=2)) print(f"dataset ready -> {out} ({counts})") if __name__ == "__main__": main()