Add JAX PPO opponent to classic GUI

tests/games/classic/test_jax_ppo_policy.py

@@ -0,0 +1,60 @@
+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

tests/games/classic/test_pygame_pvp.py

@@ -18,6 +18,8 @@ def test_gui_argparser_accepts_classic_options() -> None:
             "1024",
             "--height",
             "768",
+            "--jax-checkpoint",
+            "/tmp/final_candidate",
         ]
     )
 
@@ -26,6 +28,7 @@ def test_gui_argparser_accepts_classic_options() -> None:
     assert args.seed == 7
     assert args.width == 1024
     assert args.height == 768
+    assert args.jax_checkpoint == "/tmp/final_candidate"
 
 
 def test_gui_argparser_rejects_removed_backend_option() -> None: