""" Stage 1: self-supervised training of encoder + dynamics + decoder on random (state, action, next_state, reward) transitions and multi-step rollouts -- no labels, no oracle, works for any domain implementing `environment.py`'s `Environment` interface. (Value-head training -- "stage 2" -- lives in `verifier.py`, since ConnectX uses on-policy Monte Carlo value learning rather than an oracle-labeled regression: there's no tractable exact solver at the real 7x6 board scale.) """ import torch import torch.nn as nn DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") def sample_legal_action(env, rng, state): legal = [i for i in range(env.num_actions) if env.is_legal(state, i)] return rng.choice(legal) def generate_transitions(env, rng, n_problems=4000, walk_len=6, **problem_kwargs): """Random-walk from random problem starts, collecting (s, a, s', r) along the way -- covers realistic reachable states, not just problem starts.""" transitions = [] for _ in range(n_problems): state, _ = env.random_problem(rng, **problem_kwargs) for _ in range(walk_len): if env.is_solved(state): break a_idx = sample_legal_action(env, rng, state) next_state, reward, done = env.step(state, a_idx) transitions.append((state, a_idx, next_state, reward)) state = next_state if done: break return transitions def generate_rollout_sequences(env, rng, n_problems=4000, k=3, **problem_kwargs): """K-step (states, actions) sequences for the unrolled/open-loop training objective below.""" sequences = [] for _ in range(n_problems): state, _ = env.random_problem(rng, **problem_kwargs) states = [state] actions = [] cur = state for _ in range(k): a_idx = sample_legal_action(env, rng, cur) cur, _, _ = env.step(cur, a_idx) actions.append(a_idx) states.append(cur) sequences.append((states, actions)) return sequences class StateNormalizer: def __init__(self, states): t = torch.tensor(states, dtype=torch.float32) self.mean = t.mean(dim=0) self.std = t.std(dim=0).clamp_min(1e-3) def to(self, device): self.mean = self.mean.to(device) self.std = self.std.to(device) return self def normalize(self, state_tensor): return (state_tensor - self.mean) / self.std def denormalize(self, norm_tensor): return norm_tensor * self.std + self.mean def states_to_tensor(states): return torch.tensor([list(s) for s in states], dtype=torch.float32) def train_stage1(model, normalizer, transitions, sequences=None, k=3, epochs=400, batch_size=256, lr=1e-3): """`sequences` (see generate_rollout_sequences) adds a k-step UNROLLED loss: the dynamics model is chained k times, feeding its own predicted latent back in as input at each step (never re-encoding the real intermediate state). This matches how the model is actually used at inference (chained multi-step search) -- training on 1-step transitions alone leaves compounding rollout error unaddressed.""" state_dim = normalizer.mean.shape[0] states = states_to_tensor([t[0] for t in transitions]).to(DEVICE) actions = torch.tensor([t[1] for t in transitions], dtype=torch.long).to(DEVICE) next_states = states_to_tensor([t[2] for t in transitions]).to(DEVICE) rewards = torch.tensor([t[3] for t in transitions], dtype=torch.float32).to(DEVICE) norm_states = normalizer.normalize(states) norm_next_states = normalizer.normalize(next_states) if sequences is not None: seq_states_raw = torch.stack([states_to_tensor(s) for s, _ in sequences]).to(DEVICE) # [N, k+1, D] seq_actions = torch.tensor([a for _, a in sequences], dtype=torch.long).to(DEVICE) # [N, k] norm_seq_states = normalizer.normalize(seq_states_raw.view(-1, state_dim)).view(seq_states_raw.shape) n = states.shape[0] opt = torch.optim.Adam( list(model.encoder.parameters()) + list(model.dynamics.parameters()) + list(model.decoder.parameters()), lr=lr, ) sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs) mse = nn.MSELoss() for epoch in range(epochs): perm = torch.randperm(n, device=DEVICE) total_loss = 0.0 for i in range(0, n, batch_size): idx = perm[i:i + batch_size] s, a, s_next, r = norm_states[idx], actions[idx], norm_next_states[idx], rewards[idx] z = model.encode(s) z_next_target = model.encode(s_next) pred_next_z, pred_reward = model.imagine_step(z, a) recon = model.reconstruct(z) recon_next_from_dynamics = model.reconstruct(pred_next_z) loss = ( mse(recon, s) + mse(pred_next_z, z_next_target) + mse(recon_next_from_dynamics, s_next) + mse(pred_reward, r) ) opt.zero_grad() loss.backward() opt.step() total_loss += loss.item() * idx.shape[0] if sequences is not None: n_seq = norm_seq_states.shape[0] seq_perm = torch.randperm(n_seq, device=DEVICE) unrolled_total = 0.0 for i in range(0, n_seq, batch_size): sidx = seq_perm[i:i + batch_size] seq_s = norm_seq_states[sidx] # [B, k+1, D] seq_a = seq_actions[sidx] # [B, k] z = model.encode(seq_s[:, 0, :]) loss_unrolled = 0.0 for step in range(k): z, _pred_r = model.imagine_step(z, seq_a[:, step]) target = seq_s[:, step + 1, :] target_z = model.encode(target) loss_unrolled = loss_unrolled + mse(z, target_z) + mse(model.reconstruct(z), target) loss_unrolled = loss_unrolled / k opt.zero_grad() loss_unrolled.backward() opt.step() unrolled_total += loss_unrolled.item() * sidx.shape[0] sched.step() if (epoch + 1) % 20 == 0 or epoch == 0: msg = f" [stage1] epoch {epoch+1:3d}/{epochs} loss={total_loss/n:.4f}" if sequences is not None: msg += f" unrolled_loss={unrolled_total/n_seq:.4f}" print(msg) @torch.no_grad() def eval_stage1(model, normalizer, transitions): """Decoder reconstruction / dynamics-rollout exact-match accuracy (after rounding to the nearest integer) -- the diagnostic that checks the latent isn't collapsing to something the decoder can't read back out.""" states = states_to_tensor([t[0] for t in transitions]).to(DEVICE) actions = torch.tensor([t[1] for t in transitions], dtype=torch.long).to(DEVICE) next_states = states_to_tensor([t[2] for t in transitions]).to(DEVICE) norm_states = normalizer.normalize(states) z = model.encode(norm_states) recon = normalizer.denormalize(model.reconstruct(z)) pred_next_z, _ = model.imagine_step(z, actions) recon_next = normalizer.denormalize(model.reconstruct(pred_next_z)) recon_acc = (recon.round() == states).all(dim=1).float().mean().item() dyn_acc = (recon_next.round() == next_states).all(dim=1).float().mean().item() return {"decoder_recon_exact_acc": recon_acc, "dynamics_rollout_exact_acc": dyn_acc} @torch.no_grad() def eval_multistep_rollout(model, normalizer, sequences, k): """Chained (open-loop) rollout exact-match accuracy at each step 1..k -- exposes compounding error the way a single-step eval can't.""" state_dim = normalizer.mean.shape[0] seq_states_raw = torch.stack([states_to_tensor(s) for s, _ in sequences]).to(DEVICE) seq_actions = torch.tensor([a for _, a in sequences], dtype=torch.long).to(DEVICE) norm_seq_states = normalizer.normalize(seq_states_raw.view(-1, state_dim)).view(seq_states_raw.shape) z = model.encode(norm_seq_states[:, 0, :]) results = {} for step in range(k): z, _ = model.imagine_step(z, seq_actions[:, step]) recon = normalizer.denormalize(model.reconstruct(z)) real = seq_states_raw[:, step + 1, :] acc = (recon.round() == real).all(dim=1).float().mean().item() results[f"k={step+1}"] = acc return results