130 lines
4.3 KiB
Python
130 lines
4.3 KiB
Python
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))
|