src/idm_backbone.py

# idm_backbone.py
"""
Inverse dynamics model (IDM) -- frozen Dynamics backbone.

Predicts the action a_t that caused the transition obs_t -> obs_{t+1} from
the (frozen) Dynamics transformer's contextual hidden features at frame t and
t+1, rather than from raw tokenizer latents alone (see idm_tokenizer.py for
that cheaper alternative). The premise: the dynamics model has seen far more
training signal about how actions relate to state changes than the tokenizer
alone, so its hidden features may carry more of the inverse-dynamics signal
directly. Whether that premise pays off vs. idm_tokenizer.py's cheaper
features is an empirical question this script lets you test.

How features are extracted, and why this doesn't leak the answer:
  - The whole observed window is encoded through the frozen tokenizer to
    clean packed latents, then run through the frozen `Dynamics` transformer
    in a single forward pass with `actions=None` at every slot. Passing no
    actions means the model never sees the ground-truth transition actions
    (what we're trying to predict) -- every slot only gets the learned
    "no action" base token (see ActionEncoder in model.py).
  - Every slot is presented as a fully clean / "context" frame: step_idx =
    e_max, signal_idx = k_max. This is the exact convention already used
    elsewhere in the repo for already-known past context tokens (see
    `sample_one_timestep_packed`'s `tau_ctx=0` case in train_dynamics.py) --
    it is not a novel edge case.
  - Time-attention in `Dynamics` is causal, so the hidden state at slot t+1
    has (in addition to its own frame content) full causal access to frames
    0..t -- more context than a pairwise (t, t+1)-only encoder would have,
    which may help disambiguate transitions that depend on recent motion
    (e.g. velocity/momentum) rather than a single f
    rame pair.
  - We read off `x1_hat` (the pre-existing flow-target output, always
    available) and `h_t` (agent-token features, only when the checkpoint has
    `n_agent > 0`) at slots t and t+1, and train a small head on top.

Only a new lightweight head (`InverseDynamicsHeadBackbone`) is trained; the
tokenizer and dynamics model stay frozen throughout.

Run from inside `src/`, single-GPU:
    python idm_backbone.py \
        --tokenizer_ckpt ./logs/tokenizer_ckpts/latest.pt \
        --dynamics_ckpt ./logs/dynamics_ckpts/latest.pt \
        --data_dirs ./data/expert --frame_dirs ./data/expert-shards \
        --out_ckpt ./logs/idm_backbone_ckpts/latest.pt

or multi-GPU (DDP, same flags, launched via torchrun -- mirrors
idm_tokenizer.py; only the new head is wrapped in DDP, the frozen tokenizer
and Dynamics backbone run un-replicated on each rank):
    torchrun --nproc_per_node=8 idm_backbone.py \
        --tokenizer_ckpt ./logs/tokenizer_ckpts/latest.pt \
        --dynamics_ckpt ./logs/dynamics_ckpts/latest.pt \
        --data_dirs ./data/expert --frame_dirs ./data/expert-shards \
        --out_ckpt ./logs/idm_backbone_ckpts/latest.pt
"""
import argparse
import math
import time
from pathlib import Path
from typing import Optional

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

import wandb

from model import MLP, temporal_patchify, pack_bottleneck_to_spatial
from train_dynamics import (
    load_frozen_tokenizer_from_pt_ckpt, seed_everything, worker_init_fn,
    PerDomainAccumulator, init_distributed, is_rank0, _unwrap_model,
)
from interactive import load_dynamics_from_ckpt
from task_set import DOMAINS, task_to_domain
from wm_dataset import WMDataset, collate_batch
# Shared with the tokenizer-only arm so both hold out exactly the same tasks —
# otherwise the two eval curves are not comparable.
from idm_tokenizer import split_heldout_tasks

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


class InverseDynamicsHeadBackbone(nn.Module):
    """
    Predicts the transition action a_t from the frozen Dynamics transformer's
    hidden features at frame t and t+1: the (always-available) spatial flow
    features `x1_hat`, plus the (checkpoint-dependent) pooled agent features
    `h_t`, if the loaded Dynamics checkpoint has n_agent > 0.

    Mirrors idm_tokenizer.py's [feat_t, feat_tp1, feat_tp1-feat_t] pattern, but
    unlike the tokenizer's bottleneck latents (always tanh-bounded to [-1,1]),
    these Dynamics features have no such guarantee: `x1_hat` is a plain-Linear
    flow-matching output and `h_t` is the raw pre-final-norm residual stream
    (BlockCausalTransformer applies no closing norm), whose scale can grow
    with depth. LayerNorm each feature block before concatenation so no block
    dominates in_proj purely from an accidentally larger frozen-model scale.
    """

    def __init__(self, *, n_spatial: int, d_spatial: int, d_model: int, has_agent: bool,
                 act_dim_max: int = 16, mlp_ratio: float = 4.0, dropout: float = 0.0,
                 d_hidden: int = 0):
        super().__init__()
        self.has_agent = bool(has_agent)
        d_spatial_flat = n_spatial * d_spatial
        d_in = 3 * d_spatial_flat
        if self.has_agent:
            d_in += 2 * d_model
        # Width defaults to n_spatial*d_spatial (4096 for the released
        # checkpoints), making this "lightweight" head ~193M parameters.
        # --d_hidden decouples the two so the head can be shrunk and so both
        # IDM arms can be compared at matched capacity.
        d_hidden = int(d_hidden) if int(d_hidden) > 0 else d_spatial_flat

        self.norm_spatial = nn.LayerNorm(d_spatial_flat)
        if self.has_agent:
            self.norm_agent = nn.LayerNorm(d_model)

        self.in_proj = nn.Linear(d_in, d_hidden)
        self.mlp = MLP(d_hidden, mlp_ratio=mlp_ratio, dropout=dropout)
        self.out = nn.Linear(d_hidden, act_dim_max)
        nn.init.normal_(self.out.weight, std=0.01)
        nn.init.zeros_(self.out.bias)

    def forward(
        self,
        spatial_t: torch.Tensor,       # (B,T,n_spatial,d_spatial)
        spatial_tp1: torch.Tensor,     # (B,T,n_spatial,d_spatial)
        agent_t: Optional[torch.Tensor] = None,     # (B,T,d_model)
        agent_tp1: Optional[torch.Tensor] = None,   # (B,T,d_model)
    ) -> torch.Tensor:
        B, T = spatial_t.shape[:2]
        a = self.norm_spatial(spatial_t.reshape(B, T, -1))
        b = self.norm_spatial(spatial_tp1.reshape(B, T, -1))
        feats = [a, b, b - a]
        if self.has_agent:
            assert agent_t is not None and agent_tp1 is not None
            feats += [self.norm_agent(agent_t), self.norm_agent(agent_tp1)]
        x = torch.cat(feats, dim=-1)
        h = self.in_proj(x)
        h = h + self.mlp(h)
        return torch.tanh(self.out(h))  # (B,T,act_dim_max)


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
    ~193M-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) within the first few thousand steps.
    """
    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)


@torch.no_grad()
def extract_backbone_features(dyn, encoder, obs_u8: torch.Tensor, *, patch: int,
                               packing_factor: int, n_spatial: int, k_max: int,
                               lang_emb: Optional[torch.Tensor], use_amp: bool, device,
                               clean_signal_idx: str = "trained_max"):
    """
    obs_u8: (B,T+1,3,H,W) uint8 -> (spatial_t, spatial_tp1, agent_t, agent_tp1),
    each aligned to the T transitions obs_t -> obs_{t+1}. agent_* are None when
    the loaded Dynamics checkpoint has n_agent == 0.

    On `clean_signal_idx`: the nominal "fully clean" signal index is k_max, but
    dynamics training never samples it. `_sample_tau_for_step` yields
    tau_idx = j_idx * (k_max // K) with j_idx < K, so the largest index seen
    under gradient is k_max - 1; the only two places k_max appears are the
    bootstrap target forward (inside torch.no_grad) and the inference-time
    context path, where the flow output is discarded. `signal_embed.weight[k_max]`
    therefore receives zero gradient in every training phase and stays at
    nn.Embedding's default N(0,1) init -- and it fills half the shortcut token's
    channels. Conditioning feature extraction on it means conditioning on a
    random vector the model has never seen, so we default to k_max - 1.
    Pass "k_max" to reproduce the original behavior.
    """
    frames = obs_u8.float() / 255.0
    patches = temporal_patchify(frames, patch)
    z_btLd, _ = encoder(patches)                                            # (B,T+1,n_latents,d_b)
    z_packed = pack_bottleneck_to_spatial(z_btLd, n_spatial=n_spatial, k=packing_factor)  # (B,T+1,Sz,Dz)

    B, Tw = z_packed.shape[:2]
    emax = int(round(math.log2(k_max)))
    sig_value = k_max if clean_signal_idx == "k_max" else max(0, k_max - 1)
    step_idxs = torch.full((B, Tw), emax, device=device, dtype=torch.long)
    signal_idxs = torch.full((B, Tw), sig_value, device=device, dtype=torch.long)

    with autocast(device_type=device.type, enabled=(use_amp and device.type == "cuda"), dtype=torch.bfloat16):
        x1_hat, h_t = dyn(
            None,                       # actions=None: never see the ground-truth transition action
            step_idxs, signal_idxs, z_packed,
            act_mask=None, agent_tokens=None, lang_emb=lang_emb,
        )
    x1_hat = x1_hat.float()
    spatial_t, spatial_tp1 = x1_hat[:, :-1], x1_hat[:, 1:]

    agent_t = agent_tp1 = None
    if h_t is not None:
        h_t = h_t.float().mean(dim=2)  # mean-pool over n_agent -> (B,T+1,d_model)
        agent_t, agent_tp1 = h_t[:, :-1], h_t[:, 1:]

    return spatial_t, spatial_tp1, agent_t, agent_tp1


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

    Same reduction as the training loss, so the numbers are directly
    comparable. `head_module` must be the unwrapped module -- eval runs on
    rank0 only, and calling the DDP 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
        lang_emb = batch["lang_emb"].to(device, non_blocking=True) if lang_dim > 0 else None
        s_t, s_tp1, a_t, a_tp1 = extract_backbone_features(
            dyn, encoder, obs_u8, lang_emb=lang_emb, device=device, **feat_kwargs,
        )
        pred = head_module(s_t, s_tp1, a_t, a_tp1)
        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))

    dyn, _rew_head, _policy_head, dyn_meta = load_dynamics_from_ckpt(
        args.dynamics_ckpt, device=device,
        d_bottleneck=d_bottleneck, n_latents=n_latents, packing_factor=args.packing_factor,
    )
    k_max = dyn_meta["k_max"]
    n_spatial = dyn_meta["n_spatial"]
    d_spatial = dyn_meta["d_spatial"]
    d_model = dyn_meta["d_model"]
    lang_dim = dyn_meta["lang_dim"]
    has_agent = (dyn.n_agent > 0)
    if is_rank0():
        print(f"[idm-bb] tokenizer: patch={patch}, n_latents={n_latents}, d_bottleneck={d_bottleneck}")
        print(f"[idm-bb] dynamics: k_max={k_max}, n_spatial={n_spatial}, d_spatial={d_spatial}, "
              f"d_model={d_model}, lang_dim={lang_dim}, n_agent={dyn.n_agent} (has_agent={has_agent})")
        print(f"[idm-bb] 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-bb] 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=lang_dim,  # must match the loaded Dynamics' task_proj input dim
        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)

    # Feature-extraction kwargs, shared by the train loop and the eval pass.
    feat_kwargs = dict(
        patch=patch, packing_factor=args.packing_factor, n_spatial=n_spatial,
        k_max=k_max, use_amp=args.amp, clean_signal_idx=args.clean_signal_idx,
    )
    if is_rank0():
        sig_used = k_max if args.clean_signal_idx == "k_max" else k_max - 1
        print(f"[idm-bb] clean_signal_idx={args.clean_signal_idx} -> signal_idx={sig_used} "
              f"(k_max={k_max}; row {k_max} is untrained)")

    # Held-out eval loader (rank0 only: 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=lang_dim,
                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-bb] 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-bb] WARNING: could not build held-out eval set "
                  f"({type(e).__name__}: {e}) — eval disabled.")
            eval_loader = None

    head = InverseDynamicsHeadBackbone(
        n_spatial=n_spatial, d_spatial=d_spatial, d_model=d_model, has_agent=has_agent,
        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-bb] 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)

    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
        payload = {"model": _unwrap_model(head).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-bb] saved checkpoint -> {out_path} and {step_path}")

    step = 0
    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
        lang_emb = batch["lang_emb"].to(device, non_blocking=True) if lang_dim > 0 else None

        spatial_t, spatial_tp1, agent_t, agent_tp1 = extract_backbone_features(
            dyn, encoder, obs_u8, lang_emb=lang_emb, device=device, **feat_kwargs,
        )

        pred = head(spatial_t, spatial_tp1, agent_t, agent_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 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())
        lr_now = lr_at_step(
            step, base_lr=args.lr, warmup_steps=args.warmup_steps,
            total_steps=args.steps, 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-bb] 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-bb]   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-bb]   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), dyn, encoder, eval_loader,
                    device=device, feat_kwargs=feat_kwargs,
                    n_batches=args.eval_batches, lang_dim=lang_dim,
                )
                print(f"[idm-bb] 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("--dynamics_ckpt", type=str, default="./logs/dynamics_ckpts/latest.pt")
    p.add_argument("--packing_factor", type=int, default=2)
    p.add_argument("--seq_len", type=int, default=8,
                   help="window length T; yields T (obs_t -> obs_{t+1}) pairs per sample, "
                        "each with causal access to frames 0..t via the Dynamics time-attention")
    p.add_argument("--batch_size", type=int, default=32)
    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 "
                        "within a few thousand steps (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_spatial*d_spatial (the old default, "
                        "which makes the head ~193M params). Set e.g. 1024 for a genuinely "
                        "lightweight head, or to capacity-match against idm_tokenizer.py.")
    p.add_argument("--clean_signal_idx", type=str, default="trained_max",
                   choices=["trained_max", "k_max"],
                   help="which signal_embed row represents a fully-clean frame during feature "
                        "extraction. 'trained_max' = k_max-1, the highest index dynamics "
                        "training ever samples. 'k_max' is the nominal clean index but its "
                        "embedding row receives zero gradient in training and stays at random "
                        "init; it reproduces the original behavior.")
    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_backbone_ckpts/latest.pt")
    p.add_argument("--amp", action=argparse.BooleanOptionalAction, default=True)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--wandb_project", type=str, default="mmbench2-idm-backbone")
    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)