Add non-default interleaved traversal scheduler

Co-Authored-By: Codex <codex@openai.com>
This commit is contained in:
2026-05-07 22:49:38 +09:00
co-authored by Codex
parent 09d58159b8
commit 240ef552c6
6 changed files with 1014 additions and 87 deletions
+47 -3
View File
@@ -1,7 +1,7 @@
# Plan: Option B Per-Worker Interleaved Traversal
**Status:** Phase 1 prototype started. Production trainer wiring has not begun;
default behavior is unchanged.
**Status:** Phase 2 non-default production prototype implemented. Default
behavior is unchanged.
**Owner:** Codex for prototype design and implementation; operator for long-run
benchmarks on `home`.
**Background:** Option A, the central traversal inference server, was implemented
@@ -195,7 +195,7 @@ Prototype parity: PASS for values, RNG outputs, aggregate traversal stats, and
sample checksum within float tolerance. This proves the scheduling shape can
form large policy batches. It does **not** yet prove production Cython parity.
### Phase 2: Cython Prototype Behind Non-Default Flag
### Phase 2: Non-Default Production Prototype
- Add an interleaved traversal entry point beside the existing recursive one.
- Keep the existing recursive path untouched and default.
@@ -205,6 +205,50 @@ form large policy batches. It does **not** yet prove production Cython parity.
Success gate: `uv run pytest -q tests/games/classic/test_deep_cfr_trainer.py`
and new interleaved traversal tests pass.
### Phase 2 Result (2026-05-07)
Implemented as a Python explicit-stack production path in
`interleaved_traversal.py`, not a Cython rewrite. This is deliberate: the Phase
1 prototype proved the scheduling shape, while a full Cython state-machine
rewrite would duplicate a large fraction of `traversal.pyx` before we know that
default-config wall-clock speedup survives trainer integration. The recursive
Cython path remains untouched and default.
Config surface added:
```yaml
traversal:
scheduler: recursive # recursive | interleaved
interleave_width: 64
interleave_max_batch: 128
```
Current interleaved guardrails:
- `sampling_mode: outcome`
- `opponent_policy: network`
- `cutoff_value_mode: score_diff`
- `cutoff_rollouts: 0`
- `inference_backend: local`
Validation:
- single-traversal parity against `run_cython_traversal_batch`: PASS for
aggregate stats and sample target checksums under identical RNG seed,
- trainer smoke with `traversal.scheduler=interleaved`: PASS,
- multiprocessing worker smoke via CLI: PASS,
- emitted runtime metrics:
`interleaved/batches`, `interleaved/requests`,
`interleaved/avg_batch_size`, `interleaved/max_batch_size`,
`interleaved/scheduler_seconds`, and `interleaved/forward_seconds`.
Important parity note: multi-traversal interleaving uses per-context RNG streams
so request scheduling does not couple one traversal's random stream to another.
That means exact recursive-batch RNG ordering is intentionally not preserved
across multiple simultaneous traversals. Phase 3 must therefore gate on sample
counts, traversal stats, learning stability, and wall-clock speedup, not
byte-identical multi-traversal sample order.
### Phase 3: Benchmark
Benchmark against current `default.yaml`, eval/checkpoint disabled:
@@ -112,6 +112,9 @@ class TraversalConfig(StrictModel):
endpoint_depth_bucket_width: int = 100
endpoint_depth_bucket_max: int = 1000
inference_backend: str = "local"
scheduler: str = "recursive"
interleave_width: int = 64
interleave_max_batch: int = 128
@field_validator("sampling_mode")
@classmethod
@@ -159,6 +162,14 @@ class TraversalConfig(StrictModel):
raise ValueError("must be 'local' or 'server'")
return token
@field_validator("scheduler")
@classmethod
def _validate_scheduler(cls, value: str) -> str:
token = value.strip().lower()
if token not in {"recursive", "interleaved"}:
raise ValueError("must be 'recursive' or 'interleaved'")
return token
@model_validator(mode="after")
def _validate_external_strategy_memory_convention(self) -> TraversalConfig:
if self.sampling_mode == "external" and (
@@ -173,6 +184,26 @@ class TraversalConfig(StrictModel):
"average-policy estimate. See "
"docs/research/strategy-memory-location.md."
)
if self.scheduler == "interleaved":
if self.sampling_mode != "outcome":
raise ValueError("scheduler='interleaved' currently supports only outcome sampling")
if self.opponent_policy != "network":
raise ValueError(
"scheduler='interleaved' currently supports only opponent_policy='network'"
)
if self.cutoff_rollouts != 0 or self.cutoff_value_mode != "score_diff":
raise ValueError(
"scheduler='interleaved' currently requires cutoff_value_mode='score_diff' "
"and cutoff_rollouts=0"
)
if self.inference_backend != "local":
raise ValueError(
"scheduler='interleaved' currently requires inference_backend='local'"
)
if self.interleave_width <= 0:
raise ValueError("interleave_width must be positive")
if self.interleave_max_batch <= 0:
raise ValueError("interleave_max_batch must be positive")
return self
def resolved_num_workers(self, batches: int | None = None) -> int:
@@ -0,0 +1,579 @@
from __future__ import annotations
import statistics
import time
from dataclasses import dataclass, field
from typing import Any
import numpy as np
import torch
from coolrl_lost_cities.games.classic.deep_cfr.encoding import encode_info_state
from coolrl_lost_cities.games.classic.deep_cfr.memory import TrainingSample
from coolrl_lost_cities.games.classic.deep_cfr.traversal_stats import TraversalStats
from coolrl_lost_cities.games.classic.game import GameState
def _next_u32(state: int) -> tuple[int, int]:
state = (state * 1664525 + 1013904223) & 0xFFFFFFFF
return state, state
def _next_double(state: int) -> tuple[int, float]:
state, value = _next_u32(state)
return state, value / 4294967296.0
def _sample_policy(policy: np.ndarray, actions: list[int], random_value: float) -> int:
fallback = -1
cumulative = 0.0
r = min(max(random_value, 0.0), 0.9999999999999999)
for action in actions:
if policy[action] > 0.0:
fallback = action
cumulative += float(policy[action])
if r < cumulative:
return action
return fallback
def _regret_matching(
advantages: np.ndarray,
legal_mask: np.ndarray,
epsilon: float,
) -> tuple[np.ndarray, bool, int, bool]:
policy = np.zeros_like(advantages, dtype=np.float32)
legal = np.flatnonzero(legal_mask)
positives = np.maximum(advantages[legal], 0.0)
positive_sum = float(positives.sum())
fallback = positive_sum <= epsilon
tie_size = 0
full_tie = False
if not fallback:
policy[legal] = positives / positive_sum
return policy, fallback, tie_size, full_tie
if len(legal) == 0:
return policy, fallback, tie_size, full_tie
best = float(np.max(advantages[legal]))
tied = legal[np.flatnonzero(advantages[legal] == best)]
tie_size = int(len(tied))
full_tie = tie_size > 1 and tie_size == len(legal)
policy[legal] = 1.0 / float(len(legal))
return policy, fallback, tie_size, full_tie
def _sampling_policy(policy: np.ndarray, legal_mask: np.ndarray, epsilon: float) -> np.ndarray:
legal = np.flatnonzero(legal_mask)
out = np.zeros_like(policy, dtype=np.float32)
if len(legal) == 0:
return out
if epsilon <= 0.0:
out[:] = policy
return out
uniform = 1.0 / float(len(legal))
out[legal] = (1.0 - epsilon) * policy[legal] + epsilon * uniform
return out
def _record_endpoint(stats: TraversalStats, depth: int, width: int, max_depth: int) -> None:
stats.endpoint_depth_sum += depth
start = (depth // width) * width
key = f"{max_depth}_plus" if start >= max_depth else f"{start}_{start + width - 1}"
stats.endpoint_depth_buckets[key] = stats.endpoint_depth_buckets.get(key, 0) + 1
@dataclass
class InterleavedTraversalConfig:
action_size: int
encoding: Any
epsilon: float
outcome_sampling_epsilon: float
outcome_sampling_value_clip: float | None
max_depth: int | None
max_nodes: int | None
strategy_sample_interval: int
store_strategy_on_traverser_nodes: bool
store_strategy_on_opponent_nodes: bool
endpoint_depth_bucket_width: int
endpoint_depth_bucket_max: int
@dataclass
class PolicyResult:
info_state: np.ndarray
legal_mask: np.ndarray
policy: np.ndarray
fallback: bool
tie_size: int
full_tie: bool
player: int = -1
depth: int = -1
@dataclass
class PolicyRequest:
context_index: int
player: int
info_state: np.ndarray
legal_mask: np.ndarray
depth: int
@dataclass
class Samples:
advantage: list[TrainingSample] = field(default_factory=list)
strategy: list[TrainingSample] = field(default_factory=list)
@dataclass
class AfterChildFrame:
depth: int
player: int
action: int
action_prob: float
info_state: np.ndarray
legal_mask: np.ndarray
policy: np.ndarray
sampling_policy: np.ndarray
fallback: bool
tie_size: int
full_tie: bool
swapped_deck_index: int
@dataclass
class EnterFrame:
depth: int
Frame = EnterFrame | AfterChildFrame
class BatchedPolicy:
def __init__(
self,
networks: list[torch.nn.Module],
*,
device: torch.device,
epsilon: float,
) -> None:
self.networks = networks
self.device = device
self.epsilon = epsilon
self.batch_sizes: list[int] = []
self.forward_seconds = 0.0
def batch(self, requests: list[PolicyRequest]) -> list[PolicyResult]:
if not requests:
return []
out: list[PolicyResult | None] = [None] * len(requests)
for player in sorted({request.player for request in requests}):
indices = [idx for idx, request in enumerate(requests) if request.player == player]
states = np.stack([requests[idx].info_state for idx in indices]).astype(np.float32)
x = torch.as_tensor(states, dtype=torch.float32, device=self.device)
if self.device.type == "cuda":
torch.cuda.synchronize(self.device)
start = time.perf_counter()
with torch.inference_mode():
values = self.networks[player](x).detach().cpu().numpy().astype(np.float32)
if self.device.type == "cuda":
torch.cuda.synchronize(self.device)
self.forward_seconds += time.perf_counter() - start
self.batch_sizes.append(len(indices))
for local_idx, request_idx in enumerate(indices):
request = requests[request_idx]
policy, fallback, tie_size, full_tie = _regret_matching(
values[local_idx], request.legal_mask, self.epsilon
)
out[request_idx] = PolicyResult(
info_state=request.info_state,
legal_mask=request.legal_mask,
policy=policy,
fallback=fallback,
tie_size=tie_size,
full_tie=full_tie,
)
return [result for result in out if result is not None]
class InterleavedContext:
def __init__(
self,
state: GameState,
*,
traverser: int,
iteration: int,
rng: int,
cfg: InterleavedTraversalConfig,
) -> None:
self.state = state
self.traverser = traverser
self.iteration = iteration
self.rng = rng
self.cfg = cfg
self.stats = TraversalStats()
self.samples = Samples()
self.stack: list[Frame] = [EnterFrame(0)]
self.pending: PolicyRequest | None = None
self.last_value = 0.0
self.done = False
self.value = 0.0
def advance_until_policy(self, context_index: int) -> None:
while not self.done and self.pending is None and self.stack:
frame = self.stack.pop()
if isinstance(frame, EnterFrame):
self._enter(frame.depth, context_index)
else:
self._after_child(frame)
if not self.stack and self.pending is None and not self.done:
self.done = True
self.value = self.last_value
def apply_policy(self, result: PolicyResult) -> None:
if self.pending is None:
raise RuntimeError("context has no pending policy request")
self.pending = None
player = result.player
depth = result.depth
if depth < 0:
raise RuntimeError("policy result is missing depth")
self._record_strategy(result, player, depth)
actions = [int(action) for action in np.flatnonzero(result.legal_mask)]
if not actions:
self.stats.terminals += 1
_record_endpoint(
self.stats,
depth,
self.cfg.endpoint_depth_bucket_width,
self.cfg.endpoint_depth_bucket_max,
)
self._return_value(float(self.state.score_diff(self.traverser)))
return
sampling_policy = _sampling_policy(
result.policy, result.legal_mask, self.cfg.outcome_sampling_epsilon
)
self.rng, random_value = _next_double(self.rng)
action = _sample_policy(sampling_policy, actions, random_value)
action_prob = max(float(sampling_policy[action]), self.cfg.epsilon)
swapped_deck_index = self._sample_deck_draw_chance(action)
self.state.push_unified_action(action)
self.stack.append(
AfterChildFrame(
depth=depth,
player=player,
action=action,
action_prob=action_prob,
info_state=result.info_state,
legal_mask=result.legal_mask,
policy=result.policy,
sampling_policy=sampling_policy,
fallback=result.fallback,
tie_size=result.tie_size,
full_tie=result.full_tie,
swapped_deck_index=swapped_deck_index,
)
)
self.stack.append(EnterFrame(depth + 1))
def _enter(self, depth: int, context_index: int) -> None:
self.stats.nodes += 1
self.stats.max_depth_reached = max(self.stats.max_depth_reached, depth)
cutoff = self._cutoff(depth)
if cutoff is not None:
self._return_value(cutoff)
return
player = int(self.state.current_player)
info_state = encode_info_state(self.state, player, self.cfg.encoding)
legal_mask = np.zeros(self.cfg.action_size, dtype=bool)
legal_mask[self.state.unified_legal_actions()] = True
request = PolicyRequest(context_index, player, info_state, legal_mask, depth)
self.pending = request
def _after_child(self, frame: AfterChildFrame) -> None:
child_value = self.last_value
self.state.pop_action()
if frame.swapped_deck_index >= 0:
self.state.swap_deck_cards(frame.swapped_deck_index, len(self.state.deck) - 1)
self.stats.sampled_actions += 1
self._record_regret_matching_decision(frame)
sampled_action_value = child_value / frame.action_prob
if self.cfg.outcome_sampling_value_clip is not None:
clip = float(self.cfg.outcome_sampling_value_clip)
sampled_action_value = max(-clip, min(clip, sampled_action_value))
node_value = float(frame.policy[frame.action]) * sampled_action_value
if frame.player == self.traverser:
target = np.zeros(self.cfg.action_size, dtype=np.float32)
target[frame.legal_mask] = -node_value
target[frame.action] = sampled_action_value - node_value
self.samples.advantage.append(
TrainingSample(
info_state=frame.info_state,
target=target,
legal_mask=frame.legal_mask.copy(),
iteration=self.iteration,
player=frame.player,
)
)
self.stats.advantage_samples += 1
self._return_value(node_value)
def _sample_deck_draw_chance(self, unified_action: int) -> int:
deck_draw_action = 2 * self.state.config.hand_size
deck_len = len(self.state.deck)
if self.state.phase != "draw" or unified_action != deck_draw_action or deck_len <= 1:
return -1
self.rng, value = _next_u32(self.rng)
sampled_index = int(value % deck_len)
if sampled_index == deck_len - 1:
return -1
self.state.swap_deck_cards(sampled_index, deck_len - 1)
return sampled_index
def _record_regret_matching_decision(self, frame: AfterChildFrame) -> None:
self.stats.regret_matching_decisions += 1
if not frame.fallback:
return
self.stats.regret_fallback_count += 1
self.stats.regret_fallback_depth_sum += frame.depth
self._record_fallback_depth_bucket(frame.depth)
expeditions = self.state.expeditions
opened_colors = sum(1 for cards in expeditions[frame.player] if cards)
self.stats.regret_fallback_opened_colors_sum += opened_colors
self.stats.regret_fallback_opened_colors_buckets[opened_colors] = (
self.stats.regret_fallback_opened_colors_buckets.get(opened_colors, 0) + 1
)
if frame.tie_size > 1:
self.stats.regret_fallback_argmax_tie_count += 1
self.stats.regret_fallback_argmax_tie_size_sum += frame.tie_size
if frame.full_tie:
self.stats.regret_fallback_argmax_full_tie_count += 1
card_action_size = self.state.card_action_size
hand = self.state.hand_slots(frame.player)
legal_actions = self.state.unified_legal_actions()
self.stats.regret_fallback_legal_actions_sum += len(legal_actions)
for legal_action in legal_actions:
if legal_action < card_action_size:
if legal_action % 2 == 1:
self.stats.regret_fallback_legal_discard_sum += 1
continue
card = hand[legal_action // 2]
color = int(card.color)
if not expeditions[frame.player][color]:
self.stats.regret_fallback_legal_open_new_sum += 1
self.stats.regret_fallback_open_new_available_by_color[color] = (
self.stats.regret_fallback_open_new_available_by_color.get(color, 0) + 1
)
else:
self.stats.regret_fallback_legal_play_existing_sum += 1
continue
if legal_action == card_action_size:
self.stats.regret_fallback_legal_draw_deck_sum += 1
else:
self.stats.regret_fallback_legal_draw_pile_sum += 1
if frame.action < card_action_size:
if frame.action % 2 == 1:
self.stats.regret_fallback_action_discard += 1
return
card = hand[frame.action // 2]
color = int(card.color)
if not expeditions[frame.player][color]:
self.stats.regret_fallback_action_open_new += 1
self.stats.regret_fallback_open_new_selected_by_color[color] = (
self.stats.regret_fallback_open_new_selected_by_color.get(color, 0) + 1
)
else:
self.stats.regret_fallback_action_play_existing += 1
return
if frame.action == card_action_size:
self.stats.regret_fallback_action_draw_deck += 1
else:
self.stats.regret_fallback_action_draw_pile += 1
def _record_fallback_depth_bucket(self, depth: int) -> None:
width = 50
max_depth = 400
start = (depth // width) * width
key = f"{max_depth}_plus" if start >= max_depth else f"{start}_{start + width - 1}"
self.stats.regret_fallback_depth_buckets[key] = (
self.stats.regret_fallback_depth_buckets.get(key, 0) + 1
)
def _cutoff(self, depth: int) -> float | None:
if self.cfg.max_nodes is not None and self.stats.nodes >= self.cfg.max_nodes:
self.stats.node_limit_cutoffs += 1
elif self.state.terminal:
self.stats.terminals += 1
elif self.cfg.max_depth is not None and depth >= self.cfg.max_depth:
self.stats.depth_cutoffs += 1
else:
return None
_record_endpoint(
self.stats,
depth,
self.cfg.endpoint_depth_bucket_width,
self.cfg.endpoint_depth_bucket_max,
)
return float(self.state.score_diff(self.traverser))
def _return_value(self, value: float) -> None:
self.last_value = value
if not self.stack:
self.done = True
self.value = value
def _record_strategy(self, result: PolicyResult, player: int, depth: int) -> None:
if player == self.traverser:
if not self.cfg.store_strategy_on_traverser_nodes:
return
elif not self.cfg.store_strategy_on_opponent_nodes:
return
if depth % self.cfg.strategy_sample_interval != 0:
return
self.samples.strategy.append(
TrainingSample(
info_state=result.info_state,
target=result.policy.copy(),
legal_mask=result.legal_mask.copy(),
iteration=self.iteration,
player=player,
)
)
self.stats.strategy_samples += 1
class InterleavedTraversalScheduler:
def __init__(self, cfg: InterleavedTraversalConfig, policy: BatchedPolicy) -> None:
self.cfg = cfg
self.policy = policy
self.scheduler_seconds = 0.0
def run(
self,
states: list[GameState],
*,
traverser: int,
iteration: int,
rng_seeds: list[int],
interleave_width: int,
max_batch: int,
) -> tuple[list[float], list[int], list[TraversalStats], list[Samples], list[int]]:
contexts = [
InterleavedContext(
state,
traverser=traverser,
iteration=iteration,
rng=rng,
cfg=self.cfg,
)
for state, rng in zip(states, rng_seeds, strict=True)
]
active = list(range(len(contexts)))
batch_sizes: list[int] = []
while active:
start = time.perf_counter()
runnable = active[: max(1, interleave_width)]
for context_index in runnable:
contexts[context_index].advance_until_policy(context_index)
requests: list[PolicyRequest] = []
request_contexts: list[int] = []
for context_index in runnable:
request = contexts[context_index].pending
if request is not None:
requests.append(request)
request_contexts.append(context_index)
if len(requests) >= max_batch:
break
self.scheduler_seconds += time.perf_counter() - start
if requests:
results = self.policy.batch(requests)
batch_sizes.append(len(requests))
for context_index, request, result in zip(
request_contexts, requests, results, strict=True
):
result.player = request.player
result.depth = request.depth
contexts[context_index].apply_policy(result)
continue
active = [idx for idx in active if not contexts[idx].done]
return (
[context.value for context in contexts],
[context.rng for context in contexts],
[context.stats for context in contexts],
[context.samples for context in contexts],
batch_sizes,
)
def run_interleaved_traversal_batch(
advantage_networks: list[torch.nn.Module],
game_config: Any,
seeds: list[int],
player: int,
iteration: int,
*,
device: torch.device,
action_size: int,
encoding: Any,
epsilon: float,
strategy_sample_interval: int,
store_strategy_on_traverser_nodes: bool,
store_strategy_on_opponent_nodes: bool,
max_depth: int | None,
max_nodes: int | None,
outcome_sampling_epsilon: float,
outcome_sampling_value_clip: float | None,
endpoint_depth_bucket_width: int,
endpoint_depth_bucket_max: int,
seed: int,
interleave_width: int,
interleave_max_batch: int,
) -> tuple[TraversalStats, list[TrainingSample], list[TrainingSample], dict[str, float | int]]:
cfg = InterleavedTraversalConfig(
action_size=action_size,
encoding=encoding,
epsilon=epsilon,
outcome_sampling_epsilon=outcome_sampling_epsilon,
outcome_sampling_value_clip=outcome_sampling_value_clip,
max_depth=max_depth,
max_nodes=max_nodes,
strategy_sample_interval=strategy_sample_interval,
store_strategy_on_traverser_nodes=store_strategy_on_traverser_nodes,
store_strategy_on_opponent_nodes=store_strategy_on_opponent_nodes,
endpoint_depth_bucket_width=endpoint_depth_bucket_width,
endpoint_depth_bucket_max=endpoint_depth_bucket_max,
)
states = [GameState.new_game(game_config, seed=game_seed) for game_seed in seeds]
rng_seeds = [int(seed) + index * 1_000_003 for index in range(len(seeds))]
policy = BatchedPolicy(advantage_networks, device=device, epsilon=cfg.epsilon)
scheduler = InterleavedTraversalScheduler(cfg, policy)
_values, _rng_out, stats_rows, sample_rows, batch_sizes = scheduler.run(
states,
traverser=player,
iteration=iteration,
rng_seeds=rng_seeds,
interleave_width=max(1, int(interleave_width)),
max_batch=max(1, int(interleave_max_batch)),
)
total_stats = TraversalStats()
advantage_samples: list[TrainingSample] = []
strategy_samples: list[TrainingSample] = []
for stats, samples in zip(stats_rows, sample_rows, strict=True):
total_stats.accumulate(stats)
advantage_samples.extend(samples.advantage)
strategy_samples.extend(samples.strategy)
runtime_stats: dict[str, float | int] = {
"interleaved/batches": len(batch_sizes),
"interleaved/requests": sum(batch_sizes),
"interleaved/max_batch_size": max(batch_sizes) if batch_sizes else 0,
"interleaved/avg_batch_size": (float(statistics.mean(batch_sizes)) if batch_sizes else 0.0),
"interleaved/scheduler_seconds": scheduler.scheduler_seconds,
"interleaved/forward_seconds": policy.forward_seconds,
}
return total_stats, advantage_samples, strategy_samples, runtime_stats
@@ -27,6 +27,9 @@ from coolrl_lost_cities.games.classic.deep_cfr.inference_server import (
BatchStatsMessage,
InferenceServerController,
)
from coolrl_lost_cities.games.classic.deep_cfr.interleaved_traversal import (
run_interleaved_traversal_batch,
)
from coolrl_lost_cities.games.classic.deep_cfr.memory import ReservoirMemory, TrainingSample
from coolrl_lost_cities.games.classic.deep_cfr.networks import DeepCFRMLP
from coolrl_lost_cities.games.classic.deep_cfr.tracking import (
@@ -394,51 +397,87 @@ class DeepCFRTrainer:
self.config.run.seed + iteration * 10_000 + traversal_index * 10 + player
for traversal_index in range(traversals_per_player)
]
stats, advantage_samples, strategy_samples = run_cython_traversal_batch(
self.advantage_networks,
self.game_config,
seeds,
player,
iteration,
device=self.device,
strategy_network=(
self.strategy_network
if self.config.traversal.opponent_policy == "average_strategy"
else None
),
action_size=self.action_size,
encoding=self.config.encoding,
epsilon=self.config.traversal.regret_matching_epsilon,
strategy_sample_interval=self.config.traversal.strategy_sample_interval,
store_strategy_on_traverser_nodes=(
self.config.traversal.store_strategy_on_traverser_nodes
),
store_strategy_on_opponent_nodes=(
self.config.traversal.store_strategy_on_opponent_nodes
),
max_depth=self.config.traversal.max_depth,
max_nodes=self.config.traversal.max_nodes_per_traversal,
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,
cutoff_value_mode=self.config.traversal.cutoff_value_mode,
cutoff_rollouts=self.config.traversal.cutoff_rollouts,
cutoff_rollout_policy=self.config.traversal.cutoff_rollout_policy,
cutoff_rollout_max_steps=self.config.traversal.cutoff_rollout_max_steps,
opponent_policy=self.config.traversal.opponent_policy,
all_negative_fallback=self.config.regret_matching.all_negative_fallback,
league_advantage_networks=league_networks,
self_play_anchor_probability=self.config.self_play.anchor_probability,
self_play_current_weight=self.config.self_play.current_weight,
self_play_recent_weight=self.config.self_play.recent_weight,
self_play_older_weight=self.config.self_play.older_weight,
self_play_anchor_weight=self.config.self_play.anchor_weight,
self_play_recent_window=self.config.self_play.recent_window,
endpoint_depth_bucket_width=self.config.traversal.endpoint_depth_bucket_width,
endpoint_depth_bucket_max=self.config.traversal.endpoint_depth_bucket_max,
seed=self.config.run.seed + iteration * 1_000_003 + player,
)
if self.config.traversal.scheduler == "interleaved":
stats, advantage_samples, strategy_samples, runtime_metrics = (
run_interleaved_traversal_batch(
self.advantage_networks,
self.game_config,
seeds,
player,
iteration,
device=self.device,
action_size=self.action_size,
encoding=self.config.encoding,
epsilon=self.config.traversal.regret_matching_epsilon,
strategy_sample_interval=self.config.traversal.strategy_sample_interval,
store_strategy_on_traverser_nodes=(
self.config.traversal.store_strategy_on_traverser_nodes
),
store_strategy_on_opponent_nodes=(
self.config.traversal.store_strategy_on_opponent_nodes
),
max_depth=self.config.traversal.max_depth,
max_nodes=self.config.traversal.max_nodes_per_traversal,
outcome_sampling_epsilon=self.config.traversal.outcome_sampling_epsilon,
outcome_sampling_value_clip=(
self.config.traversal.outcome_sampling_value_clip
),
endpoint_depth_bucket_width=(
self.config.traversal.endpoint_depth_bucket_width
),
endpoint_depth_bucket_max=self.config.traversal.endpoint_depth_bucket_max,
seed=self.config.run.seed + iteration * 1_000_003 + player,
interleave_width=self.config.traversal.interleave_width,
interleave_max_batch=self.config.traversal.interleave_max_batch,
)
)
self._record_interleaved_metrics(runtime_metrics)
else:
stats, advantage_samples, strategy_samples = run_cython_traversal_batch(
self.advantage_networks,
self.game_config,
seeds,
player,
iteration,
device=self.device,
strategy_network=(
self.strategy_network
if self.config.traversal.opponent_policy == "average_strategy"
else None
),
action_size=self.action_size,
encoding=self.config.encoding,
epsilon=self.config.traversal.regret_matching_epsilon,
strategy_sample_interval=self.config.traversal.strategy_sample_interval,
store_strategy_on_traverser_nodes=(
self.config.traversal.store_strategy_on_traverser_nodes
),
store_strategy_on_opponent_nodes=(
self.config.traversal.store_strategy_on_opponent_nodes
),
max_depth=self.config.traversal.max_depth,
max_nodes=self.config.traversal.max_nodes_per_traversal,
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,
cutoff_value_mode=self.config.traversal.cutoff_value_mode,
cutoff_rollouts=self.config.traversal.cutoff_rollouts,
cutoff_rollout_policy=self.config.traversal.cutoff_rollout_policy,
cutoff_rollout_max_steps=self.config.traversal.cutoff_rollout_max_steps,
opponent_policy=self.config.traversal.opponent_policy,
all_negative_fallback=self.config.regret_matching.all_negative_fallback,
league_advantage_networks=league_networks,
self_play_anchor_probability=self.config.self_play.anchor_probability,
self_play_current_weight=self.config.self_play.current_weight,
self_play_recent_weight=self.config.self_play.recent_weight,
self_play_older_weight=self.config.self_play.older_weight,
self_play_anchor_weight=self.config.self_play.anchor_weight,
self_play_recent_window=self.config.self_play.recent_window,
endpoint_depth_bucket_width=self.config.traversal.endpoint_depth_bucket_width,
endpoint_depth_bucket_max=self.config.traversal.endpoint_depth_bucket_max,
seed=self.config.run.seed + iteration * 1_000_003 + player,
)
total_stats.accumulate(stats)
memory_add_started = time.perf_counter()
self._add_advantage_samples(advantage_samples)
@@ -514,6 +553,7 @@ class DeepCFRTrainer:
result = future.result()
completed_batches += 1
total_stats.accumulate(result.stats)
self._record_interleaved_metrics(result.runtime_metrics)
memory_add_started = time.perf_counter()
self._add_advantage_samples(result.advantage_samples)
self.strategy_memory.add_many(result.strategy_samples, self.rng)
@@ -548,6 +588,28 @@ class DeepCFRTrainer:
self._record_inference_batch_stats(self._inference_server.drain_batch_stats())
return total_stats
def _record_interleaved_metrics(self, metrics: dict[str, float | int] | None) -> None:
if not metrics:
return
for key in (
"interleaved/batches",
"interleaved/requests",
"interleaved/scheduler_seconds",
"interleaved/forward_seconds",
):
self._runtime_metrics[key] = float(self._runtime_metrics.get(key, 0.0)) + float(
metrics.get(key, 0.0)
)
self._runtime_metrics["interleaved/max_batch_size"] = max(
int(self._runtime_metrics.get("interleaved/max_batch_size", 0)),
int(metrics.get("interleaved/max_batch_size", 0)),
)
batches = float(self._runtime_metrics.get("interleaved/batches", 0.0))
requests = float(self._runtime_metrics.get("interleaved/requests", 0.0))
self._runtime_metrics["interleaved/avg_batch_size"] = (
requests / batches if batches > 0.0 else 0.0
)
def _worker_batches(self, iteration: int) -> list[TraversalWorkerBatch]:
batches: list[TraversalWorkerBatch] = []
if self.config.traversal.inference_backend == "server":
@@ -15,6 +15,9 @@ from coolrl_lost_cities.games.classic.deep_cfr.inference_client import (
InferenceClient,
NetworkProxy,
)
from coolrl_lost_cities.games.classic.deep_cfr.interleaved_traversal import (
run_interleaved_traversal_batch,
)
from coolrl_lost_cities.games.classic.deep_cfr.memory import TrainingSample
from coolrl_lost_cities.games.classic.deep_cfr.networks import DeepCFRMLP
from coolrl_lost_cities.games.classic.deep_cfr.traversal import run_cython_traversal_batch
@@ -69,6 +72,7 @@ class TraversalWorkerResult:
advantage_samples: list[TrainingSample]
strategy_samples: list[TrainingSample]
traversals: int
runtime_metrics: dict[str, float | int] | None = None
def run_traversal_worker_batch(batch: TraversalWorkerBatch) -> TraversalWorkerResult:
@@ -142,43 +146,75 @@ def run_traversal_worker_batch(batch: TraversalWorkerBatch) -> TraversalWorkerRe
strategy_network.load_state_dict(batch.strategy_network)
strategy_network.eval()
game_config = LostCitiesConfig(**batch.game_config)
total_stats, advantage_samples, strategy_samples = run_cython_traversal_batch(
networks,
game_config,
batch.seeds,
batch.player,
batch.iteration,
device=device,
strategy_network=strategy_network,
action_size=batch.action_size,
encoding=cfg.encoding,
epsilon=cfg.traversal.regret_matching_epsilon,
strategy_sample_interval=cfg.traversal.strategy_sample_interval,
store_strategy_on_traverser_nodes=cfg.traversal.store_strategy_on_traverser_nodes,
store_strategy_on_opponent_nodes=cfg.traversal.store_strategy_on_opponent_nodes,
max_depth=cfg.traversal.max_depth,
max_nodes=cfg.traversal.max_nodes_per_traversal,
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,
cutoff_value_mode=cfg.traversal.cutoff_value_mode,
cutoff_rollouts=cfg.traversal.cutoff_rollouts,
cutoff_rollout_policy=cfg.traversal.cutoff_rollout_policy,
cutoff_rollout_max_steps=cfg.traversal.cutoff_rollout_max_steps,
opponent_policy=cfg.traversal.opponent_policy,
all_negative_fallback=cfg.regret_matching.all_negative_fallback,
league_advantage_networks=league_networks,
self_play_anchor_probability=cfg.self_play.anchor_probability,
self_play_current_weight=cfg.self_play.current_weight,
self_play_recent_weight=cfg.self_play.recent_weight,
self_play_older_weight=cfg.self_play.older_weight,
self_play_anchor_weight=cfg.self_play.anchor_weight,
self_play_recent_window=cfg.self_play.recent_window,
endpoint_depth_bucket_width=cfg.traversal.endpoint_depth_bucket_width,
endpoint_depth_bucket_max=cfg.traversal.endpoint_depth_bucket_max,
seed=batch.worker_seed,
)
if cfg.traversal.scheduler == "interleaved":
total_stats, advantage_samples, strategy_samples, runtime_metrics = (
run_interleaved_traversal_batch(
networks,
game_config,
batch.seeds,
batch.player,
batch.iteration,
device=device,
action_size=batch.action_size,
encoding=cfg.encoding,
epsilon=cfg.traversal.regret_matching_epsilon,
strategy_sample_interval=cfg.traversal.strategy_sample_interval,
store_strategy_on_traverser_nodes=(
cfg.traversal.store_strategy_on_traverser_nodes
),
store_strategy_on_opponent_nodes=(
cfg.traversal.store_strategy_on_opponent_nodes
),
max_depth=cfg.traversal.max_depth,
max_nodes=cfg.traversal.max_nodes_per_traversal,
outcome_sampling_epsilon=cfg.traversal.outcome_sampling_epsilon,
outcome_sampling_value_clip=cfg.traversal.outcome_sampling_value_clip,
endpoint_depth_bucket_width=cfg.traversal.endpoint_depth_bucket_width,
endpoint_depth_bucket_max=cfg.traversal.endpoint_depth_bucket_max,
seed=batch.worker_seed,
interleave_width=cfg.traversal.interleave_width,
interleave_max_batch=cfg.traversal.interleave_max_batch,
)
)
else:
total_stats, advantage_samples, strategy_samples = run_cython_traversal_batch(
networks,
game_config,
batch.seeds,
batch.player,
batch.iteration,
device=device,
strategy_network=strategy_network,
action_size=batch.action_size,
encoding=cfg.encoding,
epsilon=cfg.traversal.regret_matching_epsilon,
strategy_sample_interval=cfg.traversal.strategy_sample_interval,
store_strategy_on_traverser_nodes=cfg.traversal.store_strategy_on_traverser_nodes,
store_strategy_on_opponent_nodes=cfg.traversal.store_strategy_on_opponent_nodes,
max_depth=cfg.traversal.max_depth,
max_nodes=cfg.traversal.max_nodes_per_traversal,
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,
cutoff_value_mode=cfg.traversal.cutoff_value_mode,
cutoff_rollouts=cfg.traversal.cutoff_rollouts,
cutoff_rollout_policy=cfg.traversal.cutoff_rollout_policy,
cutoff_rollout_max_steps=cfg.traversal.cutoff_rollout_max_steps,
opponent_policy=cfg.traversal.opponent_policy,
all_negative_fallback=cfg.regret_matching.all_negative_fallback,
league_advantage_networks=league_networks,
self_play_anchor_probability=cfg.self_play.anchor_probability,
self_play_current_weight=cfg.self_play.current_weight,
self_play_recent_weight=cfg.self_play.recent_weight,
self_play_older_weight=cfg.self_play.older_weight,
self_play_anchor_weight=cfg.self_play.anchor_weight,
self_play_recent_window=cfg.self_play.recent_window,
endpoint_depth_bucket_width=cfg.traversal.endpoint_depth_bucket_width,
endpoint_depth_bucket_max=cfg.traversal.endpoint_depth_bucket_max,
seed=batch.worker_seed,
)
runtime_metrics = None
finally:
if client is not None:
client.close()
@@ -188,4 +224,5 @@ def run_traversal_worker_batch(batch: TraversalWorkerBatch) -> TraversalWorkerRe
advantage_samples=advantage_samples,
strategy_samples=strategy_samples,
traversals=len(batch.seeds),
runtime_metrics=runtime_metrics,
)
+176 -2
View File
@@ -7,7 +7,10 @@ import numpy as np
import pytest
import torch
from coolrl_lost_cities.games.classic.deep_cfr.encoding import encode_info_state, input_dim
from coolrl_lost_cities.games.classic.deep_cfr.traversal import CythonDeepCFRTraverser
from coolrl_lost_cities.games.classic.deep_cfr.traversal import (
CythonDeepCFRTraverser,
run_cython_traversal_batch,
)
from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig
from coolrl_lost_cities.games.classic.deep_cfr.benchmark import (
@@ -23,6 +26,9 @@ from coolrl_lost_cities.games.classic.deep_cfr.cli import (
)
from coolrl_lost_cities.games.classic.deep_cfr.config import DeepCFRConfig, load_config
from coolrl_lost_cities.games.classic.deep_cfr.evaluate import evaluate_strategy_network
from coolrl_lost_cities.games.classic.deep_cfr.interleaved_traversal import (
run_interleaved_traversal_batch,
)
from coolrl_lost_cities.games.classic.deep_cfr.memory import ReservoirMemory, TrainingSample
from coolrl_lost_cities.games.classic.deep_cfr.networks import DeepCFRMLP
from coolrl_lost_cities.games.classic.deep_cfr.trainer import (
@@ -108,11 +114,44 @@ def test_deep_cfr_train_cli_accepts_run_and_traversal_config_overrides() -> None
def test_deep_cfr_config_accepts_external_sampling_mode() -> None:
config = _deep_cfr_config({"traversal": {"sampling_mode": "external"}})
config = _deep_cfr_config(
{
"traversal": {
"sampling_mode": "external",
"store_strategy_on_traverser_nodes": False,
}
}
)
assert config.traversal.sampling_mode == "external"
def test_deep_cfr_config_accepts_interleaved_scheduler() -> None:
config = _deep_cfr_config(
{"traversal": {"scheduler": "interleaved", "opponent_policy": "network"}}
)
assert config.traversal.scheduler == "interleaved"
assert config.traversal.opponent_policy == "network"
def test_deep_cfr_config_rejects_unsupported_interleaved_options() -> None:
with pytest.raises(ValueError, match="opponent_policy='network'"):
_deep_cfr_config(
{"traversal": {"scheduler": "interleaved", "opponent_policy": "self_play_league"}}
)
with pytest.raises(ValueError, match="requires inference_backend='local'"):
_deep_cfr_config(
{
"traversal": {
"scheduler": "interleaved",
"opponent_policy": "network",
"inference_backend": "server",
}
}
)
def test_deep_cfr_train_cli_checkpoint_save_overrides() -> None:
args = type(
"Args",
@@ -138,6 +177,7 @@ def test_deep_cfr_train_cli_accepts_generic_config_overrides() -> None:
{
"config_overrides": [
"traversal.sampling_mode=external",
"traversal.store_strategy_on_traverser_nodes=false",
"traversal.max_depth=null",
"optimization.advantage_batch_size=64",
"checkpoint.save_latest=true",
@@ -368,6 +408,139 @@ def test_deep_cfr_trainer_smoke_run() -> None:
assert metrics[0].strategy_loss >= 0.0
def test_deep_cfr_trainer_interleaved_scheduler_smoke_run(tmp_path) -> None:
trainer = DeepCFRTrainer(
_deep_cfr_config(
{
"run": {"max_iterations": 1, "seed": 24},
"network": {"hidden_size": 16},
"traversal": {
"scheduler": "interleaved",
"opponent_policy": "network",
"traversals_per_player": 2,
"max_depth": 3,
"max_nodes_per_traversal": 64,
"interleave_width": 4,
"interleave_max_batch": 8,
},
"optimization": {
"advantage_updates_per_iteration": 1,
"strategy_updates_per_iteration": 1,
"advantage_batch_size": 2,
"strategy_batch_size": 2,
},
"checkpoint": {"save_every": 0, "save_latest": False},
"evaluation": {"eval_every": 0},
}
),
LostCitiesConfig(seed=24),
run_dir=tmp_path / "interleaved",
)
metrics = trainer.train()
runtime = metrics[0].runtime_metrics
assert len(metrics) == 1
assert metrics[0].advantage_samples > 0
assert metrics[0].strategy_samples > 0
assert metrics[0].traversal_nodes > 0
assert runtime["interleaved/batches"] > 0
assert runtime["interleaved/requests"] > 0
assert runtime["interleaved/max_batch_size"] >= 1
assert runtime["interleaved/avg_batch_size"] >= 1.0
def test_deep_cfr_interleaved_scheduler_matches_recursive_single_traversal() -> None:
config = _deep_cfr_config(
{
"run": {"seed": 26},
"network": {"hidden_size": 16},
"traversal": {
"opponent_policy": "network",
"max_depth": 3,
"max_nodes_per_traversal": 64,
},
}
)
game_config = LostCitiesConfig(seed=26)
probe = GameState.new_game(game_config, seed=26)
action_size = game_config.action_size
torch.manual_seed(26)
networks = [
DeepCFRMLP.from_config(input_dim(probe, config.encoding), action_size, config.network)
for _ in range(2)
]
for network in networks:
network.eval()
common = {
"device": torch.device("cpu"),
"action_size": action_size,
"encoding": config.encoding,
"epsilon": config.traversal.regret_matching_epsilon,
"strategy_sample_interval": config.traversal.strategy_sample_interval,
"store_strategy_on_traverser_nodes": config.traversal.store_strategy_on_traverser_nodes,
"store_strategy_on_opponent_nodes": config.traversal.store_strategy_on_opponent_nodes,
"max_depth": config.traversal.max_depth,
"max_nodes": config.traversal.max_nodes_per_traversal,
"outcome_sampling_epsilon": config.traversal.outcome_sampling_epsilon,
"outcome_sampling_value_clip": config.traversal.outcome_sampling_value_clip,
"endpoint_depth_bucket_width": config.traversal.endpoint_depth_bucket_width,
"endpoint_depth_bucket_max": config.traversal.endpoint_depth_bucket_max,
"seed": 2601,
}
recursive_stats, recursive_advantage, recursive_strategy = run_cython_traversal_batch(
networks,
game_config,
[260],
0,
1,
**common,
strategy_network=None,
sampling_mode=config.traversal.sampling_mode,
outcome_unsampled_regret=config.traversal.outcome_unsampled_regret,
cutoff_value_mode=config.traversal.cutoff_value_mode,
cutoff_rollouts=config.traversal.cutoff_rollouts,
cutoff_rollout_policy=config.traversal.cutoff_rollout_policy,
cutoff_rollout_max_steps=config.traversal.cutoff_rollout_max_steps,
opponent_policy=config.traversal.opponent_policy,
all_negative_fallback=config.regret_matching.all_negative_fallback,
league_advantage_networks=[],
self_play_anchor_probability=config.self_play.anchor_probability,
self_play_current_weight=config.self_play.current_weight,
self_play_recent_weight=config.self_play.recent_weight,
self_play_older_weight=config.self_play.older_weight,
self_play_anchor_weight=config.self_play.anchor_weight,
self_play_recent_window=config.self_play.recent_window,
)
interleaved_stats, interleaved_advantage, interleaved_strategy, _runtime = (
run_interleaved_traversal_batch(
networks,
game_config,
[260],
0,
1,
**common,
interleave_width=4,
interleave_max_batch=8,
)
)
assert interleaved_stats.to_dict() == recursive_stats.to_dict()
assert len(interleaved_advantage) == len(recursive_advantage)
assert len(interleaved_strategy) == len(recursive_strategy)
assert np.allclose(
[sample.target.sum() for sample in interleaved_advantage],
[sample.target.sum() for sample in recursive_advantage],
atol=1.0e-5,
)
assert np.allclose(
[sample.target.sum() for sample in interleaved_strategy],
[sample.target.sum() for sample in recursive_strategy],
atol=1.0e-6,
)
def test_deep_cfr_trainer_supports_lcfr_and_dcfr_loss_weighting() -> None:
for mode in ("lcfr", "dcfr"):
trainer = DeepCFRTrainer(
@@ -578,6 +751,7 @@ def test_deep_cfr_cython_traverser_supports_external_sampling() -> None:
"traversal": {
"traversals_per_player": 1,
"sampling_mode": "external",
"store_strategy_on_traverser_nodes": False,
"max_depth": 1,
"max_nodes_per_traversal": 64,
},