From 240ef552c647f4e7d64e07e5585a3b30c81f23d2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=A0=95=EC=8B=9C=EC=9B=90?= Date: Thu, 7 May 2026 22:49:38 +0900 Subject: [PATCH] Add non-default interleaved traversal scheduler Co-Authored-By: Codex --- docs/plans/option_b_interleaved_traversal.md | 50 +- .../games/classic/deep_cfr/config.py | 31 + .../classic/deep_cfr/interleaved_traversal.py | 579 ++++++++++++++++++ .../games/classic/deep_cfr/trainer.py | 152 +++-- .../games/classic/deep_cfr/workers.py | 111 ++-- tests/games/classic/test_deep_cfr_trainer.py | 178 +++++- 6 files changed, 1014 insertions(+), 87 deletions(-) create mode 100644 src/coolrl_lost_cities/games/classic/deep_cfr/interleaved_traversal.py diff --git a/docs/plans/option_b_interleaved_traversal.md b/docs/plans/option_b_interleaved_traversal.md index 7f3a6f5..072de20 100644 --- a/docs/plans/option_b_interleaved_traversal.md +++ b/docs/plans/option_b_interleaved_traversal.md @@ -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: 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 1f7dac2..3015cab 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/config.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/config.py @@ -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: 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 new file mode 100644 index 0000000..9d0cf75 --- /dev/null +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/interleaved_traversal.py @@ -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 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 8f0395d..254d766 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py @@ -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": diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/workers.py b/src/coolrl_lost_cities/games/classic/deep_cfr/workers.py index 4eac6db..da6fb89 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/workers.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/workers.py @@ -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, ) diff --git a/tests/games/classic/test_deep_cfr_trainer.py b/tests/games/classic/test_deep_cfr_trainer.py index ff18c53..65e9e71 100644 --- a/tests/games/classic/test_deep_cfr_trainer.py +++ b/tests/games/classic/test_deep_cfr_trainer.py @@ -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, },