connectx/verifier.py

"""
On-policy Monte Carlo value-head training -- no oracle, no bootstrapping,
no target network. Walk the environment with the current value head's own
epsilon-greedy policy, and label every visited state along a walk that
actually reached a solved state with its REALIZED return (what happened,
not a network's own possibly-wrong estimate of what happens next). This is
the only value-training method used for ConnectX, since the real 7x6 board
has no tractable exact solver to regress against.

`unsolved_penalty` is what makes this work for an ADVERSARIAL domain
specifically: `env.step` can return `done=True` on a LOSS (the opponent
won) or a draw, not just our own win. Discarding every one of those walks
(the natural default for a single-agent puzzle, where "unsolved" just means
"further away") would mean the value head never sees a single labeled
example of "this leads to losing" -- exactly the signal an adversarial
domain needs to learn to avoid bad moves.
"""
import collections

import torch
import torch.nn as nn

from .train_utils import states_to_tensor, DEVICE


def _observed_states_to_tensor(env, states):
    return states_to_tensor([env.observe(s) for s in states])


def _greedy_walk_action(env, model, normalizer, state, rng, epsilon):
    """Epsilon-greedy action choice using the value head's own 1-step
    lookahead, scored the same way `search.plan_action`'s depth=1 does
    (`-reward + value`, using the REAL reward/next-state from `env.step`,
    not an imagined one) -- so training optimizes for the exact decision
    rule actually used at inference."""
    legal = [a for a in range(env.num_actions) if env.is_legal(state, a)]
    if rng.random() < epsilon:
        return rng.choice(legal)
    next_states, rewards = [], []
    for a in legal:
        ns, r, _done = env.step(state, a)
        next_states.append(ns)
        rewards.append(r)
    z = model.encode(normalizer.normalize(_observed_states_to_tensor(env, next_states).to(DEVICE)))
    values = model.evaluate(z)
    rewards_t = torch.tensor(rewards, dtype=torch.float32, device=DEVICE)
    scores = -rewards_t + values
    return legal[torch.argmin(scores).item()]


def generate_mc_walks(env, model, normalizer, rng, n_problems, max_steps, epsilon,
                       unsolved_penalty=None, **problem_kwargs):
    """Complete walks (up to max_steps) using the current value head's
    epsilon-greedy policy. A walk that reaches `is_solved` labels every
    visited state with its real steps-remaining. A walk that doesn't
    (a loss, a draw, or simply running out of steps) is discarded UNLESS
    `unsolved_penalty` is set, in which case every state along it gets that
    fixed, uniform label instead -- deliberately not scaled by how early or
    late the walk went wrong; a state that's actually fine keeps
    reappearing in OTHER (solved) walks too, so its label self-corrects
    over many rounds rather than staying pinned to one bad walk's worst
    case."""
    labeled = []
    for _ in range(n_problems):
        state, _ = env.random_problem(rng, **problem_kwargs)
        path_states = [state]
        for _ in range(max_steps):
            if env.is_solved(state):
                break
            a = _greedy_walk_action(env, model, normalizer, state, rng, epsilon)
            state, _reward, done = env.step(state, a)
            path_states.append(state)
            if done:
                break
        if env.is_solved(path_states[-1]):
            T = len(path_states) - 1
            for t, s in enumerate(path_states):
                labeled.append((s, float(T - t)))
        elif unsolved_penalty is not None:
            for s in path_states:
                labeled.append((s, float(unsolved_penalty)))
    return labeled


def train_mc_value_onpolicy(env, model, normalizer, rng, n_rounds=15, n_problems_per_round=400,
                             max_steps=10, epochs_per_round=40, lr=1e-3,
                             epsilon_start=1.0, epsilon_end=0.05, warmup_rounds=3,
                             replay_capacity=20000, min_replay_before_train=30,
                             unsolved_penalty=None, verbose_every=5, **problem_kwargs):
    """`epsilon_start=1.0` + `warmup_rounds`: a freshly-initialized value
    head's "greedy" choice is pure noise, which can be WORSE than uniform
    random at stumbling into a solved state by chance. Holding epsilon at
    1.0 (pure random walk) for the first `warmup_rounds` guarantees a
    baseline solve rate to bootstrap training from, before annealing toward
    exploitation."""
    opt = torch.optim.Adam(model.value.parameters(), lr=lr)
    buffer = collections.deque(maxlen=replay_capacity)

    for round_idx in range(n_rounds):
        if round_idx < warmup_rounds:
            epsilon = 1.0
        else:
            progress = (round_idx - warmup_rounds) / max(1, n_rounds - 1 - warmup_rounds)
            epsilon = epsilon_start + (epsilon_end - epsilon_start) * progress
        labeled = generate_mc_walks(env, model, normalizer, rng, n_problems_per_round, max_steps, epsilon,
                                     unsolved_penalty=unsolved_penalty, **problem_kwargs)
        buffer.extend(labeled)
        if len(buffer) < min_replay_before_train:
            if verbose_every:
                print(f"  [mc-onpolicy] round {round_idx+1}/{n_rounds}  epsilon={epsilon:.2f}  "
                      f"only {len(buffer)} labeled states so far (need {min_replay_before_train}) -- skipping fit")
            continue

        all_data = list(buffer)
        states_t = _observed_states_to_tensor(env, [s for s, _ in all_data]).to(DEVICE)
        returns_t = torch.tensor([r for _, r in all_data], dtype=torch.float32, device=DEVICE)

        model.value_target_mean.copy_(returns_t.mean())
        model.value_target_std.copy_(returns_t.std().clamp(min=1e-3))
        returns_norm = (returns_t - model.value_target_mean) / model.value_target_std

        with torch.no_grad():
            z_states = model.encode(normalizer.normalize(states_t))

        n = len(all_data)
        for _epoch in range(epochs_per_round):
            perm = torch.randperm(n, device=DEVICE)
            pred_norm = model.value(z_states[perm])
            loss = nn.functional.mse_loss(pred_norm, returns_norm[perm])
            opt.zero_grad()
            loss.backward()
            opt.step()

        if verbose_every and (round_idx + 1) % verbose_every == 0:
            print(f"  [mc-onpolicy] round {round_idx+1}/{n_rounds}  epsilon={epsilon:.2f}  "
                  f"buffer_size={n}  loss={loss.item():.4f}  return_mean={model.value_target_mean.item():.2f}")