예제 소스

vlm_teacher.py

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image


def _grid_to_image(obs: np.ndarray, cell_px: int = 20) -> Image.Image:
    """
    obs: (H, W, 1) uint8, 0 empty, 1..7 piece type
    Render to an RGB image for VLM input.
    """
    if obs.ndim != 3 or obs.shape[-1] != 1:
        raise ValueError(f"obs shape must be (H,W,1), got {obs.shape}")
    g = obs[:, :, 0].astype(np.uint8)
    h, w = g.shape

    # 0..7 palette (roughly matches Qt colors)
    palette = np.array(
        [
            [20, 22, 28],   # empty
            [80, 220, 250],  # I
            [250, 220, 60],  # O
            [180, 80, 220],  # T
            [90, 220, 120],  # S
            [240, 90, 90],   # Z
            [80, 120, 240],  # J
            [240, 160, 60],  # L
        ],
        dtype=np.uint8,
    )
    rgb = palette[g]  # (H,W,3)
    # upscale
    rgb = np.repeat(np.repeat(rgb, cell_px, axis=0), cell_px, axis=1)
    return Image.fromarray(rgb, mode="RGB")


@dataclass(frozen=True)
class VLMTeacherConfig:
    enabled: bool = False
    # Local directory where HF model is downloaded
    model_dir: str = "llms/Qwen2.5-VL-3B-Instruct"
    # Generation
    max_new_tokens: int = 128
    temperature: float = 0.0
    top_p: float = 0.9
    # Reduce visual tokens for speed (processor min/max pixels)
    min_pixels: int = 256 * 28 * 28
    max_pixels: int = 640 * 28 * 28
    # Call throttling
    call_every: int = 4
    # Render size
    cell_px: int = 20


class LocalQwenVLViaTransformersTeacher:
    """
    Qwen2.5-VL (Transformers) local inference.
    Input: rendered image of the grid + short text prompt
    Output: JSON with {"action": int, "score": float}
    """

    def __init__(self, cfg: VLMTeacherConfig):
        self.cfg = cfg
        self._step = 0
        self._last_score = 0.0
        self._last_action = 0

        try:
            import torch
            from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
            from qwen_vl_utils import process_vision_info
        except Exception as e:  # pragma: no cover
            raise RuntimeError(
                "VLM(Transformers) 실행에 필요한 패키지가 없습니다.\n"
                "  .venv/bin/pip install torch transformers accelerate safetensors pillow qwen-vl-utils\n"
                f"원래 에러: {type(e).__name__}: {e}"
            ) from e

        self._torch = torch
        self._process_vision_info = process_vision_info

        md = Path(cfg.model_dir).expanduser().resolve()
        if not md.exists():
            raise FileNotFoundError(
                f"VLM 모델 디렉토리가 없습니다: {md}\n"
                "먼저 모델을 llms/ 아래로 다운로드해야 합니다."
            )

        # Load
        self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
            str(md),
            torch_dtype="auto",
            device_map="auto",
        )
        self.processor = AutoProcessor.from_pretrained(
            str(md),
            min_pixels=int(cfg.min_pixels),
            max_pixels=int(cfg.max_pixels),
        )

    def _prompt(self, last_action: int, info: dict[str, Any]) -> str:
        stack_h = int(info.get("stack_height", -1))
        bump = int(info.get("bumpiness", -1))
        lines = int(info.get("lines_cleared_total", -1))
        return (
            "You are a Tetris coach and judge.\n"
            "Look at the image of the board AFTER applying the last action.\n"
            "Return ONLY JSON with:\n"
            '- "action": integer in [0..6] for the NEXT action\n'
            '- "score": float in [-1.0..1.0] scoring the LAST action outcome\n'
            "Prefer flat surface, low stack height, and avoid creating holes.\n"
            "Actions: 0 Noop, 1 Left, 2 Right, 3 RotateCW, 4 RotateCCW, 5 SoftDrop, 6 HardDrop\n"
            f"last_action={int(last_action)} stack_height={stack_h} bumpiness={bump} lines_total={lines}\n"
        )

    def judge_and_suggest(self, obs: np.ndarray, last_action: int, info: dict[str, Any]) -> tuple[int, float]:
        self._step += 1
        if self.cfg.call_every > 1 and (self._step % int(self.cfg.call_every) != 0):
            return self._last_action, self._last_score

        img = _grid_to_image(obs, cell_px=int(self.cfg.cell_px))
        prompt = self._prompt(last_action, info)

        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": img},
                    {"type": "text", "text": prompt},
                ],
            }
        ]
        text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        image_inputs, video_inputs = self._process_vision_info(messages)
        inputs = self.processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        # Move to model device
        inputs = {k: v.to(self.model.device) for k, v in inputs.items()}

        gen = self.model.generate(
            **inputs,
            max_new_tokens=int(self.cfg.max_new_tokens),
            do_sample=(float(self.cfg.temperature) > 0.0),
            temperature=float(self.cfg.temperature) if float(self.cfg.temperature) > 0.0 else None,
            top_p=float(self.cfg.top_p) if float(self.cfg.temperature) > 0.0 else None,
        )
        # Remove prompt tokens
        out_ids = gen[:, inputs["input_ids"].shape[1] :]
        out_text = self.processor.batch_decode(out_ids, skip_special_tokens=True)[0]

        m = re.search(r"\{[\s\S]*\}", out_text)
        if not m:
            return self._last_action, self._last_score
        try:
            obj = json.loads(m.group(0))
            a = int(obj.get("action", 0))
            s = float(obj.get("score", 0.0))
        except Exception:
            return self._last_action, self._last_score

        a = int(np.clip(a, 0, 6))
        s = float(np.clip(s, -1.0, 1.0))
        self._last_action, self._last_score = a, s
        return a, s


__all__ = ["VLMTeacherConfig", "********"]