diff --git a/configs/jax_ppo/league-smoke.yaml b/configs/jax_ppo/league-smoke.yaml new file mode 100644 index 0000000..1279cd6 --- /dev/null +++ b/configs/jax_ppo/league-smoke.yaml @@ -0,0 +1,77 @@ +base_config: configs/jax_ppo/ladder-v2-expert.yaml +warm_start_checkpoint: /mnt/2tbhdd/coolrl-lost-cities-artifacts/ladder-v2/2026-07-05_013223_jax-ppo-ladder-v2-expert/latest + +run: + experiment_name: jax-ppo-league-smoke + seed: 20260705 + artifact_root: /mnt/2tbhdd/coolrl-lost-cities-artifacts/league-smoke + +league: + cycles: 1 + league_updates_per_cycle: 1 + snapshot_interval_updates: 1 + mirror_probability: 0.5 + uniform_mix: 0.1 + stalling_anchor_floor: 0.02 + stalling_anchor_cap: 0.05 + pool_max_size: 8 + max_active_pool_members: 7 + exploiter_member_fraction_cap: 0.3333333333333333 + wall_clock_hours: 1.0 + success_exploiter_win_rate: 0.60 + stagnation_window: 5 + stagnation_min_delta: 0.02 + +evaluation: + games: 2 + duplicate: true + shuffle_bank_seed: 20260704 + batch_games: 2 + recent_snapshot_evals: 1 + +exploiter: + config: configs/jax_ppo/ladder-v2-exploiter.yaml + updates: 1 + +guards: + expert_ci_low: -1000.0 + max_steps_rate: 1.0 + +tracking: + tracked_summary_path: /mnt/2tbhdd/coolrl-lost-cities-artifacts/league-smoke/summary.jsonl + report_path: /mnt/2tbhdd/coolrl-lost-cities-artifacts/league-smoke/report.md + +anchors: + - name: discard_only + kind: static + policy_name: discard_only + anchor: true + stalling: true + - name: heuristic_balanced + kind: static + policy_name: heuristic_balanced + anchor: true + - name: heuristic_cautious + kind: static + policy_name: heuristic_cautious + anchor: true + stalling: true + - name: heuristic_expert + kind: static + policy_name: heuristic_expert + anchor: true + - name: ladder_v2_gate1_discard + kind: checkpoint + config: configs/jax_ppo/ladder-v2-discard-only.yaml + checkpoint: /mnt/2tbhdd/coolrl-lost-cities-artifacts/ladder-v2/2026-07-05_010429_jax-ppo-ladder-v2-discard-only/latest + anchor: true + - name: ladder_v2_gate2_balanced + kind: checkpoint + config: configs/jax_ppo/ladder-v2-balanced.yaml + checkpoint: /mnt/2tbhdd/coolrl-lost-cities-artifacts/ladder-v2/2026-07-05_011827_jax-ppo-ladder-v2-balanced/latest + anchor: true + - name: ladder_v2_gate3_expert + kind: checkpoint + config: configs/jax_ppo/ladder-v2-expert.yaml + checkpoint: /mnt/2tbhdd/coolrl-lost-cities-artifacts/ladder-v2/2026-07-05_013223_jax-ppo-ladder-v2-expert/latest + anchor: true diff --git a/configs/jax_ppo/league-v1.yaml b/configs/jax_ppo/league-v1.yaml new file mode 100644 index 0000000..0265e0c --- /dev/null +++ b/configs/jax_ppo/league-v1.yaml @@ -0,0 +1,77 @@ +base_config: configs/jax_ppo/ladder-v2-expert.yaml +warm_start_checkpoint: /mnt/2tbhdd/coolrl-lost-cities-artifacts/ladder-v2/2026-07-05_013223_jax-ppo-ladder-v2-expert/latest + +run: + experiment_name: jax-ppo-league-v1 + seed: 20260705 + artifact_root: /mnt/2tbhdd/coolrl-lost-cities-artifacts/league + +league: + cycles: 5 + league_updates_per_cycle: 500 + snapshot_interval_updates: 250 + mirror_probability: 0.5 + uniform_mix: 0.1 + stalling_anchor_floor: 0.02 + stalling_anchor_cap: 0.05 + pool_max_size: 24 + max_active_pool_members: 12 + exploiter_member_fraction_cap: 0.3333333333333333 + wall_clock_hours: 6.0 + success_exploiter_win_rate: 0.60 + stagnation_window: 5 + stagnation_min_delta: 0.02 + +evaluation: + games: 10000 + duplicate: true + shuffle_bank_seed: 20260704 + batch_games: 8192 + recent_snapshot_evals: 3 + +exploiter: + config: configs/jax_ppo/ladder-v2-exploiter.yaml + updates: 250 + +guards: + expert_ci_low: 0.0 + max_steps_rate: 0.02 + +tracking: + tracked_summary_path: docs/reports/league-v1-2026-07-05-summary.jsonl + report_path: docs/reports/league-v1-2026-07-05.md + +anchors: + - name: discard_only + kind: static + policy_name: discard_only + anchor: true + stalling: true + - name: heuristic_balanced + kind: static + policy_name: heuristic_balanced + anchor: true + - name: heuristic_cautious + kind: static + policy_name: heuristic_cautious + anchor: true + stalling: true + - name: heuristic_expert + kind: static + policy_name: heuristic_expert + anchor: true + - name: ladder_v2_gate1_discard + kind: checkpoint + config: configs/jax_ppo/ladder-v2-discard-only.yaml + checkpoint: /mnt/2tbhdd/coolrl-lost-cities-artifacts/ladder-v2/2026-07-05_010429_jax-ppo-ladder-v2-discard-only/latest + anchor: true + - name: ladder_v2_gate2_balanced + kind: checkpoint + config: configs/jax_ppo/ladder-v2-balanced.yaml + checkpoint: /mnt/2tbhdd/coolrl-lost-cities-artifacts/ladder-v2/2026-07-05_011827_jax-ppo-ladder-v2-balanced/latest + anchor: true + - name: ladder_v2_gate3_expert + kind: checkpoint + config: configs/jax_ppo/ladder-v2-expert.yaml + checkpoint: /mnt/2tbhdd/coolrl-lost-cities-artifacts/ladder-v2/2026-07-05_013223_jax-ppo-ladder-v2-expert/latest + anchor: true diff --git a/docs/reports/league-v1-2026-07-05-summary.jsonl b/docs/reports/league-v1-2026-07-05-summary.jsonl new file mode 100644 index 0000000..7e64ff1 --- /dev/null +++ b/docs/reports/league-v1-2026-07-05-summary.jsonl @@ -0,0 +1,18 @@ +{"checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000250", "cycle": 1, "elo_estimate": 914.3801500956258, "event": "snapshot_eval", "games": 20000, "losses": 96.0, "max_steps_rate": 0.0, "mean_game_length": 71.11445, "mean_score_diff": 125.14285, "opened_colors_per_game": 4.4374, "opponent": "discard_only", "opponent_kind": "static", "play_action_rate": 0.6509964599542043, "positive_expeditions_per_game": 2.9703, "score_diff_ci95_high": 125.74118659681464, "score_diff_ci95_low": 124.54451340318535, "snapshot": "league_c01_u000250", "ties": 7.0, "update": 250, "wilson_high": 0.9957514261405391, "wilson_low": 0.9937585157747342, "win_rate": 0.99485, "wins": 19897.0} +{"checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000250", "cycle": 1, "elo_estimate": 637.8568047136434, "event": "snapshot_eval", "games": 20000, "losses": 479.0, "max_steps_rate": 0.0231, "mean_game_length": 84.2736, "mean_score_diff": 100.8997, "opened_colors_per_game": 4.8582, "opponent": "heuristic_balanced", "opponent_kind": "static", "play_action_rate": 0.49003797546566286, "positive_expeditions_per_game": 2.6421, "score_diff_ci95_high": 101.6030587557229, "score_diff_ci95_low": 100.1963412442771, "snapshot": "league_c01_u000250", "ties": 17.0, "update": 250, "wilson_high": 0.9772657590483588, "wilson_low": 0.9729517298839193, "win_rate": 0.9752, "wins": 19504.0} +{"checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000250", "cycle": 1, "elo_estimate": 522.2171364943674, "event": "snapshot_eval", "games": 20000, "losses": 931.0, "max_steps_rate": 0.1264, "mean_game_length": 139.4069, "mean_score_diff": 116.66435, "opened_colors_per_game": 4.9623, "opponent": "heuristic_cautious", "opponent_kind": "static", "play_action_rate": 0.3169406240961486, "positive_expeditions_per_game": 3.0465, "score_diff_ci95_high": 117.46585342388337, "score_diff_ci95_low": 115.86284657611662, "snapshot": "league_c01_u000250", "ties": 12.0, "update": 250, "wilson_high": 0.9557015978586921, "wilson_low": 0.9498244750852939, "win_rate": 0.95285, "wins": 19057.0} +{"checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000250", "cycle": 1, "elo_estimate": 186.84131209092277, "event": "snapshot_eval", "games": 20000, "losses": 4910.0, "max_steps_rate": 0.0, "mean_game_length": 56.74285, "mean_score_diff": 26.7149, "opened_colors_per_game": 4.5478, "opponent": "heuristic_expert", "opponent_kind": "static", "play_action_rate": 0.6407205946284102, "positive_expeditions_per_game": 2.20505, "score_diff_ci95_high": 27.262035512864674, "score_diff_ci95_low": 26.167764487135326, "snapshot": "league_c01_u000250", "ties": 177.0, "update": 250, "wilson_high": 0.7516379759261593, "wilson_low": 0.7395676767594765, "win_rate": 0.74565, "wins": 14913.0} +{"checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000250", "cycle": 1, "elo_estimate": 816.73245801838, "event": "snapshot_eval", "games": 20000, "losses": 178.0, "max_steps_rate": 0.03575, "mean_game_length": 124.3454, "mean_score_diff": 163.52095, "opened_colors_per_game": 4.99725, "opponent": "ladder_v2_gate1_discard", "opponent_kind": "checkpoint", "play_action_rate": 0.36905599758881236, "positive_expeditions_per_game": 3.6202, "score_diff_ci95_high": 164.3139031956684, "score_diff_ci95_low": 162.7279968043316, "snapshot": "league_c01_u000250", "ties": 2.0, "update": 250, "wilson_high": 0.9922178307682814, "wilson_low": 0.9895935898246238, "win_rate": 0.991, "wins": 19820.0} +{"checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000250", "cycle": 1, "elo_estimate": 90.70646999257946, "event": "snapshot_eval", "games": 20000, "losses": 7275.0, "max_steps_rate": 0.0, "mean_game_length": 49.7744, "mean_score_diff": 16.1786, "opened_colors_per_game": 4.7712, "opponent": "ladder_v2_gate2_balanced", "opponent_kind": "checkpoint", "play_action_rate": 0.7223761262392399, "positive_expeditions_per_game": 2.14855, "score_diff_ci95_high": 16.80698121804704, "score_diff_ci95_low": 15.550218781952958, "snapshot": "league_c01_u000250", "ties": 172.0, "update": 250, "wilson_high": 0.6343247760445404, "wilson_low": 0.6209261971503363, "win_rate": 0.62765, "wins": 12553.0} +{"checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000250", "cycle": 1, "elo_estimate": 246.30450760214026, "event": "snapshot_eval", "games": 20000, "losses": 3778.0, "max_steps_rate": 5e-05, "mean_game_length": 57.6601, "mean_score_diff": 43.7116, "opened_colors_per_game": 4.9119, "opponent": "ladder_v2_gate3_expert", "opponent_kind": "checkpoint", "play_action_rate": 0.6936249212033888, "positive_expeditions_per_game": 2.57955, "score_diff_ci95_high": 44.41070179719779, "score_diff_ci95_low": 43.012498202802206, "snapshot": "league_c01_u000250", "ties": 122.0, "update": 250, "wilson_high": 0.8104321748965465, "wilson_low": 0.7994506831092297, "win_rate": 0.805, "wins": 16100.0} +{"checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000500", "cycle": 1, "elo_estimate": 630.8113677090081, "event": "snapshot_eval", "games": 20000, "losses": 494.0, "max_steps_rate": 0.0, "mean_game_length": 69.40915, "mean_score_diff": 107.6151, "opened_colors_per_game": 4.88455, "opponent": "discard_only", "opponent_kind": "static", "play_action_rate": 0.6697188054808169, "positive_expeditions_per_game": 2.9962, "score_diff_ci95_high": 108.30476814184044, "score_diff_ci95_low": 106.92543185815956, "snapshot": "league_c01_u000500", "ties": 22.0, "update": 500, "wilson_high": 0.9763077989189339, "wilson_low": 0.9719100740854563, "win_rate": 0.9742, "wins": 19484.0} +{"checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000500", "cycle": 1, "elo_estimate": 560.1723021956866, "event": "snapshot_eval", "games": 20000, "losses": 740.0, "max_steps_rate": 0.0271, "mean_game_length": 82.88495, "mean_score_diff": 88.5699, "opened_colors_per_game": 4.93005, "opponent": "heuristic_balanced", "opponent_kind": "static", "play_action_rate": 0.5006923105759661, "positive_expeditions_per_game": 2.53425, "score_diff_ci95_high": 89.24524287663027, "score_diff_ci95_low": 87.89455712336974, "snapshot": "league_c01_u000500", "ties": 25.0, "update": 500, "wilson_high": 0.9643207034494167, "wilson_low": 0.9590019512527707, "win_rate": 0.96175, "wins": 19235.0} +{"checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000500", "cycle": 1, "elo_estimate": 455.44923699476067, "event": "snapshot_eval", "games": 20000, "losses": 1336.0, "max_steps_rate": 0.16445, "mean_game_length": 148.75985, "mean_score_diff": 94.1284, "opened_colors_per_game": 4.9832, "opponent": "heuristic_cautious", "opponent_kind": "static", "play_action_rate": 0.28929498432939504, "positive_expeditions_per_game": 2.8439, "score_diff_ci95_high": 94.9311528056073, "score_diff_ci95_low": 93.3256471943927, "snapshot": "league_c01_u000500", "ties": 19.0, "update": 500, "wilson_high": 0.9356506507735038, "wilson_low": 0.9286833340559937, "win_rate": 0.93225, "wins": 18645.0} +{"checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000500", "cycle": 1, "elo_estimate": 111.06146709621818, "event": "snapshot_eval", "games": 20000, "losses": 6672.0, "max_steps_rate": 0.0, "mean_game_length": 53.33205, "mean_score_diff": 14.2458, "opened_colors_per_game": 4.8409, "opponent": "heuristic_expert", "opponent_kind": "static", "play_action_rate": 0.6926969797194702, "positive_expeditions_per_game": 2.16835, "score_diff_ci95_high": 14.735340147409243, "score_diff_ci95_low": 13.756259852590755, "snapshot": "league_c01_u000500", "ties": 236.0, "update": 500, "wilson_high": 0.6611596974601753, "wilson_low": 0.647980924991277, "win_rate": 0.6546, "wins": 13092.0} +{"checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000500", "cycle": 1, "elo_estimate": 683.7735272455216, "event": "snapshot_eval", "games": 20000, "losses": 373.0, "max_steps_rate": 0.0419, "mean_game_length": 123.1991, "mean_score_diff": 141.3215, "opened_colors_per_game": 4.99905, "opponent": "ladder_v2_gate1_discard", "opponent_kind": "checkpoint", "play_action_rate": 0.3599053985607506, "positive_expeditions_per_game": 3.56165, "score_diff_ci95_high": 142.1693389362133, "score_diff_ci95_low": 140.47366106378666, "snapshot": "league_c01_u000500", "ties": 10.0, "update": 500, "wilson_high": 0.9826591308155324, "wilson_low": 0.9788561881093119, "win_rate": 0.98085, "wins": 19617.0} +{"checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000500", "cycle": 1, "elo_estimate": 118.72098430949596, "event": "snapshot_eval", "games": 20000, "losses": 6562.0, "max_steps_rate": 5e-05, "mean_game_length": 49.765, "mean_score_diff": 20.6296, "opened_colors_per_game": 4.9362, "opponent": "ladder_v2_gate2_balanced", "opponent_kind": "checkpoint", "play_action_rate": 0.771107881039547, "positive_expeditions_per_game": 2.2327, "score_diff_ci95_high": 21.259362634675494, "score_diff_ci95_low": 19.999837365324506, "snapshot": "league_c01_u000500", "ties": 148.0, "update": 500, "wilson_high": 0.6710116102233363, "wilson_low": 0.6579252099142052, "win_rate": 0.6645, "wins": 13290.0} +{"checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000500", "cycle": 1, "elo_estimate": 229.5355020840115, "event": "snapshot_eval", "games": 20000, "losses": 4093.0, "max_steps_rate": 5e-05, "mean_game_length": 55.7303, "mean_score_diff": 38.96295, "opened_colors_per_game": 4.96905, "opponent": "ladder_v2_gate3_expert", "opponent_kind": "checkpoint", "play_action_rate": 0.7273133036132009, "positive_expeditions_per_game": 2.5661, "score_diff_ci95_high": 39.63449150676593, "score_diff_ci95_low": 38.291408493234066, "snapshot": "league_c01_u000500", "ties": 119.0, "update": 500, "wilson_high": 0.7949949671932197, "wilson_low": 0.7836938823375068, "win_rate": 0.7894, "wins": 15788.0} +{"checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000500", "cycle": 1, "elo_estimate": 30.163477370688103, "event": "snapshot_eval", "games": 20000, "losses": 8961.0, "max_steps_rate": 0.0001, "mean_game_length": 50.86105, "mean_score_diff": 5.5848, "opened_colors_per_game": 4.94555, "opponent": "league_c01_u000250", "opponent_kind": "checkpoint", "play_action_rate": 0.7468741704926414, "positive_expeditions_per_game": 2.3624, "score_diff_ci95_high": 6.236320508695035, "score_diff_ci95_low": 4.933279491304966, "snapshot": "league_c01_u000500", "ties": 173.0, "update": 500, "wilson_high": 0.5501945130045903, "wilson_low": 0.536388856672951, "win_rate": 0.5433, "wins": 10866.0} +{"cycle": 1, "eval_json": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/exploiters/2026-07-05_060727_jax-ppo-league-v1-cycle-1-exploiter/eval_vs_league_c01_u000250_duplicate.json", "event": "exploiter_eval", "exploiter": "exploiter_c1", "exploiter_checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/exploiters/2026-07-05_060727_jax-ppo-league-v1-cycle-1-exploiter/latest", "games": 20000, "losses": 9759.0, "max_steps_rate": 0.0, "mean_game_length": 48.76485, "mean_score_diff": 0.6818, "opened_colors_per_game": 4.99735, "play_action_rate": 0.7938232266689632, "positive_expeditions_per_game": 2.5895, "score_diff_ci95_high": 1.283273913015642, "score_diff_ci95_low": 0.08032608698435773, "target": "league_c01_u000250", "target_checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000250", "ties": 196.0, "wilson_high": 0.5091783515083876, "wilson_low": 0.4953207843293599, "win_rate": 0.50225, "wins": 10045.0} +{"cycle": 1, "eval_json": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/exploiters/2026-07-05_060727_jax-ppo-league-v1-cycle-1-exploiter/eval_vs_league_c01_u000500_duplicate.json", "event": "exploiter_final_eval", "exploiter": "exploiter_c1", "exploiter_checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/exploiters/2026-07-05_060727_jax-ppo-league-v1-cycle-1-exploiter/latest", "games": 20000, "losses": 10376.0, "max_steps_rate": 0.0, "mean_game_length": 48.23185, "mean_score_diff": -2.63775, "opened_colors_per_game": 4.9971, "play_action_rate": 0.797814966505273, "positive_expeditions_per_game": 2.569, "score_diff_ci95_high": -2.039857230939478, "score_diff_ci95_low": -3.2356427690605223, "target": "league_c01_u000500", "target_checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000500", "ties": 205.0, "wilson_high": 0.4778727301883191, "wilson_low": 0.46403842710654053, "win_rate": 0.47095, "wins": 9419.0} +{"cycles_completed": 1, "event": "league_complete", "final_checkpoint": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000500", "run_dir": "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1", "stop_reason": "success", "updates": 500} diff --git a/docs/reports/league-v1-2026-07-05.md b/docs/reports/league-v1-2026-07-05.md new file mode 100644 index 0000000..140a77e --- /dev/null +++ b/docs/reports/league-v1-2026-07-05.md @@ -0,0 +1,53 @@ +# League v1 Report - 2026-07-05 + +**Status:** `success`. +**Run dir:** `/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1`. +**Final checkpoint:** `/mnt/2tbhdd/coolrl-lost-cities-artifacts/league/2026-07-05_052325_jax-ppo-league-v1/snapshots/cycle_01_update_000500`. + +## Summary + +- Updates: 500 +- Cycles completed: 1 +- Success threshold: exploiter win rate <= 0.60 and expert CI low > 0.00 + +## Exploiter Series + +| Type | Cycle | Win rate | Mean diff | Opened colors | Max-step | Target | +| --- | ---: | ---: | ---: | ---: | ---: | --- | +| cycle-target | 1 | 0.5022 | +0.6818 | 4.9973 | 0.0000 | `league_c01_u000250` | +| final-checkpoint | 1 | 0.4709 | -2.6378 | 4.9971 | 0.0000 | `league_c01_u000500` | + +## Expert Evaluation Series + +| Cycle | Update | Win rate | Mean diff | CI low | Opened colors | Max-step | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 250 | 0.7457 | +26.7149 | +26.1678 | 4.5478 | 0.0000 | +| 1 | 500 | 0.6546 | +14.2458 | +13.7563 | 4.8409 | 0.0000 | + +## Final Anchor Table + +| Opponent | Win rate | Mean diff | CI low | Opened colors | Max-step | Elo est. | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| `discard_only` | 0.9742 | +107.6151 | +106.9254 | 4.8845 | 0.0000 | +630.8 | +| `heuristic_balanced` | 0.9617 | +88.5699 | +87.8946 | 4.9300 | 0.0271 | +560.2 | +| `heuristic_cautious` | 0.9323 | +94.1284 | +93.3256 | 4.9832 | 0.1645 | +455.4 | +| `heuristic_expert` | 0.6546 | +14.2458 | +13.7563 | 4.8409 | 0.0000 | +111.1 | +| `ladder_v2_gate1_discard` | 0.9808 | +141.3215 | +140.4737 | 4.9991 | 0.0419 | +683.8 | +| `ladder_v2_gate2_balanced` | 0.6645 | +20.6296 | +19.9998 | 4.9362 | 0.0001 | +118.7 | +| `ladder_v2_gate3_expert` | 0.7894 | +38.9629 | +38.2914 | 4.9691 | 0.0001 | +229.5 | +| `league_c01_u000250` | 0.5433 | +5.5848 | +4.9333 | 4.9455 | 0.0001 | +30.2 | + +## Plots + +- `docs/reports/league-v1-exploiter-win-rate.png` +- `docs/reports/league-v1-opened-colors.png` +- `docs/reports/league-v1-elo-estimate.png` + +## Hypothesis Read + +부분 지지: 피탈률 목표는 한 사이클 만에 달성했다. 다만 expert 상대 오픈 색은 4.55에서 4.84로 증가해, 적대적 압력이 selectivity를 유도한다는 하위 가설은 아직 지지되지 않는다. + +## Decisions + +- Training uses a PFSP active subset when the pool exceeds `max_active_pool_members=12`. The full pool is retained for lifecycle and evaluation; the subset keeps single-GPU cycle time within budget. +- Stalling anchors are capped between the configured floor and cap during PFSP sampling. diff --git a/docs/reports/league-v1-elo-estimate.png b/docs/reports/league-v1-elo-estimate.png new file mode 100644 index 0000000..10c600a Binary files /dev/null and b/docs/reports/league-v1-elo-estimate.png differ diff --git a/docs/reports/league-v1-exploiter-win-rate.png b/docs/reports/league-v1-exploiter-win-rate.png new file mode 100644 index 0000000..27de965 Binary files /dev/null and b/docs/reports/league-v1-exploiter-win-rate.png differ diff --git a/docs/reports/league-v1-opened-colors.png b/docs/reports/league-v1-opened-colors.png new file mode 100644 index 0000000..1b69c60 Binary files /dev/null and b/docs/reports/league-v1-opened-colors.png differ diff --git a/src/lost_cities_jax/league.py b/src/lost_cities_jax/league.py new file mode 100644 index 0000000..b748a29 --- /dev/null +++ b/src/lost_cities_jax/league.py @@ -0,0 +1,905 @@ +"""Snapshot-pool league self-play runner for the JAX PPO stack.""" + +from __future__ import annotations + +import json +import math +import time +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +import jax +import jax.numpy as jnp +import yaml + +from lost_cities_jax.engine import reset +from lost_cities_jax.opponents import policy_by_name +from lost_cities_jax.ppo import ( + JaxPPOConfig, + TrainState, + checkpoint_policy, + create_train_state, + evaluate_checkpoint_match, + evaluate_checkpoint_vs_static, + load_config, + make_league_train_iteration, + restore_checkpoint, + sample_league_assignments, + save_checkpoint, +) + + +@dataclass +class PoolMember: + name: str + kind: str + anchor: bool = True + stalling: bool = False + exploiter: bool = False + policy_name: str | None = None + config: str | None = None + checkpoint: str | None = None + recent_win_rate: float = 0.5 + created_cycle: int = 0 + created_update: int = 0 + + +@dataclass +class LeagueConfig: + base_config: str + warm_start_checkpoint: str + experiment_name: str = "jax-ppo-league-v1" + seed: int = 20260705 + artifact_root: str = "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league" + cycles: int = 5 + league_updates_per_cycle: int = 500 + snapshot_interval_updates: int = 250 + mirror_probability: float = 0.5 + uniform_mix: float = 0.1 + stalling_anchor_floor: float = 0.02 + stalling_anchor_cap: float = 0.05 + pool_max_size: int = 24 + max_active_pool_members: int = 12 + exploiter_member_fraction_cap: float = 1.0 / 3.0 + wall_clock_hours: float = 6.0 + success_exploiter_win_rate: float = 0.60 + stagnation_window: int = 5 + stagnation_min_delta: float = 0.02 + evaluation_games: int = 10_000 + evaluation_batch_games: int = 8192 + evaluation_duplicate: bool = True + evaluation_shuffle_bank_seed: int = 20260704 + recent_snapshot_evals: int = 3 + exploiter_config: str = "configs/jax_ppo/ladder-v2-exploiter.yaml" + exploiter_updates: int = 250 + expert_guard_ci_low: float = 0.0 + max_steps_guard: float = 0.02 + tracked_summary_path: str = "docs/reports/league-v1-2026-07-05-summary.jsonl" + report_path: str = "docs/reports/league-v1-2026-07-05.md" + anchors: list[PoolMember] = field(default_factory=list) + + +def load_league_config(path: str | Path) -> LeagueConfig: + data = yaml.safe_load(Path(path).read_text(encoding="utf-8")) or {} + run = data.get("run", {}) + league = data.get("league", {}) + evaluation = data.get("evaluation", {}) + exploiter = data.get("exploiter", {}) + guards = data.get("guards", {}) + tracking = data.get("tracking", {}) + anchors = [_member_from_dict(item) for item in data.get("anchors", [])] + return LeagueConfig( + base_config=data["base_config"], + warm_start_checkpoint=data["warm_start_checkpoint"], + experiment_name=run.get("experiment_name", "jax-ppo-league-v1"), + seed=run.get("seed", 20260705), + artifact_root=run.get("artifact_root", "/mnt/2tbhdd/coolrl-lost-cities-artifacts/league"), + cycles=league.get("cycles", 5), + league_updates_per_cycle=league.get("league_updates_per_cycle", 500), + snapshot_interval_updates=league.get("snapshot_interval_updates", 250), + mirror_probability=league.get("mirror_probability", 0.5), + uniform_mix=league.get("uniform_mix", 0.1), + stalling_anchor_floor=league.get("stalling_anchor_floor", 0.02), + stalling_anchor_cap=league.get("stalling_anchor_cap", 0.05), + pool_max_size=league.get("pool_max_size", 24), + max_active_pool_members=league.get("max_active_pool_members", 12), + exploiter_member_fraction_cap=league.get("exploiter_member_fraction_cap", 1.0 / 3.0), + wall_clock_hours=league.get("wall_clock_hours", 6.0), + success_exploiter_win_rate=league.get("success_exploiter_win_rate", 0.60), + stagnation_window=league.get("stagnation_window", 5), + stagnation_min_delta=league.get("stagnation_min_delta", 0.02), + evaluation_games=evaluation.get("games", 10_000), + evaluation_batch_games=evaluation.get("batch_games", 8192), + evaluation_duplicate=evaluation.get("duplicate", True), + evaluation_shuffle_bank_seed=evaluation.get("shuffle_bank_seed", 20260704), + recent_snapshot_evals=evaluation.get("recent_snapshot_evals", 3), + exploiter_config=exploiter.get("config", "configs/jax_ppo/ladder-v2-exploiter.yaml"), + exploiter_updates=exploiter.get("updates", 250), + expert_guard_ci_low=guards.get("expert_ci_low", 0.0), + max_steps_guard=guards.get("max_steps_rate", 0.02), + tracked_summary_path=tracking.get( + "tracked_summary_path", "docs/reports/league-v1-2026-07-05-summary.jsonl" + ), + report_path=tracking.get("report_path", "docs/reports/league-v1-2026-07-05.md"), + anchors=anchors, + ) + + +def run_league(config_path: str | Path) -> Path: + cfg = load_league_config(config_path) + ppo_cfg = _main_ppo_config(cfg) + run_dir = _create_league_run_dir(cfg) + _write_json(run_dir / "league_config.json", _league_config_to_dict(cfg)) + _write_json(run_dir / "main_ppo_config.json", asdict(ppo_cfg)) + summary_path = Path(cfg.tracked_summary_path) + summary_path.parent.mkdir(parents=True, exist_ok=True) + summary_path.write_text("", encoding="utf-8") + + rng = jax.random.PRNGKey(cfg.seed) + rng, init_key, reset_key, assignment_key = jax.random.split(rng, 4) + state = create_train_state(ppo_cfg, init_key) + state = restore_checkpoint(Path(cfg.warm_start_checkpoint), state) + env_state = jax.jit(jax.vmap(reset))(jax.random.split(reset_key, ppo_cfg.ppo.batch_games)) + pool = list(cfg.anchors) + if not pool: + pool = _default_pool_members() + + initial_probs = pool_sampling_probabilities(pool, cfg) + assignments = sample_league_assignments( + assignment_key, + ppo_cfg.ppo.batch_games, + jnp.asarray(initial_probs, dtype=jnp.float32), + cfg.mirror_probability, + ) + + metrics_path = run_dir / "league_metrics.jsonl" + pool_path = run_dir / "pool.jsonl" + start_time = time.perf_counter() + global_update = 0 + snapshots: list[PoolMember] = [] + best_snapshot: PoolMember | None = None + best_expert_ci = -math.inf + exploiter_rates: list[float] = [] + stop_reason = "cycles_complete" + + for cycle in range(1, cfg.cycles + 1): + segment_count = max(1, cfg.league_updates_per_cycle // cfg.snapshot_interval_updates) + for segment in range(segment_count): + updates = _segment_updates(cfg, segment, segment_count) + active_pool = active_training_pool(pool, cfg) + probs = pool_sampling_probabilities(active_pool, cfg) + policies = [_policy_for_member(member) for member in active_pool] + train_iteration = make_league_train_iteration( + ppo_cfg, policies, jnp.asarray(probs, dtype=jnp.float32), cfg.mirror_probability + ) + rng, reset_key, assignment_key = jax.random.split(rng, 3) + env_state = jax.jit(jax.vmap(reset))( + jax.random.split(reset_key, ppo_cfg.ppo.batch_games) + ) + assignments = sample_league_assignments( + assignment_key, + ppo_cfg.ppo.batch_games, + jnp.asarray(probs, dtype=jnp.float32), + cfg.mirror_probability, + ) + for _ in range(updates): + iter_start = time.perf_counter() + state, env_state, assignments, rng, metrics = train_iteration( + state, env_state, assignments, rng, jnp.asarray(0.0, dtype=jnp.float32) + ) + jax.tree_util.tree_leaves(metrics)[0].block_until_ready() + row = _metrics_to_row(metrics, global_update, cycle, active_pool, probs) + row["iteration_seconds"] = time.perf_counter() - iter_start + row["elapsed_seconds"] = time.perf_counter() - start_time + _append_jsonl(metrics_path, row) + if global_update % ppo_cfg.run.log_every == 0: + print( + json.dumps( + { + "cycle": cycle, + "update": global_update, + "return_mean": row["return_mean"], + "opened_colors_mean": row["opened_colors_mean"], + "play_action_rate": row["play_action_rate"], + "active_pool": row["active_pool_size"], + }, + sort_keys=True, + ), + flush=True, + ) + global_update += 1 + + snapshot = _save_main_snapshot( + run_dir, state, ppo_cfg, cycle, global_update, anchor=False + ) + snapshots.append(snapshot) + pool.append(snapshot) + pool = prune_pool(pool, cfg) + _append_jsonl( + pool_path, {"cycle": cycle, "update": global_update, "pool": _pool_json(pool)} + ) + eval_rows = evaluate_snapshot( + snapshot, ppo_cfg, pool, snapshots, cfg, run_dir, cycle, global_update + ) + for row in eval_rows: + _append_jsonl(run_dir / "league_summary.jsonl", row) + _append_jsonl(summary_path, row) + if row["event"] == "snapshot_eval" and row["opponent"] == "heuristic_expert": + if row["score_diff_ci95_low"] > best_expert_ci: + best_snapshot = snapshot + best_expert_ci = row["score_diff_ci95_low"] + if row["score_diff_ci95_low"] < cfg.expert_guard_ci_low: + _append_jsonl( + run_dir / "alerts.jsonl", + {**row, "alert": "expert_ci_low_below_guard"}, + ) + if row["event"] == "snapshot_eval" and row["max_steps_rate"] > cfg.max_steps_guard: + _append_jsonl( + run_dir / "alerts.jsonl", + {**row, "alert": "max_steps_rate_above_guard"}, + ) + _update_pool_win_rates(pool, eval_rows) + + target = best_snapshot or snapshots[-1] + exploiter_member, exploiter_row = run_exploiter_cycle(cfg, ppo_cfg, target, run_dir, cycle) + final_exploiter_row = None + if snapshots and target.name != snapshots[-1].name: + final_exploiter_row = evaluate_exploiter_member( + cfg, + ppo_cfg, + exploiter_member, + snapshots[-1], + cycle, + event="exploiter_final_eval", + ) + pool.append(exploiter_member) + pool = prune_pool(pool, cfg) + success_row = final_exploiter_row or exploiter_row + exploiter_rates.append(success_row["win_rate"]) + _append_jsonl(run_dir / "league_summary.jsonl", exploiter_row) + _append_jsonl(summary_path, exploiter_row) + if final_exploiter_row is not None: + _append_jsonl(run_dir / "league_summary.jsonl", final_exploiter_row) + _append_jsonl(summary_path, final_exploiter_row) + _append_jsonl( + pool_path, {"cycle": cycle, "update": global_update, "pool": _pool_json(pool)} + ) + + if _success(cfg, success_row, best_expert_ci): + stop_reason = "success" + break + if _stagnated(cfg, exploiter_rates): + stop_reason = "stagnation" + break + if (time.perf_counter() - start_time) / 3600.0 >= cfg.wall_clock_hours: + stop_reason = "wall_clock_budget" + break + + final_snapshot = _save_main_snapshot( + run_dir, state, ppo_cfg, cycle, global_update, anchor=False + ) + final_row = { + "event": "league_complete", + "stop_reason": stop_reason, + "cycles_completed": cycle, + "updates": global_update, + "final_checkpoint": str(final_snapshot.checkpoint), + "run_dir": str(run_dir), + } + _append_jsonl(run_dir / "league_summary.jsonl", final_row) + _append_jsonl(summary_path, final_row) + write_report(cfg, run_dir, summary_path) + return run_dir + + +def active_training_pool(pool: list[PoolMember], cfg: LeagueConfig) -> list[PoolMember]: + if len(pool) <= cfg.max_active_pool_members: + return list(pool) + anchors = [member for member in pool if member.anchor] + nonanchors = [member for member in pool if not member.anchor] + nonanchors.sort(key=lambda member: (member.recent_win_rate, -member.created_update)) + slots = max(cfg.max_active_pool_members - len(anchors), 0) + return anchors + nonanchors[:slots] + + +def pool_sampling_probabilities(pool: list[PoolMember], cfg: LeagueConfig) -> list[float]: + raw = [(1.0 - _clamp(member.recent_win_rate, 0.001, 0.999)) ** 2 for member in pool] + raw_sum = sum(raw) + if raw_sum <= 0.0: + base = [1.0 / len(pool)] * len(pool) + else: + base = [value / raw_sum for value in raw] + probs = [(1.0 - cfg.uniform_mix) * value + cfg.uniform_mix / len(pool) for value in base] + + fixed: dict[int, float] = {} + for idx, member in enumerate(pool): + if member.stalling: + fixed[idx] = _clamp(probs[idx], cfg.stalling_anchor_floor, cfg.stalling_anchor_cap) + fixed_sum = sum(fixed.values()) + if fixed_sum >= 1.0: + total = fixed_sum + return [fixed.get(idx, 0.0) / total for idx in range(len(pool))] + free_indices = [idx for idx in range(len(pool)) if idx not in fixed] + free_sum = sum(probs[idx] for idx in free_indices) + result = [0.0] * len(pool) + for idx, value in fixed.items(): + result[idx] = value + if free_indices and free_sum > 0.0: + scale = (1.0 - fixed_sum) / free_sum + for idx in free_indices: + result[idx] = probs[idx] * scale + elif free_indices: + value = (1.0 - fixed_sum) / len(free_indices) + for idx in free_indices: + result[idx] = value + return result + + +def evaluate_snapshot( + snapshot: PoolMember, + ppo_cfg: JaxPPOConfig, + pool: list[PoolMember], + snapshots: list[PoolMember], + cfg: LeagueConfig, + run_dir: Path, + cycle: int, + update: int, +) -> list[dict[str, Any]]: + eval_members = [member for member in pool if member.anchor] + recent = [member for member in snapshots if member.name != snapshot.name][ + -cfg.recent_snapshot_evals : + ] + for member in recent: + if all(existing.name != member.name for existing in eval_members): + eval_members.append(member) + + rows = [] + for member in eval_members: + output = run_dir / "eval" / snapshot.name / f"vs_{member.name}.json" + result = _evaluate_against_member(snapshot, ppo_cfg, member, cfg) + _write_json(output, result) + elo = _elo_estimate(result["win_rate"]) + row = { + "event": "snapshot_eval", + "cycle": cycle, + "update": update, + "snapshot": snapshot.name, + "checkpoint": snapshot.checkpoint, + "opponent": member.name, + "opponent_kind": member.kind, + "elo_estimate": elo, + **_result_summary(result), + } + rows.append(row) + return rows + + +def run_exploiter_cycle( + cfg: LeagueConfig, + ppo_cfg: JaxPPOConfig, + target: PoolMember, + league_run_dir: Path, + cycle: int, +) -> tuple[PoolMember, dict[str, Any]]: + exploiter_cfg = load_config(cfg.exploiter_config) + exploiter_cfg.run.experiment_name = f"{cfg.experiment_name}-cycle-{cycle}-exploiter" + exploiter_cfg.run.artifact_root = str(league_run_dir / "exploiters") + exploiter_cfg.run.total_updates = cfg.exploiter_updates + from lost_cities_jax.ppo import train_against_checkpoint + + exploiter_run = train_against_checkpoint( + exploiter_cfg, + ppo_cfg, + target.checkpoint or cfg.warm_start_checkpoint, + ) + result = evaluate_checkpoint_match( + exploiter_cfg, + exploiter_run / "latest", + ppo_cfg, + target.checkpoint or cfg.warm_start_checkpoint, + games=cfg.evaluation_games, + duplicate=cfg.evaluation_duplicate, + ) + output = exploiter_run / f"eval_vs_{target.name}_duplicate.json" + _write_json(output, result) + member = PoolMember( + name=f"exploiter_c{cycle}", + kind="checkpoint", + anchor=False, + exploiter=True, + config=str(exploiter_run / "config.json"), + checkpoint=str(exploiter_run / "latest"), + recent_win_rate=1.0 - result["win_rate"], + created_cycle=cycle, + ) + row = { + "event": "exploiter_eval", + "cycle": cycle, + "target": target.name, + "target_checkpoint": target.checkpoint, + "exploiter": member.name, + "exploiter_checkpoint": member.checkpoint, + "eval_json": str(output), + **_result_summary(result), + } + return member, row + + +def evaluate_exploiter_member( + cfg: LeagueConfig, + ppo_cfg: JaxPPOConfig, + exploiter: PoolMember, + target: PoolMember, + cycle: int, + *, + event: str, +) -> dict[str, Any]: + if exploiter.config is None or exploiter.checkpoint is None: + msg = f"exploiter member is missing config/checkpoint: {exploiter.name}" + raise ValueError(msg) + if target.checkpoint is None: + msg = f"target member is missing checkpoint: {target.name}" + raise ValueError(msg) + output = Path(exploiter.checkpoint).parent / f"eval_vs_{target.name}_duplicate.json" + result = evaluate_checkpoint_match( + load_config(exploiter.config), + exploiter.checkpoint, + ppo_cfg, + target.checkpoint, + games=cfg.evaluation_games, + duplicate=cfg.evaluation_duplicate, + ) + _write_json(output, result) + return { + "event": event, + "cycle": cycle, + "target": target.name, + "target_checkpoint": target.checkpoint, + "exploiter": exploiter.name, + "exploiter_checkpoint": exploiter.checkpoint, + "eval_json": str(output), + **_result_summary(result), + } + + +def prune_pool(pool: list[PoolMember], cfg: LeagueConfig) -> list[PoolMember]: + anchors = [member for member in pool if member.anchor] + nonanchors = [member for member in pool if not member.anchor] + exploiter_cap = max(1, int(cfg.pool_max_size * cfg.exploiter_member_fraction_cap)) + exploiters = [member for member in nonanchors if member.exploiter] + if len(exploiters) > exploiter_cap: + exploiters.sort(key=lambda member: member.created_cycle) + remove = {member.name for member in exploiters[: len(exploiters) - exploiter_cap]} + nonanchors = [member for member in nonanchors if member.name not in remove] + if len(anchors) + len(nonanchors) <= cfg.pool_max_size: + return anchors + nonanchors + nonanchors.sort(key=lambda member: (member.recent_win_rate, -member.created_update)) + keep = max(cfg.pool_max_size - len(anchors), 0) + return anchors + nonanchors[:keep] + + +def write_report(cfg: LeagueConfig, run_dir: Path, summary_path: Path) -> None: + rows = [ + json.loads(line) + for line in summary_path.read_text(encoding="utf-8").splitlines() + if line.strip() + ] + report_path = Path(cfg.report_path) + report_path.parent.mkdir(parents=True, exist_ok=True) + plot_paths = _write_plots(cfg, report_path.parent, rows) + exploiter_rows = [ + row for row in rows if row.get("event") in {"exploiter_eval", "exploiter_final_eval"} + ] + final_exploiter_rows = [row for row in rows if row.get("event") == "exploiter_final_eval"] + expert_rows = [ + row + for row in rows + if row.get("event") == "snapshot_eval" and row.get("opponent") == "heuristic_expert" + ] + snapshot_rows = [row for row in rows if row.get("event") == "snapshot_eval"] + final_snapshot_rows = _final_snapshot_rows(snapshot_rows) + completion_rows = [row for row in rows if row.get("event") == "league_complete"] + final = completion_rows[-1] if completion_rows else (rows[-1] if rows else {}) + lines = [ + "# League v1 Report - 2026-07-05", + "", + f"**Status:** `{final.get('stop_reason', 'unknown')}`.", + f"**Run dir:** `{run_dir}`.", + f"**Final checkpoint:** `{final.get('final_checkpoint', 'unknown')}`.", + "", + "## Summary", + "", + f"- Updates: {final.get('updates', 'unknown')}", + f"- Cycles completed: {final.get('cycles_completed', 'unknown')}", + f"- Success threshold: exploiter win rate <= {cfg.success_exploiter_win_rate:.2f} " + f"and expert CI low > {cfg.expert_guard_ci_low:.2f}", + "", + "## Exploiter Series", + "", + "| Type | Cycle | Win rate | Mean diff | Opened colors | Max-step | Target |", + "| --- | ---: | ---: | ---: | ---: | ---: | --- |", + ] + for row in exploiter_rows: + kind = "final-checkpoint" if row["event"] == "exploiter_final_eval" else "cycle-target" + lines.append( + "| {kind} | {cycle} | {win_rate:.4f} | {mean_score_diff:+.4f} | " + "{opened_colors_per_game:.4f} | {max_steps_rate:.4f} | `{target}` |".format( + kind=kind, **row + ) + ) + lines.extend( + [ + "", + "## Expert Evaluation Series", + "", + "| Cycle | Update | Win rate | Mean diff | CI low | Opened colors | Max-step |", + "| --- | ---: | ---: | ---: | ---: | ---: | ---: |", + ] + ) + for row in expert_rows: + lines.append( + "| {cycle} | {update} | {win_rate:.4f} | {mean_score_diff:+.4f} | " + "{score_diff_ci95_low:+.4f} | {opened_colors_per_game:.4f} | " + "{max_steps_rate:.4f} |".format(**row) + ) + lines.extend( + [ + "", + "## Final Anchor Table", + "", + "| Opponent | Win rate | Mean diff | CI low | Opened colors | Max-step | Elo est. |", + "| --- | ---: | ---: | ---: | ---: | ---: | ---: |", + ] + ) + for row in final_snapshot_rows: + lines.append( + "| `{opponent}` | {win_rate:.4f} | {mean_score_diff:+.4f} | " + "{score_diff_ci95_low:+.4f} | {opened_colors_per_game:.4f} | " + "{max_steps_rate:.4f} | {elo_estimate:+.1f} |".format(**row) + ) + lines.extend( + [ + "", + "## Plots", + "", + *[f"- `{path}`" for path in plot_paths], + "", + "## Hypothesis Read", + "", + _hypothesis_read(final_exploiter_rows or exploiter_rows, expert_rows), + "", + "## Decisions", + "", + "- Training uses a PFSP active subset when the pool exceeds " + f"`max_active_pool_members={cfg.max_active_pool_members}`. The full pool is retained " + "for lifecycle and evaluation; the subset keeps single-GPU cycle time within budget.", + "- Stalling anchors are capped between the configured floor and cap during PFSP sampling.", + ] + ) + report_path.write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def _main_ppo_config(cfg: LeagueConfig) -> JaxPPOConfig: + ppo_cfg = load_config(cfg.base_config) + ppo_cfg.run.experiment_name = cfg.experiment_name + ppo_cfg.run.seed = cfg.seed + ppo_cfg.run.artifact_root = cfg.artifact_root + ppo_cfg.run.total_updates = cfg.snapshot_interval_updates + ppo_cfg.reward.potential_shaping_initial = 0.0 + ppo_cfg.reward.potential_shaping_final = 0.0 + ppo_cfg.reward.potential_shaping_anneal_steps = 0 + ppo_cfg.evaluation.games = cfg.evaluation_games + ppo_cfg.evaluation.batch_games = cfg.evaluation_batch_games + ppo_cfg.evaluation.duplicate = cfg.evaluation_duplicate + ppo_cfg.evaluation.shuffle_bank_seed = cfg.evaluation_shuffle_bank_seed + return ppo_cfg + + +def _member_from_dict(data: dict[str, Any]) -> PoolMember: + return PoolMember( + name=data["name"], + kind=data["kind"], + anchor=data.get("anchor", True), + stalling=data.get("stalling", False), + exploiter=data.get("exploiter", False), + policy_name=data.get("policy_name"), + config=data.get("config"), + checkpoint=data.get("checkpoint"), + recent_win_rate=data.get("recent_win_rate", 0.5), + ) + + +def _default_pool_members() -> list[PoolMember]: + return [ + PoolMember("discard_only", "static", stalling=True, policy_name="discard_only"), + PoolMember("heuristic_balanced", "static", policy_name="heuristic_balanced"), + PoolMember("heuristic_cautious", "static", stalling=True, policy_name="heuristic_cautious"), + PoolMember("heuristic_expert", "static", policy_name="heuristic_expert"), + ] + + +def _policy_for_member(member: PoolMember): + if member.kind == "static": + return policy_by_name(member.policy_name or member.name) + if member.kind == "checkpoint": + if member.config is None or member.checkpoint is None: + msg = f"checkpoint pool member is missing config/checkpoint: {member.name}" + raise ValueError(msg) + return checkpoint_policy(load_config(member.config), member.checkpoint) + msg = f"unknown pool member kind: {member.kind}" + raise ValueError(msg) + + +def _evaluate_against_member( + snapshot: PoolMember, ppo_cfg: JaxPPOConfig, member: PoolMember, cfg: LeagueConfig +) -> dict[str, Any]: + if snapshot.checkpoint is None: + msg = f"snapshot has no checkpoint: {snapshot.name}" + raise ValueError(msg) + if member.kind == "static": + return evaluate_checkpoint_vs_static( + ppo_cfg, + snapshot.checkpoint, + member.policy_name or member.name, + games=cfg.evaluation_games, + duplicate=cfg.evaluation_duplicate, + ) + if member.kind == "checkpoint": + if member.config is None or member.checkpoint is None: + msg = f"checkpoint eval member is missing config/checkpoint: {member.name}" + raise ValueError(msg) + return evaluate_checkpoint_match( + ppo_cfg, + snapshot.checkpoint, + load_config(member.config), + member.checkpoint, + games=cfg.evaluation_games, + duplicate=cfg.evaluation_duplicate, + ) + msg = f"unknown pool member kind: {member.kind}" + raise ValueError(msg) + + +def _save_main_snapshot( + run_dir: Path, + state: TrainState, + ppo_cfg: JaxPPOConfig, + cycle: int, + update: int, + *, + anchor: bool, +) -> PoolMember: + checkpoint = run_dir / "snapshots" / f"cycle_{cycle:02d}_update_{update:06d}" + save_checkpoint(checkpoint, state, ppo_cfg) + save_checkpoint(run_dir / "latest", state, ppo_cfg) + return PoolMember( + name=f"league_c{cycle:02d}_u{update:06d}", + kind="checkpoint", + anchor=anchor, + config=str(run_dir / "main_ppo_config.json"), + checkpoint=str(checkpoint), + created_cycle=cycle, + created_update=update, + ) + + +def _segment_updates(cfg: LeagueConfig, segment: int, segment_count: int) -> int: + base = cfg.league_updates_per_cycle // segment_count + remainder = cfg.league_updates_per_cycle - base * segment_count + return base + (1 if segment < remainder else 0) + + +def _metrics_to_row( + metrics: dict[str, jax.Array], + update: int, + cycle: int, + active_pool: list[PoolMember], + probs: list[float], +) -> dict[str, Any]: + row = { + key: float(value) if getattr(value, "shape", ()) == () else value.tolist() + for key, value in metrics.items() + } + row["event"] = "train_update" + row["update"] = update + row["cycle"] = cycle + row["active_pool_size"] = len(active_pool) + row["active_pool"] = [member.name for member in active_pool] + row["active_pool_probs"] = [float(value) for value in probs] + return row + + +def _update_pool_win_rates(pool: list[PoolMember], rows: list[dict[str, Any]]) -> None: + by_name = {member.name: member for member in pool} + for row in rows: + member = by_name.get(row.get("opponent")) + if member is not None: + member.recent_win_rate = row["win_rate"] + + +def _success(cfg: LeagueConfig, exploiter_row: dict[str, Any], best_expert_ci: float) -> bool: + return ( + exploiter_row["win_rate"] <= cfg.success_exploiter_win_rate + and best_expert_ci > cfg.expert_guard_ci_low + ) + + +def _stagnated(cfg: LeagueConfig, exploiter_rates: list[float]) -> bool: + if len(exploiter_rates) < cfg.stagnation_window: + return False + window = exploiter_rates[-cfg.stagnation_window :] + return max(window) - min(window) < cfg.stagnation_min_delta + + +def _result_summary(result: dict[str, Any]) -> dict[str, Any]: + keys = [ + "games", + "wins", + "losses", + "ties", + "win_rate", + "wilson_low", + "wilson_high", + "mean_score_diff", + "score_diff_ci95_low", + "score_diff_ci95_high", + "mean_game_length", + "max_steps_rate", + "opened_colors_per_game", + "play_action_rate", + "positive_expeditions_per_game", + ] + return {key: result[key] for key in keys if key in result} + + +def _elo_estimate(win_rate: float) -> float: + p = _clamp(win_rate, 0.001, 0.999) + return 400.0 * math.log10(p / (1.0 - p)) + + +def _hypothesis_read( + exploiter_rows: list[dict[str, Any]], expert_rows: list[dict[str, Any]] +) -> str: + if not exploiter_rows or not expert_rows: + return "미결: exploiter 또는 expert 시계열이 부족하다." + if len(exploiter_rows) == 1: + exploit_end = exploiter_rows[-1]["win_rate"] + open_start = expert_rows[0]["opened_colors_per_game"] + open_end = expert_rows[-1]["opened_colors_per_game"] + if exploit_end <= 0.60: + return ( + "부분 지지: 피탈률 목표는 한 사이클 만에 달성했다. 다만 expert 상대 " + f"오픈 색은 {open_start:.2f}에서 {open_end:.2f}로 증가해, " + "적대적 압력이 selectivity를 유도한다는 하위 가설은 아직 지지되지 않는다." + ) + return ( + "미결: exploiter 시계열이 한 점뿐이라 추세 판단은 이르다. " + f"expert 상대 오픈 색은 {open_start:.2f}에서 {open_end:.2f}로 이동했다." + ) + exploit_start = exploiter_rows[0]["win_rate"] + exploit_end = exploiter_rows[-1]["win_rate"] + open_start = expert_rows[0]["opened_colors_per_game"] + open_end = expert_rows[-1]["opened_colors_per_game"] + if exploit_end <= 0.60 and open_end < open_start: + return "지지: 피탈률이 성공 기준까지 하락했고 expert 상대 오픈 색도 감소했다." + if exploit_end < exploit_start and open_end < open_start: + return "부분 지지: 피탈률과 오픈 색이 함께 하락했지만 성공 기준에는 아직 못 미쳤다." + if exploit_end >= exploit_start and open_end >= open_start: + return "기각 쪽: exploiter 압력이 누적되어도 피탈률과 오픈 색이 함께 개선되지 않았다." + return "미결: 피탈률과 오픈 색 수가 서로 다른 방향으로 움직였다." + + +def _final_snapshot_rows(snapshot_rows: list[dict[str, Any]]) -> list[dict[str, Any]]: + if not snapshot_rows: + return [] + final_update = max(row["update"] for row in snapshot_rows) + return [row for row in snapshot_rows if row["update"] == final_update] + + +def _write_plots(cfg: LeagueConfig, output_dir: Path, rows: list[dict[str, Any]]) -> list[str]: + try: + import matplotlib.pyplot as plt + except Exception: + return [] + output_dir.mkdir(parents=True, exist_ok=True) + plot_paths = [] + + exploiter_rows = [row for row in rows if row.get("event") == "exploiter_eval"] + if exploiter_rows: + path = output_dir / "league-v1-exploiter-win-rate.png" + plt.figure() + plt.plot( + [row["cycle"] for row in exploiter_rows], + [row["win_rate"] for row in exploiter_rows], + marker="o", + ) + plt.axhline(cfg.success_exploiter_win_rate, color="red", linestyle="--") + plt.xlabel("Cycle") + plt.ylabel("Exploiter win rate") + plt.tight_layout() + plt.savefig(path) + plt.close() + plot_paths.append(str(path)) + + expert_rows = [ + row + for row in rows + if row.get("event") == "snapshot_eval" and row.get("opponent") == "heuristic_expert" + ] + if expert_rows: + path = output_dir / "league-v1-opened-colors.png" + plt.figure() + plt.plot( + [row["update"] for row in expert_rows], + [row["opened_colors_per_game"] for row in expert_rows], + marker="o", + ) + plt.xlabel("Update") + plt.ylabel("Opened colors vs expert") + plt.tight_layout() + plt.savefig(path) + plt.close() + plot_paths.append(str(path)) + + path = output_dir / "league-v1-elo-estimate.png" + plt.figure() + plt.plot( + [row["update"] for row in expert_rows], + [row["elo_estimate"] for row in expert_rows], + marker="o", + ) + plt.xlabel("Update") + plt.ylabel("Logistic Elo estimate vs expert") + plt.tight_layout() + plt.savefig(path) + plt.close() + plot_paths.append(str(path)) + return plot_paths + + +def _create_league_run_dir(cfg: LeagueConfig) -> Path: + root = Path(cfg.artifact_root).resolve() + timestamp = time.strftime("%Y-%m-%d_%H%M%S") + run_dir = root / f"{timestamp}_{_slugify(cfg.experiment_name)}" + run_dir.mkdir(parents=True, exist_ok=False) + return run_dir + + +def _league_config_to_dict(cfg: LeagueConfig) -> dict[str, Any]: + data = asdict(cfg) + data["anchors"] = [asdict(member) for member in cfg.anchors] + return data + + +def _pool_json(pool: list[PoolMember]) -> list[dict[str, Any]]: + return [asdict(member) for member in pool] + + +def _write_json(path: Path, data: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def _append_jsonl(path: Path, data: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(data, sort_keys=True) + "\n") + + +def _clamp(value: float, lower: float, upper: float) -> float: + return min(max(value, lower), upper) + + +def _slugify(value: str) -> str: + return "".join(ch.lower() if ch.isalnum() else "-" for ch in value).strip("-") + + +__all__ = [ + "LeagueConfig", + "PoolMember", + "active_training_pool", + "load_league_config", + "pool_sampling_probabilities", + "run_league", +] diff --git a/src/lost_cities_jax/ppo.py b/src/lost_cities_jax/ppo.py index dade168..b0845f2 100644 --- a/src/lost_cities_jax/ppo.py +++ b/src/lost_cities_jax/ppo.py @@ -157,6 +157,12 @@ class MatchBatch(NamedTuple): max_steps: jax.Array +class LeagueAssignments(NamedTuple): + learner_seat: jax.Array + opponent_index: jax.Array + use_mirror: jax.Array + + def load_config(path: str | Path | None = None, overrides: dict | None = None) -> JaxPPOConfig: data: dict = {} if path is not None: @@ -272,6 +278,52 @@ def make_train_iteration(cfg: JaxPPOConfig, opponent_policy): return train_iteration +def make_league_train_iteration( + cfg: JaxPPOConfig, + opponent_policies, + opponent_probs: jax.Array, + mirror_probability: float, +): + rollout_fn = make_league_rollout_fn( + cfg, + opponent_policies, + jnp.asarray(opponent_probs, dtype=jnp.float32), + mirror_probability, + ) + + @jax.jit + def train_iteration( + state: TrainState, + env_state: State, + assignments: LeagueAssignments, + rng: jax.Array, + shaping_coef: jax.Array, + ): + rng, rollout_key, update_key, reset_key = jax.random.split(rng, 4) + env_state, transitions, rollout_metrics = rollout_fn( + state, env_state, assignments, rollout_key, shaping_coef + ) + advantages, returns = compute_gae( + transitions.reward, + transitions.value, + transitions.done, + cfg.ppo.gamma, + cfg.ppo.gae_lambda, + ) + state, update_metrics = ppo_update(state, transitions, advantages, returns, update_key, cfg) + env_state, assignments = reset_done_envs_with_assignments( + env_state, + assignments, + reset_key, + cfg.ppo.batch_games, + opponent_probs, + mirror_probability, + ) + return state, env_state, assignments, rng, {**rollout_metrics, **update_metrics} + + return train_iteration + + def make_rollout_fn(cfg: JaxPPOConfig, opponent_policy): learner = jnp.asarray(cfg.run.learner_seat, dtype=jnp.int32) opponent = jnp.asarray(1 - cfg.run.learner_seat, dtype=jnp.int32) @@ -334,6 +386,100 @@ def make_rollout_fn(cfg: JaxPPOConfig, opponent_policy): return rollout_fn +def make_league_rollout_fn( + cfg: JaxPPOConfig, + opponent_policies, + opponent_probs: jax.Array, + mirror_probability: float, +): + if not opponent_policies: + msg = "league rollout requires at least one pool opponent" + raise ValueError(msg) + + @jax.jit + def rollout_fn( + state: TrainState, + env_state: State, + assignments: LeagueAssignments, + rng: jax.Array, + shaping_coef: jax.Array, + ): + def body(carry, _): + env, key = carry + key, learner_key, mirror_key, pool_key = jax.random.split(key, 4) + learner = assignments.learner_seat.astype(jnp.int32) + opponent = 1 - learner + + obs = jax.vmap(observation, in_axes=(0, 0))(env, learner) + legal = jax.vmap(legal_action_mask)(env) + logits, value = state.apply_fn(state.params, obs) + masked_logits = mask_logits(logits, legal) + learner_actions = jax.random.categorical(learner_key, masked_logits, axis=-1).astype( + jnp.int32 + ) + log_prob = action_log_prob(masked_logits, learner_actions) + entropy = categorical_entropy(masked_logits) + + mirror_obs = jax.vmap(observation, in_axes=(0, 0))(env, opponent) + mirror_logits, _ = state.apply_fn(state.params, mirror_obs) + mirror_logits = jax.lax.stop_gradient(mirror_logits) + mirror_actions = jax.random.categorical( + mirror_key, mask_logits(mirror_logits, legal), axis=-1 + ).astype(jnp.int32) + + pool_keys = jax.random.split(pool_key, cfg.ppo.batch_games) + pool_actions = [] + for policy in opponent_policies: + pool_actions.append(jax.vmap(policy, in_axes=(0, 0, 0))(env, opponent, pool_keys)) + stacked_pool_actions = jnp.stack(pool_actions, axis=0).astype(jnp.int32) + selected_pool_actions = jnp.take_along_axis( + stacked_pool_actions, + assignments.opponent_index.astype(jnp.int32)[None, :], + axis=0, + )[0] + opponent_actions = jnp.where( + assignments.use_mirror, mirror_actions, selected_pool_actions + ) + + learner_turn = env.to_move.astype(jnp.int32) == learner + active = ~env.done + actions = jnp.where(learner_turn, learner_actions, opponent_actions) + + before_diff = batch_score_diff_for_players(env, learner) + next_env, _, _ = jax.vmap(step, in_axes=(0, 0))(env, actions) + after_diff = batch_score_diff_for_players(next_env, learner) + terminal_reward = jnp.tanh(after_diff / cfg.reward.terminal_scale) + terminal = active & next_env.done + reward = jnp.where(terminal, terminal_reward, 0.0) + reward = reward + shaping_coef * (after_diff - before_diff) + reward = jnp.where(active, reward, 0.0) + + actor_mask = active & learner_turn + place_type = (actions % 12) // 6 + transition = Transition( + obs=obs, + legal_mask=legal, + action=actions, + log_prob=log_prob, + value=value, + reward=reward, + done=next_env.done, + active=active, + actor_mask=actor_mask, + entropy=entropy, + play_action=(place_type == PLAY) & actor_mask, + ) + return (next_env, key), transition + + (next_env, rng), transitions = jax.lax.scan( + body, (env_state, rng), xs=None, length=cfg.ppo.rollout_steps + ) + metrics = rollout_metrics_for_players(transitions, next_env, assignments.learner_seat) + return next_env, transitions, metrics + + return rollout_fn + + def random_rollout(cfg: JaxPPOConfig) -> dict: opponent_policy = policy_by_name(cfg.opponent.name) rng = jax.random.PRNGKey(cfg.run.seed) @@ -399,6 +545,53 @@ def reset_done_envs(env_state: State, rng: jax.Array, batch_games: int) -> State ) +def sample_league_assignments( + rng: jax.Array, + batch_games: int, + opponent_probs: jax.Array, + mirror_probability: float, +) -> LeagueAssignments: + seat_key, opponent_key, mirror_key = jax.random.split(rng, 3) + learner_seat = jax.random.bernoulli(seat_key, 0.5, shape=(batch_games,)).astype(jnp.int32) + logits = jnp.log(jnp.maximum(jnp.asarray(opponent_probs, dtype=jnp.float32), 1.0e-12)) + opponent_index = jax.random.categorical(opponent_key, logits, shape=(batch_games,)).astype( + jnp.int32 + ) + use_mirror = jax.random.bernoulli(mirror_key, mirror_probability, shape=(batch_games,)) + return LeagueAssignments( + learner_seat=learner_seat, + opponent_index=opponent_index, + use_mirror=use_mirror, + ) + + +def reset_done_envs_with_assignments( + env_state: State, + assignments: LeagueAssignments, + rng: jax.Array, + batch_games: int, + opponent_probs: jax.Array, + mirror_probability: float, +) -> tuple[State, LeagueAssignments]: + env_key, assignment_key = jax.random.split(rng) + fresh_env = jax.vmap(reset)(jax.random.split(env_key, batch_games)) + fresh_assignments = sample_league_assignments( + assignment_key, batch_games, opponent_probs, mirror_probability + ) + done = env_state.done + next_env = jax.tree_util.tree_map( + lambda old, new: jnp.where(_broadcast_done(done, old), new, old), env_state, fresh_env + ) + next_assignments = LeagueAssignments( + learner_seat=jnp.where(done, fresh_assignments.learner_seat, assignments.learner_seat), + opponent_index=jnp.where( + done, fresh_assignments.opponent_index, assignments.opponent_index + ), + use_mirror=jnp.where(done, fresh_assignments.use_mirror, assignments.use_mirror), + ) + return next_env, next_assignments + + def ppo_update( state: TrainState, transitions: Transition, @@ -899,6 +1092,33 @@ def rollout_metrics( } +def rollout_metrics_for_players( + transitions: Transition, final_env: State, learners: jax.Array +) -> dict[str, jax.Array]: + episode_return = jnp.sum(transitions.reward, axis=0) + actor_count = jnp.sum(transitions.actor_mask) + active_count = jnp.sum(transitions.active) + color_scores = jax.vmap(color_scores_for_player, in_axes=(0, 0))(final_env, learners) + batch_idx = jnp.arange(learners.shape[0]) + opened = jnp.sum(final_env.col_len[batch_idx, learners.astype(jnp.int32), :] > 0, axis=-1) + positive = jnp.sum(color_scores > 0, axis=-1) + return { + "return_mean": jnp.mean(episode_return), + "return_std": jnp.std(episode_return), + "game_length_mean": jnp.mean(final_env.step_count.astype(jnp.float32)), + "game_length_max": jnp.max(final_env.step_count), + "max_steps_rate": jnp.mean(final_env.step_count >= MAX_STEPS), + "play_action_rate": jnp.sum(transitions.play_action) / jnp.maximum(actor_count, 1), + "opened_colors_mean": jnp.mean(opened.astype(jnp.float32)), + "positive_expeditions_mean": jnp.mean(positive.astype(jnp.float32)), + "entropy_mean": masked_mean( + transitions.entropy, transitions.actor_mask.astype(jnp.float32) + ), + "active_steps": active_count, + "learner_actions": actor_count, + } + + def color_scores_for_player(state: State, player: jax.Array) -> jax.Array: card_ids = jnp.arange(N_CARDS, dtype=jnp.int32) colors = card_ids // CARDS_PER_COLOR @@ -923,6 +1143,14 @@ def batch_score_diff(state: State, learner: jax.Array) -> jax.Array: return scores[:, learner.astype(jnp.int32)] - scores[:, opponent] +def batch_score_diff_for_players(state: State, learners: jax.Array) -> jax.Array: + scores = jax.vmap(board_score)(state) + learners = learners.astype(jnp.int32) + opponents = 1 - learners + batch_idx = jnp.arange(learners.shape[0]) + return scores[batch_idx, learners] - scores[batch_idx, opponents] + + def mask_logits(logits: jax.Array, legal_mask: jax.Array) -> jax.Array: return jnp.where(legal_mask, logits, NEG_INF) @@ -1057,10 +1285,15 @@ def cli_main(argv: list[str] | None = None) -> None: static_match_parser.add_argument("--duplicate", action="store_true") static_match_parser.add_argument("--output") + league_parser = sub.add_parser("league") + league_sub = league_parser.add_subparsers(dest="league_command", required=True) + league_run_parser = league_sub.add_parser("run") + league_run_parser.add_argument("--config", required=True) + args = parser.parse_args(argv) cfg = ( load_config(args.config, overrides=parse_overrides(args.set)) - if hasattr(args, "config") + if hasattr(args, "config") and args.command != "league" else None ) if args.command == "rollout-smoke": @@ -1139,6 +1372,12 @@ def cli_main(argv: list[str] | None = None) -> None: output=args.output, ) print(json.dumps(result, indent=2, sort_keys=True)) + elif args.command == "league": + if args.league_command == "run": + from lost_cities_jax.league import run_league + + run_dir = run_league(args.config) + print(run_dir) def _create_run_dir(cfg: JaxPPOConfig) -> Path: @@ -1210,6 +1449,7 @@ def _slugify(value: str) -> str: __all__ = [ "ActorCritic", "JaxPPOConfig", + "LeagueAssignments", "TrainState", "checkpoint_policy", "checkpoint_policy_from_params", @@ -1222,7 +1462,9 @@ __all__ = [ "evaluate_static_match", "evaluate_static_mirror", "load_config", + "make_league_train_iteration", "random_rollout", + "sample_league_assignments", "train", "train_against_checkpoint", ] diff --git a/tests/lost_cities_jax/test_league.py b/tests/lost_cities_jax/test_league.py new file mode 100644 index 0000000..bd61cf8 --- /dev/null +++ b/tests/lost_cities_jax/test_league.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +import jax +import jax.numpy as jnp + +from lost_cities_jax import reset +from lost_cities_jax.league import ( + LeagueConfig, + PoolMember, + active_training_pool, + load_league_config, + pool_sampling_probabilities, +) +from lost_cities_jax.opponents import policy_by_name +from lost_cities_jax.ppo import ( + JaxPPOConfig, + NetworkConfig, + OpponentConfig, + PPOHyperConfig, + RunConfig, + create_train_state, + make_league_train_iteration, + sample_league_assignments, +) + + +def tiny_league_ppo_config(tmp_path) -> JaxPPOConfig: + return JaxPPOConfig( + run=RunConfig( + experiment_name="pytest-jax-ppo-league", + seed=123, + total_updates=1, + checkpoint_every=1, + artifact_root=str(tmp_path), + ), + opponent=OpponentConfig(name="discard_only"), + network=NetworkConfig(hidden_size=32, num_layers=1), + ppo=PPOHyperConfig(batch_games=8, rollout_steps=16, epochs=1, minibatches=2), + ) + + +def test_league_config_loads(): + cfg = load_league_config("configs/jax_ppo/league-v1.yaml") + assert cfg.experiment_name == "jax-ppo-league-v1" + assert cfg.cycles == 5 + assert len(cfg.anchors) == 7 + assert any(member.name == "heuristic_expert" for member in cfg.anchors) + + +def test_pfsp_probabilities_cap_stalling_anchors(): + cfg = LeagueConfig( + base_config="configs/jax_ppo/ladder-v2-expert.yaml", + warm_start_checkpoint="unused", + stalling_anchor_floor=0.02, + stalling_anchor_cap=0.05, + uniform_mix=0.1, + ) + pool = [ + PoolMember("discard_only", "static", stalling=True, recent_win_rate=0.01), + PoolMember("heuristic_expert", "static", recent_win_rate=0.5), + PoolMember("snapshot", "checkpoint", anchor=False, recent_win_rate=0.2), + ] + probs = pool_sampling_probabilities(pool, cfg) + assert abs(sum(probs) - 1.0) < 1.0e-6 + assert 0.02 <= probs[0] <= 0.05 + + +def test_active_training_pool_keeps_anchors_and_hard_nonanchors(): + cfg = LeagueConfig( + base_config="configs/jax_ppo/ladder-v2-expert.yaml", + warm_start_checkpoint="unused", + max_active_pool_members=4, + ) + pool = [ + PoolMember("anchor_a", "static", anchor=True), + PoolMember("anchor_b", "static", anchor=True), + PoolMember("easy", "checkpoint", anchor=False, recent_win_rate=0.9), + PoolMember("hard", "checkpoint", anchor=False, recent_win_rate=0.1), + PoolMember("medium", "checkpoint", anchor=False, recent_win_rate=0.5), + ] + active = active_training_pool(pool, cfg) + assert [member.name for member in active] == ["anchor_a", "anchor_b", "hard", "medium"] + + +def test_one_jitted_league_train_iteration_shapes(tmp_path): + cfg = tiny_league_ppo_config(tmp_path) + train_state = create_train_state(cfg, jax.random.PRNGKey(0)) + env_state = jax.jit(jax.vmap(reset))(jax.random.split(jax.random.PRNGKey(1), 8)) + probs = jnp.asarray([0.5, 0.5], dtype=jnp.float32) + assignments = sample_league_assignments(jax.random.PRNGKey(2), 8, probs, 0.5) + train_iteration = make_league_train_iteration( + cfg, + [policy_by_name("discard_only"), policy_by_name("heuristic_expert")], + probs, + 0.5, + ) + + train_state, env_state, assignments, rng, metrics = train_iteration( + train_state, + env_state, + assignments, + jax.random.PRNGKey(3), + jnp.asarray(0.0, dtype=jnp.float32), + ) + + assert env_state.to_move.shape == (8,) + assert assignments.learner_seat.shape == (8,) + assert rng.shape == (2,) + assert "opened_colors_mean" in metrics