Train on matches: match observation, asymmetric critic, both seats, match reward

Phases 2-4 land together because they all rewrite the same rollout, and doing them in sequence would mean writing it three times. Observation adds the four things a match policy cannot play without: carry as a scalar *and* a binned one-hot (round three is a threshold problem -- "win by 41 or lose" plays nothing like "win by 39" -- and the old score_diff divided by MAX_ABS_SCORE=780, squashing a decisive 50-point lead to 0.06); the round index; whose turn it is, which the single-round observation never carried even though the critic is trained on opponent-turn states; and the deck clock, since a round ends on the last deck draw and players bend that parity by drawing from discard piles. Live points per colour are split by hand / discard pile / unseen, because a discard pile is public and recoverable. The critic is asymmetric: it gets the opponent's hand and the deck in order, on a separate trunk so none of it can reach the logits. A test pins that down -- perturbing the privileged input leaves the policy logits bit-identical while moving the value. Deal luck is what makes a match-terminal reward hard to learn from, and a state-value baseline may condition on anything action-independent. Both seats now train. Self-play ran one network on both sides and stop_gradiented the opponent, throwing away half of every game; each ply now emits a transition per seat, folded into the batch so each seat keeps an independent GAE chain. Reward is the match: rounds one and two only bank into carry, and round three pays tanh(total / terminal_scale). Potential shaping on the running total covers the early sparsity and anneals out. match_eval adds the two measurements the plan turns on: duplicate match play (same deals and coins from both seats -- self-play scores exactly 0.500 with zero mean lead, so the mirroring cancels deal luck exactly) and the carry probe. On an untrained net the probe is flat: 4.97 expeditions opened at a 60-point deficit and at a 60-point lead alike. Breaking that flat line is success criterion 1. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XBQKgvBbxbheiTF1AVy1Sh
2026-07-15 02:23:50 +09:00
parent 62b755c3f0
commit 9ba5a07a9d
4 changed files with 994 additions and 0 deletions
@@ -0,0 +1,194 @@
+"""Phases 2-4: match reward, match observation, asymmetric critic, both-seat training."""
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from lost_cities_jax.match import match_reset, match_score, match_step
+from lost_cities_jax.match_obs import (
+    MATCH_CRITIC_OBS_DIM,
+    MATCH_OBS_DIM,
+    match_critic_observation,
+    match_observation,
+)
+from lost_cities_jax.match_ppo import (
+    create_match_train_state,
+    make_match_rollout_fn,
+    make_match_train_iteration,
+)
+from lost_cities_jax.opponents import random_legal_action
+from lost_cities_jax.ppo import JaxPPOConfig
+from lost_cities_jax.types import LOC_P0_HAND, N_CARDS
+
+
+def _cfg(rollout_steps: int = 700, batch_games: int = 8) -> JaxPPOConfig:
+    cfg = JaxPPOConfig()
+    cfg.ppo.batch_games = batch_games
+    cfg.ppo.rollout_steps = rollout_steps
+    cfg.ppo.minibatches = 4
+    cfg.ppo.epochs = 1
+    cfg.network.hidden_size = 32
+    cfg.network.num_layers = 1
+    return cfg
+
+
+# --- the asymmetric critic -------------------------------------------------
+
+
+def test_the_critic_sees_the_opponents_hand_and_the_actor_does_not():
+    match = match_reset(jax.random.PRNGKey(0))
+    loc = np.asarray(match.round.card_loc)
+    opp_hand = np.flatnonzero(loc == LOC_P0_HAND + 1)
+
+    actor = np.asarray(match_observation(match, jnp.int32(0)))
+    critic = np.asarray(match_critic_observation(match, jnp.int32(0)))
+
+    # Card-major, three channels each, matching the base observation's layout.
+    hidden = critic[MATCH_OBS_DIM:].reshape(N_CARDS, 3)
+    assert np.array_equal(np.flatnonzero(hidden[:, 0]), opp_hand)
+
+    # The actor's view is a strict prefix of it and carries none of that.
+    assert np.array_equal(actor, critic[:MATCH_OBS_DIM])
+    assert critic.shape == (MATCH_CRITIC_OBS_DIM,)
+
+
+def test_privileged_input_cannot_move_the_policy_logits():
+    """Separate trunks, or the critic's view of the deck leaks into play."""
+    cfg = _cfg()
+    state = create_match_train_state(cfg, jax.random.PRNGKey(1))
+    match = match_reset(jax.random.PRNGKey(2))
+
+    obs = match_observation(match, jnp.int32(0))[None, :]
+    critic_a = match_critic_observation(match, jnp.int32(0))[None, :]
+    critic_b = critic_a.at[0, MATCH_OBS_DIM:].set(1.0)  # a totally different hidden state
+
+    logits_a, value_a = state.apply_fn(state.params, obs, critic_a)
+    logits_b, value_b = state.apply_fn(state.params, obs, critic_b)
+
+    assert np.array_equal(np.asarray(logits_a), np.asarray(logits_b))
+    # ...and the value head must actually be using it, or the critic is pointless.
+    assert not np.allclose(np.asarray(value_a), np.asarray(value_b))
+
+
+# --- the match reward ------------------------------------------------------
+
+
+def test_reward_is_zero_sum_and_only_paid_at_the_end_of_the_match():
+    cfg = _cfg()
+    state = create_match_train_state(cfg, jax.random.PRNGKey(3))
+    env = jax.jit(jax.vmap(match_reset))(
+        jax.random.split(jax.random.PRNGKey(4), cfg.ppo.batch_games)
+    )
+
+    rollout = make_match_rollout_fn(cfg)
+    # Shaping off, so any non-zero reward must be a terminal one.
+    _, transitions, _, metrics = rollout(state, env, jax.random.PRNGKey(5), jnp.asarray(0.0))
+
+    reward = np.asarray(transitions.reward)  # (T, 2 * batch)
+    seat0, seat1 = reward[:, : cfg.ppo.batch_games], reward[:, cfg.ppo.batch_games :]
+    assert np.allclose(seat0, -seat1)
+
+    done = np.asarray(transitions.done)[:, : cfg.ppo.batch_games]
+    paid = seat0 != 0.0
+    assert np.array_equal(paid, done & paid)  # never paid on a non-terminal ply
+    assert float(metrics["matches_completed"]) > 0
+    assert np.abs(seat0[paid]).max() <= 1.0  # tanh-bounded
+
+
+def test_shaping_tracks_the_running_match_total():
+    """Phi is carry + board diff, so shaping must follow the total, not the round."""
+    match = match_reset(jax.random.PRNGKey(6))
+    key = jax.random.PRNGKey(7)
+
+    prev = np.asarray(match_score(match))
+    for _ in range(200):
+        if bool(match.done):
+            break
+        key, step_key = jax.random.split(key)
+        action = random_legal_action(match.round, match.round.to_move, step_key)
+        match, _, _ = match_step(match, action)
+        total = np.asarray(match_score(match))
+        # The total never resets when a round rolls over; it only accumulates.
+        assert total.shape == (2,)
+        prev = total
+    assert prev.shape == (2,)
+
+
+# --- both seats ------------------------------------------------------------
+
+
+def test_exactly_one_seat_acts_per_ply_but_both_are_trained():
+    cfg = _cfg()
+    state = create_match_train_state(cfg, jax.random.PRNGKey(8))
+    env = jax.jit(jax.vmap(match_reset))(
+        jax.random.split(jax.random.PRNGKey(9), cfg.ppo.batch_games)
+    )
+
+    rollout = make_match_rollout_fn(cfg)
+    _, transitions, last_value, _ = rollout(state, env, jax.random.PRNGKey(10), jnp.asarray(0.0))
+
+    batch = cfg.ppo.batch_games
+    actor = np.asarray(transitions.actor_mask)
+    seat0, seat1 = actor[:, :batch], actor[:, batch:]
+
+    # One mover per ply while the match is live.
+    live = np.asarray(transitions.active)[:, :batch]
+    assert np.array_equal(seat0 ^ seat1, live)
+
+    # But the critic gets both seats: every live ply is a value target on both.
+    assert np.asarray(transitions.active).sum() == 2 * live.sum()
+    assert transitions.value.shape == (cfg.ppo.rollout_steps, 2 * batch)
+    assert last_value.shape == (2 * batch,)
+
+
+def test_self_play_is_balanced_and_a_train_step_stays_finite():
+    cfg = _cfg(rollout_steps=700, batch_games=16)
+    state = create_match_train_state(cfg, jax.random.PRNGKey(11))
+    env = jax.jit(jax.vmap(match_reset))(jax.random.split(jax.random.PRNGKey(12), 16))
+    train_iteration = make_match_train_iteration(cfg)
+
+    state, env, _, metrics = train_iteration(state, env, jax.random.PRNGKey(13), jnp.asarray(0.1))
+
+    assert bool(jnp.isfinite(metrics["loss"]))
+    assert float(metrics["matches_completed"]) > 0
+    # One network on both sides: neither seat should be favoured.
+    assert 0.2 < float(metrics["match_win_rate"]) < 0.8
+
+
+# --- the match observation -------------------------------------------------
+
+
+def test_carry_reaches_the_observation_with_a_usable_scale():
+    """The old score_diff divided by 780; a 50-point lead vanished into 0.06."""
+    match = match_reset(jax.random.PRNGKey(14))
+    behind = match._replace(carry=jnp.asarray([0, 50], dtype=jnp.int32))
+    ahead = match._replace(carry=jnp.asarray([50, 0], dtype=jnp.int32))
+
+    obs_behind = np.asarray(match_observation(behind, jnp.int32(0)))
+    obs_ahead = np.asarray(match_observation(ahead, jnp.int32(0)))
+
+    delta = np.abs(obs_ahead - obs_behind)
+    # A 100-point swing has to be plainly visible, not a rounding error.
+    assert delta.max() > 0.5
+    assert (delta > 0.01).sum() >= 2  # the scalar and at least one bin flip
+
+
+@pytest.mark.parametrize("round_idx", [0, 1, 2])
+def test_the_round_index_is_observable(round_idx):
+    match = match_reset(jax.random.PRNGKey(15))._replace(
+        round_idx=jnp.asarray(round_idx, dtype=jnp.int32)
+    )
+    obs = np.asarray(match_observation(match, jnp.int32(0)))
+    assert obs.shape == (MATCH_OBS_DIM,)
+    assert np.isfinite(obs).all()
+
+
+def test_whose_turn_it_is_is_observable():
+    """The single-round obs never said; the critic had to read it off step_count."""
+    match = match_reset(jax.random.PRNGKey(16))
+    mover = int(match.round.to_move)
+
+    from_mover = np.asarray(match_observation(match, jnp.int32(mover)))
+    from_waiter = np.asarray(match_observation(match, jnp.int32(1 - mover)))
+    assert not np.array_equal(from_mover, from_waiter)