src/idm_hallucination.py

# idm_hallucination.py
"""
Hallucination detection by inverse-dynamics rollout divergence.

A standalone predictor. It deliberately does not extend the u_r / u_f / u_s
scorers in `uncertainty.py`: those score a *single predicted latent* per
candidate action and are label-free at the level of one step, whereas this
method needs a whole observed trajectory to compare against. Their
`(predictions_KN, z_prev_K)` scorer signature cannot express that, so this
lives on its own rather than being bent to fit it.

Method
------
1. Infer the action behind every observed transition with a trained inverse
   dynamics model (`idm_tokenizer.py`):

       a_hat_t = IDM(z_t, z_{t+1})

   This needs no action labels -- only the video -- which is what makes the
   method applicable to arbitrary unlabeled trajectories.

2. Re-simulate the trajectory with the world model, open-loop, conditioned on
   a_hat. Each predicted latent is fed back in, so errors accumulate the way
   they do in a real rollout.

3. Compare the simulated latents against the real ones, step by step:

       dist_t      = RMS(z_sim_t - z_real_t)
       dist_norm_t = dist_t / RMS(z_real_t - z_real_{t-1})
       score       = mean_t dist_norm_t

A world model that has genuinely learned this region reproduces the real
trajectory under the actions that demonstrably generated it. A hallucinating
one invents its own continuation and drifts away.

Keep the horizon short
----------------------
Measured on walker-run with the released `base` checkpoint, the divergence of
a ground-truth-action rollout versus a random-action rollout is 1.55x at step
2 and decays to ~1.2x by step 6: an open-loop rollout reaches the distance
between two *unrelated* real latents (0.587 there) within about four steps,
after which every action sequence looks equally wrong and further steps only
add noise. Aggregating over a 13-step rollout collapsed that separation to
1.07x. Hence `horizon=4` by default.

Validated over 16 rollouts across 4 tasks (walker-run, cheetah-run,
mw-assembly, pygame-cowboy) with `horizon=3`:

    ground-truth actions   0.578   (1.00x)
    IDM-inferred actions   0.678   (1.17x)
    shuffled GT actions    1.073   (1.86x)
    random actions         1.148   (1.99x)

with ground-truth below random in 16/16 rollouts and IDM-inferred below random
in 16/16. The inferred actions land far closer to ground truth than to wrong
actions, which is the point: the label-free signal retains most of what the
labeled one gives you.

Usage
-----
    from idm_hallucination import IDMRolloutDivergence

    det = IDMRolloutDivergence.from_checkpoints(
        tokenizer_ckpt="./checkpoints/base/tokenizer.pt",
        dynamics_ckpt="./checkpoints/base/dynamics.pt",
        idm_ckpt="./logs/idm_tokenizer_ckpts/base_run3/step_040000.pt",
        device=torch.device("cuda"),
    )
    z = det.encode(obs_u8)                     # (1, T+1, 3, H, W) uint8 -> packed latents
    out = det(z, act_mask=mask, lang_emb=lang)
    print(out["score"], out["dist"])

Or from the command line, scoring trajectories of one task:

    python idm_hallucination.py --task walker-run --idm_ckpt <path> --n_traj 4
"""
from typing import Any, Dict, Optional

import torch

from model import (
    Dynamics, pack_bottleneck_to_spatial, temporal_patchify,
    unpack_spatial_to_bottleneck,
)
from train_dynamics import sample_one_timestep_packed


def load_idm_tokenizer_head(ckpt_path: str, *, device, n_latents: int, d_bottleneck: int):
    """Load a trained tokenizer-arm inverse dynamics head (`idm_tokenizer.py`).

    The architecture is recovered from the checkpoint's own tensor shapes, so
    any `--d_hidden` loads without being told which one was used.

    Only the tokenizer arm applies here. It maps a pair of latents
    (z_t, z_{t+1}) to an action, which is exactly what this method has at each
    step. The backbone arm (`idm_backbone.py`) reads contextual features out of
    a Dynamics forward over an observed window, which is a different input
    contract.
    """
    from idm_tokenizer import InverseDynamicsHeadTokenizer

    ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    sd = ck["model"] if isinstance(ck, dict) and "model" in ck else ck
    d_hidden, d_in3 = sd["in_proj.weight"].shape
    d_in = d_in3 // 3
    expected = n_latents * d_bottleneck
    if d_in != expected:
        raise ValueError(
            f"IDM checkpoint expects flattened latents of {d_in}, but the tokenizer "
            f"gives n_latents*d_bottleneck = {expected}. Checkpoint/tokenizer mismatch."
        )
    head = InverseDynamicsHeadTokenizer(
        n_latents=n_latents, d_bottleneck=d_bottleneck, d_hidden=int(d_hidden),
    )
    head.load_state_dict(sd)
    head.to(device).eval()
    for p in head.parameters():
        p.requires_grad_(False)
    meta = {"step": (ck.get("step") if isinstance(ck, dict) else None),
            "d_hidden": int(d_hidden), "d_in": int(d_in)}
    return head, meta


class IDMRolloutDivergence:
    """Hallucination score from inverse-dynamics rollout divergence.

    Holds the frozen world model + inverse dynamics model and the rollout
    configuration; call the instance on a trajectory's packed latents.
    """

    def __init__(
        self,
        dyn: Dynamics,
        idm,
        *,
        k_max: int,
        sched: Dict[str, Any],
        packing_factor: int,
        encoder=None,
        patch: Optional[int] = None,
        n_latents: Optional[int] = None,
        n_context: int = 8,
        horizon: Optional[int] = 4,
        max_ctx: int = 8,
        tau_ctx: float = 0.0,
        motion_eps: float = 1e-3,
        use_kv_cache: bool = False,
    ):
        self.dyn = dyn
        self.idm = idm
        self.k_max = int(k_max)
        self.sched = sched
        self.packing_factor = int(packing_factor)
        self.encoder = encoder
        self.patch = patch
        self.n_latents = n_latents
        self.n_context = int(n_context)
        self.horizon = horizon
        self.max_ctx = int(max_ctx)
        self.tau_ctx = float(tau_ctx)
        self.motion_eps = float(motion_eps)
        # The rollout is the dominant cost in a live demo: H steps x 2 action
        # sources, each re-attending over the whole context. Caching the
        # context K,V is an exact optimisation -- same numbers, less compute.
        self.use_kv_cache = bool(use_kv_cache)

    @classmethod
    def from_checkpoints(
        cls,
        *,
        tokenizer_ckpt: str,
        dynamics_ckpt: str,
        idm_ckpt: str,
        device,
        packing_factor: int = 2,
        schedule: str = "shortcut",
        eval_d: float = 0.25,
        **kwargs,
    ) -> "IDMRolloutDivergence":
        """Build everything from checkpoint paths (tokenizer, dynamics, IDM)."""
        from train_dynamics import load_frozen_tokenizer_from_pt_ckpt, make_tau_schedule
        from interactive import load_dynamics_from_ckpt

        enc, _dec, tok_args = load_frozen_tokenizer_from_pt_ckpt(tokenizer_ckpt, device=device)
        patch = int(tok_args.get("patch", 14))
        n_latents = int(tok_args.get("n_latents", 64))
        d_bottleneck = int(tok_args.get("d_bottleneck", 64))

        dyn, _rew, _pol, dyn_meta = load_dynamics_from_ckpt(
            dynamics_ckpt, device=device, d_bottleneck=d_bottleneck,
            n_latents=n_latents, packing_factor=packing_factor,
        )
        idm, _meta = load_idm_tokenizer_head(
            idm_ckpt, device=device, n_latents=n_latents, d_bottleneck=d_bottleneck,
        )
        sched = make_tau_schedule(
            k_max=dyn_meta["k_max"], schedule=schedule,
            d=(eval_d if schedule == "shortcut" else None),
        )
        return cls(
            dyn, idm, k_max=dyn_meta["k_max"], sched=sched,
            packing_factor=packing_factor, encoder=enc, patch=patch,
            n_latents=n_latents, **kwargs,
        )

    @torch.no_grad()
    def encode(self, obs_u8: torch.Tensor) -> torch.Tensor:
        """(B, T+1, 3, H, W) uint8 -> packed latents (B, T+1, Sz, Dz)."""
        if self.encoder is None or self.patch is None or self.n_latents is None:
            raise RuntimeError("encode() needs an encoder; build via from_checkpoints().")
        z = self.encoder(temporal_patchify(obs_u8.float() / 255.0, self.patch))[0]
        return pack_bottleneck_to_spatial(
            z, n_spatial=self.n_latents // self.packing_factor, k=self.packing_factor,
        )

    @torch.no_grad()
    def infer_actions(self, z_real_packed: torch.Tensor) -> torch.Tensor:
        """Packed latents (1, T+1, Sz, Dz) -> inferred actions (1, T, A).

        `a_hat[:, t]` is the action for the transition z_t -> z_{t+1}.
        """
        dtype = next(self.idm.parameters()).dtype
        zu = unpack_spatial_to_bottleneck(z_real_packed, k=self.packing_factor)
        return self.idm(zu[:, :-1].to(dtype), zu[:, 1:].to(dtype)).float()

    @torch.no_grad()
    def __call__(
        self,
        z_real_packed: torch.Tensor,                      # (1, T+1, Sz, Dz)
        *,
        act_mask: Optional[torch.Tensor] = None,
        lang_emb: Optional[torch.Tensor] = None,
        actions_override: Optional[torch.Tensor] = None,  # (1, T, A), bypass the IDM
        n_context: Optional[int] = None,
        horizon: Optional[int] = -1,                      # -1 = use the instance default
    ) -> Dict[str, torch.Tensor]:
        """Score one trajectory.

        `actions_override` runs the same rollout under known or deliberately
        perturbed actions instead of inferred ones -- the control used to
        validate the signal.

        Returns:
          "dist":        (H,)    RMS(z_sim_t - z_real_t) per simulated step
          "dist_norm":   (H,)    same, divided by the real per-step motion
          "motion":      (H,)    RMS(z_real_t - z_real_{t-1})
          "actions_hat": (T, A)  the inferred (or overridden) actions
          "score":       scalar  mean of dist_norm -- the hallucination score
        """
        assert z_real_packed.dim() == 4 and z_real_packed.shape[0] == 1, \
            "z_real_packed must be (1, T+1, Sz, Dz)"
        device = z_real_packed.device
        dtype = next(self.dyn.parameters()).dtype
        Tp1 = z_real_packed.shape[1]
        T = Tp1 - 1
        assert T >= 1, "need at least one transition"

        n_ctx = self.n_context if n_context is None else int(n_context)
        n_ctx = max(1, min(n_ctx, T))
        h = self.horizon if horizon == -1 else horizon

        # 1. Infer the action sequence from the real transitions.
        if actions_override is not None:
            a_hat = actions_override.to(device).float()
        else:
            a_hat = self.infer_actions(z_real_packed)
        A = a_hat.shape[-1]

        # led-to convention: actions_led_to[i] is the action that led to frame i,
        # so a_hat_t (for z_t -> z_{t+1}) sits at index t+1 and index 0 is zero.
        actions_led_to = torch.zeros((1, Tp1, A), device=device, dtype=torch.float32)
        actions_led_to[:, 1:] = a_hat

        z_real = z_real_packed.float()
        history = [z_real[:, i].to(dtype) for i in range(n_ctx)]

        last = Tp1 if h is None else min(Tp1, n_ctx + int(h))
        dists, norms, motions = [], [], []
        for m in range(n_ctx, last):
            z_win = history[-self.max_ctx:] if (self.max_ctx > 0 and len(history) > self.max_ctx) else history
            ctx_len = len(z_win)
            past = torch.stack(z_win, dim=1)                          # (1, ctx_len, Sz, Dz)
            acts = torch.cat(
                [actions_led_to[:, m - ctx_len:m], actions_led_to[:, m:m + 1]], dim=1,
            ).to(dtype)                                               # (1, ctx_len+1, A)

            z_next = sample_one_timestep_packed(
                self.dyn,
                past_packed=past,
                k_max=self.k_max, sched=self.sched,
                actions=acts, act_mask=act_mask, tau_ctx=self.tau_ctx,
                lang_emb=lang_emb, use_kv_cache=self.use_kv_cache,
            )                                                         # (1, Sz, Dz)

            z_sim = z_next.float()[0]
            d = (z_sim - z_real[0, m]).pow(2).mean().sqrt()
            mo = (z_real[0, m] - z_real[0, m - 1]).pow(2).mean().sqrt()
            dists.append(d)
            motions.append(mo)
            norms.append(d / mo.clamp(min=self.motion_eps))

            history.append(z_next.to(dtype)[0].unsqueeze(0))          # open-loop feedback

        dist_norm = torch.stack(norms)
        return {
            "dist": torch.stack(dists),
            "dist_norm": dist_norm,
            "motion": torch.stack(motions),
            "actions_hat": a_hat[0],
            "score": dist_norm.mean(),
        }


if __name__ == "__main__":
    import argparse

    from wm_dataset import WMDataset, collate_batch

    p = argparse.ArgumentParser(description=__doc__.split("\n")[1])
    p.add_argument("--task", type=str, default="walker-run")
    p.add_argument("--tokenizer_ckpt", type=str, default="./checkpoints/base/tokenizer.pt")
    p.add_argument("--dynamics_ckpt", type=str, default="./checkpoints/base/dynamics.pt")
    p.add_argument("--idm_ckpt", type=str,
                   default="./logs/idm_tokenizer_ckpts/base_run3/step_040000.pt")
    p.add_argument("--data_dir", type=str, default="/data5/expert")
    p.add_argument("--frame_dir", type=str, default="/data5/expert-shards")
    p.add_argument("--tasks_json", type=str, default="../tasks.json")
    p.add_argument("--seq_len", type=int, default=16)
    p.add_argument("--n_traj", type=int, default=4)
    p.add_argument("--n_context", type=int, default=8)
    p.add_argument("--horizon", type=int, default=3)
    p.add_argument("--packing_factor", type=int, default=2)
    p.add_argument("--controls", action="store_true",
                   help="also score ground-truth / shuffled / random actions as controls")
    args = p.parse_args()

    dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    det = IDMRolloutDivergence.from_checkpoints(
        tokenizer_ckpt=args.tokenizer_ckpt, dynamics_ckpt=args.dynamics_ckpt,
        idm_ckpt=args.idm_ckpt, device=dev, packing_factor=args.packing_factor,
        n_context=args.n_context, horizon=args.horizon,
    )
    lang_dim = det.dyn.lang_dim

    ds = WMDataset(
        data_dir=[args.data_dir], frames_dir=[args.frame_dir], seq_len=args.seq_len,
        img_size=224, action_dim=16, lang_dim=lang_dim, tasks_json=args.tasks_json,
        tasks=[args.task], verbose=False, cache_mb=512, ddp_partition=False,
        iid_sampling=False, strict_tasks=False,
    )
    hdr = f"{'traj':>4s} {'score':>9s}"
    if args.controls:
        hdr += f" {'GT':>9s} {'shuffled':>9s} {'random':>9s}"
    print(f"task={args.task}  n_context={args.n_context}  horizon={args.horizon}")
    print(hdr)
    print("-" * len(hdr))

    for j in range(args.n_traj):
        b = collate_batch([ds[(len(ds) // (args.n_traj + 1)) * (j + 1)]])
        obs, am = b["obs"].to(dev), b["act_mask"].to(dev)
        a_gt = (b["act"].to(dev).clamp(-1, 1)) * am
        lang = b["lang_emb"].to(dev) if lang_dim > 0 else None
        z = det.encode(obs)
        kw = dict(act_mask=am[0, 0], lang_emb=lang)
        row = f"{j:>4d} {float(det(z, **kw)['score']):>9.4f}"
        if args.controls:
            torch.manual_seed(j)
            perm = torch.randperm(a_gt.shape[1], device=dev)
            for a in (a_gt, a_gt[:, perm], (torch.rand_like(a_gt) * 2 - 1) * am):
                row += f" {float(det(z, actions_override=a, **kw)['score']):>9.4f}"
        print(row)