coorl-lost-cities/src/lost_cities_jax/match_ppo.py

"""PPO over three-round matches.

Three things differ from the single-round trainer beyond the env swap:

**Asymmetric actor-critic.** The critic never plays, so it is fed the hidden
state (opponent's hand, deck in order) while the actor keeps the masked view. A
state-value baseline may condition on anything action-independent without biasing
the policy gradient, and deal luck is exactly what makes a match-terminal reward
hard to learn from. The trunks are separate so the privileged features cannot
reach the logits.

**Both seats train.** In self-play one network picks both players' moves, but the
single-round trainer stop_gradients the opponent seat and throws those plies away
-- half of every game. Here every ply emits a transition for both seats, each
with its own seat-relative view, value and reward; the actor mask keeps the
policy loss on the seat that actually moved, while the critic learns from both.
Folding the seat axis into the batch keeps each seat's GAE chain independent.

**Reward is the match total, paid densely.** Each ply pays the points by which it
moved the running match difference; at gamma=1 that telescopes to the final total
difference, so the objective is exactly the rulebook's -- score more across the
three rounds.

This used to be ``tanh(total / terminal_scale)`` at the end of round three, to
optimise P(win match) rather than expected total. That was over-engineered. The
two differ only in risk attitude, and risk attitude is worth almost nothing here:
a policy told to gamble when it trails by 20 going into round three *loses* to a
greedy clone over 6144 duplicate matches (0.482), and one that gambles only when
it trails by 40 breaks even (0.498). Lost Cities' variance levers -- a marginal
wager, a marginal expedition -- cost more expected margin than the convexity they
buy. Measured ceiling on the whole carry-conditioning idea: under one win-rate
point. The dense linear reward also gives credit assignment every ply instead of
one bounded number per ~160.

``carry`` stays in the observation: it costs nothing, it is the state the
start-player rule keys off, and the critic reads it cleanly.
"""

from __future__ import annotations

import json
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import NamedTuple

import flax.linen as nn
import jax
import jax.numpy as jnp
import optax

from lost_cities_jax.match import (
    MatchState,
    match_legal_action_mask,
    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.obs import observation
from lost_cities_jax.ppo import (
    EpisodeEnd,
    JaxPPOConfig,
    TrainState,
    _append_jsonl,
    _broadcast_done,
    _create_run_dir,
    _write_json,
    action_log_prob,
    categorical_entropy,
    compute_gae,
    mask_logits,
    masked_mean,
    restore_checkpoint,
    save_checkpoint,
    shaping_coefficient,
)
from lost_cities_jax.types import MAX_STEPS, N_ACTIONS, OBS_DIM, PLAY

N_SEATS = 2


@dataclass
class Ablation:
    """Switches for measuring what each piece of the match stack is worth."""

    privileged_critic: bool = True
    """On: the critic also sees the opponent's hand and the deck in order.

    Its worth flips with scale, so do not trust a cheap ablation here. Ablated at
    39.3M learner actions, switching it OFF *won* 0.5160 [0.505, 0.527] -- the
    obvious reading being that a critic fitting V(full state) rather than
    E[return | masked obs] hands the actor advantage noise it cannot act on. But
    head to head at 131M, both sides trained identically, OFF *loses*: 0.4633
    [0.453, 0.474]. The privileged critic earns its keep once there is enough data
    to fit it. Defaults on for that reason, and the small-scale result is kept as
    a warning about ablating at a budget you do not intend to ship.
    """

    both_seats: bool = True
    """Off: the opponent seat's plies are discarded, as the old trainer did."""

    match_obs: bool = True
    """Off: the actor gets the plain single-round observation, carry and all the
    rest of the match features withheld."""

    def actor_obs_dim(self) -> int:
        return MATCH_OBS_DIM if self.match_obs else OBS_DIM

    def critic_obs_dim(self) -> int:
        if not self.privileged_critic:
            return self.actor_obs_dim()
        return MATCH_CRITIC_OBS_DIM

    def label(self) -> str:
        parts = []
        if self.privileged_critic:
            parts.append("privileged_critic")
        if not self.both_seats:
            parts.append("no-both_seats")
        if not self.match_obs:
            parts.append("no-match_obs")
        return "default" if not parts else "+".join(parts)


FULL = Ablation()


class MatchTransition(NamedTuple):
    obs: jax.Array
    critic_obs: jax.Array
    legal_mask: jax.Array
    action: jax.Array
    log_prob: jax.Array
    value: jax.Array
    reward: jax.Array
    done: jax.Array
    active: jax.Array
    actor_mask: jax.Array
    entropy: jax.Array
    play_action: jax.Array


class MatchActorCritic(nn.Module):
    """Separate trunks, so the critic's privileged input cannot reach the logits."""

    hidden_size: int = 512
    num_layers: int = 3

    @nn.compact
    def __call__(self, obs: jax.Array, critic_obs: jax.Array) -> tuple[jax.Array, jax.Array]:
        x = obs
        for _ in range(self.num_layers):
            x = nn.relu(nn.Dense(self.hidden_size)(x))
        logits = nn.Dense(N_ACTIONS)(x)

        v = critic_obs
        for _ in range(self.num_layers):
            v = nn.relu(nn.Dense(self.hidden_size)(v))
        value = nn.Dense(1)(v)
        return logits, jnp.squeeze(value, axis=-1)


def create_match_train_state(
    cfg: JaxPPOConfig, rng: jax.Array, ablation: Ablation = FULL
) -> TrainState:
    model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
    params = model.init(
        rng,
        jnp.zeros((1, ablation.actor_obs_dim()), dtype=jnp.float32),
        jnp.zeros((1, ablation.critic_obs_dim()), dtype=jnp.float32),
    )
    tx = optax.chain(
        optax.clip_by_global_norm(cfg.ppo.max_grad_norm),
        optax.adam(cfg.ppo.learning_rate),
    )
    return TrainState.create(apply_fn=model.apply, params=params, tx=tx)


def match_lead(env: MatchState) -> jax.Array:
    """Running match total from seat 0's side, in raw points."""
    totals = match_score(env)
    return (totals[0] - totals[1]).astype(jnp.float32)


def reset_done_matches(env: MatchState, rng: jax.Array, batch: int) -> MatchState:
    fresh = jax.vmap(match_reset)(jax.random.split(rng, batch))
    return jax.tree_util.tree_map(
        lambda old, new: jnp.where(_broadcast_done(env.done, old), new, old), env, fresh
    )


def _seat_views(env: MatchState, seats: jax.Array, ablation: Ablation = FULL):
    actor_fn = match_observation if ablation.match_obs else (lambda m, p: observation(m.round, p))
    obs = jax.vmap(lambda s: jax.vmap(actor_fn, in_axes=(0, None))(env, s))(seats)
    if not ablation.privileged_critic:
        return obs, obs
    critic = jax.vmap(lambda s: jax.vmap(match_critic_observation, in_axes=(0, None))(env, s))(
        seats
    )
    return obs, critic


def _episode_end(env: MatchState, terminal: jax.Array, episode_return: jax.Array) -> EpisodeEnd:
    lead = jax.vmap(match_lead)(env)
    opened = jnp.sum(env.round.col_len[:, 0, :] > 0, axis=-1)
    return EpisodeEnd(
        done=terminal,
        episode_return=episode_return,
        length=env.round.step_count,
        opened_colors=opened,
        positive_expeditions=lead.astype(jnp.int32),  # final match lead, in points
        hit_max_steps=env.round.step_count >= MAX_STEPS,
    )


def match_episode_metrics(transitions: MatchTransition, episodes: EpisodeEnd) -> dict:
    done = episodes.done.astype(jnp.float32)
    completed = jnp.sum(done)
    denom = jnp.maximum(completed, 1.0)
    actor_count = jnp.sum(transitions.actor_mask)
    lead = episodes.positive_expeditions.astype(jnp.float32)
    return {
        "return_mean": jnp.sum(episodes.episode_return * done) / denom,
        "match_lead_mean": jnp.sum(lead * done) / denom,
        "match_win_rate": jnp.sum((lead > 0).astype(jnp.float32) * done) / denom,
        "final_round_length_mean": jnp.sum(episodes.length.astype(jnp.float32) * done) / denom,
        "max_steps_rate": jnp.sum(episodes.hit_max_steps.astype(jnp.float32) * done) / denom,
        "play_action_rate": jnp.sum(transitions.play_action) / jnp.maximum(actor_count, 1),
        "entropy_mean": masked_mean(
            transitions.entropy, transitions.actor_mask.astype(jnp.float32)
        ),
        "learner_actions": actor_count,
        "matches_completed": completed,
    }


def make_match_rollout_fn(cfg: JaxPPOConfig, ablation: Ablation = FULL):
    batch = cfg.ppo.batch_games
    seats = jnp.arange(N_SEATS, dtype=jnp.int32)

    @jax.jit
    def rollout_fn(state: TrainState, env: MatchState, rng: jax.Array, shaping_coef: jax.Array):
        def body(carry, _):
            env, episode_return, key = carry
            key, act_key, reset_key = jax.random.split(key, 3)

            obs, critic_obs = _seat_views(env, seats, ablation)
            legal = jax.vmap(match_legal_action_mask)(env)
            legal_both = jnp.broadcast_to(legal, (N_SEATS, batch, N_ACTIONS))

            logits, value = state.apply_fn(state.params, obs, critic_obs)
            masked = mask_logits(logits, legal_both)
            sampled = jax.random.categorical(act_key, masked, axis=-1).astype(jnp.int32)
            log_prob = jax.vmap(action_log_prob)(masked, sampled)
            entropy = jax.vmap(categorical_entropy)(masked)

            to_move = env.round.to_move.astype(jnp.int32)
            active = ~env.done
            actor_mask = (seats[:, None] == to_move[None, :]) & active[None, :]
            if not ablation.both_seats:
                # The old trainer kept one seat and stop_gradiented the other.
                keep = (seats[:, None] == 0) & jnp.ones_like(active)[None, :]
                actor_mask = actor_mask & keep
            actions = jnp.take_along_axis(sampled, to_move[None, :], axis=0)[0]

            before = jax.vmap(match_lead)(env)
            next_env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, actions)
            after = jax.vmap(match_lead)(next_env)

            done = next_env.done
            terminal = active & done
            # Every ply pays the points by which it moved the match total. At
            # gamma=1 that telescopes to the final total difference, so the
            # objective is exactly "score more across the three rounds" -- just
            # handed out densely instead of once, 150 plies later.
            reward0 = jnp.where(active, (after - before) / cfg.reward.terminal_scale, 0.0)
            # Zero-sum: seat 1 sees exactly the negation.
            reward = jnp.stack([reward0, -reward0], axis=0)

            played = ((actions % 12) // 6 == PLAY)[None, :] & actor_mask
            transition = MatchTransition(
                obs=obs,
                critic_obs=critic_obs,
                legal_mask=legal_both,
                action=jnp.broadcast_to(actions, (N_SEATS, batch)),
                log_prob=log_prob,
                value=value,
                reward=reward,
                done=jnp.broadcast_to(done, (N_SEATS, batch)),
                active=(
                    jnp.broadcast_to(active, (N_SEATS, batch))
                    if ablation.both_seats
                    else jnp.broadcast_to(active, (N_SEATS, batch)) & (seats[:, None] == 0)
                ),
                actor_mask=actor_mask,
                entropy=entropy,
                play_action=played,
            )

            episode_return = episode_return + reward0
            episode = _episode_end(next_env, terminal, episode_return)
            next_env = reset_done_matches(next_env, reset_key, batch)
            episode_return = jnp.where(done, 0.0, episode_return)
            return (next_env, episode_return, key), (transition, episode)

        init_return = jnp.zeros((batch,), dtype=jnp.float32)
        (next_env, _, rng), (transitions, episodes) = jax.lax.scan(
            body, (env, init_return, rng), xs=None, length=cfg.ppo.rollout_steps
        )

        # Fold the seat axis into the batch: each column is then an independent
        # episode chain, which is exactly what GAE wants.
        transitions = jax.tree_util.tree_map(
            lambda x: x.reshape((x.shape[0], N_SEATS * batch, *x.shape[3:])), transitions
        )
        final_obs, final_critic = _seat_views(next_env, seats, ablation)
        _, last_value = state.apply_fn(state.params, final_obs, final_critic)
        last_value = last_value.reshape((N_SEATS * batch,))

        return next_env, transitions, last_value, match_episode_metrics(transitions, episodes)

    return rollout_fn


def match_ppo_update(
    state: TrainState,
    transitions: MatchTransition,
    advantages: jax.Array,
    returns: jax.Array,
    rng: jax.Array,
    cfg: JaxPPOConfig,
) -> tuple[TrainState, dict[str, jax.Array]]:
    batch_size = int(transitions.reward.size)
    minibatch_size = batch_size // cfg.ppo.minibatches
    flat = jax.tree_util.tree_map(lambda x: x.reshape((-1, *x.shape[2:])), transitions)
    flat_advantages = advantages.reshape((batch_size,))
    flat_returns = returns.reshape((batch_size,))

    actor_mask = flat.actor_mask.astype(jnp.float32)
    adv_mean = masked_mean(flat_advantages, actor_mask)
    adv_std = jnp.sqrt(masked_mean((flat_advantages - adv_mean) ** 2, actor_mask) + 1.0e-8)
    flat_advantages = (flat_advantages - adv_mean) / adv_std

    def loss_fn(params, mb, mb_adv, mb_returns):
        logits, value = state.apply_fn(params, mb.obs, mb.critic_obs)
        masked_logits = mask_logits(logits, mb.legal_mask)
        log_prob = action_log_prob(masked_logits, mb.action)
        entropy = categorical_entropy(masked_logits)
        ratio = jnp.exp(log_prob - mb.log_prob)
        actor_weight = mb.actor_mask.astype(jnp.float32)
        active_weight = mb.active.astype(jnp.float32)
        unclipped = ratio * mb_adv
        clipped = jnp.clip(ratio, 1.0 - cfg.ppo.clip_epsilon, 1.0 + cfg.ppo.clip_epsilon) * mb_adv
        actor_loss = -masked_mean(jnp.minimum(unclipped, clipped), actor_weight)
        value_loss = masked_mean((mb_returns - value) ** 2, active_weight)
        entropy_loss = masked_mean(entropy, actor_weight)
        loss = actor_loss + cfg.ppo.value_coef * value_loss - cfg.ppo.entropy_coef * entropy_loss
        return loss, {
            "loss": loss,
            "actor_loss": actor_loss,
            "value_loss": value_loss,
            "entropy_loss": entropy_loss,
            "approx_kl": masked_mean(mb.log_prob - log_prob, actor_weight),
        }

    def epoch_update(carry, key):
        train_state = carry
        permutation = jax.random.permutation(key, batch_size)

        def minibatch_update(carry, idx):
            train_state, accum = carry
            mb_idx = jax.lax.dynamic_slice(permutation, (idx * minibatch_size,), (minibatch_size,))
            mb = jax.tree_util.tree_map(lambda x: x[mb_idx], flat)
            (_, metrics), grads = jax.value_and_grad(loss_fn, has_aux=True)(
                train_state.params, mb, flat_advantages[mb_idx], flat_returns[mb_idx]
            )
            train_state = train_state.apply_gradients(grads=grads)
            accum = jax.tree_util.tree_map(lambda a, b: a + b, accum, metrics)
            return (train_state, accum), None

        zeros = {
            key: jnp.array(0.0)
            for key in ("loss", "actor_loss", "value_loss", "entropy_loss", "approx_kl")
        }
        (train_state, metrics), _ = jax.lax.scan(
            minibatch_update, (train_state, zeros), jnp.arange(cfg.ppo.minibatches)
        )
        metrics = jax.tree_util.tree_map(lambda x: x / cfg.ppo.minibatches, metrics)
        return train_state, metrics

    keys = jax.random.split(rng, cfg.ppo.epochs)
    state, metrics = jax.lax.scan(epoch_update, state, keys)
    return state, jax.tree_util.tree_map(lambda x: jnp.mean(x, axis=0), metrics)


def make_match_train_iteration(cfg: JaxPPOConfig, ablation: Ablation = FULL):
    rollout_fn = make_match_rollout_fn(cfg, ablation)

    @jax.jit
    def train_iteration(
        state: TrainState, env: MatchState, rng: jax.Array, shaping_coef: jax.Array
    ):
        rng, rollout_key, update_key = jax.random.split(rng, 3)
        env, transitions, last_value, stats = rollout_fn(state, env, rollout_key, shaping_coef)
        advantages, returns = compute_gae(
            transitions.reward,
            transitions.value,
            transitions.done,
            cfg.ppo.gamma,
            cfg.ppo.gae_lambda,
            last_value,
        )
        state, update_metrics = match_ppo_update(
            state, transitions, advantages, returns, update_key, cfg
        )
        return state, env, rng, {**stats, **update_metrics}

    return train_iteration


def match_train(cfg: JaxPPOConfig, *, resume: str | None = None, ablation: Ablation = FULL) -> Path:
    run_dir = _create_run_dir(cfg)
    _write_json(run_dir / "config.json", asdict(cfg))
    metrics_path = run_dir / "metrics.jsonl"

    rng = jax.random.PRNGKey(cfg.run.seed)
    rng, init_key, reset_key = jax.random.split(rng, 3)
    state = create_match_train_state(cfg, init_key, ablation)
    if resume:
        state = restore_checkpoint(Path(resume), state)

    env = jax.jit(jax.vmap(match_reset))(jax.random.split(reset_key, cfg.ppo.batch_games))
    train_iteration = make_match_train_iteration(cfg, ablation)

    start = time.perf_counter()
    learner_actions_seen = 0
    for update in range(cfg.run.total_updates):
        shaping_coef = jnp.asarray(shaping_coefficient(cfg, learner_actions_seen), jnp.float32)
        iter_start = time.perf_counter()
        state, env, rng, metrics = train_iteration(state, env, rng, shaping_coef)
        jax.tree_util.tree_leaves(metrics)[0].block_until_ready()

        row = {
            key: float(value) if getattr(value, "shape", ()) == () else value.tolist()
            for key, value in metrics.items()
        }
        learner_actions_seen += int(row["learner_actions"])
        row.update(
            update=update,
            shaping_coef=float(shaping_coef),
            learner_actions_total=learner_actions_seen,
            iteration_seconds=time.perf_counter() - iter_start,
            elapsed_seconds=time.perf_counter() - start,
        )
        _append_jsonl(metrics_path, row)
        if update % cfg.run.log_every == 0:
            print(
                json.dumps(
                    {
                        "update": update,
                        "match_win_rate": row["match_win_rate"],
                        "match_lead_mean": row["match_lead_mean"],
                        "matches_completed": row["matches_completed"],
                        "entropy_mean": row["entropy_mean"],
                    },
                    sort_keys=True,
                ),
                flush=True,
            )
        if cfg.run.checkpoint_every and (update + 1) % cfg.run.checkpoint_every == 0:
            save_checkpoint(run_dir / "latest", state, cfg)

    save_checkpoint(run_dir / "latest", state, cfg)
    return run_dir