Add external sampling traversal mode

This commit is contained in:
2026-05-07 15:23:40 +09:00
parent 712a1eedaf
commit 43507375fc
5 changed files with 166 additions and 6 deletions
@@ -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:
@@ -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,
@@ -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 <float>(state.total_scores[traverser] - state.total_scores[1 - traverser])
self._sampling_policy(policy, legal, sampling_policy)
action = _sample_policy_from_actions_c(sampling_policy, actions, legal_count, _next_double(&self.rng))
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,
@@ -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,
@@ -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(