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