Add first-open replay reweighting

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:

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

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

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)

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,

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

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(