coolguy
e69f3165b6
Implements a proof-of-concept single-observer IS-MCTS trainer with AlphaZero-style policy/value network, determinization, replay, self-play, CLI configs, and focused tests. Mini acceptance run reaches positive random eval while keeping play_action_rate above the Deep CFR trap threshold. 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
291 lines
11 KiB
Python
291 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import random
|
|
import time
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
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.game import GameState, LostCitiesConfig
|
|
|
|
from .config import IsMctsConfig
|
|
from .mcts import IsMctsSearcher
|
|
from .network import AlphaZeroNet
|
|
from .replay_buffer import ReplayBuffer, ReplaySample
|
|
from .self_play import play_self_play_game
|
|
|
|
|
|
@dataclass
|
|
class IterationMetrics:
|
|
iteration: int
|
|
samples_added: int
|
|
replay_size: int
|
|
policy_loss: float
|
|
value_loss: float
|
|
total_loss: float
|
|
self_play_seconds: float
|
|
train_seconds: float
|
|
eval_metrics: dict[str, float | int]
|
|
|
|
def to_dict(self) -> dict[str, float | int]:
|
|
data: dict[str, float | int] = {
|
|
"iteration": self.iteration,
|
|
"samples/added": self.samples_added,
|
|
"memory/replay": self.replay_size,
|
|
"loss/policy": self.policy_loss,
|
|
"loss/value": self.value_loss,
|
|
"loss/total": self.total_loss,
|
|
"time/self_play_seconds": self.self_play_seconds,
|
|
"time/train_seconds": self.train_seconds,
|
|
}
|
|
data.update(self.eval_metrics)
|
|
return data
|
|
|
|
|
|
class IsMctsTrainer:
|
|
def __init__(
|
|
self,
|
|
config: IsMctsConfig,
|
|
game_config: LostCitiesConfig,
|
|
*,
|
|
run_dir: str | Path,
|
|
device: torch.device | str = "cpu",
|
|
) -> None:
|
|
self.config = config
|
|
self.game_config = game_config
|
|
self.run_dir = Path(run_dir)
|
|
self.device = self._resolve_device(device)
|
|
probe = GameState.new_game(game_config, seed=config.run.seed)
|
|
self.input_dim = input_dim(probe, config.encoding)
|
|
self.action_size = probe.action_size
|
|
self.network = AlphaZeroNet.from_config(self.input_dim, self.action_size, config).to(
|
|
self.device
|
|
)
|
|
self.optimizer = torch.optim.AdamW(
|
|
self.network.parameters(),
|
|
lr=config.optimization.learning_rate,
|
|
weight_decay=max(float(config.optimization.weight_decay), 1.0e-4),
|
|
)
|
|
self.buffer = ReplayBuffer(config.training.replay_capacity, seed=config.run.seed)
|
|
self.metrics_path = self.run_dir / "metrics.jsonl"
|
|
self.rng = random.Random(config.run.seed)
|
|
|
|
def _resolve_device(self, device: torch.device | str) -> torch.device:
|
|
token = str(device)
|
|
if token == "auto":
|
|
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
return torch.device(token)
|
|
|
|
def train(self) -> list[IterationMetrics]:
|
|
self.run_dir.mkdir(parents=True, exist_ok=True)
|
|
(self.run_dir / "config.json").write_text(
|
|
json.dumps(self.config.to_dict(), indent=2, sort_keys=True),
|
|
encoding="utf-8",
|
|
)
|
|
if self.metrics_path.exists():
|
|
self.metrics_path.unlink()
|
|
metrics: list[IterationMetrics] = []
|
|
max_iterations = self.config.run.max_iterations or 1
|
|
started = time.perf_counter()
|
|
for iteration in range(1, max_iterations + 1):
|
|
if self._time_limit_reached(started):
|
|
break
|
|
item = self.run_iteration(iteration)
|
|
metrics.append(item)
|
|
self._append_metrics(item)
|
|
self._save_checkpoints(iteration, item)
|
|
print(json.dumps(item.to_dict(), sort_keys=True))
|
|
return metrics
|
|
|
|
def run_iteration(self, iteration: int) -> IterationMetrics:
|
|
self.network.eval()
|
|
sp_started = time.perf_counter()
|
|
added = 0
|
|
for _ in range(self.config.training.games_per_iter):
|
|
samples = play_self_play_game(
|
|
self.network,
|
|
self.config.mcts,
|
|
self.game_config,
|
|
self.rng,
|
|
device=self.device,
|
|
encoding=self.config.encoding,
|
|
temperature=self.config.temperature.training,
|
|
max_steps=self.config.evaluation.max_steps,
|
|
)
|
|
self.buffer.add(samples)
|
|
added += len(samples)
|
|
self_play_seconds = time.perf_counter() - sp_started
|
|
|
|
train_started = time.perf_counter()
|
|
losses = []
|
|
for _ in range(self.config.training.gradient_steps_per_iter):
|
|
batch = self.buffer.sample(self.config.training.batch_size)
|
|
losses.append(self._train_batch(batch))
|
|
train_seconds = time.perf_counter() - train_started
|
|
loss_arr = np.asarray(losses, dtype=np.float64)
|
|
eval_metrics = self._evaluate(iteration)
|
|
return IterationMetrics(
|
|
iteration=iteration,
|
|
samples_added=added,
|
|
replay_size=len(self.buffer),
|
|
policy_loss=float(loss_arr[:, 0].mean()) if len(loss_arr) else 0.0,
|
|
value_loss=float(loss_arr[:, 1].mean()) if len(loss_arr) else 0.0,
|
|
total_loss=float(loss_arr[:, 2].mean()) if len(loss_arr) else 0.0,
|
|
self_play_seconds=self_play_seconds,
|
|
train_seconds=train_seconds,
|
|
eval_metrics=eval_metrics,
|
|
)
|
|
|
|
def _train_batch(self, batch: list[ReplaySample]) -> tuple[float, float, float]:
|
|
self.network.train()
|
|
info = torch.as_tensor(
|
|
np.stack([sample.info_state for sample in batch]),
|
|
dtype=torch.float32,
|
|
device=self.device,
|
|
)
|
|
legal = torch.as_tensor(
|
|
np.stack([sample.legal_mask for sample in batch]),
|
|
dtype=torch.bool,
|
|
device=self.device,
|
|
)
|
|
pi = torch.as_tensor(
|
|
np.stack([sample.pi_target for sample in batch]),
|
|
dtype=torch.float32,
|
|
device=self.device,
|
|
)
|
|
value_target = torch.as_tensor(
|
|
[sample.v_target for sample in batch],
|
|
dtype=torch.float32,
|
|
device=self.device,
|
|
)
|
|
logits, value_pred = self.network(info, legal)
|
|
log_probs = torch.log_softmax(logits, dim=-1)
|
|
policy_loss = -(pi * log_probs).sum(dim=-1).mean()
|
|
value_loss = nn.functional.mse_loss(value_pred, value_target)
|
|
loss = policy_loss + value_loss
|
|
self.optimizer.zero_grad(set_to_none=True)
|
|
loss.backward()
|
|
if self.config.optimization.grad_clip > 0:
|
|
torch.nn.utils.clip_grad_norm_(
|
|
self.network.parameters(),
|
|
self.config.optimization.grad_clip,
|
|
)
|
|
self.optimizer.step()
|
|
return float(policy_loss.item()), float(value_loss.item()), float(loss.item())
|
|
|
|
def _evaluate(self, iteration: int) -> dict[str, float | int]:
|
|
opponents = self.config.evaluation.opponents_for_iteration(iteration)
|
|
if not opponents:
|
|
return {}
|
|
self.network.eval()
|
|
results: dict[str, float | int] = {}
|
|
for opponent in opponents:
|
|
result = evaluate_policy(
|
|
self.network,
|
|
self.game_config,
|
|
opponent=opponent,
|
|
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,
|
|
mcts_config=self.config.mcts,
|
|
)
|
|
key = opponent.replace("-", "_")
|
|
for metric_key, value in result.items():
|
|
results[f"eval/{key}/{metric_key}"] = value
|
|
return results
|
|
|
|
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")
|
|
|
|
def _save_checkpoints(self, iteration: int, metrics: IterationMetrics) -> None:
|
|
payload = {
|
|
"config": self.config.to_dict(),
|
|
"game_config": self.game_config.to_snapshot(),
|
|
"iteration": iteration,
|
|
"network": self.network.state_dict(),
|
|
"optimizer": self.optimizer.state_dict(),
|
|
"metrics": metrics.to_dict(),
|
|
}
|
|
if self.config.checkpoint.save_latest:
|
|
torch.save(payload, self.run_dir / "latest.pt")
|
|
if (
|
|
self.config.checkpoint.save_every > 0
|
|
and iteration % self.config.checkpoint.save_every == 0
|
|
):
|
|
torch.save(payload, self.run_dir / f"iteration_{iteration:05d}.pt")
|
|
|
|
def _time_limit_reached(self, started: float) -> bool:
|
|
if self.config.run.max_minutes is None:
|
|
return False
|
|
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),
|
|
}
|