# train_dynamics.py """Train the Dreamer 4 dynamics model: a block-causal Transformer over frozen tokenizer latents, trained with shortcut flow-matching and action conditioning (plus optional reward and behavior-cloning heads and coverage-aware sampling). Run from inside ``src/`` (flat imports), e.g. ``torchrun --nproc_per_node=8 train_dynamics.py``. """ import os import time import math import random import argparse from contextlib import nullcontext from pathlib import Path from typing import Optional, Dict, Any import numpy as np import torch import torch.nn.functional as F import torch.distributed as dist from torch.amp import autocast from torch.utils.data import DataLoader import wandb from task_set import TASK_SET, DOMAINS, UNSEEN_TASK_SET, task_to_domain, compute_task_weights from model import ( Encoder, Decoder, Tokenizer, temporal_patchify, temporal_unpatchify, pack_bottleneck_to_spatial, unpack_spatial_to_bottleneck, Dynamics, RewardHeadMTP, PolicyHeadMTP, symlog, dist_cross_entropy_from_symlog, EmaRms, ) torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True class PerDomainAccumulator: """Buckets per-sample losses by domain and aggregates across DDP ranks. Each rank keeps running sums/counts of losses between log flushes, bucketed by the sample's domain index (0..n_domains-1). flush() optionally all-reduces across ranks and returns per-domain means. """ def __init__(self, n_domains: int, device: torch.device): self.n_domains = int(n_domains) self.device = device self.loss_sum = torch.zeros(self.n_domains, device=device, dtype=torch.float64) self.count = torch.zeros(self.n_domains, device=device, dtype=torch.float64) @torch.no_grad() def update(self, loss_per_sample: torch.Tensor, domain_ids: torch.Tensor): # loss_per_sample: (B,) float; domain_ids: (B,) int in [0, n_domains). l = loss_per_sample.detach().to(device=self.device, dtype=torch.float64) d = domain_ids.to(device=self.device, dtype=torch.long) self.loss_sum.index_add_(0, d, l) self.count.index_add_(0, d, torch.ones_like(l)) @torch.no_grad() def flush(self, ddp: bool) -> tuple[torch.Tensor, torch.Tensor]: loss_sum = self.loss_sum.clone() count = self.count.clone() if ddp: dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM) dist.all_reduce(count, op=dist.ReduceOp.SUM) means = torch.where(count > 0, loss_sum / count.clamp_min(1.0), torch.full_like(loss_sum, float("nan"))) self.loss_sum.zero_() self.count.zero_() return means, count def get_dist_info(): rank = int(os.environ.get("RANK", "0")) world_size = int(os.environ.get("WORLD_SIZE", "1")) local_rank = int(os.environ.get("LOCAL_RANK", "0")) return rank, world_size, local_rank def is_rank0() -> bool: return int(os.environ.get("RANK", "0")) == 0 def seed_everything(seed: int): s = int(seed) % (2**32) random.seed(s) np.random.seed(s) torch.manual_seed(s) torch.cuda.manual_seed_all(s) def worker_init_fn(worker_id: int): info = torch.utils.data.get_worker_info() seed_everything(info.seed) def init_distributed() -> tuple[bool, int, int, int]: rank, world_size, local_rank = get_dist_info() ddp = world_size > 1 if ddp: # 60-min collective timeout (default is 10 min). Needed for runs that do # long single-rank work like multi-episode env evaluation, where rank 0 # stays in eval while the other ranks block at the next gradient # all-reduce. import datetime dist.init_process_group( backend="nccl", init_method="env://", timeout=datetime.timedelta(minutes=60), ) torch.cuda.set_device(local_rank) return ddp, rank, world_size, local_rank def _unwrap_model(model): """Strip DDP and torch.compile wrappers to get the raw nn.Module.""" target = model if hasattr(target, "module"): # DDP wrapper target = target.module if hasattr(target, "_orig_mod"): # torch.compile wrapper target = target._orig_mod return target def save_ckpt(path: Path, *, step: int, epoch: int, dyn_model, rew_head, policy_head, opt, args: argparse.Namespace, rms_state: dict = None): path.parent.mkdir(parents=True, exist_ok=True) rew_head_state = _unwrap_model(rew_head).state_dict() if rew_head is not None else None policy_head_state = _unwrap_model(policy_head).state_dict() if policy_head is not None else None obj = { "step": step, "epoch": epoch, "dynamics": _unwrap_model(dyn_model).state_dict(), "rew_head": rew_head_state, "policy_head": policy_head_state, "opt": opt.state_dict(), "args": vars(args), } if rms_state is not None: obj["rms_state"] = rms_state tmp = path.with_suffix(".tmp") torch.save(obj, tmp) tmp.replace(path) def load_ckpt(path: Path, *, dyn_model, rew_head, policy_head, opt, rms_objects: dict = None, strict: bool = True) -> tuple[int, int]: ckpt = torch.load(path, map_location="cpu") state = ckpt["dynamics"] result = _unwrap_model(dyn_model).load_state_dict(state, strict=strict) if not strict and is_rank0(): if result.missing_keys: print(f"[rank0][info] Missing keys (randomly initialized): {result.missing_keys}") if result.unexpected_keys: print(f"[rank0][info] Unexpected keys (ignored): {result.unexpected_keys}") # Reward head is optional (may be absent from older checkpoints) if rew_head is not None and ckpt.get("rew_head") is not None: _unwrap_model(rew_head).load_state_dict(ckpt["rew_head"], strict=True) # Policy (BC) head is optional (may be absent from older checkpoints) if policy_head is not None and ckpt.get("policy_head") is not None: _unwrap_model(policy_head).load_state_dict(ckpt["policy_head"], strict=True) # Optimizer state may be incompatible when enabling rewards/BC mid-run; fall back gracefully. try: opt.load_state_dict(ckpt["opt"]) except Exception as e: if is_rank0(): print(f"[rank0][warning] Could not load optimizer state (likely param mismatch after enabling rewards/BC): {e}") print("[rank0][warning] Continuing with freshly-initialized optimizer state.") if rms_objects is not None: for k, v in rms_objects.items(): if k in ckpt.get("rms_state", {}): v.load_state_dict(ckpt["rms_state"][k]) return int(ckpt.get("step", 0)), int(ckpt.get("epoch", 0)) @torch.no_grad() def load_frozen_tokenizer_from_pt_ckpt( ckpt_path: str, *, device: torch.device, override: Optional[Dict[str, Any]] = None, ) -> tuple[Encoder, Decoder, Dict[str, Any]]: ckpt = torch.load(ckpt_path, map_location="cpu") tok_args = dict(ckpt.get("args", {})) if override: tok_args.update(override) # Required keys (fall back to defaults if missing) H = int(tok_args.get("H", 224)) W = int(tok_args.get("W", 224)) C = int(tok_args.get("C", 3)) patch = int(tok_args.get("patch", 4)) n_patches = (H // patch) * (W // patch) d_patch = patch * patch * C enc = Encoder( patch_dim=d_patch, d_model=int(tok_args.get("d_model", 256)), n_latents=int(tok_args.get("n_latents", 16)), n_patches=n_patches, n_heads=int(tok_args.get("n_heads", 4)), depth=int(tok_args.get("depth", 8)), d_bottleneck=int(tok_args.get("d_bottleneck", 32)), dropout=0.0, mlp_ratio=float(tok_args.get("mlp_ratio", 4.0)), time_every=int(tok_args.get("time_every", 1)), latents_only_time=bool(tok_args.get("latents_only_time", True)), mae_p_min=0.0, mae_p_max=0.0, ) dec = Decoder( d_bottleneck=int(tok_args.get("d_bottleneck", 32)), d_model=int(tok_args.get("d_model", 256)), n_heads=int(tok_args.get("n_heads", 4)), depth=int(tok_args.get("depth", 8)), n_latents=int(tok_args.get("n_latents", 16)), n_patches=n_patches, d_patch=d_patch, dropout=0.0, mlp_ratio=float(tok_args.get("mlp_ratio", 4.0)), time_every=int(tok_args.get("time_every", 1)), latents_only_time=bool(tok_args.get("latents_only_time", True)), ) tok = Tokenizer(enc, dec) state = ckpt["model"] # Strip _orig_mod. prefix added by torch.compile state = {k.removeprefix("_orig_mod."): v for k, v in state.items()} tok.load_state_dict(state, strict=True) tok = tok.to(device) tok.eval() for p in tok.parameters(): p.requires_grad_(False) return tok.encoder, tok.decoder, tok_args def _emax_from_kmax(k_max: int) -> int: emax = int(round(math.log2(k_max))) assert (1 << emax) == k_max, "k_max must be power of two" return emax def _sample_step_excluding_dmin(device: torch.device, B: int, T: int, k_max: int) -> tuple[torch.Tensor, torch.Tensor]: emax = _emax_from_kmax(k_max) # step_idx in [0, emax) i.e. excludes emax (d_min) step_idx = torch.randint(low=0, high=max(1, emax), size=(B, T), device=device, dtype=torch.long) d = 1.0 / (1 << step_idx).to(torch.float32) return d, step_idx def _sample_tau_for_step(device: torch.device, B: int, T: int, k_max: int, step_idx: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: # K = 2^step_idx K = (1 << step_idx).to(torch.long) # (B,T) u = torch.rand((B, T), device=device, dtype=torch.float32) j_idx = torch.floor(u * K.to(torch.float32)).to(torch.long) # (B,T) in [0,K) tau = j_idx.to(torch.float32) / K.to(torch.float32) # (B,T) scale = torch.div(torch.tensor(k_max, device=device), K, rounding_mode="floor") # (B,T) tau_idx = j_idx * scale # (B,T) <= k_max-1 return tau, tau_idx def dynamics_pretrain_loss( dynamics: torch.nn.Module, *, z1: torch.Tensor, # (B,T,Sz,Dz) packed clean targets actions: Optional[torch.Tensor], # (B,T,A) led-to convention: actions[t] led to obs[t]; actions[:,0]=0 act_mask: Optional[torch.Tensor], # (B,T,A) per-task per-dim validity, same led-to shift as actions k_max: int, B_self: int, step: int, bootstrap_start: int, lang_emb: Optional[torch.Tensor] = None, # (B,lang_dim) task embedding rew_head: Optional[torch.nn.Module] = None, rewards: Optional[torch.Tensor] = None, # (B,T) aligned to frames/actions (see train loop) reward_weight: float = 0.0, policy_head: Optional[torch.nn.Module] = None, bc_weight: float = 0.0, rms_flow: Optional[EmaRms] = None, rms_rew: Optional[EmaRms] = None, rms_bc: Optional[EmaRms] = None, ) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]: device = z1.device B, T = z1.shape[:2] assert 0 <= B_self < B B_emp = B - B_self emax = _emax_from_kmax(k_max) # action mask slices act_mask_full = act_mask act_mask_self = None if act_mask_full is None else act_mask_full[B_emp:] # step idx: empirical rows are finest (d_min), self rows sample coarser step_idx_emp = torch.full((B_emp, T), emax, device=device, dtype=torch.long) if B_self > 0: d_self, step_idx_self = _sample_step_excluding_dmin(device, B_self, T, k_max) step_idx_full = torch.cat([step_idx_emp, step_idx_self], dim=0) else: d_self = torch.zeros((0, T), device=device, dtype=torch.float32) step_idx_self = torch.zeros((0, T), device=device, dtype=torch.long) step_idx_full = step_idx_emp # sigma/tau per row/time sigma_full, sigma_idx_full = _sample_tau_for_step(device, B, T, k_max, step_idx_full) sigma_emp = sigma_full[:B_emp] sigma_self = sigma_full[B_emp:] sigma_idx_self = sigma_idx_full[B_emp:] # Corrupt inputs z0_full = torch.randn_like(z1) z_tilde_full = (1.0 - sigma_full)[..., None, None] * z0_full + sigma_full[..., None, None] * z1 z_tilde_self = z_tilde_full[B_emp:] # Weights w_emp = 0.9 * sigma_emp + 0.1 w_self = 0.9 * sigma_self + 0.1 lang_emb_self = lang_emb[B_emp:] if lang_emb is not None else None # Main forward z1_hat_full, h_t_full = dynamics(actions, step_idx_full, sigma_idx_full, z_tilde_full, act_mask=act_mask_full, agent_tokens=None, lang_emb=lang_emb) z1_hat_emp = z1_hat_full[:B_emp] z1_hat_self = z1_hat_full[B_emp:] flow_per = (z1_hat_emp.float() - z1[:B_emp].float()).pow(2).mean(dim=(2, 3)) # (B_emp,T) loss_emp = (flow_per * w_emp).mean() boot_mse = torch.zeros((), device=device, dtype=torch.float32) loss_self = torch.zeros((), device=device, dtype=torch.float32) do_boot = (B_self > 0) and (step >= bootstrap_start) if do_boot: d_half = d_self / 2.0 step_idx_half = step_idx_self + 1 sigma_plus = sigma_self + d_half sigma_idx_plus = (sigma_idx_self + (torch.tensor(k_max, device=device, dtype=torch.float32) * d_half).to(torch.long)).clamp(0, k_max) with torch.no_grad(): z1_hat_half1, _ = dynamics(actions[B_emp:] if actions is not None else None, step_idx_half, sigma_idx_self, z_tilde_self, act_mask=act_mask_self, agent_tokens=None, lang_emb=lang_emb_self) b_prime = (z1_hat_half1.float() - z_tilde_self.float()) / (1.0 - sigma_self).clamp_min(1e-6)[..., None, None] z_prime = z_tilde_self.float() + b_prime * d_half[..., None, None] z1_hat_half2, _ = dynamics(actions[B_emp:] if actions is not None else None, step_idx_half, sigma_idx_plus, z_prime.to(z_tilde_self.dtype), act_mask=act_mask_self, agent_tokens=None, lang_emb=lang_emb_self) b_doubleprime = (z1_hat_half2.float() - z_prime.float()) / (1.0 - sigma_plus).clamp_min(1e-6)[..., None, None] vhat_sigma = (z1_hat_self.float() - z_tilde_self.float()) / (1.0 - sigma_self).clamp_min(1e-6)[..., None, None] vbar_target = (b_prime + b_doubleprime) / 2.0 boot_per = (1.0 - sigma_self).pow(2) * (vhat_sigma - vbar_target).pow(2).mean(dim=(2, 3)) # (B_self,T) loss_self = (boot_per * w_self).mean() boot_mse = boot_per.mean() # Combine losses loss_flow = ((loss_emp * (B - B_self)) + (loss_self * B_self)) / B # ----------------------- # Reward modeling loss (optional) # ----------------------- loss_rew = torch.zeros((), device=device, dtype=torch.float32) if (rew_head is not None) and (rewards is not None) and (reward_weight > 0.0): if h_t_full is None: raise RuntimeError("Reward modeling requested but dynamics returned h_t_full=None (set n_agent>0).") # logits: (B,T,L,K), centers_log: (K,) logits_btlk, centers_log = rew_head(h_t_full) K = logits_btlk.shape[-1] L = logits_btlk.shape[-2] # Align targets for MTP: # head l predicts reward at time (t + l). Mask invalid shifts. # Shift rewards to led-to convention: rewards_shift[t] = r_{t-1} = reward # received arriving at state s_t via action a_{t-1}. This matches the # led-to action convention already used by the dynamics, so head l=0 at # time t predicts a reward deterministic in h_t[t]'s causal context # (s_0..s_t, a_0..a_{t-1}). r_{-1} is undefined at t=0. rewards_shift = torch.zeros_like(rewards) rewards_shift[:, 1:] = rewards[:, :-1] rew_symlog = symlog(rewards_shift.float()) # (B,T) valid_bt = torch.ones((B, T), device=device, dtype=torch.bool) valid_bt[:, 0] = False # Build (B,T,L) shifted targets + masks with pad+unfold # pad on the right with L-1 zeros/False so unfold yields exactly L windows. rew_pad = F.pad(rew_symlog, (0, max(0, L - 1))) # (B,T+L-1) msk_pad = F.pad(valid_bt, (0, max(0, L - 1)), value=False) # (B,T+L-1) tgt_blt = rew_pad.unfold(dimension=1, size=T, step=1) # (B,L,T) msk_blt = msk_pad.unfold(dimension=1, size=T, step=1) # (B,L,T) tgt_btl = tgt_blt.permute(0, 2, 1).contiguous() # (B,T,L) msk_btl = msk_blt.permute(0, 2, 1).contiguous() # (B,T,L) loss_rew = dist_cross_entropy_from_symlog( logits=logits_btlk.float().reshape(-1, K), target_symlog=tgt_btl.reshape(-1), centers_log=centers_log, # (K,) mask=msk_btl.reshape(-1), ) # ----------------------- # BC (policy) loss (optional, gradient-isolated from dynamics) # ----------------------- # h_t.detach() severs the graph upstream: BC gradients update only # policy_head's parameters, not dynamics / agent-token init. # # Target alignment (mirror of the reward head but shifted the opposite way): # at slot t, the causal context has seen obs[0..t] + actions[0..t] (led-to # convention). The natural BC target is the NEXT action to take from obs[t], # which in the led-to array is actions[t+1]. Head l at slot t predicts # actions[t+1+l]. Invalid when t+1+l > T-1. act_mask is also shifted to the # "next position" to get the correct per-dim validity for the target action. loss_bc = torch.zeros((), device=device, dtype=torch.float32) if (policy_head is not None) and (bc_weight > 0.0) and (actions is not None) and (act_mask is not None): if h_t_full is None: raise RuntimeError("BC loss requested but dynamics returned h_t_full=None (set n_agent>0).") A = actions.shape[-1] # Shift actions / per-dim mask forward by 1: actions_next[t] = actions[t+1] actions_next = torch.zeros_like(actions) actions_next[:, :-1] = actions[:, 1:] mask_next = torch.zeros_like(act_mask) mask_next[:, :-1] = act_mask[:, 1:] # Time-position validity: actions[t+1] exists iff t < T-1. valid_bt = torch.ones((B, T), device=device, dtype=torch.bool) valid_bt[:, -1] = False L_bc = _unwrap_model(policy_head).L # pad + unfold along the T axis, producing (B, T, L, A) targets and masks. # F.pad order for (B, T, A): (last_dim_left, last_dim_right, T_left, T_right) act_pad = F.pad(actions_next, (0, 0, 0, max(0, L_bc - 1))) # (B, T+L-1, A) act_unf = act_pad.unfold(dimension=1, size=T, step=1) # (B, L, A, T) tgt_btla = act_unf.permute(0, 3, 1, 2).contiguous() # (B, T, L, A) mask_pad = F.pad(mask_next, (0, 0, 0, max(0, L_bc - 1))) # (B, T+L-1, A) mask_unf = mask_pad.unfold(dimension=1, size=T, step=1) # (B, L, A, T) dim_mask_btla = mask_unf.permute(0, 3, 1, 2).contiguous() # (B, T, L, A) valid_pad = F.pad(valid_bt, (0, max(0, L_bc - 1)), value=False) # (B, T+L-1) valid_unf = valid_pad.unfold(dimension=1, size=T, step=1) # (B, L, T) pos_mask_btl = valid_unf.permute(0, 2, 1).contiguous() # (B, T, L) # Combined per-element mask: valid timestep AND valid dim for this task. full_mask = pos_mask_btl.unsqueeze(-1).to(dim_mask_btla.dtype) * dim_mask_btla # (B, T, L, A) # Policy forward with gradient isolation. pred_btla = policy_head(h_t_full.detach()) # (B, T, L, A) diff_sq = (pred_btla.float() - tgt_btla.float()).pow(2) masked_sq = diff_sq * full_mask.float() denom = full_mask.float().sum().clamp_min(1.0) loss_bc = masked_sq.sum() / denom if rms_flow is not None: rms_flow.update(loss_flow) loss_flow_normed = rms_flow.normalize(loss_flow) else: loss_flow_normed = loss_flow if rms_rew is not None and float(loss_rew.item()) > 0.0: rms_rew.update(loss_rew) loss_rew_normed = rms_rew.normalize(loss_rew) else: loss_rew_normed = loss_rew if rms_bc is not None and float(loss_bc.item()) > 0.0: rms_bc.update(loss_bc) loss_bc_normed = rms_bc.normalize(loss_bc) else: loss_bc_normed = loss_bc loss = ( loss_flow_normed + float(reward_weight) * loss_rew_normed + float(bc_weight) * loss_bc_normed ) aux = { "flow_mse": flow_per.mean().detach(), "flow_per_sample": flow_per.mean(dim=1).detach(), # (B_emp,) — for per-domain bucketing (unweighted) "bootstrap_mse": boot_mse.detach(), "loss_emp": loss_emp.detach(), "loss_self": loss_self.detach(), "sigma_mean": sigma_full.mean().detach(), "sigma_std": sigma_full.std().detach(), "loss_rew": loss_rew.detach(), "loss_bc": loss_bc.detach(), "loss_flow": loss_flow.detach(), "rms_flow": rms_flow.rms_val if rms_flow is not None else 1.0, "rms_rew": rms_rew.rms_val if rms_rew is not None else 1.0, "rms_bc": rms_bc.rms_val if rms_bc is not None else 1.0, } return loss, aux def _is_pow2(n: int) -> bool: return (n > 0) and ((n & (n - 1)) == 0) def make_tau_schedule(*, k_max: int, schedule: str, d: Optional[float] = None) -> Dict[str, Any]: """ Returns a schedule dict: K = number of integration steps (also grid size) e = log2(K) (step_idx) scale = k_max // K tau_idx[i] = discrete signal index at step i tau[i] = i/K dt = 1/K """ assert _is_pow2(k_max), "k_max must be power of two" if schedule == "finest": K = k_max elif schedule == "shortcut": assert d is not None, "shortcut schedule requires --eval_d" inv = int(round(1.0 / float(d))) assert _is_pow2(inv), "eval_d must be 1/(power of two)" assert inv <= k_max, "eval_d must be >= 1/k_max" assert (k_max % inv) == 0, "k_max must be divisible by 1/eval_d" K = inv else: raise ValueError(f"unknown schedule: {schedule}") e = int(round(math.log2(K))) scale = k_max // K tau = [i / K for i in range(K)] + [1.0] tau_idx = [i * scale for i in range(K)] + [k_max] # allow final clean index return dict(K=K, e=e, scale=scale, tau=tau, tau_idx=tau_idx, dt=1.0 / K, schedule=schedule, d=1.0 / K) @torch.no_grad() def sample_one_timestep_packed( dyn: Dynamics, *, past_packed: torch.Tensor, # (B,t,n_spatial,d_spatial) k_max: int, sched: Dict[str, Any], actions: Optional[torch.Tensor] = None, # (B,T,A) aligned to frames or None act_mask: Optional[torch.Tensor] = None, # (B,T,A) or (A,) or None tau_ctx: float = 0.0, # context corruption level lang_emb: Optional[torch.Tensor] = None, # (B,lang_dim) or None use_amp: bool = True, # match training bf16 autocast use_kv_cache: bool = False, # cache time-attn K,V for context tokens ) -> torch.Tensor: """Generate next packed latent z_t given past length t. When `use_kv_cache=True` and `t > 0`, the time-attention K,V for the t context tokens is computed once in a prefill pass and reused across all K Euler denoising steps. Each denoising step then only runs the single new token through the transformer, attending to the cached context. This cuts per-step attention from O(t+1) to O(1) and gives roughly a 3x speedup at typical ctx_window=24 with K=4 denoising steps. Requires `Dynamics` to support the `return_kv_cache` / `kv_cache` kwargs. """ device = past_packed.device dtype = past_packed.dtype B, t = past_packed.shape[:2] n_spatial, d_spatial = past_packed.shape[2], past_packed.shape[3] K = int(sched["K"]) e = int(sched["e"]) tau = sched["tau"] tau_idx = sched["tau_idx"] dt = float(sched["dt"]) # start from noise at tau=0 z = torch.randn((B, 1, n_spatial, d_spatial), device=device, dtype=dtype) emax = int(round(math.log2(k_max))) # Slightly corrupt past context tokens for robustness to autoregressive errors. # tau_ctx is the noise fraction: 0 = fully clean, 1 = fully noisy. # Signal index convention: 0 = fully noisy (sigma=0), k_max = fully clean (sigma=1), # so ctx_sig_idx = round((1 - tau_ctx) * k_max) to match the training convention. if tau_ctx > 0.0 and t > 0: z0_ctx = torch.randn_like(past_packed) past_input = ((1.0 - tau_ctx) * past_packed.float() + tau_ctx * z0_ctx.float()).to(dtype) ctx_sig_idx = min(int(round((1.0 - tau_ctx) * k_max)), k_max) else: past_input = past_packed ctx_sig_idx = k_max # broadcast (A,) -> (B,T,A) if needed (only if actions are present) if act_mask is not None and act_mask.dim() == 1: act_mask = act_mask.view(1, 1, -1) actions_in = None if actions is None else actions[:, : t + 1] actmask_in = None if act_mask is None else act_mask[:, : t + 1] # --- KV cache: prefill context tokens once --- kv_cache = None if use_kv_cache and t > 0: ctx_step_idxs = torch.full((B, t), emax, device=device, dtype=torch.long) ctx_signal_idxs = torch.full((B, t), ctx_sig_idx, device=device, dtype=torch.long) ctx_actions = None if actions_in is None else actions_in[:, :t] ctx_actmask = None if actmask_in is None else actmask_in[:, :t] with torch.autocast(device_type=device.type, enabled=(use_amp and device.type == "cuda"), dtype=torch.bfloat16): _, _, kv_cache = dyn( ctx_actions, ctx_step_idxs, ctx_signal_idxs, past_input, act_mask=ctx_actmask, agent_tokens=None, lang_emb=lang_emb, return_kv_cache=True, ) # Full-sequence path uses these (allocated once). step_idxs_full = torch.full((B, t + 1), emax, device=device, dtype=torch.long) step_idxs_full[:, -1] = e signal_idxs_full = torch.full((B, t + 1), ctx_sig_idx, device=device, dtype=torch.long) for i in range(K): tau_i = float(tau[i]) sig_i = int(tau_idx[i]) with torch.autocast(device_type=device.type, enabled=(use_amp and device.type == "cuda"), dtype=torch.bfloat16): if kv_cache is not None: # Decode: only the single new token (position t) attends to cached ctx. new_step_idxs = torch.full((B, 1), e, device=device, dtype=torch.long) new_signal_idxs = torch.full((B, 1), sig_i, device=device, dtype=torch.long) new_actions = None if actions_in is None else actions_in[:, -1:] new_actmask = None if actmask_in is None else actmask_in[:, -1:] x1_hat, _ = dyn( new_actions, new_step_idxs, new_signal_idxs, z, act_mask=new_actmask, agent_tokens=None, lang_emb=lang_emb, kv_cache=kv_cache, ) else: # Full-sequence (no-cache / t==0) path. signal_idxs_full[:, -1] = sig_i packed_seq = torch.cat([past_input, z], dim=1) # (B,t+1,...) x1_hat_full, _ = dyn( actions_in, step_idxs_full, signal_idxs_full, packed_seq, act_mask=actmask_in, agent_tokens=None, lang_emb=lang_emb, ) x1_hat = x1_hat_full[:, -1:, :, :] # (B,1,n_spatial,d_spatial) denom = max(1e-4, 1.0 - tau_i) b = (x1_hat.float() - z.float()) / denom z = (z.float() + b * dt).to(dtype) return z[:, 0] # (B,n_spatial,d_spatial) @torch.no_grad() def sample_autoregressive_packed_sequence( dyn: Dynamics, *, z_gt_packed: torch.Tensor, # (B,T,n_spatial,d_spatial) ctx_length: int, horizon: int, k_max: int, sched: Dict[str, Any], actions: Optional[torch.Tensor] = None, # (B,T,A) or None act_mask: Optional[torch.Tensor] = None, # (B,T,A) or (A,) or None tau_ctx: float = 0.0, # context corruption level lang_emb: Optional[torch.Tensor] = None, # (B,lang_dim) or None ) -> torch.Tensor: B, T = z_gt_packed.shape[:2] L = min(T, ctx_length + horizon) ctx_length = min(ctx_length, L - 1) horizon = min(horizon, L - ctx_length) outs = [z_gt_packed[:, t] for t in range(ctx_length)] for t in range(ctx_length, ctx_length + horizon): past = torch.stack(outs, dim=1) # (B,t,...) z_next = sample_one_timestep_packed( dyn, past_packed=past, k_max=k_max, sched=sched, actions=actions, act_mask=act_mask, tau_ctx=tau_ctx, lang_emb=lang_emb, ) outs.append(z_next) return torch.stack(outs, dim=1) @torch.no_grad() def decode_packed_to_frames( decoder: Decoder, *, z_packed: torch.Tensor, # (B,T',n_spatial,d_spatial) H: int, W: int, C: int, patch: int, packing_factor: int, ) -> torch.Tensor: z_btLd = unpack_spatial_to_bottleneck(z_packed, k=packing_factor) # (B,T',L,D_b) patches_btnd = decoder(z_btLd) # (B,T',Np,Dp) in [0,1] frames = temporal_unpatchify(patches_btnd, H, W, C, patch) # (B,T',C,H,W) in [0,1] return frames.clamp(0, 1) @torch.no_grad() def log_dynamics_eval_wandb( *, gt: torch.Tensor, # (B,T,C,H,W) float [0,1] pred: torch.Tensor, # (B,T,C,H,W) float [0,1] ctx_length: int, step: int, tag: str, # wandb key prefix, e.g. "eval" or "val" max_items: int = 4, gap_px: int = 16, ): B, T, C, H, W = gt.shape Bv = min(B, max_items) def tile_time(x: torch.Tensor) -> torch.Tensor: x = x[:Bv] B_, T_, C_, H_, W_ = x.shape ctx = int(max(0, min(ctx_length, T_))) y = x.permute(0, 2, 3, 1, 4).contiguous().view(B_, C_, H_, T_ * W_) if gap_px > 0 and 0 < ctx < T_: split = ctx * W_ left = y[..., :split] right = y[..., split:] gap = torch.zeros((B_, C_, H_, gap_px), device=y.device, dtype=y.dtype) y = torch.cat([left, gap, right], dim=-1) return y gt_t = tile_time(gt) pr_t = tile_time(pred) # Stack rows: GT / Pred panel = torch.cat([gt_t, pr_t], dim=2) # (Bv,C,2H,TW+gap) big = torch.cat([panel[i] for i in range(Bv)], dim=1) # (C,Bv*2H,TW+gap) big = (big.clamp(0, 1) * 255.0).to(torch.uint8) big_hwc = big.permute(1, 2, 0).cpu().numpy() wandb.log( {f"{tag}/viz": wandb.Image(big_hwc, caption=f"rows=GT/Pred | ctx={ctx_length} | T={T}")}, step=step, ) @torch.no_grad() def run_dynamics_eval( *, encoder: Encoder, decoder: Decoder, dyn: Dynamics, frames: torch.Tensor, # (B,T,C,H,W) float [0,1] actions: Optional[torch.Tensor], # (B,T,A) or None act_mask: Optional[torch.Tensor], # (A,) or None H: int, W: int, C: int, patch: int, packing_factor: int, k_max: int, ctx_length: int, horizon: int, sched: Dict[str, Any], max_items: int, step: int, tau_ctx: float = 0.0, lang_emb: Optional[torch.Tensor] = None, # (B,lang_dim) or None tag: str = "eval", # wandb key prefix; use "val" for held-out set ): dyn_was_training = dyn.training dyn.eval() B, T = frames.shape[:2] T_eval = min(T, ctx_length + horizon) ctx_length = min(ctx_length, T_eval - 1) horizon = min(horizon, T_eval - ctx_length) frames_eval = frames[:, :T_eval] patches = temporal_patchify(frames_eval, patch) z_btLd, _ = encoder(patches) # (B,T_eval,L,D_b) assert z_btLd.shape[2] % packing_factor == 0 n_spatial = z_btLd.shape[2] // packing_factor z_gt_packed = pack_bottleneck_to_spatial(z_btLd, n_spatial=n_spatial, k=packing_factor) # (B,T_eval,Sz,Dz) actions_eval = None if actions is None else actions[:, :T_eval] act_mask_eval = None if act_mask is None else act_mask[:, :T_eval] if act_mask.dim() == 3 else act_mask z_pred_packed = sample_autoregressive_packed_sequence( dyn, z_gt_packed=z_gt_packed, ctx_length=ctx_length, horizon=horizon, k_max=k_max, sched=sched, actions=actions_eval, act_mask=act_mask_eval, tau_ctx=tau_ctx, lang_emb=lang_emb, ) pred_frames = decode_packed_to_frames( decoder, z_packed=z_pred_packed, H=H, W=W, C=C, patch=patch, packing_factor=packing_factor, ) # floor baseline: repeat last context frame over horizon floor = frames_eval.clone() if horizon > 0: floor[:, ctx_length:ctx_length + horizon] = frames_eval[:, ctx_length - 1:ctx_length].expand(-1, horizon, -1, -1, -1) # metric on horizon only gt_h = frames_eval[:, ctx_length:ctx_length + horizon] # (B,Hz,C,H,W) pred_h = pred_frames[:, ctx_length:ctx_length + horizon] floor_h = floor[:, ctx_length:ctx_length + horizon] mse_pred = (pred_h.float() - gt_h.float()).pow(2).mean() mse_floor = (floor_h.float() - gt_h.float()).pow(2).mean() psnr_pred = 10.0 * torch.log10(1.0 / mse_pred.clamp_min(1e-12)) psnr_floor = 10.0 * torch.log10(1.0 / mse_floor.clamp_min(1e-12)) mse_ratio = mse_pred / mse_floor.clamp_min(1e-12) # <1 is better psnr_gain = psnr_pred - psnr_floor # >0 is better # per-timestep horizon MSE: log first/mid/last # per_t: (Hz,) per_t_pred = (pred_h.float() - gt_h.float()).pow(2).mean(dim=(0,2,3,4)) per_t_floor = (floor_h.float() - gt_h.float()).pow(2).mean(dim=(0,2,3,4)) if horizon > 0: i0 = 0 im = (horizon - 1) // 2 i1 = horizon - 1 wandb.log( { f"{tag}/mse_pred": float(mse_pred.item()), f"{tag}/mse_floor": float(mse_floor.item()), f"{tag}/mse_ratio_pred_over_floor": float(mse_ratio.item()), f"{tag}/psnr_pred": float(psnr_pred.item()), f"{tag}/psnr_floor": float(psnr_floor.item()), f"{tag}/psnr_gain_over_floor_db": float(psnr_gain.item()), # 1-indexed step labels in the horizon f"{tag}/mse_pred_t1": float(per_t_pred[i0].item()), f"{tag}/mse_pred_tmid": float(per_t_pred[im].item()), f"{tag}/mse_pred_tend": float(per_t_pred[i1].item()), f"{tag}/mse_floor_t1": float(per_t_floor[i0].item()), f"{tag}/mse_floor_tmid": float(per_t_floor[im].item()), f"{tag}/mse_floor_tend": float(per_t_floor[i1].item()), }, step=step, ) log_dynamics_eval_wandb( gt=frames_eval, pred=pred_frames, ctx_length=ctx_length, step=step, tag=tag, max_items=max_items, ) if dyn_was_training: dyn.train() @torch.no_grad() def run_action_attribution_eval( *, encoder: Encoder, dyn: Dynamics, batches, # iterable of parsed (frames, actions, act_mask, lang_emb) patch: int, packing_factor: int, k_max: int, tau_ctx: float, step: int, tag: str, rng: torch.Generator, sched: Optional[Dict[str, Any]] = None, ): """Averages action-attribution ratios over many val batches. For each sequence, at each position t in [1, T-1], a teacher-forced one-step prediction is formed by running the shortcut sampler from pure noise with the ground-truth (packed) latents at positions 0..t-1 as context. The prediction is compared against z_gt[t] in latent space. Metric keys: {tag}/action_shuffle_ratio : pool-level batch-permuted actions {tag}/action_timeshuffle_ratio : per-sequence time-permuted actions {tag}/action_zero_ratio : zero actions `pool-level` means the batch-shuffle permutation is drawn over the concatenation of all N val batches (not within each batch), which makes the intervention far more severe for multi-task mini-batches with only a few sequences per batch. The per-(seq,t) initial noise and tau_ctx corruption are drawn once from `rng` and shared across the four intervention calls, so the three ratios differ ONLY in the action tensor. That drops a lot of sampling variance from the ratio. """ if sched is None: sched = make_tau_schedule(k_max=k_max, schedule="shortcut", d=1.0) dyn_was_training = dyn.training enc_was_training = encoder.training dyn.eval() encoder.eval() # Encode all batches; stash actions/masks/lang for the per-intervention calls. z_list, act_list, mask_list, lang_list = [], [], [], [] for (frames, actions, act_mask, lang_emb) in batches: if actions is None: continue patches = temporal_patchify(frames, patch) z_btLd, _ = encoder(patches) assert z_btLd.shape[2] % packing_factor == 0 n_spatial_local = z_btLd.shape[2] // packing_factor z_p = pack_bottleneck_to_spatial(z_btLd, n_spatial=n_spatial_local, k=packing_factor) z_list.append(z_p) act_list.append(actions) mask_list.append(act_mask) lang_list.append(lang_emb) if len(z_list) == 0: if dyn_was_training: dyn.train() if enc_was_training: encoder.train() return # Align B and T across batches (drop partial last batch if needed). B_common = min(z.shape[0] for z in z_list) T_common = min(z.shape[1] for z in z_list) z_list = [z[:B_common, :T_common] for z in z_list] act_list = [a[:B_common, :T_common] for a in act_list] mask_list = [ (m[:B_common, :T_common] if (m is not None and m.dim() >= 2) else m) for m in mask_list ] device = z_list[0].device dtype = z_list[0].dtype # Pool concatenation along batch dim: (N*B, T, S, D) z_pool = torch.cat(z_list, dim=0) a_pool = torch.cat(act_list, dim=0) if mask_list[0] is not None and mask_list[0].dim() >= 2: m_pool = torch.cat(mask_list, dim=0) else: m_pool = mask_list[0] if all(l is None for l in lang_list): l_pool = None else: ref = next(l for l in lang_list if l is not None) filled = [l if l is not None else torch.zeros_like(ref[:B_common]) for l in lang_list] l_pool = torch.cat([l[:B_common] for l in filled], dim=0) BP, TP = z_pool.shape[:2] if TP < 2: if dyn_was_training: dyn.train() if enc_was_training: encoder.train() return n_spatial = z_pool.shape[2] d_spatial = z_pool.shape[3] K = int(sched["K"]) tau_s = sched["tau"] tau_i_s = sched["tau_idx"] dt = float(sched["dt"]) e_step = int(sched["e"]) emax = int(round(math.log2(k_max))) # Shared per-(seq,t) initial noise; same tensor reused for all interventions. noise_pool = torch.randn( (BP, TP - 1, n_spatial, d_spatial), generator=rng, device=rng.device, dtype=torch.float32, ).to(device=device, dtype=dtype) # tau_ctx-corrupted past context, matching the inference-time rollout. if tau_ctx > 0.0: ctx_noise = torch.randn( z_pool.shape, generator=rng, device=rng.device, dtype=torch.float32, ).to(device=device, dtype=dtype) z_past_input = ((1.0 - tau_ctx) * z_pool.float() + tau_ctx * ctx_noise.float()).to(dtype) ctx_sig_idx = min(int(round((1.0 - tau_ctx) * k_max)), k_max) else: z_past_input = z_pool ctx_sig_idx = k_max def _pred_mse(actions_in, mask_in): """Teacher-forced one-step MSE per (seq, t). Returns (BP, TP-1).""" per_t_mses = [] for t in range(1, TP): past = z_past_input[:, :t] target_t = z_pool[:, t].float() z = noise_pool[:, t - 1:t].clone() # (BP, 1, S, D) step_idxs_full = torch.full((BP, t + 1), emax, device=device, dtype=torch.long) step_idxs_full[:, -1] = e_step signal_idxs_full = torch.full((BP, t + 1), ctx_sig_idx, device=device, dtype=torch.long) act_in = actions_in[:, :t + 1] if actions_in is not None else None if mask_in is not None and mask_in.dim() >= 2: msk_in = mask_in[:, :t + 1] else: msk_in = mask_in for i in range(K): tau_i = float(tau_s[i]) sig_i = int(tau_i_s[i]) signal_idxs_full[:, -1] = sig_i packed_seq = torch.cat([past, z], dim=1) x1_hat_full, _ = dyn( act_in, step_idxs_full, signal_idxs_full, packed_seq, act_mask=msk_in, agent_tokens=None, lang_emb=l_pool, ) x1_hat = x1_hat_full[:, -1:, :, :] denom = max(1e-4, 1.0 - tau_i) b = (x1_hat.float() - z.float()) / denom z = (z.float() + b * dt).to(dtype) z_pred_t = z[:, 0].float() diff_sq_t = (z_pred_t - target_t).pow(2).mean(dim=(1, 2)) # (BP,) per_t_mses.append(diff_sq_t) return torch.stack(per_t_mses, dim=1) # (BP, TP-1) # Pool-level batch-shuffle permutation (over N*B sequences) perm_b = torch.randperm(BP, generator=rng, device=rng.device).to(device) # Per-sequence time-shuffle (independent permutation per seq, vectorized) perm_t_per_seq = torch.stack( [torch.randperm(TP, generator=rng, device=rng.device) for _ in range(BP)], dim=0 ).to(device) a_tshuf = torch.gather( a_pool, 1, perm_t_per_seq.unsqueeze(-1).expand(-1, -1, a_pool.shape[-1]), ) if m_pool is not None and m_pool.dim() >= 2: m_tshuf = torch.gather( m_pool, 1, perm_t_per_seq.unsqueeze(-1).expand(-1, -1, m_pool.shape[-1]), ) else: m_tshuf = m_pool mse_real_bt = _pred_mse(a_pool, m_pool) mse_bshuf_bt = _pred_mse( a_pool[perm_b], m_pool[perm_b] if (m_pool is not None and m_pool.dim() >= 2) else m_pool, ) mse_tshuf_bt = _pred_mse(a_tshuf, m_tshuf) mse_zero_bt = _pred_mse(torch.zeros_like(a_pool), m_pool) mse_real = float(mse_real_bt.mean().item()) mse_bshuf = float(mse_bshuf_bt.mean().item()) mse_tshuf = float(mse_tshuf_bt.mean().item()) mse_zero = float(mse_zero_bt.mean().item()) ratio_bshuf = mse_bshuf / max(mse_real, 1e-12) ratio_tshuf = mse_tshuf / max(mse_real, 1e-12) ratio_zero = mse_zero / max(mse_real, 1e-12) # Per-sequence ratio distribution (mean over t per seq). Headline ratios # above are pool-level (mean-of-MSEs); the per-seq stats below reveal # heterogeneity across sequences. real_seq_safe = mse_real_bt.mean(dim=1).clamp_min(1e-12) seq_ratios_b = mse_bshuf_bt.mean(dim=1) / real_seq_safe seq_ratios_t = mse_tshuf_bt.mean(dim=1) / real_seq_safe seq_ratios_z = mse_zero_bt.mean(dim=1) / real_seq_safe def _ratio_stats(ratios, key): out = {} if BP <= 1: out[f"{key}_std"] = 0.0 return out out[f"{key}_std"] = float(ratios.std(unbiased=True).item()) qs = torch.tensor([0.1, 0.5, 0.9], device=ratios.device, dtype=ratios.dtype) q = torch.quantile(ratios, qs) out[f"{key}_p10"] = float(q[0].item()) out[f"{key}_p50"] = float(q[1].item()) out[f"{key}_p90"] = float(q[2].item()) return out log_payload = { f"{tag}/action_shuffle_ratio": ratio_bshuf, f"{tag}/action_timeshuffle_ratio": ratio_tshuf, f"{tag}/action_zero_ratio": ratio_zero, f"{tag}/action_mse_real": mse_real, f"{tag}/action_attr_pool_size": BP, # One-step teacher-forced prediction MSE on val (latent space). f"{tag}/pred_mse": mse_real, } log_payload.update(_ratio_stats(seq_ratios_b, f"{tag}/action_shuffle_ratio")) log_payload.update(_ratio_stats(seq_ratios_t, f"{tag}/action_timeshuffle_ratio")) log_payload.update(_ratio_stats(seq_ratios_z, f"{tag}/action_zero_ratio")) wandb.log(log_payload, step=step) if dyn_was_training: dyn.train() if enc_was_training: encoder.train() def parse_batch(batch, *, use_rewards: bool, device): """Mirror of the inline batch parsing in the train loop. Returns (frames, actions, act_mask, rewards, lang_emb), each on `device`. Kept as a helper so the validation block can reuse the exact same convention without duplicating logic.""" obs_u8 = batch["obs"].to(device, non_blocking=True) # (B,T+1,3,H,W) uint8 act = batch["act"].to(device, non_blocking=True) # (B,T,A) float mask = batch["act_mask"].to(device, non_blocking=True) # (B,T,A) float rew_in = batch.get("rew", None) rew_in = None if rew_in is None else rew_in.to(device, non_blocking=True) lang_emb = batch["lang_emb"].to(device, non_blocking=True) # (B,lang_dim) act = act.clamp(-1, 1) * mask frames = obs_u8[:, :-1].float() / 255.0 actions = torch.zeros_like(act) actions[:, 1:] = act[:, :-1] act_mask = torch.zeros_like(mask) act_mask[:, 1:] = mask[:, :-1] rewards = rew_in if use_rewards else None return frames, actions, act_mask, rewards, lang_emb 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) # Dataset and DataLoader from wm_dataset import WMDataset, collate_batch # Targeted-collection fine-tunes additionally load UNSEEN tasks; the dataset # filters to what actually has data on disk so runs without UNSEEN demos are # unaffected when none are present. full_task_list = list(TASK_SET) + list(UNSEEN_TASK_SET) tw_list = compute_task_weights( full_task_list, args.task_weighting, targeted_alpha=args.targeted_alpha, ) task_weights = None if args.task_weighting == "valid_starts" else dict(zip(full_task_list, tw_list)) 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=args.lang_dim, tasks_json=args.tasks_json, tasks=full_task_list, verbose=is_rank0(), cache_mb=args.cache_mb, ddp_partition=True, iid_sampling=True, samples_per_shard=args.samples_per_shard, task_weights=task_weights, ) # Build per-rank lookup: local task_idx (emb_id) -> domain_idx. # Used to bucket per-sample losses by domain for logging. task_idx_to_domain_idx = torch.tensor( [DOMAINS.index(task_to_domain(t)) for t in dataset.tasks], dtype=torch.long, device=device, ) # With ddp_partition + iid_sampling, tasks are already split across # ranks and __getitem__ ignores the index, so DistributedSampler is # unnecessary. Plain shuffle keeps the DataLoader happy. 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), prefetch_factor=args.prefetch_factor if args.num_workers > 0 else None, worker_init_fn=worker_init_fn, collate_fn=collate_batch, ) # ---- validation data (rank 0 only) ---- # Mirrors the train dataset construction but reads from a single held-out # directory. We only build it on rank 0 because the existing eval block # also runs on rank 0 only — other ranks idle at the next collective. val_loader = None val_iter = None if is_rank0() and args.val_every > 0 and args.val_frame_dir is not None: assert args.val_data_dir is not None, "--val_data_dir is required for validation" val_dataset = WMDataset( data_dir=args.val_data_dir, frames_dir=args.val_frame_dir, seq_len=args.seq_len, img_size=224, action_dim=16, lang_dim=args.lang_dim, tasks_json=args.tasks_json, tasks=TASK_SET, verbose=False, ddp_partition=False, iid_sampling=True, samples_per_shard=1, task_weights=task_weights, ) val_loader = DataLoader( val_dataset, batch_size=args.eval_batch_size, shuffle=True, num_workers=2, pin_memory=True, drop_last=True, persistent_workers=True, worker_init_fn=worker_init_fn, collate_fn=collate_batch, ) def _val_batch_iterator(loader): while True: for b in loader: yield b val_iter = _val_batch_iterator(val_loader) # Load frozen tokenizer tok_override = {} if args.H is not None: tok_override["H"] = args.H if args.W is not None: tok_override["W"] = args.W if args.C is not None: tok_override["C"] = args.C if args.patch is not None: tok_override["patch"] = args.patch encoder, decoder, tok_args = load_frozen_tokenizer_from_pt_ckpt( args.tokenizer_ckpt, device=device, override=tok_override ) H = int(tok_args.get("H", 224)) W = int(tok_args.get("W", 224)) C = int(tok_args.get("C", 3)) patch = int(tok_args.get("patch", 4)) n_latents = int(tok_args.get("n_latents", 16)) d_bottleneck = int(tok_args.get("d_bottleneck", 32)) assert H % patch == 0 and W % patch == 0 assert n_latents % args.packing_factor == 0 n_spatial = n_latents // args.packing_factor d_spatial = d_bottleneck * args.packing_factor # Build dynamics model dyn = Dynamics( d_model=args.d_model_dyn, d_bottleneck=d_bottleneck, d_spatial=d_spatial, n_spatial=n_spatial, n_register=args.n_register, n_agent=args.n_agent, n_heads=args.n_heads, depth=args.dyn_depth, k_max=args.k_max, dropout=args.dropout, mlp_ratio=args.mlp_ratio, time_every=args.time_every, lang_dim=args.lang_dim, ).to(device) # ----------------------- # Reward head (optional) # ----------------------- use_rewards = args.reward_weight > 0.0 if use_rewards: if args.n_agent <= 0: raise RuntimeError("Reward modeling requires --n_agent > 0 (agent token features h_t).") rew_head = RewardHeadMTP( d_model=args.d_model_dyn, L=args.reward_L, num_bins=args.reward_num_bins, dropout=args.dropout, mlp_ratio=args.reward_mlp_ratio, log_low=args.reward_log_low, log_high=args.reward_log_high, pool_agent=args.reward_pool_agent, ).to(device) else: rew_head = None # ----------------------- # BC (policy) head (optional, gradient-isolated) # ----------------------- # BC gradients are isolated from the dynamics via h_t.detach() at the call # site. use_bc = args.bc_weight > 0.0 if use_bc: if args.n_agent <= 0: raise RuntimeError("BC modeling requires --n_agent > 0 (agent token features h_t).") policy_head = PolicyHeadMTP( d_model=args.d_model_dyn, L=args.bc_L, act_dim_max=args.bc_act_dim, dropout=args.dropout, mlp_ratio=args.bc_mlp_ratio, pool_agent=args.bc_pool_agent, ).to(device) else: policy_head = None rms_flow = EmaRms().to(device) rms_rew = EmaRms().to(device) rms_bc = EmaRms().to(device) if is_rank0(): print(dyn) dyn_param_count = sum(p.numel() for p in dyn.parameters() if p.requires_grad) rew_param_count = sum(p.numel() for p in rew_head.parameters() if p.requires_grad) if rew_head is not None else 0 bc_param_count = sum(p.numel() for p in policy_head.parameters() if p.requires_grad) if policy_head is not None else 0 param_count = dyn_param_count + rew_param_count + bc_param_count if rew_head is not None: print(f"Learnable parameters (dynamics): {dyn_param_count:,}") print(f"Learnable parameters (reward head): {rew_param_count:,}") if policy_head is not None: print(f"Learnable parameters (policy/BC head, gradient-isolated): {bc_param_count:,}") print(f"Total learnable parameters: {param_count:,}") print(f"[tokenizer] H={H} W={W} C={C} patch={patch} n_lat={n_latents} d_b={d_bottleneck} packing={args.packing_factor}") if ddp: dyn = torch.nn.parallel.DistributedDataParallel( dyn, device_ids=[local_rank], output_device=local_rank, broadcast_buffers=False ) if rew_head is not None: rew_head = torch.nn.parallel.DistributedDataParallel( rew_head, device_ids=[local_rank], output_device=local_rank, broadcast_buffers=False ) if policy_head is not None: policy_head = torch.nn.parallel.DistributedDataParallel( policy_head, device_ids=[local_rank], output_device=local_rank, broadcast_buffers=False ) if args.compile: dyn = torch.compile(dyn) if rew_head is not None: rew_head = torch.compile(rew_head) if policy_head is not None: policy_head = torch.compile(policy_head) # Optimizer. BC gradients stay isolated to policy_head params via h_t.detach() # at the call site (see dynamics_pretrain_loss); having them in the same # optimizer is fine and expected. params = list(dyn.parameters()) if rew_head is not None: params += list(rew_head.parameters()) if policy_head is not None: params += list(policy_head.parameters()) opt = torch.optim.AdamW(params, lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.999)) use_amp = torch.cuda.is_available() # Initialize wandb if is_rank0(): _dyn_for_count = dyn.module if hasattr(dyn, "module") else dyn dyn_param_count_total = sum(p.numel() for p in _dyn_for_count.parameters()) dyn_param_count_trainable = sum(p.numel() for p in _dyn_for_count.parameters() if p.requires_grad) if rew_head is not None: _rew_for_count = rew_head.module if hasattr(rew_head, "module") else rew_head rew_param_count_trainable = sum(p.numel() for p in _rew_for_count.parameters() if p.requires_grad) else: rew_param_count_trainable = 0 frames_per_opt_step = args.batch_size * world_size * max(1, args.grad_accum) * args.seq_len 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, "run/effective_batch_seqs": args.batch_size * world_size * max(1, args.grad_accum), "run/frames_per_opt_step": frames_per_opt_step, "model/dyn_param_count_total": dyn_param_count_total, "model/dyn_param_count_trainable": dyn_param_count_trainable, "model/rew_param_count_trainable": rew_param_count_trainable, "model/total_param_count_trainable": dyn_param_count_trainable + rew_param_count_trainable, "model/n_spatial": n_spatial, "model/d_spatial": d_spatial, "model/n_latents": n_latents, "model/d_bottleneck": d_bottleneck, "tokenizer/H": H, "tokenizer/W": W, "tokenizer/C": C, "tokenizer/patch": patch, }, ) # Resume from checkpoint step = 0 start_epoch = 0 ckpt_dir = Path(args.ckpt_dir) if args.resume is not None: step, start_epoch = load_ckpt(Path(args.resume), dyn_model=dyn, rew_head=rew_head, policy_head=policy_head, opt=opt, rms_objects={"flow": rms_flow, "rew": rms_rew, "bc": rms_bc}, strict=args.resume_strict) if is_rank0(): print(f"[rank0] Resumed from {args.resume} (step={step}, epoch={start_epoch})") # Count LR warmup from the resume point, so fine-tuning resumes (which start # at a high step) still warm up. Fresh starts have resume_step=0 (warmup from # step 0). resume_step = step # Training loop dyn.train() if rew_head is not None: rew_head.train() t0 = time.time() grad_accum = max(1, int(args.grad_accum)) # Per-domain flow-mse accumulator (flushed at each log_every cycle). domain_acc = PerDomainAccumulator(n_domains=len(DOMAINS), device=device) while step <= args.max_steps: for epoch in range(start_epoch, 10_000_000): for batch in loader: if step > args.max_steps: break frames, actions, act_mask, rewards, lang_emb = parse_batch( batch, use_rewards=use_rewards, device=device) # Frozen encoder -> packed spatial tokens z1 with torch.no_grad(): patches = temporal_patchify(frames, patch) # (B,T,Np,Dp) z_btLd, _ = encoder(patches) # (B,T,n_latents,d_b) z1 = pack_bottleneck_to_spatial(z_btLd, n_spatial=n_spatial, k=args.packing_factor) # (B,T,Sz,Dz) B = z1.shape[0] B_self = int(round(args.self_fraction * B)) B_self = max(0, min(B - 1, B_self)) do_step = ((step + 1) % grad_accum == 0) # emb_id is already computed by the dataset; materialize it once # up-front so it can feed the per-domain accumulator below. emb_id_batch = batch["emb_id"].to(device, non_blocking=True).long() # (B,) sync_dyn = dyn.no_sync() if ddp and not do_step else nullcontext() sync_rew = rew_head.no_sync() if (rew_head is not None and ddp and not do_step) else nullcontext() sync_bc = policy_head.no_sync() if (policy_head is not None and ddp and not do_step) else nullcontext() with sync_dyn, sync_rew, sync_bc: with autocast(device_type="cuda", dtype=torch.bfloat16, enabled=use_amp): loss, aux = dynamics_pretrain_loss( dyn, z1=z1, actions=actions, act_mask=act_mask, k_max=args.k_max, B_self=B_self, step=step, bootstrap_start=args.bootstrap_start, lang_emb=lang_emb, rew_head=rew_head, rewards=rewards, reward_weight=args.reward_weight, policy_head=policy_head, bc_weight=args.bc_weight, rms_flow=rms_flow, rms_rew=rms_rew, rms_bc=rms_bc, ) if ddp: rms_flow.sync(world_size) rms_rew.sync(world_size) rms_bc.sync(world_size) # Accumulate per-domain flow_mse for the empirical portion # of the batch (first B_emp samples, aligned with flow_per_sample). flow_ps = aux["flow_per_sample"] # (B_emp,) B_emp_local = int(flow_ps.shape[0]) domain_ids = task_idx_to_domain_idx[emb_id_batch[:B_emp_local]] domain_acc.update(flow_ps, domain_ids) if not torch.isfinite(loss): raise RuntimeError(f"Non-finite loss at step {step}: loss={loss}") loss_to_backprop = loss / grad_accum loss_to_backprop.backward() grad_norm = 0.0 grad_norm_dyn = 0.0 grad_norm_rew = 0.0 grad_norm_bc = 0.0 if do_step: _dyn = dyn.module if hasattr(dyn, "module") else dyn clip_dyn = args.grad_clip_dyn if args.grad_clip_dyn > 0 else float('inf') grad_norm_dyn = float(torch.nn.utils.clip_grad_norm_( _dyn.parameters(), max_norm=clip_dyn, ).item()) if rew_head is not None: _rew = rew_head.module if hasattr(rew_head, "module") else rew_head clip_rew = args.grad_clip_rew if args.grad_clip_rew > 0 else float('inf') grad_norm_rew = float(torch.nn.utils.clip_grad_norm_( _rew.parameters(), max_norm=clip_rew, ).item()) if policy_head is not None: _bc = policy_head.module if hasattr(policy_head, "module") else policy_head clip_bc = args.grad_clip_bc if args.grad_clip_bc > 0 else float('inf') grad_norm_bc = float(torch.nn.utils.clip_grad_norm_( _bc.parameters(), max_norm=clip_bc, ).item()) grad_norm = (grad_norm_dyn ** 2 + grad_norm_rew ** 2 + grad_norm_bc ** 2) ** 0.5 # LR warmup measured from the resume point (see resume_step). steps_since_resume = step - resume_step if args.warmup_steps > 0 and steps_since_resume < args.warmup_steps: warmup_frac = (steps_since_resume + 1) / args.warmup_steps for pg in opt.param_groups: pg["lr"] = args.lr * warmup_frac opt.step() opt.zero_grad(set_to_none=True) # Evaluation / visualization if is_rank0() and args.eval_every > 0 and (step % args.eval_every == 0): # Evaluate on a small slice of the current batch B_eval = min(frames.shape[0], args.eval_batch_size) frames_eval = frames[:B_eval] actions_eval = actions[:B_eval] act_mask_eval = act_mask[:B_eval] lang_emb_eval = lang_emb[:B_eval] sched = make_tau_schedule(k_max=args.k_max, schedule=args.eval_schedule, d=args.eval_d) run_dynamics_eval( encoder=encoder, decoder=decoder, dyn=(dyn.module if hasattr(dyn, "module") else dyn), frames=frames_eval, actions=actions_eval, act_mask=act_mask_eval, H=H, W=W, C=C, patch=patch, packing_factor=args.packing_factor, k_max=args.k_max, ctx_length=args.eval_ctx, horizon=args.eval_horizon, sched=sched, max_items=args.eval_max_items, step=step, tau_ctx=args.tau_ctx, lang_emb=lang_emb_eval, tag="eval", ) # Validation on a held-out dataset (rank 0 only) if val_iter is not None and (step % args.val_every == 0): val_batch = next(val_iter) val_frames, val_actions, val_act_mask, _val_rewards, val_lang_emb = parse_batch( val_batch, use_rewards=use_rewards, device=device, ) B_val = min(val_frames.shape[0], args.eval_batch_size) val_frames = val_frames[:B_val] val_actions = val_actions[:B_val] val_act_mask = val_act_mask[:B_val] val_lang_emb = val_lang_emb[:B_val] sched_val = make_tau_schedule(k_max=args.k_max, schedule=args.eval_schedule, d=args.eval_d) run_dynamics_eval( encoder=encoder, decoder=decoder, dyn=(dyn.module if hasattr(dyn, "module") else dyn), frames=val_frames, actions=val_actions, act_mask=val_act_mask, H=H, W=W, C=C, patch=patch, packing_factor=args.packing_factor, k_max=args.k_max, ctx_length=args.eval_ctx, horizon=args.eval_horizon, sched=sched_val, max_items=args.eval_max_items, step=step, tau_ctx=args.tau_ctx, lang_emb=val_lang_emb, tag="val", ) # Action-attribution metrics averaged over many val batches # (batch-shuffle, time-shuffle, zero-action). if args.action_attr_n_batches > 0: attr_batches = [(val_frames, val_actions, val_act_mask, val_lang_emb)] for _ in range(args.action_attr_n_batches - 1): try: b = next(val_iter) except StopIteration: break f_i, a_i, m_i, _r_i, l_i = parse_batch( b, use_rewards=use_rewards, device=device, ) Bv = min(f_i.shape[0], args.eval_batch_size) attr_batches.append(( f_i[:Bv], a_i[:Bv], m_i[:Bv], l_i[:Bv], )) _attr_rng = torch.Generator(device="cpu") _attr_rng.manual_seed(12345) run_action_attribution_eval( encoder=encoder, dyn=(dyn.module if hasattr(dyn, "module") else dyn), batches=attr_batches, patch=patch, packing_factor=args.packing_factor, k_max=args.k_max, tau_ctx=args.tau_ctx, step=step, tag="val", rng=_attr_rng, ) # Flush per-domain accumulators at log cadence (all ranks # must participate in the all_reduce; only rank 0 logs). domain_flow_means = None domain_flow_counts = None if step % args.log_every == 0: domain_flow_means, domain_flow_counts = domain_acc.flush(ddp=ddp) # Logging if is_rank0() and (step % args.log_every == 0): # Action shuffle loss ratio _dyn_unwrapped = dyn.module if hasattr(dyn, "module") else dyn with torch.no_grad(): loss_real, _ = dynamics_pretrain_loss( _dyn_unwrapped, z1=z1, actions=actions, act_mask=act_mask, k_max=args.k_max, B_self=B_self, step=step, bootstrap_start=args.bootstrap_start, lang_emb=lang_emb, rew_head=None, rewards=None, reward_weight=0.0, ) perm = torch.randperm(actions.shape[0], device=actions.device) loss_shuffled, _ = dynamics_pretrain_loss( _dyn_unwrapped, z1=z1, actions=actions[perm], act_mask=act_mask[perm] if act_mask is not None else None, k_max=args.k_max, B_self=B_self, step=step, bootstrap_start=args.bootstrap_start, lang_emb=lang_emb, rew_head=None, rewards=None, reward_weight=0.0, ) action_shuffle_loss_ratio = loss_shuffled / loss_real.clamp_min(1e-8) boot_over_flow = ( float(aux["bootstrap_mse"].item()) / max(float(aux["flow_mse"].item()), 1e-8) if float(aux["flow_mse"].item()) > 0.0 else 0.0 ) # Weight norms _dyn_log = dyn.module if hasattr(dyn, "module") else dyn weight_norm_dyn = float(sum(p.float().norm().item() ** 2 for p in _dyn_log.parameters()) ** 0.5) weight_norm_rew = 0.0 if rew_head is not None: _rew_log = rew_head.module if hasattr(rew_head, "module") else rew_head weight_norm_rew = float(sum(p.float().norm().item() ** 2 for p in _rew_log.parameters()) ** 0.5) weight_norm_bc = 0.0 if policy_head is not None: _bc_log = policy_head.module if hasattr(policy_head, "module") else policy_head weight_norm_bc = float(sum(p.float().norm().item() ** 2 for p in _bc_log.parameters()) ** 0.5) # Log to wandb wandb.log( { "loss/total": float(loss.item()), "loss/flow_total": float(aux["loss_flow"].item()), "loss/flow_mse": float(aux["flow_mse"].item()), "loss/bootstrap_mse": float(aux["bootstrap_mse"].item()), "loss/bootstrap_over_flow": boot_over_flow, "loss/loss_emp": float(aux["loss_emp"].item()), "loss/loss_self": float(aux["loss_self"].item()), "loss/reward_ce": float(aux["loss_rew"].item()), "loss/bc_mse": float(aux["loss_bc"].item()), "stats/action_shuffle_loss_ratio": float(action_shuffle_loss_ratio.item()), "stats/sigma_mean": float(aux["sigma_mean"].item()), "stats/sigma_std": float(aux["sigma_std"].item()), "stats/B_self": float(B_self), "stats/grad_norm": grad_norm, "stats/grad_norm_dyn": grad_norm_dyn, "stats/grad_norm_rew": grad_norm_rew, "stats/grad_norm_bc": grad_norm_bc, "stats/weight_norm_dyn": weight_norm_dyn, "stats/weight_norm_rew": weight_norm_rew, "stats/weight_norm_bc": weight_norm_bc, "lr": float(opt.param_groups[0]["lr"]), "time/hrs": (time.time() - t0) / 3600.0, "rms/flow": aux["rms_flow"], "rms/rew": aux["rms_rew"], "rms/bc": aux["rms_bc"], }, step=step, ) # Per-domain flow_mse, aggregated over all ranks across the # log_every window. if domain_flow_means is not None: dom_payload = {} for di, dname in enumerate(DOMAINS): m = float(domain_flow_means[di].item()) if not math.isnan(m): dom_payload[f"domain/{dname}/flow_mse"] = m if dom_payload: wandb.log(dom_payload, step=step) # Log to console print( f"step {step:07d} | loss={loss.item():.6f} " f"| flow_mse={aux['flow_mse'].item():.6f} " f"| boot_mse={aux['bootstrap_mse'].item():.6f} " f"| rew_ce={aux['loss_rew'].item():.6f} " f"| sigma={aux['sigma_mean'].item():.3f} | gnorm={grad_norm:.3f} | B_self={B_self}" ) # Checkpointing if is_rank0() and args.save_every > 0 and (step % args.save_every == 0) and step > 0 and do_step: rms_state = {"flow": rms_flow.state_dict(), "rew": rms_rew.state_dict(), "bc": rms_bc.state_dict()} ckpt_path = ckpt_dir / f"step_{step:07d}.pt" save_ckpt(ckpt_path, step=step, epoch=epoch, dyn_model=dyn, rew_head=rew_head, policy_head=policy_head, opt=opt, args=args, rms_state=rms_state) latest = ckpt_dir / "latest.pt" save_ckpt(latest, step=step, epoch=epoch, dyn_model=dyn, rew_head=rew_head, policy_head=policy_head, opt=opt, args=args, rms_state=rms_state) # log checkpoint as a wandb artifact (versioned, easy to retrieve from the UI) artifact = wandb.Artifact( name=f"{wandb.run.id}-dynamics", type="model", metadata={"step": step, "epoch": epoch}, ) artifact.add_file(str(ckpt_path), name=f"step_{step:07d}.pt") wandb.log_artifact(artifact, aliases=["latest", f"step-{step}"]) step += 1 start_epoch = epoch + 1 if step > args.max_steps: break # exit the epoch loop too; otherwise the dataloader keeps spinning idle if ddp: dist.barrier() dist.destroy_process_group() if __name__ == "__main__": p = argparse.ArgumentParser() # data (if using multiple datasets, make sure they align in order) p.add_argument("--data_dirs", type=str, nargs="+", default=[ "./data/expert", "./data/mixed-large", "./data/mixed-small", "./data/zeros", "./data/collected", ]) p.add_argument("--frame_dirs", type=str, nargs="+", default=[ "./data/expert-shards", "./data/mixed-large-shards", "./data/mixed-small-shards", "./data/zeros-shards", "./data/collected-shards", ]) p.add_argument("--tasks_json", type=str, default="../tasks.json") # task metadata # validation data (held-out, single directory each) p.add_argument("--val_data_dir", type=str, default="./data/val", help="optional single raw-data dir for validation") p.add_argument("--val_frame_dir", type=str, default="./data/val-shards", help="optional single preprocessed-frames dir for validation") p.add_argument("--val_every", type=int, default=2_000, help="run validation rollouts every N steps (0 disables)") p.add_argument("--seq_len", type=int, default=24) p.add_argument("--num_workers", type=int, default=4) p.add_argument("--batch_size", type=int, default=64) p.add_argument("--prefetch_factor", type=int, default=4, help="batches prefetched per worker") p.add_argument("--cache_mb", type=int, default=13312, help="per-worker LRU shard cache size in MB (WMDataset)") p.add_argument("--samples_per_shard", type=int, default=24, help="sequences drawn from a shard before switching " "(1 = pure iid, larger = less I/O)") p.add_argument("--task_weighting", type=str, default="valid_starts", choices=["valid_starts", "uniform", "targeted"], help="how to weight tasks during iid sampling. " "'valid_starts' (default): P(task) ∝ #valid windows, " "favors long-trajectory tasks. " "'uniform': every task drawn with equal probability. " "'targeted': split α uniformly over the targeted-collection " "tasks (SEEN_TASK_SET ∪ UNSEEN_TASK_SET) and (1-α) uniformly " "over the rest — use for targeted-collection fine-tunes.") p.add_argument("--targeted_alpha", type=float, default=0.5, help="(--task_weighting=targeted only) total per-batch " "probability mass allocated to the targeted-collection " "tasks (split uniformly among them); the remaining 1-α is " "spread uniformly over all other tasks.") # tokenizer restore p.add_argument("--tokenizer_ckpt", type=str, default="./logs/tokenizer_ckpts/latest.pt") p.add_argument("--H", type=int, default=None) p.add_argument("--W", type=int, default=None) p.add_argument("--C", type=int, default=None) p.add_argument("--patch", type=int, default=None) # dynamics arch p.add_argument("--d_model_dyn", type=int, default=1024) p.add_argument("--dyn_depth", type=int, default=16) p.add_argument("--n_heads", type=int, default=8) p.add_argument("--dropout", type=float, default=0.0) p.add_argument("--mlp_ratio", type=float, default=4.0) p.add_argument("--time_every", type=int, default=2) p.add_argument("--packing_factor", type=int, default=2) p.add_argument("--n_register", type=int, default=4) p.add_argument("--n_agent", type=int, default=0) # shortcut / schedule p.add_argument("--k_max", type=int, default=64) p.add_argument("--bootstrap_start", type=int, default=10_000) p.add_argument("--self_fraction", type=float, default=0.25) # actions p.add_argument("--lang_dim", type=int, default=512) # rewards p.add_argument("--reward_weight", type=float, default=0) # set >0 to enable p.add_argument("--reward_L", type=int, default=8) p.add_argument("--reward_num_bins", type=int, default=255) p.add_argument("--reward_log_low", type=float, default=-10.0) p.add_argument("--reward_log_high", type=float, default=10.0) p.add_argument("--reward_mlp_ratio", type=float, default=2.0) p.add_argument("--reward_pool_agent", type=str, default="attn", choices=["attn", "mean", "first"]) # BC policy head (gradient-isolated from dynamics) p.add_argument("--bc_weight", type=float, default=0) # set >0 to enable p.add_argument("--bc_L", type=int, default=8) p.add_argument("--bc_act_dim", type=int, default=16) p.add_argument("--bc_mlp_ratio", type=float, default=2.0) p.add_argument("--bc_pool_agent", type=str, default="attn", choices=["attn", "mean", "first"]) # optim p.add_argument("--lr", type=float, default=1e-4) p.add_argument("--weight_decay", type=float, default=1e-2) p.add_argument("--max_steps", type=int, default=10_000_000) p.add_argument("--grad_accum", type=int, default=1) p.add_argument("--grad_clip_dyn", type=float, default=1.0) p.add_argument("--grad_clip_rew", type=float, default=1.0) p.add_argument("--grad_clip_bc", type=float, default=1.0) p.add_argument("--warmup_steps", type=int, default=1000) # eval / viz p.add_argument("--eval_every", type=int, default=1_000) p.add_argument("--eval_batch_size", type=int, default=4) p.add_argument("--eval_max_items", type=int, default=4) p.add_argument("--eval_ctx", type=int, default=8) p.add_argument("--eval_horizon", type=int, default=16) p.add_argument("--eval_schedule", type=str, default="shortcut", choices=["finest", "shortcut"]) p.add_argument("--eval_d", type=float, default=0.25) p.add_argument("--tau_ctx", type=float, default=0.1) # context corruption at inference p.add_argument("--action_attr_n_batches", type=int, default=32, help="N val batches to average the action-attribution metrics over (0 disables).") # logging p.add_argument("--log_every", type=int, default=200) # wandb p.add_argument("--wandb_project", type=str, default="mmbench2-dynamics") p.add_argument("--wandb_run_name", type=str, default="default") p.add_argument("--wandb_entity", type=str, default=None) # ckpt p.add_argument("--ckpt_dir", type=str, default="./logs/dynamics_ckpts") p.add_argument("--save_every", type=int, default=5_000) p.add_argument("--resume", type=str, default=None, help="Path to a dynamics checkpoint to resume from. Defaults to None (fresh run); " "pass e.g. ./logs/dynamics_ckpts/latest.pt to resume.") p.add_argument("--resume_strict", action="store_true", default=True, help="Use strict=True when loading checkpoint (default). Use --no-resume_strict for finetuning with architecture changes (e.g. adding agent tokens).") p.add_argument("--no-resume_strict", dest="resume_strict", action="store_false") # misc p.add_argument("--seed", type=int, default=0) p.add_argument("--compile", action="store_true") args = p.parse_args() train(args)