예제 소스

dataset_invader.py

#!/usr/bin/env python3
"""invader_vlm_server.py / INVL 클라이언트 로그(JSONL) → 행동 레이블 학습 샘플.

예:
  python dataset_invader.py --jsonl logs/invader.jsonl --out-dir datasets/invader_sft --save-images

--won-only: 마지막 틱이 Won 인 세션만 사용(레거시). 무한 웨이브 빌드에선 Won 이 안 나오므로 거의 비어 있을 수 있음.
"""

from __future__ import annotations

import argparse
import json
import os
from typing import Iterator, List

from PIL import Image, ImageDraw

INV_ROWS = 14
INV_COLS = 11

ACTION_NAMES = ("left", "right", "fire", "none")


def _read_jsonl(path: str) -> Iterator[dict]:
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                yield json.loads(line)
            except json.JSONDecodeError:
                continue


def group_by_game(records: List[dict]) -> Iterator[List[dict]]:
    """상태가 없는 레코드(서버 predict 전용 JSONL 등)는 한 줄이 한 샘플이 되게 그대로 내보냄."""
    cur: List[dict] = []
    for rec in records:
        if "status" not in rec:
            if cur:
                yield cur
                cur = []
            yield [rec]
            continue
        cur.append(rec)
        st = rec.get("status")
        if st in ("Won", "Lost"):
            yield cur
            cur = []
    if cur:
        yield cur


def render_board(board: List[int], cell: int = 22, pad: int = 8) -> Image.Image:
    w = pad * 2 + cell * INV_COLS
    h = pad * 2 + cell * INV_ROWS
    img = Image.new("RGB", (w, h), (10, 12, 28))
    draw = ImageDraw.Draw(img)
    colors = {
        0: (10, 12, 28),
        1: (80, 220, 120),
        2: (220, 220, 240),
        3: (255, 230, 80),
        4: (120, 200, 255),
    }
    for r in range(INV_ROWS):
        for c in range(INV_COLS):
            v = board[r * INV_COLS + c]
            x0 = pad + c * cell
            y0 = pad + r * cell
            draw.rectangle([x0, y0, x0 + cell - 1, y0 + cell - 1], outline=(40, 50, 90))
            if v:
                draw.rectangle(
                    [x0 + 2, y0 + 2, x0 + cell - 3, y0 + cell - 3],
                    fill=colors.get(v, (128, 128, 128)),
                )
    return img


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--jsonl", required=True)
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--save-images", action="store_true")
    ap.add_argument("--won-only", action="store_true")
    args = ap.parse_args()

    def is_sample(r: dict) -> bool:
        if r.get("game") != "invader" or "action" not in r or "board" not in r:
            return False
        k = r.get("kind")
        if k is not None and k not in ("move", "predict"):
            return False
        return True

    moves = [r for r in _read_jsonl(args.jsonl) if is_sample(r)]
    os.makedirs(args.out_dir, exist_ok=True)
    items: List[dict] = []
    for game in group_by_game(moves):
        if not game:
            continue
        if args.won_only and game[-1].get("status") != "Won":
            continue
        items.extend(game)

    train_path = os.path.join(args.out_dir, "train.jsonl")
    n = 0
    with open(train_path, "w", encoding="utf-8") as out:
        for rec in items:
            board = rec.get("board")
            if not board or len(board) != INV_ROWS * INV_COLS:
                continue
            act = rec.get("action")
            if act is None or not (0 <= int(act) <= 3):
                continue
            prompt = (
                "Space Invaders grid 14x11. Symbols: A=alien P=you b=your bullet B=enemy bullet. "
                "Output exactly one digit: 0=left 1=right 2=fire 3=wait."
            )
            row = {"messages": [{"role": "user", "content": prompt + "\n(see image)"}], "action": int(act)}
            if args.save_images:
                img = render_board(board)
                img_path = os.path.join(args.out_dir, f"frame_{n:06d}.png")
                img.save(img_path)
                row["image"] = os.path.basename(img_path)
            out.write(json.dumps(row, ensure_ascii=False) + "\n")
            n += 1
    print(f"wrote {n} examples → {train_path}")


if __name__ == "__main__":
    main()