scripts/plot_trajectories.py

#!/usr/bin/env python
"""Plot the manipulated object's world-frame motion from exported trajectories.

    python scripts/plot_trajectories.py --episode <uuid>

Reads the ``*_trajectory.npz`` files written by the pipeline and renders three
panels that answer "how is the object moving":

    A. path over the workbench (world X-Y, viewed from above)
    B. height above the robot base (world Z) against time
    C. speed against time

Series are capped at three clips. That is a colour constraint, not an arbitrary
one: the reference categorical palette's first three slots are the set validated
for all-pairs separation under colour-vision deficiency, and a fourth slot would
put yellow beside orange and fail that floor.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

REPO_ROOT = Path(__file__).resolve().parents[1]

# Reference categorical palette, light mode, slots 1-3 (blue / orange / aqua).
SERIES = ("#2a78d6", "#eb6834", "#1baf7a")
SURFACE = "#fcfcfb"
INK = "#0b0b0b"
INK_SECONDARY = "#52514e"
GRID = "#e4e3df"
MAX_SERIES = 3


def load(path: Path) -> dict:
    with np.load(path, allow_pickle=False) as z:
        return {k: z[k] for k in z.files}


def style_axes(ax: plt.Axes, title: str, xlabel: str, ylabel: str) -> None:
    """Recessive grid and axes; the data is the only assertive thing on screen."""
    ax.set_facecolor(SURFACE)
    ax.set_title(title, color=INK, fontsize=11, loc="left", pad=10)
    ax.set_xlabel(xlabel, color=INK_SECONDARY, fontsize=9)
    ax.set_ylabel(ylabel, color=INK_SECONDARY, fontsize=9)
    ax.grid(True, color=GRID, linewidth=0.8, zorder=0)
    ax.set_axisbelow(True)
    for side in ("top", "right"):
        ax.spines[side].set_visible(False)
    for side in ("left", "bottom"):
        ax.spines[side].set_color(GRID)
    ax.tick_params(colors=INK_SECONDARY, labelsize=8, length=0)


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--episode", required=True, help="episode uuid")
    ap.add_argument("--outputs", type=Path, default=REPO_ROOT / "outputs")
    ap.add_argument("--out", type=Path, default=None)
    args = ap.parse_args()

    ep_dir = args.outputs / args.episode
    files = sorted(ep_dir.glob("*_trajectory.npz"))
    if not files:
        raise SystemExit(f"no trajectory files in {ep_dir}")

    runs = []
    for f in files:
        d = load(f)
        speed = d["object_speed_mps"]
        if not np.isfinite(speed).any():
            continue
        d["_clip"] = f.stem.split("_")[1]
        d["_peak"] = float(np.nanmax(speed))
        runs.append(d)

    if not runs:
        raise SystemExit("no clip produced a finite speed")

    # Show the clips where the object actually moved -- a panel of flat lines
    # tells the reader nothing.
    runs.sort(key=lambda d: d["_peak"], reverse=True)
    shown, dropped = runs[:MAX_SERIES], runs[MAX_SERIES:]

    task = str(runs[0]["task"])
    fig, axes = plt.subplots(1, 3, figsize=(15, 4.6), facecolor=SURFACE)
    fig.suptitle(
        f"Object motion in the world (robot-base) frame\n{task}",
        color=INK,
        fontsize=13,
        x=0.011,
        ha="left",
        y=1.06,
    )

    for i, d in enumerate(shown):
        color = SERIES[i]
        c = d["centroid_xyz_world_m"]
        t = d["timestamps_s"]
        speed = d["object_speed_mps"]
        ok = np.isfinite(c).all(axis=1)
        label = f"clip {d['_clip']}"

        axes[0].plot(c[ok, 0], c[ok, 1], color=color, lw=2.0, zorder=3, label=label)
        # Start hollow, end filled: direction is readable without an arrow legend.
        axes[0].plot(
            c[ok, 0][0], c[ok, 1][0], "o", ms=8, mfc=SURFACE, mec=color, mew=2, zorder=4
        )
        axes[0].plot(c[ok, 0][-1], c[ok, 1][-1], "o", ms=8, color=color, zorder=4)

        axes[1].plot(t[ok], c[ok, 2], color=color, lw=2.0, zorder=3, label=label)
        fin = np.isfinite(speed)
        axes[2].plot(t[fin], speed[fin], color=color, lw=2.0, zorder=3, label=label)

    style_axes(axes[0], "A · Path from above", "world X (m)", "world Y (m)")
    axes[0].set_aspect("equal", adjustable="datalim")
    style_axes(axes[1], "B · Height", "time (s)", "world Z (m)")
    style_axes(axes[2], "C · Speed", "time (s)", "speed (m/s)")

    axes[0].legend(
        frameon=False, fontsize=8, labelcolor=INK_SECONDARY, loc="best"
    )

    note = "hollow = start, filled = end"
    if dropped:
        note += f"  ·  {len(dropped)} slower clip(s) omitted: " + ", ".join(
            f"{d['_clip']} ({d['_peak']:.3f} m/s)" for d in dropped
        )
    fig.text(0.011, -0.06, note, color=INK_SECONDARY, fontsize=8, ha="left")

    out = args.out or (args.outputs / f"{args.episode}_motion.png")
    fig.savefig(out, dpi=150, bbox_inches="tight", facecolor=SURFACE)
    print(f"wrote {out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())