eval_checkpoint.py

#!/usr/bin/env python
"""Evaluate the old PolicyValueNet checkpoint (simple CNN, 15-channel input).

Usage:
    uv run python eval_checkpoint.py checkpoints/sl_test.eqx --games 200
    uv run python eval_checkpoint.py checkpoints/sl_test.eqx --games 100 --opponent both
"""

import argparse
import time

import jax
import jax.numpy as jnp
import jax.random as jrandom
import numpy as np
import equinox as eqx

from generals.core.env import GeneralsEnv
from generals_bot.decision.trainer import observation_to_tensor


# ═══════════════════════════════════════════════════════════════════════
# Old PolicyValueNet (simple CNN, matched to checkpoint structure)
# ═══════════════════════════════════════════════════════════════════════

class PolicyValueNet(eqx.Module):
    """Old-style CNN policy-value network (pre-Padded23UNet)."""

    conv1: eqx.nn.Conv2d
    conv2: eqx.nn.Conv2d
    conv3: eqx.nn.Conv2d
    conv4: eqx.nn.Conv2d
    conv5: eqx.nn.Conv2d
    conv6: eqx.nn.Conv2d
    conv7: eqx.nn.Conv2d
    conv8: eqx.nn.Conv2d
    conv9: eqx.nn.Conv2d
    conv10: eqx.nn.Conv2d
    dir_head: eqx.nn.Conv2d
    split_head: eqx.nn.Conv2d
    pass_head_fc1: eqx.nn.Linear
    pass_head_fc2: eqx.nn.Linear
    value_conv: eqx.nn.Conv2d
    value_fc1: eqx.nn.Linear
    value_fc2: eqx.nn.Linear

    def __init__(self, key, in_channels: int = 15, latent_channels: int = 32):
        keys = jrandom.split(key, 17)
        ch = (48, 48, 48, 48, 48, 48, 32, 32, 16, 16)

        self.conv1 = eqx.nn.Conv2d(in_channels, ch[0], kernel_size=3, padding=1, key=keys[0])
        self.conv2 = eqx.nn.Conv2d(ch[0], ch[1], kernel_size=3, padding=1, key=keys[1])
        self.conv3 = eqx.nn.Conv2d(ch[1], ch[2], kernel_size=3, padding=1, key=keys[2])
        self.conv4 = eqx.nn.Conv2d(ch[2], ch[3], kernel_size=3, padding=1, key=keys[3])
        self.conv5 = eqx.nn.Conv2d(ch[3], ch[4], kernel_size=3, padding=1, key=keys[4])
        self.conv6 = eqx.nn.Conv2d(ch[4], ch[5], kernel_size=3, padding=1, key=keys[5])
        self.conv7 = eqx.nn.Conv2d(ch[5], ch[6], kernel_size=3, padding=1, key=keys[6])
        self.conv8 = eqx.nn.Conv2d(ch[6], ch[7], kernel_size=3, padding=1, key=keys[7])
        self.conv9 = eqx.nn.Conv2d(ch[7], ch[8], kernel_size=3, padding=1, key=keys[8])
        self.conv10 = eqx.nn.Conv2d(ch[8], ch[9], kernel_size=3, padding=1, key=keys[9])

        self.dir_head = eqx.nn.Conv2d(ch[9], 4, kernel_size=1, key=keys[10])
        self.split_head = eqx.nn.Conv2d(ch[9], 1, kernel_size=1, key=keys[11])
        self.pass_head_fc1 = eqx.nn.Linear(ch[9], 64, key=keys[12])
        self.pass_head_fc2 = eqx.nn.Linear(64, 1, key=keys[13])

        self.value_conv = eqx.nn.Conv2d(ch[9], latent_channels, kernel_size=1, key=keys[14])
        self.value_fc1 = eqx.nn.Linear(latent_channels, 64, key=keys[15])
        self.value_fc2 = eqx.nn.Linear(64, 1, key=keys[16])

    def __call__(self, obs, mask, key):
        """obs: (C, H, W), mask: (H, W, 4)"""
        H, W = obs.shape[1], obs.shape[2]
        grid_cells = H * W

        x = jax.nn.relu(self.conv1(obs))
        x = jax.nn.relu(self.conv2(x))
        x = jax.nn.relu(self.conv3(x))
        x = jax.nn.relu(self.conv4(x))
        x = jax.nn.relu(self.conv5(x))
        x = jax.nn.relu(self.conv6(x))
        x = jax.nn.relu(self.conv7(x))
        x = jax.nn.relu(self.conv8(x))
        x = jax.nn.relu(self.conv9(x))
        x = jax.nn.relu(self.conv10(x))

        # Value head
        v = jax.nn.relu(self.value_conv(x))
        v_pooled = v.mean(axis=(1, 2))
        value = self.value_fc2(jax.nn.relu(self.value_fc1(v_pooled)))[0]

        # Direction head
        dir_logits = self.dir_head(x)
        mask_t = jnp.transpose(mask, (2, 0, 1))
        big_penalty = -1e9
        mask_penalty = (1.0 - mask_t) * big_penalty
        dir_logits = dir_logits + mask_penalty

        # Split head
        split_logits = self.split_head(x).squeeze(0)

        # Pass head
        pooled = x.mean(axis=(1, 2))
        pass_logit = self.pass_head_fc2(jax.nn.relu(self.pass_head_fc1(pooled)))[0]

        # Early-game pass mask
        t_norm = obs[12, 0, 0]
        is_early = t_norm < 0.04
        pass_allowed = is_early | (~jnp.any(mask))
        pass_logit = jnp.where(pass_allowed, pass_logit, -1e9)

        # Joint action space: 4·H·W directions + 1 pass
        n_dirs = 4 * grid_cells
        dir_flat = dir_logits.reshape(-1)
        action_logits = jnp.concatenate([dir_flat, pass_logit[None]])

        # Sample
        key, split_key = jrandom.split(key)
        idx = jrandom.categorical(key, action_logits)

        is_pass = idx >= n_dirs
        move_idx = jnp.where(is_pass, 0, idx)
        direction = move_idx // grid_cells
        position = move_idx % grid_cells
        row, col = position // W, position % W

        split_logit_cell = split_logits[row, col]
        split_sample = jrandom.bernoulli(
            split_key, jax.nn.sigmoid(split_logit_cell)
        ).astype(jnp.int32)
        is_split = jnp.where(is_pass, 0, split_sample)

        action = jnp.array(
            [is_pass.astype(jnp.int32), row, col, direction, is_split],
            dtype=jnp.int32,
        )

        # log_prob and entropy
        log_probs = jax.nn.log_softmax(action_logits)
        probs = jax.nn.softmax(action_logits)
        log_p_a = log_probs[idx]

        is_pass_bool = is_pass > 0
        split_logit_cell = split_logits[row, col]
        log_split_1 = jax.nn.log_sigmoid(split_logit_cell)
        log_split_0 = jax.nn.log_sigmoid(-split_logit_cell)
        log_p_split = is_split * log_split_1 + (1 - is_split) * log_split_0
        log_p_split = jnp.where(is_pass_bool, 0.0, log_p_split)
        log_prob = log_p_a + log_p_split

        # Entropy
        cat_entropy = -jnp.sum(probs * log_probs)
        dir_probs = probs[:n_dirs].reshape(4, H, W)
        cell_probs = dir_probs.sum(axis=0)
        p_split = jax.nn.sigmoid(split_logits)
        log_p_split_cell = jax.nn.log_sigmoid(split_logits)
        log_1m_p_split_cell = jax.nn.log_sigmoid(-split_logits)
        bern_entropy = -p_split * log_p_split_cell - (1 - p_split) * log_1m_p_split_cell
        expected_bern = jnp.sum(cell_probs * bern_entropy)
        total_entropy = cat_entropy + expected_bern

        return action, value, log_prob, total_entropy


# ═══════════════════════════════════════════════════════════════════════
# Old-style opponent functions (5-element actions)
# ═══════════════════════════════════════════════════════════════════════

def old_random_action(key, obs):
    """Return 5-element random action [pass, row, col, dir, split]."""
    from generals.core.action import compute_valid_move_mask as jax_move_mask

    mask = jax_move_mask(obs.armies, obs.owned_cells, obs.mountains)
    H, W = obs.armies.shape
    valid = jnp.argwhere(mask, size=H * W * 4, fill_value=-1)
    num_valid = jnp.sum(jnp.all(valid >= 0, axis=-1))

    k1, k2 = jrandom.split(key)
    should_pass = num_valid == 0
    idx = jnp.minimum(
        jrandom.randint(k1, (), 0, jnp.maximum(num_valid, 1)),
        num_valid - 1,
    )
    move = valid[idx]
    is_half = jrandom.randint(k2, (), 0, 2)

    return jnp.array(
        [should_pass, move[0], move[1], move[2], is_half], dtype=jnp.int32
    )


def old_expander_action(key, obs):
    """Return 5-element expander action [pass, row, col, dir, split]."""
    from generals.core.action import compute_valid_move_mask as jax_move_mask

    mask = jax_move_mask(obs.armies, obs.owned_cells, obs.mountains)
    H, W = obs.armies.shape
    valid = jnp.argwhere(mask, size=H * W * 4, fill_value=-1)
    armies = obs.armies
    opponent = obs.opponent_cells

    def score_move(m):
        i, j, d = m[0], m[1], m[2]
        ti = jnp.clip(i + (d == 1).astype(jnp.int32) - (d == 0).astype(jnp.int32), 0, armies.shape[0] - 1)
        tj = jnp.clip(j + (d == 3).astype(jnp.int32) - (d == 2).astype(jnp.int32), 0, armies.shape[1] - 1)
        s = armies[i, j]
        s = jnp.where(opponent[ti, tj], s * 10, s)
        return s

    scores = jax.vmap(score_move)(valid)
    best_idx = jnp.argmax(scores)
    move = valid[best_idx]
    has_valid = jnp.any(mask)
    return jnp.array([
        0,
        jnp.where(has_valid, move[0], 0).astype(jnp.int32),
        jnp.where(has_valid, move[1], 0).astype(jnp.int32),
        jnp.where(has_valid, move[2], 0).astype(jnp.int32),
        0,
    ], dtype=jnp.int32)


# ═══════════════════════════════════════════════════════════════════════
# Evaluation loop
# ═══════════════════════════════════════════════════════════════════════

def evaluate(model, env, num_games, seed, opponent_fn):
    key = jrandom.PRNGKey(seed)

    from generals.core.action import compute_valid_move_mask as jax_move_mask

    @jax.jit
    def _compute_mask(obs):
        return jax_move_mask(obs.armies, obs.owned_cells, obs.mountains)

    @jax.jit
    def _build_tensor(obs):
        return observation_to_tensor(obs)

    stats = {"wins": 0, "losses": 0, "draws": 0, "total": 0}
    num_envs = min(256, num_games)
    games_done = 0
    t_start = time.perf_counter()

    while games_done < num_games:
        n = min(num_envs, num_games - games_done)

        key, k_pool, k_init = jrandom.split(key, 3)
        pool, _ = env.reset(k_pool)
        init_keys = jrandom.split(k_init, n)
        states = jax.vmap(env.init_state)(init_keys)

        step_vmap = jax.vmap(env.step, in_axes=(0, 0, None))

        for _ in range(env.truncation):
            obs_0 = jax.vmap(lambda s: _get_obs(s, 0))(states)
            obs_1 = jax.vmap(lambda s: _get_obs(s, 1))(states)

            # Model actions for player 0
            tensor_0 = jax.vmap(_build_tensor)(obs_0)
            masks_0 = jax.vmap(_compute_mask)(obs_0)
            key, *nks = jrandom.split(key, n + 1)
            nk_arr = jnp.stack(nks)
            actions_0 = jax.vmap(lambda t, m, k: model(t, m, k)[0])(
                tensor_0, masks_0, nk_arr,
            )

            # Opponent actions for player 1
            key, *eks = jrandom.split(key, n + 1)
            ek_arr = jnp.stack(eks)
            actions_1 = jax.vmap(opponent_fn)(ek_arr, obs_1)

            actions = jnp.stack([actions_0, actions_1], axis=1)
            timesteps, states = step_vmap(states, actions, pool)

            done = timesteps.terminated | timesteps.truncated
            new_done = int(jnp.sum(done))
            winners = timesteps.info.winner

            if new_done > 0:
                stats["total"] += new_done
                for w in np.asarray(winners)[np.asarray(done)]:
                    if w == 0:
                        stats["wins"] += 1
                    elif w == 1:
                        stats["losses"] += 1
                    elif w == -1:
                        stats["draws"] += 1
                    games_done += 1

                if games_done >= num_games:
                    break

        if games_done >= num_games:
            break

    elapsed = time.perf_counter() - t_start
    stats["elapsed"] = elapsed
    stats["games_per_sec"] = stats["total"] / max(elapsed, 0.001)
    return stats


def _get_obs(states, player_idx: int):
    from generals.core.game import get_observation
    return get_observation(states, player_idx)


def main():
    parser = argparse.ArgumentParser(
        description="Evaluate old PolicyValueNet checkpoint (simple CNN, 15-channel)"
    )
    parser.add_argument("model_path", type=str, help="Path to .eqx checkpoint")
    parser.add_argument("--opponent", choices=["random", "expander", "both"],
                        default="expander")
    parser.add_argument("--games", type=int, default=200)
    parser.add_argument("--grid", type=str, default="10x10")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--truncation", type=int, default=500)
    args = parser.parse_args()

    h, w = (int(x) for x in args.grid.split("x"))

    # Load checkpoint
    print(f"[Eval] Loading checkpoint from {args.model_path}")
    print(f"[Eval] Grid: {h}x{w}, Games: {args.games}, Truncation: {args.truncation}")

    key = jrandom.PRNGKey(args.seed)
    model = PolicyValueNet(key, in_channels=15, latent_channels=32)
    model = eqx.tree_deserialise_leaves(args.model_path, model)

    param_count = sum(p.size for p in jax.tree.leaves(model) if hasattr(p, 'size'))
    print(f"[Eval] Model loaded ({param_count:,} params)")

    env = GeneralsEnv(grid_dims=(h, w), truncation=args.truncation)
    print(f"[Eval] Environment: {env}")

    opponents = ["expander", "random"] if args.opponent == "both" else [args.opponent]
    opp_fns = {"expander": old_expander_action, "random": old_random_action}

    for opp_name in opponents:
        print(f"\n{'='*55}")
        print(f"  Evaluating vs {opp_name.upper()} ...")
        print(f"{'='*55}")

        result = evaluate(model, env, args.games, args.seed, opp_fns[opp_name])

        total = result["total"]
        wr = 100 * result["wins"] / max(total, 1)
        lr = 100 * result["losses"] / max(total, 1)
        dr = 100 * result["draws"] / max(total, 1)
        print(f"  Model wins:  {result['wins']:5d}  ({wr:.1f}%)")
        print(f"  Opp.  wins:  {result['losses']:5d}  ({lr:.1f}%)")
        print(f"  Draws:       {result['draws']:5d}  ({dr:.1f}%)")
        print(f"  Total:       {total:5d}")
        print(f"  Throughput:  {result['games_per_sec']:.0f} g/s  ({result['elapsed']:.1f}s)")
        print()


if __name__ == "__main__":
    main()