Deep CFR traversal을 Cython 엔진으로 교체

docs/deep-cfr-v0-gap-vs-coolrl.md

@@ -30,14 +30,12 @@ Implemented:
 
 Important implementation note:
 
-- The active training path still uses the Python recursive implementation in
-  `deep_cfr/traverser.py`.
-- `deep_cfr/traversal.pyx` exists, but it currently provides Cython traversal
-  primitives for random rollouts and root action value smoke runs. It is not yet
-  the full Deep CFR traversal backend used by the trainer.
-- The rules engine (`game.pyx`), encoding (`encoding.pyx`), and regret-matching
-  math (`cfr_math.pyx`) are Cythonized, but the Deep CFR tree-walking hot loop is
-  still Python.
+- The active training path now calls `deep_cfr/traversal.pyx`.
+- The old Python recursive `deep_cfr/traverser.py` path has been removed from
+  mainline code.
+- The rules engine (`game.pyx`), encoding (`encoding.pyx`), regret-matching math
+  (`cfr_math.pyx`), and Deep CFR tree-walking loop now have Cython
+  implementations.
 
 ### Training And Memory
 
@@ -130,19 +128,16 @@ Current state:
 1. `GameState` mutation, legal-action generation, apply/undo, and cached scoring
    are implemented in Cython.
 2. Information-state encoding and regret matching have Cython modules.
-3. Full Deep CFR traversal is still Python recursive.
-4. The trainer calls `DeepCFRTraverser` from `deep_cfr/traverser.py`, not a full
-   Cython Deep CFR traversal backend.
-5. A recursion-limit guard is present so full-depth runs with node budgets can
-   execute, but this is an execution safety guard rather than a performance
-   solution.
+3. Full Deep CFR traversal now runs through `traversal.pyx`.
+4. PyTorch policy inference and reservoir memory sample materialization still
+   cross the Python boundary.
+5. Traversal is still recursive inside Cython. The Python recursion-limit guard
+   is no longer the main execution path, but an explicit iterative scheduler is
+   still a future optimization.
 
 Recommended performance roadmap:
 
-1. Add `traversal.backend: python | cython` to config.
-2. Keep `traverser.py` as the reference/debug Python traversal.
-3. Implement a Cython Deep CFR traversal entrypoint in `traversal.pyx`.
-4. Move the tree-walking hot path to Cython:
+1. Continue moving the traversal hot path away from Python object boundaries:
    - C-level legal action enumeration
    - C-level push/pop undo
    - terminal, depth cutoff, and node-budget cutoff
@@ -152,26 +147,24 @@ Recommended performance roadmap:
    - instantaneous regret calculation
    - strategy sample collection
    - traversal stats collection
-5. Initially allow Cython traversal to call Python/PyTorch policy callbacks and
-   Python reservoir memories.
-6. Then reduce Python boundary costs with batched memory writes.
-7. After that, evaluate batched network inference for policy calls.
-8. Consider an explicit Cython iterative stack only after the Cython recursive
-   backend is correct and benchmarked.
+2. Reduce Python boundary costs with batched memory writes.
+3. Add batched network inference for policy calls.
+4. Replace the recursive Cython DFS with an explicit Cython traversal scheduler.
+5. Run multiple traversal contexts concurrently so policy-needed states can be
+   encoded and evaluated in batches.
 
 Python iterative traversal is not the preferred performance path. It would
 remove Python recursion-limit risk, but it would keep most Python object and
 callback overhead in the hot loop. For performance, the next serious step is a
-Cython Deep CFR traversal backend.
+Cython batched iterative traversal scheduler.
 
 ## Suggested Next Steps
 
-1. Add a Cython Deep CFR traversal backend behind a config switch.
-2. Add Python-vs-Cython traversal parity tests for state restoration and sample
-   shape/count invariants.
-3. Add benchmark output that directly compares Python and Cython traversal
-   throughput.
-4. Add batched memory writes after the Cython backend is correct.
+1. Add batched memory writes from the Cython traversal engine.
+2. Add benchmark output for recursive Cython traversal vs batched iterative
+   traversal once the scheduler exists.
+3. Add batched policy inference.
+4. Add an explicit Cython iterative traversal scheduler.
 5. Add a status/plot command that reads `metrics.jsonl`.
 6. Add worker progress logging and hotspot timing profile.
 7. Add W&B checkpoint artifacts after checkpoint quality is stable.

src/coolrl_lost_cities/games/classic/deep_cfr/encoding.pxd

@@ -2,5 +2,16 @@ from coolrl_lost_cities.games.classic.game cimport GameState
 
 
 cdef int input_dim_c(GameState state) noexcept
+cdef int _input_dim_with_flags_c(
+    GameState state,
+    bint derived_playability,
+    bint slot_aware_playability,
+) noexcept
 cdef int encode_info_state_c(GameState state, int player, float* out) except -1
-
+cdef int _encode_info_state_with_flags_c(
+    GameState state,
+    int player,
+    float* out,
+    bint derived_playability,
+    bint slot_aware_playability,
+) except -1

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

@@ -41,6 +41,13 @@ class ReservoirMemory:
             self._samples[index] = sample
 
     def extend(self, samples: list[TrainingSample], rng: np.random.Generator | None = None) -> None:
+        self.add_many(samples, rng)
+
+    def add_many(
+        self,
+        samples: list[TrainingSample],
+        rng: np.random.Generator | None = None,
+    ) -> None:
         for sample in samples:
             self.add(sample, rng)
 

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

@@ -26,7 +26,8 @@ from coolrl_lost_cities.games.classic.deep_cfr.tracking import (
     FileRunTracker,
     RunTracker,
 )
-from coolrl_lost_cities.games.classic.deep_cfr.traverser import DeepCFRTraverser, TraversalStats
+from coolrl_lost_cities.games.classic.deep_cfr.traversal import run_cython_traversal_batch
+from coolrl_lost_cities.games.classic.deep_cfr.traversal_stats import TraversalStats
 from coolrl_lost_cities.games.classic.deep_cfr.workers import (
     TraversalWorkerBatch,
     run_traversal_worker_batch,
@@ -243,16 +244,35 @@ class DeepCFRTrainer:
 
     def _run_traversals_single_process(self, iteration: int) -> TraversalStats:
         total_stats = TraversalStats()
-        traverser = DeepCFRTraverser(
+        for network in self.advantage_networks:
+            network.eval()
+        league_networks = self._materialize_league_networks()
+        progress_every = int(self.config.traversal.progress_every_traversals)
+        completed = 0
+        progress_started = time.perf_counter()
+        traversals_per_player = self.config.traversal.resolved_traversals_per_player()
+        for player in range(2):
+            seeds = [
+                self.config.run.seed + iteration * 10_000 + traversal_index * 10 + player
+                for traversal_index in range(traversals_per_player)
+            ]
+            stats, advantage_samples, strategy_samples = run_cython_traversal_batch(
                 self.advantage_networks,
-            self.advantage_memory,
-            self.strategy_memory,
+                self.game_config,
+                seeds,
+                player,
+                iteration,
                 device=self.device,
                 action_size=self.action_size,
+                encoding=self.config.encoding,
                 epsilon=self.config.traversal.regret_matching_epsilon,
                 strategy_sample_interval=self.config.traversal.strategy_sample_interval,
-            store_strategy_on_traverser_nodes=self.config.traversal.store_strategy_on_traverser_nodes,
-            store_strategy_on_opponent_nodes=self.config.traversal.store_strategy_on_opponent_nodes,
+                store_strategy_on_traverser_nodes=(
+                    self.config.traversal.store_strategy_on_traverser_nodes
+                ),
+                store_strategy_on_opponent_nodes=(
+                    self.config.traversal.store_strategy_on_opponent_nodes
+                ),
                 max_depth=self.config.traversal.max_depth,
                 max_nodes=self.config.traversal.resolved_max_nodes(),
                 outcome_sampling_epsilon=self.config.traversal.outcome_sampling_epsilon,
@@ -263,7 +283,7 @@ class DeepCFRTrainer:
                 cutoff_rollout_policy=self.config.traversal.cutoff_rollout_policy,
                 cutoff_rollout_max_steps=self.config.traversal.cutoff_rollout_max_steps,
                 opponent_policy=self.config.traversal.opponent_policy,
-            league_advantage_networks=self._materialize_league_networks(),
+                league_advantage_networks=league_networks,
                 self_play_anchor_probability=self.config.self_play.anchor_probability,
                 self_play_current_weight=self.config.self_play.current_weight,
                 self_play_recent_weight=self.config.self_play.recent_weight,
@@ -272,22 +292,13 @@ class DeepCFRTrainer:
                 self_play_recent_window=self.config.self_play.recent_window,
                 endpoint_depth_bucket_width=self.config.traversal.endpoint_depth_bucket_width,
                 endpoint_depth_bucket_max=self.config.traversal.endpoint_depth_bucket_max,
-            encoding=self.config.encoding,
-            rng=self.rng,
+                seed=self.config.run.seed + iteration * 1_000_003 + player,
             )
-        for network in self.advantage_networks:
-            network.eval()
-        progress_every = int(self.config.traversal.progress_every_traversals)
-        completed = 0
-        progress_started = time.perf_counter()
-        for traversal_index in range(self.config.traversal.resolved_traversals_per_player()):
-            for player in range(2):
-                seed = self.config.run.seed + iteration * 10_000 + traversal_index * 10 + player
-                state = GameState.new_game(self.game_config, seed=seed)
-                _, stats = traverser.traverse(state, player, iteration)
             total_stats.accumulate(stats)
-                completed += 1
-                if progress_every > 0 and completed % progress_every == 0:
+            self.advantage_memory.add_many(advantage_samples, self.rng)
+            self.strategy_memory.add_many(strategy_samples, self.rng)
+            completed += len(seeds)
+            if progress_every > 0 and completed >= progress_every:
                 elapsed = time.perf_counter() - progress_started
                 self.tracker.log_event(
                     f"Traversal progress iteration={iteration} completed={completed} "
@@ -328,8 +339,8 @@ class DeepCFRTrainer:
             for completed_batches, future in enumerate(as_completed(futures), start=1):
                 result = future.result()
                 total_stats.accumulate(result.stats)
-                self.advantage_memory.extend(result.advantage_samples, self.rng)
-                self.strategy_memory.extend(result.strategy_samples, self.rng)
+                self.advantage_memory.add_many(result.advantage_samples, self.rng)
+                self.strategy_memory.add_many(result.strategy_samples, self.rng)
                 progress_nodes += result.stats.nodes
                 progress_traversals += result.traversals
                 if next_progress_at is not None and progress_traversals >= next_progress_at:

src/coolrl_lost_cities/games/classic/deep_cfr/traversal.pyx

@@ -1,14 +1,706 @@
 # cython: language_level=3, boundscheck=False, wraparound=False, cdivision=True, initializedcheck=False
-"""Cython traversal primitives for Deep CFR smoke runs."""
+"""Cython traversal engine and rollout primitives for Deep CFR."""
 
+from libc.stdlib cimport free, malloc
+
+import numpy as np
+import torch
+
+from coolrl_lost_cities.games.classic.bots import SafeHeuristicBot
+from coolrl_lost_cities.games.classic.deep_cfr.cfr_math cimport regret_matching_c
+from coolrl_lost_cities.games.classic.deep_cfr.encoding cimport (
+    _encode_info_state_with_flags_c,
+    _input_dim_with_flags_c,
+)
+from coolrl_lost_cities.games.classic.deep_cfr.memory import TrainingSample
+from coolrl_lost_cities.games.classic.deep_cfr.traversal_stats import TraversalStats
 from coolrl_lost_cities.games.classic.game cimport GameState
 
 
+DEF MAX_ACTIONS = 64
+
+
 cdef unsigned int _next_u32(unsigned int* state) noexcept:
     state[0] = state[0] * 1664525 + 1013904223
     return state[0]
 
 
+cdef double _next_double(unsigned int* state) noexcept:
+    return <double>_next_u32(state) / 4294967296.0
+
+
+cdef int _sample_policy_from_actions_c(
+    const float* policy,
+    const int* actions,
+    int count,
+    double random_value,
+) noexcept:
+    cdef int i
+    cdef int fallback = -1
+    cdef double cumulative = 0.0
+    cdef double r = random_value
+    if count <= 0:
+        return -1
+    if r < 0.0:
+        r = 0.0
+    elif r >= 1.0:
+        r = 0.9999999999999999
+    for i in range(count):
+        if policy[actions[i]] > 0.0:
+            fallback = actions[i]
+            cumulative += policy[actions[i]]
+            if r < cumulative:
+                return actions[i]
+    return fallback
+
+
+cdef int _depth_bucket_start(int depth, int width, int max_depth) noexcept:
+    cdef int start = (depth // width) * width
+    if start >= max_depth:
+        return max_depth
+    return start
+
+
+cdef class CythonDeepCFRTraverser:
+    cdef object advantage_networks
+    cdef object advantage_samples
+    cdef object strategy_samples
+    cdef object device
+    cdef object encoding
+    cdef object league_advantage_networks
+    cdef object safe_heuristic_rollout_bot
+    cdef object safe_heuristic_opponent_bot
+    cdef int action_size
+    cdef int input_dim
+    cdef float epsilon
+    cdef int strategy_sample_interval
+    cdef bint store_strategy_on_traverser_nodes
+    cdef bint store_strategy_on_opponent_nodes
+    cdef bint has_max_depth
+    cdef int max_depth
+    cdef bint has_max_nodes
+    cdef int max_nodes
+    cdef float outcome_sampling_epsilon
+    cdef bint has_value_clip
+    cdef float outcome_sampling_value_clip
+    cdef bint unsampled_regret_zero
+    cdef bint cutoff_random_rollout
+    cdef int cutoff_rollouts
+    cdef int cutoff_rollout_max_steps
+    cdef int opponent_policy_id
+    cdef float self_play_anchor_probability
+    cdef float self_play_current_weight
+    cdef float self_play_recent_weight
+    cdef float self_play_older_weight
+    cdef float self_play_anchor_weight
+    cdef int self_play_recent_window
+    cdef int endpoint_depth_bucket_width
+    cdef int endpoint_depth_bucket_max
+    cdef bint derived_playability
+    cdef bint slot_aware_playability
+    cdef unsigned int rng
+
+    def __init__(
+        self,
+        object advantage_networks,
+        *,
+        object device,
+        int action_size,
+        object encoding=None,
+        float epsilon=1.0e-8,
+        int strategy_sample_interval=1,
+        bint store_strategy_on_traverser_nodes=True,
+        bint store_strategy_on_opponent_nodes=True,
+        object max_depth=None,
+        object max_nodes=None,
+        float outcome_sampling_epsilon=0.0,
+        object outcome_sampling_value_clip=None,
+        str outcome_unsampled_regret="negative_node_value",
+        str cutoff_value_mode="score_diff",
+        int cutoff_rollouts=0,
+        str cutoff_rollout_policy="random",
+        int cutoff_rollout_max_steps=10000,
+        str opponent_policy="network",
+        object league_advantage_networks=None,
+        float self_play_anchor_probability=0.0,
+        float self_play_current_weight=0.5,
+        float self_play_recent_weight=0.3,
+        float self_play_older_weight=0.2,
+        float self_play_anchor_weight=0.0,
+        int self_play_recent_window=5,
+        int endpoint_depth_bucket_width=10,
+        int endpoint_depth_bucket_max=100,
+        unsigned int seed=1,
+    ):
+        self.advantage_networks = advantage_networks
+        self.advantage_samples = []
+        self.strategy_samples = []
+        self.device = device
+        self.action_size = action_size
+        if action_size > MAX_ACTIONS:
+            raise ValueError("action_size exceeds fixed traversal action buffer")
+        self.encoding = encoding
+        self.derived_playability = False
+        self.slot_aware_playability = False
+        if encoding is not None:
+            self.derived_playability = bool(encoding.derived_playability)
+            self.slot_aware_playability = bool(encoding.slot_aware_playability)
+        self.input_dim = -1
+        self.epsilon = epsilon
+        self.strategy_sample_interval = max(1, 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.has_max_depth = max_depth is not None
+        self.max_depth = 0 if max_depth is None else int(max_depth)
+        self.has_max_nodes = max_nodes is not None
+        self.max_nodes = 0 if max_nodes is None else int(max_nodes)
+        self.outcome_sampling_epsilon = min(1.0, max(0.0, outcome_sampling_epsilon))
+        self.has_value_clip = outcome_sampling_value_clip is not None
+        self.outcome_sampling_value_clip = (
+            0.0 if outcome_sampling_value_clip is None else max(1.0e-9, float(outcome_sampling_value_clip))
+        )
+        if outcome_unsampled_regret not in {"negative_node_value", "zero"}:
+            raise ValueError("outcome_unsampled_regret must be 'negative_node_value' or 'zero'")
+        self.unsampled_regret_zero = outcome_unsampled_regret == "zero"
+        if cutoff_value_mode not in {"score_diff", "random_rollout"}:
+            raise ValueError("cutoff_value_mode must be 'score_diff' or 'random_rollout'")
+        self.cutoff_random_rollout = cutoff_value_mode == "random_rollout"
+        self.cutoff_rollouts = max(0, cutoff_rollouts)
+        if cutoff_rollout_policy not in {"random", "safe_heuristic"}:
+            raise ValueError("cutoff_rollout_policy must be 'random' or 'safe_heuristic'")
+        self.cutoff_rollout_max_steps = max(1, cutoff_rollout_max_steps)
+        self.safe_heuristic_rollout_bot = (
+            SafeHeuristicBot() if cutoff_rollout_policy == "safe_heuristic" else None
+        )
+        if opponent_policy == "network":
+            self.opponent_policy_id = 0
+        elif opponent_policy == "safe_heuristic":
+            self.opponent_policy_id = 1
+        elif opponent_policy == "self_play_league":
+            self.opponent_policy_id = 2
+        else:
+            raise ValueError("opponent_policy must be 'network', 'safe_heuristic', or 'self_play_league'")
+        self.league_advantage_networks = [] if league_advantage_networks is None else league_advantage_networks
+        self.self_play_anchor_probability = min(1.0, max(0.0, self_play_anchor_probability))
+        self.self_play_current_weight = max(0.0, self_play_current_weight)
+        self.self_play_recent_weight = max(0.0, self_play_recent_weight)
+        self.self_play_older_weight = max(0.0, self_play_older_weight)
+        self.self_play_anchor_weight = max(0.0, self_play_anchor_weight)
+        self.self_play_recent_window = max(0, self_play_recent_window)
+        self.safe_heuristic_opponent_bot = (
+            SafeHeuristicBot()
+            if self.opponent_policy_id == 1 or self.self_play_anchor_probability > 0.0
+            else None
+        )
+        self.endpoint_depth_bucket_width = max(1, endpoint_depth_bucket_width)
+        self.endpoint_depth_bucket_max = max(1, endpoint_depth_bucket_max)
+        self.rng = seed if seed != 0 else 1
+
+    cpdef tuple traverse(self, GameState state, int traverser, int iteration):
+        cdef object stats = TraversalStats()
+        cdef float value
+        if self.input_dim < 0:
+            self.input_dim = _input_dim_with_flags_c(
+                state, self.derived_playability, self.slot_aware_playability
+            )
+        value = self._traverse(state, traverser, iteration, 0, stats)
+        return value, stats
+
+    cdef float _traverse(
+        self,
+        GameState state,
+        int traverser,
+        int iteration,
+        int depth,
+        object stats,
+    ) except *:
+        cdef int player
+        cdef int fixed_action
+        cdef int fixed_unified_action
+        cdef int swapped_deck_index
+        cdef int actions[MAX_ACTIONS]
+        cdef int legal_count
+        cdef int i
+        cdef int action
+        cdef int local_action
+        cdef float child_value
+        cdef float action_prob
+        cdef float sampled_action_value
+        cdef float node_value
+        cdef float policy[MAX_ACTIONS]
+        cdef float sampling_policy[MAX_ACTIONS]
+        cdef unsigned char legal[MAX_ACTIONS]
+        cdef object info_state
+
+        stats.nodes += 1
+        if depth > stats.max_depth_reached:
+            stats.max_depth_reached = depth
+
+        if self.has_max_nodes and stats.nodes >= self.max_nodes:
+            stats.node_limit_cutoffs += 1
+            self._record_endpoint(stats, depth)
+            return self._cutoff_value(state, traverser, stats)
+        if state.terminal:
+            stats.terminals += 1
+            self._record_endpoint(stats, depth)
+            return <float>(state.total_scores[traverser] - state.total_scores[1 - traverser])
+        if self.has_max_depth and depth >= self.max_depth:
+            stats.depth_cutoffs += 1
+            self._record_endpoint(stats, depth)
+            return self._cutoff_value(state, traverser, stats)
+
+        player = state.current_player
+        fixed_action = self._fixed_opponent_action(state, player, traverser)
+        if fixed_action >= 0:
+            fixed_unified_action = self._to_unified_action_c(state, fixed_action)
+            swapped_deck_index = self._sample_deck_draw_chance(state, fixed_unified_action)
+            state._push_action_c(fixed_action)
+            try:
+                return self._traverse(state, traverser, iteration, depth + 1, stats)
+            finally:
+                state._pop_action_c()
+                if swapped_deck_index >= 0:
+                    state._swap_deck_cards_c(swapped_deck_index, state.deck_len - 1)
+
+        info_state = self._policy(state, player, legal, policy)
+        self._record_strategy(info_state, legal, policy, player, traverser, iteration, depth, stats)
+
+        legal_count = 0
+        for i in range(self.action_size):
+            if legal[i] != 0:
+                actions[legal_count] = i
+                legal_count += 1
+        if legal_count <= 0:
+            stats.terminals += 1
+            self._record_endpoint(stats, depth)
+            return <float>(state.total_scores[traverser] - state.total_scores[1 - traverser])
+
+        self._sampling_policy(policy, legal, sampling_policy)
+        action = _sample_policy_from_actions_c(sampling_policy, actions, legal_count, _next_double(&self.rng))
+        local_action = self._from_unified_action_c(state, action)
+        swapped_deck_index = self._sample_deck_draw_chance(state, action)
+        state._push_action_c(local_action)
+        try:
+            child_value = self._traverse(state, traverser, iteration, depth + 1, stats)
+        finally:
+            state._pop_action_c()
+            if swapped_deck_index >= 0:
+                state._swap_deck_cards_c(swapped_deck_index, state.deck_len - 1)
+
+        stats.sampled_actions += 1
+        action_prob = sampling_policy[action]
+        if action_prob < self.epsilon:
+            action_prob = self.epsilon
+        sampled_action_value = child_value / action_prob
+        if self.has_value_clip:
+            if sampled_action_value > self.outcome_sampling_value_clip:
+                sampled_action_value = self.outcome_sampling_value_clip
+            elif sampled_action_value < -self.outcome_sampling_value_clip:
+                sampled_action_value = -self.outcome_sampling_value_clip
+        node_value = policy[action] * sampled_action_value
+
+        if player == traverser:
+            self._record_advantage(
+                info_state,
+                legal,
+                action,
+                sampled_action_value,
+                node_value,
+                iteration,
+                player,
+                stats,
+            )
+        return node_value
+
+    cdef object _policy(
+        self,
+        GameState state,
+        int player,
+        unsigned char* legal,
+        float* policy,
+    ):
+        return self._policy_from_networks(self.advantage_networks, state, player, legal, policy)
+
+    cdef object _policy_from_networks(
+        self,
+        object networks,
+        GameState state,
+        int player,
+        unsigned char* legal,
+        float* policy,
+    ):
+        cdef float[::1] info_view
+        cdef float[::1] adv_view
+        cdef int actions[MAX_ACTIONS]
+        cdef int action_count
+        cdef int i
+
+        info_state = np.empty(self.input_dim, dtype=np.float32)
+        info_view = info_state
+        _encode_info_state_with_flags_c(
+            state,
+            player,
+            &info_view[0],
+            self.derived_playability,
+            self.slot_aware_playability,
+        )
+        for i in range(self.action_size):
+            legal[i] = 0
+        action_count = state._unified_legal_actions_c(actions)
+        for i in range(action_count):
+            legal[actions[i]] = 1
+        with torch.inference_mode():
+            x = torch.as_tensor(info_state, dtype=torch.float32, device=self.device).unsqueeze(0)
+            advantages = networks[player](x).squeeze(0).detach().cpu().numpy().astype(np.float32)
+        adv_view = advantages
+        regret_matching_c(&adv_view[0], legal, self.action_size, self.epsilon, policy)
+        return info_state
+
+    cdef void _sampling_policy(
+        self,
+        const float* policy,
+        const unsigned char* legal,
+        float* out_policy,
+    ) noexcept:
+        cdef int i
+        cdef int legal_count = 0
+        cdef float uniform
+        for i in range(self.action_size):
+            if legal[i] != 0:
+                legal_count += 1
+        if legal_count <= 0:
+            for i in range(self.action_size):
+                out_policy[i] = 0.0
+            return
+        if self.outcome_sampling_epsilon <= 0.0:
+            for i in range(self.action_size):
+                out_policy[i] = policy[i]
+            return
+        uniform = 1.0 / <float>legal_count
+        for i in range(self.action_size):
+            if legal[i] != 0:
+                out_policy[i] = (
+                    (1.0 - self.outcome_sampling_epsilon) * policy[i]
+                    + self.outcome_sampling_epsilon * uniform
+                )
+            else:
+                out_policy[i] = 0.0
+
+    cdef int _fixed_opponent_action(self, GameState state, int player, int traverser) except *:
+        cdef int bucket
+        cdef object networks
+        cdef unsigned char legal[MAX_ACTIONS]
+        cdef float policy[MAX_ACTIONS]
+        cdef int actions[MAX_ACTIONS]
+        cdef int count = 0
+        cdef int i
+        cdef int unified_action
+        if player == traverser or self.opponent_policy_id == 0:
+            return -1
+        if self.opponent_policy_id == 1:
+            if self.safe_heuristic_opponent_bot is None:
+                self.safe_heuristic_opponent_bot = SafeHeuristicBot()
+            return int(self.safe_heuristic_opponent_bot.act(state))
+        bucket = self._self_play_bucket()
+        if bucket == 0:
+            return -1
+        if bucket == 3:
+            if self.safe_heuristic_opponent_bot is None:
+                self.safe_heuristic_opponent_bot = SafeHeuristicBot()
+            return int(self.safe_heuristic_opponent_bot.act(state))
+        networks = self._self_play_snapshot_networks(bucket)
+        if networks is None:
+            return -1
+        self._policy_from_networks(networks, state, player, legal, policy)
+        for i in range(self.action_size):
+            if legal[i] != 0:
+                actions[count] = i
+                count += 1
+        if count <= 0:
+            return -1
+        unified_action = _sample_policy_from_actions_c(policy, actions, count, _next_double(&self.rng))
+        return self._from_unified_action_c(state, unified_action)
+
+    cdef int _self_play_bucket(self) noexcept:
+        cdef int recent_count
+        cdef int older_count
+        cdef double weights[4]
+        cdef double total
+        cdef double pick
+        if self.self_play_anchor_probability > 0.0 and _next_double(&self.rng) < self.self_play_anchor_probability:
+            return 3
+        recent_count = min(len(self.league_advantage_networks), self.self_play_recent_window)
+        older_count = max(0, len(self.league_advantage_networks) - recent_count)
+        weights[0] = self.self_play_current_weight
+        weights[1] = self.self_play_recent_weight if recent_count > 0 else 0.0
+        weights[2] = self.self_play_older_weight if older_count > 0 else 0.0
+        weights[3] = self.self_play_anchor_weight
+        total = weights[0] + weights[1] + weights[2] + weights[3]
+        if total <= 0.0:
+            return 0
+        pick = _next_double(&self.rng) * total
+        if pick < weights[0]:
+            return 0
+        pick -= weights[0]
+        if pick < weights[1]:
+            return 1
+        pick -= weights[1]
+        if pick < weights[2]:
+            return 2
+        return 3
+
+    cdef object _self_play_snapshot_networks(self, int bucket):
+        cdef int recent_count = min(len(self.league_advantage_networks), self.self_play_recent_window)
+        cdef object candidates
+        cdef int index
+        if len(self.league_advantage_networks) == 0:
+            return None
+        if bucket == 1 and recent_count > 0:
+            candidates = self.league_advantage_networks[-recent_count:]
+        elif bucket == 2:
+            candidates = self.league_advantage_networks[:max(0, len(self.league_advantage_networks) - recent_count)]
+        else:
+            candidates = self.league_advantage_networks
+        if len(candidates) == 0:
+            return None
+        index = <int>(_next_u32(&self.rng) % <unsigned int>len(candidates))
+        return candidates[index]
+
+    cdef float _cutoff_value(self, GameState state, int traverser, object stats) except *:
+        cdef int i
+        cdef float total = 0.0
+        if not self.cutoff_random_rollout or self.cutoff_rollouts <= 0:
+            return <float>(state.total_scores[traverser] - state.total_scores[1 - traverser])
+        for i in range(self.cutoff_rollouts):
+            total += self._rollout_value(state, traverser, stats)
+        return total / <float>self.cutoff_rollouts
+
+    cdef float _rollout_value(self, GameState state, int traverser, object stats) except *:
+        cdef int steps = 0
+        cdef int actions[MAX_ACTIONS]
+        cdef int count
+        cdef int unified_action
+        cdef int local_action
+        cdef int swapped_deck_index
+        cdef int* swapped_indices = <int*>malloc(self.cutoff_rollout_max_steps * sizeof(int))
+        cdef float value
+        if swapped_indices == NULL:
+            raise MemoryError()
+        while not state.terminal and steps < self.cutoff_rollout_max_steps:
+            if self.safe_heuristic_rollout_bot is not None:
+                local_action = int(self.safe_heuristic_rollout_bot.act(state))
+                unified_action = self._to_unified_action_c(state, local_action)
+            else:
+                count = state._unified_legal_actions_c(actions)
+                if count <= 0:
+                    break
+                unified_action = actions[_next_u32(&self.rng) % <unsigned int>count]
+                local_action = self._from_unified_action_c(state, unified_action)
+            swapped_deck_index = self._sample_deck_draw_chance(state, unified_action)
+            state._push_action_c(local_action)
+            swapped_indices[steps] = swapped_deck_index
+            steps += 1
+        stats.cutoff_rollouts += 1
+        stats.cutoff_rollout_steps += steps
+        if not state.terminal:
+            stats.cutoff_rollout_timeouts += 1
+        value = <float>(state.total_scores[traverser] - state.total_scores[1 - traverser])
+        while steps > 0:
+            steps -= 1
+            state._pop_action_c()
+            if swapped_indices[steps] >= 0:
+                state._swap_deck_cards_c(swapped_indices[steps], state.deck_len - 1)
+        free(swapped_indices)
+        return value
+
+    cdef int _sample_deck_draw_chance(self, GameState state, int unified_action) except *:
+        cdef int deck_draw_action = 2 * state.hand_size
+        cdef int sampled_index
+        if state.phase_id != 1 or unified_action != deck_draw_action or state.deck_len <= 1:
+            return -1
+        sampled_index = <int>(_next_u32(&self.rng) % <unsigned int>state.deck_len)
+        if sampled_index == state.deck_len - 1:
+            return -1
+        state._swap_deck_cards_c(sampled_index, state.deck_len - 1)
+        return sampled_index
+
+    cdef void _record_strategy(
+        self,
+        object info_state,
+        const unsigned char* legal,
+        const float* policy,
+        int player,
+        int traverser,
+        int iteration,
+        int depth,
+        object stats,
+    ):
+        cdef int i
+        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
+        target = np.empty(self.action_size, dtype=np.float32)
+        legal_mask = np.empty(self.action_size, dtype=np.bool_)
+        cdef float[::1] target_view = target
+        cdef unsigned char[::1] legal_view = legal_mask.view(np.uint8)
+        for i in range(self.action_size):
+            target_view[i] = policy[i]
+            legal_view[i] = legal[i]
+        self.strategy_samples.append(
+            TrainingSample(
+                info_state=info_state,
+                target=target,
+                legal_mask=legal_mask,
+                iteration=iteration,
+                player=player,
+            )
+        )
+        stats.strategy_samples += 1
+
+    cdef void _record_advantage(
+        self,
+        object info_state,
+        const unsigned char* legal,
+        int sampled_action,
+        float sampled_action_value,
+        float node_value,
+        int iteration,
+        int player,
+        object stats,
+    ):
+        cdef int i
+        target = np.empty(self.action_size, dtype=np.float32)
+        legal_mask = np.empty(self.action_size, dtype=np.bool_)
+        cdef float[::1] target_view = target
+        cdef unsigned char[::1] legal_view = legal_mask.view(np.uint8)
+        for i in range(self.action_size):
+            legal_view[i] = legal[i]
+            if legal[i] == 0 or self.unsampled_regret_zero:
+                target_view[i] = 0.0
+            else:
+                target_view[i] = -node_value
+        target_view[sampled_action] = sampled_action_value - node_value
+        self.advantage_samples.append(
+            TrainingSample(
+                info_state=info_state,
+                target=target,
+                legal_mask=legal_mask,
+                iteration=iteration,
+                player=player,
+            )
+        )
+        stats.advantage_samples += 1
+
+    cdef void _record_endpoint(self, object stats, int depth):
+        cdef int width = self.endpoint_depth_bucket_width
+        cdef int max_depth = self.endpoint_depth_bucket_max
+        cdef int start = _depth_bucket_start(depth, width, max_depth)
+        cdef str key
+        stats.endpoint_depth_sum += depth
+        if start >= max_depth:
+            key = f"{max_depth}_plus"
+        else:
+            key = f"{start}_{start + width - 1}"
+        stats.endpoint_depth_buckets[key] = stats.endpoint_depth_buckets.get(key, 0) + 1
+
+    cdef int _from_unified_action_c(self, GameState state, int action_id) noexcept:
+        if state.phase_id == 0:
+            return action_id
+        return action_id - 2 * state.hand_size
+
+    cdef int _to_unified_action_c(self, GameState state, int action_id) noexcept:
+        if state.phase_id == 0:
+            return action_id
+        return 2 * state.hand_size + action_id
+
+    def drain_samples(self):
+        advantage = self.advantage_samples
+        strategy = self.strategy_samples
+        self.advantage_samples = []
+        self.strategy_samples = []
+        return advantage, strategy
+
+
+def run_cython_traversal_batch(
+    object advantage_networks,
+    object game_config,
+    list seeds,
+    int player,
+    int iteration,
+    *,
+    object device,
+    int action_size,
+    object encoding=None,
+    float epsilon=1.0e-8,
+    int strategy_sample_interval=1,
+    bint store_strategy_on_traverser_nodes=True,
+    bint store_strategy_on_opponent_nodes=True,
+    object max_depth=None,
+    object max_nodes=None,
+    float outcome_sampling_epsilon=0.0,
+    object outcome_sampling_value_clip=None,
+    str outcome_unsampled_regret="negative_node_value",
+    str cutoff_value_mode="score_diff",
+    int cutoff_rollouts=0,
+    str cutoff_rollout_policy="random",
+    int cutoff_rollout_max_steps=10000,
+    str opponent_policy="network",
+    object league_advantage_networks=None,
+    float self_play_anchor_probability=0.0,
+    float self_play_current_weight=0.5,
+    float self_play_recent_weight=0.3,
+    float self_play_older_weight=0.2,
+    float self_play_anchor_weight=0.0,
+    int self_play_recent_window=5,
+    int endpoint_depth_bucket_width=10,
+    int endpoint_depth_bucket_max=100,
+    unsigned int seed=1,
+):
+    cdef object stats = TraversalStats()
+    cdef object local_stats
+    cdef object value
+    cdef GameState state
+    cdef int game_seed
+    traverser = CythonDeepCFRTraverser(
+        advantage_networks,
+        device=device,
+        action_size=action_size,
+        encoding=encoding,
+        epsilon=epsilon,
+        strategy_sample_interval=strategy_sample_interval,
+        store_strategy_on_traverser_nodes=store_strategy_on_traverser_nodes,
+        store_strategy_on_opponent_nodes=store_strategy_on_opponent_nodes,
+        max_depth=max_depth,
+        max_nodes=max_nodes,
+        outcome_sampling_epsilon=outcome_sampling_epsilon,
+        outcome_sampling_value_clip=outcome_sampling_value_clip,
+        outcome_unsampled_regret=outcome_unsampled_regret,
+        cutoff_value_mode=cutoff_value_mode,
+        cutoff_rollouts=cutoff_rollouts,
+        cutoff_rollout_policy=cutoff_rollout_policy,
+        cutoff_rollout_max_steps=cutoff_rollout_max_steps,
+        opponent_policy=opponent_policy,
+        league_advantage_networks=league_advantage_networks,
+        self_play_anchor_probability=self_play_anchor_probability,
+        self_play_current_weight=self_play_current_weight,
+        self_play_recent_weight=self_play_recent_weight,
+        self_play_older_weight=self_play_older_weight,
+        self_play_anchor_weight=self_play_anchor_weight,
+        self_play_recent_window=self_play_recent_window,
+        endpoint_depth_bucket_width=endpoint_depth_bucket_width,
+        endpoint_depth_bucket_max=endpoint_depth_bucket_max,
+        seed=seed,
+    )
+    for game_seed in seeds:
+        state = GameState.new_game(game_config, seed=game_seed)
+        value, local_stats = traverser.traverse(state, player, iteration)
+        stats.accumulate(local_stats)
+    advantage_samples, strategy_samples = traverser.drain_samples()
+    return stats, advantage_samples, strategy_samples
 cdef float random_rollout_value_c(
     GameState state,
     int player,

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

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@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
+    cutoff_rollouts: int = 0
+    cutoff_rollout_steps: int = 0
+    cutoff_rollout_timeouts: 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.cutoff_rollouts += other.cutoff_rollouts
+        self.cutoff_rollout_steps += other.cutoff_rollout_steps
+        self.cutoff_rollout_timeouts += other.cutoff_rollout_timeouts
+        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_cutoff_rollouts": self.cutoff_rollouts,
+            "traversal_cutoff_rollout_steps": self.cutoff_rollout_steps,
+            "traversal_cutoff_rollout_timeouts": self.cutoff_rollout_timeouts,
+            "traversal_endpoint_depth_sum": self.endpoint_depth_sum,
+            "traversal_endpoints": self.endpoints,
+            "traversal_avg_endpoint_depth": self.avg_endpoint_depth,
+            **{
+                f"traversal_endpoint_depth_bucket_{key}": value
+                for key, value in self.endpoint_depth_buckets.items()
+            },
+        }

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

@@ -1,474 +0,0 @@
-from __future__ import annotations
-
-import sys
-from dataclasses import dataclass, field
-
-import numpy as np
-import torch
-
-from coolrl_lost_cities.games.classic.bots import SafeHeuristicBot
-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
-    cutoff_rollouts: int = 0
-    cutoff_rollout_steps: int = 0
-    cutoff_rollout_timeouts: 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.cutoff_rollouts += other.cutoff_rollouts
-        self.cutoff_rollout_steps += other.cutoff_rollout_steps
-        self.cutoff_rollout_timeouts += other.cutoff_rollout_timeouts
-        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_cutoff_rollouts": self.cutoff_rollouts,
-            "traversal_cutoff_rollout_steps": self.cutoff_rollout_steps,
-            "traversal_cutoff_rollout_timeouts": self.cutoff_rollout_timeouts,
-            "traversal_endpoint_depth_sum": self.endpoint_depth_sum,
-            "traversal_endpoints": self.endpoints,
-            "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,
-        outcome_sampling_epsilon: float = 0.0,
-        outcome_sampling_value_clip: float | None = None,
-        outcome_unsampled_regret: str = "negative_node_value",
-        cutoff_value_mode: str = "score_diff",
-        cutoff_rollouts: int = 0,
-        cutoff_rollout_policy: str = "random",
-        cutoff_rollout_max_steps: int = 10_000,
-        opponent_policy: str = "network",
-        league_advantage_networks: list[list[torch.nn.Module]] | None = None,
-        self_play_anchor_probability: float = 0.0,
-        self_play_current_weight: float = 0.5,
-        self_play_recent_weight: float = 0.3,
-        self_play_older_weight: float = 0.2,
-        self_play_anchor_weight: float = 0.0,
-        self_play_recent_window: int = 5,
-        endpoint_depth_bucket_width: int = 10,
-        endpoint_depth_bucket_max: int = 100,
-        encoding=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.outcome_sampling_epsilon = min(1.0, max(0.0, float(outcome_sampling_epsilon)))
-        self.outcome_sampling_value_clip = (
-            None
-            if outcome_sampling_value_clip is None
-            else max(1.0e-9, float(outcome_sampling_value_clip))
-        )
-        self.outcome_unsampled_regret = outcome_unsampled_regret
-        if self.outcome_unsampled_regret not in {"negative_node_value", "zero"}:
-            raise ValueError("outcome_unsampled_regret must be 'negative_node_value' or 'zero'")
-        self.cutoff_value_mode = cutoff_value_mode
-        if self.cutoff_value_mode not in {"score_diff", "random_rollout"}:
-            raise ValueError("cutoff_value_mode must be 'score_diff' or 'random_rollout'")
-        self.cutoff_rollouts = max(0, int(cutoff_rollouts))
-        self.cutoff_rollout_policy = cutoff_rollout_policy
-        if self.cutoff_rollout_policy not in {"random", "safe_heuristic"}:
-            raise ValueError("cutoff_rollout_policy must be 'random' or 'safe_heuristic'")
-        self.cutoff_rollout_max_steps = max(1, int(cutoff_rollout_max_steps))
-        self.opponent_policy = opponent_policy
-        if self.opponent_policy not in {"network", "safe_heuristic", "self_play_league"}:
-            raise ValueError(
-                "opponent_policy must be 'network', 'safe_heuristic', or 'self_play_league'"
-            )
-        self.league_advantage_networks = league_advantage_networks or []
-        self.self_play_anchor_probability = min(1.0, max(0.0, float(self_play_anchor_probability)))
-        self.self_play_current_weight = max(0.0, float(self_play_current_weight))
-        self.self_play_recent_weight = max(0.0, float(self_play_recent_weight))
-        self.self_play_older_weight = max(0.0, float(self_play_older_weight))
-        self.self_play_anchor_weight = max(0.0, float(self_play_anchor_weight))
-        self.self_play_recent_window = max(0, int(self_play_recent_window))
-        self.endpoint_depth_bucket_width = max(1, int(endpoint_depth_bucket_width))
-        self.endpoint_depth_bucket_max = max(1, int(endpoint_depth_bucket_max))
-        self.encoding = encoding
-        self.rng = rng or np.random.default_rng()
-        self._safe_heuristic_rollout_bot = (
-            SafeHeuristicBot() if self.cutoff_rollout_policy == "safe_heuristic" else None
-        )
-        self._safe_heuristic_opponent_bot = (
-            SafeHeuristicBot()
-            if self.opponent_policy == "safe_heuristic" or self.self_play_anchor_probability > 0.0
-            else None
-        )
-        self._ensure_recursion_limit()
-
-    def _ensure_recursion_limit(self) -> None:
-        target_depth = self.max_depth if self.max_depth is not None else self.max_nodes
-        if target_depth is None:
-            target_depth = 10_000
-        desired_limit = min(max(int(target_depth) + 1_000, 2_000), 200_000)
-        if sys.getrecursionlimit() < desired_limit:
-            sys.setrecursionlimit(desired_limit)
-
-    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 self._cutoff_value(state, traverser, stats)
-        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 self._cutoff_value(state, traverser, stats)
-
-        player = state.current_player
-        fixed_action = self._fixed_opponent_action(state, player, traverser)
-        if fixed_action is not None:
-            unified_action = state.to_unified_action(fixed_action)
-            swapped_deck_index = self._sample_deck_draw_chance(state, unified_action)
-            state.push_action(fixed_action)
-            try:
-                return self._traverse(
-                    state,
-                    traverser,
-                    iteration,
-                    depth=depth + 1,
-                    stats=stats,
-                )
-            finally:
-                state.pop_action()
-                if swapped_deck_index is not None:
-                    state.swap_deck_cards(swapped_deck_index, len(state.deck) - 1)
-
-        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))
-
-        sampling_policy = self._sampling_policy(policy, legal)
-        action = self._sample_action(sampling_policy, legal_actions)
-        local_action = state.from_unified_action(int(action))
-        swapped_deck_index = self._sample_deck_draw_chance(state, int(action))
-        state.push_action(local_action)
-        try:
-            child_value = self._traverse(
-                state,
-                traverser,
-                iteration,
-                depth=depth + 1,
-                stats=stats,
-            )
-        finally:
-            state.pop_action()
-            if swapped_deck_index is not None:
-                state.swap_deck_cards(swapped_deck_index, len(state.deck) - 1)
-
-        stats.sampled_actions += 1
-        action_prob = max(float(sampling_policy[action]), self.epsilon)
-        sampled_action_value = child_value / action_prob
-        if self.outcome_sampling_value_clip is not None:
-            sampled_action_value = float(
-                np.clip(
-                    sampled_action_value,
-                    -self.outcome_sampling_value_clip,
-                    self.outcome_sampling_value_clip,
-                )
-            )
-        node_value = float(policy[action]) * sampled_action_value
-
-        if player == traverser:
-            if self.outcome_unsampled_regret == "zero":
-                regrets = np.zeros_like(policy, dtype=np.float32)
-            else:
-                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,
-                ),
-                self.rng,
-            )
-            stats.advantage_samples += 1
-
-        return node_value
-
-    def _policy(self, state: GameState, player: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-        return self._policy_from_networks(self.advantage_networks, state, player)
-
-    def _policy_from_networks(
-        self,
-        networks: list[torch.nn.Module],
-        state: GameState,
-        player: int,
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-        info_state = encode_info_state(state, player, self.encoding)
-        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 = 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 _fixed_opponent_action(
-        self,
-        state: GameState,
-        player: int,
-        traverser: int,
-    ) -> int | None:
-        if player == traverser or self.opponent_policy == "network":
-            return None
-        if self.opponent_policy == "safe_heuristic":
-            if self._safe_heuristic_opponent_bot is None:
-                self._safe_heuristic_opponent_bot = SafeHeuristicBot()
-            return self._safe_heuristic_opponent_bot.act(state)
-        bucket = self._self_play_bucket()
-        if bucket == "current":
-            return None
-        if bucket == "anchor":
-            if self._safe_heuristic_opponent_bot is None:
-                self._safe_heuristic_opponent_bot = SafeHeuristicBot()
-            return self._safe_heuristic_opponent_bot.act(state)
-        networks = self._self_play_snapshot_networks(bucket)
-        if networks is None:
-            return None
-        _, legal, policy = self._policy_from_networks(networks, state, player)
-        legal_actions = np.flatnonzero(legal)
-        if len(legal_actions) == 0:
-            return None
-        unified_action = self._sample_action(policy, legal_actions)
-        return state.from_unified_action(unified_action)
-
-    def _self_play_bucket(self) -> str:
-        if (
-            self.self_play_anchor_probability > 0.0
-            and self.rng.random() < self.self_play_anchor_probability
-        ):
-            return "anchor"
-        recent_count = min(len(self.league_advantage_networks), self.self_play_recent_window)
-        older_count = max(0, len(self.league_advantage_networks) - recent_count)
-        labels = ["current", "recent", "older", "anchor"]
-        weights = np.asarray(
-            [
-                self.self_play_current_weight,
-                self.self_play_recent_weight if recent_count > 0 else 0.0,
-                self.self_play_older_weight if older_count > 0 else 0.0,
-                self.self_play_anchor_weight,
-            ],
-            dtype=np.float64,
-        )
-        total = float(weights.sum())
-        if total <= 0.0:
-            return "current"
-        weights /= total
-        return str(self.rng.choice(labels, p=weights))
-
-    def _self_play_snapshot_networks(self, bucket: str) -> list[torch.nn.Module] | None:
-        if not self.league_advantage_networks:
-            return None
-        recent_count = min(len(self.league_advantage_networks), self.self_play_recent_window)
-        if bucket == "recent" and recent_count > 0:
-            candidates = self.league_advantage_networks[-recent_count:]
-        elif bucket == "older":
-            candidates = self.league_advantage_networks[
-                : max(0, len(self.league_advantage_networks) - recent_count)
-            ]
-        else:
-            candidates = self.league_advantage_networks
-        if not candidates:
-            return None
-        return candidates[int(self.rng.integers(0, len(candidates)))]
-
-    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 _sampling_policy(self, policy: np.ndarray, legal: np.ndarray) -> np.ndarray:
-        legal_count = int(np.count_nonzero(legal))
-        if legal_count <= 0:
-            return np.zeros_like(policy, dtype=np.float32)
-        if self.outcome_sampling_epsilon <= 0.0:
-            return policy.astype(np.float32)
-        uniform = legal.astype(np.float32) / float(legal_count)
-        return (
-            (1.0 - self.outcome_sampling_epsilon) * policy + self.outcome_sampling_epsilon * uniform
-        ).astype(np.float32)
-
-    def _cutoff_value(self, state: GameState, traverser: int, stats: TraversalStats) -> float:
-        if self.cutoff_value_mode == "score_diff" or self.cutoff_rollouts <= 0:
-            return float(state.score_diff(traverser))
-        total = 0.0
-        for _ in range(self.cutoff_rollouts):
-            total += self._rollout_value(state, traverser, stats)
-        return total / float(self.cutoff_rollouts)
-
-    def _rollout_value(self, state: GameState, traverser: int, stats: TraversalStats) -> float:
-        rollout_state = state.clone()
-        steps = 0
-        while not rollout_state.terminal and steps < self.cutoff_rollout_max_steps:
-            action = self._rollout_action(rollout_state)
-            if action is None:
-                break
-            unified_action = rollout_state.to_unified_action(action)
-            self._sample_deck_draw_chance(rollout_state, unified_action)
-            rollout_state.apply_action(action)
-            steps += 1
-        stats.cutoff_rollouts += 1
-        stats.cutoff_rollout_steps += steps
-        if not rollout_state.terminal:
-            stats.cutoff_rollout_timeouts += 1
-        return float(rollout_state.score_diff(traverser))
-
-    def _rollout_action(self, state: GameState) -> int | None:
-        if self.cutoff_rollout_policy == "safe_heuristic":
-            if self._safe_heuristic_rollout_bot is None:
-                self._safe_heuristic_rollout_bot = SafeHeuristicBot()
-            return self._safe_heuristic_rollout_bot.act(state)
-        legal_actions = state.unified_legal_actions()
-        if not legal_actions:
-            return None
-        return state.from_unified_action(int(self.rng.choice(legal_actions)))
-
-    def _sample_deck_draw_chance(self, state: GameState, unified_action: int) -> int | None:
-        deck_draw_action = 2 * state.config.hand_size
-        if state.phase != "draw" or unified_action != deck_draw_action or len(state.deck) <= 1:
-            return None
-        sampled_index = int(self.rng.integers(0, len(state.deck)))
-        if sampled_index == len(state.deck) - 1:
-            return None
-        state.swap_deck_cards(sampled_index, len(state.deck) - 1)
-        return sampled_index
-
-    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,
-            ),
-            self.rng,
-        )
-        stats.strategy_samples += 1
-
-    def _record_endpoint(self, stats: TraversalStats, depth: int) -> None:
-        stats.endpoint_depth_sum += depth
-        width = self.endpoint_depth_bucket_width
-        max_depth = self.endpoint_depth_bucket_max
-        start = min(depth // width * width, max_depth)
-        key = f"{max_depth}_plus" if start >= max_depth else f"{start}_{start + width - 1}"
-        stats.endpoint_depth_buckets[key] = stats.endpoint_depth_buckets.get(key, 0) + 1

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

@@ -4,14 +4,14 @@ import os
 from dataclasses import dataclass
 from typing import Any
 
-import numpy as np
 import torch
 
 from coolrl_lost_cities.games.classic.deep_cfr.config import config_from_dict
-from coolrl_lost_cities.games.classic.deep_cfr.memory import ReservoirMemory, TrainingSample
+from coolrl_lost_cities.games.classic.deep_cfr.memory import TrainingSample
 from coolrl_lost_cities.games.classic.deep_cfr.networks import DeepCFRMLP
-from coolrl_lost_cities.games.classic.deep_cfr.traverser import DeepCFRTraverser, TraversalStats
-from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig
+from coolrl_lost_cities.games.classic.deep_cfr.traversal import run_cython_traversal_batch
+from coolrl_lost_cities.games.classic.deep_cfr.traversal_stats import TraversalStats
+from coolrl_lost_cities.games.classic.game import LostCitiesConfig
 
 _TORCH_THREADS_CONFIGURED = False
 
@@ -76,14 +76,16 @@ def run_traversal_worker_batch(batch: TraversalWorkerBatch) -> TraversalWorkerRe
             network.load_state_dict(state_dict)
             network.eval()
         league_networks.append(snapshot_networks)
-    advantage_memory = ReservoirMemory()
-    strategy_memory = ReservoirMemory()
-    traverser = DeepCFRTraverser(
+    game_config = LostCitiesConfig(**batch.game_config)
+    total_stats, advantage_samples, strategy_samples = run_cython_traversal_batch(
         networks,
-        advantage_memory,
-        strategy_memory,
+        game_config,
+        batch.seeds,
+        batch.player,
+        batch.iteration,
         device=device,
         action_size=batch.action_size,
+        encoding=cfg.encoding,
         epsilon=cfg.traversal.regret_matching_epsilon,
         strategy_sample_interval=cfg.traversal.strategy_sample_interval,
         store_strategy_on_traverser_nodes=cfg.traversal.store_strategy_on_traverser_nodes,
@@ -107,19 +109,12 @@ def run_traversal_worker_batch(batch: TraversalWorkerBatch) -> TraversalWorkerRe
         self_play_recent_window=cfg.self_play.recent_window,
         endpoint_depth_bucket_width=cfg.traversal.endpoint_depth_bucket_width,
         endpoint_depth_bucket_max=cfg.traversal.endpoint_depth_bucket_max,
-        encoding=cfg.encoding,
-        rng=np.random.default_rng(batch.worker_seed),
+        seed=batch.worker_seed,
     )
-    game_config = LostCitiesConfig(**batch.game_config)
-    total_stats = TraversalStats()
-    for seed in batch.seeds:
-        state = GameState.new_game(game_config, seed=seed)
-        _, stats = traverser.traverse(state, batch.player, batch.iteration)
-        total_stats.accumulate(stats)
     return TraversalWorkerResult(
         player=batch.player,
         stats=total_stats,
-        advantage_samples=advantage_memory.all(),
-        strategy_samples=strategy_memory.all(),
+        advantage_samples=advantage_samples,
+        strategy_samples=strategy_samples,
         traversals=len(batch.seeds),
     )

tests/games/classic/test_deep_cfr_trainer.py

@@ -4,6 +4,7 @@ import re
 
 import numpy as np
 from coolrl_lost_cities.games.classic.deep_cfr.encoding import encode_info_state, input_dim
+from coolrl_lost_cities.games.classic.deep_cfr.traversal import CythonDeepCFRTraverser
 from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig
 
 from coolrl_lost_cities.games.classic.deep_cfr.benchmark import (
@@ -20,7 +21,6 @@ from coolrl_lost_cities.games.classic.deep_cfr.cli import (
 from coolrl_lost_cities.games.classic.deep_cfr.config import DeepCFRConfig, load_config
 from coolrl_lost_cities.games.classic.deep_cfr.memory import ReservoirMemory, TrainingSample
 from coolrl_lost_cities.games.classic.deep_cfr.trainer import DeepCFRTrainer
-from coolrl_lost_cities.games.classic.deep_cfr.traverser import DeepCFRTraverser
 
 
 def _deep_cfr_config(data: dict) -> DeepCFRConfig:
@@ -225,7 +225,7 @@ def test_deep_cfr_trainer_smoke_run() -> None:
     assert metrics[0].strategy_loss >= 0.0
 
 
-def test_deep_cfr_recursive_traverser_restores_state_and_collects_samples() -> None:
+def test_deep_cfr_cython_traverser_restores_state_and_collects_samples() -> None:
     trainer = DeepCFRTrainer(
         _deep_cfr_config(
             {
@@ -244,18 +244,17 @@ def test_deep_cfr_recursive_traverser_restores_state_and_collects_samples() -> N
     )
     state = GameState.new_game(LostCitiesConfig(seed=29), seed=29)
     before = state.to_snapshot()
-    traverser = DeepCFRTraverser(
+    traverser = CythonDeepCFRTraverser(
         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),
+        seed=29,
     )
 
     value, stats = traverser.traverse(state, traverser=0, iteration=1)
+    advantage_samples, strategy_samples = traverser.drain_samples()
 
     assert isinstance(value, float)
     assert state.to_snapshot() == before
@@ -263,14 +262,14 @@ def test_deep_cfr_recursive_traverser_restores_state_and_collects_samples() -> N
     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 len(strategy_samples) == stats.strategy_samples
+    assert len(advantage_samples) == stats.advantage_samples
+    sample = advantage_samples[0]
     assert sample.legal_mask.dtype == bool
     assert sample.target.shape == sample.legal_mask.shape
 
 
-def test_deep_cfr_traverser_supports_outcome_sampling_and_rollout_cutoffs() -> None:
+def test_deep_cfr_cython_traverser_supports_outcome_sampling_and_rollout_cutoffs() -> None:
     trainer = DeepCFRTrainer(
         _deep_cfr_config(
             {
@@ -296,10 +295,8 @@ def test_deep_cfr_traverser_supports_outcome_sampling_and_rollout_cutoffs() -> N
     )
     state = GameState.new_game(LostCitiesConfig(seed=31), seed=31)
     before = state.to_snapshot()
-    traverser = DeepCFRTraverser(
+    traverser = CythonDeepCFRTraverser(
         trainer.advantage_networks,
-        trainer.advantage_memory,
-        trainer.strategy_memory,
         device=trainer.device,
         action_size=trainer.action_size,
         max_depth=1,
@@ -311,16 +308,17 @@ def test_deep_cfr_traverser_supports_outcome_sampling_and_rollout_cutoffs() -> N
         cutoff_rollouts=2,
         cutoff_rollout_policy="random",
         cutoff_rollout_max_steps=16,
-        rng=np.random.default_rng(31),
+        seed=31,
     )
 
     _, stats = traverser.traverse(state, traverser=0, iteration=1)
+    advantage_samples, _ = traverser.drain_samples()
 
     assert state.to_snapshot() == before
     assert stats.depth_cutoffs > 0
     assert stats.cutoff_rollouts == stats.depth_cutoffs * 2
     assert stats.cutoff_rollout_steps > 0
-    sample = trainer.advantage_memory.all()[0]
+    sample = advantage_samples[0]
     unsampled_legal = sample.legal_mask.copy()
     unsampled_legal[np.nonzero(sample.target)[0]] = False
     assert np.all(sample.target[unsampled_legal] == 0.0)