main_physics.py

import os
import numpy as np
import matplotlib.pyplot as plt
import argparse
import torch
from lib.physics import PushPhysics
from helpers.utils import load_data, prepare_dataloaders_split
from helpers.config import load_config
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="Physics push planner evaluation")
    parser.add_argument("--config", type=str, default=None, help="Path to config file")
    return parser.parse_args()


def main():
    args = parse_args()

    # Load configuration
    config = load_config(args.config)
    print_info(f"Using device: cpu (physics model runs on CPU)")

    # 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}")

    # Use the same 70/15/15 split as main.py — keep only the test portion
    _, _, test_loader = prepare_dataloaders_split(
        x_data, y_data, config, val_split=0.15, test_split=0.15
    )
    x_test_list, y_test_list = [], []
    for xb, yb in test_loader:
        x_test_list.append(xb.numpy())
        y_test_list.append(yb.numpy())
    x_test = np.concatenate(x_test_list)
    y_test = np.concatenate(y_test_list)
    print_info(f"Test split shapes: x={x_test.shape}, y={y_test.shape}")

    # Build physics model from config
    print_header("Running Physics Model")
    physics = PushPhysics.from_config(config.model["physics"])
    print_info(
        f"Physics params — mass={physics.mass}, size={physics.size}, "
        f"inertia={physics.inertia:.6f}, duration={physics._push_duration}s, "
        f"steps={physics._simulation_steps}"
    )

    # Convert inputs to tensors:  x_data = [N, 3] -> [rotation, side, distance]
    x_tensor = torch.FloatTensor(x_test)   # push parameters
    y_tensor = torch.FloatTensor(y_test)   # ground truth [x, y, theta]

    # Analytically compute final states via physics simulation
    # compute_motion returns [N, 3, simulation_steps]
    with torch.no_grad():
        motion = physics.compute_motion(x_tensor)   # [N, 3, T]

    # Take the final timestep as the predicted final state → [N, 3]
    # preds = motion[:, :, -1]
    preds = motion

    # MSE loss
    mse_loss = torch.mean((preds - y_tensor) ** 2).item()

    # Position error (Euclidean distance on x, y)
    pos_error = torch.norm(preds[:, :2] - y_tensor[:, :2], dim=1).mean().item()

    # Angle error (wrapped to [-pi, pi])
    angle_diff = preds[:, 2] - y_tensor[:, 2]
    angle_error = torch.abs(
        torch.atan2(torch.sin(angle_diff), torch.cos(angle_diff))
    ).mean().item()

    print_success(f"\nPhysics Model Results on Test Set ({x_tensor.shape[0]} samples):")
    print_success(f"  MSE Loss:              {mse_loss:.6f}")
    print_info(   f"  Mean Position Error:   {pos_error:.6f} m")
    print_info(   f"  Mean Angle Error:      {angle_error:.6f} rad  ({angle_error * 180 / 3.14159:.4f} deg)")

    # Visualise predicted vs ground truth for x and y
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    labels = ["x (m)", "y (m)", "theta (rad)"]
    for i, ax in enumerate(axes):
        ax.scatter(y_tensor[:, i].numpy(), preds[:, i].numpy(), alpha=0.3, s=5)
        lims = [
            min(y_tensor[:, i].min(), preds[:, i].min()).item(),
            max(y_tensor[:, i].max(), preds[:, i].max()).item(),
        ]
        ax.plot(lims, lims, "r--", linewidth=1, label="ideal")
        ax.set_xlabel(f"Ground Truth {labels[i]}")
        ax.set_ylabel(f"Physics Pred {labels[i]}")
        ax.set_title(labels[i])
        ax.legend()
        ax.grid(True)

    plt.suptitle(f"Physics Model — MSE={mse_loss:.4f}", fontsize=12)
    plt.tight_layout()
    plt.savefig("results/physics_predictions.png")
    plt.show()
    print_success("Plot saved to results/physics_predictions.png")

    # Trajectory comparison for 10 pushes
    print_header("Plotting Trajectory Comparison (10 pushes)")
    num_pushes = 10
    indices = np.arange(num_pushes)

    gt_np  = y_tensor.numpy()
    pr_np  = preds.numpy()

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

    pr_x  = pr_np[indices, 0]
    pr_y  = pr_np[indices, 1]
    pr_th = pr_np[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")

    # Physics model predictions
    ax.plot(pr_x, pr_y, "s--", color="tomato", label="Physics Model", 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.text(xi, yi - 0.005, str(i), fontsize=7, ha="center", color="tomato")

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


if __name__ == "__main__":
    main()