scripts/collect_chunk.py

"""Chunked collector driver: one game per invocation, resumable.

The local sandbox reaps background processes, so the sweep is driven as
repeated short foreground runs (< 9 min each):

    python3 scripts/collect_chunk.py --game ar25
    python3 scripts/collect_chunk.py --all --from ft09   # remaining games

Each episode file is checked before running: a file with >= max_actions
lines is kept (resumable), otherwise the episode is re-collected.
"""
from __future__ import annotations

import argparse
import json
import random
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import yaml

from data.collect_selfplay import PUBLIC_GAMES_FALLBACK, make_env, grids_from, meta_from, run_episode
from qwenjev.hybrid import HybridAgent, HybridConfig


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="configs/collect_real25.yaml")
    ap.add_argument("--game", default=None, help="single game to collect")
    ap.add_argument("--from", dest="start", default=None,
                    help="collect all games starting from this one")
    args = ap.parse_args()

    with open(args.config) as f:
        cfg = yaml.safe_load(f)
    games = cfg.get("games") or PUBLIC_GAMES_FALLBACK
    if args.game:
        games = [args.game]
    elif args.start:
        games = games[games.index(args.start):]

    out_root = Path(cfg["out_dir"])
    episodes = int(cfg.get("episodes_per_game", 8))
    max_steps = int(cfg.get("max_actions_per_episode", 400))
    brain = HybridAgent(config=HybridConfig(epsilon=0.0,
                                            system2_max_calls_per_game=0))

    for game_id in games:
        gdir = out_root / game_id
        gdir.mkdir(parents=True, exist_ok=True)
        t0 = time.time()
        done_eps = 0
        for ep in range(episodes):
            ep_file = gdir / f"ep{ep:03d}.jsonl"
            if ep_file.exists():
                n = sum(1 for _ in open(ep_file))
                if n > 0:
                    # non-empty file = finished episode (some games end early
                    # via GAME_OVER/WIN — that is a valid, shorter episode)
                    done_eps += 1
                    continue
            rng = random.Random(int(cfg.get("seed", 7)) + ep * 977)
            run_episode(game_id, ep, cfg, brain, out_root / game_id, rng)
        print(f"[chunk] {game_id}: {done_eps}/{episodes} episodes ready "
              f"({time.time()-t0:.0f}s this call)", flush=True)


if __name__ == "__main__":
    main()