train.py

"""
Training script for the learned dynamics forecaster (converted from Data_ITransformer.ipynb).

Loads logged data, builds the DynamicForecaster from training/src, trains with ADE/FDE
monitoring, and saves the best checkpoint.

Example:
  python training/train.py --train-files data/full_state1.p data/full_state2.p --test-file data/stanley.p --save-path models/small.pth
"""
from __future__ import annotations

import argparse
import random
from pathlib import Path

import numpy as np
import torch
from tqdm import tqdm

from src.dataset import loadLog, create_dataset
from src.models import DynamicForecaster
from src.training_utils import WeightedMSE, train_epoch
from src.geometry_utils import body_frame_to_world


def ade_fde(pred_xy: torch.Tensor, gt_xy: torch.Tensor) -> tuple[float, float]:
    d = torch.sqrt(((pred_xy - gt_xy) ** 2).sum(-1))
    return d.mean().item(), d[:, -1].mean().item()


@torch.no_grad()
def evaluate(model, past, fut_states, fut_ctrl, anchors, batch=128, device="cpu"):
    model.eval()
    N = past.size(0)
    ade_sum = fde_sum = 0.0

    for i in range(0, N, batch):
        bp = past[i : i + batch].to(device)
        bfs = fut_states[i : i + batch].to(device)
        bfc = fut_ctrl[i : i + batch].to(device)
        ba = anchors[i : i + batch].to(device)

        pred = model(bp, bfc)  # [B, Tf, 7]
        pred_w = body_frame_to_world(pred[:, :, :2], ba[:, :2], ba[:, 4])
        gt_w = body_frame_to_world(bfs[:, :, :2], ba[:, :2], ba[:, 4])

        ade, fde = ade_fde(pred_w, gt_w)
        ade_sum += ade * bp.size(0)
        fde_sum += fde * bp.size(0)

    return ade_sum / N, fde_sum / N


def parse_args():
    parser = argparse.ArgumentParser(description="Train the learned dynamics forecaster.")
    parser.add_argument(
        "--train-files",
        nargs="+",
        default=[
            "data/full_state1.p",
            "data/full_state2.p",
            "data/full_state3.p",
            "data/full_state4.p",
            "data/full_state5.p",
            "data/drifting.p",
            "data/control_effort_penalty.p",
        ],
        help="List of pickle log files for training.",
    )
    parser.add_argument("--test-file", default="data/stanley.p", help="Pickle log file for validation/testing.")
    parser.add_argument("--Tp", type=int, default=30, help="Past horizon.")
    parser.add_argument("--Tf", type=int, default=20, help="Future horizon.")
    parser.add_argument("--epochs", type=int, default=30)
    parser.add_argument("--batch", type=int, default=256)
    parser.add_argument("--lr", type=float, default=2e-4)
    parser.add_argument("--weight-decay", type=float, default=1e-2)
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--d-model", type=int, default=120)
    parser.add_argument("--nhead", type=int, default=8)
    parser.add_argument("--d-control", type=int, default=8)
    parser.add_argument("--depth", type=int, default=4)
    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--save-path", default="small.pth", help="Where to save the best checkpoint.")
    parser.add_argument("--seed", type=int, default=0)
    return parser.parse_args()


def main():
    args = parse_args()

    torch.manual_seed(args.seed)
    random.seed(args.seed)
    np.random.seed(args.seed)

    device = args.device

    # Load data
    train_logs = []
    for f in tqdm(args.train_files, desc="Loading train logs"):
        train_logs.append(loadLog(f))
    train_set = np.concatenate(train_logs)
    test_set = loadLog(args.test_file)

    print("Creating trajectory dataset...")

    past_train, fut_train, ctrl_train, anc_train = create_dataset(train_set, args.Tp, args.Tf)
    past_val, fut_val, ctrl_val, anc_val = create_dataset(test_set, args.Tp, args.Tf)

    past_train = past_train.to(device)
    fut_train = fut_train.to(device)
    ctrl_train = ctrl_train.to(device)
    anc_train = anc_train.to(device)

    past_val = past_val.to(device)
    fut_val = fut_val.to(device)
    ctrl_val = ctrl_val.to(device)
    anc_val = anc_val.to(device)

    model = DynamicForecaster(
        Tp=args.Tp,
        Tf=args.Tf,
        d_model=args.d_model,
        nhead=args.nhead,
        d_control=args.d_control,
        depth=args.depth,
        dropout=args.dropout,
    ).to(device)

    crit = WeightedMSE([2.0, 2.0, 0.5, 0.5, 0.1, 0.1, 0.1])
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)

    best_fde = float("inf")
    save_path = Path(args.save_path)
    save_path.parent.mkdir(parents=True, exist_ok=True)

    pbar = tqdm(range(1, args.epochs + 1), desc="Training epochs")

    for ep in pbar:
        train_loss = train_epoch(model, opt, crit, past_train, fut_train, ctrl_train, batch=args.batch, device=device)
        ade, fde = evaluate(model, past_val, fut_val, ctrl_val, anc_val, batch=128, device=device)
        pbar.set_postfix({"train": train_loss, "ade": ade, "fde": fde})

        if fde < best_fde:
            best_fde = fde
            torch.save(model.state_dict(), save_path)

    print(f"Training complete. Best FDE: {best_fde:.3f}. Checkpoint saved to {save_path}")


if __name__ == "__main__":
    main()