validateNN.py

import argparse
from pathlib import Path

import torch
import torch.nn as nn


def double_integrator(x: torch.Tensor, u: torch.Tensor) -> torch.Tensor:
    return torch.stack([x[..., 1], u[..., 0]], dim=-1)


def control_barrier_function(x: torch.Tensor, c: float) -> torch.Tensor:
    # Matches the legacy MATLAB sign convention.
    return 0.5 * torch.sum(x * x, dim=-1, keepdim=True) - c


class CBFRegressor(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(2, 64)
        self.fc2 = nn.Linear(64, 64)
        self.fc3 = nn.Linear(64, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)


def load_model(path: str):
    payload = torch.load(path, map_location="cpu")
    required = [
        "model_state_dict",
        "input_mean",
        "input_std",
        "output_mean",
        "output_std",
    ]
    for key in required:
        if key not in payload:
            raise ValueError(f"Model checkpoint missing key: {key}")

    model = CBFRegressor()
    model.load_state_dict(payload["model_state_dict"])
    model.eval()

    stats = {
        "x_mean": payload["input_mean"].float(),
        "x_std": payload["input_std"].float(),
        "y_mean": payload["output_mean"].float(),
        "y_std": payload["output_std"].float(),
    }
    return model, stats


def predict_h(model: nn.Module, stats: dict, x: torch.Tensor) -> torch.Tensor:
    x_norm = (x - stats["x_mean"]) / stats["x_std"]
    y_norm = model(x_norm)
    y = y_norm * stats["y_std"] + stats["y_mean"]
    return y


def simulate(model: nn.Module, stats: dict, x0, T: float, dt: float):
    steps = int(round(T / dt))
    t = torch.linspace(0.0, T, steps + 1, dtype=torch.float32)

    x = torch.zeros(steps + 1, 2, dtype=torch.float32)
    u = torch.zeros(steps, 1, dtype=torch.float32)
    h_pred = torch.zeros(steps, 1, dtype=torch.float32)
    x[0] = torch.tensor(x0, dtype=torch.float32)

    for k in range(steps):
        u_k = -x[k, 0] - x[k, 1]

        h_k = predict_h(model, stats, x[k : k + 1]).item()
        h_pred[k, 0] = h_k
        if h_k < 0.0:
            u_k = -0.1 * x[k, 0] - 0.1 * x[k, 1]

        u[k, 0] = u_k
        dx = double_integrator(x[k : k + 1], u[k : k + 1]).squeeze(0)
        x[k + 1] = x[k] + dt * dx

    return t, x, u, h_pred


def maybe_plot(t: torch.Tensor, x: torch.Tensor, h_pred: torch.Tensor):
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        print("matplotlib not installed; skipping plot.")
        return

    t_list = t.tolist()
    x1 = x[:, 0].tolist()
    x2 = x[:, 1].tolist()
    h_list = h_pred[:, 0].tolist()

    fig, axes = plt.subplots(3, 1, figsize=(10, 8))
    axes[0].plot(t_list, x1, color="r", linewidth=1.5)
    axes[0].set_xlabel("Time (s)")
    axes[0].set_ylabel("Position")
    axes[0].set_title("Double Integrator with Learned CBF")

    axes[1].plot(t_list, x2, color="b", linewidth=1.5)
    axes[1].set_xlabel("Time (s)")
    axes[1].set_ylabel("Velocity")

    axes[2].plot(t_list[:-1], h_list, color="g", linewidth=1.5)
    axes[2].axhline(0.0, color="k", linestyle="--", linewidth=1.0)
    axes[2].set_xlabel("Time (s)")
    axes[2].set_ylabel("Predicted h(x)")

    fig.tight_layout()
    plt.show()


def main():
    parser = argparse.ArgumentParser(description="Validate learned CBF model in simulation (Python replacement for validateNN.m).")
    parser.add_argument(
        "--model",
        type=str,
        default="artifacts/models/cbf_model.pth",
        help="Model checkpoint generated by trainNN.py.",
    )
    parser.add_argument("--c", type=float, default=1.0, help="Safe set constant for reporting true barrier values.")
    parser.add_argument("--x0", type=float, nargs=2, default=[1.0, 0.0], help="Initial state [x1, x2].")
    parser.add_argument("--T", type=float, default=10.0, help="Simulation horizon.")
    parser.add_argument("--dt", type=float, default=0.01, help="Time step.")
    parser.add_argument(
        "--output",
        type=str,
        default="artifacts/data/safe_trajectories_learned.pt",
        help="Output file (.pt) with rollout states, controls, and predicted barrier values.",
    )
    parser.add_argument("--plot", action="store_true", help="Plot rollout results.")
    args = parser.parse_args()

    model, stats = load_model(args.model)
    t, x, u, h_pred = simulate(model, stats, x0=args.x0, T=args.T, dt=args.dt)
    h_true = control_barrier_function(x, args.c)

    payload = {
        "time": t,
        "states": x,
        "controls": u,
        "predicted_h": h_pred,
        "true_h": h_true,
        "safe_trajectories": x.T.contiguous(),
        "c": args.c,
        "dt": args.dt,
    }
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(payload, output_path)

    unsafe_ratio = (h_true < 0.0).float().mean().item()
    print(f"Saved validation rollout to {output_path}")
    print(f"unsafe_state_ratio (true h < 0): {unsafe_ratio:.6f}")

    if args.plot:
        maybe_plot(t, x, h_pred)


if __name__ == "__main__":
    main()