diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/config.py b/src/coolrl_lost_cities/games/classic/deep_cfr/config.py index fc87bcc..d1b6ac5 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/config.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/config.py @@ -7,8 +7,12 @@ from dataclasses import dataclass class DeepCFRConfig: iterations: int = 1 traversals_per_iteration: int = 2 - rollouts_per_action: int = 1 - max_rollout_steps: int = 512 + max_traversal_depth: int | None = 8 + max_nodes_per_traversal: int | None = 10_000 + regret_matching_epsilon: float = 1.0e-8 + strategy_sample_interval: int = 1 + store_strategy_on_traverser_nodes: bool = True + store_strategy_on_opponent_nodes: bool = True advantage_train_steps: int = 1 strategy_train_steps: int = 1 batch_size: int = 32 diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/memory.py b/src/coolrl_lost_cities/games/classic/deep_cfr/memory.py index fb78ba2..567b692 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/memory.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/memory.py @@ -9,6 +9,7 @@ import numpy as np class TrainingSample: info_state: np.ndarray target: np.ndarray + legal_mask: np.ndarray iteration: int player: int diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py b/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py index 90b29a9..0f6055e 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py @@ -6,12 +6,11 @@ import numpy as np import torch from torch import nn -from coolrl_lost_cities.games.classic.deep_cfr.cfr_math import regret_matching from coolrl_lost_cities.games.classic.deep_cfr.config import DeepCFRConfig -from coolrl_lost_cities.games.classic.deep_cfr.encoding import encode_info_state, input_dim +from coolrl_lost_cities.games.classic.deep_cfr.encoding import input_dim from coolrl_lost_cities.games.classic.deep_cfr.memory import ReservoirMemory, TrainingSample from coolrl_lost_cities.games.classic.deep_cfr.networks import DeepCFRMLP -from coolrl_lost_cities.games.classic.deep_cfr.traversal import root_action_values +from coolrl_lost_cities.games.classic.deep_cfr.traverser import DeepCFRTraverser, TraversalStats from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig @@ -22,6 +21,11 @@ class IterationMetrics: strategy_samples: int advantage_loss: float strategy_loss: float + traversal_nodes: int + traversal_terminals: int + traversal_depth_cutoffs: int + traversal_node_limit_cutoffs: int + traversal_max_depth_reached: int class DeepCFRTrainer: @@ -57,37 +61,32 @@ class DeepCFRTrainer: ) self.advantage_memory = ReservoirMemory() self.strategy_memory = ReservoirMemory() + self.rng = np.random.default_rng(self.config.seed + 101) def run_iteration(self, iteration: int) -> IterationMetrics: + total_stats = TraversalStats() + traverser = DeepCFRTraverser( + self.advantage_networks, + self.advantage_memory, + self.strategy_memory, + device=self.device, + action_size=self.action_size, + epsilon=self.config.regret_matching_epsilon, + strategy_sample_interval=self.config.strategy_sample_interval, + store_strategy_on_traverser_nodes=self.config.store_strategy_on_traverser_nodes, + store_strategy_on_opponent_nodes=self.config.store_strategy_on_opponent_nodes, + max_depth=self.config.max_traversal_depth, + max_nodes=self.config.max_nodes_per_traversal, + rng=self.rng, + ) + for network in self.advantage_networks: + network.eval() for traversal_index in range(self.config.traversals_per_iteration): for player in range(2): seed = self.config.seed + iteration * 10_000 + traversal_index * 10 + player state = GameState.new_game(self.game_config, seed=seed) - info_state = encode_info_state(state, player) - values, legal = root_action_values( - state, - player, - seed=seed, - rollouts_per_action=self.config.rollouts_per_action, - max_steps=self.config.max_rollout_steps, - ) - legal_bool = legal.astype(bool) - advantages = np.zeros_like(values, dtype=np.float32) - if np.any(legal_bool): - baseline = float(np.mean(values[legal_bool])) - advantages[legal_bool] = values[legal_bool] - baseline - policy = regret_matching(advantages, legal) - - self.advantage_memory.add( - TrainingSample( - info_state=info_state, target=advantages, iteration=iteration, player=player - ) - ) - self.strategy_memory.add( - TrainingSample( - info_state=info_state, target=policy, iteration=iteration, player=player - ) - ) + _, stats = traverser.traverse(state, player, iteration) + total_stats.accumulate(stats) advantage_loss = self._train_advantage_networks() strategy_loss = self._train_strategy_network() @@ -97,6 +96,11 @@ class DeepCFRTrainer: strategy_samples=len(self.strategy_memory), advantage_loss=advantage_loss, strategy_loss=strategy_loss, + traversal_nodes=total_stats.nodes, + traversal_terminals=total_stats.terminals, + traversal_depth_cutoffs=total_stats.depth_cutoffs, + traversal_node_limit_cutoffs=total_stats.node_limit_cutoffs, + traversal_max_depth_reached=total_stats.max_depth_reached, ) def train(self) -> list[IterationMetrics]: @@ -109,12 +113,7 @@ class DeepCFRTrainer: if not samples: continue losses.append( - self._train_supervised( - network, - self.advantage_optimizers[player], - samples, - self.config.advantage_train_steps, - ) + self._train_advantage(network, self.advantage_optimizers[player], samples) ) return float(np.mean(losses)) if losses else 0.0 @@ -122,39 +121,70 @@ class DeepCFRTrainer: samples = self.strategy_memory.all() if not samples: return 0.0 - return self._train_supervised( - self.strategy_network, - self.strategy_optimizer, - samples, - self.config.strategy_train_steps, - ) + return self._train_strategy(self.strategy_network, self.strategy_optimizer, samples) - def _train_supervised( + def _batch(self, samples: list[TrainingSample], step: int) -> list[TrainingSample]: + batch_size = min(self.config.batch_size, len(samples)) + offset = (step * batch_size) % len(samples) + batch = samples[offset : offset + batch_size] + if len(batch) < batch_size: + batch = batch + samples[0 : batch_size - len(batch)] + return batch + + def _batch_tensors( + self, + batch: list[TrainingSample], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + x = torch.as_tensor( + np.stack([sample.info_state for sample in batch]), + dtype=torch.float32, + device=self.device, + ) + y = torch.as_tensor( + np.stack([sample.target 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, + ) + return x, y, legal + + def _train_advantage( self, network: nn.Module, optimizer: torch.optim.Optimizer, samples: list[TrainingSample], - steps: int, ) -> float: last_loss = 0.0 - batch_size = min(self.config.batch_size, len(samples)) - for step in range(max(steps, 0)): - offset = (step * batch_size) % len(samples) - batch = samples[offset : offset + batch_size] - if len(batch) < batch_size: - batch = batch + samples[0 : batch_size - len(batch)] - x = torch.as_tensor( - np.stack([sample.info_state for sample in batch]), - dtype=torch.float32, - device=self.device, - ) - y = torch.as_tensor( - np.stack([sample.target for sample in batch]), - dtype=torch.float32, - device=self.device, - ) + network.train() + for step in range(max(self.config.advantage_train_steps, 0)): + x, y, legal = self._batch_tensors(self._batch(samples, step)) + pred = network(x) + diff = (pred - y).masked_fill(~legal, 0.0) + loss = diff.square().sum() / legal.sum().clamp_min(1) + optimizer.zero_grad(set_to_none=True) + loss.backward() + optimizer.step() + last_loss = float(loss.detach().cpu()) + return last_loss + + def _train_strategy( + self, + network: nn.Module, + optimizer: torch.optim.Optimizer, + samples: list[TrainingSample], + ) -> float: + last_loss = 0.0 + network.train() + for step in range(max(self.config.strategy_train_steps, 0)): + x, y, legal = self._batch_tensors(self._batch(samples, step)) + logits = network(x).masked_fill(~legal, torch.finfo(torch.float32).min) + log_probs = nn.functional.log_softmax(logits, dim=-1).masked_fill(~legal, 0.0) + loss = -(y * log_probs).sum(dim=-1).mean() optimizer.zero_grad(set_to_none=True) - loss = nn.functional.mse_loss(network(x), y) loss.backward() optimizer.step() last_loss = float(loss.detach().cpu()) diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/traverser.py b/src/coolrl_lost_cities/games/classic/deep_cfr/traverser.py new file mode 100644 index 0000000..e789272 --- /dev/null +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/traverser.py @@ -0,0 +1,231 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + +import numpy as np +import torch + +from coolrl_lost_cities.games.classic.deep_cfr.cfr_math import regret_matching +from coolrl_lost_cities.games.classic.deep_cfr.encoding import encode_info_state +from coolrl_lost_cities.games.classic.deep_cfr.memory import ReservoirMemory, TrainingSample +from coolrl_lost_cities.games.classic.game import GameState + + +@dataclass +class TraversalStats: + nodes: int = 0 + terminals: int = 0 + depth_cutoffs: int = 0 + node_limit_cutoffs: int = 0 + max_depth_reached: int = 0 + advantage_samples: int = 0 + strategy_samples: int = 0 + sampled_actions: int = 0 + endpoint_depth_sum: int = 0 + endpoint_depth_buckets: dict[str, int] = field(default_factory=dict) + + def accumulate(self, other: TraversalStats) -> None: + self.nodes += other.nodes + self.terminals += other.terminals + self.depth_cutoffs += other.depth_cutoffs + self.node_limit_cutoffs += other.node_limit_cutoffs + self.max_depth_reached = max(self.max_depth_reached, other.max_depth_reached) + self.advantage_samples += other.advantage_samples + self.strategy_samples += other.strategy_samples + self.sampled_actions += other.sampled_actions + self.endpoint_depth_sum += other.endpoint_depth_sum + for key, value in other.endpoint_depth_buckets.items(): + self.endpoint_depth_buckets[key] = self.endpoint_depth_buckets.get(key, 0) + value + + @property + def endpoints(self) -> int: + return self.terminals + self.depth_cutoffs + self.node_limit_cutoffs + + @property + def avg_endpoint_depth(self) -> float: + return self.endpoint_depth_sum / max(1, self.endpoints) + + def to_dict(self) -> dict[str, float | int]: + return { + "traversal_nodes": self.nodes, + "traversal_terminals": self.terminals, + "traversal_depth_cutoffs": self.depth_cutoffs, + "traversal_node_limit_cutoffs": self.node_limit_cutoffs, + "traversal_max_depth_reached": self.max_depth_reached, + "traversal_advantage_samples": self.advantage_samples, + "traversal_strategy_samples": self.strategy_samples, + "traversal_sampled_actions": self.sampled_actions, + "traversal_avg_endpoint_depth": self.avg_endpoint_depth, + **{ + f"traversal_endpoint_depth_bucket_{key}": value + for key, value in self.endpoint_depth_buckets.items() + }, + } + + +class DeepCFRTraverser: + def __init__( + self, + advantage_networks: list[torch.nn.Module], + advantage_memory: ReservoirMemory, + strategy_memory: ReservoirMemory, + *, + device: torch.device, + action_size: int, + epsilon: float = 1.0e-8, + strategy_sample_interval: int = 1, + store_strategy_on_traverser_nodes: bool = True, + store_strategy_on_opponent_nodes: bool = True, + max_depth: int | None = None, + max_nodes: int | None = None, + rng: np.random.Generator | None = None, + ) -> None: + self.advantage_networks = advantage_networks + self.advantage_memory = advantage_memory + self.strategy_memory = strategy_memory + self.device = device + self.action_size = action_size + self.epsilon = float(epsilon) + self.strategy_sample_interval = max(1, int(strategy_sample_interval)) + self.store_strategy_on_traverser_nodes = store_strategy_on_traverser_nodes + self.store_strategy_on_opponent_nodes = store_strategy_on_opponent_nodes + self.max_depth = max_depth + self.max_nodes = max_nodes + self.rng = rng or np.random.default_rng() + + def traverse( + self, state: GameState, traverser: int, iteration: int + ) -> tuple[float, TraversalStats]: + stats = TraversalStats() + value = self._traverse(state, traverser, iteration, depth=0, stats=stats) + return value, stats + + def _traverse( + self, + state: GameState, + traverser: int, + iteration: int, + *, + depth: int, + stats: TraversalStats, + ) -> float: + stats.nodes += 1 + stats.max_depth_reached = max(stats.max_depth_reached, depth) + + if self.max_nodes is not None and stats.nodes >= self.max_nodes: + stats.node_limit_cutoffs += 1 + self._record_endpoint(stats, depth) + return float(state.score_diff(traverser)) + if state.terminal: + stats.terminals += 1 + self._record_endpoint(stats, depth) + return float(state.score_diff(traverser)) + if self.max_depth is not None and depth >= self.max_depth: + stats.depth_cutoffs += 1 + self._record_endpoint(stats, depth) + return float(state.score_diff(traverser)) + + player = state.current_player + info_state, legal, policy = self._policy(state, player) + self._record_strategy(info_state, legal, policy, player, traverser, iteration, depth, stats) + + legal_actions = np.flatnonzero(legal) + if len(legal_actions) == 0: + stats.terminals += 1 + self._record_endpoint(stats, depth) + return float(state.score_diff(traverser)) + + action = self._sample_action(policy, legal_actions) + local_action = state.from_unified_action(int(action)) + state.push_action(local_action) + try: + child_value = self._traverse( + state, + traverser, + iteration, + depth=depth + 1, + stats=stats, + ) + finally: + state.pop_action() + + stats.sampled_actions += 1 + action_prob = max(float(policy[action]), self.epsilon) + sampled_action_value = child_value / action_prob + node_value = float(policy[action]) * sampled_action_value + + if player == traverser: + regrets = np.where(legal, -node_value, 0.0).astype(np.float32) + regrets[action] = np.float32(sampled_action_value - node_value) + self.advantage_memory.add( + TrainingSample( + info_state=info_state, + target=regrets, + legal_mask=legal, + iteration=iteration, + player=player, + ) + ) + stats.advantage_samples += 1 + + return node_value + + def _policy(self, state: GameState, player: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + info_state = encode_info_state(state, player) + legal = np.asarray(state.unified_legal_mask(), dtype=bool) + with torch.inference_mode(): + x = torch.as_tensor(info_state, dtype=torch.float32, device=self.device).unsqueeze(0) + advantages = ( + self.advantage_networks[player](x) + .squeeze(0) + .detach() + .cpu() + .numpy() + .astype(np.float32) + ) + policy = regret_matching(advantages, legal, self.epsilon).astype(np.float32) + return info_state, legal, policy + + def _sample_action(self, policy: np.ndarray, legal_actions: np.ndarray) -> int: + probs = policy[legal_actions].astype(np.float64) + total = float(probs.sum()) + if total <= 0.0: + probs = np.full(len(legal_actions), 1.0 / len(legal_actions), dtype=np.float64) + else: + probs /= total + return int(self.rng.choice(legal_actions, p=probs)) + + def _record_strategy( + self, + info_state: np.ndarray, + legal: np.ndarray, + policy: np.ndarray, + player: int, + traverser: int, + iteration: int, + depth: int, + stats: TraversalStats, + ) -> None: + if player == traverser: + if not self.store_strategy_on_traverser_nodes: + return + elif not self.store_strategy_on_opponent_nodes: + return + if depth % self.strategy_sample_interval != 0: + return + self.strategy_memory.add( + TrainingSample( + info_state=info_state, + target=policy, + legal_mask=legal, + iteration=iteration, + player=player, + ) + ) + stats.strategy_samples += 1 + + def _record_endpoint(self, stats: TraversalStats, depth: int) -> None: + stats.endpoint_depth_sum += depth + start = min(depth // 10 * 10, 100) + key = "100_plus" if start >= 100 else f"{start}_{start + 9}" + stats.endpoint_depth_buckets[key] = stats.endpoint_depth_buckets.get(key, 0) + 1 diff --git a/tests/games/classic/test_deep_cfr_trainer.py b/tests/games/classic/test_deep_cfr_trainer.py index 8185f4e..df79cd4 100644 --- a/tests/games/classic/test_deep_cfr_trainer.py +++ b/tests/games/classic/test_deep_cfr_trainer.py @@ -1,9 +1,11 @@ from __future__ import annotations -from coolrl_lost_cities.games.classic.game import LostCitiesConfig +import numpy as np +from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig from coolrl_lost_cities.games.classic.deep_cfr.config import DeepCFRConfig from coolrl_lost_cities.games.classic.deep_cfr.trainer import DeepCFRTrainer +from coolrl_lost_cities.games.classic.deep_cfr.traverser import DeepCFRTraverser def test_deep_cfr_trainer_smoke_run() -> None: @@ -11,8 +13,8 @@ def test_deep_cfr_trainer_smoke_run() -> None: DeepCFRConfig( iterations=1, traversals_per_iteration=1, - rollouts_per_action=1, - max_rollout_steps=64, + max_traversal_depth=3, + max_nodes_per_traversal=64, advantage_train_steps=1, strategy_train_steps=1, batch_size=2, @@ -25,7 +27,50 @@ def test_deep_cfr_trainer_smoke_run() -> None: metrics = trainer.train() assert len(metrics) == 1 - assert metrics[0].advantage_samples == 2 - assert metrics[0].strategy_samples == 2 + assert metrics[0].advantage_samples > 0 + assert metrics[0].strategy_samples > 0 + assert metrics[0].traversal_nodes > 0 + assert metrics[0].traversal_max_depth_reached <= 3 assert metrics[0].advantage_loss >= 0.0 assert metrics[0].strategy_loss >= 0.0 + + +def test_deep_cfr_recursive_traverser_restores_state_and_collects_samples() -> None: + trainer = DeepCFRTrainer( + DeepCFRConfig( + iterations=1, + traversals_per_iteration=1, + max_traversal_depth=2, + max_nodes_per_traversal=32, + batch_size=2, + hidden_size=16, + seed=29, + ), + LostCitiesConfig(seed=29), + ) + state = GameState.new_game(LostCitiesConfig(seed=29), seed=29) + before = state.to_snapshot() + traverser = DeepCFRTraverser( + trainer.advantage_networks, + trainer.advantage_memory, + trainer.strategy_memory, + device=trainer.device, + action_size=trainer.action_size, + max_depth=2, + max_nodes=32, + rng=np.random.default_rng(29), + ) + + value, stats = traverser.traverse(state, traverser=0, iteration=1) + + assert isinstance(value, float) + assert state.to_snapshot() == before + assert stats.nodes > 0 + assert stats.depth_cutoffs + stats.terminals + stats.node_limit_cutoffs > 0 + assert stats.strategy_samples > 0 + assert stats.advantage_samples > 0 + assert len(trainer.strategy_memory) == stats.strategy_samples + assert len(trainer.advantage_memory) == stats.advantage_samples + sample = trainer.advantage_memory.all()[0] + assert sample.legal_mask.dtype == bool + assert sample.target.shape == sample.legal_mask.shape