예제 소스

llm_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


def _grid_to_ascii(obs: np.ndarray) -> str:
    """
    obs: (H, W, 1) uint8, 0 empty, 1..7 piece type
    """
    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]
    # 0=., 1..7=digits (compact)
    rows = []
    for y in range(g.shape[0]):
        row = "".join("." if int(v) == 0 else str(int(v)) for v in g[y])
        rows.append(row)
    return "\n".join(rows)


@dataclass(frozen=True)
class LLMTeacherConfig:
    enabled: bool = False
    model_path: str = "llms/qwen2.5-3b-instruct-q4_k_m.gguf"
    n_ctx: int = 2048
    n_threads: int = 8
    n_gpu_layers: int = 0
    temperature: float = 0.0
    top_p: float = 0.9
    max_tokens: int = 128
    # DQN 보상에 더할 가중치
    reward_alpha: float = 0.2
    # 매 스텝 호출은 너무 느리므로 기본은 N스텝마다 1회(1이면 매 스텝)
    call_every: int = 4


class LocalGGUFTeacher:
    def __init__(self, cfg: LLMTeacherConfig):
        self.cfg = cfg
        self._step = 0
        self._last_score = 0.0
        self._last_action = 0

        try:
            from llama_cpp import Llama  # type: ignore
        except Exception as e:  # pragma: no cover
            raise RuntimeError(
                "llama-cpp-python이 필요합니다.\n"
                "  .venv/bin/pip install llama-cpp-python\n"
                f"원래 에러: {type(e).__name__}: {e}"
            ) from e

        mp = Path(cfg.model_path).expanduser()
        if not mp.exists():
            raise FileNotFoundError(f"GGUF 모델 파일이 없습니다: {mp}")

        self.llm = Llama(
            model_path=str(mp),
            n_ctx=int(cfg.n_ctx),
            n_threads=int(cfg.n_threads),
            n_gpu_layers=int(cfg.n_gpu_layers),
            logits_all=False,
            verbose=False,
        )

    def _prompt(self, obs_ascii: str, 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"
            "Given the board (top to bottom) after applying the last action, do two things:\n"
            "1) Propose the NEXT action as an integer in [0..6].\n"
            "2) Score the LAST action outcome as a float in [-1.0..1.0] (higher is better).\n"
            "Prefer flat surface, low stack height, and avoid creating holes.\n"
            "\n"
            "Actions: 0 Noop, 1 Left, 2 Right, 3 RotateCW, 4 RotateCCW, 5 SoftDrop, 6 HardDrop\n"
            f"last_action={int(last_action)}\n"
            f"stack_height={stack_h} bumpiness={bump} lines_total={lines}\n"
            "board:\n"
            f"{obs_ascii}\n"
            "\n"
            "Return ONLY JSON like:\n"
            '{"action": 6, "score": 0.1}\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

        obs_ascii = _grid_to_ascii(obs)
        prompt = self._prompt(obs_ascii, last_action, info)

        out = self.llm.create_completion(
            prompt=prompt,
            max_tokens=int(self.cfg.max_tokens),
            temperature=float(self.cfg.temperature),
            top_p=float(self.cfg.top_p),
            stop=["\n\n"],
        )
        text = (out.get("choices") or [{}])[0].get("text", "")
        # JSON만 뽑기
        m = re.search(r"\{[\s\S]*\}", 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__ = ["LLMTeacherConfig", "LocalGGUFTeacher"]