"""Train FACT from scratch on every valid K-block anchor of one episode.""" from __future__ import annotations import argparse import json import math import random import time from pathlib import Path import numpy as np import torch from torch.utils.data import DataLoader from fact import ( FactAnchorDataset, FactConfig, FactEpisode, FactPolicy, collate, count_parameters, ) def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser() p.add_argument("--dataset-root", required=True) p.add_argument("--episode", type=int, default=0) p.add_argument("--cache", required=True) p.add_argument("--out", required=True) p.add_argument("--steps", type=int, default=3000) p.add_argument("--batch", type=int, default=8) p.add_argument("--blocks", type=int, default=4) p.add_argument("--lr", type=float, default=2e-4) p.add_argument("--min-lr-ratio", type=float, default=0.1) p.add_argument("--weight-decay", type=float, default=0.01) p.add_argument("--warmup", type=int, default=150) p.add_argument("--clip", type=float, default=1.0) p.add_argument("--seed", type=int, default=42) p.add_argument("--workers", type=int, default=4) p.add_argument("--log-every", type=int, default=25) p.add_argument("--save-every", type=int, default=1000) p.add_argument("--bf16", action="store_true") p.add_argument("--config-json", default=None, help="JSON overrides for FactConfig") # common config shortcuts p.add_argument("--dim", type=int, default=None) p.add_argument("--layers", type=int, default=None) p.add_argument("--heads", type=int, default=None) p.add_argument("--ffn", type=int, default=None) p.add_argument("--no-future-state", action="store_true") p.add_argument("--no-force", action="store_true") p.add_argument("--train-fraction", type=float, default=None, help="use only anchors whose K-block span ends before this fraction of the episode") return p.parse_args() def build_config(args: argparse.Namespace) -> FactConfig: cfg = FactConfig() if args.config_json: for key, value in json.loads(args.config_json).items(): if not hasattr(cfg, key): raise KeyError(key) setattr(cfg, key, value) if args.dim: cfg.dim = args.dim if args.layers: cfg.num_layers = args.layers if args.heads: cfg.num_heads = args.heads if args.ffn: cfg.ffn_dim = args.ffn if args.no_future_state: cfg.predict_future_state = False if args.no_force: cfg.use_force = False return cfg def lr_at(step: int, args: argparse.Namespace) -> float: if step < args.warmup: return args.lr * (step + 1) / args.warmup progress = (step - args.warmup) / max(1, args.steps - args.warmup) cosine = 0.5 * (1 + math.cos(math.pi * min(1.0, progress))) return args.lr * (args.min_lr_ratio + (1 - args.min_lr_ratio) * cosine) def main() -> None: args = parse_args() torch.manual_seed(args.seed) np.random.seed(args.seed) random.seed(args.seed) device = torch.device("cuda") out = Path(args.out) out.mkdir(parents=True, exist_ok=True) cfg = build_config(args) model = FactPolicy(cfg).to(device) n_total = count_parameters(model) n_coarse = count_parameters(model.coarse) n_force = count_parameters(model.force) if model.force is not None else 0 print(f"params: total={n_total/1e6:.2f}M coarse={n_coarse/1e6:.2f}M force={n_force/1e6:.2f}M", flush=True) (out / "fact_config.json").write_text(cfg.to_json()) (out / "train_args.json").write_text(json.dumps(vars(args), indent=1)) episode = FactEpisode(args.dataset_root, args.episode, latent_cache=args.cache) dataset = FactAnchorDataset(episode, blocks=args.blocks, train_fraction=args.train_fraction) print(f"episode {args.episode}: {episode.length} frames, {len(dataset)} anchors (K={args.blocks})", flush=True) generator = torch.Generator().manual_seed(args.seed) loader = DataLoader( dataset, batch_size=args.batch, shuffle=True, drop_last=True, num_workers=args.workers, collate_fn=collate, generator=generator, persistent_workers=args.workers > 0, ) decay, no_decay = [], [] for name, param in model.named_parameters(): if not param.requires_grad: continue (no_decay if param.ndim < 2 or "embedding" in name or "modulation" in name else decay).append(param) optimizer = torch.optim.AdamW( [{"params": decay, "weight_decay": args.weight_decay}, {"params": no_decay, "weight_decay": 0.0}], lr=args.lr, betas=(0.9, 0.95), ) log_path = out / "loss_log.jsonl" log_file = log_path.open("w") model.train() step = 0 start = time.time() running: dict[str, float] = {} while step < args.steps: for batch in loader: if step >= args.steps: break lr = lr_at(step, args) for group in optimizer.param_groups: group["lr"] = lr with torch.autocast("cuda", dtype=torch.bfloat16, enabled=args.bf16): losses = model.forward_core(batch) loss = losses["loss"] optimizer.zero_grad(set_to_none=True) loss.backward() grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip) optimizer.step() step += 1 record = {k: float(v) for k, v in losses.items()} record.update({"step": step, "lr": lr, "grad_norm": float(grad_norm)}) log_file.write(json.dumps(record) + "\n") for key, value in record.items(): running[key] = running.get(key, 0.0) + value if step % args.log_every == 0: avg = {k: v / args.log_every for k, v in running.items()} running = {} elapsed = time.time() - start print( f"step {step}/{args.steps} loss {avg['loss']:.4f} act {avg['action_loss']:.4f} " f"state {avg['state_loss']:.4f} force {avg['force_loss']:.4f} vq {avg['vq_loss']:.4f} " f"gn {avg['grad_norm']:.2f} lr {lr:.2e} {elapsed/step:.2f}s/it", flush=True, ) log_file.flush() if step % args.save_every == 0 or step == args.steps: model.save(out / f"checkpoint-{step}", extra={"step": step, "params_total": n_total}) log_file.close() model.save(out / "final", extra={"step": step, "params_total": n_total, "params_coarse": n_coarse, "params_force": n_force}) print(f"done in {(time.time() - start)/60:.1f} min; saved {out/'final'}", flush=True) if __name__ == "__main__": main()