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

Use Deep CFR diagnostics for IS-MCTS eval

Browse Source 
Wrap AlphaZeroNet with a logits-only view so IS-MCTS training evaluation can call evaluate_strategy_network and emit the same full diagnostic metric set as Deep CFR. Adds root prior capture and per-iteration MCTS entropy, value error, and policy-vs-search KL metrics.

Tests: uv run python -m pytest tests/games/classic/ismcts/ -x; uv run python -m pytest tests/games/classic/test_deep_cfr_trainer.py -x; uv run lost-cities-ismcts train --config configs/ismcts/mini.yaml --set run.experiment_name=ismcts-metrics-smoke --set run.max_iterations=2 --set training.games_per_iter=2
This commit is contained in:
coolguy
2026-05-10 23:04:05 +09:00
parent bec59dfc3c
commit 25a3fba53f
5 changed files with 184 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/coolrl_lost_cities/games/classic/ismcts/network.py
+12
View File
@@ -70,3 +70,15 @@ class AlphaZeroNet(nn.Module):
probs = torch.softmax(logits, dim=-1).masked_fill(~legal_mask.bool(), 0.0)
normalizer = probs.sum(dim=-1, keepdim=True).clamp_min(1.0e-12)
return probs / normalizer
class AlphaZeroLogitsView(nn.Module):
"""Expose AlphaZeroNet's policy logits as a one-argument module."""
def __init__(self, net: AlphaZeroNet) -> None:
super().__init__()
self.net = net
def forward(self, x: torch.Tensor) -> torch.Tensor:
logits, _value = self.net(x, legal_mask=None)
return logits
src/coolrl_lost_cities/games/classic/ismcts/replay_buffer.py
+1
View File
@@ -14,6 +14,7 @@ class ReplaySample:
pi_target: np.ndarray
v_target: float
player: int
prior: np.ndarray | None = None
class ReplayBuffer:
src/coolrl_lost_cities/games/classic/ismcts/self_play.py
+17 -5
View File
@@ -60,10 +60,23 @@ def play_self_play_game(
max_steps: int = 10_000,
) -> list[ReplaySample]:
state = GameState.new_game(game_config, seed=rng.randrange(2**31))
pending: list[tuple[np.ndarray, np.ndarray, np.ndarray, int]] = []
pending: list[tuple[np.ndarray, np.ndarray, np.ndarray, int, np.ndarray]] = []
steps = 0
while not state.terminal and steps < max_steps:
player = int(state.current_player)
legal_mask = np.asarray(state.unified_legal_mask(), dtype=bool)
info = encode_info_state(state, player, encoding)
with torch.inference_mode():
x = torch.as_tensor(info[None, :], dtype=torch.float32, device=device)
mask = torch.as_tensor(legal_mask[None, :], dtype=torch.bool, device=device)
prior = (
network.policy_distribution(x, mask)
.squeeze(0)
.detach()
.cpu()
.numpy()
.astype(np.float32)
)
searcher = IsMctsSearcher(
network,
mcts_config,
@@ -72,20 +85,18 @@ def play_self_play_game(
rng=random.Random(rng.randrange(2**31)),
)
visits = searcher.search(state, player)
legal_mask = np.asarray(state.unified_legal_mask(), dtype=bool)
pi = visit_distribution(visits, state.action_size, temperature=temperature)
if pi.sum() <= 0:
legal_actions = np.flatnonzero(legal_mask)
pi[legal_actions] = 1.0 / len(legal_actions)
info = encode_info_state(state, player, encoding)
pending.append((info.astype(np.float32), legal_mask, pi, player))
pending.append((info.astype(np.float32), legal_mask, pi, player, prior))
action = select_from_distribution(pi, rng)
state.apply_unified_action(action)
steps += 1
final_diff0 = float(state.score_diff(0))
samples: list[ReplaySample] = []
for info, legal_mask, pi, player in pending:
for info, legal_mask, pi, player, prior in pending:
value = final_diff0 if player == 0 else -final_diff0
samples.append(
ReplaySample(
@@ -94,6 +105,7 @@ def play_self_play_game(
pi_target=pi.astype(np.float32),
v_target=value,
player=player,
prior=prior,
)
)
return samples
src/coolrl_lost_cities/games/classic/ismcts/trainer.py
+74 -66
View File
@@ -10,13 +10,12 @@ import numpy as np
import torch
from torch import nn
from coolrl_lost_cities.games.classic.bots import build_bot
from coolrl_lost_cities.games.classic.deep_cfr.encoding import input_dim
from coolrl_lost_cities.games.classic.deep_cfr.evaluate import evaluate_strategy_network
from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig
from .config import IsMctsConfig
from .mcts import IsMctsSearcher
from .network import AlphaZeroNet
from .network import AlphaZeroLogitsView, AlphaZeroNet
from .replay_buffer import ReplayBuffer, ReplaySample
from .self_play import play_self_play_game
@@ -32,6 +31,7 @@ class IterationMetrics:
self_play_seconds: float
train_seconds: float
eval_metrics: dict[str, float | int]
mcts_metrics: dict[str, float]
def to_dict(self) -> dict[str, float | int]:
data: dict[str, float | int] = {
@@ -44,6 +44,7 @@ class IterationMetrics:
"time/self_play_seconds": self.self_play_seconds,
"time/train_seconds": self.train_seconds,
}
data.update(self.mcts_metrics)
data.update(self.eval_metrics)
return data
@@ -107,6 +108,7 @@ class IsMctsTrainer:
self.network.eval()
sp_started = time.perf_counter()
added = 0
iteration_samples: list[ReplaySample] = []
for _ in range(self.config.training.games_per_iter):
samples = play_self_play_game(
self.network,
@@ -119,8 +121,10 @@ class IsMctsTrainer:
max_steps=self.config.evaluation.max_steps,
)
self.buffer.add(samples)
iteration_samples.extend(samples)
added += len(samples)
self_play_seconds = time.perf_counter() - sp_started
mcts_metrics = self._compute_mcts_metrics(iteration_samples)
train_started = time.perf_counter()
losses = []
@@ -140,6 +144,7 @@ class IsMctsTrainer:
self_play_seconds=self_play_seconds,
train_seconds=train_seconds,
eval_metrics=eval_metrics,
mcts_metrics=mcts_metrics,
)
def _train_batch(self, batch: list[ReplaySample]) -> tuple[float, float, float]:
@@ -181,27 +186,72 @@ class IsMctsTrainer:
def _evaluate(self, iteration: int) -> dict[str, float | int]:
opponents = self.config.evaluation.opponents_for_iteration(iteration)
if (
not opponents
and self.config.evaluation.eval_every > 0
and self.config.run.max_iterations is not None
and iteration >= self.config.run.max_iterations
):
opponents = self.config.evaluation.opponents
if not opponents:
return {}
self.network.eval()
results: dict[str, float | int] = {}
logits_view = AlphaZeroLogitsView(self.network)
for opponent in opponents:
result = evaluate_policy(
self.network,
result = evaluate_strategy_network(
logits_view,
self.game_config,
opponent=opponent,
games=self.config.evaluation.games,
seed=self.config.run.seed + iteration * 1000,
opponent=opponent,
device=self.device,
encoding=self.config.encoding,
max_steps=self.config.evaluation.max_steps,
mcts_config=self.config.mcts,
batch_size=self.config.evaluation.batch_size,
)
key = opponent.replace("-", "_")
for metric_key, value in result.items():
results[f"eval/{key}/{metric_key}"] = value
return results
def _compute_mcts_metrics(self, samples: list[ReplaySample]) -> dict[str, float]:
if not samples:
return {
"mcts/avg_visit_entropy": 0.0,
"mcts/value_prediction_error": 0.0,
"mcts/policy_mcts_kl": 0.0,
}
entropies = [_entropy(sample.pi_target) for sample in samples]
policy_kls = [
_kl_divergence(sample.pi_target, sample.prior)
for sample in samples
if sample.prior is not None
]
info = torch.as_tensor(
np.stack([sample.info_state for sample in samples]),
dtype=torch.float32,
device=self.device,
)
legal = torch.as_tensor(
np.stack([sample.legal_mask for sample in samples]),
dtype=torch.bool,
device=self.device,
)
target = torch.as_tensor(
[sample.v_target for sample in samples],
dtype=torch.float32,
device=self.device,
)
with torch.inference_mode():
_logits, value_pred = self.network(info, legal)
value_error = nn.functional.mse_loss(value_pred, target)
return {
"mcts/avg_visit_entropy": float(np.mean(entropies)) if entropies else 0.0,
"mcts/value_prediction_error": float(value_error.item()),
"mcts/policy_mcts_kl": float(np.mean(policy_kls)) if policy_kls else 0.0,
}
def _append_metrics(self, metrics: IterationMetrics) -> None:
with self.metrics_path.open("a", encoding="utf-8") as handle:
handle.write(json.dumps(metrics.to_dict(), sort_keys=True) + "\n")
@@ -229,62 +279,20 @@ class IsMctsTrainer:
return (time.perf_counter() - started) / 60.0 >= self.config.run.max_minutes
def evaluate_policy(
network: AlphaZeroNet,
config: LostCitiesConfig,
*,
opponent: str,
games: int,
seed: int,
device: torch.device | str,
encoding=None,
max_steps: int = 10_000,
mcts_config=None,
) -> dict[str, float | int]:
rng = random.Random(seed)
score_diffs: list[int] = []
wins = losses = draws = 0
policy_actions = play_actions = 0
for game_index in range(games):
policy_player = game_index % 2
policies = [
build_bot(opponent, seed=seed + game_index),
build_bot(opponent, seed=seed + game_index),
]
state = GameState.new_game(config, seed=seed + game_index)
for _ in range(max_steps):
if state.terminal:
break
current = int(state.current_player)
if current == policy_player:
searcher = IsMctsSearcher(
network,
mcts_config or IsMctsConfig().mcts,
device=device,
encoding=encoding,
rng=random.Random(rng.randrange(2**31)),
)
visits = searcher.search(state, current)
unified = max(visits, key=visits.get)
action = state.from_unified_action(unified)
else:
action = policies[current].act(state)
if current == policy_player and state.phase == "card":
policy_actions += 1
if action % 2 == 0:
play_actions += 1
state.apply_action(action)
diff = state.score_diff(policy_player)
score_diffs.append(diff)
wins += int(diff > 0)
losses += int(diff < 0)
draws += int(diff == 0)
return {
"games": games,
"wins0": wins,
"wins1": losses,
"draws": draws,
"win_rate0": wins / max(1, games),
"avg_score_diff0": float(np.mean(score_diffs)) if score_diffs else 0.0,
"play_action_rate": play_actions / max(1, policy_actions),
}
def _entropy(distribution: np.ndarray) -> float:
probs = np.asarray(distribution, dtype=np.float64)
probs = probs[probs > 0.0]
if len(probs) == 0:
return 0.0
return float(-(probs * np.log(probs)).sum())
def _kl_divergence(target: np.ndarray, prior: np.ndarray | None) -> float:
if prior is None:
return 0.0
pi = np.asarray(target, dtype=np.float64)
p = np.asarray(prior, dtype=np.float64)
mask = pi > 0.0
if not np.any(mask):
return 0.0
return float((pi[mask] * (np.log(pi[mask]) - np.log(np.clip(p[mask], 1.0e-12, 1.0)))).sum())
tests/games/classic/ismcts/test_ismcts.py
+80 -1
View File
@@ -11,7 +11,7 @@ from coolrl_lost_cities.games.classic.ismcts.config import IsMctsConfig, MctsCon
from coolrl_lost_cities.games.classic.ismcts.determinization import sample_determinization
from coolrl_lost_cities.games.classic.ismcts.info_set import canonical_info_set_key
from coolrl_lost_cities.games.classic.ismcts.mcts import IsMctsSearcher
from coolrl_lost_cities.games.classic.ismcts.network import AlphaZeroNet
from coolrl_lost_cities.games.classic.ismcts.network import AlphaZeroLogitsView, AlphaZeroNet
from coolrl_lost_cities.games.classic.ismcts.replay_buffer import ReplayBuffer, ReplaySample
from coolrl_lost_cities.games.classic.ismcts.self_play import play_self_play_game
from coolrl_lost_cities.games.classic.ismcts.trainer import IsMctsTrainer
@@ -69,6 +69,16 @@ def test_network_shapes_and_mask() -> None:
assert torch.all(probs[~mask] == 0)
def test_logits_view_adapter() -> None:
state = GameState.new_game(mini_config(), seed=6)
dim = input_dim(state)
net = AlphaZeroNet(dim, state.action_size, hidden_size=16, num_layers=1)
logits_view = AlphaZeroLogitsView(net)
x = torch.as_tensor(encode_info_state(state, 0)[None, :], dtype=torch.float32)
logits = logits_view(x)
assert logits.shape == (1, state.action_size)
def test_mcts_prior_drives_visits() -> None:
state = GameState.new_game(mini_config(), seed=7)
dim = input_dim(state)
@@ -114,6 +124,7 @@ def test_self_play_game_returns_signed_targets() -> None:
assert samples
assert {sample.player for sample in samples} <= {0, 1}
assert all(sample.pi_target.sum() > 0 for sample in samples)
assert all(sample.prior is not None for sample in samples)
def test_trainer_one_iteration_smoke(tmp_path) -> None:
@@ -142,3 +153,71 @@ def test_trainer_one_iteration_smoke(tmp_path) -> None:
metrics = trainer.train()
assert len(metrics) == 1
assert (tmp_path / "metrics.jsonl").exists()
def test_trainer_emits_full_eval_metrics(tmp_path) -> None:
config = IsMctsConfig.model_validate(
{
"run": {"max_iterations": 1, "seed": 12, "device": "cpu"},
"rules": {
"n_colors": 3,
"n_ranks": 5,
"n_handshakes": 1,
"hand_size": 4,
"bonus_threshold": 4,
},
"network": {"hidden_size": 16, "num_layers": 1},
"mcts": {"n_simulations": 2},
"training": {"games_per_iter": 1, "gradient_steps_per_iter": 1, "batch_size": 8},
"checkpoint": {"save_every": 0},
"evaluation": {
"eval_every": 1,
"games": 2,
"opponents": ["random"],
"num_workers": 1,
"max_steps": 80,
},
}
)
trainer = IsMctsTrainer(
config,
config.rules.to_lost_cities_config(seed=config.run.seed),
run_dir=tmp_path,
)
metrics = trainer.train()[0].to_dict()
assert "eval/random/avg_opened_colors" in metrics
assert "eval/random/bad_open_rate" in metrics
assert "eval/random/per_game_negative_expeditions" in metrics
def test_trainer_emits_mcts_metrics(tmp_path) -> None:
config = IsMctsConfig.model_validate(
{
"run": {"max_iterations": 1, "seed": 13, "device": "cpu"},
"rules": {
"n_colors": 3,
"n_ranks": 5,
"n_handshakes": 1,
"hand_size": 4,
"bonus_threshold": 4,
},
"network": {"hidden_size": 16, "num_layers": 1},
"mcts": {"n_simulations": 2},
"training": {"games_per_iter": 1, "gradient_steps_per_iter": 1, "batch_size": 8},
"checkpoint": {"save_every": 0},
"evaluation": {"eval_every": 0, "num_workers": 1, "max_steps": 80},
}
)
trainer = IsMctsTrainer(
config,
config.rules.to_lost_cities_config(seed=config.run.seed),
run_dir=tmp_path,
)
metrics = trainer.train()[0].to_dict()
for key in (
"mcts/avg_visit_entropy",
"mcts/value_prediction_error",
"mcts/policy_mcts_kl",
):
assert key in metrics
assert np.isfinite(metrics[key])