code/tools/eval_lam_controls.py

"""Controls for the LaWAM-VGGT Stage-1 probes -- the numbers `eval_lam_probes.py`
reports are not interpretable on their own.

Two gaps in the original evaluation, both of which could make a dead `z` look alive:

A. `z2action_r2 = 0.55` has no baseline. LIBERO actions are strongly predictable
   from the current frame alone (the arm's position implies where it is going).
   If a 32-d PCA of `u_t` scores the same R^2, `z` contributes nothing and the
   "z encodes action" claim collapses. We match dimensionality exactly (32 vs 32)
   so the comparison is about *content*, not capacity, and fit the PCA on train
   episodes only.

B. `delta_z` shuffles `z` across the whole batch, which mixes two effects: wrong
   *motion* and wrong *scene/task*. A `z` that only encoded "which LIBERO suite
   is this" would still produce a large delta_z. Shuffling *within* an episode
   holds scene and task fixed, so only the motion component can explain the gap.

Also characterizes what the decoder does with a wrong `z` (collapse to identity,
or confidently wrong motion?) and adds z=0 / random-z references.
"""

import argparse
import importlib.machinery
import sys
import types
from collections import defaultdict

import numpy as np
import torch
import torch.nn.functional as F

REPO = "/home/ma-user/work/dataset/xxd-dataset/dataset_yhw/WAM/LaWAM_official"
sys.path.insert(0, REPO)


def _stub(name, **attrs):
    mod = types.ModuleType(name)
    mod.__spec__ = importlib.machinery.ModuleSpec(name, None)
    mod.__path__ = []
    for k, v in attrs.items():
        setattr(mod, k, v)
    sys.modules[name] = mod
    return mod


class _Callback:
    pass


_stub("lightning", LightningModule=torch.nn.Module)
_stub("lightning.pytorch", Callback=_Callback)
_stub("lightning.pytorch.callbacks", Callback=_Callback)
sys.modules["lightning"].pytorch = sys.modules["lightning.pytorch"]
_stub("wandb", Image=lambda *a, **k: None, log=lambda *a, **k: None)

from latent_action_model.core.lam_model import load_latent_action_model  # noqa: E402
from latent_action_model.data_loader.lerobot_dataset import LeRobotLAMDataset  # noqa: E402
from latent_action_model.data_loader.collate import lam_collate  # noqa: E402
from latent_action_model.data_loader.video_aug import gpu_two_view_video_aug  # noqa: E402


def standardize(x, mu, sd):
    return (x - mu) / sd


def ridge_r2(X, Y, tr, te, tag):
    """Episode-split ridge with a lambda sweep. Returns (r2, lambda)."""
    Xm, Xs = X[tr].mean(0), X[tr].std(0) + 1e-8
    Ym = Y[tr].mean(0)
    Xtr, Xte = (X[tr] - Xm) / Xs, (X[te] - Xm) / Xs
    Ytr, Yte = Y[tr] - Ym, Y[te] - Ym
    best = (-1e9, None)
    for lam in [1e-3, 1e-2, 1e-1, 1.0, 10.0, 100.0, 1e3]:
        A = Xtr.T @ Xtr + lam * np.eye(Xtr.shape[1])
        W = np.linalg.solve(A, Xtr.T @ Ytr)
        P = Xte @ W
        r2 = 1 - ((Yte - P) ** 2).sum() / max((Yte ** 2).sum(), 1e-12)
        if r2 > best[0]:
            best = (r2, lam)
    print(f"  {tag:36s} R^2 = {best[0]:.4f}   (lambda={best[1]})")
    return best


@torch.no_grad()
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="/home/ma-user/work/lam_runs/vggt_vae_libero/checkpoints/epoch=39.ckpt")
    ap.add_argument("--config", default="/home/ma-user/work/lam_runs/vggt_vae_libero/version_0/config.yaml")
    ap.add_argument("--n-batches", type=int, default=60)
    ap.add_argument("--batch-size", type=int, default=16)
    ap.add_argument("--val-tail-ratio", type=float, default=0.05)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    dev = "cuda"

    model = load_latent_action_model(args.ckpt, args.config).to(dev).eval()
    print(f"[ctrl] encoder={type(model.vision_encoder).__name__} input_dim={model.input_dim}")

    ds = LeRobotLAMDataset(
        data_root_dir="/home/ma-user/work/lam_datasets",
        data_mix="libero",
        num_frames=2,
        mode="val",
        val_tail_ratio=args.val_tail_ratio,
        video_backend="pyav",
        image_hw=(256, 256),
        frame_dt_sec=1.6,
        debug_repeat_batch=False,
    )
    loader = torch.utils.data.DataLoader(
        ds, batch_size=args.batch_size, num_workers=8,
        collate_fn=lambda b: lam_collate(b, max_state_dim=32), drop_last=True,
    )

    all_dec_in, all_tgt, all_recon, all_z, all_traj, all_base = [], [], [], [], [], []
    it = iter(loader)
    for i in range(args.n_batches):
        try:
            batch = next(it)
        except StopIteration:
            print(f"[ctrl] loader exhausted after {i} batches")
            break
        videos = batch["videos"].to(dev)
        v1, v2 = gpu_two_view_video_aug(videos, output_size=(256, 256), training=False, dual_view_aug=False)
        states = batch["states"].to(dev)
        state_mask = batch["state_mask"].to(dev)
        emb = batch["embodiment_ids"].to(dev)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            recon, dec_in, tgt, _, _, _, _, z, _, _ = model.inference(
                v1, states, v2, state_mask=state_mask, embodiment_ids=emb
            )
        all_dec_in.append(dec_in.float().cpu())
        all_tgt.append(tgt.float().cpu())
        all_recon.append(recon.float().cpu())
        all_z.append(z.float().reshape(z.shape[0], -1).cpu())
        all_traj.extend(batch["trajectory_ids"])
        all_base.append(batch["base_indices"])
        if (i + 1) % 20 == 0:
            print(f"[ctrl] pass1 {i+1}/{args.n_batches}")

    dec_in = torch.cat(all_dec_in)
    tgt = torch.cat(all_tgt)
    recon = torch.cat(all_recon)
    Z = torch.cat(all_z)
    base_idx = torch.cat(all_base).numpy()
    eps_all = np.array([int(t) for t in all_traj])
    N = dec_in.shape[0]
    print(f"\n[ctrl] N={N} samples | z dim={Z.shape[1]} | {len(np.unique(eps_all))} episodes")

    flat = tgt.reshape(-1, tgt.shape[-1])
    mu, sd = flat.mean(0), flat.std(0).clamp_min(1e-6)
    T_ = standardize(tgt, mu, sd)
    R_ = standardize(recon, mu, sd)
    I_ = standardize(dec_in, mu, sd)
    mse = lambda a, b: (a - b).pow(2).mean().item()
    mse_pred, mse_identity = mse(R_, T_), mse(I_, T_)
    print(f"[ctrl] mse_pred={mse_pred:.4f}  mse_identity={mse_identity:.4f}  mse_mean={mse(torch.zeros_like(T_), T_):.4f}")

    def decode_with(zs):
        """Re-run only the decoder with the given latents. Returns standardized pred."""
        outs = []
        for s in range(0, N, args.batch_size):
            e = min(s + args.batch_size, N)
            with torch.autocast("cuda", dtype=torch.bfloat16):
                r = model.decoder(dec_in[s:e].to(dev), zs[s:e].to(dev).unsqueeze(1)).float().cpu()
            outs.append(standardize(r, mu, sd))
        return torch.cat(outs)

    # ---- B. delta_z variants ----------------------------------------------
    print("\n=== 对照 B: delta_z 的四种 z 替换 ===")
    g = torch.Generator().manual_seed(args.seed)

    # (b1) shuffle across the whole set -- the original probe
    R_glob = decode_with(Z[torch.randperm(N, generator=g)])

    # (b2) shuffle WITHIN episode -- holds scene/task fixed, only motion differs
    by_ep = defaultdict(list)
    for i, e in enumerate(eps_all):
        by_ep[int(e)].append(i)
    perm_w = np.arange(N)
    n_swappable = 0
    for e, idxs in by_ep.items():
        if len(idxs) < 2:
            continue
        a = np.array(idxs)
        b = a.copy()
        # derangement-ish: roll by 1 within the episode
        perm_w[a] = np.roll(b, 1)
        n_swappable += len(idxs)
    R_within = decode_with(Z[torch.from_numpy(perm_w)])

    # (b3) z = 0  and  (b4) z ~ N(0, I)
    R_zero = decode_with(torch.zeros_like(Z))
    R_rand = decode_with(torch.randn(Z.shape, generator=g) * Z.std())

    rows = [
        ("真实 z",              mse_pred,            R_),
        ("跨全集打乱 z",         mse(R_glob, T_),     R_glob),
        ("同 episode 内打乱 z",  mse(R_within, T_),   R_within),
        ("z = 0",              mse(R_zero, T_),     R_zero),
        ("z ~ 高斯噪声",         mse(R_rand, T_),     R_rand),
        ("抄袭 u_t (identity)",  mse_identity,        I_),
    ]
    print(f"  {'替换方式':22s} {'mse':>8s} {'delta_z':>9s} {'与u_t距离':>10s}")
    for name, m, arr in rows:
        d = m - mse_pred
        di = mse(arr, I_)
        print(f"  {name:22s} {m:8.4f} {d:9.4f} {di:10.4f}")
    print(f"  可同-episode 交换的样本: {n_swappable}/{N}")
    print("  注: '与u_t距离'=0 表示解码器退化成原样抄袭输入")

    # ---- A. z vs u_t as action predictors ---------------------------------
    print("\n=== 对照 A: z 与 u_t 谁能解出动作 (同为32维, 同 episode 切分) ===")
    import pyarrow.parquet as pq

    tbl = pq.read_table(
        "/home/ma-user/work/lam_datasets/libero_merged_no_noops_20hz/data/chunk-000/file-000.parquet",
        columns=["episode_index", "frame_index", "action"],
    )
    ep_arr = np.asarray(tbl.column("episode_index"))
    fr_arr = np.asarray(tbl.column("frame_index"))
    act_arr = np.stack(tbl.column("action").to_numpy(zero_copy_only=False))
    key2row = {(int(e), int(f)): i for i, (e, f) in enumerate(zip(ep_arr, fr_arr))}

    STRIDE = 32
    keep, Y_rows = [], []
    for i in range(N):
        e, b = int(eps_all[i]), int(base_idx[i])
        chunk, ok = [], True
        for k in range(0, STRIDE, 4):
            r = key2row.get((e, b + k))
            if r is None:
                ok = False
                break
            chunk.append(act_arr[r])
        if ok:
            keep.append(i)
            Y_rows.append(np.concatenate(chunk))
    keep = np.array(keep)
    Y = np.asarray(Y_rows, dtype=np.float64)
    eps_k = eps_all[keep]
    print(f"  对齐样本 {len(keep)} | action {Y.shape[1]}维")

    uniq = np.unique(eps_k)
    rng = np.random.RandomState(0)
    rng.shuffle(uniq)
    tr_eps = set(uniq[: max(1, int(len(uniq) * 0.7))].tolist())
    tr = np.array([e in tr_eps for e in eps_k])
    te = ~tr
    print(f"  episode 级切分: train {tr.sum()} / test {te.sum()}, {len(uniq)} episodes")

    Xz = Z[keep].numpy().astype(np.float64)

    # u_t -> 32 dims via PCA fit on TRAIN episodes only (no leakage)
    U = dec_in[keep, 0].mean(dim=1).numpy().astype(np.float64)   # [n, 2048] token-mean
    Um = U[tr].mean(0)
    Uc = U - Um
    _, _, Vt = np.linalg.svd(Uc[tr], full_matrices=False)
    Xu = Uc @ Vt[:32].T

    # also the delta feature u_t itself at full token resolution, PCA-32
    Uf = dec_in[keep, 0].reshape(len(keep), -1).numpy().astype(np.float64)
    Ufm = Uf[tr].mean(0)
    Ufc = Uf - Ufm
    _, _, Vt2 = np.linalg.svd(Ufc[tr], full_matrices=False)
    Xuf = Ufc @ Vt2[:32].T

    r2_z, _ = ridge_r2(Xz, Y, tr, te, "z (32维, 本模型的latent action)")
    r2_u, _ = ridge_r2(Xu, Y, tr, te, "PCA32(u_t token均值)  [对照]")
    r2_uf, _ = ridge_r2(Xuf, Y, tr, te, "PCA32(u_t 全token)   [对照]")
    r2_cat, _ = ridge_r2(np.concatenate([Xz, Xuf], 1), Y, tr, te, "z + PCA32(u_t) 拼接")

    print("\n=== 汇总 ===")
    print(f"  z 单独            R^2 = {r2_z:.4f}")
    print(f"  u_t 单独 (最好)    R^2 = {max(r2_u, r2_uf):.4f}")
    print(f"  增量 (z - u_t)          = {r2_z - max(r2_u, r2_uf):+.4f}")
    print(f"  拼接              R^2 = {r2_cat:.4f}")
    print(f"  delta_z 跨全集          = {mse(R_glob, T_) - mse_pred:.4f}")
    print(f"  delta_z 同episode内      = {mse(R_within, T_) - mse_pred:.4f}  <- 排除场景/任务混淆")


if __name__ == "__main__":
    main()