예제 소스

qt_rl_client.py

import argparse
import json
import socket
import time
from pathlib import Path

import numpy as np


def _resolve_model_zip(model: str | Path) -> Path:
    """
    SB3 DQN 모델 zip 경로를 안전하게 해석합니다.
    - 이미 .zip로 끝나면 중복으로 .zip을 붙이지 않습니다.
    - 디렉토리가 주어지면(예: model/dqn_tetris/) 그 안의 *.zip 중 가장 최근 수정 파일을 선택합니다.
    """
    p = Path(model).expanduser()
    candidates: list[Path] = []

    if p.suffix.lower() == ".zip":
        candidates.append(p)
    else:
        candidates.append(p)
        candidates.append(p.with_suffix(".zip"))
        candidates.append(Path(str(p) + ".zip"))

    for c in candidates:
        if c.exists() and c.is_file():
            return c.resolve()

    # 디렉토리 입력 지원
    if p.exists() and p.is_dir():
        zips = sorted(p.glob("*.zip"), key=lambda x: x.stat().st_mtime, reverse=True)
        if zips:
            return zips[0].resolve()

    tried = "\n".join(f"  - {c}" for c in candidates)
    raise FileNotFoundError(
        "\n".join(
            [
                f"모델 zip을 찾을 수 없습니다: {model}",
                "시도한 경로:",
                tried,
                "",
                "예) python qt_rl_client.py --model model/dqn_tetris/tetris_policy.zip --port 5555",
            ]
        )
    )


def _connect_with_retry(host: str, port: int, timeout_s: float, retry_interval_s: float) -> socket.socket:
    deadline = time.time() + timeout_s
    last_err: Exception | None = 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)
    msg = "\n".join(
        [
            f"Qt 테트리스 서버({host}:{port})에 연결할 수 없습니다.",
            f"마지막 에러: {last_err}",
            "",
            "아래처럼 Qt 앱을 먼저 실행했는지 확인하세요:",
            f"  ./tetris/build/tetris_rl --rl-port {port}",
            "",
            "또는 포트가 이미 사용 중이면 다른 포트를 사용하세요.",
        ]
    )
    raise ConnectionError(msg)


def main():
    p = argparse.ArgumentParser(description="Qt Tetris RL client (TCP)")
    p.add_argument("--host", default="127.0.0.1")
    p.add_argument("--port", type=int, required=True)
    p.add_argument("--model", required=True, help="SB3 DQN .zip path")
    p.add_argument("--deterministic", action="store_true")
    p.add_argument("--auto-reset", action="store_true", help="game_over면 reset 명령 전송")
    p.add_argument("--connect-timeout", type=float, default=15.0, help="서버 연결 재시도 시간(초)")
    args = p.parse_args()

    from stable_baselines3 import DQN

    model_path = _resolve_model_zip(args.model)
    with open(model_path, "rb") as _f:
        model = DQN.load(_f)

    s = _connect_with_retry(
        host=args.host, port=args.port, timeout_s=args.connect_timeout, retry_interval_s=0.3
    )
    f_in = s.makefile("r", encoding="utf-8", newline="\n")
    f_out = s.makefile("w", encoding="utf-8", newline="\n")

    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)

        # 모델/환경 observation shape 확인 (초기 1회만)
        if not hasattr(main, "_shape_checked"):
            main._shape_checked = True  # type: ignore[attr-defined]
            model_shape = getattr(getattr(model, "observation_space", None), "shape", None)
            env_shape = grid.shape
            if model_shape is not None and tuple(model_shape) != tuple(env_shape):
                raise RuntimeError(
                    "\n".join(
                        [
                            "모델이 기대하는 관측 shape과 Qt 테트리스 관측 shape이 달라 적용할 수 없습니다.",
                            f"  - model.observation_space.shape = {model_shape}",
                            f"  - qt observation shape         = {env_shape}",
                            "",
                            "해결 방법:",
                            "  1) 로컬 tetris 엔진(LocalTetrisEnv)으로 다시 학습한 모델을 사용하세요.",
                            "     예) python train_tetris_rl.py --env-id tetris/build/libtetris_c_api.dylib --timesteps 300000",
                            "  2) 또는 같은 관측/액션 정의로 재학습된 모델을 준비해야 합니다.",
                        ]
                    )
                )

        status = state.get("status", "running")
        if status == "game_over":
            if args.auto_reset:
                # Qt 서버는 restart가 리셋+재개까지 포함 (구버전 reset도 호환)
                f_out.write("restart\n")
                f_out.flush()
            else:
                f_out.write("0\n")
                f_out.flush()
            continue
        if status == "paused":
            if args.auto_reset:
                f_out.write("resume\n")
                f_out.flush()
            continue

        action, _ = model.predict(grid, deterministic=args.deterministic)
        f_out.write(f"{int(action)}\n")
        f_out.flush()


if __name__ == "__main__":
    main()