Deep CFR 재귀 traversal 구현

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

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
 

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())

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

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