Add deterministic Deep CFR traversal mode

configs/deep_cfr/default.yaml

@@ -5,6 +5,7 @@ run:
   max_minutes: null
   device: cuda
   use_amp: false
+  deterministic: false
 
 rules:
   n_colors: 5

configs/deep_cfr/default_server.yaml

@@ -5,6 +5,7 @@ run:
   max_minutes: null
   device: cuda
   use_amp: false
+  deterministic: false
 
 rules:
   n_colors: 5

src/coolrl_lost_cities/games/classic/deep_cfr/config.py

@@ -23,6 +23,7 @@ class RunConfig(StrictModel):
     seed: int = 1
     device: str = "auto"
     use_amp: bool = False
+    deterministic: bool = False
 
     @field_validator("device")
     @classmethod

src/coolrl_lost_cities/games/classic/deep_cfr/evaluate.py

@@ -182,12 +182,14 @@ class StrategyNetPolicy(LostCitiesPolicy):
         sample: bool = False,
         seed: int | None = None,
         encoding: EncodingConfig | None = None,
+        deterministic: bool = False,
     ) -> None:
         self.strategy_network = strategy_network
         self.device = torch.device(device)
         self.sample = sample
         self.rng = np.random.default_rng(seed)
         self.encoding = encoding
+        self.deterministic = deterministic
         self.runtime = EvalRuntimeCounters()
 
     def select_actions_batch(self, states: list[GameState]) -> list[tuple[int, float]]:
@@ -209,6 +211,17 @@ class StrategyNetPolicy(LostCitiesPolicy):
 
         network_started = time.perf_counter()
         with torch.inference_mode():
+            if self.deterministic:
+                logits = torch.cat(
+                    [
+                        self.strategy_network(
+                            torch.as_tensor(info[None, :], dtype=torch.float32, device=self.device)
+                        )
+                        for info in infos
+                    ],
+                    dim=0,
+                )
+            else:
                 x = torch.as_tensor(np.stack(infos), dtype=torch.float32, device=self.device)
                 logits = self.strategy_network(x)
         self.runtime.policy_network_seconds += time.perf_counter() - network_started
@@ -297,6 +310,7 @@ def evaluate_strategy_network(
     encoding: EncodingConfig | None = None,
     batch_size: int = 64,
     save_games_path: str | None = None,
+    deterministic: bool = False,
 ) -> dict[str, float | int]:
     strategy_network.eval()
     return _evaluate_strategy_network_with_diagnostics(
@@ -310,6 +324,7 @@ def evaluate_strategy_network(
         encoding=encoding,
         batch_size=batch_size,
         save_games_path=save_games_path,
+        deterministic=deterministic,
     )
 
 
@@ -325,6 +340,7 @@ def _evaluate_strategy_network_with_diagnostics(
     encoding: EncodingConfig | None,
     batch_size: int,
     save_games_path: str | None = None,
+    deterministic: bool = False,
 ) -> dict[str, float | int]:
     if games <= 0:
         raise ValueError(f"games must be positive, got {games}")
@@ -335,6 +351,7 @@ def _evaluate_strategy_network_with_diagnostics(
         device=device,
         seed=seed * 2,
         encoding=encoding,
+        deterministic=deterministic,
     )
     started = time.perf_counter()
     active_games: list[_EvalGame] = []

src/coolrl_lost_cities/games/classic/deep_cfr/interleaved_traversal.py

@@ -114,6 +114,7 @@ class InterleavedTraversalConfig:
     opponent_policy: str
     endpoint_depth_bucket_width: int
     endpoint_depth_bucket_max: int
+    deterministic: bool = False
 
 
 @dataclass
@@ -182,11 +183,13 @@ class BatchedPolicy:
         device: torch.device,
         epsilon: float,
         strategy_network: torch.nn.Module | None = None,
+        deterministic: bool = False,
     ) -> None:
         self.networks = networks
         self.strategy_network = strategy_network
         self.device = device
         self.epsilon = epsilon
+        self.deterministic = deterministic
         self.batch_sizes: list[int] = []
         self.forward_seconds = 0.0
 
@@ -194,13 +197,27 @@ class BatchedPolicy:
         if not requests:
             return []
         out: list[PolicyResult | None] = [None] * len(requests)
+        if self.deterministic:
+            groups = [
+                (request.network_kind, request.player, [index])
+                for index, request in enumerate(requests)
+            ]
+        else:
             group_keys = sorted({(request.network_kind, request.player) for request in requests})
-        for network_kind, player in group_keys:
+            groups = [
-            indices = [
+                (
+                    network_kind,
+                    player,
+                    [
                         idx
                         for idx, request in enumerate(requests)
                         if (request.network_kind, request.player) == (network_kind, player)
+                    ],
+                )
+                for network_kind, player in group_keys
             ]
+        for network_kind, player, indices in groups:
+            indices = [idx for idx in indices]
             states = np.stack([requests[idx].info_state for idx in indices]).astype(np.float32)
             x = torch.as_tensor(states, dtype=torch.float32, device=self.device)
             if self.device.type == "cuda":
@@ -605,6 +622,8 @@ def run_interleaved_traversal_batch(
     seed: int,
     interleave_width: int,
     interleave_max_batch: int,
+    traversal_start_index: int = 0,
+    deterministic: bool = False,
 ) -> tuple[TraversalStats, list[TrainingSample], list[TrainingSample], dict[str, float | int]]:
     cfg = InterleavedTraversalConfig(
         action_size=action_size,
@@ -620,14 +639,18 @@ def run_interleaved_traversal_batch(
         opponent_policy=opponent_policy,
         endpoint_depth_bucket_width=endpoint_depth_bucket_width,
         endpoint_depth_bucket_max=endpoint_depth_bucket_max,
+        deterministic=deterministic,
     )
     states = [GameState.new_game(game_config, seed=game_seed) for game_seed in seeds]
-    rng_seeds = [int(seed) + index * 1_000_003 for index in range(len(seeds))]
+    rng_seeds = [
+        int(seed) + (int(traversal_start_index) + index) * 1_000_003 for index in range(len(seeds))
+    ]
     policy = BatchedPolicy(
         advantage_networks,
         device=device,
         epsilon=cfg.epsilon,
         strategy_network=strategy_network,
+        deterministic=cfg.deterministic,
     )
     scheduler = InterleavedTraversalScheduler(cfg, policy)
     _values, _rng_out, stats_rows, sample_rows, batch_sizes = scheduler.run(

src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py

@@ -69,9 +69,22 @@ class EvaluationWorkerJob:
     device: str
     max_steps: int
     batch_size: int
+    deterministic: bool = False
+
+
+def configure_torch_determinism(enabled: bool) -> None:
+    torch.use_deterministic_algorithms(bool(enabled))
+    if not enabled:
+        return
+    if torch.backends.cudnn.is_available():
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.allow_tf32 = False
+    if hasattr(torch.backends, "cuda") and hasattr(torch.backends.cuda, "matmul"):
+        torch.backends.cuda.matmul.allow_tf32 = False
 
 
 def run_evaluation_worker(job: EvaluationWorkerJob) -> tuple[str, dict[str, float | int]]:
+    configure_torch_determinism(job.deterministic)
     device = _resolve_torch_device(job.device)
     game_config = LostCitiesConfig(**job.game_config)
     network_config = NetworkConfig.model_validate(job.network_config)
@@ -89,6 +102,7 @@ def run_evaluation_worker(job: EvaluationWorkerJob) -> tuple[str, dict[str, floa
         max_steps=job.max_steps,
         encoding=encoding_config,
         batch_size=job.batch_size,
+        deterministic=job.deterministic,
     )
     return job.opponent, result
 
@@ -201,6 +215,7 @@ class DeepCFRTrainer:
         extra_trackers: list[RunTracker] | None = None,
     ) -> None:
         self.config = config or DeepCFRConfig()
+        configure_torch_determinism(self.config.run.deterministic)
         self.game_config = game_config or self.config.rules.to_lost_cities_config(
             seed=self.config.run.seed
         )
@@ -327,9 +342,11 @@ class DeepCFRTrainer:
             self._ensure_inference_server()
             self._maybe_sync_inference_server(iteration)
         traversal_started = time.perf_counter()
+        worker_count = self.config.traversal.resolved_num_workers()
         if (
-            self.config.traversal.resolved_num_workers() > 1
+            worker_count > 1
             or self.config.traversal.inference_backend == "server"
+            or (self.config.run.deterministic and worker_count > 0)
         ):
             total_stats = self._run_traversals_parallel(iteration)
         else:
@@ -435,6 +452,7 @@ class DeepCFRTrainer:
                         seed=self.config.run.seed + iteration * 1_000_003 + player,
                         interleave_width=self.config.traversal.interleave_width,
                         interleave_max_batch=self.config.traversal.interleave_max_batch,
+                        deterministic=self.config.run.deterministic,
                     )
                 )
                 self._record_interleaved_metrics(runtime_metrics)
@@ -545,6 +563,7 @@ class DeepCFRTrainer:
                 for _, batch in zip(range(in_flight_limit), batch_iter, strict=False)
             }
             completed_batches = 0
+            completed_results = []
             while futures:
                 done, futures = wait(futures, timeout=5.0, return_when=FIRST_COMPLETED)
                 if not done:
@@ -557,17 +576,8 @@ class DeepCFRTrainer:
                     continue
                 for future in done:
                     result = future.result()
+                    completed_results.append(result)
                     completed_batches += 1
-                    total_stats.accumulate(result.stats)
-                    self._record_interleaved_metrics(result.runtime_metrics)
-                    memory_add_started = time.perf_counter()
-                    self._add_advantage_samples(result.advantage_samples)
-                    self.strategy_memory.add_many(result.strategy_samples, self.rng)
-                    self._runtime_metrics["time/memory_add_seconds"] = (
-                        float(self._runtime_metrics.get("time/memory_add_seconds", 0.0))
-                        + time.perf_counter()
-                        - memory_add_started
-                    )
                     progress_nodes += result.stats.nodes
                     progress_traversals += result.traversals
                     if next_progress_at is not None and progress_traversals >= next_progress_at:
@@ -587,6 +597,20 @@ class DeepCFRTrainer:
                     next_batch = next(batch_iter, None)
                     if next_batch is not None:
                         futures.add(executor.submit(run_traversal_worker_batch, next_batch))
+            memory_add_started = time.perf_counter()
+            for result in sorted(
+                completed_results,
+                key=lambda item: (item.player, item.traversal_start_index, item.batch_index),
+            ):
+                total_stats.accumulate(result.stats)
+                self._record_interleaved_metrics(result.runtime_metrics)
+                self._add_advantage_samples(result.advantage_samples)
+                self.strategy_memory.add_many(result.strategy_samples, self.rng)
+            self._runtime_metrics["time/memory_add_seconds"] = (
+                float(self._runtime_metrics.get("time/memory_add_seconds", 0.0))
+                + time.perf_counter()
+                - memory_add_started
+            )
         if (
             self.config.traversal.inference_backend == "server"
             and self._inference_server is not None
@@ -647,8 +671,10 @@ class DeepCFRTrainer:
                 chunk = seeds[start : start + chunk_size]
                 batches.append(
                     TraversalWorkerBatch(
+                        batch_index=batch_index,
                         player=player,
                         iteration=iteration,
+                        traversal_start_index=start,
                         seeds=chunk,
                         config=self.config.to_dict(),
                         game_config=self.game_config.to_snapshot(),
@@ -656,7 +682,7 @@ class DeepCFRTrainer:
                         action_size=self.action_size,
                         advantage_networks=network_payloads,
                         league_advantage_networks=league_payloads,
-                        worker_seed=self.config.run.seed + iteration * 1_000_003 + batch_index,
+                        worker_seed=self.config.run.seed + iteration * 1_000_003 + player,
                         strategy_network=strategy_payload,
                         inference_handles=None,
                     )
@@ -884,6 +910,7 @@ class DeepCFRTrainer:
                 max_steps=self.config.evaluation.max_steps,
                 encoding=self.config.encoding,
                 batch_size=self.config.evaluation.batch_size,
+                deterministic=self.config.run.deterministic,
             )
             for key, value in result.items():
                 results[f"eval/{opponent}/{key}"] = value

src/coolrl_lost_cities/games/classic/deep_cfr/workers.py

@@ -43,6 +43,17 @@ def _configure_worker_torch_threads() -> None:
     _TORCH_THREADS_CONFIGURED = True
 
 
+def _configure_worker_torch_determinism(enabled: bool) -> None:
+    torch.use_deterministic_algorithms(bool(enabled))
+    if not enabled:
+        return
+    if torch.backends.cudnn.is_available():
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.allow_tf32 = False
+    if hasattr(torch.backends, "cuda") and hasattr(torch.backends.cuda, "matmul"):
+        torch.backends.cuda.matmul.allow_tf32 = False
+
+
 def initialize_traversal_worker(inference_handles: InferenceClientHandles | None = None) -> None:
     global _INFERENCE_HANDLES
     _INFERENCE_HANDLES = inference_handles
@@ -51,8 +62,10 @@ def initialize_traversal_worker(inference_handles: InferenceClientHandles | None
 
 @dataclass(frozen=True)
 class TraversalWorkerBatch:
+    batch_index: int
     player: int
     iteration: int
+    traversal_start_index: int
     seeds: list[int]
     config: dict[str, Any]
     game_config: dict[str, Any]
@@ -67,7 +80,9 @@ class TraversalWorkerBatch:
 
 @dataclass(frozen=True)
 class TraversalWorkerResult:
+    batch_index: int
     player: int
+    traversal_start_index: int
     stats: TraversalStats
     advantage_samples: list[TrainingSample]
     strategy_samples: list[TrainingSample]
@@ -79,6 +94,7 @@ def run_traversal_worker_batch(batch: TraversalWorkerBatch) -> TraversalWorkerRe
     _configure_worker_torch_threads()
 
     cfg = config_from_dict(batch.config)
+    _configure_worker_torch_determinism(cfg.run.deterministic)
     device = torch.device("cpu")
     client: InferenceClient | None = None
     try:
@@ -176,6 +192,8 @@ def run_traversal_worker_batch(batch: TraversalWorkerBatch) -> TraversalWorkerRe
                     seed=batch.worker_seed,
                     interleave_width=cfg.traversal.interleave_width,
                     interleave_max_batch=cfg.traversal.interleave_max_batch,
+                    traversal_start_index=batch.traversal_start_index,
+                    deterministic=cfg.run.deterministic,
                 )
             )
         else:
@@ -221,7 +239,9 @@ def run_traversal_worker_batch(batch: TraversalWorkerBatch) -> TraversalWorkerRe
         if client is not None:
             client.close()
     return TraversalWorkerResult(
+        batch_index=batch.batch_index,
         player=batch.player,
+        traversal_start_index=batch.traversal_start_index,
         stats=total_stats,
         advantage_samples=advantage_samples,
         strategy_samples=strategy_samples,