diff --git a/configs/ismcts/default.yaml b/configs/ismcts/default.yaml new file mode 100644 index 0000000..5885b38 --- /dev/null +++ b/configs/ismcts/default.yaml @@ -0,0 +1,46 @@ +run: + experiment_name: ismcts-default + max_iterations: 100 + seed: 1 + device: auto +rules: + n_colors: 5 + n_ranks: 9 + min_rank: 2 + n_handshakes: 3 + hand_size: 8 + expedition_penalty: -20 + bonus_threshold: 8 + bonus_amount: 20 +encoding: + derived_playability: true + slot_aware_playability: true +network: + kind: mlp + hidden_size: 512 + num_layers: 3 + activation: relu +mcts: + n_simulations: 50 + c_puct: 1.5 + max_depth: 200 +temperature: + training: 1.0 + eval: 0.0 +training: + games_per_iter: 10 + gradient_steps_per_iter: 10 + batch_size: 128 + replay_capacity: 100000 +optimization: + learning_rate: 0.0003 + grad_clip: 5.0 +checkpoint: + save_every: 10 + save_latest: true +evaluation: + eval_every: 10 + games: 20 + opponents: [random, discard-only, heuristic-cautious] + max_steps: 10000 + num_workers: 1 diff --git a/configs/ismcts/mini.yaml b/configs/ismcts/mini.yaml new file mode 100644 index 0000000..4844ff4 --- /dev/null +++ b/configs/ismcts/mini.yaml @@ -0,0 +1,46 @@ +run: + experiment_name: ismcts-mini + max_iterations: 50 + seed: 1 + device: cpu +rules: + n_colors: 3 + n_ranks: 5 + min_rank: 2 + n_handshakes: 1 + hand_size: 4 + expedition_penalty: -20 + bonus_threshold: 4 + bonus_amount: 20 +encoding: + derived_playability: true + slot_aware_playability: true +network: + kind: mlp + hidden_size: 128 + num_layers: 2 + activation: relu +mcts: + n_simulations: 50 + c_puct: 1.5 + max_depth: 100 +temperature: + training: 1.0 + eval: 0.0 +training: + games_per_iter: 10 + gradient_steps_per_iter: 10 + batch_size: 128 + replay_capacity: 50000 +optimization: + learning_rate: 0.001 + grad_clip: 5.0 +checkpoint: + save_every: 0 + save_latest: true +evaluation: + eval_every: 5 + games: 20 + opponents: [random, discard-only, heuristic-cautious] + max_steps: 300 + num_workers: 1 diff --git a/pyproject.toml b/pyproject.toml index b27daf4..cf705a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,7 @@ lost-cities-classic = "coolrl_lost_cities.games.classic:main" lost-cities-eval = "coolrl_lost_cities.games.classic.evaluation:main" lost-cities-classic-gui = "coolrl_lost_cities.games.classic.pygame_pvp:main" lost-cities-deep-cfr = "coolrl_lost_cities.games.classic.deep_cfr.cli:main" +lost-cities-ismcts = "coolrl_lost_cities.games.classic.ismcts.cli:main" [dependency-groups] dev = [ diff --git a/src/coolrl_lost_cities/games/classic/ismcts/__init__.py b/src/coolrl_lost_cities/games/classic/ismcts/__init__.py new file mode 100644 index 0000000..c1e7a47 --- /dev/null +++ b/src/coolrl_lost_cities/games/classic/ismcts/__init__.py @@ -0,0 +1 @@ +"""Single-observer IS-MCTS AlphaZero-style training for Lost Cities.""" diff --git a/src/coolrl_lost_cities/games/classic/ismcts/cli.py b/src/coolrl_lost_cities/games/classic/ismcts/cli.py new file mode 100644 index 0000000..cbd3639 --- /dev/null +++ b/src/coolrl_lost_cities/games/classic/ismcts/cli.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +import argparse +import re +from datetime import datetime +from pathlib import Path +from typing import Any + +import yaml + +from .config import IsMctsConfig, load_config +from .trainer import IsMctsTrainer + +_SLUG_RE = re.compile(r"[^a-z0-9]+") + + +def _kebab_slug(value: str) -> str: + return _SLUG_RE.sub("-", value.strip().lower()).strip("-") or "run" + + +def _deep_update(base: dict[str, Any], patch: dict[str, Any]) -> None: + for key, value in patch.items(): + if isinstance(value, dict) and isinstance(base.get(key), dict): + _deep_update(base[key], value) + else: + base[key] = value + + +def _set_path_override(overrides: dict[str, Any], assignment: str) -> None: + if "=" not in assignment: + raise ValueError(f"config override must be PATH=VALUE: {assignment}") + path, raw_value = assignment.split("=", 1) + keys = path.split(".") + value = yaml.safe_load(raw_value) + cursor = overrides + for key in keys[:-1]: + cursor = cursor.setdefault(key, {}) + cursor[keys[-1]] = value + + +def _with_overrides(config: IsMctsConfig, assignments: list[str]) -> IsMctsConfig: + overrides: dict[str, Any] = {} + for assignment in assignments: + _set_path_override(overrides, assignment) + data = config.model_dump(mode="python") + _deep_update(data, overrides) + return IsMctsConfig.model_validate(data) + + +def _resolve_run_dir(config: IsMctsConfig, *, keep: bool) -> Path: + parent = Path("runs") if keep else Path("runs/tmp") + timestamp = datetime.now().strftime("%Y-%m-%d_%H%M%S") + return parent / f"{timestamp}_{_kebab_slug(config.run.experiment_name)}" + + +def train_command(args: argparse.Namespace) -> None: + config = load_config(args.config) if args.config else IsMctsConfig() + config = _with_overrides(config, args.config_overrides) + run_dir = _resolve_run_dir(config, keep=args.keep) + trainer = IsMctsTrainer( + config, + config.rules.to_lost_cities_config(seed=config.run.seed), + run_dir=run_dir, + device=config.run.device, + ) + trainer.train() + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description="Lost Cities SO-ISMCTS tools.") + subparsers = parser.add_subparsers(dest="command", required=True) + train = subparsers.add_parser("train") + train.add_argument("--config") + train.add_argument("--keep", action="store_true") + train.add_argument( + "--set", + action="append", + default=[], + dest="config_overrides", + metavar="PATH=VALUE", + ) + train.add_argument("--wandb", action="store_true", help="Accepted for CLI parity; ignored.") + train.add_argument("--wandb-project", default="coolrl-lost-cities") + train.add_argument("--wandb-name") + train.add_argument("--wandb-mode", choices=("online", "offline", "disabled"), default="online") + train.add_argument("--wandb-group") + train.add_argument("--wandb-job-type") + train.add_argument("--wandb-tag", action="append", default=[]) + train.add_argument("--wandb-notes") + train.set_defaults(func=train_command) + args = parser.parse_args(argv) + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/src/coolrl_lost_cities/games/classic/ismcts/config.py b/src/coolrl_lost_cities/games/classic/ismcts/config.py new file mode 100644 index 0000000..d5c6fb4 --- /dev/null +++ b/src/coolrl_lost_cities/games/classic/ismcts/config.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +import json +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import yaml +from pydantic import Field, field_validator + +from coolrl_lost_cities.games.classic.deep_cfr.config import ( + CheckpointConfig, + EncodingConfig, + EvaluationConfig, + NetworkConfig, + OptimizationConfig, + RulesConfig, + RunConfig, + StrictModel, +) + + +class MctsConfig(StrictModel): + n_simulations: int = 50 + c_puct: float = 1.5 + max_depth: int = 200 + use_rollout_value: bool = True + + @field_validator("n_simulations", "max_depth") + @classmethod + def _positive_int(cls, value: int) -> int: + if value <= 0: + raise ValueError("must be positive") + return value + + +class TemperatureConfig(StrictModel): + training: float = 1.0 + eval: float = 0.0 + + +class TrainingConfig(StrictModel): + games_per_iter: int = 10 + gradient_steps_per_iter: int = 10 + batch_size: int = 128 + replay_capacity: int = 100_000 + + @field_validator("games_per_iter", "gradient_steps_per_iter", "batch_size", "replay_capacity") + @classmethod + def _positive_int(cls, value: int) -> int: + if value <= 0: + raise ValueError("must be positive") + return value + + +class IsMctsConfig(StrictModel): + run: RunConfig = Field(default_factory=lambda: RunConfig(experiment_name="ismcts")) + rules: RulesConfig = Field(default_factory=RulesConfig) + encoding: EncodingConfig = Field(default_factory=EncodingConfig) + network: NetworkConfig = Field( + default_factory=lambda: NetworkConfig(hidden_size=512, num_layers=3) + ) + mcts: MctsConfig = Field(default_factory=MctsConfig) + temperature: TemperatureConfig = Field(default_factory=TemperatureConfig) + training: TrainingConfig = Field(default_factory=TrainingConfig) + optimization: OptimizationConfig = Field(default_factory=OptimizationConfig) + checkpoint: CheckpointConfig = Field(default_factory=CheckpointConfig) + evaluation: EvaluationConfig = Field(default_factory=EvaluationConfig) + + def to_dict(self) -> dict[str, Any]: + return self.model_dump(mode="json") + + +def config_from_dict(data: Mapping[str, Any]) -> IsMctsConfig: + return IsMctsConfig.model_validate(data) + + +def load_config(path: str | Path) -> IsMctsConfig: + config_path = Path(path) + text = config_path.read_text(encoding="utf-8") + if config_path.suffix.lower() in {".yaml", ".yml"}: + data = yaml.safe_load(text) or {} + else: + data = json.loads(text) + return config_from_dict(data) diff --git a/src/coolrl_lost_cities/games/classic/ismcts/determinization.py b/src/coolrl_lost_cities/games/classic/ismcts/determinization.py new file mode 100644 index 0000000..1eba494 --- /dev/null +++ b/src/coolrl_lost_cities/games/classic/ismcts/determinization.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +import random + +from coolrl_lost_cities.games.classic.game import GameState + +from .info_set import unseen_cards + + +def sample_determinization(state: GameState, player: int, rng: random.Random) -> GameState: + """Sample a concrete state uniformly from ``player``'s current information set.""" + p = int(player) + opponent = 1 - p + snapshot = state.to_snapshot() + unseen = unseen_cards(state, p) + rng.shuffle(unseen) + opponent_hand_size = len(state.hands[opponent]) + deck_size = len(state.deck) + if len(unseen) != opponent_hand_size + deck_size: + raise ValueError( + "information set card count mismatch: " + f"unseen={len(unseen)} opponent_hand={opponent_hand_size} deck={deck_size}" + ) + snapshot["hands"][opponent] = [card.to_snapshot() for card in unseen[:opponent_hand_size]] + snapshot["deck"] = [card.to_snapshot() for card in unseen[opponent_hand_size:]] + det = GameState.from_snapshot(snapshot) + det.validate_invariants() + return det diff --git a/src/coolrl_lost_cities/games/classic/ismcts/info_set.py b/src/coolrl_lost_cities/games/classic/ismcts/info_set.py new file mode 100644 index 0000000..008dd25 --- /dev/null +++ b/src/coolrl_lost_cities/games/classic/ismcts/info_set.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +import json +from collections import Counter + +from coolrl_lost_cities.games.classic.game import Card, GameState, build_deck + + +def _card_tuple(card: Card) -> tuple[int, int]: + return int(card.color), int(card.rank) + + +def _sorted_cards(cards: list[Card]) -> list[tuple[int, int]]: + return sorted(_card_tuple(card) for card in cards) + + +def canonical_info_set_key(state: GameState, player: int) -> bytes: + """Deterministic key for observable information from ``player``'s POV.""" + p = int(player) + payload = { + "config": state.config.to_snapshot(), + "player": p, + "current_player": int(state.current_player), + "phase": state.phase, + "pending_discarded_color": ( + None if state.pending_discarded_color < 0 else int(state.pending_discarded_color) + ), + "turn_count": int(state.turn_count), + "terminal": bool(state.terminal), + "deck_size": len(state.deck), + "hand": _sorted_cards(state.hands[p]), + "hand_size_opp": len(state.hands[1 - p]), + "expeditions": [ + [[_card_tuple(card) for card in expedition] for expedition in player_expeditions] + for player_expeditions in state.expeditions + ], + "discards": [[_card_tuple(card) for card in discard] for discard in state.discards], + "legal_mask": list(map(bool, state.unified_legal_mask())), + } + return json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8") + + +def visible_cards(state: GameState, player: int) -> list[Card]: + cards: list[Card] = [] + cards.extend(state.hands[int(player)]) + for player_expeditions in state.expeditions: + for expedition in player_expeditions: + cards.extend(expedition) + for discard in state.discards: + cards.extend(discard) + return cards + + +def unseen_cards(state: GameState, player: int) -> list[Card]: + remaining = Counter(_card_tuple(card) for card in build_deck(state.config)) + for card in visible_cards(state, player): + remaining[_card_tuple(card)] -= 1 + cards: list[Card] = [] + for (color, rank), count in remaining.items(): + cards.extend(Card(color, rank) for _ in range(count)) + return cards diff --git a/src/coolrl_lost_cities/games/classic/ismcts/mcts.py b/src/coolrl_lost_cities/games/classic/ismcts/mcts.py new file mode 100644 index 0000000..6ae2f2a --- /dev/null +++ b/src/coolrl_lost_cities/games/classic/ismcts/mcts.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +import math +import random +from dataclasses import dataclass, field + +import numpy as np +import torch + +from coolrl_lost_cities.games.classic.deep_cfr.encoding import encode_info_state +from coolrl_lost_cities.games.classic.game import GameState + +from .config import MctsConfig +from .determinization import sample_determinization +from .info_set import canonical_info_set_key +from .network import AlphaZeroNet + + +@dataclass +class MctsNode: + info_set_key: bytes + player: int + priors: dict[int, float] = field(default_factory=dict) + visits: dict[int, int] = field(default_factory=dict) + value_sum: dict[int, float] = field(default_factory=dict) + children: dict[int, bytes] = field(default_factory=dict) + terminal: bool = False + + def is_expanded(self) -> bool: + return self.terminal or bool(self.priors) + + def q(self, action: int) -> float: + n = self.visits.get(action, 0) + if n <= 0: + return 0.0 + return self.value_sum.get(action, 0.0) / n + + +class MctsTree: + def __init__(self) -> None: + self.nodes: dict[bytes, MctsNode] = {} + + def get_or_create(self, key: bytes, *, player: int, terminal: bool = False) -> MctsNode: + node = self.nodes.get(key) + if node is None: + node = MctsNode(key, player=player, terminal=terminal) + self.nodes[key] = node + return node + + +class IsMctsSearcher: + def __init__( + self, + network: AlphaZeroNet, + config: MctsConfig, + *, + device: torch.device | str = "cpu", + encoding=None, + rng: random.Random | None = None, + ) -> None: + self.network = network + self.config = config + self.device = torch.device(device) + self.encoding = encoding + self.rng = rng or random.Random() + self.tree = MctsTree() + + def search( + self, + state: GameState, + traverser: int, + n_sims: int | None = None, + ) -> dict[int, int]: + root_key = canonical_info_set_key(state, state.current_player) + root = self.tree.get_or_create( + root_key, player=state.current_player, terminal=state.terminal + ) + sims = int(n_sims or self.config.n_simulations) + for _ in range(sims): + det = sample_determinization(state, traverser, self.rng) + self._simulate(det, depth=0) + legal = state.unified_legal_actions() + return {action: root.visits.get(action, 0) for action in legal} + + def _simulate(self, state: GameState, *, depth: int) -> float: + player = int(state.current_player) + if state.terminal or depth >= self.config.max_depth: + return float(state.score_diff(player)) + + key = canonical_info_set_key(state, player) + node = self.tree.get_or_create(key, player=player, terminal=state.terminal) + if not node.is_expanded(): + value = self._expand_and_evaluate(node, state, player) + return value + + action = self._select_action(node, state.unified_legal_actions()) + child = state.clone() + child.apply_unified_action(action) + child_value = self._simulate(child, depth=depth + 1) + value = child_value if child.current_player == player else -child_value + node.visits[action] = node.visits.get(action, 0) + 1 + node.value_sum[action] = node.value_sum.get(action, 0.0) + value + child_key = canonical_info_set_key(child, child.current_player) + node.children[action] = child_key + return value + + def _expand_and_evaluate(self, node: MctsNode, state: GameState, player: int) -> float: + legal_actions = state.unified_legal_actions() + if not legal_actions: + node.terminal = True + return float(state.score_diff(player)) + info = encode_info_state(state, player, self.encoding) + legal_mask = np.asarray(state.unified_legal_mask(), dtype=bool) + with torch.inference_mode(): + x = torch.as_tensor(info[None, :], dtype=torch.float32, device=self.device) + mask = torch.as_tensor(legal_mask[None, :], dtype=torch.bool, device=self.device) + probs = self.network.policy_distribution(x, mask).squeeze(0).detach().cpu().numpy() + _logits, network_value = self.network(x, mask) + for action in legal_actions: + node.priors[action] = float(probs[action]) + node.visits.setdefault(action, 0) + node.value_sum.setdefault(action, 0.0) + rollout_value = ( + self._rollout_value(state, player) if self.config.use_rollout_value else None + ) + if rollout_value is None: + return float(network_value.item()) + return rollout_value + + def _select_action(self, node: MctsNode, legal_actions: list[int]) -> int: + total_visits = sum(node.visits.get(action, 0) for action in legal_actions) + sqrt_total = math.sqrt(max(1, total_visits)) + best_score = -float("inf") + best_action = legal_actions[0] + for action in legal_actions: + n = node.visits.get(action, 0) + prior = node.priors.get(action, 0.0) + score = node.q(action) + self.config.c_puct * prior * sqrt_total / (1 + n) + if score > best_score: + best_score = score + best_action = action + return int(best_action) + + def _rollout_value(self, state: GameState, player: int) -> float | None: + rollout = state.clone() + steps = 0 + while not rollout.terminal and steps < self.config.max_depth: + legal = rollout.unified_legal_actions() + if not legal: + break + action = self.rng.choice(legal) + rollout.apply_unified_action(action) + steps += 1 + return float(rollout.score_diff(player)) diff --git a/src/coolrl_lost_cities/games/classic/ismcts/network.py b/src/coolrl_lost_cities/games/classic/ismcts/network.py new file mode 100644 index 0000000..00ff0cc --- /dev/null +++ b/src/coolrl_lost_cities/games/classic/ismcts/network.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +import torch +from torch import nn + +from coolrl_lost_cities.games.classic.deep_cfr.networks import _activation + +from .config import IsMctsConfig + + +class AlphaZeroNet(nn.Module): + def __init__( + self, + input_dim: int, + action_size: int, + hidden_size: int = 512, + *, + num_layers: int = 3, + activation: str = "relu", + value_scale: float = 100.0, + ) -> None: + super().__init__() + self.input_dim = int(input_dim) + self.action_size = int(action_size) + self.value_scale = float(value_scale) + layers: list[nn.Module] = [] + last_dim = self.input_dim + for _ in range(max(0, int(num_layers))): + layers.append(nn.Linear(last_dim, hidden_size)) + layers.append(_activation(activation)) + last_dim = hidden_size + self.backbone = nn.Sequential(*layers) + self.policy_head = nn.Linear(last_dim, self.action_size) + self.value_head = nn.Linear(last_dim, 1) + + @classmethod + def from_config( + cls, + input_dim: int, + action_size: int, + config: IsMctsConfig | object, + ) -> AlphaZeroNet: + network_config = config.network if hasattr(config, "network") else config + return cls( + input_dim, + action_size, + network_config.hidden_size, + num_layers=network_config.num_layers, + activation=network_config.activation, + ) + + def forward( + self, + info_state_tensor: torch.Tensor, + legal_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + hidden = self.backbone(info_state_tensor) + logits = self.policy_head(hidden) + value = torch.tanh(self.value_head(hidden)).squeeze(-1) * self.value_scale + if legal_mask is not None: + logits = logits.masked_fill(~legal_mask.bool(), torch.finfo(logits.dtype).min) + return logits, value + + def policy_distribution( + self, + info_state_tensor: torch.Tensor, + legal_mask: torch.Tensor, + ) -> torch.Tensor: + logits, _value = self.forward(info_state_tensor, legal_mask) + probs = torch.softmax(logits, dim=-1).masked_fill(~legal_mask.bool(), 0.0) + normalizer = probs.sum(dim=-1, keepdim=True).clamp_min(1.0e-12) + return probs / normalizer diff --git a/src/coolrl_lost_cities/games/classic/ismcts/policy.py b/src/coolrl_lost_cities/games/classic/ismcts/policy.py new file mode 100644 index 0000000..3d7c2a4 --- /dev/null +++ b/src/coolrl_lost_cities/games/classic/ismcts/policy.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import numpy as np +import torch + +from coolrl_lost_cities.games.classic.deep_cfr.encoding import encode_info_state +from coolrl_lost_cities.games.classic.game import GameState +from coolrl_lost_cities.games.classic.policy import LostCitiesPolicy, PolicyInput + +from .network import AlphaZeroNet + + +class AlphaZeroPolicy(LostCitiesPolicy): + def __init__( + self, + network: AlphaZeroNet, + *, + device: torch.device | str = "cpu", + encoding=None, + sample: bool = False, + seed: int | None = None, + ) -> None: + self.network = network + self.device = torch.device(device) + self.encoding = encoding + self.sample = sample + self.rng = np.random.default_rng(seed) + + def act(self, obs_or_state: PolicyInput) -> int: + if not isinstance(obs_or_state, GameState): + legal = np.asarray(obs_or_state["legal_mask"], dtype=bool) + return int(np.flatnonzero(legal)[0]) + state = obs_or_state + legal = np.asarray(state.unified_legal_mask(), dtype=bool) + info = encode_info_state(state, state.current_player, self.encoding) + with torch.inference_mode(): + x = torch.as_tensor(info[None, :], dtype=torch.float32, device=self.device) + mask = torch.as_tensor(legal[None, :], dtype=torch.bool, device=self.device) + probs = self.network.policy_distribution(x, mask).squeeze(0).cpu().numpy() + legal_actions = np.flatnonzero(legal) + if self.sample: + unified = int(self.rng.choice(legal_actions, p=probs[legal_actions])) + else: + unified = int(legal_actions[int(np.argmax(probs[legal_actions]))]) + return state.from_unified_action(unified) diff --git a/src/coolrl_lost_cities/games/classic/ismcts/replay_buffer.py b/src/coolrl_lost_cities/games/classic/ismcts/replay_buffer.py new file mode 100644 index 0000000..dbec70b --- /dev/null +++ b/src/coolrl_lost_cities/games/classic/ismcts/replay_buffer.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from collections import deque +from collections.abc import Iterable +from dataclasses import dataclass + +import numpy as np + + +@dataclass +class ReplaySample: + info_state: np.ndarray + legal_mask: np.ndarray + pi_target: np.ndarray + v_target: float + player: int + + +class ReplayBuffer: + def __init__(self, capacity: int, *, seed: int | None = None) -> None: + self.capacity = int(capacity) + self._items: deque[ReplaySample] = deque(maxlen=self.capacity) + self.rng = np.random.default_rng(seed) + + def __len__(self) -> int: + return len(self._items) + + def add(self, samples: Iterable[ReplaySample]) -> None: + self._items.extend(samples) + + def sample(self, batch_size: int) -> list[ReplaySample]: + if not self._items: + raise ValueError("cannot sample from an empty replay buffer") + size = min(int(batch_size), len(self._items)) + indices = self.rng.choice(len(self._items), size=size, replace=False) + items = list(self._items) + return [items[int(index)] for index in indices] diff --git a/src/coolrl_lost_cities/games/classic/ismcts/self_play.py b/src/coolrl_lost_cities/games/classic/ismcts/self_play.py new file mode 100644 index 0000000..c092ded --- /dev/null +++ b/src/coolrl_lost_cities/games/classic/ismcts/self_play.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import random + +import numpy as np +import torch + +from coolrl_lost_cities.games.classic.deep_cfr.encoding import encode_info_state +from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig + +from .config import MctsConfig +from .mcts import IsMctsSearcher +from .network import AlphaZeroNet +from .replay_buffer import ReplaySample + + +def visit_distribution( + visit_counts: dict[int, int], + action_size: int, + *, + temperature: float, +) -> np.ndarray: + pi = np.zeros(action_size, dtype=np.float32) + if not visit_counts: + return pi + actions = np.asarray(list(visit_counts), dtype=np.int64) + counts = np.asarray([visit_counts[int(action)] for action in actions], dtype=np.float64) + if temperature <= 0.0: + best = int(actions[int(np.argmax(counts))]) + pi[best] = 1.0 + return pi + adjusted = np.power(np.maximum(counts, 1.0e-12), 1.0 / temperature) + adjusted /= adjusted.sum() + pi[actions] = adjusted.astype(np.float32) + return pi + + +def select_from_distribution(pi: np.ndarray, rng: random.Random) -> int: + total = float(pi.sum()) + if total <= 0.0: + raise RuntimeError("empty action distribution") + threshold = rng.random() * total + cumsum = 0.0 + for action, prob in enumerate(pi): + cumsum += float(prob) + if cumsum >= threshold: + return action + return int(len(pi) - 1) + + +def play_self_play_game( + network: AlphaZeroNet, + mcts_config: MctsConfig, + game_config: LostCitiesConfig, + rng: random.Random, + *, + device: torch.device | str = "cpu", + encoding=None, + temperature: float = 1.0, + max_steps: int = 10_000, +) -> list[ReplaySample]: + state = GameState.new_game(game_config, seed=rng.randrange(2**31)) + pending: list[tuple[np.ndarray, np.ndarray, np.ndarray, int]] = [] + steps = 0 + while not state.terminal and steps < max_steps: + player = int(state.current_player) + searcher = IsMctsSearcher( + network, + mcts_config, + device=device, + encoding=encoding, + rng=random.Random(rng.randrange(2**31)), + ) + visits = searcher.search(state, player) + legal_mask = np.asarray(state.unified_legal_mask(), dtype=bool) + pi = visit_distribution(visits, state.action_size, temperature=temperature) + if pi.sum() <= 0: + legal_actions = np.flatnonzero(legal_mask) + pi[legal_actions] = 1.0 / len(legal_actions) + info = encode_info_state(state, player, encoding) + pending.append((info.astype(np.float32), legal_mask, pi, player)) + action = select_from_distribution(pi, rng) + state.apply_unified_action(action) + steps += 1 + + final_diff0 = float(state.score_diff(0)) + samples: list[ReplaySample] = [] + for info, legal_mask, pi, player in pending: + value = final_diff0 if player == 0 else -final_diff0 + samples.append( + ReplaySample( + info_state=info, + legal_mask=legal_mask.astype(bool), + pi_target=pi.astype(np.float32), + v_target=value, + player=player, + ) + ) + return samples diff --git a/src/coolrl_lost_cities/games/classic/ismcts/trainer.py b/src/coolrl_lost_cities/games/classic/ismcts/trainer.py new file mode 100644 index 0000000..f57de03 --- /dev/null +++ b/src/coolrl_lost_cities/games/classic/ismcts/trainer.py @@ -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), + } diff --git a/tests/games/classic/ismcts/test_ismcts.py b/tests/games/classic/ismcts/test_ismcts.py new file mode 100644 index 0000000..65193be --- /dev/null +++ b/tests/games/classic/ismcts/test_ismcts.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +import random + +import numpy as np +import torch +from coolrl_lost_cities.games.classic.deep_cfr.encoding import encode_info_state, input_dim +from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig + +from coolrl_lost_cities.games.classic.ismcts.config import IsMctsConfig, MctsConfig +from coolrl_lost_cities.games.classic.ismcts.determinization import sample_determinization +from coolrl_lost_cities.games.classic.ismcts.info_set import canonical_info_set_key +from coolrl_lost_cities.games.classic.ismcts.mcts import IsMctsSearcher +from coolrl_lost_cities.games.classic.ismcts.network import AlphaZeroNet +from coolrl_lost_cities.games.classic.ismcts.replay_buffer import ReplayBuffer, ReplaySample +from coolrl_lost_cities.games.classic.ismcts.self_play import play_self_play_game +from coolrl_lost_cities.games.classic.ismcts.trainer import IsMctsTrainer + + +def mini_config(seed: int = 1) -> LostCitiesConfig: + return LostCitiesConfig( + n_colors=3, + n_ranks=5, + n_handshakes=1, + hand_size=4, + bonus_threshold=4, + seed=seed, + ) + + +def test_canonical_key_ignores_opponent_hand_identity() -> None: + config = mini_config() + state = GameState.new_game(config, seed=3) + snap = state.to_snapshot() + snap["hands"][1] = list(reversed(snap["hands"][1])) + other = GameState.from_snapshot(snap) + assert canonical_info_set_key(state, 0) == canonical_info_set_key(other, 0) + + +def test_determinization_consistent_with_info_set() -> None: + config = mini_config() + state = GameState.new_game(config, seed=4) + rng = random.Random(5) + key = canonical_info_set_key(state, 0) + for _ in range(100): + det = sample_determinization(state, 0, rng) + assert canonical_info_set_key(det, 0) == key + assert len(det.hands[1]) == config.hand_size + all_cards = det.hands[0] + det.hands[1] + det.deck + for player_expeditions in det.expeditions: + for expedition in player_expeditions: + all_cards.extend(expedition) + for discard in det.discards: + all_cards.extend(discard) + assert len(all_cards) == config.deck_size + + +def test_network_shapes_and_mask() -> None: + state = GameState.new_game(mini_config(), seed=6) + dim = input_dim(state) + net = AlphaZeroNet(dim, state.action_size, hidden_size=16, num_layers=1) + x = torch.as_tensor(encode_info_state(state, 0)[None, :], dtype=torch.float32) + mask = torch.as_tensor(np.asarray(state.unified_legal_mask(), dtype=bool)[None, :]) + logits, value = net(x, mask) + probs = net.policy_distribution(x, mask) + assert logits.shape == (1, state.action_size) + assert value.shape == (1,) + assert torch.allclose(probs.sum(dim=-1), torch.ones(1)) + assert torch.all(probs[~mask] == 0) + + +def test_mcts_prior_drives_visits() -> None: + state = GameState.new_game(mini_config(), seed=7) + dim = input_dim(state) + net = AlphaZeroNet(dim, state.action_size, hidden_size=8, num_layers=0) + for param in net.parameters(): + param.data.zero_() + legal = state.unified_legal_actions() + favored = legal[0] + net.policy_head.bias.data[favored] = 5.0 + searcher = IsMctsSearcher( + net, + MctsConfig(n_simulations=12, c_puct=2.0, use_rollout_value=False), + rng=random.Random(8), + ) + visits = searcher.search(state, state.current_player) + assert visits[favored] == max(visits.values()) + + +def test_replay_buffer_capacity_and_sample() -> None: + sample = ReplaySample( + info_state=np.zeros(4, dtype=np.float32), + legal_mask=np.ones(3, dtype=bool), + pi_target=np.array([1.0, 0.0, 0.0], dtype=np.float32), + v_target=1.0, + player=0, + ) + buffer = ReplayBuffer(2, seed=1) + buffer.add([sample, sample, sample]) + assert len(buffer) == 2 + assert len(buffer.sample(2)) == 2 + + +def test_self_play_game_returns_signed_targets() -> None: + config = mini_config() + state = GameState.new_game(config, seed=9) + net = AlphaZeroNet(input_dim(state), state.action_size, hidden_size=8, num_layers=1) + samples = play_self_play_game( + net, + MctsConfig(n_simulations=2), + config, + random.Random(10), + ) + assert samples + assert {sample.player for sample in samples} <= {0, 1} + assert all(sample.pi_target.sum() > 0 for sample in samples) + + +def test_trainer_one_iteration_smoke(tmp_path) -> None: + config = IsMctsConfig.model_validate( + { + "run": {"max_iterations": 1, "seed": 11, "device": "cpu"}, + "rules": { + "n_colors": 3, + "n_ranks": 5, + "n_handshakes": 1, + "hand_size": 4, + "bonus_threshold": 4, + }, + "network": {"hidden_size": 16, "num_layers": 1}, + "mcts": {"n_simulations": 2}, + "training": {"games_per_iter": 1, "gradient_steps_per_iter": 1, "batch_size": 8}, + "checkpoint": {"save_every": 0}, + "evaluation": {"eval_every": 0, "num_workers": 1, "max_steps": 80}, + } + ) + trainer = IsMctsTrainer( + config, + config.rules.to_lost_cities_config(seed=config.run.seed), + run_dir=tmp_path, + ) + metrics = trainer.train() + assert len(metrics) == 1 + assert (tmp_path / "metrics.jsonl").exists()