Add JAX PPO league self-play v1

2026-07-05 06:33:22 +09:00
parent 6037b650f3
commit 94e9ac1854
10 changed files with 1482 additions and 1 deletions
@@ -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