diff --git a/docs/plans/deep-cfr-selectivity.md b/docs/plans/deep-cfr-selectivity.md index 6632df6..2099992 100644 --- a/docs/plans/deep-cfr-selectivity.md +++ b/docs/plans/deep-cfr-selectivity.md @@ -222,7 +222,31 @@ more directly: Success criterion: the target distribution should show a usable separation between good and bad opens. If it does not, the training target is the blocker. -### 2. Short open-selectivity ablation +### 2. First-open replay reweighting + +Implemented option: + +- `optimization.advantage_first_open_fraction` + +Meaning: reserve this fraction of each advantage minibatch for samples where a +first-open action was legally available at the traverser decision. The target is +not changed; only replay sampling frequency changes. This keeps the experiment +on the pure self-play side more than hand-written target shaping. + +Initial planned run: + +- `run.experiment_name=first-open-reweight-50-512x3-det-500` +- `optimization.advantage_first_open_fraction=0.5` +- `traversal.outcome_sampling_epsilon=0.05` +- `traversal.outcome_unsampled_regret=zero` +- `run.deterministic=true` +- `run.max_iterations=500` + +Primary comparison is the `eps=0.05` confirmation run and the pure external +sampling run. Success requires bad-open rate and score/opened color to improve +together without a large score-diff regression. + +### 3. Short open-selectivity ablation Run a 200-300 iteration ablation only after the target audit identifies a specific change. Candidate changes include: 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 9e61627..3b5021f 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/config.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/config.py @@ -264,10 +264,19 @@ class OptimizationConfig(StrictModel): strategy_batch_size: int = 256 advantage_updates_per_iteration: int = 64 strategy_updates_per_iteration: int = 64 + advantage_first_open_fraction: float = 0.0 learning_rate: float = 1.0e-3 weight_decay: float = 0.0 grad_clip: float = 0.0 + @field_validator("advantage_first_open_fraction") + @classmethod + def _validate_advantage_first_open_fraction(cls, value: float) -> float: + fraction = float(value) + if not 0.0 <= fraction <= 1.0: + raise ValueError("must be between 0.0 and 1.0") + return fraction + class MemoryConfig(StrictModel): advantage_capacity: int = 2_000_000 diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/interleaved_traversal.py b/src/coolrl_lost_cities/games/classic/deep_cfr/interleaved_traversal.py index 7b68f73..ea50ca1 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/interleaved_traversal.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/interleaved_traversal.py @@ -92,6 +92,20 @@ def _masked_softmax(logits: np.ndarray, legal_mask: np.ndarray) -> np.ndarray: return policy +def _has_legal_first_open(state: GameState, player: int, legal_mask: np.ndarray) -> bool: + card_action_size = state.config.hand_size * 2 + hand = state.hand_slots(player) + expeditions = state.expeditions[player] + for unified_action in np.flatnonzero(legal_mask): + action = int(unified_action) + if action >= card_action_size or action % 2 == 1: + continue + card = hand[action // 2] + if card is not None and not expeditions[int(card.color)]: + return True + return False + + def _record_endpoint(stats: TraversalStats, depth: int, width: int, max_depth: int) -> None: stats.endpoint_depth_sum += depth start = (depth // width) * width @@ -395,6 +409,7 @@ class InterleavedContext: legal_mask=frame.legal_mask.copy(), iteration=self.iteration, player=frame.player, + is_first_open=_has_legal_first_open(self.state, frame.player, frame.legal_mask), ) ) self.stats.advantage_samples += 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 bdc1c9c..b2b59c1 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/memory.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/memory.py @@ -12,6 +12,7 @@ class TrainingSample: legal_mask: np.ndarray iteration: int player: int + is_first_open: bool = False class ReservoirMemory: @@ -31,6 +32,7 @@ class ReservoirMemory: legal_mask=np.asarray(sample.legal_mask, dtype=bool).copy(), iteration=int(sample.iteration), player=int(sample.player), + is_first_open=bool(sample.is_first_open), ) if self.capacity is None or len(self._samples) < self.capacity: self._samples.append(sample) @@ -60,14 +62,23 @@ class ReservoirMemory: rng: np.random.Generator, *, player: int | None = None, + first_open_only: bool = False, ) -> list[TrainingSample]: - candidates = ( - self._samples - if player is None - else [sample for sample in self._samples if sample.player == player] - ) + candidates = self._samples + if player is not None: + candidates = [sample for sample in candidates if sample.player == player] + if first_open_only: + candidates = [sample for sample in candidates if sample.is_first_open] if not candidates: raise ValueError("cannot sample from empty memory") size = min(int(batch_size), len(candidates)) indices = rng.choice(len(candidates), size=size, replace=len(candidates) < size) return [candidates[int(index)] for index in indices] + + def count(self, *, player: int | None = None, first_open_only: bool = False) -> int: + candidates = self._samples + if player is not None: + candidates = [sample for sample in candidates if sample.player == player] + if first_open_only: + candidates = [sample for sample in candidates if sample.is_first_open] + return len(candidates) 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 b72ec13..6b7b123 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py @@ -1045,10 +1045,7 @@ class DeepCFRTrainer: network.train() for _step in range(self.config.optimization.advantage_updates_per_iteration): sample_started = time.perf_counter() - batch = self.advantage_memories[player].sample( - self.config.optimization.advantage_batch_size, - self.rng, - ) + batch = self._sample_advantage_batch(player) self._runtime_metrics[f"time/advantage_player_{player}_sample_seconds"] = ( float( self._runtime_metrics.get(f"time/advantage_player_{player}_sample_seconds", 0.0) @@ -1105,6 +1102,29 @@ class DeepCFRTrainer: losses.append(float(loss.detach().float().cpu())) return float(np.mean(losses)) if losses else 0.0 + def _sample_advantage_batch(self, player: int) -> list[TrainingSample]: + memory = self.advantage_memories[player] + batch_size = int(self.config.optimization.advantage_batch_size) + first_open_fraction = float(self.config.optimization.advantage_first_open_fraction) + first_open_count = memory.count(first_open_only=True) + self._runtime_metrics[f"samples/advantage_player_{player}_first_open"] = first_open_count + if first_open_fraction <= 0.0 or first_open_count <= 0: + return memory.sample(batch_size, self.rng) + + first_open_batch_size = min(batch_size, max(1, round(batch_size * first_open_fraction))) + first_open_batch = memory.sample( + first_open_batch_size, + self.rng, + first_open_only=True, + ) + remaining = batch_size - len(first_open_batch) + if remaining <= 0: + self.rng.shuffle(first_open_batch) + return first_open_batch + batch = first_open_batch + memory.sample(remaining, self.rng) + self.rng.shuffle(batch) + return batch + def _train_strategy( self, network: nn.Module, 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 75bbda2..38205c3 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/traversal.pyx +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/traversal.pyx @@ -358,6 +358,7 @@ cdef class CythonDeepCFRTraverser: policy_argmax_full_tie, ) self._record_external_advantage( + state, info_state, legal, action_values, @@ -423,6 +424,7 @@ cdef class CythonDeepCFRTraverser: node_value, iteration, player, + self._has_legal_first_open(state, player, legal), stats, ) return node_value @@ -920,6 +922,7 @@ cdef class CythonDeepCFRTraverser: float node_value, int iteration, int player, + bint is_first_open, object stats, ): cdef int i @@ -941,12 +944,14 @@ cdef class CythonDeepCFRTraverser: legal_mask=legal_mask, iteration=iteration, player=player, + is_first_open=bool(is_first_open), ) ) stats.advantage_samples += 1 cdef void _record_external_advantage( self, + GameState state, object info_state, const unsigned char* legal, const float* action_values, @@ -973,10 +978,34 @@ cdef class CythonDeepCFRTraverser: legal_mask=legal_mask, iteration=iteration, player=player, + is_first_open=self._has_legal_first_open(state, player, legal), ) ) stats.advantage_samples += 1 + cdef bint _has_legal_first_open( + self, + GameState state, + int player, + const unsigned char* legal, + ) noexcept: + cdef int action + cdef int slot + cdef int card + cdef int color + cdef int card_action_size = state.hand_size * 2 + for action in range(card_action_size): + if legal[action] == 0 or action % 2 == 1: + continue + slot = action // 2 + if slot >= state.hand_lens[player]: + continue + card = state.hand_cards[state._hand_index(player, slot)] + color = state._card_color(card) + if state.expedition_lens[state._expedition_len_index(player, color)] == 0: + return True + return False + 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 diff --git a/tests/games/classic/test_deep_cfr_trainer.py b/tests/games/classic/test_deep_cfr_trainer.py index c360469..9432bcf 100644 --- a/tests/games/classic/test_deep_cfr_trainer.py +++ b/tests/games/classic/test_deep_cfr_trainer.py @@ -973,6 +973,63 @@ def test_reservoir_memory_caps_samples_and_filters_player_batches() -> None: assert all(sample.player == 1 for sample in player_one) +def test_reservoir_memory_filters_first_open_batches() -> None: + memory = ReservoirMemory() + rng = np.random.default_rng(37) + for index in range(6): + memory.add( + TrainingSample( + info_state=np.asarray([index], dtype=np.float32), + target=np.asarray([index], dtype=np.float32), + legal_mask=np.asarray([True]), + iteration=index, + player=0, + is_first_open=index % 2 == 0, + ), + rng, + ) + + first_open = memory.sample(8, rng, first_open_only=True) + + assert len(first_open) == 3 + assert all(sample.is_first_open for sample in first_open) + assert memory.count(first_open_only=True) == 3 + + +def test_deep_cfr_trainer_can_oversample_first_open_advantage_batches(tmp_path) -> None: + trainer = DeepCFRTrainer( + _deep_cfr_config( + { + "run": {"seed": 41}, + "network": {"hidden_size": 16}, + "optimization": { + "advantage_batch_size": 4, + "advantage_first_open_fraction": 0.5, + }, + } + ), + run_dir=tmp_path, + device="cpu", + ) + for index in range(8): + trainer.advantage_memories[0].add( + TrainingSample( + info_state=np.asarray([index], dtype=np.float32), + target=np.asarray([index], dtype=np.float32), + legal_mask=np.asarray([True]), + iteration=index, + player=0, + is_first_open=index in {1, 3}, + ), + trainer.rng, + ) + + batch = trainer._sample_advantage_batch(player=0) + + assert len(batch) == 4 + assert sum(sample.is_first_open for sample in batch) >= 2 + + def test_deep_cfr_trainer_saves_loads_and_evaluates_checkpoint(tmp_path) -> None: checkpoint_dir = tmp_path / "deep_cfr" trainer = DeepCFRTrainer(