| |
| """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 |
| from src.dualrate.data import DualRateDataset, collate_dualrate |
| from src.dualrate.model import DualRatePhysicsModel |
| from src.dualrate.normalization import StateNormalizer |
|
|
|
|
| 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: |
| 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() |
|
|
