From 9fdfa88b2308db867076a5470c3f576683693a36 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=A0=95=EC=8B=9C=EC=9B=90?= Date: Mon, 11 May 2026 13:36:28 +0900 Subject: [PATCH] Add lost-cities-ismcts pretrain: behavior-clone heuristic into network --- .../games/classic/ismcts/cli.py | 9 + .../games/classic/ismcts/eval_checkpoint.py | 12 +- .../games/classic/ismcts/pretrain.py | 279 ++++++++++++++++++ 3 files changed, 299 insertions(+), 1 deletion(-) create mode 100644 src/coolrl_lost_cities/games/classic/ismcts/pretrain.py diff --git a/src/coolrl_lost_cities/games/classic/ismcts/cli.py b/src/coolrl_lost_cities/games/classic/ismcts/cli.py index 45e2f9b..8f1427a 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/cli.py +++ b/src/coolrl_lost_cities/games/classic/ismcts/cli.py @@ -140,6 +140,15 @@ def main(argv: list[str] | None = None) -> None: add_eval_args(eval_cmd) eval_cmd.set_defaults(func=lambda a: run_eval(a)) + from .pretrain import add_pretrain_args, run_pretrain + + pretrain_cmd = subparsers.add_parser( + "pretrain", + help="Behavior-clone a heuristic bot into the AlphaZero network as a warm start.", + ) + add_pretrain_args(pretrain_cmd) + pretrain_cmd.set_defaults(func=lambda a: run_pretrain(a)) + args = parser.parse_args(argv) args.func(args) diff --git a/src/coolrl_lost_cities/games/classic/ismcts/eval_checkpoint.py b/src/coolrl_lost_cities/games/classic/ismcts/eval_checkpoint.py index 093d24a..feea2ab 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/eval_checkpoint.py +++ b/src/coolrl_lost_cities/games/classic/ismcts/eval_checkpoint.py @@ -33,6 +33,7 @@ class _WorkerJob: device: str worker_index: int verbose: bool + n_sims_override: int = 0 # 0 means use checkpoint's eval_n_simulations @dataclass @@ -85,7 +86,9 @@ def _run_games(job: _WorkerJob) -> _WorkerResult: net.eval() eval_mcts_cfg = cfg.mcts.model_copy() - if cfg.mcts.eval_n_simulations > 0: + if job.n_sims_override > 0: + eval_mcts_cfg = eval_mcts_cfg.model_copy(update={"n_simulations": job.n_sims_override}) + elif cfg.mcts.eval_n_simulations > 0: eval_mcts_cfg = eval_mcts_cfg.model_copy( update={"n_simulations": cfg.mcts.eval_n_simulations} ) @@ -258,6 +261,12 @@ def add_eval_args(parser: argparse.ArgumentParser) -> None: action="store_true", help="Print per-game result lines (turns/score/PA) in addition to per-worker summaries.", ) + parser.add_argument( + "--n-sims", + type=int, + default=0, + help="Override n_simulations at eval time (0 = use checkpoint's eval_n_simulations).", + ) def run_eval(args: argparse.Namespace) -> None: @@ -294,6 +303,7 @@ def run_eval(args: argparse.Namespace) -> None: device=args.device, worker_index=i, verbose=args.verbose, + n_sims_override=int(args.n_sims), ) for i in range(len(slices)) ] diff --git a/src/coolrl_lost_cities/games/classic/ismcts/pretrain.py b/src/coolrl_lost_cities/games/classic/ismcts/pretrain.py new file mode 100644 index 0000000..2fccedf --- /dev/null +++ b/src/coolrl_lost_cities/games/classic/ismcts/pretrain.py @@ -0,0 +1,279 @@ +"""Behavior cloning warm-start for the SO-ISMCTS network. + +Generates games from a fixed heuristic policy vs itself, then trains the +AlphaZero-style network's policy + value heads in supervised fashion: + +- policy loss: cross-entropy between network logits and the heuristic's chosen + action (one-hot target, masked to legal actions). +- value loss: MSE on the final game score diff from each decision-maker's + perspective, normalized by value_scale (same convention as the trainer). + +The resulting checkpoint can be passed to `lost-cities-ismcts train +--resume-from` to start self-play with a heuristic-level prior instead of a +random init. This addresses the self-play "weak-equilibrium" problem: starting +from random, MCTS visit distributions converge to a mutually mediocre policy +that has near-zero win rate against the heuristic. Warm-starting at heuristic +level gives self-play a meaningful baseline to improve from. + +Invoked via ``lost-cities-ismcts pretrain``. +""" + +from __future__ import annotations + +import argparse +import random +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn + +from coolrl_lost_cities.games.classic.bots.registry import build_bot +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 .config import IsMctsConfig, load_config +from .network import AlphaZeroNet + + +def _collect_samples( + game_config: LostCitiesConfig, + encoding, + n_games: int, + bot_name: str = "heuristic-balanced", + seed: int = 0, + max_turns: int = 500, +) -> list[tuple[np.ndarray, np.ndarray, int, float]]: + """Roll out n_games of bot vs bot, returning per-decision samples. + + Each sample: (info_state, legal_mask, action_idx, value). + value is the player-perspective score diff at game end. + """ + samples: list[tuple[np.ndarray, np.ndarray, int, float]] = [] + for game_idx in range(n_games): + bots = [ + build_bot(bot_name, seed=seed + game_idx * 2), + build_bot(bot_name, seed=seed + game_idx * 2 + 1), + ] + state = GameState.new_game(game_config, seed=seed + game_idx) + decisions: list[tuple[np.ndarray, np.ndarray, int, int]] = [] + turns = 0 + while not state.terminal and turns < max_turns: + player = int(state.current_player) + info_state = encode_info_state(state, player, encoding) + legal_mask = np.asarray(state.unified_legal_mask(), dtype=bool) + phase_action = bots[player].act(state) + unified = state.to_unified_action(phase_action) + decisions.append((info_state, legal_mask, int(unified), player)) + state.apply_unified_action(unified) + turns += 1 + final_diff0 = float(state.score_diff(0)) + for info, mask, act, player in decisions: + value = final_diff0 if player == 0 else -final_diff0 + samples.append((info, mask, act, value)) + return samples + + +def _train_supervised( + network: AlphaZeroNet, + samples: list, + device: torch.device, + *, + epochs: int, + batch_size: int, + lr: float, + weight_decay: float, + grad_clip: float, + value_loss_weight: float, +) -> None: + network.train() + optimizer = torch.optim.AdamW(network.parameters(), lr=lr, weight_decay=max(weight_decay, 1e-4)) + rng = random.Random(0) + indices = list(range(len(samples))) + v_scale = float(network.value_scale) + + for epoch in range(1, epochs + 1): + rng.shuffle(indices) + n_batches = (len(indices) + batch_size - 1) // batch_size + epoch_pl = 0.0 + epoch_vl = 0.0 + epoch_acc = 0.0 + epoch_n = 0 + for b in range(n_batches): + batch_idx = indices[b * batch_size : (b + 1) * batch_size] + infos = np.stack([samples[i][0] for i in batch_idx]) + masks = np.stack([samples[i][1] for i in batch_idx]) + actions = np.array([samples[i][2] for i in batch_idx], dtype=np.int64) + values = np.array([samples[i][3] for i in batch_idx], dtype=np.float32) + + info_t = torch.as_tensor(infos, dtype=torch.float32, device=device) + mask_t = torch.as_tensor(masks, dtype=torch.bool, device=device) + action_t = torch.as_tensor(actions, device=device) + value_t = torch.as_tensor(values, device=device) + + logits, value_pred = network(info_t, mask_t) + policy_loss = F.cross_entropy(logits, action_t) + value_loss = F.mse_loss(value_pred / v_scale, value_t / v_scale) + loss = policy_loss + value_loss_weight * value_loss + + optimizer.zero_grad(set_to_none=True) + loss.backward() + if grad_clip > 0: + nn.utils.clip_grad_norm_(network.parameters(), grad_clip) + optimizer.step() + + with torch.no_grad(): + preds = logits.argmax(dim=-1) + acc = (preds == action_t).float().mean().item() + epoch_pl += float(policy_loss.item()) * len(batch_idx) + epoch_vl += float(value_loss.item()) * len(batch_idx) + epoch_acc += acc * len(batch_idx) + epoch_n += len(batch_idx) + + print( + f" epoch {epoch:3d}: policy_loss={epoch_pl / epoch_n:.4f} " + f"value_loss={epoch_vl / epoch_n:.4f} " + f"top1_match={epoch_acc / epoch_n:.3f}", + flush=True, + ) + + network.eval() + return optimizer + + +def add_pretrain_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--config", default=None, help="ISMCTS config YAML (controls network shape + rules)." + ) + parser.add_argument( + "--set", + action="append", + default=[], + dest="config_overrides", + metavar="PATH=VALUE", + ) + parser.add_argument( + "--bot", + default="heuristic-balanced", + help="Bot to clone (default: heuristic-balanced).", + ) + parser.add_argument( + "--games", type=int, default=2000, help="Number of bot-vs-bot games to roll out." + ) + parser.add_argument("--epochs", type=int, default=10, help="Supervised training epochs.") + parser.add_argument("--batch-size", type=int, default=256) + parser.add_argument("--lr", type=float, default=3.0e-4) + parser.add_argument("--weight-decay", type=float, default=1.0e-4) + parser.add_argument("--grad-clip", type=float, default=5.0) + parser.add_argument( + "--value-loss-weight", + type=float, + default=50.0, + help="Multiplier on value MSE (raw_MSE / value_scale^2). Default 50 to " + "make value loss magnitude comparable to policy CE.", + ) + parser.add_argument( + "--out", + default="runs/pretrain/heuristic_clone.pt", + help="Output checkpoint path (compatible with --resume-from).", + ) + parser.add_argument("--device", default="cuda") + parser.add_argument("--seed", type=int, default=12345) + + +def _apply_config_overrides(config: IsMctsConfig, assignments: list[str]) -> IsMctsConfig: + import yaml + + def deep_update(base: dict, patch: dict) -> None: + for k, v in patch.items(): + if isinstance(v, dict) and isinstance(base.get(k), dict): + deep_update(base[k], v) + else: + base[k] = v + + overrides: dict = {} + for assignment in assignments: + if "=" not in assignment: + raise ValueError(f"override must be PATH=VALUE: {assignment}") + path, raw_value = assignment.split("=", 1) + value = yaml.safe_load(raw_value) + cursor = overrides + keys = path.split(".") + for k in keys[:-1]: + cursor = cursor.setdefault(k, {}) + cursor[keys[-1]] = value + data = config.model_dump(mode="python") + deep_update(data, overrides) + return IsMctsConfig.model_validate(data) + + +def run_pretrain(args: argparse.Namespace) -> None: + config = load_config(args.config) if args.config else IsMctsConfig() + config = _apply_config_overrides(config, args.config_overrides) + game_config = config.rules.to_lost_cities_config(seed=config.run.seed) + device = torch.device(args.device) + + probe = GameState.new_game(game_config, seed=config.run.seed) + in_dim = input_dim(probe, config.encoding) + network = AlphaZeroNet.from_config(in_dim, probe.action_size, config).to(device) + print( + f"network: input_dim={in_dim} action_size={probe.action_size} " + f"hidden={config.network.hidden_size} layers={config.network.num_layers}", + flush=True, + ) + + print(f"rolling out {args.games} games of {args.bot} vs {args.bot}...", flush=True) + t0 = time.perf_counter() + samples = _collect_samples( + game_config, + config.encoding, + n_games=args.games, + bot_name=args.bot, + seed=args.seed, + ) + rollout_secs = time.perf_counter() - t0 + print( + f"collected {len(samples)} decisions from {args.games} games " + f"in {rollout_secs:.1f}s ({len(samples) / args.games:.1f} decisions/game)", + flush=True, + ) + + print( + f"training: {args.epochs} epochs, batch={args.batch_size}, lr={args.lr}, " + f"value_weight={args.value_loss_weight}", + flush=True, + ) + t1 = time.perf_counter() + optimizer = _train_supervised( + network, + samples, + device, + epochs=args.epochs, + batch_size=args.batch_size, + lr=args.lr, + weight_decay=args.weight_decay, + grad_clip=args.grad_clip, + value_loss_weight=args.value_loss_weight, + ) + train_secs = time.perf_counter() - t1 + print(f"training done in {train_secs:.1f}s", flush=True) + + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + payload = { + "config": config.to_dict(), + "game_config": game_config.to_snapshot(), + "iteration": 0, + "network": network.state_dict(), + "optimizer": optimizer.state_dict(), + "metrics": { + "pretrain/games": args.games, + "pretrain/samples": len(samples), + "pretrain/epochs": args.epochs, + "pretrain/bot": args.bot, + }, + } + torch.save(payload, out_path) + print(f"saved pretrained checkpoint to {out_path}", flush=True)