Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
dafb7a6a1d | ||
|
|
1ee329250e |
@@ -36,6 +36,10 @@ training:
|
||||
replay_capacity: 100000
|
||||
interleave_games: 8
|
||||
interleave_max_batch: 64
|
||||
use_central_scheduler: false
|
||||
use_inference_server: false
|
||||
inference_server_max_batch: 128
|
||||
inference_server_batch_timeout_ms: 10.0
|
||||
optimization:
|
||||
learning_rate: 0.0003
|
||||
grad_clip: 5.0
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -60,6 +60,10 @@ class TrainingConfig(StrictModel):
|
||||
interleave_max_batch: int = 64
|
||||
num_workers: int = 1
|
||||
worker_device: str = "cpu"
|
||||
use_central_scheduler: bool = False
|
||||
use_inference_server: bool = False
|
||||
inference_server_max_batch: int = 128
|
||||
inference_server_batch_timeout_ms: float = 10.0
|
||||
|
||||
@field_validator(
|
||||
"games_per_iter",
|
||||
@@ -69,6 +73,7 @@ class TrainingConfig(StrictModel):
|
||||
"interleave_games",
|
||||
"interleave_max_batch",
|
||||
"num_workers",
|
||||
"inference_server_max_batch",
|
||||
)
|
||||
@classmethod
|
||||
def _positive_int(cls, value: int) -> int:
|
||||
@@ -76,6 +81,13 @@ class TrainingConfig(StrictModel):
|
||||
raise ValueError("must be positive")
|
||||
return value
|
||||
|
||||
@field_validator("inference_server_batch_timeout_ms")
|
||||
@classmethod
|
||||
def _positive_float(cls, value: float) -> float:
|
||||
if value <= 0:
|
||||
raise ValueError("must be positive")
|
||||
return value
|
||||
|
||||
|
||||
class IsMctsConfig(StrictModel):
|
||||
run: RunConfig = Field(default_factory=lambda: RunConfig(experiment_name="ismcts"))
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -13,22 +13,37 @@ 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, config_from_dict
|
||||
from .inference_server import InferenceClient
|
||||
from .info_set import canonical_info_set_key
|
||||
from .mcts import IsMctsSearcher
|
||||
from .network import AlphaZeroNet
|
||||
|
||||
_INFERENCE_REQUEST_QUEUE: Any | None = None
|
||||
_INFERENCE_RESPONSE_QUEUES: list[Any] | None = None
|
||||
|
||||
|
||||
def init_eval_inference_queues(request_queue: Any, response_queues: list[Any]) -> None:
|
||||
global _INFERENCE_REQUEST_QUEUE, _INFERENCE_RESPONSE_QUEUES
|
||||
_INFERENCE_REQUEST_QUEUE = request_queue
|
||||
_INFERENCE_RESPONSE_QUEUES = response_queues
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class EvalWorkerBatch:
|
||||
worker_index: int
|
||||
config: dict[str, Any]
|
||||
game_config: dict[str, Any]
|
||||
network_state: dict[str, Any]
|
||||
network_state: dict[str, Any] | None
|
||||
mcts_config: dict[str, Any]
|
||||
opponent: str
|
||||
game_indices: list[int]
|
||||
seed: int
|
||||
device: str
|
||||
max_steps: int
|
||||
tasks: list[tuple[str, int]] | None = None
|
||||
use_inference_server: bool = False
|
||||
request_queue: Any | None = None
|
||||
response_queue: Any | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -41,6 +56,7 @@ class EvalWorkerResult:
|
||||
policy_turns: int
|
||||
play_actions: int
|
||||
timeouts: int
|
||||
by_opponent: dict[str, dict[str, Any]] | None = None
|
||||
|
||||
|
||||
def run_eval_worker(batch: EvalWorkerBatch) -> EvalWorkerResult:
|
||||
@@ -51,12 +67,30 @@ def run_eval_worker(batch: EvalWorkerBatch) -> EvalWorkerResult:
|
||||
torch.set_num_threads(1)
|
||||
cfg: IsMctsConfig = config_from_dict(batch.config)
|
||||
game_config = LostCitiesConfig(**batch.game_config)
|
||||
device = torch.device(batch.device)
|
||||
probe = GameState.new_game(game_config, seed=batch.seed)
|
||||
in_dim = input_dim(probe, cfg.encoding)
|
||||
network = AlphaZeroNet.from_config(in_dim, probe.action_size, cfg).to(device)
|
||||
network.load_state_dict(batch.network_state)
|
||||
network.eval()
|
||||
if batch.use_inference_server:
|
||||
device = torch.device("cpu")
|
||||
network = _NetworkShape(probe.action_size)
|
||||
request_queue = batch.request_queue or _INFERENCE_REQUEST_QUEUE
|
||||
response_queue = batch.response_queue
|
||||
if response_queue is None and _INFERENCE_RESPONSE_QUEUES is not None:
|
||||
response_queue = _INFERENCE_RESPONSE_QUEUES[batch.worker_index]
|
||||
if request_queue is None or response_queue is None:
|
||||
raise RuntimeError("inference server queues are required")
|
||||
inference_client = InferenceClient(
|
||||
batch.worker_index,
|
||||
request_queue,
|
||||
response_queue,
|
||||
)
|
||||
else:
|
||||
device = torch.device(batch.device)
|
||||
network = AlphaZeroNet.from_config(in_dim, probe.action_size, cfg).to(device)
|
||||
if batch.network_state is None:
|
||||
raise RuntimeError("network_state is required without inference server")
|
||||
network.load_state_dict(batch.network_state)
|
||||
network.eval()
|
||||
inference_client = None
|
||||
from .config import MctsConfig
|
||||
|
||||
mcts_config = MctsConfig.model_validate(batch.mcts_config)
|
||||
@@ -67,11 +101,27 @@ def run_eval_worker(batch: EvalWorkerBatch) -> EvalWorkerResult:
|
||||
policy_turns = 0
|
||||
play_actions = 0
|
||||
timeouts = 0
|
||||
for game_index in batch.game_indices:
|
||||
tasks = batch.tasks or [(batch.opponent, game_index) for game_index in batch.game_indices]
|
||||
by_opponent: dict[str, dict[str, Any]] = {}
|
||||
for opponent_name, game_index in tasks:
|
||||
bucket = by_opponent.setdefault(
|
||||
opponent_name,
|
||||
{
|
||||
"score_diffs": [],
|
||||
"wins0": 0,
|
||||
"wins1": 0,
|
||||
"draws": 0,
|
||||
"policy_turns": 0,
|
||||
"play_actions": 0,
|
||||
"timeouts": 0,
|
||||
},
|
||||
)
|
||||
game_policy_turns = 0
|
||||
game_play_actions = 0
|
||||
policy_player = game_index % 2
|
||||
opponents = [
|
||||
build_bot(batch.opponent, seed=batch.seed + game_index),
|
||||
build_bot(batch.opponent, seed=batch.seed + game_index + 1),
|
||||
build_bot(opponent_name, seed=batch.seed + game_index),
|
||||
build_bot(opponent_name, seed=batch.seed + game_index + 1),
|
||||
]
|
||||
state = GameState.new_game(game_config, seed=batch.seed + game_index)
|
||||
steps = 0
|
||||
@@ -89,15 +139,21 @@ def run_eval_worker(batch: EvalWorkerBatch) -> EvalWorkerResult:
|
||||
encoding=cfg.encoding,
|
||||
rng=random.Random(rng.randrange(2**31)),
|
||||
)
|
||||
visits = searcher.search(state, current)
|
||||
visits = (
|
||||
_search_with_inference_server(searcher, state, current, inference_client)
|
||||
if inference_client is not None
|
||||
else searcher.search(state, current)
|
||||
)
|
||||
if visits:
|
||||
unified = max(visits, key=visits.get)
|
||||
else:
|
||||
unified = state.unified_legal_actions()[0]
|
||||
if state.phase == "card":
|
||||
policy_turns += 1
|
||||
game_policy_turns += 1
|
||||
if unified % 2 == 0:
|
||||
play_actions += 1
|
||||
game_play_actions += 1
|
||||
state.apply_unified_action(unified)
|
||||
else:
|
||||
action = opponents[current].act(state)
|
||||
@@ -105,14 +161,21 @@ def run_eval_worker(batch: EvalWorkerBatch) -> EvalWorkerResult:
|
||||
steps += 1
|
||||
if not terminated:
|
||||
timeouts += 1
|
||||
bucket["timeouts"] += 1
|
||||
diff = float(state.score_diff(policy_player))
|
||||
score_diffs.append(diff)
|
||||
bucket["score_diffs"].append(diff)
|
||||
if diff > 0:
|
||||
wins0 += 1
|
||||
bucket["wins0"] += 1
|
||||
elif diff < 0:
|
||||
wins1 += 1
|
||||
bucket["wins1"] += 1
|
||||
else:
|
||||
draws += 1
|
||||
bucket["draws"] += 1
|
||||
bucket["policy_turns"] += game_policy_turns
|
||||
bucket["play_actions"] += game_play_actions
|
||||
return EvalWorkerResult(
|
||||
worker_index=batch.worker_index,
|
||||
score_diffs=score_diffs,
|
||||
@@ -122,4 +185,47 @@ def run_eval_worker(batch: EvalWorkerBatch) -> EvalWorkerResult:
|
||||
policy_turns=policy_turns,
|
||||
play_actions=play_actions,
|
||||
timeouts=timeouts,
|
||||
by_opponent=by_opponent,
|
||||
)
|
||||
|
||||
|
||||
def _search_with_inference_server(
|
||||
searcher: IsMctsSearcher,
|
||||
state: GameState,
|
||||
traverser: int,
|
||||
inference_client: InferenceClient,
|
||||
) -> dict[int, int]:
|
||||
from .interleaved_self_play import _evaluate_global_batch
|
||||
|
||||
root_key = canonical_info_set_key(state, state.current_player)
|
||||
root = searcher.tree.get_or_create(
|
||||
root_key,
|
||||
player=state.current_player,
|
||||
terminal=state.terminal,
|
||||
)
|
||||
completed = 0
|
||||
sims = int(searcher.config.n_simulations)
|
||||
while completed < sims:
|
||||
quota = min(int(searcher.config.parallel_simulations), sims - completed)
|
||||
pending = searcher.prepare_simulation_batch(state, traverser, quota)
|
||||
if not pending:
|
||||
break
|
||||
jobs = [(_SearchProxy(searcher), item) for item in pending]
|
||||
_evaluate_global_batch(
|
||||
searcher.network,
|
||||
jobs,
|
||||
searcher.device,
|
||||
inference_client=inference_client,
|
||||
)
|
||||
completed += len(pending)
|
||||
return {action: root.visits.get(action, 0) for action in state.unified_legal_actions()}
|
||||
|
||||
|
||||
@dataclass
|
||||
class _SearchProxy:
|
||||
searcher: IsMctsSearcher
|
||||
|
||||
|
||||
class _NetworkShape:
|
||||
def __init__(self, action_size: int) -> None:
|
||||
self.action_size = int(action_size)
|
||||
|
||||
@@ -11,7 +11,9 @@ from coolrl_lost_cities.games.classic.bots.registry import build_bot
|
||||
from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig
|
||||
|
||||
from .config import IsMctsConfig, MctsConfig
|
||||
from .eval_worker import EvalWorkerBatch, run_eval_worker
|
||||
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
|
||||
|
||||
@@ -174,6 +176,312 @@ def _evaluate_parallel(
|
||||
play_actions += res.play_actions
|
||||
timeouts += res.timeouts
|
||||
n = len(score_diffs)
|
||||
return _evaluation_metrics(
|
||||
score_diffs=score_diffs,
|
||||
wins0=wins0,
|
||||
wins1=wins1,
|
||||
draws=draws,
|
||||
policy_turns=policy_turns,
|
||||
play_actions=play_actions,
|
||||
timeouts=timeouts,
|
||||
elapsed_seconds=time.perf_counter() - started,
|
||||
n=n,
|
||||
)
|
||||
|
||||
|
||||
def evaluate_opponents_with_mcts_parallel(
|
||||
network: AlphaZeroNet,
|
||||
game_config: LostCitiesConfig,
|
||||
mcts_config: MctsConfig,
|
||||
*,
|
||||
config: IsMctsConfig,
|
||||
opponents: tuple[str, ...],
|
||||
games: int,
|
||||
seed: int,
|
||||
num_workers: int,
|
||||
max_steps: int,
|
||||
request_queue=None,
|
||||
response_queues=None,
|
||||
) -> dict[str, dict[str, float | int]]:
|
||||
started = time.perf_counter()
|
||||
tasks = [(opponent, game_index) for opponent in opponents for game_index in range(games)]
|
||||
if not tasks:
|
||||
return {}
|
||||
effective_workers = min(max(1, int(num_workers)), len(tasks))
|
||||
tasks_per_worker = [tasks[i::effective_workers] for i in range(effective_workers)]
|
||||
use_inference_server = request_queue is not None and response_queues is not None
|
||||
cpu_state = (
|
||||
None
|
||||
if use_inference_server
|
||||
else {name: tensor.detach().cpu() for name, tensor in network.state_dict().items()}
|
||||
)
|
||||
config_dict = config.to_dict()
|
||||
game_snapshot = game_config.to_snapshot()
|
||||
mcts_dict = mcts_config.model_dump(mode="json")
|
||||
worker_device = str(config.training.worker_device)
|
||||
batches = [
|
||||
EvalWorkerBatch(
|
||||
worker_index=i,
|
||||
config=config_dict,
|
||||
game_config=game_snapshot,
|
||||
network_state=cpu_state,
|
||||
mcts_config=mcts_dict,
|
||||
opponent=tasks_per_worker[i][0][0] if tasks_per_worker[i] else "",
|
||||
game_indices=[],
|
||||
seed=seed,
|
||||
device=worker_device,
|
||||
max_steps=max_steps,
|
||||
tasks=tasks_per_worker[i],
|
||||
use_inference_server=use_inference_server,
|
||||
request_queue=None,
|
||||
response_queue=None,
|
||||
)
|
||||
for i in range(effective_workers)
|
||||
]
|
||||
ctx = mp.get_context("spawn")
|
||||
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
|
||||
}
|
||||
executor_kwargs = (
|
||||
{
|
||||
"initializer": init_eval_inference_queues,
|
||||
"initargs": (request_queue, response_queues),
|
||||
}
|
||||
if use_inference_server
|
||||
else {}
|
||||
)
|
||||
with ProcessPoolExecutor(
|
||||
max_workers=effective_workers,
|
||||
mp_context=ctx,
|
||||
**executor_kwargs,
|
||||
) as executor:
|
||||
for result in executor.map(run_eval_worker, batches):
|
||||
for opponent, bucket in (result.by_opponent or {}).items():
|
||||
dest = aggregate[opponent]
|
||||
dest["score_diffs"].extend(bucket["score_diffs"])
|
||||
dest["wins0"] += int(bucket["wins0"])
|
||||
dest["wins1"] += int(bucket["wins1"])
|
||||
dest["draws"] += int(bucket["draws"])
|
||||
dest["policy_turns"] += int(bucket["policy_turns"])
|
||||
dest["play_actions"] += int(bucket["play_actions"])
|
||||
dest["timeouts"] += int(bucket["timeouts"])
|
||||
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()
|
||||
}
|
||||
|
||||
|
||||
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],
|
||||
wins0: int,
|
||||
wins1: int,
|
||||
draws: int,
|
||||
policy_turns: int,
|
||||
play_actions: int,
|
||||
timeouts: int,
|
||||
elapsed_seconds: float,
|
||||
n: int,
|
||||
) -> dict[str, float | int]:
|
||||
avg_diff = sum(score_diffs) / n if n else 0.0
|
||||
return {
|
||||
"games": n,
|
||||
@@ -186,5 +494,5 @@ def _evaluate_parallel(
|
||||
"policy_turns": policy_turns,
|
||||
"play_action_rate": play_actions / policy_turns if policy_turns else 0.0,
|
||||
"max_step_timeouts": timeouts,
|
||||
"elapsed_seconds": time.perf_counter() - started,
|
||||
"elapsed_seconds": elapsed_seconds,
|
||||
}
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .network import AlphaZeroNet
|
||||
|
||||
InferenceRequest = tuple[int, int, np.ndarray, np.ndarray] | None
|
||||
InferenceResponse = tuple[int, np.ndarray, np.ndarray]
|
||||
|
||||
|
||||
class InferenceClient:
|
||||
def __init__(self, worker_id: int, request_queue: Any, response_queue: Any) -> None:
|
||||
self.worker_id = int(worker_id)
|
||||
self.request_queue = request_queue
|
||||
self.response_queue = response_queue
|
||||
self._ids = itertools.count()
|
||||
|
||||
def infer(self, infos: np.ndarray, masks: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
||||
request_id = next(self._ids)
|
||||
self.request_queue.put(
|
||||
(
|
||||
self.worker_id,
|
||||
request_id,
|
||||
np.asarray(infos, dtype=np.float32),
|
||||
np.asarray(masks, dtype=bool),
|
||||
)
|
||||
)
|
||||
while True:
|
||||
response_id, priors, values = self.response_queue.get()
|
||||
if response_id == request_id:
|
||||
return priors, values
|
||||
raise RuntimeError(
|
||||
f"inference response id mismatch: expected {request_id}, got {response_id}"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class InferenceServer:
|
||||
network: AlphaZeroNet
|
||||
device: torch.device
|
||||
request_queue: Any
|
||||
response_queues: list[Any]
|
||||
max_batch: int = 64
|
||||
batch_timeout_seconds: float = 0.001
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self._thread: threading.Thread | None = None
|
||||
self._stop = threading.Event()
|
||||
self.forward_batches = 0
|
||||
self.forward_requests = 0
|
||||
self.forward_positions = 0
|
||||
|
||||
def start(self) -> None:
|
||||
if self._thread is not None:
|
||||
return
|
||||
self.network.eval()
|
||||
self._thread = threading.Thread(target=self._run, name="ismcts-inference-server")
|
||||
self._thread.start()
|
||||
|
||||
def stop(self) -> None:
|
||||
self._stop.set()
|
||||
self.request_queue.put(None)
|
||||
if self._thread is not None:
|
||||
self._thread.join()
|
||||
self._thread = None
|
||||
|
||||
def __enter__(self) -> InferenceServer:
|
||||
self.start()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb) -> None:
|
||||
self.stop()
|
||||
|
||||
def _run(self) -> None:
|
||||
while not self._stop.is_set():
|
||||
try:
|
||||
first = self.request_queue.get(timeout=0.1)
|
||||
except queue.Empty:
|
||||
continue
|
||||
if first is None:
|
||||
break
|
||||
batch: list[tuple[int, int, np.ndarray, np.ndarray]] = [first]
|
||||
rows = _request_rows(first)
|
||||
deadline = time.perf_counter() + self.batch_timeout_seconds
|
||||
while rows < self.max_batch:
|
||||
remaining = deadline - time.perf_counter()
|
||||
if remaining <= 0:
|
||||
break
|
||||
try:
|
||||
item = self.request_queue.get(timeout=remaining)
|
||||
except queue.Empty:
|
||||
break
|
||||
if item is None:
|
||||
self._stop.set()
|
||||
break
|
||||
batch.append(item)
|
||||
rows += _request_rows(item)
|
||||
self._serve(batch)
|
||||
|
||||
def _serve(self, batch: list[tuple[int, int, np.ndarray, np.ndarray]]) -> None:
|
||||
infos = np.concatenate([_ensure_2d(item[2]) for item in batch], axis=0)
|
||||
masks = np.concatenate([_ensure_2d(item[3]) for item in batch], axis=0)
|
||||
with torch.inference_mode():
|
||||
x = torch.as_tensor(infos, dtype=torch.float32, device=self.device)
|
||||
legal = torch.as_tensor(masks, dtype=torch.bool, device=self.device)
|
||||
logits, values = self.network(x, legal)
|
||||
probs = torch.softmax(logits, dim=-1).masked_fill(~legal, 0.0)
|
||||
normalizer = probs.sum(dim=-1, keepdim=True).clamp_min(1.0e-12)
|
||||
priors = (probs / normalizer).detach().cpu().numpy()
|
||||
values_np = values.detach().cpu().numpy()
|
||||
cursor = 0
|
||||
for worker_id, request_id, request_infos, _request_masks in batch:
|
||||
size = _ensure_2d(request_infos).shape[0]
|
||||
self.response_queues[worker_id].put(
|
||||
(
|
||||
request_id,
|
||||
priors[cursor : cursor + size].astype(np.float32, copy=False),
|
||||
values_np[cursor : cursor + size].astype(np.float32, copy=False),
|
||||
)
|
||||
)
|
||||
cursor += size
|
||||
self.forward_batches += 1
|
||||
self.forward_requests += len(batch)
|
||||
self.forward_positions += int(infos.shape[0])
|
||||
|
||||
|
||||
def _ensure_2d(array: np.ndarray) -> np.ndarray:
|
||||
array = np.asarray(array)
|
||||
if array.ndim == 1:
|
||||
return array[None, :]
|
||||
return array
|
||||
|
||||
|
||||
def _request_rows(item: tuple[int, int, np.ndarray, np.ndarray]) -> int:
|
||||
return int(_ensure_2d(item[2]).shape[0])
|
||||
@@ -10,6 +10,7 @@ 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, TrainingConfig
|
||||
from .inference_server import InferenceClient
|
||||
from .info_set import canonical_info_set_key
|
||||
from .mcts import IsMctsSearcher, PendingSimulation
|
||||
from .network import AlphaZeroNet
|
||||
@@ -55,6 +56,7 @@ def play_self_play_iteration(
|
||||
encoding=None,
|
||||
temperature: float = 1.0,
|
||||
max_steps: int = 10_000,
|
||||
inference_client: InferenceClient | None = None,
|
||||
) -> list[ReplaySample]:
|
||||
device = torch.device(device)
|
||||
completed: list[list[ReplaySample]] = []
|
||||
@@ -102,7 +104,13 @@ def play_self_play_iteration(
|
||||
|
||||
active = still_active
|
||||
if jobs:
|
||||
_run_search_jobs(network, jobs, training_config.interleave_max_batch, device)
|
||||
_run_search_jobs(
|
||||
network,
|
||||
jobs,
|
||||
training_config.interleave_max_batch,
|
||||
device,
|
||||
inference_client=inference_client,
|
||||
)
|
||||
for job in jobs:
|
||||
_finish_decision(job, mcts_config, encoding, temperature)
|
||||
|
||||
@@ -119,6 +127,8 @@ def _run_search_jobs(
|
||||
jobs: list[_SearchJob],
|
||||
max_batch: int,
|
||||
device: torch.device,
|
||||
*,
|
||||
inference_client: InferenceClient | None = None,
|
||||
) -> None:
|
||||
while any(job.remaining > 0 for job in jobs):
|
||||
pending: list[tuple[_SearchJob, PendingSimulation]] = []
|
||||
@@ -141,13 +151,15 @@ def _run_search_jobs(
|
||||
break
|
||||
if not pending:
|
||||
break
|
||||
_evaluate_global_batch(network, pending, device)
|
||||
_evaluate_global_batch(network, pending, device, inference_client=inference_client)
|
||||
|
||||
|
||||
def _evaluate_global_batch(
|
||||
network: AlphaZeroNet,
|
||||
pending: list[tuple[_SearchJob, PendingSimulation]],
|
||||
device: torch.device,
|
||||
*,
|
||||
inference_client: InferenceClient | None = None,
|
||||
) -> None:
|
||||
network_pending = [(job, item) for job, item in pending if item.terminal_value is None]
|
||||
values_by_id: dict[int, float] = {}
|
||||
@@ -159,12 +171,17 @@ def _evaluate_global_batch(
|
||||
masks = np.stack(
|
||||
[item.legal_mask for _job, item in network_pending if item.legal_mask is not None]
|
||||
)
|
||||
with torch.inference_mode():
|
||||
x = torch.as_tensor(infos, dtype=torch.float32, device=device)
|
||||
mask = torch.as_tensor(masks, dtype=torch.bool, device=device)
|
||||
probs = network.policy_distribution(x, mask).detach().cpu().numpy()
|
||||
_logits, values = network(x, mask)
|
||||
values_np = values.detach().cpu().numpy()
|
||||
if inference_client is None:
|
||||
with torch.inference_mode():
|
||||
x = torch.as_tensor(infos, dtype=torch.float32, device=device)
|
||||
mask = torch.as_tensor(masks, dtype=torch.bool, device=device)
|
||||
logits, values = network(x, mask)
|
||||
policy = torch.softmax(logits, dim=-1).masked_fill(~mask, 0.0)
|
||||
normalizer = policy.sum(dim=-1, keepdim=True).clamp_min(1.0e-12)
|
||||
probs = (policy / normalizer).detach().cpu().numpy()
|
||||
values_np = values.detach().cpu().numpy()
|
||||
else:
|
||||
probs, values_np = inference_client.infer(infos, masks)
|
||||
for index, (_job, item) in enumerate(network_pending):
|
||||
priors_by_id[id(item)] = probs[index]
|
||||
values_by_id[id(item)] = float(values_np[index])
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -17,11 +17,16 @@ 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_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
|
||||
from .replay_buffer import ReplayBuffer, ReplaySample
|
||||
from .workers import SelfPlayWorkerBatch, run_self_play_worker
|
||||
from .workers import SelfPlayWorkerBatch, init_inference_queues, run_self_play_worker
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -122,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(
|
||||
@@ -182,10 +199,14 @@ class IsMctsTrainer:
|
||||
base = total_games // effective_workers
|
||||
remainder = total_games % effective_workers
|
||||
per_worker = [base + (1 if i < remainder else 0) for i in range(effective_workers)]
|
||||
# Move network state dict to CPU for cross-process transfer.
|
||||
cpu_state = {
|
||||
name: tensor.detach().cpu() for name, tensor in self.network.state_dict().items()
|
||||
}
|
||||
use_inference_server = bool(training_cfg.use_inference_server and effective_workers > 1)
|
||||
# Move network state dict to CPU for cross-process transfer when workers
|
||||
# run local inference. In server mode, workers never deserialize the model.
|
||||
cpu_state = (
|
||||
None
|
||||
if use_inference_server
|
||||
else {name: tensor.detach().cpu() for name, tensor in self.network.state_dict().items()}
|
||||
)
|
||||
config_dict = self.config.to_dict()
|
||||
game_snapshot = self.game_config.to_snapshot()
|
||||
max_steps = self.config.evaluation.max_steps
|
||||
@@ -205,6 +226,7 @@ class IsMctsTrainer:
|
||||
temperature=temperature,
|
||||
max_steps=max_steps,
|
||||
device=worker_device,
|
||||
use_inference_server=use_inference_server,
|
||||
)
|
||||
)
|
||||
samples: list[ReplaySample] = []
|
||||
@@ -215,19 +237,60 @@ class IsMctsTrainer:
|
||||
flush=True,
|
||||
)
|
||||
spawn_started = time.perf_counter()
|
||||
with ProcessPoolExecutor(max_workers=effective_workers, mp_context=ctx) as executor:
|
||||
futures = [executor.submit(run_self_play_worker, batch) for batch in batches]
|
||||
print(
|
||||
f" workers submitted in {time.perf_counter() - spawn_started:.1f}s, waiting for results...",
|
||||
flush=True,
|
||||
request_queue = ctx.Queue() if use_inference_server else None
|
||||
response_queues = (
|
||||
[ctx.Queue() for _ in range(effective_workers)] if use_inference_server else None
|
||||
)
|
||||
server = (
|
||||
InferenceServer(
|
||||
self.network,
|
||||
self.device,
|
||||
request_queue,
|
||||
response_queues,
|
||||
max_batch=int(training_cfg.inference_server_max_batch),
|
||||
batch_timeout_seconds=float(training_cfg.inference_server_batch_timeout_ms)
|
||||
/ 1000.0,
|
||||
)
|
||||
results = []
|
||||
for future in futures:
|
||||
res = future.result()
|
||||
results.append(res)
|
||||
if use_inference_server
|
||||
else None
|
||||
)
|
||||
try:
|
||||
if server is not None:
|
||||
server.start()
|
||||
executor_kwargs = (
|
||||
{
|
||||
"initializer": init_inference_queues,
|
||||
"initargs": (request_queue, response_queues),
|
||||
}
|
||||
if use_inference_server
|
||||
else {}
|
||||
)
|
||||
with ProcessPoolExecutor(
|
||||
max_workers=effective_workers,
|
||||
mp_context=ctx,
|
||||
**executor_kwargs,
|
||||
) as executor:
|
||||
futures = [executor.submit(run_self_play_worker, batch) for batch in batches]
|
||||
print(
|
||||
f" worker {res.worker_index} done ({len(res.samples)} samples, "
|
||||
f"elapsed {time.perf_counter() - spawn_started:.1f}s)",
|
||||
f" workers submitted in {time.perf_counter() - spawn_started:.1f}s, waiting for results...",
|
||||
flush=True,
|
||||
)
|
||||
results = []
|
||||
for future in futures:
|
||||
res = future.result()
|
||||
results.append(res)
|
||||
print(
|
||||
f" worker {res.worker_index} done ({len(res.samples)} samples, "
|
||||
f"elapsed {time.perf_counter() - spawn_started:.1f}s)",
|
||||
flush=True,
|
||||
)
|
||||
finally:
|
||||
if server is not None:
|
||||
server.stop()
|
||||
print(
|
||||
" inference server "
|
||||
f"batches={server.forward_batches} requests={server.forward_requests} "
|
||||
f"positions={server.forward_positions}",
|
||||
flush=True,
|
||||
)
|
||||
for result in sorted(results, key=lambda item: item.worker_index):
|
||||
@@ -285,6 +348,122 @@ 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),
|
||||
int(self.config.training.num_workers),
|
||||
)
|
||||
use_eval_server = bool(
|
||||
self.config.training.use_inference_server
|
||||
and self.config.mcts.eval_with_mcts
|
||||
and eval_workers > 1
|
||||
)
|
||||
if self.config.mcts.eval_with_mcts and eval_workers > 1:
|
||||
print(
|
||||
f" eval vs {', '.join(opponents)} (workers={eval_workers}, "
|
||||
f"inference_server={use_eval_server})...",
|
||||
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}
|
||||
)
|
||||
ctx = mp.get_context("spawn")
|
||||
request_queue = ctx.Queue() if use_eval_server else None
|
||||
response_queues = (
|
||||
[ctx.Queue() for _ in range(eval_workers)] if use_eval_server else None
|
||||
)
|
||||
server = (
|
||||
InferenceServer(
|
||||
self.network,
|
||||
self.device,
|
||||
request_queue,
|
||||
response_queues,
|
||||
max_batch=int(self.config.training.inference_server_max_batch),
|
||||
batch_timeout_seconds=float(
|
||||
self.config.training.inference_server_batch_timeout_ms
|
||||
)
|
||||
/ 1000.0,
|
||||
)
|
||||
if use_eval_server
|
||||
else None
|
||||
)
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
if server is not None:
|
||||
server.start()
|
||||
eval_results = evaluate_opponents_with_mcts_parallel(
|
||||
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,
|
||||
num_workers=eval_workers,
|
||||
max_steps=self.config.evaluation.max_steps,
|
||||
request_queue=request_queue,
|
||||
response_queues=response_queues,
|
||||
)
|
||||
finally:
|
||||
if server is not None:
|
||||
server.stop()
|
||||
print(
|
||||
" eval inference server "
|
||||
f"batches={server.forward_batches} requests={server.forward_requests} "
|
||||
f"positions={server.forward_positions}",
|
||||
flush=True,
|
||||
)
|
||||
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,
|
||||
)
|
||||
print(f" eval opponents done in {time.perf_counter() - started:.1f}s", flush=True)
|
||||
return results
|
||||
for opponent in opponents:
|
||||
print(f" eval vs {opponent}...", flush=True)
|
||||
opp_started = time.perf_counter()
|
||||
|
||||
@@ -20,11 +20,20 @@ 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, config_from_dict
|
||||
from .inference_server import InferenceClient
|
||||
from .interleaved_self_play import play_self_play_iteration
|
||||
from .network import AlphaZeroNet
|
||||
from .replay_buffer import ReplaySample
|
||||
|
||||
_TORCH_THREADS_CONFIGURED = False
|
||||
_INFERENCE_REQUEST_QUEUE: Any | None = None
|
||||
_INFERENCE_RESPONSE_QUEUES: list[Any] | None = None
|
||||
|
||||
|
||||
def init_inference_queues(request_queue: Any, response_queues: list[Any]) -> None:
|
||||
global _INFERENCE_REQUEST_QUEUE, _INFERENCE_RESPONSE_QUEUES
|
||||
_INFERENCE_REQUEST_QUEUE = request_queue
|
||||
_INFERENCE_RESPONSE_QUEUES = response_queues
|
||||
|
||||
|
||||
def _configure_worker_torch_threads() -> None:
|
||||
@@ -49,10 +58,13 @@ class SelfPlayWorkerBatch:
|
||||
base_seed: int
|
||||
config: dict[str, Any]
|
||||
game_config: dict[str, Any]
|
||||
network_state: dict[str, Any]
|
||||
network_state: dict[str, Any] | None
|
||||
temperature: float
|
||||
max_steps: int
|
||||
device: str
|
||||
use_inference_server: bool = False
|
||||
request_queue: Any | None = None
|
||||
response_queue: Any | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -69,13 +81,31 @@ def run_self_play_worker(batch: SelfPlayWorkerBatch) -> SelfPlayWorkerResult:
|
||||
_configure_worker_torch_threads()
|
||||
cfg: IsMctsConfig = config_from_dict(batch.config)
|
||||
game_config = LostCitiesConfig(**batch.game_config)
|
||||
device = torch.device(batch.device)
|
||||
probe = GameState.new_game(game_config, seed=batch.base_seed)
|
||||
in_dim = input_dim(probe, cfg.encoding)
|
||||
action_size = probe.action_size
|
||||
network = AlphaZeroNet.from_config(in_dim, action_size, cfg).to(device)
|
||||
network.load_state_dict(batch.network_state)
|
||||
network.eval()
|
||||
if batch.use_inference_server:
|
||||
device = torch.device("cpu")
|
||||
network = _NetworkShape(action_size)
|
||||
request_queue = batch.request_queue or _INFERENCE_REQUEST_QUEUE
|
||||
response_queue = batch.response_queue
|
||||
if response_queue is None and _INFERENCE_RESPONSE_QUEUES is not None:
|
||||
response_queue = _INFERENCE_RESPONSE_QUEUES[batch.worker_index]
|
||||
if request_queue is None or response_queue is None:
|
||||
raise RuntimeError("inference server queues are required")
|
||||
inference_client = InferenceClient(
|
||||
batch.worker_index,
|
||||
request_queue,
|
||||
response_queue,
|
||||
)
|
||||
else:
|
||||
device = torch.device(batch.device)
|
||||
network = AlphaZeroNet.from_config(in_dim, action_size, cfg).to(device)
|
||||
if batch.network_state is None:
|
||||
raise RuntimeError("network_state is required without inference server")
|
||||
network.load_state_dict(batch.network_state)
|
||||
network.eval()
|
||||
inference_client = None
|
||||
print(
|
||||
f" [worker {batch.worker_index}] init done in {_time.perf_counter() - _t0:.1f}s, self-play start",
|
||||
flush=True,
|
||||
@@ -96,9 +126,15 @@ def run_self_play_worker(batch: SelfPlayWorkerBatch) -> SelfPlayWorkerResult:
|
||||
encoding=cfg.encoding,
|
||||
temperature=batch.temperature,
|
||||
max_steps=batch.max_steps,
|
||||
inference_client=inference_client,
|
||||
)
|
||||
print(
|
||||
f" [worker {batch.worker_index}] self-play done in {_time.perf_counter() - _sp_t0:.1f}s ({len(samples)} samples)",
|
||||
flush=True,
|
||||
)
|
||||
return SelfPlayWorkerResult(worker_index=batch.worker_index, samples=samples)
|
||||
|
||||
|
||||
class _NetworkShape:
|
||||
def __init__(self, action_size: int) -> None:
|
||||
self.action_size = int(action_size)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import multiprocessing as mp
|
||||
import random
|
||||
import sys
|
||||
from pathlib import Path
|
||||
@@ -16,6 +17,14 @@ 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 (
|
||||
InferenceClient,
|
||||
InferenceServer,
|
||||
)
|
||||
from coolrl_lost_cities.games.classic.ismcts.info_set import canonical_info_set_key
|
||||
from coolrl_lost_cities.games.classic.ismcts.interleaved_self_play import (
|
||||
play_self_play_iteration,
|
||||
@@ -450,3 +459,216 @@ def test_smoke_iter_with_batching(tmp_path) -> None:
|
||||
assert "mcts/avg_visit_entropy" in metrics
|
||||
assert "mcts/value_prediction_error" in metrics
|
||||
assert "mcts/policy_mcts_kl" in metrics
|
||||
|
||||
|
||||
def test_inference_server_roundtrip_shapes() -> None:
|
||||
state = GameState.new_game(mini_config(), seed=41)
|
||||
net = AlphaZeroNet(input_dim(state), state.action_size, hidden_size=8, num_layers=1)
|
||||
ctx = mp.get_context("spawn")
|
||||
manager = ctx.Manager()
|
||||
try:
|
||||
request_queue = manager.Queue()
|
||||
response_queues = [manager.Queue()]
|
||||
server = InferenceServer(
|
||||
net,
|
||||
torch.device("cpu"),
|
||||
request_queue,
|
||||
response_queues,
|
||||
max_batch=4,
|
||||
)
|
||||
server.start()
|
||||
try:
|
||||
client = InferenceClient(0, request_queue, response_queues[0])
|
||||
infos = []
|
||||
masks = []
|
||||
for player in (0, 1, 0):
|
||||
infos.append(encode_info_state(state, player))
|
||||
masks.append(np.asarray(state.unified_legal_mask(), dtype=bool))
|
||||
priors, values = client.infer(np.stack(infos), np.stack(masks))
|
||||
finally:
|
||||
server.stop()
|
||||
finally:
|
||||
manager.shutdown()
|
||||
|
||||
assert priors.shape == (3, state.action_size)
|
||||
assert values.shape == (3,)
|
||||
assert np.allclose(priors.sum(axis=1), 1.0)
|
||||
assert np.all(priors[:, ~np.asarray(state.unified_legal_mask(), dtype=bool)] == 0.0)
|
||||
|
||||
|
||||
def test_parallel_self_play_server_matches_sample_count(tmp_path) -> None:
|
||||
base_config = {
|
||||
"run": {"max_iterations": 1, "seed": 42, "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, "parallel_simulations": 2, "use_rollout_value": False},
|
||||
"training": {
|
||||
"games_per_iter": 2,
|
||||
"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},
|
||||
}
|
||||
off = IsMctsConfig.model_validate(
|
||||
{
|
||||
**base_config,
|
||||
"training": {**base_config["training"], "use_inference_server": False},
|
||||
}
|
||||
)
|
||||
on = IsMctsConfig.model_validate(
|
||||
{
|
||||
**base_config,
|
||||
"training": {**base_config["training"], "use_inference_server": True},
|
||||
}
|
||||
)
|
||||
torch.manual_seed(45)
|
||||
off_trainer = IsMctsTrainer(
|
||||
off,
|
||||
off.rules.to_lost_cities_config(seed=off.run.seed),
|
||||
run_dir=tmp_path / "off",
|
||||
)
|
||||
torch.manual_seed(45)
|
||||
on_trainer = IsMctsTrainer(
|
||||
on,
|
||||
on.rules.to_lost_cities_config(seed=on.run.seed),
|
||||
run_dir=tmp_path / "on",
|
||||
)
|
||||
|
||||
off_metrics = off_trainer.train()[0].to_dict()
|
||||
on_metrics = on_trainer.train()[0].to_dict()
|
||||
|
||||
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(
|
||||
{
|
||||
"run": {"seed": 43, "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, "parallel_simulations": 2, "use_rollout_value": False},
|
||||
"training": {"num_workers": 2},
|
||||
"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)
|
||||
net = AlphaZeroNet(input_dim(state), state.action_size, hidden_size=16, num_layers=1)
|
||||
local = evaluate_opponents_with_mcts_parallel(
|
||||
net,
|
||||
game_config,
|
||||
config.mcts,
|
||||
config=config,
|
||||
opponents=("random",),
|
||||
games=2,
|
||||
seed=44,
|
||||
num_workers=2,
|
||||
max_steps=80,
|
||||
)
|
||||
|
||||
ctx = mp.get_context("spawn")
|
||||
manager = ctx.Manager()
|
||||
try:
|
||||
request_queue = manager.Queue()
|
||||
response_queues = [manager.Queue() for _ in range(2)]
|
||||
server = InferenceServer(
|
||||
net,
|
||||
torch.device("cpu"),
|
||||
request_queue,
|
||||
response_queues,
|
||||
max_batch=8,
|
||||
)
|
||||
server.start()
|
||||
try:
|
||||
server_result = evaluate_opponents_with_mcts_parallel(
|
||||
net,
|
||||
game_config,
|
||||
config.mcts,
|
||||
config=config,
|
||||
opponents=("random",),
|
||||
games=2,
|
||||
seed=44,
|
||||
num_workers=2,
|
||||
max_steps=80,
|
||||
request_queue=request_queue,
|
||||
response_queues=response_queues,
|
||||
)
|
||||
finally:
|
||||
server.stop()
|
||||
finally:
|
||||
manager.shutdown()
|
||||
|
||||
for key in ("games", "wins0", "wins1", "draws", "policy_turns", "max_step_timeouts"):
|
||||
assert server_result["random"][key] == local["random"][key]
|
||||
assert np.isclose(
|
||||
server_result["random"]["avg_score_diff0"],
|
||||
local["random"]["avg_score_diff0"],
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user