diff --git a/configs/ismcts/default.yaml b/configs/ismcts/default.yaml index 4ed8153..c23b991 100644 --- a/configs/ismcts/default.yaml +++ b/configs/ismcts/default.yaml @@ -36,7 +36,8 @@ training: replay_capacity: 100000 interleave_games: 8 interleave_max_batch: 64 - use_inference_server: true + use_central_scheduler: false + use_inference_server: false inference_server_max_batch: 128 inference_server_batch_timeout_ms: 10.0 optimization: diff --git a/src/coolrl_lost_cities/games/classic/game.pxd b/src/coolrl_lost_cities/games/classic/game.pxd index c1e0898..8077d29 100644 --- a/src/coolrl_lost_cities/games/classic/game.pxd +++ b/src/coolrl_lost_cities/games/classic/game.pxd @@ -57,6 +57,7 @@ cdef class GameState: cdef void _clear(self) noexcept cpdef GameState clone(self) + cpdef GameState determinize_for_player(self, int player, object rng) cpdef list legal_card_mask(self) cpdef list legal_draw_mask(self) cpdef list legal_mask(self) diff --git a/src/coolrl_lost_cities/games/classic/game.pyx b/src/coolrl_lost_cities/games/classic/game.pyx index 130e0bb..8df81c2 100644 --- a/src/coolrl_lost_cities/games/classic/game.pyx +++ b/src/coolrl_lost_cities/games/classic/game.pyx @@ -595,6 +595,29 @@ cdef class GameState: other.terminal = self.terminal return other + cpdef GameState determinize_for_player(self, int player, object rng): + """Clone and reshuffle hidden opponent hand/deck cards for ``player``.""" + cdef int p = int(player) + cdef int opponent = 1 - p + cdef int opponent_hand_len = self.hand_lens[opponent] + cdef int unseen_len = opponent_hand_len + self.deck_len + cdef int i + cdef list unseen = [0] * unseen_len + cdef GameState other + if p < 0 or p > 1: + raise ValueError(f"player must be 0 or 1, got {player}") + for i in range(opponent_hand_len): + unseen[i] = self.hand_cards[self._hand_index(opponent, i)] + for i in range(self.deck_len): + unseen[opponent_hand_len + i] = self.deck_cards[i] + rng.shuffle(unseen) + other = self.clone() + for i in range(opponent_hand_len): + other.hand_cards[other._hand_index(opponent, i)] = unseen[i] + for i in range(self.deck_len): + other.deck_cards[i] = unseen[opponent_hand_len + i] + return other + cpdef list legal_card_mask(self): cdef list mask = [False] * (2 * self.hand_size) cdef int slot diff --git a/src/coolrl_lost_cities/games/classic/ismcts/config.py b/src/coolrl_lost_cities/games/classic/ismcts/config.py index bafb0bf..b49f7b9 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/config.py +++ b/src/coolrl_lost_cities/games/classic/ismcts/config.py @@ -60,7 +60,8 @@ class TrainingConfig(StrictModel): interleave_max_batch: int = 64 num_workers: int = 1 worker_device: str = "cpu" - use_inference_server: bool = True + use_central_scheduler: bool = False + use_inference_server: bool = False inference_server_max_batch: int = 128 inference_server_batch_timeout_ms: float = 10.0 diff --git a/src/coolrl_lost_cities/games/classic/ismcts/determinization.py b/src/coolrl_lost_cities/games/classic/ismcts/determinization.py index 1eba494..eb3621e 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/determinization.py +++ b/src/coolrl_lost_cities/games/classic/ismcts/determinization.py @@ -9,6 +9,8 @@ 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.""" + if hasattr(state, "determinize_for_player"): + return state.determinize_for_player(int(player), rng) p = int(player) opponent = 1 - p snapshot = state.to_snapshot() diff --git a/src/coolrl_lost_cities/games/classic/ismcts/evaluate.py b/src/coolrl_lost_cities/games/classic/ismcts/evaluate.py index 19c32d9..ea76cce 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/evaluate.py +++ b/src/coolrl_lost_cities/games/classic/ismcts/evaluate.py @@ -12,6 +12,8 @@ from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig from .config import IsMctsConfig, MctsConfig from .eval_worker import EvalWorkerBatch, init_eval_inference_queues, run_eval_worker +from .info_set import canonical_info_set_key +from .interleaved_self_play import _run_search_jobs, _SearchJob from .mcts import IsMctsSearcher from .network import AlphaZeroNet @@ -289,6 +291,185 @@ def evaluate_opponents_with_mcts_parallel( } +class _EvalContext: + def __init__( + self, + *, + opponent: str, + game_index: int, + state: GameState, + policy_player: int, + opponents, + rng: random.Random, + ) -> None: + self.opponent = opponent + self.game_index = int(game_index) + self.state = state + self.policy_player = int(policy_player) + self.opponents = opponents + self.rng = rng + self.steps = 0 + self.terminated = False + self.policy_turns = 0 + self.play_actions = 0 + self.timeout = False + + +def evaluate_opponents_with_mcts_central( + network: AlphaZeroNet, + game_config: LostCitiesConfig, + mcts_config: MctsConfig, + *, + config: IsMctsConfig, + opponents: tuple[str, ...], + games: int, + seed: int, + device: torch.device | str, + encoding=None, + max_steps: int, +) -> dict[str, dict[str, float | int]]: + started = time.perf_counter() + rng = random.Random(seed) + active: list[_EvalContext] = [] + for opponent in opponents: + for game_index in range(games): + active.append( + _EvalContext( + opponent=opponent, + game_index=game_index, + state=GameState.new_game(game_config, seed=seed + game_index), + policy_player=game_index % 2, + opponents=[ + build_bot(opponent, seed=seed + game_index), + build_bot(opponent, seed=seed + game_index + 1), + ], + rng=random.Random(rng.randrange(2**31)), + ) + ) + aggregate: dict[str, dict[str, object]] = { + opponent: { + "score_diffs": [], + "wins0": 0, + "wins1": 0, + "draws": 0, + "policy_turns": 0, + "play_actions": 0, + "timeouts": 0, + } + for opponent in opponents + } + torch_device = torch.device(device) + width = max(1, int(config.training.interleave_games)) + max_batch = max(1, int(config.training.interleave_max_batch)) + while active: + jobs: list[_SearchJob] = [] + job_contexts: list[_EvalContext] = [] + still_active: list[_EvalContext] = [] + for context in active[:width]: + if context.state.terminal: + context.terminated = True + _record_eval_context(context, aggregate) + continue + if context.steps >= max_steps: + context.timeout = True + _record_eval_context(context, aggregate) + continue + current = int(context.state.current_player) + if current != context.policy_player: + action = context.opponents[current].act(context.state) + context.state.apply_action(action) + context.steps += 1 + still_active.append(context) + continue + searcher = IsMctsSearcher( + network, + mcts_config, + device=torch_device, + encoding=encoding, + rng=random.Random(context.rng.randrange(2**31)), + ) + jobs.append( + _SearchJob( + context=context, + searcher=searcher, + traverser=current, + remaining=mcts_config.n_simulations, + ) + ) + job_contexts.append(context) + still_active.append(context) + still_active.extend(active[width:]) + active = still_active + if jobs: + _run_search_jobs(network, jobs, max_batch, torch_device) + for job, context in zip(jobs, job_contexts, strict=True): + root_key = canonical_info_set_key(context.state, context.state.current_player) + root = job.searcher.tree.get_or_create( + root_key, + player=context.state.current_player, + terminal=context.state.terminal, + ) + visits = { + action: root.visits.get(action, 0) + for action in context.state.unified_legal_actions() + } + unified = ( + max(visits, key=visits.get) + if visits + else context.state.unified_legal_actions()[0] + ) + if context.state.phase == "card": + context.policy_turns += 1 + if unified % 2 == 0: + context.play_actions += 1 + context.state.apply_unified_action(unified) + context.steps += 1 + active = [ + context + for context in active + if not context.state.terminal and context.steps < max_steps + ] + for context in still_active: + if context not in active and (context.state.terminal or context.steps >= max_steps): + if context.steps >= max_steps and not context.state.terminal: + context.timeout = True + _record_eval_context(context, aggregate) + elapsed = time.perf_counter() - started + return { + opponent: _evaluation_metrics( + score_diffs=list(bucket["score_diffs"]), + wins0=int(bucket["wins0"]), + wins1=int(bucket["wins1"]), + draws=int(bucket["draws"]), + policy_turns=int(bucket["policy_turns"]), + play_actions=int(bucket["play_actions"]), + timeouts=int(bucket["timeouts"]), + elapsed_seconds=elapsed, + n=len(bucket["score_diffs"]), + ) + for opponent, bucket in aggregate.items() + } + + +def _record_eval_context( + context: _EvalContext, + aggregate: dict[str, dict[str, object]], +) -> None: + bucket = aggregate[context.opponent] + diff = float(context.state.score_diff(context.policy_player)) + bucket["score_diffs"].append(diff) + if diff > 0: + bucket["wins0"] += 1 + elif diff < 0: + bucket["wins1"] += 1 + else: + bucket["draws"] += 1 + bucket["policy_turns"] += context.policy_turns + bucket["play_actions"] += context.play_actions + if context.timeout: + bucket["timeouts"] += 1 + + def _evaluation_metrics( *, score_diffs: list[float], diff --git a/src/coolrl_lost_cities/games/classic/ismcts/mcts.pyx b/src/coolrl_lost_cities/games/classic/ismcts/mcts.pyx index f142a70..c67130a 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/mcts.pyx +++ b/src/coolrl_lost_cities/games/classic/ismcts/mcts.pyx @@ -12,7 +12,6 @@ from coolrl_lost_cities.games.classic.bots.heuristic import HeuristicBot as PyHe from coolrl_lost_cities.games.classic.deep_cfr.encoding import encode_info_state from coolrl_lost_cities.games.classic.game cimport GameState -from .determinization import sample_determinization from .info_set import canonical_info_set_key @@ -371,7 +370,7 @@ cdef class IsMctsSearcher: return pending cpdef PendingSimulation prepare_simulation(self, GameState root_state, int traverser): - cdef GameState state = sample_determinization(root_state, traverser, self.rng) + cdef GameState state = root_state.determinize_for_player(traverser, self.rng) cdef list path = [] cdef int depth = 0 cdef object cached_key = None diff --git a/src/coolrl_lost_cities/games/classic/ismcts/trainer.py b/src/coolrl_lost_cities/games/classic/ismcts/trainer.py index d2dfe4c..d06db18 100644 --- a/src/coolrl_lost_cities/games/classic/ismcts/trainer.py +++ b/src/coolrl_lost_cities/games/classic/ismcts/trainer.py @@ -17,7 +17,11 @@ from coolrl_lost_cities.games.classic.deep_cfr.evaluate import evaluate_strategy from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig from .config import IsMctsConfig -from .evaluate import evaluate_opponents_with_mcts_parallel, evaluate_with_mcts +from .evaluate import ( + evaluate_opponents_with_mcts_central, + evaluate_opponents_with_mcts_parallel, + evaluate_with_mcts, +) from .inference_server import InferenceServer from .interleaved_self_play import play_self_play_iteration from .network import AlphaZeroLogitsView, AlphaZeroNet @@ -123,7 +127,19 @@ class IsMctsTrainer: ) self.network.eval() sp_started = time.perf_counter() - if self.config.training.num_workers > 1: + if self.config.training.use_central_scheduler: + iteration_samples = play_self_play_iteration( + self.network, + self.config.mcts, + self.config.training, + self.game_config, + self.rng, + device=self.device, + encoding=self.config.encoding, + temperature=self.config.temperature.training, + max_steps=self.config.evaluation.max_steps, + ) + elif self.config.training.num_workers > 1: iteration_samples = self._run_self_play_parallel(iteration) else: iteration_samples = play_self_play_iteration( @@ -332,6 +348,41 @@ class IsMctsTrainer: return {} self.network.eval() results: dict[str, float | int] = {} + if self.config.training.use_central_scheduler and self.config.mcts.eval_with_mcts: + print( + f" eval vs {', '.join(opponents)} (central scheduler)...", + flush=True, + ) + eval_mcts_cfg = self.config.mcts.model_copy() + if self.config.mcts.eval_n_simulations > 0: + eval_mcts_cfg = eval_mcts_cfg.model_copy( + update={"n_simulations": self.config.mcts.eval_n_simulations} + ) + eval_results = evaluate_opponents_with_mcts_central( + self.network, + self.game_config, + eval_mcts_cfg, + config=self.config, + opponents=tuple(opponents), + 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, + ) + for opponent, result in eval_results.items(): + key = opponent.replace("-", "_") + for metric_key, value in result.items(): + results[f"eval/{key}/{metric_key}"] = value + par = result.get("play_action_rate", 0.0) + sd = result.get("avg_score_diff0", 0.0) + wr = result.get("win_rate0", 0.0) + print( + f" eval vs {opponent} done in {result.get('elapsed_seconds', 0.0):.1f}s " + f"PA={par:.2f} W={wr:.2f} S={sd:.1f}", + flush=True, + ) + return results eval_workers = max( 1, int(self.config.evaluation.num_workers), diff --git a/tests/games/classic/ismcts/test_ismcts.py b/tests/games/classic/ismcts/test_ismcts.py index a509e70..5457c09 100644 --- a/tests/games/classic/ismcts/test_ismcts.py +++ b/tests/games/classic/ismcts/test_ismcts.py @@ -18,6 +18,7 @@ from coolrl_lost_cities.games.classic.bots.heuristic_py import ( 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.evaluate import ( + evaluate_opponents_with_mcts_central, evaluate_opponents_with_mcts_parallel, ) from coolrl_lost_cities.games.classic.ismcts.inference_server import ( @@ -512,6 +513,7 @@ def test_parallel_self_play_server_matches_sample_count(tmp_path) -> None: "gradient_steps_per_iter": 1, "batch_size": 8, "num_workers": 2, + "use_central_scheduler": False, }, "checkpoint": {"save_every": 0}, "evaluation": {"eval_every": 0, "num_workers": 1, "max_steps": 80}, @@ -547,6 +549,59 @@ def test_parallel_self_play_server_matches_sample_count(tmp_path) -> None: assert on_metrics["samples/added"] == off_metrics["samples/added"] +def test_central_eval_matches_parallel_stats() -> None: + config = IsMctsConfig.model_validate( + { + "run": {"seed": 46, "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": 1, "parallel_simulations": 1, "use_rollout_value": False}, + "training": {"num_workers": 2, "interleave_games": 2, "interleave_max_batch": 8}, + "evaluation": {"games": 2, "opponents": ["random"], "num_workers": 2, "max_steps": 80}, + } + ) + game_config = config.rules.to_lost_cities_config(seed=config.run.seed) + state = GameState.new_game(game_config, seed=config.run.seed) + torch.manual_seed(47) + net = AlphaZeroNet(input_dim(state), state.action_size, hidden_size=16, num_layers=1) + parallel = evaluate_opponents_with_mcts_parallel( + net, + game_config, + config.mcts, + config=config, + opponents=("random",), + games=2, + seed=48, + num_workers=2, + max_steps=80, + ) + central = evaluate_opponents_with_mcts_central( + net, + game_config, + config.mcts, + config=config, + opponents=("random",), + games=2, + seed=48, + device="cpu", + max_steps=80, + ) + + assert central["random"]["games"] == parallel["random"]["games"] + assert central["random"]["policy_turns"] == parallel["random"]["policy_turns"] + assert central["random"]["max_step_timeouts"] == parallel["random"]["max_step_timeouts"] + assert np.isclose( + central["random"]["avg_score_diff0"], + parallel["random"]["avg_score_diff0"], + ) + + def test_parallel_eval_server_matches_local_stats() -> None: config = IsMctsConfig.model_validate( {