v5.2: spatial policy head, 6x6 action grid aligned with FSQ token grid, advantage normalization removed (entropy-collapse root cause), grad clip 1.0, standard entropy bonus; smoke-validated end-to-end on CPU

train_imagination_v5.py

@@ -1,638 +1 @@
-#!/usr/bin/env python3
-"""
-train_imagination_v5.py — imagination model for computer use, v5 (actor-representation fix).
-
-History this build answers:
-  v1/fix1b: RL-in-imagination only -> actor had no path from instruction to click location;
-            real-env success 0.00, entropy collapse.
-  v4:       BC added -> still 0.10 success; actor inputs (25 coarse FSQ codes + hashed word
-            buckets) cannot represent "where is the 'Ok' button".
-  v5 fix:   actor sees frozen DINOv2-small patch/CLS features (semantic, position-preserving)
-            + frozen SmolLM2-135M embedding of the REAL instruction text. World model stays a
-            small transformer over FSQ codes; reward ensemble is pessimistic (min of 2 heads).
-  Also fixes the WM copy flaw: CE now predicts codes[:,t+1] from output position t
-  (v2/v3 compared a position's output to the codes fed as input at that same position).
-  v5.1: fix classifier device crash (H on cpu vs clf on cuda), fix dream-vs-real label mixing
-        (real embeddings now come from the pre-roll context window, not post-imagination),
-        quadratic entropy floor, BC-stage checkpoint pushed before RL so a crash loses nothing.
-
-Pipeline: A collect (DOM oracle + random) -> B FSQ tokenizer on frozen DINOv2 patch features
--> C world model (val-split early stop) -> D BC on oracle demos -> E RL in imagination
-(pessimistic min-of-2-heads reward, REINFORCE + critic, entropy target) -> real-env evals ->
-F dream-vs-real classifier (AUC).
-
-Env facts (sandbox-verified 2026-09-19, miniwob 1.1.0):
-  obs keys: dom_elements, fields, screenshot, utterance; screenshot (210,160,3) uint8 (H210,W160)
-  action {"action_type": 2, "coords": float32[x,y]} = CLICK_COORDS over Box(0,[160,210])
-  dom elements: flat left/top/width/height (1-elem float32 arrays), text match is case-sensitive
-Usage:
-  python train_imagination_v5.py --smoke            # tiny CPU run, all phases
-  python train_imagination_v5.py --device cuda      # real run (t4-small, ~1.5h)
-"""
-import argparse, math, os, random, time
-import numpy as np
-import torch, torch.nn as nn, torch.nn.functional as F
-
-os.environ.setdefault("MINIWOB_CHROME_BINARY", "/usr/bin/chromium")
-os.environ.setdefault("MINIWOB_CHROMEDRIVER", "/usr/bin/chromedriver")
-
-import gymnasium, miniwob
-from PIL import Image
-from transformers import AutoModel, AutoTokenizer
-
-gymnasium.register_envs(miniwob)
-
-# ── constants ────────────────────────────────────────────────────────────────
-SCREEN_W, SCREEN_H = 160, 210          # full screenshot, interactive page
-IMG = 224                              # DINOv2 input (square resize)
-PP = 14                                # DINOv2 patch size
-GP = IMG // PP                         # 16x16 = 256 patch tokens
-POOL = 2                               # avg-pool patch grid -> 8x8 = 64 tokens/frame
-NT = (GP // POOL) ** 2                 # 64
-TD, LEVELS = 6, 5                      # FSQ dims/levels -> vocab 5^6 = 15625
-VOCAB = LEVELS ** TD
-GRID = 16                              # 16x16 = 256 click cells
-N_ACTIONS = GRID * GRID
-TEXT_DIM = 576                         # SmolLM2-135M hidden size (verified 576)
-MAX_EP_STEPS = 30
-L_CTX = 3                              # real context frames before imagining
-DM, LAYERS, HEADS = 256, 4, 4
-W_LEN = 16
-PUNCT = str.maketrans("", "", '"\'.,!?:;()')
-
-TASKS = ["click-button-v1", "click-test-2-v1", "click-link-v1"]
-HF_REPO = "rootxhacker/miniwob-imagination"
-
-
-def coords_of(cell):
-    col, row = cell % GRID, cell // GRID
-    return np.array([(col + 0.5) * SCREEN_W / GRID, (row + 0.5) * SCREEN_H / GRID], np.float32)
-
-
-def screen_of(raw):
-    img = raw["screenshot"]  # (210,160,3)
-    return np.asarray(Image.fromarray(img).resize((IMG, IMG)), np.uint8)
-
-
-def instr_text(raw):
-    words = [str(raw["utterance"]).translate(PUNCT)]
-    words += [str(f) for _, f in (raw["fields"] or ()) if f]
-    return " ".join(words)[:256]
-
-
-def oracle_from_raw(raw):
-    """(grid_cell, click_action) for the element named by fields['target'], else None.
-    Exact text match first, then case-insensitive fallback."""
-    fields = dict(raw["fields"]) if raw["fields"] else {}
-    target = str(fields.get("target", "")).strip()
-    if not target:
-        return None
-    best = None
-    for el in raw["dom_elements"]:
-        text = str(el.get("text", "")).strip()
-        if text != target and text.lower() != target.lower():
-            continue
-        l, t = el.get("left"), el.get("top")
-        w, h = el.get("width"), el.get("height")
-        if l is None or t is None:
-            continue
-        l, t, w, h = (float(np.ravel(v)[0]) if v is not None else 0.0 for v in (l, t, w, h))
-        cx, cy = l + w / 2, t + h / 2
-        if not (0 <= cx < SCREEN_W and 0 <= cy < SCREEN_H):
-            continue
-        col = min(GRID - 1, max(0, int(cx * GRID / SCREEN_W)))
-        row = min(GRID - 1, max(0, int(cy * GRID / SCREEN_H)))
-        exact = text == target
-        cand = (row * GRID + col, exact)
-        if best is None or (exact and not best[1]):
-            best = cand
-    if best is None:
-        return None
-    cell = best[0]
-    return cell, {"action_type": 2, "coords": coords_of(cell)}
-
-
-def collect(tasks, n_demo, n_rand, seed):
-    rng = random.Random(seed)
-    trajs = []
-    for task in tasks:
-        env = gymnasium.make(f"miniwob/{task}")
-        for kind, n in (("oracle", n_demo), ("random", n_rand)):
-            rets = []
-            for ep in range(n):
-                raw, _ = env.reset(seed=rng.randrange(2 ** 31))
-                frames, acts, rews, dones = [screen_of(raw)], [], [], []
-                ret = 0.0
-                for t in range(MAX_EP_STEPS):
-                    oa = oracle_from_raw(raw) if kind == "oracle" else None
-                    if oa is not None:
-                        cell, act = oa
-                    else:
-                        cell = rng.randrange(N_ACTIONS)
-                        act = {"action_type": 2, "coords": coords_of(cell)}
-                    raw2, r, term, trunc, _ = env.step(act)
-                    r = float(r); ret += r
-                    acts.append(cell); rews.append(r); dones.append(float(term or trunc))
-                    frames.append(screen_of(raw2)); raw = raw2
-                    if term or trunc:
-                        break
-                rets.append(ret)
-                trajs.append(dict(task=task, instr=instr_text(raw),
-                                  frames=np.stack(frames), actions=np.array(acts, np.int64),
-                                  rewards=np.array(rews, np.float32), dones=np.array(dones, np.float32)))
-            print(f"[collect] {task} {kind}: n={n} avg_return={np.mean(rets):.3f}", flush=True)
-        env.close()
-    return trajs
-
-
-# ── frozen encoders ──────────────────────────────────────────────────────────
-class FrozenEncoders(nn.Module):
-    def __init__(self, device):
-        super().__init__()
-        self.dino = AutoModel.from_pretrained("facebook/dinov2-small").to(device).eval()
-        self.lm = AutoModel.from_pretrained("HuggingFaceTB/SmolLM2-135M").to(device).eval()
-        self.tok = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM2-135M")
-        self.tok.pad_token = self.tok.eos_token  # required for padding on CPU
-        for p in self.parameters():
-            p.requires_grad_(False)
-
-    @torch.no_grad()
-    def patches(self, screens_uint8):
-        """(B,IMG,IMG,3) uint8 -> patch features (B,NT,384) and CLS (B,384)."""
-        x = screens_uint8.float().permute(0, 3, 1, 2) / 127.5 - 1.0
-        out = self.dino(pixel_values=x).last_hidden_state          # (B,257,384)
-        cls = out[:, 0]
-        p = out[:, 1:].reshape(-1, GP, GP, 384)
-        p = p.reshape(-1, GP // POOL, POOL, GP // POOL, POOL, 384).mean((2, 4))    # (B,8,8,384)
-        return p.reshape(-1, NT, 384), cls
-
-    @torch.no_grad()
-    def text(self, texts):
-        texts = [t if t and t.strip() else "click the button" for t in texts]
-        enc = self.lm_tok_batch(texts)
-        out = self.lm(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"]).last_hidden_state
-        m = enc["attention_mask"].unsqueeze(-1).float()
-        return (out * m).sum(1) / m.sum(1).clamp(min=1)            # (B,576) mean-pooled
-
-    def lm_tok_batch(self, texts):
-        enc = self.tok(texts, padding=True, truncation=True, max_length=W_LEN, return_tensors="pt")
-        return {k: v.to(next(self.parameters()).device) for k, v in enc.items()}
-
-
-# ── FSQ tokenizer over frozen patch features ────────────────────────────────
-class FSQTok(nn.Module):
-    """trainable projection 384->TD quantized to 5 levels, + decoder back to 384 (feature recon)."""
-    def __init__(self):
-        super().__init__()
-        self.proj = nn.Linear(384, TD)
-        self.dec = nn.Sequential(nn.Linear(TD, 128), nn.ReLU(), nn.Linear(128, 384))
-
-    def encode(self, feats):                                   # (B,NT,384)
-        z = torch.tanh(self.proj(feats))
-        half = (LEVELS - 1) / 2.0
-        q = (torch.round(z * half) / half).clamp(-1, 1)
-        dims = ((q + 1) * half).round().long()
-        flat = torch.zeros(dims.shape[:2], dtype=torch.long, device=z.device)
-        for d in range(TD):
-            flat = flat * LEVELS + dims[..., d]
-        zq = z + (q - z).detach()
-        return zq, flat
-
-    def forward(self, feats):
-        zq, codes = self.encode(feats)
-        recon = self.dec(zq)
-        return recon, codes, zq
-
-
-# ── world model (small causal transformer over code pairs) ──────────────────
-class Attn(nn.Module):
-    def __init__(self, dm, heads):
-        super().__init__()
-        self.h, self.dk = heads, dm // heads
-        self.qkv, self.proj = nn.Linear(dm, 3 * dm), nn.Linear(dm, dm)
-
-    def forward(self, x, kv=None):
-        B, T, D = x.shape
-        q, k, v = self.qkv(x).split(D, -1)
-        q = q.view(B, T, self.h, self.dk).transpose(1, 2)
-        k = k.view(B, T, self.h, self.dk).transpose(1, 2)
-        v = v.view(B, T, self.h, self.dk).transpose(1, 2)
-        if kv is not None:
-            k = torch.cat([kv[0], k], 2); v = torch.cat([kv[1], v], 2)
-        att = q @ k.transpose(-2, -1) / math.sqrt(self.dk)
-        if kv is None:
-            mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=x.device), 1)
-            att = att.masked_fill(mask[None, None], float("-inf"))
-        out = (att.softmax(-1) @ v).transpose(1, 2).reshape(B, T, D)
-        return self.proj(out), (k, v)
-
-
-class Block(nn.Module):
-    def __init__(self, dm, heads):
-        super().__init__()
-        self.ln1, self.ln2 = nn.LayerNorm(dm), nn.LayerNorm(dm)
-        self.attn = Attn(dm, heads)
-        self.mlp = nn.Sequential(nn.Linear(dm, 4 * dm), nn.GELU(), nn.Linear(4 * dm, dm))
-
-    def forward(self, x, kv=None):
-        a, kv = self.attn(self.ln1(x), kv)
-        x = x + a
-        return x + self.mlp(self.ln2(x)), kv
-
-
-class WorldModel(nn.Module):
-    """Layout per step: [frame tokens (NT), action (1)]; frames t>=1 predicted from context."""
-    def __init__(self):
-        super().__init__()
-        self.tok = nn.Embedding(VOCAB, DM)
-        self.act = nn.Embedding(N_ACTIONS, DM)
-        self.pos = nn.Embedding(2048, DM)
-        self.blocks = nn.ModuleList([Block(DM, HEADS) for _ in range(LAYERS)])
-        self.ln = nn.LayerNorm(DM)
-        self.head = nn.Linear(DM, VOCAB)
-        self.r1 = nn.Sequential(nn.Linear(DM, DM), nn.ReLU(), nn.Linear(DM, 1))
-        self.r2 = nn.Sequential(nn.Linear(DM, DM), nn.ReLU(), nn.Linear(DM, 1))
-        self.cnt = nn.Sequential(nn.Linear(DM, DM), nn.ReLU(), nn.Linear(DM, 1))
-
-    def _seq(self, codes, actx):
-        B, L = codes.shape[0], codes.shape[1]
-        segs = []
-        for t in range(L - 1):
-            segs.append(self.tok(codes[:, t]))
-            segs.append(self.act(actx[:, t]).unsqueeze(1))
-        segs.append(self.tok(codes[:, L - 1]))
-        return torch.cat(segs, 1), L
-
-    def forward_full(self, codes, actx):
-        x, L = self._seq(codes, actx)
-        P = x.shape[1]
-        x = x + self.pos(torch.arange(P, device=x.device))[None]
-        for blk in self.blocks:
-            x, _ = blk(x)
-        h = self.ln(x)
-        pairs = h[:, : (L - 1) * (NT + 1)].reshape(-1, L - 1, NT + 1, DM)
-        f_h = torch.cat([pairs[:, :, :NT], h[:, (L - 1) * (NT + 1):].unsqueeze(1)], 1)  # (B,L,NT,DM)
-        a_h = pairs[:, :, NT]                                                            # (B,L-1,DM)
-        logits = self.head(f_h[:, :-1])                                                  # predict codes[:,1:]
-        r1, r2 = self.r1(a_h).squeeze(-1), self.r2(a_h).squeeze(-1)
-        c = self.cnt(a_h).squeeze(-1)
-        return logits, r1, r2, c, f_h
-
-    def wm_loss(self, codes, actx, rews, dones):
-        logits, r1, r2, c, _ = self.forward_full(codes, actx)
-        B, L = codes.shape[0], codes.shape[1]
-        ce = F.cross_entropy(logits.reshape(-1, VOCAB), codes[:, 1:].reshape(-1),
-                             reduction="none").view(B, L - 1, NT).mean(-1).mean()
-        rt = torch.sign(rews) * torch.log1p(torch.abs(rews))
-        rew = 0.5 * (F.mse_loss(r1, rt) + F.mse_loss(r2, rt))
-        cont = F.binary_cross_entropy_with_logits(c, 1.0 - dones)
-        return ce, rew, cont
-
-    def _cache(self, codes, actx):
-        """KV cache after full context ending on the last frame (no trailing action)."""
-        x, L = self._seq(codes, actx)
-        P = x.shape[1]
-        x = x + self.pos(torch.arange(P, device=x.device))[None]
-        kvs = []
-        for blk in self.blocks:
-            x, kv = blk(x)
-            kvs.append(kv)
-        return {"kv": kvs, "pos": P}
-
-    @torch.no_grad()
-    def imagine_step(self, cache, action):
-        """feed action -> (next codes (B,NT), r_pess (B,), c (B,)); generates NT codes."""
-        x = self.act(action).unsqueeze(1)
-        pos = cache["pos"]
-        x = x + self.pos(torch.arange(pos, pos + 1, device=x.device))[None]
-        for i, blk in enumerate(self.blocks):
-            a, kv = blk.attn(blk.ln1(x), cache["kv"][i])
-            cache["kv"][i] = kv
-            x = x + a
-            x = x + blk.mlp(blk.ln2(x))
-        cache["pos"] = pos + 1
-        a_h = self.ln(x[:, -1])
-        r = torch.minimum(self.r1(a_h), self.r2(a_h)).squeeze(-1)   # pessimistic ensemble
-        c = self.cnt(a_h).squeeze(-1)
-        codes = []
-        h = x
-        for k in range(NT):
-            nxt = self.head(self.ln(h))[:, -1].argmax(-1)
-            codes.append(nxt)
-            if k < NT - 1:
-                e = self.tok(nxt).unsqueeze(1) + self.pos(torch.arange(cache["pos"], cache["pos"] + 1, device=x.device))[None]
-                hh = e
-                for i, blk in enumerate(self.blocks):
-                    a2, kv2 = blk.attn(blk.ln1(hh), cache["kv"][i])
-                    cache["kv"][i] = kv2
-                    hh = hh + a2
-                    hh = hh + blk.mlp(blk.ln2(hh))
-                cache["pos"] += 1
-                h = hh
-        return torch.stack(codes, 1), r, c
-
-
-# ── actor / critic on semantic features ──────────────────────────────────────
-class Actor(nn.Module):
-    def __init__(self):
-        super().__init__()
-        d_in = 384 + TEXT_DIM
-        self.pi = nn.Sequential(nn.Linear(d_in, 256), nn.ReLU(), nn.Linear(256, N_ACTIONS))
-        self.v = nn.Sequential(nn.Linear(d_in, 256), nn.ReLU(), nn.Linear(256, 1))
-
-    def forward(self, cls, txt):
-        x = torch.cat([cls, txt], -1)
-        return self.pi(x), self.v(x).squeeze(-1)
-
-
-# ── classifier (dream vs real) ───────────────────────────────────────────────
-def auc_score(scores, labels):
-    order = np.argsort(scores)
-    ranks = np.empty(len(scores)); ranks[order] = np.arange(len(scores))
-    pos = labels.astype(bool)
-    n_pos, n_neg = pos.sum(), (~pos).sum()
-    if n_pos == 0 or n_neg == 0:
-        return float("nan")
-    return float((ranks[pos].sum() - n_pos * (n_pos - 1) / 2) / (n_pos * n_neg))
-
-
-# ── batching helpers ─────────────────────────────────────────────────────────
-def batch_windows(trajs, idxs, L, enc, tok, device):
-    """Random windows of length L from trajectories. Returns codes (B,L,NT),
-    actx (B,L-1), rews (B,L-1), dones (B,L-1), ctx screens/acts for BC/RL."""
-    B = len(idxs)
-    screens = np.zeros((B, L + 1, IMG, IMG, 3), np.uint8)
-    codes = np.zeros((B, L, NT), np.int64)
-    actx = np.zeros((B, L - 1), np.int64)
-    rews = np.zeros((B, L - 1), np.float32)
-    dones = np.zeros((B, L - 1), np.float32)
-    for b, (ti, s) in enumerate(idxs):
-        tr = trajs[ti]
-        T = len(tr["actions"])
-        s = max(0, min(s, T - (L - 1)))
-        fr = tr["frames"][s:s + L + 1]
-        if len(fr) < L + 1:                            # pad by repeating last frame
-            fr = np.concatenate([fr] + [fr[-1:]] * (L + 1 - len(fr)), 0)
-        screens[b] = fr
-        ac = tr["actions"][s:s + L - 1]
-        actx[b, : len(ac)] = ac
-        rw = tr["rewards"][s:s + L - 1]
-        rews[b, : len(rw)] = rw
-        dn = tr["dones"][s:s + L - 1]
-        dones[b, : len(dn)] = dn
-    scr_t = torch.from_numpy(screens).to(device)
-    with torch.no_grad():
-        flat = scr_t.reshape(-1, IMG, IMG, 3)          # (B*(L+1),IMG,IMG,3)
-        feats = []
-        for i in range(0, flat.shape[0], 16):
-            p, _ = enc.patches(flat[i:i + 16])
-            feats.append(p)
-        feats = torch.cat(feats).reshape(B, (L + 1) * NT, 384)
-    for b in range(B):
-        _, cd = tok.encode(feats[b].view((L + 1) * NT, 384).unsqueeze(0))
-        codes[b] = cd[0][: L * NT].view(L, NT).cpu().numpy()   # only the L context frames are coded; the (L+1)th frame is for BC/RL features
-    return (torch.from_numpy(codes).to(device), torch.from_numpy(actx).to(device),
-            torch.from_numpy(rews).to(device), torch.from_numpy(dones).to(device),
-            scr_t, feats)
-
-
-# ── main ─────────────────────────────────────────────────────────────────────
-def main():
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--smoke", action="store_true")
-    ap.add_argument("--device", default="cpu")
-    args = ap.parse_args()
-    device = args.device
-    S = args.smoke
-    n_demo, n_rand = (4, 2) if S else (400, 150)
-    tok_steps, wm_steps, bc_steps, rl_steps = ((30, 60, 60, 50) if S else (800, 1500, 1500, 2000))
-    n_eval = 2 if S else 4
-
-    trackio = None
-    if not S:
-        try:
-            import trackio
-            trackio.init(project="miniwob-imagination",
-                         space_id=os.environ.get("TRACKIO_SPACE_ID", "rootxhacker/miniwob-imagination-trackio"))
-        except Exception as e:
-            print(f"[trackio] disabled: {e}", flush=True)
-
-    def tlog(metrics, step=0):
-        if trackio is not None:
-            try:
-                trackio.log(metrics, step=step)
-            except Exception:
-                pass
-
-    t0 = time.time()
-    trajs = collect(TASKS[:1] if S else TASKS, n_demo, n_rand, seed=7)
-    n_train = int(0.9 * len(trajs)); tr_va = trajs[n_train:]; trajs = trajs[:n_train]
-    print(f"[data] {len(trajs)} train / {len(tr_va)} val trajs in {time.time()-t0:.0f}s", flush=True)
-    tlog({"data/train_trajs": len(trajs), "data/val_trajs": len(tr_va)})
-
-    enc = FrozenEncoders(device)
-    tok = FSQTok().to(device)
-
-    # ── B: tokenizer ──
-    opt = torch.optim.Adam(tok.parameters(), lr=1e-3)
-    for step in range(tok_steps):
-        b = random.sample(trajs, min(16, len(trajs)))
-        scr = torch.from_numpy(np.stack([t["frames"][random.randrange(len(t["frames"]))] for t in b])).to(device)
-        with torch.no_grad():
-            feats, _ = enc.patches(scr)
-        recon, _, _ = tok(feats)
-        loss = F.mse_loss(recon, feats)
-        opt.zero_grad(); loss.backward(); opt.step()
-    print(f"[tok] recon {loss.item():.4f}", flush=True)
-    tlog({"tok/recon": loss.item()}, step=tok_steps)
-
-    # ── C: world model with val early stop ──
-    wm = WorldModel().to(device)
-    opt = torch.optim.Adam(wm.parameters(), lr=3e-4)
-    def val_loss():
-        with torch.no_grad():
-            idxs = [(random.randrange(len(tr_va)), random.randrange(max(1, len(tr_va[0]["actions"]) - L_CTX))) for _ in range(8)]
-            codes, actx, rews, dones, _, _ = batch_windows(tr_va, idxs, L_CTX, enc, tok, device)
-            lg, r1, r2, c, _ = wm.forward_full(codes, actx)
-            ce = F.cross_entropy(lg.reshape(-1, VOCAB), codes[:, 1:].reshape(-1)).item()
-            rt = torch.sign(rews) * torch.log1p(torch.abs(rews))
-            rew = 0.5 * (F.mse_loss(r1, rt).item() + F.mse_loss(r2, rt).item())
-            cont = F.binary_cross_entropy_with_logits(c, 1.0 - dones).item()
-        return ce + rew + cont
-    best_v, patience, since = float("inf"), 3, 0
-    for step in range(wm_steps):
-        idxs = [(random.randrange(len(trajs)), random.randrange(max(1, len(trajs[0]["actions"]) - L_CTX))) for _ in range(16 if S else 32)]
-        codes, actx, rews, dones, _, _ = batch_windows(trajs, idxs, L_CTX, enc, tok, device)
-        ce, rew, cont = wm.wm_loss(codes, actx, rews, dones)
-        loss = ce + rew + cont
-        opt.zero_grad(); loss.backward(); opt.step()
-        if (step + 1) % 50 == 0:
-            v = val_loss()
-            print(f"[wm] step {step+1} train {loss.item():.4f} val {v:.4f}", flush=True)
-            tlog({"wm/train": loss.item(), "wm/val": v}, step=step + 1)
-            if v < best_v - 1e-4:
-                best_v, since = v, 0
-            else:
-                since += 1
-                if since >= patience:
-                    print(f"[wm] early stop at step {step+1}, best val {best_v:.4f}", flush=True)
-                    break
-    print(f"[wm] final train {loss.item():.4f} val {best_v:.4f}", flush=True)
-    tlog({"wm/final_val": best_v}, step=wm_steps)
-
-    # feature extractor for actor: CLS + text
-    def actor_feats(screens_np, texts):
-        scr = torch.from_numpy(np.stack(screens_np)).to(device)
-        _, cls = enc.patches(scr)
-        txt = enc.text(texts)
-        return torch.cat([cls, txt], -1)
-
-    actor = Actor().to(device)
-    opt = torch.optim.Adam(actor.parameters(), lr=1e-3)
-
-    # ── D: BC on oracle demos ──
-    for step in range(bc_steps):
-        b = random.sample(trajs, min(8 if S else 32, len(trajs)))
-        xs, ys, txts = [], [], []
-        for tr in b:
-            T = len(tr["actions"])
-            s = random.randrange(T)
-            xs.append(tr["frames"][s]); ys.append(int(tr["actions"][s])); txts.append(tr["instr"])
-        feats = actor_feats(xs, txts)
-        logits, v = actor(feats[:, :384], feats[:, 384:])
-        loss = F.cross_entropy(logits, torch.tensor(ys, device=device))
-        opt.zero_grad(); loss.backward(); opt.step()
-    print(f"[bc] CE {loss.item():.4f}", flush=True)
-    tlog({"bc/ce": loss.item()}, step=bc_steps)
-
-    # real-env eval of BC policy
-    def real_eval(policy, n_ep, greedy=True):
-        env = gymnasium.make(f"miniwob/{TASKS[0]}")
-        succ, rets = 0, []
-        for ep in range(n_ep):
-            raw, _ = env.reset(seed=1000 + ep)
-            ret = 0.0
-            for t in range(MAX_EP_STEPS):
-                cell = policy(raw)
-                raw2, r, term, trunc, _ = env.step({"action_type": 2, "coords": coords_of(cell)})
-                r = float(r); ret += r; raw = raw2
-                if term or trunc:
-                    break
-            rets.append(ret); succ += ret > 0.5
-        env.close()
-        return succ / max(1, n_ep), float(np.mean(rets))
-
-    def bc_policy(raw):
-        f = actor_feats([screen_of(raw)], [instr_text(raw)])
-        with torch.no_grad():
-            logits, _ = actor(f[:, :384], f[:, 384:])
-        return int(logits.argmax(-1))
-    s0, r0 = real_eval(bc_policy, n_ep=n_eval)
-    print(f"[bc] real-env success {s0:.2f} avg_return {r0:.3f}", flush=True)
-    tlog({"eval/bc_success": s0, "eval/bc_return": r0}, step=bc_steps)
-
-    # v5.1: push BC-stage checkpoint immediately so an RL crash loses nothing
-    os.makedirs("/tmp/ckpt", exist_ok=True)
-    torch.save({"tok": tok.state_dict(), "wm": wm.state_dict(), "actor": actor.state_dict()},
-               "/tmp/ckpt/bc_stage.pt")
-    if os.environ.get("HF_TOKEN"):
-        try:
-            from huggingface_hub import HfApi
-            HfApi().upload_file(path_or_fileobj="/tmp/ckpt/bc_stage.pt",
-                                path_in_repo="v5/bc_stage.pt", repo_id=HF_REPO, repo_type="model")
-            print("[push] BC checkpoint uploaded", flush=True)
-        except Exception as e:
-            print(f"[push] BC upload failed: {e}", flush=True)
-
-    # ── E: RL in imagination (REINFORCE + critic, pessimistic reward) ──
-    clf = nn.Linear(DM, 1).to(device)
-    ent_target, ent_coef = 4.0, 0.05
-    h_real, h_imag, lab = [], [], []
-    for step in range(rl_steps):
-        idxs = [(random.randrange(len(trajs)), random.randrange(max(1, len(trajs[0]["actions"]) - L_CTX))) for _ in range(8 if S else 24)]
-        codes, actx, rews, dones, scr, _ = batch_windows(trajs, idxs, L_CTX, enc, tok, device)
-        B = codes.shape[0]
-        codes0, actx0 = codes, actx          # v5.1: snapshot of the REAL window for classifier labels
-        # actor features from the last context frame + instruction
-        with torch.no_grad():
-            p, cl = enc.patches(scr[:, L_CTX])
-        txts = [trajs[ti]["instr"] for ti, _ in idxs]
-        txt = enc.text(txts)
-        logits, values = actor(cl, txt)
-        logp_all = F.log_softmax(logits, -1)
-        probs = logp_all.exp()
-        ent = -(probs * logp_all).sum(-1).mean()
-        actions = torch.multinomial(probs, 1).squeeze(-1)
-        logp = logp_all.gather(1, actions[:, None]).squeeze(1)
-        # imagine H=2 steps
-        cache = wm._cache(codes[:, :L_CTX - 1] if L_CTX > 1 else codes[:, :1],
-                          actx[:, :L_CTX - 1] if L_CTX > 1 else actx[:, :0])
-        rets = torch.zeros(B, device=device)
-        for hi in range(2 if not S else 1):
-            ac = actions if hi == 0 else torch.multinomial(probs, 1).squeeze(-1)
-            nc, r, c = wm.imagine_step(cache, ac)
-            rets = rets + r * (0.99 ** hi)
-            # build next context by rolling: drop oldest frame, append imagined codes
-            codes = torch.cat([codes[:, 1:], nc.unsqueeze(1)], 1)
-            actx = torch.cat([actx[:, 1:], ac[:, None]], 1)
-        # advantage from imagined return vs critic on context features
-        adv = (rets - values).detach()
-        adv = (adv - adv.mean()) / (adv.std() + 1e-6)
-        pg = -(logp * adv).mean() - ent_coef * F.relu(ent_target - ent) ** 2   # v5.1: quadratic floor
-        vloss = F.mse_loss(values, rets)
-        (pg + 0.5 * vloss).backward()
-        torch.nn.utils.clip_grad_norm_(actor.parameters(), 1.0)
-        opt.step(); opt.zero_grad()
-        if (step + 1) % 100 == 0:
-            print(f"[rl] step {step+1} ent {ent.item():.2f} R {rets.mean().item():.3f} v {values.mean().item():.3f}", flush=True)
-            tlog({"rl/entropy": ent.item(), "rl/imagined_R": rets.mean().item(),
-                  "rl/v_mean": values.mean().item()}, step=step + 1)
-        if (step + 1) % 500 == 0 or (S and step + 1 == rl_steps):
-            s, rr = real_eval(bc_policy, n_ep=n_eval)
-            print(f"[rl-eval] step {step+1} real-env success {s:.2f} avg_return {rr:.3f} entropy {ent.item():.2f}", flush=True)
-            tlog({"eval/success": s, "eval/return": rr}, step=step + 1)
-        # v5.1: classifier data — real window (pre-roll) vs post-imagination window
-        with torch.no_grad():
-            _, _, _, _, fh_r = wm.forward_full(codes0, actx0)
-            _, _, _, _, fh_i = wm.forward_full(codes, actx)
-        h_real.append(fh_r[:, -1].mean(1).cpu()); lab.append(torch.ones(B))
-        h_imag.append(fh_i[:, -1].mean(1).cpu()); lab.append(torch.zeros(B))
-
-    # ── F: dream-vs-real classifier ──
-    H = torch.cat(h_real + h_imag).to(device)                  # v5.1: was left on cpu -> crash on cuda
-    Y = torch.cat(lab).to(device)
-    H = (H - H.mean(0)) / (H.std(0) + 1e-6)
-    n = len(H); ntr = int(0.8 * n)
-    optc = torch.optim.Adam(clf.parameters(), lr=1e-2)
-    for step in range(200 if not S else 50):
-        i = torch.randperm(ntr)[: min(64, ntr)]
-        loss = F.binary_cross_entropy_with_logits(clf(H[i]).squeeze(-1), Y[i])
-        opt.zero_grad(); loss.backward(); opt.step()
-    with torch.no_grad():
-        sc = clf(H[ntr:]).squeeze(-1).cpu().numpy()
-        yy = Y[ntr:].cpu().numpy()
-    auc = auc_score(sc, yy)
-    print(f"[clf] dream-vs-real AUC {auc:.3f}", flush=True)
-    tlog({"clf/auc": auc}, step=rl_steps)
-
-    # ── final eval + push ──
-    s, rr = real_eval(bc_policy, n_ep=max(8, n_eval))
-    print(f"[final] real-env success {s:.2f} avg_return {rr:.3f}", flush=True)
-    tlog({"eval/final_success": s, "eval/final_return": rr}, step=rl_steps)
-    ckpt = {"tok": tok.state_dict(), "wm": wm.state_dict(), "actor": actor.state_dict(), "clf": clf.state_dict()}
-    os.makedirs("/tmp/ckpt", exist_ok=True)
-    torch.save(ckpt, "/tmp/ckpt/imagination_model_v5.pt")
-    if os.environ.get("HF_TOKEN"):
-        from huggingface_hub import HfApi
-        HfApi().upload_file(path_or_fileobj="/tmp/ckpt/imagination_model_v5.pt",
-                            path_in_repo="v5/imagination_model_v5.pt", repo_id=HF_REPO, repo_type="model")
-        print("[push] checkpoint uploaded", flush=True)
-    if trackio is not None:
-        try:
-            trackio.finish()
-        except Exception:
-            pass
-
-
-if __name__ == "__main__":
-    main()
+IyEvdXNyL2Jpbi9lbnYgcHl0aG9uMwoiIiIKdHJhaW5faW1hZ2luYXRpb25fdjUucHkg4oCUIGltYWdpbmF0aW9uIG1vZGVsIGZvciBjb21wdXRlciB1c2UsIHY1IChhY3Rvci1yZXByZXNlbnRhdGlvbiBmaXgpLgoKSGlzdG9yeSB0aGlzIGJ1aWxkIGFuc3dlcnM6CiAgdjEvZml4MWI6IFJMLWluLWltYWdpbmF0aW9uIG9ubHkgLT4gYWN0b3IgaGFkIG5vIHBhdGggZnJvbSBpbnN0cnVjdGlvbiB0byBjbGljayBsb2NhdGlvbjsKICAgICAgICAgICAgcmVhbC1lbnYgc3VjY2VzcyAwLjAwLCBlbnRyb3B5IGNvbGxhcHNlLgogIHY0OiAgICAgICBCQyBhZGRlZCAtPiBzdGlsbCAwLjEwIHN1Y2Nlc3M7IGFjdG9yIGlucHV0cyAoMjUgY29hcnNlIEZTUSBjb2RlcyArIGhhc2hlZCB3b3JkCiAgICAgICAgICAgIGJ1Y2tldHMpIGNhbm5vdCByZXByZXNlbnQgIndoZXJlIGlzIHRoZSAnT2snIGJ1dHRvbiIuCiAgdjUgZml4OiAgIGFjdG9yIHNlZXMgZnJvemVuIERJTk92Mi1zbWFsbCBwYXRjaC9DTFMgZmVhdHVyZXMgKHNlbWFudGljLCBwb3NpdGlvbi1wcmVzZXJ2aW5nKQogICAgICAgICAgICArIGZyb3plbiBTbW9sTE0yLTEzNU0gZW1iZWRkaW5nIG9mIHRoZSBSRUFMIGluc3RydWN0aW9uIHRleHQuIFdvcmxkIG1vZGVsIHN0YXlzIGEKICAgICAgICAgICAgc21hbGwgdHJhbnNmb3JtZXIgb3ZlciBGU1EgY29kZXM7IHJld2FyZCBlbnNlbWJsZSBpcyBwZXNzaW1pc3RpYyAobWluIG9mIDIgaGVhZHMpLgogIEFsc28gZml4ZXMgdGhlIFdNIGNvcHkgZmxhdzogQ0Ugbm93IHByZWRpY3RzIGNvZGVzWzosdCsxXSBmcm9tIG91dHB1dCBwb3NpdGlvbiB0CiAgKHYyL3YzIGNvbXBhcmVkIGEgcG9zaXRpb24ncyBvdXRwdXQgdG8gdGhlIGNvZGVzIGZlZCBhcyBpbnB1dCBhdCB0aGF0IHNhbWUgcG9zaXRpb24pLgogIHY1LjE6IGZpeCBjbGFzc2lmaWVyIGRldmljZSBjcmFzaCAoSCBvbiBjcHUgdnMgY2xmIG9uIGN1ZGEpLCBmaXggZHJlYW0tdnMtcmVhbCBsYWJlbCBtaXhpbmcKICAgICAgICAocmVhbCBlbWJlZGRpbmdzIG5vdyBjb21lIGZyb20gdGhlIHByZS1yb2xsIGNvbnRleHQgd2luZG93LCBub3QgcG9zdC1pbWFnaW5hdGlvbiksCiAgICAgICAgcXVhZHJhdGljIGVudHJvcHkgZmxvb3IsIEJDLXN0YWdlIGNoZWNrcG9pbnQgcHVzaGVkIGJlZm9yZSBSTCBzbyBhIGNyYXNoIGxvc2VzIG5vdGhpbmcuCiAgdjUuMjogYWN0aW9uIGdyaWQgNng2PTM2IGFsaWduZWQgMToxIHdpdGggdGhlIEZTUSB0b2tlbiBncmlkIChOVCA9IEdSSUQqR1JJRDsgRElOT3YyIHBhdGNoCiAgICAgICAgZmVhdHVyZXMgYWRhcHRpdmUtcG9vbGVkIHRvIEdSSUQgeCBHUklEKTsgb3JhY2xlIGNsaWNrcyB0aGUgQ0VMTCBDRU5URVIgc28gZXZlcnkKICAgICAgICBzdG9yZWQgbGFiZWwgaXMgZXhhY3RseSB0aGUgZXhlY3V0ZWQgYWN0aW9uOyBhY3RvciByZXBsYWNlZCBieSBhIHBlci1wb3NpdGlvbgogICAgICAgIGJpbGluZWFyIHNwYXRpYWwgaGVhZDogZV9rID0gTUxQKFtwZXItcG9zaXRpb24gRElOT3YyIGZlYXR1cmUsIGxlYXJuZWQgMkQgcG9zaXRpb24KICAgICAgICBlbWJlZGRpbmcsIGluc3RydWN0aW9uIGVtYmVkZGluZ10pICsgYmlsaW5lYXIgeCp5IHRlcm0sIGxvZ2l0X2sgPSAoVyBlX2spLmcvc3FydChkKSwKICAgICAgICBwb2xpY3kgPSBDYXRlZ29yaWNhbCBvdmVyIHRoZSBOVCBncmlkIHBvc2l0aW9ucyAoc2FtZSBpbmRleGluZyBhcyBjb29yZHNfb2YpLCBjcml0aWMKICAgICAgICBvbiB0aGUgcG9zaXRpb24gbWVhbi4gUkwgZml4OiBERUxFVEVEIGFkdmFudGFnZSBub3JtYWxpemF0aW9uIChpdCBhbXBsaWZpZWQKICAgICAgICB6ZXJvLXNpZ25hbCBub2lzZSB+MWU4eCBhbmQgY29sbGFwc2VkIGVudHJvcHkgdG8gMCB3aXRoaW4gMTAwIHN0ZXBzKSAtPgogICAgICAgIGFkdiA9IChSIC0gdmFscy5kZXRhY2goKSkuY2xhbXAoLTEwLDEwKTsgc3RhbmRhcmQgZW50cm9weSBib251cyAtMC4wMiplbnQ7CiAgICAgICAgcG9saWN5IGdyYWQgY2xpcCAxLjAuIC0tbWVhc3VyZS1ncmlkIE4gcnVucyB0aGUgb3JhY2xlLXJldHVybiBtZWFzdXJlbWVudC4KClBpcGVsaW5lOiBBIGNvbGxlY3QgKERPTSBvcmFjbGUgKyByYW5kb20pIC0+IEIgRlNRIHRva2VuaXplciBvbiBmcm96ZW4gRElOT3YyIHBhdGNoIGZlYXR1cmVzCi0+IEMgd29ybGQgbW9kZWwgKHZhbC1zcGxpdCBlYXJseSBzdG9wKSAtPiBEIEJDIG9uIG9yYWNsZSBkZW1vcyAtPiBFIFJMIGluIGltYWdpbmF0aW9uCihwZXNzaW1pc3RpYyBtaW4tb2YtMi1oZWFkcyByZXdhcmQsIFJFSU5GT1JDRSArIGNyaXRpY2FsLCBlbnRyb3B5IHRhcmdldCkgLT4gcmVhbC1lbnYgZXZhbHMgLTI+CkYgZHJlYW0tdnMtcmVhbCBjbGFzc2lmaWVyIChBVUMpLgoKRW52IGZhY3RzIChzYW5kYm94LXZlcmlmaWVkIDIwMjYtMDktMTksIG1pbml3b2IgMS4xLjApOgogIG9icyBrZXlzOiBkb21fZWxlbWVudHMsIGZpZWxkcywgc2NyZWVuc2hvdCwgdXR0ZXJhbmNlOyBzY3JlZW5zaG90ICgyMTAsMTYwLDMpIHVpbnQ4IChIMjEwLFcxNjApCiAgYWN0aW9uIHsiYWN0aW9uX3R5cGUiOiAyLCAiY29vcmRzIjogZmxvYXQzMlt4LHldfSA9IENMSUNLX0NPT1JEUyBvdmVyIEJveCgwLFsxNjAsMjEwXSkKICBkb20gZWxlbWVudHM6IGZsYXQgbGVmdC90b3Avd2lkdGgvaGVpZ2h0ICgxLWVsZW0gZmxvYXQzMiBhcnJheXMpLCB0ZXh0IG1hdGNoIGlzIGNhc2Utc2Vuc2l0aXZlClVzYWdlOgogIHB5dGhvbiB0cmFpbl9pbWFnaW5hdGlvbl92NS5weSAtLXNtb2tlICAgICAgICAgICAgIyB0aW55IENQVSBydW4sIGFsbCBwaGFzZXMKICBweXRob24gdHJhaW5faW1hZ2luYXRpb25fdjUucHkgLS1kZXZpY2UgY3VkYSAgICAgICMgcmVhbCBydW4gKHQ0LXNtYWxsLCB+MS41aCkKIiIiCmltcG9ydCBhcmdwYXJzZSwgbWF0aCwgb3MsIHJhbmRvbSwgdGltZQppbXBvcnQgbnVtcHkgYXMgbnAKaW1wb3J0IHRvcmNoLCB0b3JjaC5ubiBhcyBubiwgdG9yY2gubm4uZnVuY3Rpb25hbCBhcyBGCgpvcy5lbnZpcm9uLnNldGRlZmF1bHQoIk1JTklXT0JfQ0hST01FX0JJTkFSWSIsICIvdXNyL2Jpbi9jaHJvbWl1bSIpCm9zLmVudmlyb24uc2V0ZGVmYXVsdCgiTUlOSVdPQl9DSFJPTUVEUklWRVIiLCAiL3Vzci9iaW4vY2hyb21lZHJpdmVyIikKCmltcG9ydCBneW1uYXNpdX1LCBtaW5pd29iCmZyb20gUElMIGltcG9ydCBJbWFnZQpmcm9tIHRyYW5zZm9ybWVycyBpbXBvcnQgQXV0b01vZGVsLCBBdXRvVG9rZW5pemVyCgpneW1uYXNpdW0ucmVnaXN0ZXJfZW52cyhtaW5pd29iKQo=