예제 소스

vlm_macro_client.py

import argparse
import json
import socket
import time
from pathlib import Path
import re
import numpy as np
from PIL import Image

# Micro actions
ACT_NOOP = 0
ACT_LEFT = 1
ACT_RIGHT = 2
ACT_ROT_CW = 3
ACT_ROT_CCW = 4
ACT_SOFT = 5
ACT_HARD = 6

def _grid_to_image(obs: np.ndarray, cell_px: int = 20) -> Image.Image:
    """Render the (H, W, 1) grid to an RGB image."""
    g = obs[:, :, 0].astype(np.uint8)
    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]
    rgb = np.repeat(np.repeat(rgb, cell_px, axis=0), cell_px, axis=1)
    return Image.fromarray(rgb, mode="RGB")


def next_micro_for_target(piece_type: int, cur_rot: int, cur_x: int, target_rot: int, target_x: int) -> int:
    """Convert target macro state to next micro action."""
    if piece_type == 0:
        return ACT_NOOP
    
    # 1. Rotate to target rotation
    dr = (target_rot - cur_rot) % 4
    if dr == 0:
        pass
    elif dr <= 2:
        return ACT_ROT_CW
    else:
        return ACT_ROT_CCW

    # 2. Move to target x
    if cur_x < target_x:
        return ACT_RIGHT
    if cur_x > target_x:
        return ACT_LEFT
    
    # 3. If rotation and x are matched, hard drop
    return ACT_HARD


class VLMMacroAgent:
    def __init__(self, model_dir: str):
        try:
            import torch
            from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
            from qwen_vl_utils import process_vision_info
        except ImportError as e:
            raise RuntimeError("Missing required packages for VLM: torch, transformers, qwen-vl-utils") from e

        self._torch = torch
        self._process_vision_info = process_vision_info

        md = Path(model_dir).expanduser().resolve()
        if not md.exists():
            raise FileNotFoundError(f"VLM model dir not found: {md}")

        print(f"Loading model from {md}...")
        self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
            str(md),
            torch_dtype="auto",
            device_map="auto",
        )
        self.processor = AutoProcessor.from_pretrained(
            str(md),
            min_pixels=256 * 28 * 28,
            max_pixels=640 * 28 * 28,
        )
        print("Model loaded successfully.")

    def get_macro_action(self, obs: np.ndarray, current_piece_name: str) -> tuple[int, int]:
        img = _grid_to_image(obs)
        prompt = (
            "You are an expert Tetris AI.\n"
            f"A new piece ({current_piece_name}) has just appeared at the top of the board.\n"
            "Look at the current board state and decide the best final resting position for this piece.\n"
            "Return your decision strictly as JSON with two keys:\n"
            '- "target_rot": an integer from 0 to 3 representing the number of clockwise rotations.\n'
            '- "target_x": an integer representing the target column (0-9 for a standard board, from left to right).\n'
            "Try to keep the stack flat, avoid creating holes, and clear lines if possible.\n"
            "Return ONLY the JSON string. Do not output anything else."
        )

        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",
        )
        inputs = {k: v.to(self.model.device) for k, v in inputs.items()}

        gen = self.model.generate(
            **inputs,
            max_new_tokens=128,
            do_sample=False,
        )
        out_ids = gen[:, inputs["input_ids"].shape[1] :]
        out_text = self.processor.batch_decode(out_ids, skip_special_tokens=True)[0]

        print(f"VLM Raw Output: {out_text}")

        m = re.search(r"\{[\s\S]*\}", out_text)
        if m:
            try:
                obj = json.loads(m.group(0))
                rot = int(obj.get("target_rot", 0))
                x = int(obj.get("target_x", 4))
                return rot % 4, x
            except Exception as e:
                print(f"JSON Parse Error: {e}")
        return 0, 4

def _connect_with_retry(host: str, port: int, timeout_s: float, retry_interval_s: float) -> socket.socket:
    deadline = time.time() + timeout_s
    last_err = None
    while time.time() < deadline:
        try:
            return socket.create_connection((host, port), timeout=2.0)
        except OSError as e:
            last_err = e
            time.sleep(retry_interval_s)
    raise ConnectionError(f"Could not connect to {host}:{port}. Last error: {last_err}")


def name_piece(ptype: int) -> str:
    mapping = {0: "-", 1: "I", 2: "O", 3: "T", 4: "S", 5: "Z", 6: "J", 7: "L"}
    return mapping.get(ptype, "?")

def main():
    p = argparse.ArgumentParser(description="VLM Macro Action Client for Qt Tetris")
    p.add_argument("--host", default="127.0.0.1")
    p.add_argument("--port", type=int, default=5555)
    p.add_argument("--model-dir", default="llms/Qwen2.5-VL-3B-Instruct")
    args = p.parse_args()

    agent = VLMMacroAgent(args.model_dir)

    print(f"Connecting to Qt server at {args.host}:{args.port}...")
    s = _connect_with_retry(args.host, args.port, 15.0, 0.5)
    print("Connected to Qt Tetris Server.")
    
    # Enable agent mode on Qt server
    f_out = s.makefile("w", encoding="utf-8", newline="\n")
    f_out.write("agent\n")
    f_out.flush()

    f_in = s.makefile("r", encoding="utf-8", newline="\n")

    current_piece_index = -1
    target_rot = 0
    target_x = 4

    while True:
        line = f_in.readline()
        if not line:
            break
        state = json.loads(line)
        w = int(state["w"])
        visible_h = int(state["visible_h"])
        grid = np.asarray(state["grid"], dtype=np.uint8).reshape(visible_h, w, 1)

        status = state.get("status", "running")
        if status == "game_over":
            print("Game Over. Sending restart command...")
            f_out.write("restart\n")
            f_out.flush()
            continue
        if status == "paused":
            print("Paused. Sending resume command...")
            f_out.write("resume\n")
            f_out.flush()
            continue

        piece_index = int(state.get("piece", -1))
        piece_type = int(state.get("piece_type", 0))
        cur_x = int(state.get("piece_x", 0))
        cur_rot = int(state.get("piece_rot", 0))

        if piece_type == 0:
            # No active piece
            f_out.write(f"{ACT_NOOP}\n")
            f_out.flush()
            continue

        if piece_index > current_piece_index:
            # New piece spawned. Ask VLM for macro action
            piece_name = name_piece(piece_type)
            print(f"\n--- New Piece Spawned: {piece_name} (Index: {piece_index}) ---")
            current_piece_index = piece_index
            target_rot, target_x = agent.get_macro_action(grid, piece_name)
            target_x = max(0, min(target_x, w - 1))
            print(f"VLM Decision -> Target Rot: {target_rot}, Target X: {target_x}")

        # Translate target macro state to next micro action
        micro = next_micro_for_target(piece_type, cur_rot, cur_x, target_rot, target_x)
        
        # Send action to Qt Server
        f_out.write(f"{micro}\n")
        f_out.flush()

if __name__ == "__main__":
    main()