src/idm_tokenizer.py

# idm_tokenizer.py
"""
Inverse dynamics model (IDM) -- Option A: tokenizer-only.

Predicts the action a_t that caused the transition obs_t -> obs_{t+1} using
only a pair of frozen-tokenizer latents (z_t, z_{t+1}). No dynamics
transformer is involved: this is the cheap end of the spectrum, useful as a
baseline, for action-labeling unlabeled trajectories, or to see how much
inverse-dynamics signal is already recoverable from the tokenizer's per-frame
representation alone. See idm_backbone.py for the richer alternative that
reads off the (frozen) Dynamics transformer's contextual features instead.

Only a new lightweight head (`InverseDynamicsHeadTokenizer`) is trained; the
tokenizer stays frozen throughout (as in train_dynamics.py).

Run from inside `src/`, single-GPU:
    python idm_tokenizer.py \
        --tokenizer_ckpt ./logs/tokenizer_ckpts/latest.pt \
        --data_dirs /data5/expert --frame_dirs /data5/expert-shards \
        --out_ckpt ./logs/idm_tokenizer_ckpts/latest.pt

or multi-GPU (DDP, same flags, launched via torchrun):
    torchrun --nproc_per_node=8 idm_tokenizer.py \
        --tokenizer_ckpt ./logs/tokenizer_ckpts/latest.pt \
        --data_dirs /data5/expert --frame_dirs /data5/expert-shards \
        --out_ckpt ./logs/idm_tokenizer_ckpts/latest.pt
"""
import argparse
import math
import time
from pathlib import Path
from typing import List, Optional, Tuple

import torch
import torch.nn as nn
import torch.distributed as dist
from torch.utils.data import DataLoader, Subset

import wandb

from model import MLP, temporal_patchify
from train_dynamics import (
    load_frozen_tokenizer_from_pt_ckpt, seed_everything, worker_init_fn,
    PerDomainAccumulator, init_distributed, is_rank0, _unwrap_model,
)
from task_set import TASK_SET, DOMAINS, task_to_domain
from wm_dataset import WMDataset, collate_batch

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True


class InverseDynamicsHeadTokenizer(nn.Module):
    """
    Predicts the transition action a_t from a pair of adjacent frozen-tokenizer
    latents (z_t, z_{t+1}), each (B,T,n_latents,d_bottleneck).

    Features are [z_t, z_{t+1}, z_{t+1}-z_t] flattened over the latent grid
    and concatenated -- the explicit delta term gives the head a cheap,
    direct "what changed" signal on top of the two raw codes.
    """

    def __init__(self, *, n_latents: int, d_bottleneck: int, act_dim_max: int = 16,
                 mlp_ratio: float = 4.0, dropout: float = 0.0, d_hidden: int = 0):
        super().__init__()
        d_in = n_latents * d_bottleneck
        # Width defaults to d_in (n_latents*d_bottleneck = 4096 for the released
        # tokenizer), which makes this "lightweight" head ~185M parameters --
        # comparable to the frozen dynamics model itself. --d_hidden decouples
        # the two so the head can be shrunk, and so the tokenizer/backbone arms
        # can be compared at matched capacity rather than matched feature width.
        d_h = int(d_hidden) if int(d_hidden) > 0 else d_in
        self.in_proj = nn.Linear(3 * d_in, d_h)
        self.mlp = MLP(d_h, mlp_ratio=mlp_ratio, dropout=dropout)
        self.out = nn.Linear(d_h, act_dim_max)
        # Small-normal + zero bias: predictions start near 0 (small initial
        # loss) while still letting gradient flow from step 1 -- same
        # convention as PolicyHeadMTP in model.py.
        nn.init.normal_(self.out.weight, std=0.01)
        nn.init.zeros_(self.out.bias)

    def forward(self, z_t: torch.Tensor, z_tp1: torch.Tensor) -> torch.Tensor:
        B, T = z_t.shape[:2]
        a = z_t.reshape(B, T, -1)
        b = z_tp1.reshape(B, T, -1)
        x = torch.cat([a, b, b - a], dim=-1)
        h = self.in_proj(x)
        h = h + self.mlp(h)
        return torch.tanh(self.out(h))  # (B,T,act_dim_max), matches the [-1,1] action convention


def masked_action_mse(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    diff_sq = (pred.float() - target.float()).pow(2) * mask.float()
    return diff_sq.sum() / mask.float().sum().clamp_min(1.0)


def masked_action_mse_per_sample(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """Same loss as masked_action_mse, but reduced only over (T,A), keeping the
    batch dim -- lets the caller bucket per-sample loss by task/domain."""
    diff_sq = (pred.float() - target.float()).pow(2) * mask.float()
    return diff_sq.sum(dim=(1, 2)) / mask.float().sum(dim=(1, 2)).clamp_min(1.0)  # (B,)


def lr_at_step(step: int, *, base_lr: float, warmup_steps: int, total_steps: int,
               schedule: str, min_frac: float) -> float:
    """Linear warmup, then constant or cosine decay to `min_frac * base_lr`.

    Cosine is the default because a constant post-warmup LR diverges here: on a
    ~185M-parameter head, 3e-4 held flat drives the tanh output into saturation
    and the loss settles at a degenerate ~1.5 plateau (worse than predicting
    zero) partway through training.
    """
    if warmup_steps > 0 and step < warmup_steps:
        return base_lr * (step + 1) / warmup_steps
    if schedule == "constant":
        return base_lr
    progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
    progress = min(max(progress, 0.0), 1.0)
    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
    return base_lr * (min_frac + (1.0 - min_frac) * cosine)


def split_heldout_tasks(n_heldout: int, explicit: Optional[List[str]] = None) -> Tuple[List[str], List[str]]:
    """Split TASK_SET into (train_tasks, heldout_tasks).

    Held-out tasks are picked round-robin across domains (domains sorted, then
    index within domain) so the eval set spans several domains rather than
    landing entirely in one. Deterministic: no RNG, so both IDM arms hold out
    exactly the same tasks and their eval curves are directly comparable.

    UNSEEN_TASK_SET is never used for training here. It is the paper's transfer
    split, and none of its 10 tasks ship in the standard partitions anyway --
    with strict_tasks=False they were silently dropped, so excluding them
    explicitly changes nothing in practice while removing the contamination
    risk if a targeted-collection dir that does contain them is passed later.
    """
    if explicit:
        seen = set(TASK_SET)
        heldout = [t for t in explicit if t in seen]
    else:
        by_domain: dict = {}
        for t in TASK_SET:
            by_domain.setdefault(task_to_domain(t), []).append(t)
        domains = sorted(by_domain)
        heldout, i = [], 0
        while len(heldout) < n_heldout:
            progressed = False
            for d in domains:
                if i < len(by_domain[d]):
                    heldout.append(by_domain[d][i])
                    progressed = True
                    if len(heldout) >= n_heldout:
                        break
            if not progressed:
                break
            i += 1
    hs = set(heldout)
    return [t for t in TASK_SET if t not in hs], heldout


@torch.no_grad()
def encode_window(encoder, obs_u8: torch.Tensor, patch: int) -> torch.Tensor:
    """obs_u8: (B,T+1,3,H,W) uint8 -> z_btLd: (B,T+1,n_latents,d_bottleneck)."""
    frames = obs_u8.float() / 255.0
    patches = temporal_patchify(frames, patch)
    z_btLd, _ = encoder(patches)
    return z_btLd


@torch.no_grad()
def evaluate(head_module, encoder, loader, *, device, patch: int, n_batches: int) -> float:
    """Mask-weighted action MSE over up to n_batches held-out batches.

    Uses the same sum-of-squares / sum-of-mask reduction as the training loss,
    so the two numbers are directly comparable. `head_module` must be the
    unwrapped module (not the DDP wrapper) -- eval runs on rank0 only, and
    calling the wrapper would involve the other ranks.
    """
    was_training = head_module.training
    head_module.eval()
    tot_sq, tot_mask = 0.0, 0.0
    it = iter(loader)
    for _ in range(max(1, n_batches)):
        try:
            batch = next(it)
        except StopIteration:
            break
        obs_u8 = batch["obs"].to(device, non_blocking=True)
        act = batch["act"].to(device, non_blocking=True).clamp(-1, 1)
        act_mask = batch["act_mask"].to(device, non_blocking=True)
        act = act * act_mask
        z_btLd = encode_window(encoder, obs_u8, patch)
        pred = head_module(z_btLd[:, :-1], z_btLd[:, 1:])
        diff_sq = (pred.float() - act.float()).pow(2) * act_mask.float()
        tot_sq += float(diff_sq.sum().item())
        tot_mask += float(act_mask.float().sum().item())
    if was_training:
        head_module.train()
    return tot_sq / max(tot_mask, 1.0)


def train(args):
    ddp, rank, world_size, local_rank = init_distributed()
    device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
    seed_everything(args.seed + rank)

    encoder, _decoder, tok_args = load_frozen_tokenizer_from_pt_ckpt(args.tokenizer_ckpt, device=device)
    patch = int(tok_args.get("patch", 14))
    n_latents = int(tok_args.get("n_latents", 16))
    d_bottleneck = int(tok_args.get("d_bottleneck", 32))
    if is_rank0():
        print(f"[idm-tok] tokenizer: patch={patch}, n_latents={n_latents}, d_bottleneck={d_bottleneck}")
        print(f"[idm-tok] world_size={world_size}")

    train_tasks, heldout_tasks = split_heldout_tasks(
        args.n_heldout_tasks, explicit=(args.heldout_tasks or None),
    )
    if is_rank0():
        print(f"[idm-tok] train tasks: {len(train_tasks)}  held-out: {len(heldout_tasks)} "
              f"-> {heldout_tasks if heldout_tasks else '(none — eval disabled)'}")

    dataset = WMDataset(
        data_dir=args.data_dirs,
        frames_dir=args.frame_dirs,
        seq_len=args.seq_len,
        img_size=224,
        action_dim=16,
        lang_dim=0,  # unused by this head
        tasks_json=args.tasks_json,
        tasks=train_tasks,
        verbose=is_rank0(),
        cache_mb=args.cache_mb,
        ddp_partition=True,
        iid_sampling=True,
        samples_per_shard=args.samples_per_shard,
        strict_tasks=False,  # tolerate partitions that hold only a subset of tasks
    )
    loader = DataLoader(
        dataset, batch_size=args.batch_size, shuffle=True,
        num_workers=args.num_workers, pin_memory=True, drop_last=True,
        persistent_workers=(args.num_workers > 0),
        worker_init_fn=worker_init_fn, collate_fn=collate_batch,
    )

    # Per-sample -> domain lookup, for the per-domain loss breakdown below.
    # dataset.tasks[i] is the task name for local task index i; batch["emb_id"]
    # gives each sample's local task index.
    task_idx_to_domain_idx = torch.tensor(
        [DOMAINS.index(task_to_domain(t)) for t in dataset.tasks],
        dtype=torch.long, device=device,
    )
    domain_acc = PerDomainAccumulator(n_domains=len(DOMAINS), device=device)

    # Held-out eval loader (rank0 only: it is a small extra pass and building the
    # dataset on every rank would duplicate the shard scan and its cache).
    eval_loader = None
    if is_rank0() and heldout_tasks and args.eval_every > 0:
        try:
            eval_ds = WMDataset(
                data_dir=args.data_dirs,
                frames_dir=args.frame_dirs,
                seq_len=args.seq_len,
                img_size=224,
                action_dim=16,
                lang_dim=0,
                tasks_json=args.tasks_json,
                tasks=heldout_tasks,
                verbose=False,
                cache_mb=min(args.cache_mb, 1024),
                ddp_partition=False,
                # iid_sampling=False is required for a fixed eval set: under
                # iid_sampling __getitem__ ignores idx and draws a random shard
                # and random start, so shuffle=False would still score different
                # windows on every pass and the eval curve would be pure noise.
                iid_sampling=False,
                strict_tasks=False,
            )
            # Deterministic strided subset: indices are ordered by task, so a
            # contiguous head would only cover the first held-out task. Striding
            # spans all of them, and with shuffle=False every eval scores the
            # identical windows.
            n_want = max(1, args.eval_batches * args.batch_size)
            total = len(eval_ds)
            stride = max(1, total // n_want)
            idxs = list(range(0, total, stride))[:n_want]
            eval_loader = DataLoader(
                Subset(eval_ds, idxs), batch_size=args.batch_size, shuffle=False,
                num_workers=min(2, args.num_workers), pin_memory=True, drop_last=False,
                worker_init_fn=worker_init_fn, collate_fn=collate_batch,
            )
            print(f"[idm-tok] eval: {len(idxs)} fixed windows (strided from {total}) "
                  f"across {len(heldout_tasks)} held-out tasks, every {args.eval_every} steps")
        except Exception as e:
            print(f"[idm-tok] WARNING: could not build held-out eval set "
                  f"({type(e).__name__}: {e}) — eval disabled.")
            eval_loader = None

    head = InverseDynamicsHeadTokenizer(
        n_latents=n_latents, d_bottleneck=d_bottleneck,
        act_dim_max=16, mlp_ratio=args.mlp_ratio, dropout=args.dropout,
        d_hidden=args.d_hidden,
    ).to(device)
    if is_rank0():
        n_params = sum(p.numel() for p in head.parameters())
        print(f"[idm-tok] head parameters: {n_params/1e6:.1f}M")
    if ddp:
        head = torch.nn.parallel.DistributedDataParallel(
            head, device_ids=[local_rank], output_device=local_rank, broadcast_buffers=False,
        )
    opt = torch.optim.AdamW(head.parameters(), lr=args.lr, weight_decay=args.weight_decay)

    # Resume. Weights always; optimizer state only if the checkpoint carries it
    # (checkpoints written before --resume existed do not, so AdamW's moments
    # restart from zero there -- the post-resume warmup below cushions that).
    start_step = 0
    if args.resume:
        ck = torch.load(args.resume, map_location="cpu", weights_only=False)
        _unwrap_model(head).load_state_dict(ck["model"])
        has_opt = isinstance(ck, dict) and ck.get("optim") is not None
        if has_opt:
            opt.load_state_dict(ck["optim"])
        if args.resume_step:
            start_step = int(ck.get("step", 0))
        if is_rank0():
            opt_note = ("optimizer state restored" if has_opt else
                        "no optimizer state in checkpoint — AdamW moments restart from "
                        "zero, expect a transient loss bump")
            print(f"[idm-tok] resumed from {args.resume} (step={ck.get('step')}); {opt_note}")
            print(f"[idm-tok] continuing {start_step} -> {args.steps} "
                  f"({max(0, args.steps - start_step)} new steps)")
        if start_step >= args.steps and is_rank0():
            print(f"[idm-tok] WARNING: --steps {args.steps} is not beyond the resumed step "
                  f"{start_step}; nothing to train. Raise --steps.")

    out_path = Path(args.out_ckpt)
    if is_rank0():
        out_path.parent.mkdir(parents=True, exist_ok=True)
        wandb.init(
            project=args.wandb_project,
            name=args.wandb_run_name,
            entity=args.wandb_entity,
            mode="online",
            config={**vars(args), "run/world_size": world_size},
        )

    def save(step):
        if not is_rank0():
            return
        # Optimizer state travels with the checkpoint so a later --resume is exact.
        payload = {"model": _unwrap_model(head).state_dict(), "optim": opt.state_dict(),
                   "args": vars(args), "step": step}
        torch.save(payload, out_path)
        step_path = out_path.with_name(f"step_{step:06d}.pt")
        torch.save(payload, step_path)
        print(f"[idm-tok] saved checkpoint -> {out_path} and {step_path}")

    step = start_step
    t0 = time.time()
    data_iter = iter(loader)
    while step < args.steps:
        try:
            batch = next(data_iter)
        except StopIteration:
            data_iter = iter(loader)
            batch = next(data_iter)

        obs_u8 = batch["obs"].to(device, non_blocking=True)             # (B,T+1,3,H,W)
        act = batch["act"].to(device, non_blocking=True).clamp(-1, 1)   # (B,T,A)
        act_mask = batch["act_mask"].to(device, non_blocking=True)      # (B,T,A)
        act = act * act_mask

        z_btLd = encode_window(encoder, obs_u8, patch)  # (B,T+1,n_latents,d_b), no_grad internally
        z_t = z_btLd[:, :-1]
        z_tp1 = z_btLd[:, 1:]

        pred = head(z_t, z_tp1)
        loss = masked_action_mse(pred, act, act_mask)

        opt.zero_grad(set_to_none=True)
        loss.backward()
        # Grad clip + LR warmup: the head's `out` layer is deliberately small-
        # initialized (see InverseDynamicsHeadTokenizer) so it starts near a
        # sane [-1,1] prediction, but AdamW's adaptive step size ignores that
        # init scale and can push the tanh into saturation within the first
        # few dozen steps, permanently killing its gradient. Both guard
        # against that; mirrors train_dynamics.py's grad-clip convention.
        clip = args.grad_clip if args.grad_clip > 0 else float("inf")
        grad_norm = float(torch.nn.utils.clip_grad_norm_(head.parameters(), max_norm=clip).item())
        # Schedule spans the steps this invocation actually runs. On a resume that
        # means a fresh warmup (which cushions a restarted optimizer) then cosine
        # over the added steps — rather than re-entering the old cosine midway,
        # which would jump the LR back up.
        lr_now = lr_at_step(
            step - start_step, base_lr=args.lr, warmup_steps=args.warmup_steps,
            total_steps=args.steps - start_step,
            schedule=args.lr_schedule, min_frac=args.lr_min_frac,
        )
        for pg in opt.param_groups:
            pg["lr"] = lr_now
        opt.step()

        with torch.no_grad():
            per_sample_loss = masked_action_mse_per_sample(pred, act, act_mask)  # (B,)
            emb_id_batch = batch["emb_id"].to(device, non_blocking=True).long()  # (B,)
            domain_ids = task_idx_to_domain_idx[emb_id_batch]
            domain_acc.update(per_sample_loss, domain_ids)

        do_log = (step % args.log_every == 0)
        domain_means = domain_counts = None
        if do_log:
            # flush() all-reduces across ranks -- every rank must call it,
            # even though only rank0 goes on to print/log the result.
            domain_means, domain_counts = domain_acc.flush(ddp=ddp)

        if do_log and is_rank0():
            dt = time.time() - t0
            print(f"[idm-tok] step {step:6d}  loss {loss.item():.6f}  ({dt:.1f}s)")

            wandb_payload = {
                "loss/step_loss": loss.item(),
                "stats/grad_norm": grad_norm,
                "stats/lr": opt.param_groups[0]["lr"],
            }
            all_mask = domain_counts > 0
            if all_mask.any():
                all_domain_loss = float(
                    (domain_means[all_mask] * domain_counts[all_mask]).sum()
                    / domain_counts[all_mask].sum()
                )
                print(f"[idm-tok]   action_loss/all_domains = {all_domain_loss:.6f}")
                wandb_payload["loss/all_domains"] = all_domain_loss
                for di, dname in enumerate(DOMAINS):
                    if domain_counts[di] > 0:
                        m = float(domain_means[di].item())
                        print(f"[idm-tok]   action_loss/{dname} = {m:.6f} "
                              f"(n={int(domain_counts[di].item())})")
                        wandb_payload[f"domain/{dname}/action_loss"] = m
            wandb.log(wandb_payload, step=step)

        # Held-out eval. Only rank0 has a loader; the barrier keeps the other
        # ranks from racing ahead into the next DDP allreduce while it runs.
        if args.eval_every > 0 and step % args.eval_every == 0:
            if eval_loader is not None:
                eval_loss = evaluate(
                    _unwrap_model(head), encoder, eval_loader,
                    device=device, patch=patch, n_batches=args.eval_batches,
                )
                print(f"[idm-tok] step {step:6d}  eval/action_loss {eval_loss:.6f}")
                wandb.log({"eval/action_loss": eval_loss}, step=step)
            if ddp:
                dist.barrier()

        step += 1
        if args.save_every > 0 and step % args.save_every == 0:
            save(step)

    save(step)

    if ddp:
        dist.barrier()
        dist.destroy_process_group()


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--data_dirs", type=str, nargs="+", default=["/data5/expert"],
                   help="raw partition dir(s) (holds per-task action/reward .pt files)")
    p.add_argument("--frame_dirs", type=str, nargs="+", default=["/data5/expert-shards"],
                   help="preprocessed shard dir(s) (holds per-task frame shards)")
    p.add_argument("--tasks_json", type=str, default="../tasks.json")
    p.add_argument("--tokenizer_ckpt", type=str, default="./logs/tokenizer_ckpts/latest.pt")
    p.add_argument("--seq_len", type=int, default=8,
                   help="window length T; yields T (obs_t -> obs_{t+1}) pairs per sample")
    p.add_argument("--batch_size", type=int, default=64)
    p.add_argument("--num_workers", type=int, default=4)
    p.add_argument("--samples_per_shard", type=int, default=16)
    p.add_argument("--cache_mb", type=int, default=4096)
    p.add_argument("--mlp_ratio", type=float, default=4.0)
    p.add_argument("--dropout", type=float, default=0.0)
    p.add_argument("--lr", type=float, default=3e-4)
    p.add_argument("--weight_decay", type=float, default=0.0)
    p.add_argument("--warmup_steps", type=int, default=1000,
                   help="linear LR warmup from 0 -> --lr; guards against AdamW's early "
                        "adaptive step blowing through the head's small tanh-output init "
                        "and saturating it before it has learned anything")
    p.add_argument("--grad_clip", type=float, default=1.0, help="0 disables clipping")
    p.add_argument("--lr_schedule", type=str, default="cosine", choices=["cosine", "constant"],
                   help="post-warmup LR schedule. 'cosine' decays to --lr_min_frac * --lr by "
                        "--steps; 'constant' reproduces the old behavior, which diverged "
                        "partway through training (tanh saturation, loss plateau ~1.5).")
    p.add_argument("--lr_min_frac", type=float, default=0.1,
                   help="cosine floor as a fraction of --lr")
    p.add_argument("--d_hidden", type=int, default=0,
                   help="head hidden width; 0 = n_latents*d_bottleneck (the old default, "
                        "which makes the head ~185M params). Set e.g. 1024 for a genuinely "
                        "lightweight head, or to capacity-match against idm_backbone.py.")
    p.add_argument("--n_heldout_tasks", type=int, default=10,
                   help="number of TASK_SET tasks held out of training for eval, picked "
                        "round-robin across domains; 0 disables the held-out split")
    p.add_argument("--heldout_tasks", type=str, nargs="*", default=None,
                   help="explicit held-out task names (overrides --n_heldout_tasks)")
    p.add_argument("--eval_every", type=int, default=1000,
                   help="run held-out eval every N steps; 0 disables")
    p.add_argument("--eval_batches", type=int, default=20,
                   help="number of held-out batches per eval")
    p.add_argument("--steps", type=int, default=20_000)
    p.add_argument("--log_every", type=int, default=50)
    p.add_argument("--save_every", type=int, default=1000)
    p.add_argument("--out_ckpt", type=str, default="./logs/idm_tokenizer_ckpts/latest.pt")
    p.add_argument("--resume", type=str, default=None,
                   help="checkpoint to resume the head from. Optimizer state is restored too "
                        "when the checkpoint carries it; older ones do not, and AdamW's "
                        "moments then restart from zero (the post-resume warmup cushions it).")
    p.add_argument("--resume_step", action=argparse.BooleanOptionalAction, default=True,
                   help="continue the global step counter from the checkpoint, so --steps is a "
                        "TOTAL and new checkpoints do not overwrite old ones. "
                        "--no-resume_step restarts counting at 0.")
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--wandb_project", type=str, default="mmbench2-idm-tokenizer")
    p.add_argument("--wandb_run_name", type=str, default="default")
    p.add_argument("--wandb_entity", type=str, default=None)
    args = p.parse_args()
    train(args)