diff --git a/docs/deep-cfr-v0-gap-vs-coolrl.md b/docs/deep-cfr-v0-gap-vs-coolrl.md index d03494a..cc7c74a 100644 --- a/docs/deep-cfr-v0-gap-vs-coolrl.md +++ b/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. diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/encoding.pxd b/src/coolrl_lost_cities/games/classic/deep_cfr/encoding.pxd index 98b7af8..fc87272 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/encoding.pxd +++ b/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 diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/memory.py b/src/coolrl_lost_cities/games/classic/deep_cfr/memory.py index 0078f25..bdc1c9c 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/memory.py +++ b/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) diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py b/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py index 700c175..ae894b8 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py +++ b/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,57 +244,67 @@ class DeepCFRTrainer: def _run_traversals_single_process(self, iteration: int) -> TraversalStats: 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.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, - max_depth=self.config.traversal.max_depth, - max_nodes=self.config.traversal.resolved_max_nodes(), - outcome_sampling_epsilon=self.config.traversal.outcome_sampling_epsilon, - outcome_sampling_value_clip=self.config.traversal.outcome_sampling_value_clip, - outcome_unsampled_regret=self.config.traversal.outcome_unsampled_regret, - cutoff_value_mode=self.config.traversal.cutoff_value_mode, - cutoff_rollouts=self.config.traversal.cutoff_rollouts, - 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(), - 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, - self_play_older_weight=self.config.self_play.older_weight, - self_play_anchor_weight=self.config.self_play.anchor_weight, - 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, - ) 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() - 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: - elapsed = time.perf_counter() - progress_started - self.tracker.log_event( - f"Traversal progress iteration={iteration} completed={completed} " - f"elapsed_seconds={elapsed:.2f} total_nodes={total_stats.nodes} " - f"nodes_per_second={total_stats.nodes / max(elapsed, 1.0e-12):.1f}" - ) + 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.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 + ), + max_depth=self.config.traversal.max_depth, + max_nodes=self.config.traversal.resolved_max_nodes(), + outcome_sampling_epsilon=self.config.traversal.outcome_sampling_epsilon, + outcome_sampling_value_clip=self.config.traversal.outcome_sampling_value_clip, + outcome_unsampled_regret=self.config.traversal.outcome_unsampled_regret, + cutoff_value_mode=self.config.traversal.cutoff_value_mode, + cutoff_rollouts=self.config.traversal.cutoff_rollouts, + 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=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, + self_play_older_weight=self.config.self_play.older_weight, + self_play_anchor_weight=self.config.self_play.anchor_weight, + 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, + seed=self.config.run.seed + iteration * 1_000_003 + player, + ) + total_stats.accumulate(stats) + 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} " + f"elapsed_seconds={elapsed:.2f} total_nodes={total_stats.nodes} " + f"nodes_per_second={total_stats.nodes / max(elapsed, 1.0e-12):.1f}" + ) return total_stats def _run_traversals_parallel(self, iteration: int) -> TraversalStats: @@ -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: diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/traversal.pyx b/src/coolrl_lost_cities/games/classic/deep_cfr/traversal.pyx index efa2d2d..2937103 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/traversal.pyx +++ b/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 _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 (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 (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 / 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 = (_next_u32(&self.rng) % 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 (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 / 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 = 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) % 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 = (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 = (_next_u32(&self.rng) % 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, diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/traversal_stats.py b/src/coolrl_lost_cities/games/classic/deep_cfr/traversal_stats.py new file mode 100644 index 0000000..a724410 --- /dev/null +++ b/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() + }, + } diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/traverser.py b/src/coolrl_lost_cities/games/classic/deep_cfr/traverser.py deleted file mode 100644 index 82334bf..0000000 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/traverser.py +++ /dev/null @@ -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 diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/workers.py b/src/coolrl_lost_cities/games/classic/deep_cfr/workers.py index 6d031eb..7f75f2b 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/workers.py +++ b/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), ) diff --git a/tests/games/classic/test_deep_cfr_trainer.py b/tests/games/classic/test_deep_cfr_trainer.py index 1f16da1..97cd192 100644 --- a/tests/games/classic/test_deep_cfr_trainer.py +++ b/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)