407 lines
13 KiB
Python
407 lines
13 KiB
Python
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 (
|
|
HeuristicExpertConfig,
|
|
discard_only_action,
|
|
heuristic_balanced_action,
|
|
heuristic_cautious_action,
|
|
heuristic_expert_action,
|
|
make_heuristic_expert_policy,
|
|
)
|
|
from lost_cities_jax.ppo import (
|
|
JaxPPOConfig,
|
|
NetworkConfig,
|
|
OpponentConfig,
|
|
PPOHyperConfig,
|
|
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:
|
|
return JaxPPOConfig(
|
|
run=RunConfig(
|
|
experiment_name="pytest-jax-ppo",
|
|
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_static_opponents_return_legal_actions():
|
|
state = reset(jax.random.PRNGKey(0))
|
|
mask = np.asarray(legal_action_mask(state), dtype=bool)
|
|
for fn in [
|
|
discard_only_action,
|
|
heuristic_balanced_action,
|
|
heuristic_cautious_action,
|
|
heuristic_expert_action,
|
|
]:
|
|
action = int(fn(state, jnp.int32(0), jax.random.PRNGKey(1)))
|
|
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
|
|
|
|
|
|
def test_expert_opens_strong_ev_hand():
|
|
strong_cards = [
|
|
_hs_card(0, 0),
|
|
_rank_card(0, 7),
|
|
_rank_card(0, 8),
|
|
_rank_card(0, 9),
|
|
_rank_card(0, 10),
|
|
_rank_card(1, 2),
|
|
_rank_card(2, 2),
|
|
_rank_card(3, 2),
|
|
]
|
|
state = _manual_state(p0_hand=strong_cards, to_move=0)
|
|
action = int(heuristic_expert_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] // CARDS_PER_COLOR == 0
|
|
|
|
|
|
def test_expert_rejects_weak_new_open():
|
|
weak_cards = [
|
|
_rank_card(0, 2),
|
|
_rank_card(0, 3),
|
|
_rank_card(0, 4),
|
|
_rank_card(1, 2),
|
|
_rank_card(2, 2),
|
|
_rank_card(3, 2),
|
|
_rank_card(4, 2),
|
|
_rank_card(4, 3),
|
|
]
|
|
state = _manual_state(p0_hand=weak_cards, to_move=0)
|
|
action = int(heuristic_expert_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_expert_avoids_discarding_immediately_useful_opponent_card():
|
|
dangerous = _rank_card(0, 8)
|
|
safe = _rank_card(4, 2)
|
|
state = _manual_state(
|
|
p0_hand=[
|
|
dangerous,
|
|
safe,
|
|
_rank_card(1, 2),
|
|
_rank_card(1, 3),
|
|
_rank_card(2, 2),
|
|
_rank_card(2, 3),
|
|
_rank_card(3, 2),
|
|
_rank_card(3, 3),
|
|
],
|
|
p1_board=[_rank_card(0, 7)],
|
|
to_move=0,
|
|
)
|
|
action = int(heuristic_expert_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)]
|
|
if place_type != PLAY:
|
|
assert hand[hand_slot] != dangerous
|
|
|
|
|
|
def test_expert_cap2_blocks_third_open_color():
|
|
state = _manual_state(
|
|
p0_hand=[
|
|
_hs_card(2),
|
|
_rank_card(2, 7),
|
|
_rank_card(2, 8),
|
|
_rank_card(2, 9),
|
|
_rank_card(2, 10),
|
|
_rank_card(3, 2),
|
|
_rank_card(3, 3),
|
|
_rank_card(4, 2),
|
|
],
|
|
p0_board=[_rank_card(0, 2), _rank_card(1, 2)],
|
|
to_move=0,
|
|
)
|
|
policy = make_heuristic_expert_policy(HeuristicExpertConfig(max_open_colors=2))
|
|
action = int(policy(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)]
|
|
opened_new_color = (
|
|
place_type == PLAY and state.col_len[0, hand[hand_slot] // CARDS_PER_COLOR] == 0
|
|
)
|
|
assert not bool(opened_new_color)
|
|
|
|
|
|
def test_expert_cap3_allows_third_but_blocks_fourth_open_color():
|
|
cap3 = make_heuristic_expert_policy(HeuristicExpertConfig(max_open_colors=3))
|
|
third_state = _manual_state(
|
|
p0_hand=[
|
|
_hs_card(2),
|
|
_rank_card(2, 7),
|
|
_rank_card(2, 8),
|
|
_rank_card(2, 9),
|
|
_rank_card(2, 10),
|
|
_rank_card(3, 2),
|
|
_rank_card(3, 3),
|
|
_rank_card(4, 2),
|
|
],
|
|
p0_board=[_rank_card(0, 2), _rank_card(1, 2)],
|
|
to_move=0,
|
|
)
|
|
action = int(cap3(third_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(third_state, 0)]
|
|
assert place_type == PLAY
|
|
assert third_state.col_len[0, hand[hand_slot] // CARDS_PER_COLOR] == 0
|
|
|
|
fourth_state = _manual_state(
|
|
p0_hand=[
|
|
_hs_card(3),
|
|
_rank_card(3, 7),
|
|
_rank_card(3, 8),
|
|
_rank_card(3, 9),
|
|
_rank_card(3, 10),
|
|
_rank_card(4, 2),
|
|
_rank_card(4, 3),
|
|
_rank_card(4, 4),
|
|
],
|
|
p0_board=[_rank_card(0, 2), _rank_card(1, 2), _rank_card(2, 2)],
|
|
to_move=0,
|
|
)
|
|
action = int(cap3(fourth_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(fourth_state, 0)]
|
|
opened_new_color = (
|
|
place_type == PLAY and fourth_state.col_len[0, hand[hand_slot] // CARDS_PER_COLOR] == 0
|
|
)
|
|
assert not bool(opened_new_color)
|
|
|
|
|
|
@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_expert_duplicate_mirror_regression_stats():
|
|
result = evaluate_static_mirror(
|
|
"heuristic_expert",
|
|
games=1000,
|
|
duplicate=True,
|
|
shuffle_bank_seed=20260704,
|
|
batch_games=1000,
|
|
)
|
|
assert result["games"] == 2000
|
|
assert abs(result["mean_score_diff"]) <= 1.0e-6
|
|
assert result["wins"] == result["losses"]
|
|
assert result["max_steps_rate"] < 0.05
|
|
assert 2.0 <= result["opened_colors_per_game"] <= 3.5
|
|
assert 0.30 <= result["play_action_rate"] <= 0.60
|
|
assert result["mean_game_length"] <= 70.0
|
|
|
|
|
|
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
|
|
assert 0.0 <= row["play_action_rate"] <= 1.0
|
|
assert row["game_length_mean"] > 0.0
|
|
|
|
|
|
def test_train_checkpoint_and_eval_smoke(tmp_path):
|
|
cfg = tiny_config(tmp_path)
|
|
run_dir = train(cfg)
|
|
assert (run_dir / "latest").exists()
|
|
result = evaluate(cfg, run_dir / "latest", games=8, duplicate=True)
|
|
assert result["games"] == 16
|
|
assert 0.0 <= result["win_rate"] <= 1.0
|
|
|
|
|
|
def test_one_jitted_train_iteration_shapes(tmp_path):
|
|
cfg = tiny_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))
|
|
train_iteration = make_train_iteration(cfg, policy_by_name("discard_only"))
|
|
|
|
train_state, env_state, rng, metrics = train_iteration(
|
|
train_state, env_state, jax.random.PRNGKey(2), jnp.asarray(1.0)
|
|
)
|
|
|
|
assert env_state.to_move.shape == (8,)
|
|
assert "loss" in metrics
|
|
assert jnp.isfinite(metrics["loss"])
|
|
|
|
|
|
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 _hs_card(color: int, slot: int = 0) -> int:
|
|
return color * CARDS_PER_COLOR + slot
|
|
|
|
|
|
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)
|