from __future__ import annotations import argparse import json import random import statistics import time from pathlib import Path from typing import Any import numpy as np 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 run_cython_traversal_batch from coolrl_lost_cities.games.classic.game import GameState from coolrl_lost_cities.games.classic.deep_cfr.config import load_config from coolrl_lost_cities.games.classic.deep_cfr.networks import DeepCFRMLP def _timer(device: torch.device | None = None) -> float: if device is not None and device.type == "cuda": torch.cuda.synchronize(device) return time.perf_counter() def _median(values: list[float]) -> float: return float(statistics.median(values)) def _quantile(values: list[float], q: float) -> float: ordered = sorted(values) if not ordered: return 0.0 idx = min(len(ordered) - 1, max(0, round((len(ordered) - 1) * q))) return float(ordered[idx]) def _build_networks(cfg: Any, input_dim_value: int, action_size: int, device: torch.device): networks = [ DeepCFRMLP.from_config(input_dim_value, action_size, cfg.network).to(device).eval() for _ in range(2) ] strategy = DeepCFRMLP.from_config(input_dim_value, action_size, cfg.network).to(device).eval() return networks, strategy def _make_corpus(cfg: Any, count: int, seed: int) -> list[GameState]: rng = random.Random(seed) game_config = cfg.rules.to_lost_cities_config(seed=seed) states: list[GameState] = [] state = GameState.new_game(game_config, seed=seed) while len(states) < count: if state.terminal: state = GameState.new_game(game_config, seed=rng.randrange(1, 2**31)) continue states.append(state.clone()) legal = state.unified_legal_actions() if not legal: state = GameState.new_game(game_config, seed=rng.randrange(1, 2**31)) continue state.push_unified_action(int(legal[rng.randrange(len(legal))])) return states def _regret_matching(values: np.ndarray, legal: list[int], epsilon: float) -> np.ndarray: policy = np.zeros_like(values, dtype=np.float32) positives = [max(float(values[action]), 0.0) for action in legal] positive_sum = sum(positives) if positive_sum <= epsilon: uniform = 1.0 / max(1, len(legal)) for action in legal: policy[action] = uniform return policy for action, positive in zip(legal, positives, strict=True): policy[action] = positive / positive_sum return policy def bench_full_traversal( cfg: Any, *, traversals: int, runs: int, warmup: int, device: torch.device, input_dim_value: int, action_size: int, ) -> dict[str, Any]: timings: list[float] = [] stat_rows: list[dict[str, Any]] = [] seed_base = int(cfg.run.seed) game_config = cfg.rules.to_lost_cities_config(seed=seed_base) for run_idx in range(runs + warmup): networks, strategy = _build_networks(cfg, input_dim_value, action_size, device) seeds = [seed_base + run_idx * 100_000 + idx for idx in range(traversals)] start = _timer(device) stats, advantage_samples, strategy_samples = run_cython_traversal_batch( networks, game_config, seeds, 0, 1, device=device, action_size=action_size, strategy_network=strategy, 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=None, 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=seed_base + run_idx, ) elapsed = _timer(device) - start if run_idx >= warmup: stats_dict = stats.to_dict() stats_dict["advantage_samples_returned"] = len(advantage_samples) stats_dict["strategy_samples_returned"] = len(strategy_samples) timings.append(elapsed) stat_rows.append(stats_dict) del networks, strategy policy_calls = [int(row["traversal_regret_matching_decisions"]) for row in stat_rows] nodes = [int(row["traversal_nodes"]) for row in stat_rows] return { "seconds": timings, "median_seconds": _median(timings), "median_nodes": _median([float(value) for value in nodes]), "median_policy_calls": _median([float(value) for value in policy_calls]), "median_us_per_node": _median( [sec * 1_000_000.0 / max(1, node) for sec, node in zip(timings, nodes, strict=True)] ), "median_us_per_policy_call": _median( [ sec * 1_000_000.0 / max(1, calls) for sec, calls in zip(timings, policy_calls, strict=True) ] ), "stats": stat_rows, } def bench_encode_legal(corpus: list[GameState], cfg: Any, repeats: int) -> dict[str, Any]: durations: list[float] = [] counts: list[int] = [] total = 0.0 for _ in range(repeats): for state in corpus: start = time.perf_counter() _ = encode_info_state(state, state.current_player, cfg.encoding) legal = state.unified_legal_actions() total += len(legal) durations.append(time.perf_counter() - start) counts.append(len(legal)) return { "median_us": _median(durations) * 1_000_000.0, "p95_us": _quantile(durations, 0.95) * 1_000_000.0, "mean_legal_actions": float(statistics.mean(counts)), "checksum": total, } def bench_push_pop(corpus: list[GameState], repeats: int) -> dict[str, Any]: durations: list[float] = [] total = 0 for _ in range(repeats): for state in corpus: legal = state.unified_legal_actions() if not legal: continue action = int(legal[total % len(legal)]) start = time.perf_counter() state.push_unified_action(action) state.pop_action() durations.append(time.perf_counter() - start) total += action return { "median_us": _median(durations) * 1_000_000.0, "p95_us": _quantile(durations, 0.95) * 1_000_000.0, "checksum": total, } def bench_policy_boundary( corpus: list[GameState], cfg: Any, network: torch.nn.Module, device: torch.device, action_size: int, repeats: int, ) -> dict[str, Any]: durations: list[float] = [] checksum = 0.0 with torch.inference_mode(): for _ in range(repeats): for state in corpus: start = _timer(device) info_state = encode_info_state(state, state.current_player, cfg.encoding) legal = state.unified_legal_actions() x = torch.as_tensor(info_state, dtype=torch.float32, device=device).unsqueeze(0) advantages = network(x).squeeze(0).detach().cpu().numpy().astype(np.float32) policy = _regret_matching( advantages, legal, float(cfg.traversal.regret_matching_epsilon), ) checksum += float(policy.sum()) + float(advantages[0]) durations.append(_timer(device) - start) return { "median_us": _median(durations) * 1_000_000.0, "p95_us": _quantile(durations, 0.95) * 1_000_000.0, "checksum": checksum, } def bench_torch_forward( corpus: list[GameState], cfg: Any, network: torch.nn.Module, device: torch.device, batch_sizes: list[int], repeats: int, ) -> dict[str, Any]: states = np.stack( [encode_info_state(state, state.current_player, cfg.encoding) for state in corpus] ).astype(np.float32) results: dict[str, Any] = {} with torch.inference_mode(): for batch_size in batch_sizes: durations: list[float] = [] checksum = 0.0 for _ in range(repeats): for offset in range(0, len(states), batch_size): chunk = states[offset : offset + batch_size] if len(chunk) != batch_size: continue x = torch.as_tensor(chunk, dtype=torch.float32, device=device) start = _timer(device) out = network(x) checksum += float(out.detach().sum().cpu()) durations.append((_timer(device) - start) / batch_size) if not durations: continue results[str(batch_size)] = { "median_us_per_call": _median(durations) * 1_000_000.0, "p95_us_per_call": _quantile(durations, 0.95) * 1_000_000.0, "checksum": checksum, } return results def _print_table(result: dict[str, Any]) -> None: full = result["full_traversal"] encode = result["encode_legal"] push_pop = result["push_pop"] boundary = result["policy_boundary_bs1"] forward = result["torch_forward"] print("Traversal policy-boundary microbench") print( f"config={result['config']} device={result['device']} torch_threads={result['torch_threads']}" ) print( f"full traversal: {full['median_seconds']:.3f}s, " f"policy_calls={full['median_policy_calls']:.0f}, nodes={full['median_nodes']:.0f}" ) print() print("component median us/call p95 us/call") print( f"encode+legal {encode['median_us']:9.2f} {encode['p95_us']:9.2f}" ) print( f"push+pop {push_pop['median_us']:9.2f} {push_pop['p95_us']:9.2f}" ) print( f"policy boundary bs=1 {boundary['median_us']:9.2f} {boundary['p95_us']:9.2f}" ) for batch_size, row in forward.items(): print( f"torch forward bs={batch_size:<3} " f"{row['median_us_per_call']:9.2f} {row['p95_us_per_call']:9.2f}" ) print() print("derived") print(f"full traversal us/node {full['median_us_per_node']:.2f}") print(f"full traversal us/policy_call {full['median_us_per_policy_call']:.2f}") if boundary["median_us"] > 0: print( "full/policy-boundary ratio " f"{full['median_us_per_policy_call'] / boundary['median_us']:.2f}x" ) bs64 = forward.get("64") if bs64 and bs64["median_us_per_call"] > 0: print( "policy bs=1 vs forward bs=64 " f"{boundary['median_us'] / bs64['median_us_per_call']:.2f}x" ) def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--config", default="configs/deep_cfr/default.yaml") parser.add_argument("--device", default="cpu") parser.add_argument("--traversals", type=int, default=32) parser.add_argument("--runs", type=int, default=5) parser.add_argument("--warmup", type=int, default=1) parser.add_argument("--corpus-size", type=int, default=512) parser.add_argument("--component-repeats", type=int, default=4) parser.add_argument("--forward-repeats", type=int, default=64) parser.add_argument("--torch-threads", type=int, default=1) parser.add_argument("--output", default="") args = parser.parse_args() torch.set_num_threads(max(1, int(args.torch_threads))) cfg = load_config(args.config) device = torch.device(args.device) if device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError("--device cuda requested but CUDA is not available") probe = GameState.new_game( cfg.rules.to_lost_cities_config(seed=cfg.run.seed), seed=cfg.run.seed ) input_dim_value = input_dim(probe, cfg.encoding) action_size = 2 * probe.config.hand_size + 1 + probe.config.n_colors networks, _strategy = _build_networks(cfg, input_dim_value, action_size, device) corpus = _make_corpus(cfg, args.corpus_size, int(cfg.run.seed)) # Run a tiny warmup through the component path before timed loops. _ = encode_info_state(corpus[0], corpus[0].current_player, cfg.encoding) with torch.inference_mode(): _ = networks[0](torch.as_tensor(_, dtype=torch.float32, device=device).unsqueeze(0)) result = { "config": args.config, "device": str(device), "torch_threads": torch.get_num_threads(), "traversals": args.traversals, "runs": args.runs, "warmup": args.warmup, "corpus_size": args.corpus_size, "input_dim": input_dim_value, "action_size": action_size, "full_traversal": bench_full_traversal( cfg, traversals=args.traversals, runs=args.runs, warmup=args.warmup, device=device, input_dim_value=input_dim_value, action_size=action_size, ), "encode_legal": bench_encode_legal(corpus, cfg, args.component_repeats), "push_pop": bench_push_pop(corpus, args.component_repeats), "policy_boundary_bs1": bench_policy_boundary( corpus, cfg, networks[0], device, action_size, args.component_repeats, ), "torch_forward": bench_torch_forward( corpus, cfg, networks[0], device, [1, 4, 8, 64, 256], args.forward_repeats, ), } _print_table(result) output = Path(args.output) if args.output else Path(__file__).with_name("results.json") output.write_text(json.dumps(result, indent=2), encoding="utf-8") print(f"\nwrote {output}") if __name__ == "__main__": main()