# 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)