61 lines
2.2 KiB
Python
61 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
from coolrl_lost_cities.games.classic.game import GameState, classic_config
|
|
|
|
from coolrl_lost_cities.games.classic.jax_ppo_policy import (
|
|
JaxPPOPolicy,
|
|
snapshot_to_jax_state,
|
|
)
|
|
from coolrl_lost_cities.games.classic.snapshots import snapshot_from_state
|
|
from lost_cities_jax.engine import board_score, current_hand_sorted, legal_action_mask
|
|
from lost_cities_jax.human_play import PolicyEval
|
|
|
|
|
|
def test_snapshot_to_jax_state_preserves_initial_public_state() -> None:
|
|
state = GameState.new_game(classic_config(), seed=7)
|
|
snapshot = snapshot_from_state(state)
|
|
|
|
converted, hand_slot_map = snapshot_to_jax_state(snapshot)
|
|
|
|
assert int(converted.to_move) == snapshot.current_player
|
|
assert int(converted.draw_ptr) == 16
|
|
assert hand_slot_map == sorted(
|
|
hand_slot_map,
|
|
key=lambda slot: (
|
|
snapshot.hands[0][slot].color,
|
|
snapshot.hands[0][slot].rank,
|
|
),
|
|
)
|
|
assert np.asarray(current_hand_sorted(converted)).shape == (8,)
|
|
assert np.asarray(legal_action_mask(converted)).any()
|
|
assert np.asarray(board_score(converted)).tolist() == [0.0, 0.0]
|
|
|
|
|
|
def test_jax_policy_splits_atomic_action_across_classic_phases(monkeypatch) -> None:
|
|
state = GameState.new_game(classic_config(), seed=11)
|
|
converted, hand_slot_map = snapshot_to_jax_state(snapshot_from_state(state))
|
|
legal = np.flatnonzero(np.asarray(legal_action_mask(converted), dtype=bool))
|
|
atomic_action = int(legal[0])
|
|
expected_slot = hand_slot_map[atomic_action // 12]
|
|
expected_place = (atomic_action % 12) // 6
|
|
expected_draw = atomic_action % 6
|
|
|
|
monkeypatch.setattr(
|
|
"coolrl_lost_cities.games.classic.jax_ppo_policy.evaluate_agent_policy",
|
|
lambda *_args: PolicyEval(action=atomic_action, top3=[], value=0.0),
|
|
)
|
|
policy = object.__new__(JaxPPOPolicy)
|
|
policy.cfg = object()
|
|
policy.params = object()
|
|
policy.model = object()
|
|
policy.pending_draw = None
|
|
policy.last_evaluation = None
|
|
|
|
card_action = policy.act(state)
|
|
assert card_action == 2 * expected_slot + expected_place
|
|
state.apply_action(card_action)
|
|
assert state.phase == "draw"
|
|
assert policy.act(state) == expected_draw
|
|
assert policy.pending_draw is None
|