diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/config.py b/src/coolrl_lost_cities/games/classic/deep_cfr/config.py index 55b084c..263fc8b 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/config.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/config.py @@ -97,6 +97,7 @@ class NetworkConfig(StrictModel): class TraversalConfig(StrictModel): traversals_per_iteration: int = 2 traversals_per_player: int | None = None + sampling_mode: str = "outcome" max_depth: int | None = 8 max_nodes: int | None = 10_000 max_nodes_per_traversal: int | None = None @@ -119,6 +120,14 @@ class TraversalConfig(StrictModel): endpoint_depth_bucket_width: int = 100 endpoint_depth_bucket_max: int = 1000 + @field_validator("sampling_mode") + @classmethod + def _validate_sampling_mode(cls, value: str) -> str: + token = value.strip().lower() + if token not in {"outcome", "external"}: + raise ValueError("must be 'outcome' or 'external'") + return token + @field_validator("outcome_unsampled_regret") @classmethod def _validate_unsampled_regret(cls, value: str) -> str: 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 bd17b67..70609a1 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py @@ -354,6 +354,7 @@ class DeepCFRTrainer: ), max_depth=self.config.traversal.max_depth, max_nodes=self.config.traversal.resolved_max_nodes(), + sampling_mode=self.config.traversal.sampling_mode, 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, 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 6ef482c..75bbda2 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/traversal.pyx +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/traversal.pyx @@ -82,6 +82,7 @@ cdef class CythonDeepCFRTraverser: cdef int max_depth cdef bint has_max_nodes cdef int max_nodes + cdef int sampling_mode_id cdef float outcome_sampling_epsilon cdef bint has_value_clip cdef float outcome_sampling_value_clip @@ -122,6 +123,7 @@ cdef class CythonDeepCFRTraverser: bint store_strategy_on_opponent_nodes=True, object max_depth=None, object max_nodes=None, + str sampling_mode="outcome", float outcome_sampling_epsilon=0.0, object outcome_sampling_value_clip=None, str outcome_unsampled_regret="negative_node_value", @@ -165,6 +167,12 @@ cdef class CythonDeepCFRTraverser: 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) + if sampling_mode == "outcome": + self.sampling_mode_id = 0 + elif sampling_mode == "external": + self.sampling_mode_id = 1 + else: + raise ValueError("sampling_mode must be 'outcome' or 'external'") 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 = ( @@ -269,6 +277,7 @@ cdef class CythonDeepCFRTraverser: cdef float action_prob cdef float sampled_action_value cdef float node_value + cdef float action_values[MAX_ACTIONS] cdef float policy[MAX_ACTIONS] cdef float sampling_policy[MAX_ACTIONS] cdef unsigned char legal[MAX_ACTIONS] @@ -323,8 +332,47 @@ cdef class CythonDeepCFRTraverser: 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)) + if self.sampling_mode_id == 1 and player == traverser: + node_value = 0.0 + for i in range(legal_count): + action = actions[i] + 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) + action_values[action] = child_value + node_value += policy[action] * child_value + self._record_regret_matching_decision( + stats, + state, + player, + actions[0], + depth, + policy_regret_fallback, + policy_argmax_tie_size, + policy_argmax_full_tie, + ) + self._record_external_advantage( + info_state, + legal, + action_values, + node_value, + iteration, + player, + stats, + ) + return node_value + + if self.sampling_mode_id == 1: + action = _sample_policy_from_actions_c(policy, actions, legal_count, _next_double(&self.rng)) + else: + 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) @@ -346,16 +394,25 @@ cdef class CythonDeepCFRTraverser: policy_argmax_tie_size, policy_argmax_full_tie, ) - action_prob = sampling_policy[action] + if self.sampling_mode_id == 1: + action_prob = policy[action] + else: + 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 self.sampling_mode_id == 1: + sampled_action_value = child_value + else: + sampled_action_value = child_value / action_prob + if self.sampling_mode_id == 0 and 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 self.sampling_mode_id == 1: + node_value = sampled_action_value + else: + node_value = policy[action] * sampled_action_value if player == traverser: self._record_advantage( @@ -888,6 +945,38 @@ cdef class CythonDeepCFRTraverser: ) stats.advantage_samples += 1 + cdef void _record_external_advantage( + self, + object info_state, + const unsigned char* legal, + const float* action_values, + 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: + target_view[i] = 0.0 + else: + target_view[i] = action_values[i] - 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 @@ -935,6 +1024,7 @@ def run_cython_traversal_batch( bint store_strategy_on_opponent_nodes=True, object max_depth=None, object max_nodes=None, + str sampling_mode="outcome", float outcome_sampling_epsilon=0.0, object outcome_sampling_value_clip=None, str outcome_unsampled_regret="negative_node_value", @@ -972,6 +1062,7 @@ def run_cython_traversal_batch( store_strategy_on_opponent_nodes=store_strategy_on_opponent_nodes, max_depth=max_depth, max_nodes=max_nodes, + sampling_mode=sampling_mode, outcome_sampling_epsilon=outcome_sampling_epsilon, outcome_sampling_value_clip=outcome_sampling_value_clip, outcome_unsampled_regret=outcome_unsampled_regret, 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 fa358de..4f94407 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/workers.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/workers.py @@ -101,6 +101,7 @@ def run_traversal_worker_batch(batch: TraversalWorkerBatch) -> TraversalWorkerRe store_strategy_on_opponent_nodes=cfg.traversal.store_strategy_on_opponent_nodes, max_depth=cfg.traversal.max_depth, max_nodes=cfg.traversal.resolved_max_nodes(), + sampling_mode=cfg.traversal.sampling_mode, outcome_sampling_epsilon=cfg.traversal.outcome_sampling_epsilon, outcome_sampling_value_clip=cfg.traversal.outcome_sampling_value_clip, outcome_unsampled_regret=cfg.traversal.outcome_unsampled_regret, diff --git a/tests/games/classic/test_deep_cfr_trainer.py b/tests/games/classic/test_deep_cfr_trainer.py index ea42b81..795f06c 100644 --- a/tests/games/classic/test_deep_cfr_trainer.py +++ b/tests/games/classic/test_deep_cfr_trainer.py @@ -51,6 +51,7 @@ def test_deep_cfr_loads_mapped_legacy_reproduction_config() -> None: assert config.network.hidden_size == 256 assert config.network.num_layers == 3 assert config.traversal.resolved_traversals_per_player() == 70 + assert config.traversal.sampling_mode == "outcome" assert config.traversal.max_depth is None assert config.traversal.resolved_max_nodes() == 1000 assert config.traversal.resolved_worker_chunk_size() == 8 @@ -111,6 +112,12 @@ def test_deep_cfr_train_cli_count_overrides_disable_duration_limits() -> None: assert overridden.checkpoint.save_latest is False +def test_deep_cfr_config_accepts_external_sampling_mode() -> None: + config = _deep_cfr_config({"traversal": {"sampling_mode": "external"}}) + + assert config.traversal.sampling_mode == "external" + + def test_deep_cfr_train_cli_checkpoint_save_overrides() -> None: args = type( "Args", @@ -448,6 +455,57 @@ def test_deep_cfr_cython_traverser_supports_outcome_sampling_and_rollout_cutoffs assert np.all(sample.target[unsampled_legal] == 0.0) +def test_deep_cfr_cython_traverser_supports_external_sampling() -> None: + trainer = DeepCFRTrainer( + _deep_cfr_config( + { + "run": {"iterations": 1, "seed": 33}, + "network": {"hidden_size": 16}, + "traversal": { + "traversals_per_iteration": 1, + "sampling_mode": "external", + "max_depth": 1, + "max_nodes": 64, + }, + "optimization": {"batch_size": 2}, + "checkpoint": {"save_every_iteration": False}, + } + ), + LostCitiesConfig(seed=33), + ) + metrics = trainer.train() + + assert len(metrics) == 1 + assert metrics[0].advantage_samples > 0 + state = GameState.new_game(LostCitiesConfig(seed=33), seed=33) + before = state.to_snapshot() + traverser = CythonDeepCFRTraverser( + trainer.advantage_networks, + device=trainer.device, + action_size=trainer.action_size, + sampling_mode="external", + max_depth=1, + max_nodes=64, + seed=33, + ) + + 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 + assert stats.nodes > 0 + assert stats.depth_cutoffs > 0 + assert stats.advantage_samples > 0 + assert stats.strategy_samples > 0 + assert len(advantage_samples) == stats.advantage_samples + assert len(strategy_samples) == stats.strategy_samples + sample = advantage_samples[0] + assert sample.legal_mask.dtype == bool + assert sample.target.shape == sample.legal_mask.shape + assert np.count_nonzero(sample.target[sample.legal_mask]) > 1 + + def test_deep_cfr_cython_traverser_records_regret_fallback_metrics() -> None: trainer = DeepCFRTrainer( _deep_cfr_config(