File size: 6,838 Bytes
3e04895 9bd7482 3e04895 09774a9 3e04895 9bd7482 3e04895 9bd7482 3e04895 9bd7482 3e04895 9bd7482 3e04895 9bd7482 3e04895 | 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 | """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()
|