Add external sampling traversal mode
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user