Add JAX PPO static-opponent trainer

tests/lost_cities_jax/test_ppo_stack.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from lost_cities_jax import legal_action_mask, reset
+from lost_cities_jax.opponents import (
+    discard_only_action,
+    heuristic_balanced_action,
+    heuristic_cautious_action,
+)
+from lost_cities_jax.ppo import (
+    JaxPPOConfig,
+    NetworkConfig,
+    OpponentConfig,
+    PPOHyperConfig,
+    RunConfig,
+    create_train_state,
+    evaluate,
+    load_config,
+    make_train_iteration,
+    policy_by_name,
+    random_rollout,
+    train,
+)
+
+
+def tiny_config(tmp_path) -> JaxPPOConfig:
+    return JaxPPOConfig(
+        run=RunConfig(
+            experiment_name="pytest-jax-ppo",
+            seed=123,
+            total_updates=1,
+            checkpoint_every=1,
+            artifact_root=str(tmp_path),
+        ),
+        opponent=OpponentConfig(name="discard_only"),
+        network=NetworkConfig(hidden_size=32, num_layers=1),
+        ppo=PPOHyperConfig(batch_games=8, rollout_steps=16, epochs=1, minibatches=2),
+    )
+
+
+def test_static_opponents_return_legal_actions():
+    state = reset(jax.random.PRNGKey(0))
+    mask = np.asarray(legal_action_mask(state), dtype=bool)
+    for fn in [discard_only_action, heuristic_balanced_action, heuristic_cautious_action]:
+        action = int(fn(state, jnp.int32(0), jax.random.PRNGKey(1)))
+        assert mask[action]
+
+
+def test_random_rollout_smoke(tmp_path):
+    row = random_rollout(tiny_config(tmp_path))
+    assert row["env_steps"] == 8 * 16
+    assert 0.0 <= row["play_action_rate"] <= 1.0
+    assert row["game_length_mean"] > 0.0
+
+
+def test_train_checkpoint_and_eval_smoke(tmp_path):
+    cfg = tiny_config(tmp_path)
+    run_dir = train(cfg)
+    assert (run_dir / "latest").exists()
+    result = evaluate(cfg, run_dir / "latest", games=8, duplicate=True)
+    assert result["games"] == 16
+    assert 0.0 <= result["win_rate"] <= 1.0
+
+
+def test_one_jitted_train_iteration_shapes(tmp_path):
+    cfg = tiny_config(tmp_path)
+    train_state = create_train_state(cfg, jax.random.PRNGKey(0))
+    env_state = jax.jit(jax.vmap(reset))(jax.random.split(jax.random.PRNGKey(1), 8))
+    train_iteration = make_train_iteration(cfg, policy_by_name("discard_only"))
+
+    train_state, env_state, rng, metrics = train_iteration(
+        train_state, env_state, jax.random.PRNGKey(2), jnp.asarray(1.0)
+    )
+
+    assert env_state.to_move.shape == (8,)
+    assert "loss" in metrics
+    assert jnp.isfinite(metrics["loss"])
+
+
+def test_load_config_file():
+    cfg = load_config("configs/jax_ppo/smoke.yaml")
+    assert cfg.opponent.name == "discard_only"
+    assert cfg.ppo.batch_games == 64