quick_plot_app.py

"""
Streamlit app for sampling learned-dynamics trajectories with adjustable control bounds,
and benchmarking against a simple torch-based dynamic bicycle rollout.

Run locally with:
  streamlit run training/quick_plot_app.py
"""

from __future__ import annotations

import sys
from pathlib import Path
from typing import Tuple

import matplotlib.pyplot as plt
import numpy as np
import streamlit as st
import torch
from time import perf_counter

# Repo roots for imports
ROOT = Path(__file__).resolve().parents[0]

from src.dataset import create_dataset, loadLog  # type: ignore
from src.models import DynamicForecaster  # type: ignore
from src.utils import body_frame_to_world  # type: ignore
from src.benchmark.torch_dynamics import get_trajectory_torch_with_defaults  # type: ignore

@st.cache_data(show_spinner=True)
def load_data(tp: int = 30, tf: int = 20):
    """Load validation/test tensors from disk."""
    data_dir = ROOT / "data"
    training_files = [
        "full_state1.p",
        "full_state2.p",
        "full_state3.p",
        "full_state4.p",
        "full_state5.p",
        "drifting.p",
        "control_effort_penalty.p",
    ]
    test_file = "stanley.p"

    _ = [loadLog(str(data_dir / f)) for f in training_files]  # training set not used for viz
    test_set = loadLog(str(data_dir / test_file))

    past_val, fut_val, ctrl_val, anc_val = create_dataset(test_set, tp, tf)
    return past_val, fut_val, ctrl_val, anc_val


@st.cache_resource(show_spinner=True)
def load_model(model_path: str, tp: int, tf: int, device: str):
    """Instantiate and load the trained DynamicForecaster."""
    # DynamicForecaster(Tp=Tp, Tf=Tf, d_model=120, nhead=8, d_control=8, depth=4, dropout=0.1).to(device)

    model = DynamicForecaster(Tp=tp, Tf=tf, d_model=120, nhead=8, d_control=8, depth=4, dropout=0.1).to(device)
    state = torch.load(model_path, map_location=torch.device(device))
    model.load_state_dict(state)
    model.eval()
    return model


def sample_plot(
    model: torch.nn.Module,
    past: torch.Tensor,
    fut_states: torch.Tensor,
    fut_ctrl: torch.Tensor,
    anchors: torch.Tensor,
    index: int,
    n_samples: int,
    throttle_bounds: Tuple[float, float],
    steering_bounds: Tuple[float, float],
    device: str,
    dt: float,
):
    """Generate a matplotlib figure of sampled predictions and timing stats."""
    bp = past[index : index + 1].to(device)
    bfs = fut_states[index : index + 1].to(device)
    bfc = fut_ctrl[index : index + 1].to(device)
    ba = anchors[index : index + 1].to(device)

    # Ground truth / past in world frame
    past_xy = bp[0, :, :2]
    gt_xy = bfs[0, :, :2]
    past_w = body_frame_to_world(past_xy.unsqueeze(0), ba[:, :2], ba[:, 4])[0].cpu()
    gt_w = body_frame_to_world(gt_xy.unsqueeze(0), ba[:, :2], ba[:, 4])[0].cpu()

    fig, ax = plt.subplots(figsize=(6, 6))
    ax.plot(past_w[:, 0], past_w[:, 1], "k.--", label="Past")

    mins = torch.tensor([throttle_bounds[0], steering_bounds[0]], device=device)
    maxs = torch.tensor([throttle_bounds[1], steering_bounds[1]], device=device)
    ranges = maxs - mins

    # Vectorized sampling of controls and inference
    t0 = perf_counter()
    controls = torch.rand((n_samples, bfc.shape[1], 2), device=device) * ranges + mins  # [B, Tf, 2]
    state_encoding = model.past_enc(bp).expand(controls.shape[0], -1)
    ctrl_encoding = model.fc_enc(controls)
    pred = model.dec(state_encoding, ctrl_encoding)[:, :, :2]  # [B, Tf, 2]
    pred_w = body_frame_to_world(pred, ba[:, :2], ba[:, 4]).cpu()  # [B, Tf, 2]
    t1 = perf_counter()

    # Physics baseline with same controls
    controls_phys = controls.permute(2, 0, 1)  # [2, B, Tf]
    traj_phys = get_trajectory_torch_with_defaults(
        controls_phys,
        dt=dt,
        x0=float(ba[0, 0]),
        y0=float(ba[0, 1]),
        heading0=float(ba[0, 4]),
        vx0=float(ba[0, 2]),
        vy0=float(ba[0, 3]),
        n_samples=n_samples,
        time_lookahead=controls.shape[1],
    )
    phys_xy = traj_phys[0:2].permute(1, 2, 0).cpu()  # [B, Tf, 2]
    t2 = perf_counter()

    pred_w_np = pred_w.cpu().detach().numpy()

    for i in range(n_samples):
        ax.plot(
            pred_w_np[i, :, 0],
            pred_w_np[i, :, 1],
            color="b",
            alpha=0.15,
            linewidth=1,
            label="Learned" if i == 0 else None,
        )
        ax.plot(
            phys_xy[i, :, 0],
            phys_xy[i, :, 1],
            color="g",
            alpha=0.15,
            linewidth=1,
            label="Physics" if i == 0 else None,
        )

    ade = torch.sqrt(((pred_w - phys_xy) ** 2).sum(-1)).mean().item()
    ax.text(0.02, 0.98, f"ADE (learned vs physics): {ade:.3f} m", transform=ax.transAxes, va="top")

    ax.plot(gt_w[:, 0], gt_w[:, 1], "r-", label="Future GT")
    ax.axis("equal")
    ax.grid(True)
    ax.set_xlabel("x [m]")
    ax.set_ylabel("y [m]")
    ax.set_title(f"{n_samples} sampled predictions vs physics")
    ax.legend(loc="lower left")
    timings = {
        "learned_ms": (t1 - t0) * 1000.0,
        "physics_ms": (t2 - t1) * 1000.0,
        "total_ms": (t2 - t0) * 1000.0,
    }

    return fig, timings


def main():
    st.title("Learned Dynamics Quick Plot")
    st.write(
        "Sample predictions from the learned dynamics model while adjusting throttle/steering bounds "
        "to see how trajectories change, and compare against a simple physics rollout."
    )

    tp, tf = 30, 20
    data_load_state = st.info("Loading dataset...", icon="⏳")
    past_val, fut_val, ctrl_val, anc_val = load_data(tp, tf)
    data_load_state.empty()

    total_samples = past_val.shape[0]
    st.write(f"Loaded validation set with **{total_samples}** sequences (Tp={tp}, Tf={tf}).")

    default_model_path = ROOT / "small.pth"
    model_path = st.text_input("Model checkpoint path", value=str(default_model_path))

    device_opts = ["cuda"] if torch.cuda.is_available() else []
    device_opts.append("cpu")
    device = st.selectbox("Device", options=device_opts, index=0)

    model = load_model(model_path, tp, tf, device)

    col1, col2 = st.columns(2)
    with col1:
        n_samples = st.slider("Number of samples", min_value=10, max_value=1000, value=200, step=10)
        sample_idx = st.slider("Sample index", min_value=0, max_value=total_samples - 1, value=0, step=1)
    with col2:
        throttle_bounds = st.slider("Throttle bounds", 0.0, 1.0, (0.0, 1.0), step=0.01)
        steering_bounds = st.slider("Steering bounds", -0.455, 0.455, (-0.455, 0.455), step=0.01)
        dt = st.slider("Physics dt", 0.005, 0.05, 0.01, step=0.005)

    with st.spinner("Sampling trajectories..."):
        fig, timings = sample_plot(
            model=model,
            past=past_val,
            fut_states=fut_val,
            fut_ctrl=ctrl_val,
            anchors=anc_val,
            index=sample_idx,
            n_samples=n_samples,
            throttle_bounds=throttle_bounds,
            steering_bounds=steering_bounds,
            device=device,
            dt=dt,
        )
    st.pyplot(fig)

    st.caption(
        f"Timing — Learned: {timings['learned_ms']:.1f} ms | Physics: {timings['physics_ms']:.1f} ms | Total: {timings['total_ms']:.1f} ms"
    )


if __name__ == "__main__":
    main()