coolguyandClaude Opus 4.8
d1c2cf628e
rollout_steps was pinned to MAX_STEPS (400) while a round actually runs ~50 plies, and finished envs were only reset between updates. 83% of every rollout was spent stepping already-done envs to produce masked-out zeros. Measured at rollout_steps=400: active steps go 17.3% -> 100%, i.e. 5.8x the learner actions per update for the same compute. Resetting in-scan exposes three things that were previously benign: - compute_gae bootstrapped truncated episodes from zero. That was safe only because every episode used to terminate inside the scan; now episodes cross the boundary, so thread V(s_T) through. - rollout_metrics read final_env and summed rewards along the scan axis, both of which assume one episode per slot. With several episodes per slot that silently produces garbage, so aggregate at done boundaries instead. - league assignments were redrawn only between updates, which would pin a slot to one seat/opponent across every episode in a scan. Redraw them on reset. Also anneal potential shaping against learner actions rather than padded scan steps: the old accounting counted the dead steps, so a 5M-step anneal expired within two updates of 250. Any earlier evidence that shaping does not help was gathered with it effectively off. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XBQKgvBbxbheiTF1AVy1Sh
189 lines
7.0 KiB
Python
189 lines
7.0 KiB
Python
"""Phase 0a: GAE bootstrap, in-scan auto-reset, and episode-boundary metrics."""
|
|
|
|
import jax
|
|
import jax.numpy as jnp
|
|
import pytest
|
|
|
|
from lost_cities_jax.engine import reset
|
|
from lost_cities_jax.opponents import policy_by_name
|
|
from lost_cities_jax.ppo import (
|
|
EpisodeEnd,
|
|
JaxPPOConfig,
|
|
Transition,
|
|
compute_gae,
|
|
create_train_state,
|
|
episode_metrics,
|
|
make_rollout_fn,
|
|
shaping_coefficient,
|
|
)
|
|
|
|
|
|
def _cfg(**overrides) -> JaxPPOConfig:
|
|
cfg = JaxPPOConfig()
|
|
cfg.ppo.batch_games = 8
|
|
cfg.ppo.rollout_steps = 120
|
|
cfg.network.hidden_size = 16
|
|
cfg.network.num_layers = 1
|
|
for key, value in overrides.items():
|
|
setattr(cfg.run, key, value)
|
|
return cfg
|
|
|
|
|
|
# --- GAE bootstrap ---------------------------------------------------------
|
|
|
|
|
|
def test_gae_bootstraps_truncated_episode_from_last_value():
|
|
"""A rollout that ends mid-episode must carry V(s_T), not zero."""
|
|
rewards = jnp.zeros((3, 1))
|
|
values = jnp.zeros((3, 1))
|
|
dones = jnp.zeros((3, 1), dtype=bool) # never terminates inside the window
|
|
|
|
_, returns_zero = compute_gae(rewards, values, dones, gamma=1.0, gae_lambda=1.0)
|
|
_, returns_boot = compute_gae(
|
|
rewards, values, dones, gamma=1.0, gae_lambda=1.0, last_value=jnp.array([5.0])
|
|
)
|
|
|
|
# Without a bootstrap the truncated episode looks worthless.
|
|
assert jnp.allclose(returns_zero, 0.0)
|
|
# With it, every step inherits the tail value.
|
|
assert jnp.allclose(returns_boot, 5.0)
|
|
|
|
|
|
def test_gae_ignores_last_value_when_episode_terminates():
|
|
"""A terminal step must not bootstrap past the episode boundary."""
|
|
rewards = jnp.array([[0.0], [1.0]])
|
|
values = jnp.zeros((2, 1))
|
|
dones = jnp.array([[False], [True]])
|
|
|
|
_, returns = compute_gae(
|
|
rewards, values, dones, gamma=1.0, gae_lambda=1.0, last_value=jnp.array([99.0])
|
|
)
|
|
# Terminal step sees only its own reward; the step before it sees that too.
|
|
assert jnp.allclose(returns, jnp.array([[1.0], [1.0]]))
|
|
|
|
|
|
def test_gae_done_flag_cuts_credit_between_episodes():
|
|
"""With in-scan resets two episodes share a slot; credit must not leak."""
|
|
rewards = jnp.array([[1.0], [7.0]])
|
|
values = jnp.zeros((2, 1))
|
|
dones = jnp.array([[True], [False]]) # first step ends episode A
|
|
|
|
_, returns = compute_gae(rewards, values, dones, gamma=1.0, gae_lambda=1.0)
|
|
# Episode A keeps its own reward; episode B's +7 must not flow backwards.
|
|
assert jnp.allclose(returns[0], 1.0)
|
|
assert jnp.allclose(returns[1], 7.0)
|
|
|
|
|
|
# --- episode-boundary metrics ---------------------------------------------
|
|
|
|
|
|
def _episodes(done, episode_return, length) -> EpisodeEnd:
|
|
shape = jnp.asarray(done).shape
|
|
return EpisodeEnd(
|
|
done=jnp.asarray(done),
|
|
episode_return=jnp.asarray(episode_return, dtype=jnp.float32),
|
|
length=jnp.asarray(length, dtype=jnp.int32),
|
|
opened_colors=jnp.zeros(shape, dtype=jnp.int32),
|
|
positive_expeditions=jnp.zeros(shape, dtype=jnp.int32),
|
|
hit_max_steps=jnp.zeros(shape, dtype=bool),
|
|
)
|
|
|
|
|
|
def _transitions(steps: int, games: int) -> Transition:
|
|
zeros = jnp.zeros((steps, games), dtype=jnp.float32)
|
|
false = jnp.zeros((steps, games), dtype=bool)
|
|
return Transition(
|
|
obs=jnp.zeros((steps, games, 1)),
|
|
legal_mask=jnp.zeros((steps, games, 1), dtype=bool),
|
|
action=jnp.zeros((steps, games), dtype=jnp.int32),
|
|
log_prob=zeros,
|
|
value=zeros,
|
|
reward=zeros,
|
|
done=false,
|
|
active=~false,
|
|
actor_mask=~false,
|
|
entropy=zeros,
|
|
play_action=false,
|
|
)
|
|
|
|
|
|
def test_episode_metrics_only_counts_finished_episodes():
|
|
"""Unfinished episodes carry a partial return; they must not be averaged in."""
|
|
# Slot 0 finishes twice (returns 10 and 20); slot 1 never finishes.
|
|
done = jnp.array([[True, False], [False, False], [True, False]])
|
|
episode_return = jnp.array([[10.0, 3.0], [0.0, 6.0], [20.0, 9.0]])
|
|
length = jnp.array([[50, 1], [0, 2], [60, 3]])
|
|
|
|
metrics = episode_metrics(_transitions(3, 2), _episodes(done, episode_return, length))
|
|
|
|
assert float(metrics["episodes_completed"]) == 2.0
|
|
assert float(metrics["return_mean"]) == pytest.approx(15.0)
|
|
assert float(metrics["game_length_mean"]) == pytest.approx(55.0)
|
|
assert int(metrics["game_length_max"]) == 60
|
|
|
|
|
|
def test_episode_metrics_survives_a_rollout_with_no_completions():
|
|
done = jnp.zeros((2, 2), dtype=bool)
|
|
metrics = episode_metrics(
|
|
_transitions(2, 2), _episodes(done, jnp.zeros((2, 2)), jnp.zeros((2, 2)))
|
|
)
|
|
assert float(metrics["episodes_completed"]) == 0.0
|
|
assert float(metrics["return_mean"]) == 0.0 # guarded denominator, not a NaN
|
|
|
|
|
|
# --- in-scan auto-reset ----------------------------------------------------
|
|
|
|
|
|
def test_rollout_never_returns_a_done_env_and_completes_many_episodes():
|
|
"""The whole point of in-scan reset: no dead steps, several games per slot."""
|
|
cfg = _cfg()
|
|
rng = jax.random.PRNGKey(0)
|
|
rng, init_key, reset_key, roll_key = jax.random.split(rng, 4)
|
|
state = create_train_state(cfg, init_key)
|
|
env_state = jax.jit(jax.vmap(reset))(jax.random.split(reset_key, cfg.ppo.batch_games))
|
|
|
|
rollout_fn = make_rollout_fn(cfg, policy_by_name("discard_only"))
|
|
next_env, transitions, last_value, metrics = rollout_fn(
|
|
state, env_state, roll_key, jnp.asarray(0.0, dtype=jnp.float32)
|
|
)
|
|
|
|
# Every slot is mid-episode, never parked on a finished game.
|
|
assert not bool(jnp.any(next_env.done))
|
|
# 120 steps at ~50 plies a game means each of the 8 slots finished at least one.
|
|
assert float(metrics["episodes_completed"]) >= cfg.ppo.batch_games
|
|
# No step is wasted stepping an already-done env.
|
|
assert bool(jnp.all(transitions.active))
|
|
assert int(metrics["active_steps"]) == cfg.ppo.rollout_steps * cfg.ppo.batch_games
|
|
assert last_value.shape == (cfg.ppo.batch_games,)
|
|
|
|
|
|
def test_rollout_game_length_matches_the_rules_clock():
|
|
"""A round is 44 deck draws plus one turn per discard-pile draw."""
|
|
cfg = _cfg()
|
|
rng = jax.random.PRNGKey(1)
|
|
rng, init_key, reset_key, roll_key = jax.random.split(rng, 4)
|
|
state = create_train_state(cfg, init_key)
|
|
env_state = jax.jit(jax.vmap(reset))(jax.random.split(reset_key, cfg.ppo.batch_games))
|
|
|
|
rollout_fn = make_rollout_fn(cfg, policy_by_name("discard_only"))
|
|
_, _, _, metrics = rollout_fn(state, env_state, roll_key, jnp.asarray(0.0, dtype=jnp.float32))
|
|
|
|
# 44 deck draws is the floor; discard-pile draws only ever extend a round.
|
|
assert float(metrics["game_length_mean"]) >= 44.0
|
|
assert float(metrics["max_steps_rate"]) == 0.0
|
|
|
|
|
|
# --- shaping anneal accounting --------------------------------------------
|
|
|
|
|
|
def test_shaping_anneal_tracks_learner_actions_not_padded_steps():
|
|
cfg = _cfg()
|
|
cfg.reward.potential_shaping_initial = 1.0
|
|
cfg.reward.potential_shaping_final = 0.0
|
|
cfg.reward.potential_shaping_anneal_steps = 1_000
|
|
|
|
assert shaping_coefficient(cfg, 0) == pytest.approx(1.0)
|
|
assert shaping_coefficient(cfg, 500) == pytest.approx(0.5)
|
|
assert shaping_coefficient(cfg, 1_000) == pytest.approx(0.0)
|
|
assert shaping_coefficient(cfg, 10_000) == pytest.approx(0.0)
|