Add polished Lost Cities table GUI

pyproject.toml

@@ -33,6 +33,7 @@ wandb = [
 lost-cities-classic = "coolrl_lost_cities.games.classic:main"
 lost-cities-eval = "coolrl_lost_cities.games.classic.evaluation:main"
 lost-cities-classic-gui = "coolrl_lost_cities.games.classic.pygame_pvp:main"
+lost-cities-play = "coolrl_lost_cities.games.classic.pygame_table:main"
 lost-cities-deep-cfr = "coolrl_lost_cities.games.classic.deep_cfr.cli:main"
 lost-cities-ismcts = "coolrl_lost_cities.games.classic.ismcts.cli:main"
 lost-cities-jax-ppo = "lost_cities_jax.ppo_cli:main"

src/coolrl_lost_cities/games/classic/jax_ppo_policy.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+from collections.abc import Mapping, Sequence
 from pathlib import Path
 from typing import Any
 
@@ -36,7 +37,13 @@ def _jax_card_id(card: Card, handshake_copy: int = 0) -> int:
     return card.color * 12 + slot
 
 
-def snapshot_to_jax_state(snapshot: Snapshot) -> tuple[State, list[int]]:
+PublicHandCounts = Sequence[Mapping[tuple[int, int], int]]
+
+
+def snapshot_to_jax_state(
+    snapshot: Snapshot,
+    public_hand_counts: PublicHandCounts | None = None,
+) -> tuple[State, list[int]]:
     """Convert a classic card-phase snapshot into the equivalent JAX state.
 
     The returned list maps each JAX sorted hand slot back to the classic hand
@@ -67,16 +74,22 @@ def snapshot_to_jax_state(snapshot: Snapshot) -> tuple[State, list[int]]:
         return _jax_card_id(card, copy)
 
     card_loc = np.full(N_CARDS, LOC_DECK, dtype=np.int8)
+    hand_public = np.zeros(N_CARDS, dtype=bool)
     zone_ids: list[int] = []
     hand_pairs: list[tuple[int, int]] = []
 
     deck_ids = [allocate(card) for card in snapshot.deck]
     zone_ids.extend(deck_ids)
     for player, hand in enumerate(snapshot.hands):
+        public_remaining = dict(public_hand_counts[player]) if public_hand_counts else {}
         for classic_slot, card in enumerate(hand):
             card_id = allocate(card)
             card_loc[card_id] = LOC_P0_HAND + player
             zone_ids.append(card_id)
+            face = (card.color, 0 if card.is_handshake else card.numeric_value(2))
+            if public_remaining.get(face, 0) > 0:
+                hand_public[card_id] = True
+                public_remaining[face] -= 1
             if player == snapshot.current_player:
                 hand_pairs.append((card_id, classic_slot))
 
@@ -120,7 +133,7 @@ def snapshot_to_jax_state(snapshot: Snapshot) -> tuple[State, list[int]]:
         deck_order=jnp.asarray(deck_order),
         draw_ptr=jnp.asarray(draw_ptr, dtype=jnp.int32),
         card_loc=jnp.asarray(card_loc),
-        hand_public=jnp.zeros((N_CARDS,), dtype=jnp.bool_),
+        hand_public=jnp.asarray(hand_public),
         col_top=jnp.asarray(col_top),
         col_hs=jnp.asarray(col_hs),
         col_len=jnp.asarray(col_len),
@@ -149,7 +162,11 @@ class JaxPPOPolicy:
         self.pending_draw: int | None = None
         self.last_evaluation: Any | None = None
 
-    def act(self, obs_or_state: Any) -> int:
+    def act(
+        self,
+        obs_or_state: Any,
+        public_hand_counts: PublicHandCounts | None = None,
+    ) -> int:
         if not isinstance(obs_or_state, GameState):
             raise TypeError("JAX PPO GUI policy requires a GameState")
         if obs_or_state.phase == "draw":
@@ -160,7 +177,7 @@ class JaxPPOPolicy:
             return draw
 
         snapshot = snapshot_from_state(obs_or_state)
-        jax_state, hand_slot_map = snapshot_to_jax_state(snapshot)
+        jax_state, hand_slot_map = snapshot_to_jax_state(snapshot, public_hand_counts)
         result = evaluate_agent_policy(
             self.cfg,
             self.params,
@@ -176,4 +193,4 @@ class JaxPPOPolicy:
         return 2 * hand_slot_map[jax_hand_slot] + place
 
 
-__all__ = ["JaxPPOPolicy", "snapshot_to_jax_state"]
+__all__ = ["JaxPPOPolicy", "PublicHandCounts", "snapshot_to_jax_state"]

src/lost_cities_jax/ppo.py

@@ -1216,7 +1216,10 @@ def save_checkpoint(path: Path, state: TrainState, cfg: JaxPPOConfig) -> None:
 
 
 def restore_checkpoint(path: Path, state: TrainState) -> TrainState:
-    return ocp.PyTreeCheckpointer().restore(path.resolve(), item=state)
+    # Explicit restore_args keep per-leaf shardings; the bare item= form loses
+    # them on this orbax version and fails with "sharding ... Got None".
+    restore_args = ocp.checkpoint_utils.construct_restore_args(state)
+    return ocp.PyTreeCheckpointer().restore(path.resolve(), item=state, restore_args=restore_args)
 
 
 def cli_main(argv: list[str] | None = None) -> None:

tests/games/classic/test_jax_ppo_policy.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from collections import Counter
+
 import numpy as np
 from coolrl_lost_cities.games.classic.game import GameState, classic_config
 
@@ -58,3 +60,20 @@ def test_jax_policy_splits_atomic_action_across_classic_phases(monkeypatch) -> N
     assert state.phase == "draw"
     assert policy.act(state) == expected_draw
     assert policy.pending_draw is None
+
+
+def test_snapshot_conversion_marks_known_discard_draw_in_public_hand() -> None:
+    state = GameState.new_game(classic_config(), seed=19)
+    snapshot = snapshot_from_state(state)
+    card = snapshot.hands[1][0]
+    face = (card.color, 0 if card.is_handshake else card.numeric_value(2))
+
+    converted, _ = snapshot_to_jax_state(
+        snapshot,
+        [Counter(), Counter({face: 1})],
+    )
+
+    public = np.asarray(converted.hand_public, dtype=bool)
+    locations = np.asarray(converted.card_loc)
+    assert public.sum() == 1
+    assert np.all(locations[public] == 2)

tests/games/classic/test_pygame_table.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+import pytest
+from coolrl_lost_cities.games.classic.game import GameState, classic_config
+
+pytest.importorskip("pygame")
+
+from coolrl_lost_cities.games.classic import pygame_table  # noqa: E402
+
+
+def test_argparser_defaults_to_final_candidate() -> None:
+    args = pygame_table.build_argparser().parse_args([])
+
+    assert args.checkpoint == str(pygame_table.DEFAULT_CHECKPOINT)
+    assert args.seed is None
+    assert args.width == 1600
+    assert args.height == 1000
+    assert args.offline is False
+
+
+def test_argparser_accepts_overrides() -> None:
+    args = pygame_table.build_argparser().parse_args(
+        ["--checkpoint", "/tmp/ckpt", "--seed", "3", "--offline"]
+    )
+
+    assert args.checkpoint == "/tmp/ckpt"
+    assert args.seed == 3
+    assert args.offline is True
+
+
+def test_one_suit_per_classic_color() -> None:
+    config = classic_config()
+    assert len(pygame_table.SUITS) == config.n_colors
+    assert len(pygame_table.SUIT_PAINTERS) == config.n_colors
+
+
+def test_card_key_distinguishes_wagers_from_numbers() -> None:
+    config = classic_config()
+    state = GameState.new_game(config, seed=0)
+    keys = [pygame_table.card_key(card, config.min_rank) for card in state.deck]
+    wagers = [key for key in keys if key[1] == 0]
+    numbers = [key for key in keys if key[1] != 0]
+    assert all(config.min_rank <= value <= config.max_rank for _, value in numbers)
+    assert all(0 <= color < config.n_colors for color, _ in keys)
+    assert wagers  # a shuffled deck slice still contains wager cards
+
+
+def test_layout_scales_between_window_sizes() -> None:
+    small = pygame_table.Layout(1024, 720, 5)
+    large = pygame_table.Layout(1920, 1200, 5)
+
+    for layout in (small, large):
+        assert layout.board_top < layout.discard_cy < layout.board_bottom
+        assert len(layout.col_x) == 5
+        slots = layout.hand_slots(8)
+        assert len(slots) == 8
+        assert slots[0][0] >= 0
+        assert slots[-1][0] + layout.hand_w <= layout.w
+    assert large.hand_w >= small.hand_w
+
+
+def _settle(app: pygame_table.TableApp) -> None:
+    now = pygame_table.pygame.time.get_ticks() + 10_000
+    app._assign_targets(now)
+    for sprite in app.sprites:
+        sprite.pos.update(sprite.target)
+        sprite.release_at = 0
+        sprite.pending_flip_at = None
+
+
+def test_discard_draw_stays_public_and_survives_undo_redo() -> None:
+    app = pygame_table.TableApp(seed=5, offline=True, headless=True)
+    try:
+        app.input_locked_until = 0
+        # Human discards, draws from deck, then the rival takes that public card.
+        human_discard = next(i for i in app.state.unified_legal_actions() if i % 2 == 1)
+        discarded = app.state.hands[app.human_seat][human_discard // 2]
+        app.apply_unified(human_discard)
+        app.apply_unified(app.state.card_action_size)
+
+        rival_discard = next(
+            i
+            for i in app.state.unified_legal_actions()
+            if i % 2 == 1 and app.state.hands[app.ai_seat][i // 2].color != discarded.color
+        )
+        app.apply_unified(rival_discard)
+        draw_discard = app.state.card_action_size + 1 + discarded.color
+        assert app._legal(draw_discard)
+        app.apply_unified(draw_discard)
+
+        face = pygame_table.card_key(discarded, app.min_rank)
+        assert app.public_hand_counts[app.ai_seat][face] == 1
+        assert any(
+            sprite.public and (sprite.color, sprite.value) == face
+            for sprite in app.hand_zones[app.ai_seat]
+        )
+
+        before = app._copy_public_counts()
+        app.undo()
+        app.redo()
+        assert app.public_hand_counts == before
+        assert any(
+            sprite.public and sprite.face_up and (sprite.color, sprite.value) == face
+            for sprite in app.hand_zones[app.ai_seat]
+        )
+    finally:
+        app.opponent.shutdown()
+        pygame_table.pygame.quit()
+
+
+def test_missing_checkpoint_is_visible_failure_not_heuristic_fallback(tmp_path) -> None:
+    opponent = pygame_table.Opponent(tmp_path / "missing", None, seed=1, offline=False)
+    try:
+        opponent.load_future.result(timeout=5)
+        assert opponent.policy is None
+        assert opponent.load_error is not None
+        assert opponent.label == "MODEL LOAD FAILED"
+    finally:
+        opponent.shutdown()
+
+
+def test_offline_mode_explicitly_uses_heuristic() -> None:
+    opponent = pygame_table.Opponent(None, None, seed=1, offline=True)
+    try:
+        opponent.load_future.result(timeout=5)
+        assert opponent.policy is not None
+        assert opponent.load_error is None
+        assert "offline" in opponent.label
+    finally:
+        opponent.shutdown()