Add JAX PPO ladder verification pass

2026-07-05 00:17:49 +09:00
parent bb52ef9ec1
commit 7fe0bfdcfe
6 changed files with 14243 additions and 3 deletions
@@ -1,10 +1,14 @@
 from __future__ import annotations

+from pathlib import Path
+
 import jax
 import jax.numpy as jnp
 import numpy as np
+import pytest

 from lost_cities_jax import legal_action_mask, reset
+from lost_cities_jax.engine import current_hand_sorted, decode_action, reset_from_order
 from lost_cities_jax.opponents import (
    discard_only_action,
    heuristic_balanced_action,
@@ -18,12 +22,29 @@ from lost_cities_jax.ppo import (
    RunConfig,
    create_train_state,
    evaluate,
+    evaluate_checkpoint_match,
+    evaluate_static_mirror,
    load_config,
    make_train_iteration,
    policy_by_name,
    random_rollout,
    train,
 )
+from lost_cities_jax.types import (
+    CARDS_PER_COLOR,
+    HAND_SIZE,
+    LOC_DECK,
+    LOC_P0_BOARD,
+    LOC_P0_HAND,
+    N_CARDS,
+    N_COLORS,
+    PLAY,
+)
+
+GATE3_CHECKPOINT = (
+    "/mnt/2tbhdd/coolrl-lost-cities-artifacts/jax-ppo-static-opponents/"
+    "2026-07-04_230150_jax-ppo-cautious/latest"
+)


 def tiny_config(tmp_path) -> JaxPPOConfig:
@@ -49,6 +70,84 @@ def test_static_opponents_return_legal_actions():
        assert mask[action]


+def test_cautious_discards_weak_unopened_hand():
+    low_cards = [_rank_card(0, rank) for rank in range(2, 7)]
+    state = _manual_state(p0_hand=low_cards, to_move=0)
+    action = int(heuristic_cautious_action(state, jnp.int32(0), jax.random.PRNGKey(1)))
+    _, place_type, _ = [int(x) for x in decode_action(jnp.asarray(action))]
+    assert place_type != PLAY
+
+
+def test_balanced_opens_strong_unopened_rank():
+    strong = _rank_card(0, 10)
+    state = _manual_state(p0_hand=[strong, _rank_card(1, 2), _rank_card(2, 3)], to_move=0)
+    action = int(heuristic_balanced_action(state, jnp.int32(0), jax.random.PRNGKey(1)))
+    hand_slot, place_type, _ = [int(x) for x in decode_action(jnp.asarray(action))]
+    hand = [int(x) for x in current_hand_sorted(state, 0)]
+    assert place_type == PLAY
+    assert hand[hand_slot] == strong
+
+
+def test_cautious_uses_own_board_when_playing_as_p1():
+    existing = _rank_card(0, 5)
+    playable = _rank_card(0, 6)
+    state = _manual_state(
+        p1_hand=[
+            playable,
+            _rank_card(1, 2),
+            _rank_card(1, 3),
+            _rank_card(2, 2),
+            _rank_card(2, 3),
+            _rank_card(3, 2),
+            _rank_card(3, 3),
+            _rank_card(4, 2),
+        ],
+        p1_board=[existing],
+        to_move=1,
+    )
+    action = int(heuristic_cautious_action(state, jnp.int32(1), jax.random.PRNGKey(1)))
+    hand_slot, place_type, _ = [int(x) for x in decode_action(jnp.asarray(action))]
+    hand = [int(x) for x in current_hand_sorted(state, 1)]
+    assert place_type == PLAY
+    assert hand[hand_slot] == playable
+
+
+@pytest.mark.parametrize(
+    "policy_name", ["discard_only", "heuristic_balanced", "heuristic_cautious"]
+)
+def test_static_policy_duplicate_mirror_score_diff_is_zero(policy_name):
+    result = evaluate_static_mirror(
+        policy_name,
+        games=16,
+        duplicate=True,
+        shuffle_bank_seed=20260704,
+        batch_games=16,
+    )
+    assert result["games"] == 32
+    assert result["mean_score_diff"] == 0.0
+    assert result["wins"] == result["losses"]
+
+
+def test_gate3_checkpoint_duplicate_self_mirror_score_diff_is_zero():
+    checkpoint = Path(GATE3_CHECKPOINT)
+    if not checkpoint.exists():
+        pytest.skip(f"gate-3 checkpoint not available: {checkpoint}")
+    if not any(device.platform == "gpu" for device in jax.local_devices()):
+        pytest.skip("gate-3 checkpoint was saved with CUDA sharding; run under CUDA JAX")
+    cfg = load_config("configs/jax_ppo/cautious.yaml")
+    result = evaluate_checkpoint_match(
+        cfg,
+        checkpoint,
+        cfg,
+        checkpoint,
+        games=4,
+        duplicate=True,
+    )
+    assert result["games"] == 8
+    assert result["mean_score_diff"] == 0.0
+    assert result["wins"] == result["losses"]
+
+
 def test_random_rollout_smoke(tmp_path):
    row = random_rollout(tiny_config(tmp_path))
    assert row["env_steps"] == 8 * 16
@@ -84,3 +183,65 @@ def test_load_config_file():
    cfg = load_config("configs/jax_ppo/smoke.yaml")
    assert cfg.opponent.name == "discard_only"
    assert cfg.ppo.batch_games == 64
+
+
+def _rank_card(color: int, rank: int) -> int:
+    return color * CARDS_PER_COLOR + rank + 1
+
+
+def _manual_state(
+    *,
+    p0_hand: list[int] | None = None,
+    p1_hand: list[int] | None = None,
+    p0_board: list[int] | None = None,
+    p1_board: list[int] | None = None,
+    to_move: int = 0,
+):
+    p0_hand = _fill_hand(p0_hand or [], set((p0_board or []) + (p1_board or [])))
+    p1_hand = _fill_hand(p1_hand or [], set(p0_hand + (p0_board or []) + (p1_board or [])))
+    rest = [card for card in range(N_CARDS) if card not in set(p0_hand + p1_hand)]
+    deck_order = jnp.asarray(p0_hand + p1_hand + rest, dtype=jnp.int8)
+    state = reset_from_order(deck_order)
+    card_loc = jnp.full((N_CARDS,), LOC_DECK, dtype=jnp.int8)
+    card_loc = card_loc.at[jnp.asarray(p0_hand)].set(LOC_P0_HAND)
+    card_loc = card_loc.at[jnp.asarray(p1_hand)].set(LOC_P0_HAND + 1)
+    if p0_board:
+        card_loc = card_loc.at[jnp.asarray(p0_board)].set(LOC_P0_BOARD)
+    if p1_board:
+        card_loc = card_loc.at[jnp.asarray(p1_board)].set(LOC_P0_BOARD + 1)
+    col_top, col_hs, col_len = _columns(p0_board or [], p1_board or [])
+    return state._replace(
+        card_loc=card_loc,
+        col_top=col_top,
+        col_hs=col_hs,
+        col_len=col_len,
+        to_move=jnp.asarray(to_move, dtype=jnp.int8),
+    )
+
+
+def _fill_hand(cards: list[int], reserved: set[int]) -> list[int]:
+    result = list(cards)
+    blocked = set(result) | reserved
+    for card in range(N_CARDS):
+        if len(result) == HAND_SIZE:
+            return sorted(result)
+        if card not in blocked:
+            result.append(card)
+            blocked.add(card)
+    raise AssertionError("could not fill hand")
+
+
+def _columns(p0_board: list[int], p1_board: list[int]):
+    col_top = np.zeros((2, N_COLORS), dtype=np.int8)
+    col_hs = np.zeros((2, N_COLORS), dtype=np.int8)
+    col_len = np.zeros((2, N_COLORS), dtype=np.int8)
+    for player, cards in enumerate([p0_board, p1_board]):
+        for card in cards:
+            color = card // CARDS_PER_COLOR
+            slot = card % CARDS_PER_COLOR
+            col_len[player, color] += 1
+            if slot < 3:
+                col_hs[player, color] += 1
+            else:
+                col_top[player, color] = max(col_top[player, color], slot - 1)
+    return jnp.asarray(col_top), jnp.asarray(col_hs), jnp.asarray(col_len)