#!/usr/bin/env python3 """Train the dual-rate physics model: frozen Qwen3.5 semantic path + flow expert.""" import argparse import json import math import sys import time from functools import partial from pathlib import Path import numpy as np import torch from torch.utils.data import DataLoader sys.path.insert(0, str(Path(__file__).parent.parent)) from src.dualrate.config import DualRateConfig # noqa: E402 from src.dualrate.data import DualRateDataset, collate_dualrate # noqa: E402 from src.dualrate.model import DualRatePhysicsModel # noqa: E402 from src.dualrate.normalization import StateNormalizer # noqa: E402 def lr_lambda(step, cfg): if step < cfg.warmup_steps: return step / max(cfg.warmup_steps, 1) p = (step - cfg.warmup_steps) / max(cfg.total_steps - cfg.warmup_steps, 1) p = min(p, 1.0) floor = cfg.min_lr / cfg.lr return floor + (1 - floor) * 0.5 * (1 + math.cos(math.pi * p)) @torch.no_grad() def validate(model, val_batches, norm, cfg, device): """Single-step position error (px, Euclidean) vs copy-current-frame baseline.""" model.eval() errs, copy_errs, losses = [], [], [] for batch in val_batches: b = {k: v.to(device) for k, v in batch.items()} latent = model.compute_latent(b["input_ids"], b["attention_mask"]) losses.append(model.flow_loss(latent, b["cur_state"], b["static"], b["obj_mask"], b["target_delta"], b["gravity_dt"]).item()) gen = torch.Generator(device=device) gen.manual_seed(777) delta_norm = model.sample_delta(latent, b["cur_state"], b["static"], b["obj_mask"], b["gravity_dt"], generator=gen) scale = torch.as_tensor(norm.delta_scale[:2], device=device) pred_dpos = delta_norm[..., :2] * scale true_dpos = b["target_delta"][..., :2] * scale mask = b["obj_mask"] err = ((pred_dpos - true_dpos).norm(dim=-1) * mask).sum() / mask.sum() copy_err = (true_dpos.norm(dim=-1) * mask).sum() / mask.sum() errs.append(err.item()) copy_errs.append(copy_err.item()) model.train() return float(np.mean(errs)), float(np.mean(copy_errs)), float(np.mean(losses)) def main(): ap = argparse.ArgumentParser() ap.add_argument("--out", default="runs/dualrate_v1") ap.add_argument("--smoke", action="store_true", help="2k scenes / 2k steps") ap.add_argument("--max-scenes", type=int, default=None) ap.add_argument("--steps", type=int, default=None) ap.add_argument("--batch-size", type=int, default=None) ap.add_argument("--lr", type=float, default=None) ap.add_argument("--tick-k", type=int, default=None) ap.add_argument("--denoise-steps", type=int, default=None) ap.add_argument("--workers", type=int, default=8) ap.add_argument("--no-wandb", action="store_true") ap.add_argument("--val-every", type=int, default=500) ap.add_argument("--save-every", type=int, default=5000) ap.add_argument("--resume", default=None) args = ap.parse_args() cfg = DualRateConfig() if args.steps: cfg.total_steps = args.steps if args.batch_size: cfg.batch_size = args.batch_size if args.lr: cfg.lr = args.lr if args.tick_k: cfg.tick_k = args.tick_k if args.denoise_steps: cfg.denoise_steps = args.denoise_steps max_scenes = args.max_scenes if args.smoke: cfg.total_steps = args.steps or 2000 cfg.warmup_steps = 100 max_scenes = max_scenes or 2000 out_dir = Path(args.out) out_dir.mkdir(parents=True, exist_ok=True) device = "cuda" torch.manual_seed(cfg.seed) from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(cfg.model_name) norm = StateNormalizer.from_config(cfg) train_ds = DualRateDataset("data_dualrate/cache", "train", tokenizer, cfg, norm, max_scenes=max_scenes, seed=cfg.seed) val_ds = DualRateDataset("data_dualrate/cache", "val", tokenizer, cfg, norm, max_scenes=256, seed=cfg.seed) val_ds.deterministic = True collate = partial(collate_dualrate, pad_token_id=tokenizer.pad_token_id) train_dl = DataLoader(train_ds, batch_size=cfg.batch_size, shuffle=True, num_workers=args.workers, collate_fn=collate, pin_memory=True, drop_last=True, persistent_workers=args.workers > 0) val_batches = list(DataLoader(val_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=2, collate_fn=collate)) print(f"train scenes {len(train_ds)}, val scenes {len(val_ds)}, " f"steps {cfg.total_steps}, batch {cfg.batch_size}x{cfg.grad_accum}") model = DualRatePhysicsModel(cfg, readout_token_id=tokenizer.pad_token_id).to(device) trainable = model.trainable_parameters() n_train = sum(p.numel() for p in trainable) print(f"trainable params: {n_train / 1e6:.1f}M") opt = torch.optim.AdamW(trainable, lr=cfg.lr, betas=cfg.betas, weight_decay=cfg.weight_decay) sched = torch.optim.lr_scheduler.LambdaLR(opt, partial(lr_lambda, cfg=cfg)) start_step = 0 if args.resume: ck = torch.load(args.resume, map_location=device, weights_only=False) model.load_state_dict(ck["model"], strict=False) opt.load_state_dict(ck["opt"]) sched.load_state_dict(ck["sched"]) start_step = ck["step"] print(f"resumed from {args.resume} @ step {start_step}") use_wandb = not args.no_wandb if use_wandb: try: import wandb wandb.init(project="physics-dualrate", name=out_dir.name, mode="offline", config={**cfg.to_dict(), "max_scenes": max_scenes}) except Exception as e: # noqa: BLE001 print(f"wandb disabled: {e}") use_wandb = False def save(step, tag): state = {k: v for k, v in model.state_dict().items() if not k.startswith("semantic.lm.")} torch.save({"model": state, "opt": opt.state_dict(), "sched": sched.state_dict(), "step": step, "config": cfg.to_dict(), "normalizer": norm.state_dict()}, out_dir / f"ckpt_{tag}.pt") (out_dir / "config.json").write_text(json.dumps(cfg.to_dict(), indent=2)) model.train() step = start_step t0 = time.time() loss_acc, log_every = [], 50 data_iter = iter(train_dl) while step < cfg.total_steps: opt.zero_grad(set_to_none=True) for _ in range(cfg.grad_accum): try: batch = next(data_iter) except StopIteration: data_iter = iter(train_dl) batch = next(data_iter) b = {k: v.to(device, non_blocking=True) for k, v in batch.items()} latent = model.compute_latent(b["input_ids"], b["attention_mask"]) loss = model.flow_loss(latent, b["cur_state"], b["static"], b["obj_mask"], b["target_delta"], b["gravity_dt"]) (loss / cfg.grad_accum).backward() loss_acc.append(loss.item()) torch.nn.utils.clip_grad_norm_(trainable, cfg.grad_clip) opt.step() sched.step() step += 1 if step % log_every == 0: l = float(np.mean(loss_acc[-log_every * cfg.grad_accum:])) sps = log_every / (time.time() - t0) t0 = time.time() print(f"step {step}/{cfg.total_steps} loss {l:.4f} " f"lr {sched.get_last_lr()[0]:.2e} {sps:.2f} it/s", flush=True) if use_wandb: wandb.log({"train/flow_loss": l, "train/lr": sched.get_last_lr()[0], "train/steps_per_sec": sps}, step=step) if step % args.val_every == 0 or step == cfg.total_steps: pos_err, copy_err, vloss = validate(model, val_batches, norm, cfg, device) print(f" VAL step {step}: dpos_err {pos_err:.3f}px (copy-frame {copy_err:.3f}px) " f"flow_loss {vloss:.4f}", flush=True) if use_wandb: wandb.log({"val/dpos_err_px": pos_err, "val/copy_frame_err_px": copy_err, "val/flow_loss": vloss}, step=step) if step % args.save_every == 0: save(step, f"step{step}") save(step, "final") print(f"done, saved to {out_dir}") if use_wandb: wandb.finish() if __name__ == "__main__": main()