Add JAX engine verification tests
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""Reference implementations used by tests."""
|
||||
@@ -0,0 +1,247 @@
|
||||
"""Independent pure-Python Lost Cities rules reference.
|
||||
|
||||
This module intentionally uses ordinary Python containers rather than the JAX
|
||||
state representation. It is deterministic under an explicit ``deck_order`` and
|
||||
the shared flat action encoding.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from random import Random
|
||||
|
||||
N_PLAYERS = 2
|
||||
N_COLORS = 5
|
||||
CARDS_PER_COLOR = 12
|
||||
N_CARDS = N_COLORS * CARDS_PER_COLOR
|
||||
HAND_SIZE = 8
|
||||
INITIAL_DEAL = N_PLAYERS * HAND_SIZE
|
||||
MAX_STEPS = 400
|
||||
N_ACTIONS = HAND_SIZE * 2 * 6
|
||||
|
||||
LOC_DECK = 0
|
||||
LOC_P0_HAND = 1
|
||||
LOC_P1_HAND = 2
|
||||
LOC_P0_BOARD = 3
|
||||
LOC_P1_BOARD = 4
|
||||
LOC_DISCARD = 5
|
||||
|
||||
NO_CARD = -1
|
||||
PLAY = 0
|
||||
DISCARD = 1
|
||||
DRAW_DECK = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class RefState:
|
||||
deck_order: list[int]
|
||||
draw_ptr: int
|
||||
card_loc: list[int]
|
||||
hand_public: list[bool]
|
||||
board: list[list[list[int]]]
|
||||
piles: list[list[int]]
|
||||
to_move: int
|
||||
just_discarded: int
|
||||
step_count: int
|
||||
done: bool
|
||||
|
||||
|
||||
def reset(seed: int | None = None) -> RefState:
|
||||
rng = Random(seed)
|
||||
order = list(range(N_CARDS))
|
||||
rng.shuffle(order)
|
||||
return reset_from_order(order)
|
||||
|
||||
|
||||
def reset_from_order(deck_order: list[int] | tuple[int, ...]) -> RefState:
|
||||
order = [int(card) for card in deck_order]
|
||||
card_loc = [LOC_DECK] * N_CARDS
|
||||
for card in order[:HAND_SIZE]:
|
||||
card_loc[card] = LOC_P0_HAND
|
||||
for card in order[HAND_SIZE:INITIAL_DEAL]:
|
||||
card_loc[card] = LOC_P1_HAND
|
||||
return RefState(
|
||||
deck_order=order,
|
||||
draw_ptr=INITIAL_DEAL,
|
||||
card_loc=card_loc,
|
||||
hand_public=[False] * N_CARDS,
|
||||
board=[[[] for _ in range(N_COLORS)] for _ in range(N_PLAYERS)],
|
||||
piles=[[] for _ in range(N_COLORS)],
|
||||
to_move=0,
|
||||
just_discarded=NO_CARD,
|
||||
step_count=0,
|
||||
done=False,
|
||||
)
|
||||
|
||||
|
||||
def clone_state(state: RefState) -> RefState:
|
||||
return RefState(
|
||||
deck_order=list(state.deck_order),
|
||||
draw_ptr=state.draw_ptr,
|
||||
card_loc=list(state.card_loc),
|
||||
hand_public=list(state.hand_public),
|
||||
board=[[list(col) for col in player] for player in state.board],
|
||||
piles=[list(pile) for pile in state.piles],
|
||||
to_move=state.to_move,
|
||||
just_discarded=state.just_discarded,
|
||||
step_count=state.step_count,
|
||||
done=state.done,
|
||||
)
|
||||
|
||||
|
||||
def decode_action(action: int) -> tuple[int, int, int]:
|
||||
action = int(action)
|
||||
hand_slot = action // 12
|
||||
rem = action % 12
|
||||
place_type = rem // 6
|
||||
draw_source = rem % 6
|
||||
return hand_slot, place_type, draw_source
|
||||
|
||||
|
||||
def hand_cards(state: RefState, player: int | None = None) -> list[int]:
|
||||
if player is None:
|
||||
player = state.to_move
|
||||
hand_loc = LOC_P0_HAND + player
|
||||
return sorted(card for card, loc in enumerate(state.card_loc) if loc == hand_loc)
|
||||
|
||||
|
||||
def legal_action_mask(state: RefState) -> list[bool]:
|
||||
mask = [False] * N_ACTIONS
|
||||
if state.done:
|
||||
return mask
|
||||
|
||||
player = state.to_move
|
||||
hand = hand_cards(state, player)
|
||||
for hand_slot in range(HAND_SIZE):
|
||||
if hand_slot >= len(hand):
|
||||
continue
|
||||
card = hand[hand_slot]
|
||||
for place_type in (PLAY, DISCARD):
|
||||
if place_type == PLAY:
|
||||
place_ok = can_play_card(state, player, card)
|
||||
else:
|
||||
place_ok = True
|
||||
if not place_ok:
|
||||
continue
|
||||
for draw_source in range(6):
|
||||
if _can_draw_after_place(state, card, place_type, draw_source):
|
||||
mask[hand_slot * 12 + place_type * 6 + draw_source] = True
|
||||
return mask
|
||||
|
||||
|
||||
def step(state: RefState, action: int) -> tuple[RefState, list[float], bool]:
|
||||
if state.done or action < 0 or action >= N_ACTIONS:
|
||||
return clone_state(state), [0.0, 0.0], state.done
|
||||
if not legal_action_mask(state)[action]:
|
||||
return clone_state(state), [0.0, 0.0], state.done
|
||||
|
||||
next_state = clone_state(state)
|
||||
player = next_state.to_move
|
||||
hand_slot, place_type, draw_source = decode_action(action)
|
||||
card = hand_cards(next_state, player)[hand_slot]
|
||||
color = card_color(card)
|
||||
|
||||
next_state.hand_public[card] = False
|
||||
if place_type == PLAY:
|
||||
next_state.board[player][color].append(card)
|
||||
next_state.card_loc[card] = LOC_P0_BOARD + player
|
||||
next_state.just_discarded = NO_CARD
|
||||
else:
|
||||
next_state.piles[color].append(card)
|
||||
next_state.card_loc[card] = LOC_DISCARD
|
||||
next_state.just_discarded = card
|
||||
|
||||
if draw_source == DRAW_DECK:
|
||||
drawn = next_state.deck_order[next_state.draw_ptr]
|
||||
next_state.draw_ptr += 1
|
||||
public = False
|
||||
else:
|
||||
src = draw_source - 1
|
||||
drawn = next_state.piles[src].pop()
|
||||
public = True
|
||||
|
||||
next_state.card_loc[drawn] = LOC_P0_HAND + player
|
||||
next_state.hand_public[drawn] = public
|
||||
next_state.just_discarded = NO_CARD
|
||||
next_state.step_count += 1
|
||||
next_state.done = (draw_source == DRAW_DECK and next_state.draw_ptr >= N_CARDS) or (
|
||||
next_state.step_count >= MAX_STEPS
|
||||
)
|
||||
next_state.to_move = 1 - player
|
||||
|
||||
reward = board_score(next_state) if next_state.done else [0.0, 0.0]
|
||||
return next_state, reward, next_state.done
|
||||
|
||||
|
||||
def can_play_card(state: RefState, player: int, card: int) -> bool:
|
||||
column = state.board[player][card_color(card)]
|
||||
if is_handshake(card):
|
||||
return not any(not is_handshake(played) for played in column)
|
||||
top_rank = 0
|
||||
for played in column:
|
||||
if not is_handshake(played):
|
||||
top_rank = rank(played)
|
||||
return rank(card) > top_rank
|
||||
|
||||
|
||||
def board_score(state: RefState) -> list[float]:
|
||||
return [float(sum(score_column(column) for column in player)) for player in state.board]
|
||||
|
||||
|
||||
def score(state: RefState) -> list[float]:
|
||||
return board_score(state)
|
||||
|
||||
|
||||
def score_column(column: list[int]) -> int:
|
||||
if not column:
|
||||
return 0
|
||||
handshakes = sum(1 for card in column if is_handshake(card))
|
||||
rank_sum = sum(rank(card) for card in column if not is_handshake(card))
|
||||
value = (rank_sum - 20) * (1 + handshakes)
|
||||
if len(column) >= 8:
|
||||
value += 20
|
||||
return value
|
||||
|
||||
|
||||
def card_color(card: int) -> int:
|
||||
return int(card) // CARDS_PER_COLOR
|
||||
|
||||
|
||||
def card_slot(card: int) -> int:
|
||||
return int(card) % CARDS_PER_COLOR
|
||||
|
||||
|
||||
def is_handshake(card: int) -> bool:
|
||||
return card_slot(card) < 3
|
||||
|
||||
|
||||
def rank(card: int) -> int:
|
||||
return card_slot(card) - 1
|
||||
|
||||
|
||||
def _can_draw_after_place(state: RefState, card: int, place_type: int, draw_source: int) -> bool:
|
||||
if draw_source == DRAW_DECK:
|
||||
return True
|
||||
src = draw_source - 1
|
||||
same_discard_pile = place_type == DISCARD and card_color(card) == src
|
||||
after_len = len(state.piles[src]) + int(same_discard_pile)
|
||||
if after_len == 0:
|
||||
return False
|
||||
after_top = card if same_discard_pile else state.piles[src][-1]
|
||||
just_discarded = card if place_type == DISCARD else state.just_discarded
|
||||
return after_top != just_discarded
|
||||
|
||||
|
||||
__all__ = [
|
||||
"N_ACTIONS",
|
||||
"RefState",
|
||||
"board_score",
|
||||
"clone_state",
|
||||
"decode_action",
|
||||
"hand_cards",
|
||||
"legal_action_mask",
|
||||
"reset",
|
||||
"reset_from_order",
|
||||
"score",
|
||||
"step",
|
||||
]
|
||||
@@ -0,0 +1 @@
|
||||
"""Tests for the standalone JAX Lost Cities engine."""
|
||||
@@ -0,0 +1,26 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
|
||||
def pytest_addoption(parser):
|
||||
parser.addoption(
|
||||
"--full",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="run full Lost Cities JAX differential profile",
|
||||
)
|
||||
|
||||
|
||||
def pytest_configure(config):
|
||||
config.addinivalue_line("markers", "slow: longer randomized Lost Cities JAX checks")
|
||||
|
||||
|
||||
def game_count(config, env_name: str, local: int, ci: int, full: int) -> int:
|
||||
if env_name in os.environ:
|
||||
return int(os.environ[env_name])
|
||||
if config.getoption("--full"):
|
||||
return full
|
||||
if os.environ.get("CI"):
|
||||
return ci
|
||||
return local
|
||||
@@ -0,0 +1,129 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import numpy as np
|
||||
|
||||
from lost_cities_jax import reset_from_order
|
||||
from lost_cities_jax.types import (
|
||||
CARDS_PER_COLOR,
|
||||
LOC_DISCARD,
|
||||
LOC_P0_BOARD,
|
||||
LOC_P0_HAND,
|
||||
MAX_PILE_SIZE,
|
||||
N_CARDS,
|
||||
N_COLORS,
|
||||
NO_CARD,
|
||||
State,
|
||||
)
|
||||
|
||||
|
||||
def shuffled_order(seed: int) -> list[int]:
|
||||
order = list(range(N_CARDS))
|
||||
random.Random(seed).shuffle(order)
|
||||
return order
|
||||
|
||||
|
||||
def hs(color: int, index: int = 0) -> int:
|
||||
return color * CARDS_PER_COLOR + index
|
||||
|
||||
|
||||
def num(color: int, rank: int) -> int:
|
||||
return color * CARDS_PER_COLOR + rank + 1
|
||||
|
||||
|
||||
def state_with_columns(
|
||||
p0_columns: list[list[int]],
|
||||
p1_columns: list[list[int]] | None = None,
|
||||
) -> State:
|
||||
if p1_columns is None:
|
||||
p1_columns = [[] for _ in range(N_COLORS)]
|
||||
|
||||
state = reset_from_order(jnp.arange(N_CARDS, dtype=jnp.int8))
|
||||
card_loc = jnp.zeros((N_CARDS,), dtype=jnp.int8)
|
||||
col_top = jnp.zeros((2, N_COLORS), dtype=jnp.int8)
|
||||
col_hs = jnp.zeros((2, N_COLORS), dtype=jnp.int8)
|
||||
col_len = jnp.zeros((2, N_COLORS), dtype=jnp.int8)
|
||||
|
||||
for player, columns in enumerate([p0_columns, p1_columns]):
|
||||
for color, cards in enumerate(columns):
|
||||
if not cards:
|
||||
continue
|
||||
idx = jnp.asarray(cards, dtype=jnp.int32)
|
||||
card_loc = card_loc.at[idx].set(jnp.int8(LOC_P0_BOARD + player))
|
||||
slots = [card % CARDS_PER_COLOR for card in cards]
|
||||
ranks = [slot - 1 for slot in slots if slot >= 3]
|
||||
col_top = col_top.at[player, color].set(max(ranks, default=0))
|
||||
col_hs = col_hs.at[player, color].set(sum(slot < 3 for slot in slots))
|
||||
col_len = col_len.at[player, color].set(len(cards))
|
||||
|
||||
return state._replace(
|
||||
card_loc=card_loc,
|
||||
col_top=col_top,
|
||||
col_hs=col_hs,
|
||||
col_len=col_len,
|
||||
)
|
||||
|
||||
|
||||
def first_deck_draw_action(mask) -> int:
|
||||
values = np.asarray(mask, dtype=bool)
|
||||
for action, legal in enumerate(values):
|
||||
if legal and action % 6 == 0:
|
||||
return action
|
||||
raise AssertionError("no legal deck-draw action")
|
||||
|
||||
|
||||
def assert_state_consistent(state: State) -> None:
|
||||
loc = np.asarray(state.card_loc)
|
||||
public = np.asarray(state.hand_public)
|
||||
pile = np.asarray(state.pile)
|
||||
pile_len = np.asarray(state.pile_len)
|
||||
col_top = np.asarray(state.col_top)
|
||||
col_hs = np.asarray(state.col_hs)
|
||||
col_len = np.asarray(state.col_len)
|
||||
|
||||
assert loc.shape == (N_CARDS,)
|
||||
assert np.all((0 <= loc) & (loc <= LOC_DISCARD))
|
||||
assert sum(np.bincount(loc, minlength=LOC_DISCARD + 1)) == N_CARDS
|
||||
assert int(np.asarray(state.just_discarded)) == NO_CARD
|
||||
|
||||
for player in range(2):
|
||||
assert int(np.sum(loc == LOC_P0_HAND + player)) == 8
|
||||
for color in range(N_COLORS):
|
||||
cards = [
|
||||
card
|
||||
for card in range(N_CARDS)
|
||||
if loc[card] == LOC_P0_BOARD + player and card // CARDS_PER_COLOR == color
|
||||
]
|
||||
slots = [card % CARDS_PER_COLOR for card in cards]
|
||||
ranks = [slot - 1 for slot in slots if slot >= 3]
|
||||
assert int(col_len[player, color]) == len(cards)
|
||||
assert int(col_hs[player, color]) == sum(slot < 3 for slot in slots)
|
||||
assert int(col_top[player, color]) == max(ranks, default=0)
|
||||
assert int(col_hs[player, color]) + len(ranks) == int(col_len[player, color])
|
||||
|
||||
for color in range(N_COLORS):
|
||||
length = int(pile_len[color])
|
||||
assert 0 <= length < MAX_PILE_SIZE
|
||||
used = list(map(int, pile[color, :length]))
|
||||
assert all(card != NO_CARD for card in used)
|
||||
assert len(set(used)) == len(used)
|
||||
assert np.all(pile[color, length:] == NO_CARD)
|
||||
discard_cards = [
|
||||
card
|
||||
for card in range(N_CARDS)
|
||||
if loc[card] == LOC_DISCARD and card // CARDS_PER_COLOR == color
|
||||
]
|
||||
assert sorted(used) == sorted(discard_cards)
|
||||
|
||||
in_hand = (loc == LOC_P0_HAND) | (loc == LOC_P0_HAND + 1)
|
||||
assert not np.any(public & ~in_hand)
|
||||
|
||||
|
||||
def assert_states_equal(left: State, right: State) -> None:
|
||||
for left_leaf, right_leaf in zip(
|
||||
jax.tree_util.tree_leaves(left), jax.tree_util.tree_leaves(right), strict=False
|
||||
):
|
||||
np.testing.assert_array_equal(np.asarray(left_leaf), np.asarray(right_leaf))
|
||||
@@ -0,0 +1,85 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import numpy as np
|
||||
|
||||
from lost_cities_jax import (
|
||||
OBS_DIM,
|
||||
batched_legal_mask,
|
||||
batched_obs,
|
||||
batched_reset,
|
||||
batched_step,
|
||||
step,
|
||||
)
|
||||
|
||||
|
||||
def test_batched_step_jits_without_retrace_and_matches_vectorized_step():
|
||||
batch_size = int(os.environ.get("LOST_CITIES_JAX_JIT_BATCH", "8192"))
|
||||
keys = jax.random.split(jax.random.PRNGKey(123), batch_size)
|
||||
states = batched_reset(keys)
|
||||
masks = batched_legal_mask(states)
|
||||
actions = jnp.argmax(masks, axis=1).astype(jnp.int32)
|
||||
|
||||
trace_count = {"value": 0}
|
||||
|
||||
def counted_vmap_step(batch_state, batch_action):
|
||||
trace_count["value"] += 1
|
||||
return jax.vmap(step, in_axes=(0, 0))(batch_state, batch_action)
|
||||
|
||||
counted = jax.jit(counted_vmap_step)
|
||||
first = counted(states, actions)
|
||||
first[1].block_until_ready()
|
||||
assert trace_count["value"] == 1
|
||||
|
||||
second = counted(states, actions)
|
||||
second[1].block_until_ready()
|
||||
assert trace_count["value"] == 1
|
||||
|
||||
batched = batched_step(states, actions)
|
||||
direct = jax.vmap(step, in_axes=(0, 0))(states, actions)
|
||||
_assert_step_outputs_equal(batched, direct)
|
||||
|
||||
scalar_step = jax.jit(step)
|
||||
sample_count = min(batch_size, 128)
|
||||
for idx in range(sample_count):
|
||||
scalar_state = jax.tree_util.tree_map(lambda leaf, i=idx: leaf[i], states)
|
||||
scalar = scalar_step(scalar_state, actions[idx])
|
||||
_assert_scalar_matches_batch(scalar, batched, idx)
|
||||
|
||||
|
||||
def test_batched_observation_shape():
|
||||
batch_size = 32
|
||||
states = batched_reset(jax.random.split(jax.random.PRNGKey(456), batch_size))
|
||||
players = jnp.arange(batch_size, dtype=jnp.int32) % 2
|
||||
obs = batched_obs(states, players)
|
||||
assert obs.shape == (batch_size, OBS_DIM)
|
||||
assert obs.dtype == jnp.float32
|
||||
|
||||
|
||||
def _assert_step_outputs_equal(left, right) -> None:
|
||||
left_state, left_reward, left_done = left
|
||||
right_state, right_reward, right_done = right
|
||||
for left_leaf, right_leaf in zip(
|
||||
jax.tree_util.tree_leaves(left_state),
|
||||
jax.tree_util.tree_leaves(right_state),
|
||||
strict=False,
|
||||
):
|
||||
np.testing.assert_array_equal(np.asarray(left_leaf), np.asarray(right_leaf))
|
||||
np.testing.assert_array_equal(np.asarray(left_reward), np.asarray(right_reward))
|
||||
np.testing.assert_array_equal(np.asarray(left_done), np.asarray(right_done))
|
||||
|
||||
|
||||
def _assert_scalar_matches_batch(scalar, batch, idx: int) -> None:
|
||||
scalar_state, scalar_reward, scalar_done = scalar
|
||||
batch_state, batch_reward, batch_done = batch
|
||||
for scalar_leaf, batch_leaf in zip(
|
||||
jax.tree_util.tree_leaves(scalar_state),
|
||||
jax.tree_util.tree_leaves(batch_state),
|
||||
strict=False,
|
||||
):
|
||||
np.testing.assert_array_equal(np.asarray(scalar_leaf), np.asarray(batch_leaf[idx]))
|
||||
np.testing.assert_array_equal(np.asarray(scalar_reward), np.asarray(batch_reward[idx]))
|
||||
np.testing.assert_array_equal(np.asarray(scalar_done), np.asarray(batch_done[idx]))
|
||||
@@ -0,0 +1,94 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import random
|
||||
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import pytest
|
||||
|
||||
from lost_cities_jax import legal_action_mask, reset_from_order, score, step
|
||||
from reference import lost_cities_ref as ref
|
||||
from tests.lost_cities_jax.conftest import game_count
|
||||
from tests.lost_cities_jax.helpers import shuffled_order
|
||||
|
||||
JIT_MASK = jax.jit(legal_action_mask)
|
||||
JIT_STEP = jax.jit(step)
|
||||
JIT_SCORE = jax.jit(score)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_jax_matches_reference_random_legal_policy(pytestconfig, tmp_path):
|
||||
games = game_count(
|
||||
pytestconfig,
|
||||
env_name="LOST_CITIES_JAX_DIFF_GAMES",
|
||||
local=100,
|
||||
ci=100_000,
|
||||
full=1_000_000,
|
||||
)
|
||||
policy_rng = random.Random(20260704)
|
||||
|
||||
for game_idx in range(games):
|
||||
deck_order = shuffled_order(10_000_000 + game_idx)
|
||||
jax_state = reset_from_order(jnp.asarray(deck_order, dtype=jnp.int8))
|
||||
ref_state = ref.reset_from_order(deck_order)
|
||||
actions: list[int] = []
|
||||
|
||||
for step_idx in range(500):
|
||||
jax_mask = list(map(bool, JIT_MASK(jax_state).tolist()))
|
||||
ref_mask = ref.legal_action_mask(ref_state)
|
||||
if jax_mask != ref_mask:
|
||||
_dump_failure(tmp_path, game_idx, step_idx, deck_order, actions)
|
||||
diff = [
|
||||
idx
|
||||
for idx, pair in enumerate(zip(jax_mask, ref_mask, strict=False))
|
||||
if pair[0] != pair[1]
|
||||
]
|
||||
pytest.fail(f"legal mask mismatch game={game_idx} step={step_idx} diff={diff[:20]}")
|
||||
|
||||
if bool(jax_state.done) != ref_state.done:
|
||||
_dump_failure(tmp_path, game_idx, step_idx, deck_order, actions)
|
||||
pytest.fail(f"done mismatch game={game_idx} step={step_idx}")
|
||||
|
||||
if ref_state.done:
|
||||
jax_score = [float(value) for value in JIT_SCORE(jax_state).tolist()]
|
||||
ref_score = ref.score(ref_state)
|
||||
if jax_score != ref_score:
|
||||
_dump_failure(tmp_path, game_idx, step_idx, deck_order, actions)
|
||||
pytest.fail(
|
||||
f"score mismatch game={game_idx} step={step_idx}: {jax_score} != {ref_score}"
|
||||
)
|
||||
break
|
||||
|
||||
legal_actions = [idx for idx, legal in enumerate(ref_mask) if legal]
|
||||
action = legal_actions[policy_rng.randrange(len(legal_actions))]
|
||||
actions.append(action)
|
||||
jax_state, jax_reward, jax_done = JIT_STEP(jax_state, jnp.int32(action))
|
||||
ref_state, ref_reward, ref_done = ref.step(ref_state, action)
|
||||
if bool(jax_done) != ref_done or [float(v) for v in jax_reward.tolist()] != ref_reward:
|
||||
_dump_failure(tmp_path, game_idx, step_idx, deck_order, actions)
|
||||
pytest.fail(f"step result mismatch game={game_idx} step={step_idx}")
|
||||
else:
|
||||
_dump_failure(tmp_path, game_idx, 500, deck_order, actions)
|
||||
pytest.fail(f"game did not finish within guard loop: game={game_idx}")
|
||||
|
||||
|
||||
def _dump_failure(
|
||||
tmp_path,
|
||||
game_idx: int,
|
||||
step_idx: int,
|
||||
deck_order: list[int],
|
||||
actions: list[int],
|
||||
) -> None:
|
||||
path = tmp_path / f"lost_cities_jax_diff_failure_{game_idx}_{step_idx}.json"
|
||||
path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"game_idx": game_idx,
|
||||
"step_idx": step_idx,
|
||||
"deck_order": deck_order,
|
||||
"actions": actions,
|
||||
},
|
||||
indent=2,
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,62 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import jax.numpy as jnp
|
||||
import numpy as np
|
||||
|
||||
from lost_cities_jax import legal_action_mask, reset_from_order, score, step
|
||||
from lost_cities_jax.engine import current_hand_sorted
|
||||
from lost_cities_jax.types import CARDS_PER_COLOR, N_CARDS
|
||||
from tests.lost_cities_jax.helpers import assert_states_equal, first_deck_draw_action
|
||||
|
||||
|
||||
def test_discarded_card_cannot_be_drawn_from_same_pile():
|
||||
state = reset_from_order(jnp.arange(N_CARDS, dtype=jnp.int8))
|
||||
hand = np.asarray(current_hand_sorted(state))
|
||||
first_card_color = int(hand[0]) // CARDS_PER_COLOR
|
||||
discard_then_same_pile = 0 * 12 + 1 * 6 + (first_card_color + 1)
|
||||
discard_then_deck = 0 * 12 + 1 * 6
|
||||
|
||||
mask = np.asarray(legal_action_mask(state), dtype=bool)
|
||||
|
||||
assert not mask[discard_then_same_pile]
|
||||
assert mask[discard_then_deck]
|
||||
|
||||
|
||||
def test_illegal_action_is_noop_with_zero_reward():
|
||||
state = reset_from_order(jnp.arange(N_CARDS, dtype=jnp.int8))
|
||||
hand = np.asarray(current_hand_sorted(state))
|
||||
color = int(hand[0]) // CARDS_PER_COLOR
|
||||
illegal_action = 0 * 12 + 1 * 6 + (color + 1)
|
||||
|
||||
next_state, reward, done = step(state, jnp.int32(illegal_action))
|
||||
|
||||
assert_states_equal(next_state, state)
|
||||
np.testing.assert_array_equal(np.asarray(reward), np.zeros(2, dtype=np.float32))
|
||||
assert bool(done) is False
|
||||
|
||||
|
||||
def test_done_state_step_is_noop_and_has_no_repeated_terminal_reward():
|
||||
state = reset_from_order(jnp.arange(N_CARDS, dtype=jnp.int8))
|
||||
while not bool(state.done):
|
||||
action = first_deck_draw_action(legal_action_mask(state))
|
||||
state, _, _ = step(state, jnp.int32(action))
|
||||
|
||||
after_done, reward, done = step(state, jnp.int32(0))
|
||||
|
||||
assert_states_equal(after_done, state)
|
||||
np.testing.assert_array_equal(np.asarray(reward), np.zeros(2, dtype=np.float32))
|
||||
assert bool(done) is True
|
||||
|
||||
|
||||
def test_deck_only_draws_finish_in_exactly_44_plies():
|
||||
state = reset_from_order(jnp.arange(N_CARDS, dtype=jnp.int8))
|
||||
plies = 0
|
||||
while not bool(state.done):
|
||||
action = first_deck_draw_action(legal_action_mask(state))
|
||||
state, reward, done = step(state, jnp.int32(action))
|
||||
plies += 1
|
||||
|
||||
assert plies == 44
|
||||
assert int(state.draw_ptr) == N_CARDS
|
||||
assert bool(done) is True
|
||||
np.testing.assert_array_equal(np.asarray(reward), np.asarray(score(state)))
|
||||
@@ -0,0 +1,104 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from lost_cities_jax import legal_action_mask, reset_from_order, step
|
||||
from lost_cities_jax.types import (
|
||||
CARDS_PER_COLOR,
|
||||
DISCARD,
|
||||
DRAW_DECK,
|
||||
HAND_SIZE,
|
||||
LOC_P0_HAND,
|
||||
N_ACTIONS,
|
||||
N_CARDS,
|
||||
PLAY,
|
||||
)
|
||||
from tests.lost_cities_jax.conftest import game_count
|
||||
from tests.lost_cities_jax.helpers import shuffled_order
|
||||
|
||||
JIT_MASK = jax.jit(legal_action_mask)
|
||||
JIT_STEP = jax.jit(step)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_legal_mask_matches_slow_bruteforce_checker(pytestconfig):
|
||||
state_count = game_count(
|
||||
pytestconfig,
|
||||
env_name="LOST_CITIES_JAX_MASK_STATES",
|
||||
local=1_000,
|
||||
ci=10_000,
|
||||
full=10_000,
|
||||
)
|
||||
rng = random.Random(31337)
|
||||
state = reset_from_order(jnp.asarray(shuffled_order(30_000_000), dtype=jnp.int8))
|
||||
next_seed = 30_000_001
|
||||
|
||||
for _ in range(state_count):
|
||||
jax_mask = np.asarray(JIT_MASK(state), dtype=bool)
|
||||
slow_mask = _slow_mask(state)
|
||||
np.testing.assert_array_equal(jax_mask, slow_mask)
|
||||
|
||||
if bool(state.done):
|
||||
state = reset_from_order(jnp.asarray(shuffled_order(next_seed), dtype=jnp.int8))
|
||||
next_seed += 1
|
||||
continue
|
||||
|
||||
legal_actions = np.flatnonzero(jax_mask)
|
||||
action = int(legal_actions[rng.randrange(len(legal_actions))])
|
||||
state, _, _ = JIT_STEP(state, jnp.int32(action))
|
||||
|
||||
|
||||
def _slow_mask(state) -> np.ndarray:
|
||||
mask = np.zeros(N_ACTIONS, dtype=bool)
|
||||
if bool(state.done):
|
||||
return mask
|
||||
|
||||
loc = np.asarray(state.card_loc)
|
||||
pile = np.asarray(state.pile)
|
||||
pile_len = np.asarray(state.pile_len)
|
||||
col_top = np.asarray(state.col_top)
|
||||
player = int(state.to_move)
|
||||
just_discarded = int(state.just_discarded)
|
||||
hand_loc = LOC_P0_HAND + player
|
||||
hand = sorted(card for card in range(N_CARDS) if loc[card] == hand_loc)
|
||||
|
||||
for hand_slot in range(HAND_SIZE):
|
||||
if hand_slot >= len(hand):
|
||||
continue
|
||||
card = hand[hand_slot]
|
||||
color = card // CARDS_PER_COLOR
|
||||
slot = card % CARDS_PER_COLOR
|
||||
is_handshake = slot < 3
|
||||
rank = slot - 1
|
||||
|
||||
for place_type in (PLAY, DISCARD):
|
||||
if place_type == PLAY:
|
||||
place_ok = (
|
||||
col_top[player, color] == 0 if is_handshake else rank > col_top[player, color]
|
||||
)
|
||||
else:
|
||||
place_ok = True
|
||||
if not place_ok:
|
||||
continue
|
||||
|
||||
for draw_source in range(6):
|
||||
if draw_source == DRAW_DECK:
|
||||
draw_ok = True
|
||||
else:
|
||||
src = draw_source - 1
|
||||
same_discard_pile = place_type == DISCARD and src == color
|
||||
after_len = int(pile_len[src]) + int(same_discard_pile)
|
||||
if after_len == 0:
|
||||
draw_ok = False
|
||||
else:
|
||||
after_top = card if same_discard_pile else int(pile[src, after_len - 1])
|
||||
just = card if place_type == DISCARD else just_discarded
|
||||
draw_ok = after_top != just
|
||||
mask[hand_slot * 12 + place_type * 6 + draw_source] = draw_ok
|
||||
|
||||
return mask
|
||||
@@ -0,0 +1,66 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from lost_cities_jax import legal_action_mask, reset_from_order, score, step
|
||||
from tests.lost_cities_jax.conftest import game_count
|
||||
from tests.lost_cities_jax.helpers import (
|
||||
assert_state_consistent,
|
||||
assert_states_equal,
|
||||
shuffled_order,
|
||||
)
|
||||
|
||||
JIT_MASK = jax.jit(legal_action_mask)
|
||||
JIT_STEP = jax.jit(step)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_random_trajectories_preserve_state_invariants(pytestconfig):
|
||||
games = game_count(
|
||||
pytestconfig,
|
||||
env_name="LOST_CITIES_JAX_PROPERTY_GAMES",
|
||||
local=50,
|
||||
ci=1_000,
|
||||
full=20_000,
|
||||
)
|
||||
rng = random.Random(777)
|
||||
|
||||
for game_idx in range(games):
|
||||
state = reset_from_order(jnp.asarray(shuffled_order(20_000_000 + game_idx), dtype=jnp.int8))
|
||||
assert_state_consistent(state)
|
||||
|
||||
for _ in range(500):
|
||||
if bool(state.done):
|
||||
after_done, reward, done = JIT_STEP(state, jnp.int32(0))
|
||||
assert_states_equal(after_done, state)
|
||||
np.testing.assert_array_equal(np.asarray(reward), np.zeros(2, dtype=np.float32))
|
||||
assert bool(done) is True
|
||||
break
|
||||
|
||||
mask = np.asarray(JIT_MASK(state), dtype=bool)
|
||||
legal_actions = np.flatnonzero(mask)
|
||||
assert len(legal_actions) > 0
|
||||
action = int(legal_actions[rng.randrange(len(legal_actions))])
|
||||
state, reward, done = JIT_STEP(state, jnp.int32(action))
|
||||
assert_state_consistent(state)
|
||||
if bool(done):
|
||||
np.testing.assert_array_equal(np.asarray(reward), np.asarray(score(state)))
|
||||
else:
|
||||
pytest.fail(f"game did not terminate within guard loop: game={game_idx}")
|
||||
|
||||
|
||||
def test_forced_termination_scores_current_board():
|
||||
state = reset_from_order(jnp.arange(60, dtype=jnp.int8))._replace(
|
||||
step_count=jnp.asarray(399, dtype=jnp.int32)
|
||||
)
|
||||
action = int(np.flatnonzero(np.asarray(legal_action_mask(state), dtype=bool))[0])
|
||||
next_state, reward, done = step(state, jnp.int32(action))
|
||||
|
||||
assert bool(done) is True
|
||||
assert int(next_state.step_count) == 400
|
||||
np.testing.assert_array_equal(np.asarray(reward), np.asarray(score(next_state)))
|
||||
@@ -0,0 +1,31 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from lost_cities_jax import board_score
|
||||
from reference import lost_cities_ref as ref
|
||||
from tests.lost_cities_jax.helpers import hs, num, state_with_columns
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("cards", "expected"),
|
||||
[
|
||||
([num(0, 2), num(0, 3), num(0, 7), num(0, 8), num(0, 10)], 10),
|
||||
(
|
||||
[hs(0, 0), hs(0, 1), num(0, 4), num(0, 5), num(0, 6), num(0, 7), num(0, 8), num(0, 10)],
|
||||
80,
|
||||
),
|
||||
([hs(0), num(0, 4), num(0, 6), num(0, 7)], -6),
|
||||
([hs(0)], -40),
|
||||
([], 0),
|
||||
([hs(0, 0), hs(0, 1), hs(0, 2), *[num(0, rank) for rank in range(2, 11)]], 156),
|
||||
([hs(0, 0), hs(0, 1), hs(0, 2), num(0, 2)], -72),
|
||||
],
|
||||
)
|
||||
def test_board_score_examples(cards, expected):
|
||||
state = state_with_columns([cards, [], [], [], []])
|
||||
assert float(board_score(state)[0]) == expected
|
||||
|
||||
ref_state = ref.reset_from_order(list(range(60)))
|
||||
ref_state.board[0][0] = list(cards)
|
||||
assert ref.board_score(ref_state)[0] == expected
|
||||
Reference in New Issue
Block a user