scripts/train.py

"""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()