예제 소스

policy_backend.py

"""GO19 서버용 간단 정책 백엔드 (NumPy MLP).

목표:
- packed 보드(361바이트 0/1/2)가 정확하다는 전제에서,
  모델(VLM)이 보드를 다시 읽지 않도록 "수 선택"만 담당.
- 주말 내내 self-play(서버1/서버2) 돌릴 때, "게임 리셋" 감지 시 온라인 업데이트 1회.

주의:
- 현재는 패(코) 규칙 없음. 서버 `_MiniGo` 기반 마스크/검증과 동일한 수준.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Sequence, Tuple

import numpy as np

from rl_env.board import estimate_area_score  # simple territory estimate (python)

BOARD_SIZE = 19
N_SQ = BOARD_SIZE * BOARD_SIZE
PASS_A = N_SQ
N_ACT = N_SQ + 1
IN_DIM = N_SQ * 2


def encode_obs(packed: bytes, turn: int) -> np.ndarray:
    """현재 차례 기준 2채널(내 돌/상대 돌) → (IN_DIM,) float64."""
    cells = np.frombuffer(packed, dtype=np.uint8).astype(np.float64)
    opp = 3 - int(turn)
    own_plane = (cells == int(turn)).astype(np.float64)
    opp_plane = (cells == opp).astype(np.float64)
    return np.concatenate([own_plane, opp_plane])


def masked_logits(logits: np.ndarray, mask: Sequence[int]) -> np.ndarray:
    m = np.asarray(mask, dtype=np.float64)
    out = logits.copy()
    out[m < 0.5] = -1e9
    return out


def restrict_pass_to_forced(mask: List[int], pass_action: int = PASS_A) -> bool:
    """둘 수 있는 좌표가 있으면 PASS를 마스킹한다.

    반환값은 PASS만 남은 상태인지 여부다. 온라인 정책이 불리한 전술을
    PASS로 회피해 학습하지 않도록 하고, 좌표 행동이 하나도 없을 때만
    PASS를 허용한다.
    """
    has_board_move = any(mask[:pass_action])
    mask[pass_action] = 0 if has_board_move else 1
    return not has_board_move


def softmax(x: np.ndarray) -> np.ndarray:
    x = x - np.max(x)
    e = np.exp(x)
    return e / (np.sum(e) + 1e-12)


def outcome_reward(packed: bytes, side: int) -> float:
    """거친 종국 보상: (내 area − 상대 area)를 tanh로 스케일.

    area = 돌 수 + (단순 flood-fill로 추정한 영토).
    """
    b_area, w_area, _d = estimate_area_score(packed, BOARD_SIZE)
    if int(side) == 1:
        diff = b_area - w_area
    else:
        diff = w_area - b_area
    return float(np.tanh(diff / 128.0))


@dataclass
class PolicyStats:
    games_seen: int = 0
    ema_reward: float = 0.0


class PolicyMLP:
    """2층 MLP 정책. 저장 형식은 rl_env 와 호환(npz)."""

    def __init__(self, hidden: int = 256, seed: int = 0, lr: float = 5e-3) -> None:
        rng = np.random.default_rng(seed)
        self.W1 = rng.normal(0.0, 0.05, (IN_DIM, hidden)).astype(np.float64)
        self.b1 = np.zeros(hidden, dtype=np.float64)
        self.W2 = rng.normal(0.0, 0.05, (hidden, N_ACT)).astype(np.float64)
        self.b2 = np.zeros(N_ACT, dtype=np.float64)
        self.hidden = hidden
        self.lr = float(lr)
        self.stats = PolicyStats()

    def forward(self, x: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        h = np.maximum(0.0, x @ self.W1 + self.b1)
        logits = h @ self.W2 + self.b2
        return h, logits

    def greedy_action(self, x: np.ndarray, mask: Sequence[int]) -> int:
        _, logits = self.forward(x)
        lm = masked_logits(logits, mask)
        return int(np.argmax(lm))

    def sample_action(self, x: np.ndarray, mask: Sequence[int], rng: np.random.Generator) -> int:
        _, logits = self.forward(x)
        lm = masked_logits(logits, mask)
        p = softmax(lm)
        return int(rng.choice(N_ACT, p=p))

    def update_transitions(self, transitions: List[Tuple[np.ndarray, Sequence[int], int, float]]) -> None:
        """REINFORCE: (x, mask, action, advantage) 리스트 1회 SGD."""
        if not transitions:
            return
        dW1 = np.zeros_like(self.W1)
        db1 = np.zeros_like(self.b1)
        dW2 = np.zeros_like(self.W2)
        db2 = np.zeros_like(self.b2)
        for x, mask, a, adv in transitions:
            h, logits = self.forward(x)
            lm = masked_logits(logits, mask)
            p = softmax(lm)
            grad_logits = -float(adv) * ((np.arange(N_ACT) == int(a)).astype(np.float64) - p)
            db2 += grad_logits
            dW2 += np.outer(h, grad_logits)
            grad_h = grad_logits @ self.W2.T
            grad_h *= (h > 0).astype(np.float64)
            db1 += grad_h
            dW1 += np.outer(x, grad_h)
        n = float(max(1, len(transitions)))
        self.W2 -= self.lr * dW2 / n
        self.b2 -= self.lr * db2 / n
        self.W1 -= self.lr * dW1 / n
        self.b1 -= self.lr * db1 / n

    def save(self, path: str) -> None:
        np.savez(
            path,
            W1=self.W1,
            b1=self.b1,
            W2=self.W2,
            b2=self.b2,
            hidden=np.array([self.hidden]),
        )

    @classmethod
    def load(cls, path: str, lr: float = 5e-3) -> "PolicyMLP":
        d = np.load(path)
        hidden = int(d["hidden"][0])
        p = cls(hidden=hidden, lr=lr)
        p.W1 = np.asarray(d["W1"], dtype=np.float64)
        p.b1 = np.asarray(d["b1"], dtype=np.float64)
        W2 = np.asarray(d["W2"], dtype=np.float64)
        b2 = np.asarray(d["b2"], dtype=np.float64)
        if W2.shape[1] == N_SQ:
            W2n = np.zeros((W2.shape[0], N_ACT), dtype=np.float64)
            W2n[:, :N_SQ] = W2
            b2n = np.zeros(N_ACT, dtype=np.float64)
            b2n[:N_SQ] = b2
            b2n[PASS_A] = -6.0
            p.W2, p.b2 = W2n, b2n
        else:
            p.W2, p.b2 = W2, b2
        return p