105 lines
3.2 KiB
Python
105 lines
3.2 KiB
Python
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
|