"""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)