# Written by yams_any4d agent, 2026-07-11. """ Standalone fp32 trainer for the scene-only V4Head on CACHED DiT features. Vanilla PyTorch loop (no FSDP / peft / bf16 casts / imaginaire) — v4 production hyperparameters: AdamW lr 3e-4, warmup 500, cosine, losses vol+grip+rot at 1:1:1 (exactly v4's multiview_losses weighting). Held-out = episodes >= --holdout_ep. Usage (yukon, a4d-bw, single GPU): PYTHONPATH=$PWD:$PWD/externals/AnyData python custom/eval/train_v4head_cached.py \ --cache_dir=$HOME/any4d_work/feat_cache --out_dir=$HOME/any4d_work/v4head_cached \ --prep=$HOME/any4d_work/v4head_prep.npz --iters 15000 """ import argparse import glob import math import os import random import numpy as np import torch from custom.any4d.v4_head import V4Head # NOTE(yams_any4d 2026-07-12): raw DiT features carry a large per-channel DC bias plus # a few massive-activation "sink" channels (e.g. ch 1404 |mean|~1000, up to 5000 on some # windows) that dominate norm/PCA and make features from DIFFERENT videos ~0.97 cosine. # Per-channel standardization drops cross-video cosine to ~0.58, exposing scene content to # the head. FEAT_MEAN/FEAT_STD are loaded from --feat_norm (per-channel, shape (D,)). FEAT_MEAN = None FEAT_STD = None def _norm_feats(f): # f: (..., D). Standardize per channel if stats loaded; else identity. if FEAT_MEAN is None: return f return (f - FEAT_MEAN) / FEAT_STD class FeatCache(torch.utils.data.Dataset): def __init__(self, files): self.files = files def __len__(self): return len(self.files) def __getitem__(self, i): d = torch.load(self.files[i], map_location='cpu', weights_only=False) return (_norm_feats(d['feats0'].float()), d['action'].float(), d['rgb0_thumb'], d['ep']) def collate(items): feats = torch.stack([x[0] for x in items]) acts = torch.stack([x[1] for x in items]) thumbs = np.stack([x[2] for x in items]) eps = [x[3] for x in items] return (feats, acts, thumbs, eps) @torch.no_grad() def evaluate(head, loader, device, max_batches=40): head.eval() agg = {'v4h_vol': 0.0, 'v4h_grip': 0.0, 'v4h_rot': 0.0, 'ee_cm': 0.0} n = 0 for bi, (feats, acts, _, _) in enumerate(loader): if bi >= max_batches: break out = head({0: feats.to(device)}, (feats.shape[2] * 16, feats.shape[3] * 16), acts.to(device)) for k in ('v4h_vol', 'v4h_grip', 'v4h_rot'): agg[k] += float(out['losses'][k]) pred = head.decode_pred_xyz(out['volume_logits'].float(), out['voxel_positions']) agg['ee_cm'] += float((pred - out['gt_xyz']).norm(dim=-1).mean()) * 100.0 n += 1 head.train() return {k: v / max(n, 1) for (k, v) in agg.items()} def main(args): device = 'cuda' global FEAT_MEAN, FEAT_STD if args.feat_norm: _n = np.load(args.feat_norm) FEAT_MEAN = torch.from_numpy(_n['mean']).float() FEAT_STD = torch.from_numpy(_n['std']).float().clamp_min(1e-3) print(f'feature standardization ON from {args.feat_norm} ' f'(mean|max|={FEAT_MEAN.abs().max():.1f} std|max|={FEAT_STD.max():.1f})') files = sorted(glob.glob(f'{args.cache_dir}/ep*.pt')) assert files, f'no cache files in {args.cache_dir}' eps = sorted({int(os.path.basename(f)[2:6]) for f in files}) train_f = [f for f in files if int(os.path.basename(f)[2:6]) < args.holdout_ep] hold_f = [f for f in files if int(os.path.basename(f)[2:6]) >= args.holdout_ep] print(f'episodes {eps[0]}..{eps[-1]} | train windows {len(train_f)} | held-out {len(hold_f)}') assert train_f and hold_f, 'bad split' dl_train = torch.utils.data.DataLoader( FeatCache(train_f), batch_size=args.batch_size, shuffle=True, num_workers=args.workers, collate_fn=collate, drop_last=True, pin_memory=True) dl_hold = torch.utils.data.DataLoader( FeatCache(hold_f), batch_size=args.batch_size, shuffle=False, num_workers=2, collate_fn=collate, pin_memory=True) head = V4Head(prep_path=args.prep, use_wrist=False, temporal_mode=args.temporal_mode, pred_size=args.pred_size, t_out=args.t_out, n_lat_frames=args.n_lat_frames, input_channel_norm=args.input_channel_norm).to(device).float() n_params = sum(p.numel() for p in head.parameters()) print(f'V4Head scene-only: {n_params/1e6:.2f}M params') opt = torch.optim.AdamW(head.parameters(), lr=args.lr) def lr_at(it): if it < args.warmup: return args.lr * (it + 1) / args.warmup p = (it - args.warmup) / max(1, args.iters - args.warmup) return args.lr * (0.05 + 0.475 * (1 + math.cos(math.pi * min(p, 1.0)))) wandb_run = None if args.wandb: import wandb wandb_run = wandb.init(entity='cameronsmithbusiness', project='yam_any4d', name=args.run_name, config=vars(args)) os.makedirs(args.out_dir, exist_ok=True) it, best_hold = 0, 1e9 data_iter = iter(dl_train) while it < args.iters: try: (feats, acts, thumbs, _) = next(data_iter) except StopIteration: data_iter = iter(dl_train) continue for g in opt.param_groups: g['lr'] = lr_at(it) out = head({0: feats.to(device)}, (feats.shape[2] * 16, feats.shape[3] * 16), acts.to(device)) loss = out['losses']['v4h_vol'] + out['losses']['v4h_grip'] + out['losses']['v4h_rot'] opt.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(head.parameters(), 1.0) opt.step() if it % 50 == 0: logs = {f'cached/{k}': float(v) for (k, v) in out['losses'].items()} logs['cached/total'] = float(loss) logs['cached/lr'] = lr_at(it) print(f'it {it} | ' + ' '.join(f'{k.split("/")[1]}={v:.3f}' for k, v in logs.items())) if wandb_run: wandb_run.log(logs, step=it) if it % args.eval_every == 0 and it > 0: hold = evaluate(head, dl_hold, device) print(f' HELD-OUT @ {it}: ' + ' '.join(f'{k}={v:.3f}' for k, v in hold.items())) if wandb_run: wandb_run.log({f'heldout/{k}': v for (k, v) in hold.items()}, step=it) if hold['v4h_vol'] < best_hold: best_hold = hold['v4h_vol'] torch.save(head.state_dict(), f'{args.out_dir}/best.pt') # panels on one held-out sample (reuse the training viz) try: import cv2 from custom.any4d.v4_head_viz import make_v4head_visuals # random window each eval (train or held-out — variety over purity, # per Cameron 2026-07-11) so panels show different scenes/frames rf = random.choice(files) dr = torch.load(rf, map_location='cpu', weights_only=False) feats_h = _norm_feats(dr['feats0'][None].float()) acts_h = dr['action'][None].float() thumbs_h = np.stack([dr['rgb0_thumb']]) print(f' panel window: {os.path.basename(rf)}') with torch.no_grad(): out_h = head({0: feats_h[:1].to(device)}, (feats_h.shape[2] * 16, feats_h.shape[3] * 16), acts_h[:1].to(device)) fake_batch = {'a4d_raw': {'rgb0': torch.from_numpy( thumbs_h[0:1].transpose(0, 3, 1, 2)[:, :, None].repeat(2, axis=2) / 127.5 - 1 ).float()}} panels = make_v4head_visuals(head, out_h, fake_batch, f'{args.out_dir}/panels', it) if wandb_run and panels: import wandb as _wb wandb_run.log({k: _wb.Image(v) for (k, v) in panels.items()}, step=it) except Exception as e: print(f' panel render failed: {e}') it += 1 torch.save(head.state_dict(), f'{args.out_dir}/last.pt') hold = evaluate(head, dl_hold, device, max_batches=1000) print('FINAL HELD-OUT: ' + ' '.join(f'{k}={v:.4f}' for (k, v) in hold.items())) if wandb_run: wandb_run.log({f'heldout_final/{k}': v for (k, v) in hold.items()}) wandb_run.finish() print('TRAIN_CACHED EXIT OK') if __name__ == '__main__': p = argparse.ArgumentParser() p.add_argument('--cache_dir', required=True) p.add_argument('--out_dir', required=True) p.add_argument('--prep', required=True) p.add_argument('--iters', type=int, default=15000) p.add_argument('--batch_size', type=int, default=8) p.add_argument('--lr', type=float, default=3e-4) p.add_argument('--warmup', type=int, default=500) p.add_argument('--eval_every', type=int, default=1000) p.add_argument('--holdout_ep', type=int, default=62) p.add_argument('--workers', type=int, default=6) p.add_argument('--wandb', action='store_true') p.add_argument('--run_name', default='v4head_cached1') p.add_argument('--temporal_mode', default='per_step', choices=['per_step', 'stack_mlp']) p.add_argument('--pred_size', type=int, default=64) p.add_argument('--t_out', type=int, default=41, help='action horizon (timesteps) the head predicts — MUST match the cache action ' 'length (17 for the 640ft/headlora cache, 41 for the original 41-frame cache)') p.add_argument('--n_lat_frames', type=int, default=11, help='# latent frames the stack_mlp head stacks — MUST match the cached feats Tl ' '(11 for 41-frame, 5 for 17-frame/state_t=5, 3 for 9-frame)') p.add_argument('--feat_norm', default=None, help='npz with per-channel mean/std to standardize DiT features (default: off)') p.add_argument('--input_channel_norm', action='store_true', help='online per-sample per-channel standardization inside the head (matches the ' 'joint-LoRA run; use instead of --feat_norm)') main(p.parse_args())