train.py

import os
import datetime
import numpy as np
import matplotlib.pyplot as plt
import argparse
import torch
from lib.models import PushPlanner, NNModel, NNPhysicsModel, PushNetFactory
from helpers.utils import (
    load_data,
    prepare_dataloaders,
    prepare_dataloaders_split,
    prepare_dataloader,
    evaluate_planner,
    save_checkpoint,
    load_checkpoint,
)
from helpers.config import load_config
from tqdm import tqdm
from colorama import init, Fore, Style

# Initialize colorama
init()


def print_header(text: str):
    print(f"\n{Fore.CYAN}{Style.BRIGHT}{text}{Style.RESET_ALL}")


def print_success(text: str):
    print(f"{Fore.GREEN}{text}{Style.RESET_ALL}")


def print_info(text: str):
    print(f"{Fore.YELLOW}{text}{Style.RESET_ALL}")


def print_error(text: str):
    print(f"{Fore.RED}{text}{Style.RESET_ALL}")


def parse_args():
    parser = argparse.ArgumentParser(description="Train push planning model")
    parser.add_argument("--config", type=str, default=None, help="Path to config file")
    parser.add_argument(
        "--checkpoint",
        type=str,
        default=None,
        help="Path to checkpoint file to resume training",
    )
    return parser.parse_args()


def main():
    # Parse command line arguments
    args = parse_args()

    # Load configuration
    config = load_config(args.config)
    device = config.get_device()
    print_info(f"Using device: {device}")

    # Create results directory
    os.makedirs("results", exist_ok=True)

    # Load data
    print_header("Loading Data")
    x_data, y_data = load_data(config)
    print_info(f"Loaded data shapes: x={x_data.shape}, y={y_data.shape}")

    # Split dataset into train / validation / test (70 / 15 / 15)
    train_loader, val_loader, test_loader = prepare_dataloaders_split(
        x_data, y_data, config, val_split=0.15, test_split=0.15
    )
    print_info(
        f"Train batches: {len(train_loader)}, "
        f"Val batches: {len(val_loader)}, "
        f"Test batches: {len(test_loader)}"
    )

    # ToDO: Call Physics Push Planner
    # planner = PushPlanner(config.model, config.physics_sampling)

    # net_cfg = config.model["network"]
    # model = NNModel(
    #     input_dim=net_cfg["input_dim"],
    #     output_dim=net_cfg["task_dim"],
    #     hidden_dims=net_cfg["hidden_dims"],
    # ).to(device)

    model = PushNetFactory.create(config.model).to(device)

    # lr is under model.optimizer in config
    optimizer = torch.optim.Adam(
        model.parameters(), lr=config.model["optimizer"]["learning_rate"]
    )

    print_header("Starting Training")
    pbar = tqdm(range(config.training["num_epochs"]), desc="Training Progress")

    # ToDO: Implement training loop
    train_losses = []
    val_losses = []
    for epoch in pbar:
        # Training
        model.train()
        epoch_loss = 0.0
        for x_batch, y_batch in train_loader:
            x_batch, y_batch = x_batch.to(device), y_batch.to(device)
            optimizer.zero_grad()
            preds = model(x_batch)
            loss = model.loss(preds, y_batch)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()

        avg_train_loss = epoch_loss / len(train_loader)
        train_losses.append(avg_train_loss)

        # Validation
        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for x_batch, y_batch in val_loader:
                x_batch, y_batch = x_batch.to(device), y_batch.to(device)
                preds = model(x_batch)
                val_loss += model.loss(preds, y_batch).item()

        avg_val_loss = val_loss / len(val_loader)
        val_losses.append(avg_val_loss)

        pbar.set_postfix({"train_loss": f"{avg_train_loss:.4f}", "val_loss": f"{avg_val_loss:.4f}"})

    print_success("\nTraining completed!")

    # Visualize loss curves
    plt.figure(figsize=(10, 5))
    plt.plot(train_losses, label="Train Loss", color="blue")
    plt.plot(val_losses, label="Validation Loss", color="orange")
    plt.xlabel("Epoch")
    plt.ylabel("MSE Loss")
    plt.title("Training vs Validation Loss")
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig("results/training_curves.png")
    plt.show()
    print_success("Loss curve saved to results/training_curves.png")

    # Test prediction on held-out test set
    print_header("Testing on Test Set")

    model.eval()
    test_loss_total = 0.0
    all_preds, all_targets = [], []
    with torch.no_grad():
        for x_batch, y_batch in test_loader:
            x_batch, y_batch = x_batch.to(device), y_batch.to(device)
            p = model(x_batch)
            test_loss_total += model.loss(p, y_batch).item()
            all_preds.append(p.cpu())
            all_targets.append(y_batch.cpu())

    avg_test_loss = test_loss_total / len(test_loader)
    preds_cat   = torch.cat(all_preds)
    targets_cat = torch.cat(all_targets)
    metrics = model.accuracy(preds_cat, targets_cat)

    print_success(f"Test MSE Loss:      {avg_test_loss:.4f}")
    print_info(   f"Position Error:     {metrics['position_error']:.4f}")
    print_info(   f"Angle Error (rad):  {metrics['angle_error_rad']:.4f}")

    # Full-dataset tensors (used for full-dataset error)
    x_tensor = torch.FloatTensor(x_data).to(device)
    y_tensor = torch.FloatTensor(y_data).to(device)
    with torch.no_grad():
        preds = model(x_tensor)
        error = model.loss(preds, y_tensor)
    print_info(f"Full-dataset Prediction Error: {error.item():.4f}")

    # Collect test-split arrays (same split as main_physics.py)
    x_test = np.concatenate([xb.numpy() for xb, _ in test_loader])
    y_test = np.concatenate([yb.numpy() for _, yb in test_loader])

    # Plot trajectory comparison for 10 pushes (test set)
    print_header("Plotting Trajectory Comparison")
    x_test_tensor = torch.FloatTensor(x_test).to(device)
    with torch.no_grad():
        preds_all = model(x_test_tensor).cpu().numpy()
    gt_all = y_test  # [N_test, 3]: x, y, theta

    num_pushes = 10
    indices = np.arange(num_pushes)  # first 10 samples

    gt_x  = gt_all[indices, 0]
    gt_y  = gt_all[indices, 1]
    gt_th = gt_all[indices, 2]

    pr_x  = preds_all[indices, 0]
    pr_y  = preds_all[indices, 1]
    pr_th = preds_all[indices, 2]

    arrow_len = 0.015

    fig, ax = plt.subplots(figsize=(9, 7))

    # Ground truth
    ax.plot(gt_x, gt_y, "o-", color="royalblue", label="Ground Truth", zorder=2)
    for i, (xi, yi, ti) in enumerate(zip(gt_x, gt_y, gt_th)):
        ax.annotate(
            "", xy=(xi + arrow_len * np.cos(ti), yi + arrow_len * np.sin(ti)),
            xytext=(xi, yi),
            arrowprops=dict(arrowstyle="-|>", color="royalblue", lw=1.5),
            zorder=3,
        )
        ax.text(xi, yi + 0.003, str(i), fontsize=7, ha="center", color="royalblue")

    # Predicted
    ax.plot(pr_x, pr_y, "s--", color="tomato", label="Predicted", zorder=2)
    for i, (xi, yi, ti) in enumerate(zip(pr_x, pr_y, pr_th)):
        ax.annotate(
            "", xy=(xi + arrow_len * np.cos(ti), yi + arrow_len * np.sin(ti)),
            xytext=(xi, yi),
            arrowprops=dict(arrowstyle="-|>", color="tomato", lw=1.5),
            zorder=3,
        )

    ax.set_xlabel("X position (m)")
    ax.set_ylabel("Y position (m)")
    ax.set_title("Push Trajectory: Ground Truth vs Predicted (10 test-set pushes)\nArrows indicate object orientation (θ)")
    ax.legend()
    ax.set_aspect("equal")
    ax.grid(True, linestyle="--", alpha=0.5)
    plt.tight_layout()
    plt.savefig("results/trajectory_comparison.png", dpi=150)
    plt.show()
    print_success("Trajectory plot saved to results/trajectory_comparison.png")

    # Save model
    print_header("Saving Model")
    os.makedirs("models", exist_ok=True)
    net_type = config.model["network"]["type"]
    model_label = "nn" if net_type == "NNModel" else "hybrid"
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    model_path = f"models/{model_label}_{timestamp}.pt"
    torch.save(model.state_dict(), model_path)
    print_success(f"Model saved to {model_path}")


if __name__ == "__main__":
    main()