scripts/train.py

"""
Trains a WorldModel checkpoint on the real Kaggle ConnectX board (7x6,
win_len=4). No oracle exists at this scale (see env.py's BFS_MAX_CELLS), so
this is a genuinely no-oracle recipe throughout:

1. Stage 1 (dynamics + decoder): self-supervised, `train_utils.train_stage1`
   on random transitions/rollouts -- never depends on an oracle for any
   domain.
2. Stage 2 (value head): `verifier.train_mc_value_onpolicy`, on-policy
   Monte Carlo, `unsolved_penalty=max_steps` -- this is an adversarial
   domain, where a training walk can end in a LOSS, not just run out of
   steps; without this penalty every losing walk is silently discarded and
   the value head never learns to avoid losing moves at all (confirmed as
   the direct cause of a real 0%, worse-than-random win rate before this
   was added).
3. Self-play fine-tune: each round freezes the current model as an
   opponent (viewed from the opponent's side, via `_swap_agent_opponent`)
   and trains further against a small, capped POOL of past frozen
   snapshots (not just the latest one) mixed with the original fixed
   heuristic -- every opponent used above is fixed and non-learning, which
   is the actual ceiling on how strong a non-self-play policy can get.

Evaluation has no oracle-ceiling comparison to report against (none exists
at this scale) -- instead: win rate against the ORIGINAL fixed opponent
(deterministic, `opponent_epsilon=0`) the model is ultimately graded
against, plus a random-legal-play win rate as the required floor.
"""
import copy
import random

import torch

from connectx.env import ConnectXEnv, AGENT, OPPONENT, EMPTY, _encode_board
from connectx.model import WorldModel
from connectx.train_utils import DEVICE, StateNormalizer, generate_transitions, generate_rollout_sequences, \
    train_stage1, eval_stage1, eval_multistep_rollout
from connectx.verifier import train_mc_value_onpolicy
from connectx.search import evaluate
from connectx.adversarial_search import real_adversarial_plan_action


def _swap_agent_opponent(cells):
    """AGENT<->OPPONENT relabeling, EMPTY unchanged -- lets a model that
    was only ever trained to play as AGENT evaluate a position from the
    OTHER side's perspective, by pretending that side is AGENT instead."""
    return [c if c == EMPTY else (OPPONENT if c == AGENT else AGENT) for c in cells]


def make_selfplay_pool_opponent_fn(frozen_pool, frozen_normalizer, opponent_env):
    """Builds an `opponent_policy_fn(cells) -> column` (see env.py's
    `opponent_policy_fn` extension point) that plays using a frozen
    snapshot of this same architecture's own trained judgment, not a
    hand-written heuristic. Each call picks a snapshot uniformly at random
    from `frozen_pool` (a small population of past snapshots, not just the
    latest one -- a coarse approximation of real population-based self-
    play/fictitious play, so the live policy can't narrowly overfit to
    counter-play against whatever the single latest snapshot happens to
    do) and uses the real adversarial search (rounds=1 -- this function is
    called on the order of 100K+ times across a training run, so a slower
    multi-round search would balloon total training time)."""

    def opponent_policy_fn(cells):
        frozen_model = random.choice(frozen_pool)
        swapped_state = tuple(_encode_board(_swap_agent_opponent(cells)))
        with torch.no_grad():
            a = real_adversarial_plan_action(opponent_env, frozen_model, frozen_normalizer, swapped_state, rounds=1)
        return a

    return opponent_policy_fn


def random_baseline_win_rate(env, problems, max_steps, rng):
    """Required control: an agent picking uniformly among its legal
    non-PASS columns. The floor the trained model needs to beat -- not a
    certified ceiling (no oracle exists at this scale), just the honest
    "did training do anything at all" check."""
    wins = 0
    for state, _ in problems:
        cur = state
        for _ in range(max_steps):
            if env.is_solved(cur):
                break
            legal = [a for a in range(env.num_actions) if env.is_legal(cur, a)]
            non_pass = [a for a in legal if a != env.width]
            a = rng.choice(non_pass or legal)
            cur, _r, done = env.step(cur, a)
            if done:
                break
        if env.is_solved(cur):
            wins += 1
    n = len(problems)
    print(f"{'Random-legal-play baseline':30s} win_rate={wins/n:.3f} ({wins}/{n})")
    return wins / n


def main(seed=0, ckpt_path="checkpoints/connectx_checkpoint.pt", latent_dim=96, hidden_dim=256,
         n_problems=2000, walk_len=8, mc_rounds=25, mc_problems_per_round=400,
         opponent_epsilon=0.15, opponent_strong_epsilon=0.0,
         selfplay_rounds=5, selfplay_epsilon=0.4, selfplay_mc_rounds_per_iter=15,
         selfplay_pool_size=5):
    # Two env instances, same board, different opponent determinism: `env`
    # (opponent_epsilon=0, the pure deterministic opponent) is what
    # evaluation is graded against. `train_env` (opponent_epsilon>0) is
    # used ONLY for generating training data -- a perfectly deterministic
    # opponent means every training walk from a matching starting side is
    # the SAME exact game, a real, diagnosed weakness (the model
    # reproducibly lost as first player against this exact opponent while
    # winning as second player).
    env = ConnectXEnv(width=7, height=6, win_len=4)
    train_env = ConnectXEnv(width=7, height=6, win_len=4, opponent_epsilon=opponent_epsilon,
                             opponent_strong_epsilon=opponent_strong_epsilon)
    rng = random.Random(seed)
    torch.manual_seed(seed)
    # Also seed the GLOBAL random module: env.py's opponent_epsilon/
    # opponent_strong_epsilon rolls read `random.random()`/`random.choice()`
    # directly, not this function's own seeded `rng` -- without this,
    # "same seed" runs are silently not reproducible whenever opponent
    # stochasticity is enabled.
    random.seed(seed)
    max_steps = (env.width * env.height) // 2 + 2

    print(f"Domain: ConnectX (real board), {env.width}x{env.height}, win_len={env.win_len}, "
          f"num_actions={env.num_actions}, state_dim={env.state_dim}, "
          f"train opponent_epsilon={opponent_epsilon}, opponent_strong_epsilon={opponent_strong_epsilon}\n")
    assert env.bfs_solve(env.random_problem(rng)[0]) is None, \
        "expected no oracle at real-board scale -- see env.py's BFS_MAX_CELLS"

    print("Generating stage-1 (dynamics) data...")
    train_transitions = generate_transitions(train_env, rng, n_problems=n_problems, walk_len=walk_len)
    val_transitions = generate_transitions(env, rng, n_problems=300, walk_len=walk_len)
    print(f"  {len(train_transitions)} train transitions, {len(val_transitions)} val transitions")

    unroll_k = 4
    train_sequences = generate_rollout_sequences(train_env, rng, n_problems=n_problems, k=unroll_k)
    val_sequences = generate_rollout_sequences(env, rng, n_problems=300, k=unroll_k)
    print(f"  {len(train_sequences)} train sequences, {len(val_sequences)} val sequences (k={unroll_k})")

    all_states_for_norm = [t[0] for t in train_transitions] + [t[2] for t in train_transitions]
    normalizer = StateNormalizer(all_states_for_norm).to(DEVICE)

    model = WorldModel(env.state_dim, env.num_actions, latent_dim=latent_dim, hidden_dim=hidden_dim).to(DEVICE)

    print("\nStage 1: training encoder + dynamics + decoder...")
    train_stage1(model, normalizer, train_transitions, sequences=train_sequences, k=unroll_k)

    print("\nStage-1 val metrics:")
    print(" ", eval_stage1(model, normalizer, val_transitions))
    print(" ", eval_multistep_rollout(model, normalizer, val_sequences, k=unroll_k))

    print("\nStage 2: no-oracle value training (on-policy Monte Carlo, bfs_solve never called)...")
    train_mc_value_onpolicy(train_env, model, normalizer, rng, n_rounds=mc_rounds,
                             n_problems_per_round=mc_problems_per_round, max_steps=max_steps,
                             unsolved_penalty=max_steps)

    # Self-play fine-tune: only the value head trains during
    # train_mc_value_onpolicy (encoder/dynamics/decoder stay fixed from
    # stage 1), so each frozen snapshot's encoder/dynamics are identical
    # to the live model's -- only the value judgment (and therefore the
    # self-play opponent's move choices) differs round to round.
    # `selfplay_pool_size` keeps a capped, small population of past
    # snapshots (drops the oldest once full) rather than only the single
    # latest one, or an unbounded pool that would let early, still-weak
    # snapshots dominate forever.
    frozen_pool = []
    for sp_round in range(selfplay_rounds):
        frozen_model = copy.deepcopy(model).eval()
        for p in frozen_model.parameters():
            p.requires_grad_(False)
        frozen_pool.append(frozen_model)
        if len(frozen_pool) > selfplay_pool_size:
            frozen_pool.pop(0)
        print(f"\nSelf-play fine-tune round {sp_round + 1}/{selfplay_rounds} "
              f"(opponent_selfplay_epsilon={selfplay_epsilon}, pool_size={len(frozen_pool)})...")
        opponent_env = ConnectXEnv(width=7, height=6, win_len=4)  # plain -- used only for is_legal/num_actions
        selfplay_fn = make_selfplay_pool_opponent_fn(frozen_pool, normalizer, opponent_env)
        selfplay_train_env = ConnectXEnv(width=7, height=6, win_len=4,
                                          opponent_epsilon=opponent_epsilon,
                                          opponent_selfplay_epsilon=selfplay_epsilon,
                                          opponent_policy_fn=selfplay_fn)
        train_mc_value_onpolicy(selfplay_train_env, model, normalizer, rng, n_rounds=selfplay_mc_rounds_per_iter,
                                 n_problems_per_round=mc_problems_per_round, max_steps=max_steps,
                                 unsolved_penalty=max_steps)

    if ckpt_path:
        torch.save({
            "model_state": model.state_dict(),
            "norm_mean": normalizer.mean.cpu(),
            "norm_std": normalizer.std.cpu(),
            "state_dim": env.state_dim,
            "num_actions": env.num_actions,
            "latent_dim": latent_dim,
            "hidden_dim": hidden_dim,
            "board_width": env.width,
            "board_height": env.height,
            "win_len": env.win_len,
        }, ckpt_path)
        print(f"\nSaved checkpoint to {ckpt_path}")

    print("\n" + "=" * 20 + " EVALUATION (no oracle -- vs. random-legal-play baseline only) " + "=" * 20)
    eval_rng = random.Random(999)
    problems = [env.random_problem(eval_rng) for _ in range(150)]
    random_baseline_win_rate(env, problems, max_steps, random.Random(1000))
    evaluate(env, model, normalizer, problems, depth=1, beam_width=8, max_total_steps=max_steps,
             label="Baseline A (model, depth=1)")
    evaluate(env, model, normalizer, problems, depth=3, beam_width=8, max_total_steps=max_steps,
             label="Search (model, depth=3)")

    return model, normalizer


if __name__ == "__main__":
    main()