| """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 |
| from qwenjev.state_codec import canonical_state, grid_to_text |
|
|
| 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 |
| |
| |
| 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], |
| 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), |
| |
| |
| "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}) |
| |
| |
| 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() |
|
|
