110 lines
3.6 KiB
Python
110 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
import jax
|
|
import jax.numpy as jnp
|
|
|
|
from lost_cities_jax import reset
|
|
from lost_cities_jax.league import (
|
|
LeagueConfig,
|
|
PoolMember,
|
|
active_training_pool,
|
|
load_league_config,
|
|
pool_sampling_probabilities,
|
|
)
|
|
from lost_cities_jax.opponents import policy_by_name
|
|
from lost_cities_jax.ppo import (
|
|
JaxPPOConfig,
|
|
NetworkConfig,
|
|
OpponentConfig,
|
|
PPOHyperConfig,
|
|
RunConfig,
|
|
create_train_state,
|
|
make_league_train_iteration,
|
|
sample_league_assignments,
|
|
)
|
|
|
|
|
|
def tiny_league_ppo_config(tmp_path) -> JaxPPOConfig:
|
|
return JaxPPOConfig(
|
|
run=RunConfig(
|
|
experiment_name="pytest-jax-ppo-league",
|
|
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_league_config_loads():
|
|
cfg = load_league_config("configs/jax_ppo/league-v1.yaml")
|
|
assert cfg.experiment_name == "jax-ppo-league-v1"
|
|
assert cfg.cycles == 5
|
|
assert len(cfg.anchors) == 7
|
|
assert any(member.name == "heuristic_expert" for member in cfg.anchors)
|
|
|
|
|
|
def test_pfsp_probabilities_cap_stalling_anchors():
|
|
cfg = LeagueConfig(
|
|
base_config="configs/jax_ppo/ladder-v2-expert.yaml",
|
|
warm_start_checkpoint="unused",
|
|
stalling_anchor_floor=0.02,
|
|
stalling_anchor_cap=0.05,
|
|
uniform_mix=0.1,
|
|
)
|
|
pool = [
|
|
PoolMember("discard_only", "static", stalling=True, recent_win_rate=0.01),
|
|
PoolMember("heuristic_expert", "static", recent_win_rate=0.5),
|
|
PoolMember("snapshot", "checkpoint", anchor=False, recent_win_rate=0.2),
|
|
]
|
|
probs = pool_sampling_probabilities(pool, cfg)
|
|
assert abs(sum(probs) - 1.0) < 1.0e-6
|
|
assert 0.02 <= probs[0] <= 0.05
|
|
|
|
|
|
def test_active_training_pool_keeps_anchors_and_hard_nonanchors():
|
|
cfg = LeagueConfig(
|
|
base_config="configs/jax_ppo/ladder-v2-expert.yaml",
|
|
warm_start_checkpoint="unused",
|
|
max_active_pool_members=4,
|
|
)
|
|
pool = [
|
|
PoolMember("anchor_a", "static", anchor=True),
|
|
PoolMember("anchor_b", "static", anchor=True),
|
|
PoolMember("easy", "checkpoint", anchor=False, recent_win_rate=0.9),
|
|
PoolMember("hard", "checkpoint", anchor=False, recent_win_rate=0.1),
|
|
PoolMember("medium", "checkpoint", anchor=False, recent_win_rate=0.5),
|
|
]
|
|
active = active_training_pool(pool, cfg)
|
|
assert [member.name for member in active] == ["anchor_a", "anchor_b", "hard", "medium"]
|
|
|
|
|
|
def test_one_jitted_league_train_iteration_shapes(tmp_path):
|
|
cfg = tiny_league_ppo_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))
|
|
probs = jnp.asarray([0.5, 0.5], dtype=jnp.float32)
|
|
assignments = sample_league_assignments(jax.random.PRNGKey(2), 8, probs, 0.5)
|
|
train_iteration = make_league_train_iteration(
|
|
cfg,
|
|
[policy_by_name("discard_only"), policy_by_name("heuristic_expert")],
|
|
probs,
|
|
0.5,
|
|
)
|
|
|
|
train_state, env_state, assignments, rng, metrics = train_iteration(
|
|
train_state,
|
|
env_state,
|
|
assignments,
|
|
jax.random.PRNGKey(3),
|
|
jnp.asarray(0.0, dtype=jnp.float32),
|
|
)
|
|
|
|
assert env_state.to_move.shape == (8,)
|
|
assert assignments.learner_seat.shape == (8,)
|
|
assert rng.shape == (2,)
|
|
assert "opened_colors_mean" in metrics
|