coolguy/coorl-lost-cities
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add central ISMCTS scheduler prototype

Browse Source
This commit is contained in:
coolguy
2026-05-11 03:24:54 +09:00
parent 1ee329250e
commit dafb7a6a1d
9 changed files with 320 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
configs/ismcts/default.yaml
+2 -1
View File
@@ -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:
src/coolrl_lost_cities/games/classic/game.pxd
+1
View File
@@ -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)
src/coolrl_lost_cities/games/classic/game.pyx
+23
View File
@@ -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)] = <int>unseen[i]
for i in range(self.deck_len):
other.deck_cards[i] = <int>unseen[opponent_hand_len + i]
return other
cpdef list legal_card_mask(self):
cdef list mask = [False] * (2 * self.hand_size)
cdef int slot
src/coolrl_lost_cities/games/classic/ismcts/config.py
+2 -1
View File
@@ -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
src/coolrl_lost_cities/games/classic/ismcts/determinization.py
+2
View File
@@ -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()
src/coolrl_lost_cities/games/classic/ismcts/evaluate.py
+181
View File
@@ -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],
src/coolrl_lost_cities/games/classic/ismcts/mcts.pyx
+1 -2
View File
@@ -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
src/coolrl_lost_cities/games/classic/ismcts/trainer.py
+53 -2
View File
@@ -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),
tests/games/classic/ismcts/test_ismcts.py
+55
View File
@@ -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(
{