Add SO-ISMCTS mini trainer

Implements a proof-of-concept single-observer IS-MCTS trainer with AlphaZero-style policy/value network, determinization, replay, self-play, CLI configs, and focused tests. Mini acceptance run reaches positive random eval while keeping play_action_rate above the Deep CFR trap threshold. Tests: uv run python -m pytest tests/games/classic/ismcts/ -x; uv run python -m pytest tests/games/classic/test_deep_cfr_trainer.py -x; uv run lost-cities-ismcts train --config configs/ismcts/mini.yaml
2026-05-10 22:46:22 +09:00
parent 5acda3f272
commit e69f3165b6
15 changed files with 1205 additions and 0 deletions
@@ -0,0 +1,290 @@
+from __future__ import annotations
+
+import json
+import random
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch import nn
+
+from coolrl_lost_cities.games.classic.bots import build_bot
+from coolrl_lost_cities.games.classic.deep_cfr.encoding import input_dim
+from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig
+
+from .config import IsMctsConfig
+from .mcts import IsMctsSearcher
+from .network import AlphaZeroNet
+from .replay_buffer import ReplayBuffer, ReplaySample
+from .self_play import play_self_play_game
+
+
+@dataclass
+class IterationMetrics:
+    iteration: int
+    samples_added: int
+    replay_size: int
+    policy_loss: float
+    value_loss: float
+    total_loss: float
+    self_play_seconds: float
+    train_seconds: float
+    eval_metrics: dict[str, float | int]
+
+    def to_dict(self) -> dict[str, float | int]:
+        data: dict[str, float | int] = {
+            "iteration": self.iteration,
+            "samples/added": self.samples_added,
+            "memory/replay": self.replay_size,
+            "loss/policy": self.policy_loss,
+            "loss/value": self.value_loss,
+            "loss/total": self.total_loss,
+            "time/self_play_seconds": self.self_play_seconds,
+            "time/train_seconds": self.train_seconds,
+        }
+        data.update(self.eval_metrics)
+        return data
+
+
+class IsMctsTrainer:
+    def __init__(
+        self,
+        config: IsMctsConfig,
+        game_config: LostCitiesConfig,
+        *,
+        run_dir: str | Path,
+        device: torch.device | str = "cpu",
+    ) -> None:
+        self.config = config
+        self.game_config = game_config
+        self.run_dir = Path(run_dir)
+        self.device = self._resolve_device(device)
+        probe = GameState.new_game(game_config, seed=config.run.seed)
+        self.input_dim = input_dim(probe, config.encoding)
+        self.action_size = probe.action_size
+        self.network = AlphaZeroNet.from_config(self.input_dim, self.action_size, config).to(
+            self.device
+        )
+        self.optimizer = torch.optim.AdamW(
+            self.network.parameters(),
+            lr=config.optimization.learning_rate,
+            weight_decay=max(float(config.optimization.weight_decay), 1.0e-4),
+        )
+        self.buffer = ReplayBuffer(config.training.replay_capacity, seed=config.run.seed)
+        self.metrics_path = self.run_dir / "metrics.jsonl"
+        self.rng = random.Random(config.run.seed)
+
+    def _resolve_device(self, device: torch.device | str) -> torch.device:
+        token = str(device)
+        if token == "auto":
+            return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        return torch.device(token)
+
+    def train(self) -> list[IterationMetrics]:
+        self.run_dir.mkdir(parents=True, exist_ok=True)
+        (self.run_dir / "config.json").write_text(
+            json.dumps(self.config.to_dict(), indent=2, sort_keys=True),
+            encoding="utf-8",
+        )
+        if self.metrics_path.exists():
+            self.metrics_path.unlink()
+        metrics: list[IterationMetrics] = []
+        max_iterations = self.config.run.max_iterations or 1
+        started = time.perf_counter()
+        for iteration in range(1, max_iterations + 1):
+            if self._time_limit_reached(started):
+                break
+            item = self.run_iteration(iteration)
+            metrics.append(item)
+            self._append_metrics(item)
+            self._save_checkpoints(iteration, item)
+            print(json.dumps(item.to_dict(), sort_keys=True))
+        return metrics
+
+    def run_iteration(self, iteration: int) -> IterationMetrics:
+        self.network.eval()
+        sp_started = time.perf_counter()
+        added = 0
+        for _ in range(self.config.training.games_per_iter):
+            samples = play_self_play_game(
+                self.network,
+                self.config.mcts,
+                self.game_config,
+                self.rng,
+                device=self.device,
+                encoding=self.config.encoding,
+                temperature=self.config.temperature.training,
+                max_steps=self.config.evaluation.max_steps,
+            )
+            self.buffer.add(samples)
+            added += len(samples)
+        self_play_seconds = time.perf_counter() - sp_started
+
+        train_started = time.perf_counter()
+        losses = []
+        for _ in range(self.config.training.gradient_steps_per_iter):
+            batch = self.buffer.sample(self.config.training.batch_size)
+            losses.append(self._train_batch(batch))
+        train_seconds = time.perf_counter() - train_started
+        loss_arr = np.asarray(losses, dtype=np.float64)
+        eval_metrics = self._evaluate(iteration)
+        return IterationMetrics(
+            iteration=iteration,
+            samples_added=added,
+            replay_size=len(self.buffer),
+            policy_loss=float(loss_arr[:, 0].mean()) if len(loss_arr) else 0.0,
+            value_loss=float(loss_arr[:, 1].mean()) if len(loss_arr) else 0.0,
+            total_loss=float(loss_arr[:, 2].mean()) if len(loss_arr) else 0.0,
+            self_play_seconds=self_play_seconds,
+            train_seconds=train_seconds,
+            eval_metrics=eval_metrics,
+        )
+
+    def _train_batch(self, batch: list[ReplaySample]) -> tuple[float, float, float]:
+        self.network.train()
+        info = torch.as_tensor(
+            np.stack([sample.info_state for sample in batch]),
+            dtype=torch.float32,
+            device=self.device,
+        )
+        legal = torch.as_tensor(
+            np.stack([sample.legal_mask for sample in batch]),
+            dtype=torch.bool,
+            device=self.device,
+        )
+        pi = torch.as_tensor(
+            np.stack([sample.pi_target for sample in batch]),
+            dtype=torch.float32,
+            device=self.device,
+        )
+        value_target = torch.as_tensor(
+            [sample.v_target for sample in batch],
+            dtype=torch.float32,
+            device=self.device,
+        )
+        logits, value_pred = self.network(info, legal)
+        log_probs = torch.log_softmax(logits, dim=-1)
+        policy_loss = -(pi * log_probs).sum(dim=-1).mean()
+        value_loss = nn.functional.mse_loss(value_pred, value_target)
+        loss = policy_loss + value_loss
+        self.optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        if self.config.optimization.grad_clip > 0:
+            torch.nn.utils.clip_grad_norm_(
+                self.network.parameters(),
+                self.config.optimization.grad_clip,
+            )
+        self.optimizer.step()
+        return float(policy_loss.item()), float(value_loss.item()), float(loss.item())
+
+    def _evaluate(self, iteration: int) -> dict[str, float | int]:
+        opponents = self.config.evaluation.opponents_for_iteration(iteration)
+        if not opponents:
+            return {}
+        self.network.eval()
+        results: dict[str, float | int] = {}
+        for opponent in opponents:
+            result = evaluate_policy(
+                self.network,
+                self.game_config,
+                opponent=opponent,
+                games=self.config.evaluation.games,
+                seed=self.config.run.seed + iteration * 1000,
+                device=self.device,
+                encoding=self.config.encoding,
+                max_steps=self.config.evaluation.max_steps,
+                mcts_config=self.config.mcts,
+            )
+            key = opponent.replace("-", "_")
+            for metric_key, value in result.items():
+                results[f"eval/{key}/{metric_key}"] = value
+        return results
+
+    def _append_metrics(self, metrics: IterationMetrics) -> None:
+        with self.metrics_path.open("a", encoding="utf-8") as handle:
+            handle.write(json.dumps(metrics.to_dict(), sort_keys=True) + "\n")
+
+    def _save_checkpoints(self, iteration: int, metrics: IterationMetrics) -> None:
+        payload = {
+            "config": self.config.to_dict(),
+            "game_config": self.game_config.to_snapshot(),
+            "iteration": iteration,
+            "network": self.network.state_dict(),
+            "optimizer": self.optimizer.state_dict(),
+            "metrics": metrics.to_dict(),
+        }
+        if self.config.checkpoint.save_latest:
+            torch.save(payload, self.run_dir / "latest.pt")
+        if (
+            self.config.checkpoint.save_every > 0
+            and iteration % self.config.checkpoint.save_every == 0
+        ):
+            torch.save(payload, self.run_dir / f"iteration_{iteration:05d}.pt")
+
+    def _time_limit_reached(self, started: float) -> bool:
+        if self.config.run.max_minutes is None:
+            return False
+        return (time.perf_counter() - started) / 60.0 >= self.config.run.max_minutes
+
+
+def evaluate_policy(
+    network: AlphaZeroNet,
+    config: LostCitiesConfig,
+    *,
+    opponent: str,
+    games: int,
+    seed: int,
+    device: torch.device | str,
+    encoding=None,
+    max_steps: int = 10_000,
+    mcts_config=None,
+) -> dict[str, float | int]:
+    rng = random.Random(seed)
+    score_diffs: list[int] = []
+    wins = losses = draws = 0
+    policy_actions = play_actions = 0
+    for game_index in range(games):
+        policy_player = game_index % 2
+        policies = [
+            build_bot(opponent, seed=seed + game_index),
+            build_bot(opponent, seed=seed + game_index),
+        ]
+        state = GameState.new_game(config, seed=seed + game_index)
+        for _ in range(max_steps):
+            if state.terminal:
+                break
+            current = int(state.current_player)
+            if current == policy_player:
+                searcher = IsMctsSearcher(
+                    network,
+                    mcts_config or IsMctsConfig().mcts,
+                    device=device,
+                    encoding=encoding,
+                    rng=random.Random(rng.randrange(2**31)),
+                )
+                visits = searcher.search(state, current)
+                unified = max(visits, key=visits.get)
+                action = state.from_unified_action(unified)
+            else:
+                action = policies[current].act(state)
+            if current == policy_player and state.phase == "card":
+                policy_actions += 1
+                if action % 2 == 0:
+                    play_actions += 1
+            state.apply_action(action)
+        diff = state.score_diff(policy_player)
+        score_diffs.append(diff)
+        wins += int(diff > 0)
+        losses += int(diff < 0)
+        draws += int(diff == 0)
+    return {
+        "games": games,
+        "wins0": wins,
+        "wins1": losses,
+        "draws": draws,
+        "win_rate0": wins / max(1, games),
+        "avg_score_diff0": float(np.mean(score_diffs)) if score_diffs else 0.0,
+        "play_action_rate": play_actions / max(1, policy_actions),
+    }