coorl-lost-cities/tests/lost_cities_jax/test_match_ppo.py

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