87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
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
|