record_eval.py

#!/usr/bin/env python3
"""Record sim video and generate latency charts for π₀ local vs cloud.

Runs gym-aloha episodes, captures every frame, overlays live stats,
then stitches a side-by-side comparison video and a multi-panel latency chart.

Outputs (all in robodal/):
  videos/local_h100.mp4          — local H100 episodes
  videos/cloud_h100.mp4          — Modal H100 episodes  (if --url given)
  videos/comparison.mp4          — side-by-side composite
  charts/latency_breakdown.png   — waterfall + stacked bar + timeline

Usage (from openpi root):
  # Local only:
  MUJOCO_GL=egl uv run robodal/record_eval.py --episodes 3

  # Local + cloud comparison:
  MUJOCO_GL=egl uv run robodal/record_eval.py --url <modal-url> --episodes 3

  # Cloud only (skip local to save time):
  MUJOCO_GL=egl uv run robodal/record_eval.py --url <modal-url> --skip-local --episodes 3
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import pathlib
import sys
import threading
import time
from dataclasses import dataclass, field, asdict
from typing import Optional

import cv2
import imageio
import numpy as np

_DIR = pathlib.Path(__file__).parent
sys.path.insert(0, str(_DIR))
from protocol_pi0 import pack_obs, unpack_obs, pack_response, unpack_response  # noqa

VIDEOS_DIR = _DIR / "videos"
CHARTS_DIR = _DIR / "charts"
CHECKPOINT  = "gs://openpi-assets/checkpoints/pi0_aloha_sim"
TASK        = "gym_aloha/AlohaTransferCube-v0"
FPS         = 30

# ── colours ─────────────────────────────────────────────────────────────────
C_LOCAL   = "#00e676"   # green
C_A10G    = "#ff9100"   # orange
C_H100    = "#40c4ff"   # cyan
C_NET     = "#ef5350"   # red
C_GPU     = "#7c4dff"   # purple
C_PREP    = "#26a69a"   # teal
C_DESER   = "#ffd740"   # amber

# ── timing dataclass ─────────────────────────────────────────────────────────

@dataclass
class CallTiming:
    step: int = 0
    total_ms: float = 0.0
    preprocess_ms: float = 0.0   # obs → policy format
    gpu_ms: float = 0.0          # GPU-only (policy_timing or server report)
    serialize_ms: float = 0.0    # cloud: pack_obs
    network_ms: float = 0.0      # cloud: RTT minus GPU
    deserialize_ms: float = 0.0  # cloud: unpack_response
    overhead_ms: float = 0.0     # everything else


# ── timed policy wrappers ────────────────────────────────────────────────────

class TimedLocalPolicy:
    def __init__(self):
        from openpi.training.config import get_config
        from openpi.policies.policy_config import create_trained_policy
        from openpi_client import image_tools

        print("  Loading π₀ ALOHA sim policy …")
        t0 = time.perf_counter()
        config = get_config("pi0_aloha_sim")
        self._policy = create_trained_policy(config, CHECKPOINT)
        self._image_tools = image_tools
        print(f"  Policy loaded in {time.perf_counter() - t0:.1f}s")

        # Warmup — triggers JAX JIT so the first real episode starts fast.
        print("  Warming up JIT …")
        dummy_obs = {
            "pixels": {"top": np.zeros((480, 640, 3), dtype=np.uint8)},
            "agent_pos": np.zeros(14, dtype=np.float32),
        }
        self.infer(dummy_obs)
        print("  Ready ✓")

    def infer(self, gym_obs: dict, step: int = 0) -> tuple[np.ndarray, CallTiming]:
        from openpi_client import image_tools

        t0 = time.perf_counter()
        img = gym_obs["pixels"]["top"]
        img = image_tools.convert_to_uint8(image_tools.resize_with_pad(img, 224, 224))
        cam_high = np.transpose(img, (2, 0, 1))
        state = np.asarray(gym_obs["agent_pos"], dtype=np.float32)
        obs = {"images": {"cam_high": cam_high}, "state": state}
        t_preproc = time.perf_counter()

        result = self._policy.infer(obs)
        t_infer = time.perf_counter()

        actions = np.asarray(result["actions"])
        gpu_ms = result["policy_timing"]["infer_ms"]
        total_ms = (t_infer - t0) * 1000
        preprocess_ms = (t_preproc - t0) * 1000
        overhead_ms = total_ms - preprocess_ms - gpu_ms

        timing = CallTiming(
            step=step,
            total_ms=total_ms,
            preprocess_ms=preprocess_ms,
            gpu_ms=gpu_ms,
            overhead_ms=max(0, overhead_ms),
        )
        return actions, timing


class TimedCloudPolicyWS:
    def __init__(self, url: str, label: str = "cloud"):
        import websockets  # noqa

        self.label = label
        ws_url = url.replace("https://", "wss://").replace("http://", "ws://").rstrip("/") + "/ws"
        self._ws_url = ws_url
        self._loop = asyncio.new_event_loop()
        self._ws = None
        self._thread = threading.Thread(target=self._loop.run_forever, daemon=True)
        self._thread.start()

        print(f"  Connecting WebSocket [{label}] to {ws_url} …")
        asyncio.run_coroutine_threadsafe(self._connect(), self._loop).result(timeout=120)

        print("  Warming up …")
        dummy = {"cam_high": np.zeros((3, 224, 224), dtype=np.uint8),
                 "state": np.zeros(14, dtype=np.float32)}
        self._send_recv_timed(dummy)
        print(f"  Cloud server ready ✓ [{label}]")

    async def _connect(self):
        import websockets
        self._ws = await websockets.connect(self._ws_url, max_size=10 * 1024 * 1024)

    async def _send_recv_async(self, payload: bytes) -> bytes:
        await self._ws.send(payload)
        return await self._ws.recv()

    def _send_recv_timed(self, wire_obs: dict) -> tuple[np.ndarray, float, float, float, float]:
        """Returns (actions, serialize_ms, network_ms, gpu_ms, deserialize_ms)."""
        t0 = time.perf_counter()
        payload = pack_obs(wire_obs)
        t1 = time.perf_counter()

        fut = asyncio.run_coroutine_threadsafe(self._send_recv_async(payload), self._loop)
        data = fut.result(timeout=120)
        t2 = time.perf_counter()

        action, server_gpu_ms = unpack_response(data)
        t3 = time.perf_counter()

        serialize_ms   = (t1 - t0) * 1000
        rtt_ms         = (t2 - t1) * 1000
        network_ms     = max(0, rtt_ms - server_gpu_ms)
        deserialize_ms = (t3 - t2) * 1000
        return np.asarray(action), serialize_ms, network_ms, server_gpu_ms, deserialize_ms

    def infer(self, gym_obs: dict, step: int = 0) -> tuple[np.ndarray, CallTiming]:
        from openpi_client import image_tools

        t_start = time.perf_counter()
        img = gym_obs["pixels"]["top"]
        img = image_tools.convert_to_uint8(image_tools.resize_with_pad(img, 224, 224))
        cam_high = np.transpose(img, (2, 0, 1))
        state = np.asarray(gym_obs["agent_pos"], dtype=np.float32)
        wire_obs = {"cam_high": cam_high, "state": state}
        t_prep = time.perf_counter()

        actions, ser_ms, net_ms, gpu_ms, deser_ms = self._send_recv_timed(wire_obs)
        t_end = time.perf_counter()

        preprocess_ms = (t_prep - t_start) * 1000
        total_ms = (t_end - t_start) * 1000
        overhead_ms = max(0, total_ms - preprocess_ms - ser_ms - net_ms - gpu_ms - deser_ms)

        timing = CallTiming(
            step=step,
            total_ms=total_ms,
            preprocess_ms=preprocess_ms,
            gpu_ms=gpu_ms,
            serialize_ms=ser_ms,
            network_ms=net_ms,
            deserialize_ms=deser_ms,
            overhead_ms=overhead_ms,
        )
        return actions, timing

    def close(self):
        async def _c():
            if self._ws:
                await self._ws.close()
        asyncio.run_coroutine_threadsafe(_c(), self._loop).result(timeout=10)
        self._loop.call_soon_threadsafe(self._loop.stop)
        self._thread.join(timeout=10)


class TimedCloudPolicyHTTP:
    """Per-request HTTP POST transport — highest latency baseline."""

    def __init__(self, url: str, label: str = "cloud-http"):
        import httpx
        self.label = label
        self._url = url.rstrip("/") + "/infer"
        self._client = httpx.Client(timeout=120.0)
        print(f"  [{label}] HTTP endpoint: {self._url}")
        # warmup
        dummy = {"cam_high": np.zeros((3, 224, 224), dtype=np.uint8),
                 "state": np.zeros(14, dtype=np.float32)}
        self._timed_call(dummy)
        print(f"  [{label}] HTTP ready ✓")

    def _timed_call(self, wire_obs: dict):
        t0 = time.perf_counter()
        payload = pack_obs(wire_obs)
        t1 = time.perf_counter()
        resp = self._client.post(self._url, content=payload,
                                 headers={"Content-Type": "application/x-msgpack"})
        resp.raise_for_status()
        t2 = time.perf_counter()
        action, gpu_ms = unpack_response(resp.content)
        t3 = time.perf_counter()
        return (np.asarray(action),
                (t1 - t0) * 1000,
                max(0, (t2 - t1) * 1000 - gpu_ms),
                gpu_ms,
                (t3 - t2) * 1000)

    def infer(self, gym_obs: dict, step: int = 0) -> tuple[np.ndarray, CallTiming]:
        from openpi_client import image_tools
        t_start = time.perf_counter()
        img = gym_obs["pixels"]["top"]
        img = image_tools.convert_to_uint8(image_tools.resize_with_pad(img, 224, 224))
        wire_obs = {"cam_high": np.transpose(img, (2, 0, 1)),
                    "state": np.asarray(gym_obs["agent_pos"], dtype=np.float32)}
        preprocess_ms = (time.perf_counter() - t_start) * 1000
        actions, ser_ms, net_ms, gpu_ms, deser_ms = self._timed_call(wire_obs)
        total_ms = (time.perf_counter() - t_start) * 1000
        overhead = max(0, total_ms - preprocess_ms - ser_ms - net_ms - gpu_ms - deser_ms)
        return actions, CallTiming(step=step, total_ms=total_ms,
                                   preprocess_ms=preprocess_ms, gpu_ms=gpu_ms,
                                   serialize_ms=ser_ms, network_ms=net_ms,
                                   deserialize_ms=deser_ms, overhead_ms=overhead)

    def close(self):
        self._client.close()


class TimedCloudPolicyTCP:
    """Direct TCP via modal.forward() portal — Pi-style, lowest latency."""

    def __init__(self, portal_dict: str = "pi0-portal", label: str = "TCP"):
        import modal
        self.label = label
        print(f"  [{label}] Reading portal from Modal Dict …")
        d = modal.Dict.from_name(portal_dict)
        endpoint = d["tcp_endpoint"]
        host, port_str = endpoint.rsplit(":", 1)
        self._host, self._port = host, int(port_str)
        print(f"  [{label}] Connecting to {host}:{port_str} …")

        self._loop = asyncio.new_event_loop()
        self._reader: asyncio.StreamReader | None = None
        self._writer: asyncio.StreamWriter | None = None
        self._thread = threading.Thread(target=self._loop.run_forever, daemon=True)
        self._thread.start()
        asyncio.run_coroutine_threadsafe(self._connect(), self._loop).result(timeout=30)

        dummy = {"cam_high": np.zeros((3, 224, 224), dtype=np.uint8),
                 "state": np.zeros(14, dtype=np.float32)}
        self._timed_call(dummy)
        print(f"  [{label}] TCP ready ✓")

    async def _connect(self):
        self._reader, self._writer = await asyncio.open_connection(self._host, self._port)
        self._lock = asyncio.Lock()

    async def _send_recv_async(self, payload: bytes) -> bytes:
        async with self._lock:
            self._writer.write(len(payload).to_bytes(4, "big") + payload)
            await self._writer.drain()
            hdr = await self._reader.readexactly(4)
            return await self._reader.readexactly(int.from_bytes(hdr, "big"))

    def _timed_call(self, wire_obs: dict):
        t0 = time.perf_counter()
        payload = pack_obs(wire_obs)
        t1 = time.perf_counter()
        data = asyncio.run_coroutine_threadsafe(
            self._send_recv_async(payload), self._loop
        ).result(timeout=60)
        t2 = time.perf_counter()
        action, gpu_ms = unpack_response(data)
        t3 = time.perf_counter()
        return (np.asarray(action),
                (t1 - t0) * 1000,
                max(0, (t2 - t1) * 1000 - gpu_ms),
                gpu_ms,
                (t3 - t2) * 1000)

    def infer(self, gym_obs: dict, step: int = 0) -> tuple[np.ndarray, CallTiming]:
        from openpi_client import image_tools
        t_start = time.perf_counter()
        img = gym_obs["pixels"]["top"]
        img = image_tools.convert_to_uint8(image_tools.resize_with_pad(img, 224, 224))
        wire_obs = {"cam_high": np.transpose(img, (2, 0, 1)),
                    "state": np.asarray(gym_obs["agent_pos"], dtype=np.float32)}
        preprocess_ms = (time.perf_counter() - t_start) * 1000
        actions, ser_ms, net_ms, gpu_ms, deser_ms = self._timed_call(wire_obs)
        total_ms = (time.perf_counter() - t_start) * 1000
        overhead = max(0, total_ms - preprocess_ms - ser_ms - net_ms - gpu_ms - deser_ms)
        return actions, CallTiming(step=step, total_ms=total_ms,
                                   preprocess_ms=preprocess_ms, gpu_ms=gpu_ms,
                                   serialize_ms=ser_ms, network_ms=net_ms,
                                   deserialize_ms=deser_ms, overhead_ms=overhead)

    def close(self):
        async def _c():
            if self._writer:
                self._writer.close()
                await self._writer.wait_closed()
        asyncio.run_coroutine_threadsafe(_c(), self._loop).result(timeout=10)
        self._loop.call_soon_threadsafe(self._loop.stop)
        self._thread.join(timeout=10)


# ── video overlay ─────────────────────────────────────────────────────────────

def _overlay(frame: np.ndarray, label: str, step: int, total_steps: int,
             last_ms: float, success: bool, inferring: bool) -> np.ndarray:
    f = frame.copy()
    h, w = f.shape[:2]

    def box_text(text, x, y, scale=0.65, fg=(255, 255, 255), bg=(0, 0, 0)):
        font = cv2.FONT_HERSHEY_SIMPLEX
        thick = 2
        (tw, th), base = cv2.getTextSize(text, font, scale, thick)
        cv2.rectangle(f, (x - 4, y - th - 4), (x + tw + 4, y + base + 4), bg, -1)
        cv2.putText(f, text, (x, y), font, scale, fg, thick, cv2.LINE_AA)

    # Mode label (top-left)
    color = (0, 230, 118) if "LOCAL" in label else (64, 196, 255)
    box_text(label, 10, 28, scale=0.75, fg=color)

    # Step counter (top-right)
    step_txt = f"step {step:3d}/{total_steps}"
    box_text(step_txt, w - 180, 28)

    # Latency bar (bottom)
    if last_ms > 0:
        bar_w = int(min(last_ms / 800.0, 1.0) * (w - 20))
        bar_col = (0, 200, 80) if last_ms < 60 else (255, 150, 0) if last_ms < 300 else (220, 50, 50)
        cv2.rectangle(f, (10, h - 22), (10 + bar_w, h - 8), bar_col, -1)
        box_text(f"{last_ms:.0f} ms", 10, h - 28, scale=0.55)

    # Inferring indicator
    if inferring:
        box_text("INFERRING...", w // 2 - 80, h // 2, scale=0.8, fg=(255, 80, 80), bg=(30, 0, 0))

    # Success/fail badge
    if success:
        box_text("SUCCESS", w - 130, h - 28, scale=0.7, fg=(0, 255, 120))

    return f


# ── episode runner with recording ─────────────────────────────────────────────

def run_episode_record(env, policy, seed: int, label: str,
                       *, max_steps: int = 400, action_horizon: int = 10,
                       fps: int = FPS) -> dict:
    """Run one episode, recording at real wall-clock time.

    During inference the sim is frozen — we pad duplicate frames so the
    video shows the actual freeze duration.  This makes the latency difference
    between local (48ms) and cloud (330ms) immediately visible.
    """
    obs, _ = env.reset(seed=seed)
    action_chunk: Optional[np.ndarray] = None
    chunk_step = 0
    success = False
    frames: list[np.ndarray] = []
    timings: list[CallTiming] = []
    last_ms = 0.0

    ms_per_frame = 1000.0 / fps

    for step in range(max_steps):
        raw_frame = env.render()  # (H, W, 3)

        if action_chunk is None or chunk_step >= action_horizon:
            # Show a "INFERRING..." frozen frame, then block on the network/GPU call.
            frozen = _overlay(raw_frame, label, step, max_steps,
                              last_ms, success, inferring=True)
            frames.append(frozen)

            action_chunk, timing = policy.infer(obs, step=step)
            last_ms = timing.total_ms
            timings.append(timing)
            chunk_step = 0

            # Pad duplicate frames for the wall-clock time inference actually took.
            # e.g. 330ms at 30fps = 9 extra frozen frames — visually shows the freeze.
            extra = max(0, int(last_ms / ms_per_frame) - 1)
            for _ in range(extra):
                frames.append(frozen)
        else:
            overlay_frame = _overlay(raw_frame, label, step, max_steps,
                                     last_ms, success, inferring=False)
            frames.append(overlay_frame)

        action = action_chunk[chunk_step]
        chunk_step += 1
        obs, reward, terminated, truncated, _ = env.step(action)

        if reward > 0:
            success = True
        if terminated or truncated:
            break

    # One final frame
    raw_frame = env.render()
    frames.append(_overlay(raw_frame, label, step + 1, max_steps, last_ms, success, False))

    return {"frames": frames, "timings": timings, "success": success, "steps": step + 1}


# ── video writer ──────────────────────────────────────────────────────────────

def save_video(frames: list[np.ndarray], path: pathlib.Path, fps: int = FPS):
    path.parent.mkdir(parents=True, exist_ok=True)
    with imageio.get_writer(str(path), fps=fps, codec="h264", quality=8,
                            macro_block_size=None) as w:
        for f in frames:
            w.append_data(f)
    print(f"  Saved: {path}  ({len(frames)} frames, {len(frames)/fps:.1f}s)")


def make_comparison_video(local_episodes: list[dict], cloud_episodes: list[dict],
                          path: pathlib.Path, fps: int = FPS):
    """Stitch local (left) and cloud (right) frames side-by-side."""
    path.parent.mkdir(parents=True, exist_ok=True)
    all_local = [f for ep in local_episodes for f in ep["frames"]]
    all_cloud = [f for ep in cloud_episodes for f in ep["frames"]]

    # Pad shorter sequence by repeating its last frame
    n = max(len(all_local), len(all_cloud))
    if len(all_local) < n:
        all_local += [all_local[-1]] * (n - len(all_local))
    if len(all_cloud) < n:
        all_cloud += [all_cloud[-1]] * (n - len(all_cloud))

    h, w = all_local[0].shape[:2]
    divider = np.zeros((h, 4, 3), dtype=np.uint8)
    divider[:, :] = [255, 255, 255]

    with imageio.get_writer(str(path), fps=fps, codec="h264", quality=8,
                            macro_block_size=None) as w_vid:
        for fl, fc in zip(all_local, all_cloud):
            combined = np.concatenate([fl, divider, fc], axis=1)
            w_vid.append_data(combined)
    print(f"  Saved comparison: {path}  ({n} frames, {n/fps:.1f}s)")


# ── charts ────────────────────────────────────────────────────────────────────

def make_charts(results: dict[str, list[dict]], path: pathlib.Path):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import matplotlib.patches as mpatches
    from matplotlib.gridspec import GridSpec

    path.parent.mkdir(parents=True, exist_ok=True)

    # Flatten all timings per mode
    mode_timings: dict[str, list[CallTiming]] = {}
    for label, episodes in results.items():
        all_t = [t for ep in episodes for t in ep["timings"]]
        mode_timings[label] = all_t

    mode_colors = {}
    for label in results:
        if "LOCAL" in label:
            mode_colors[label] = C_LOCAL
        elif "H100" in label:
            mode_colors[label] = C_H100
        else:
            mode_colors[label] = C_A10G

    plt.style.use("dark_background")
    fig = plt.figure(figsize=(18, 14), facecolor="#0d1117")
    gs = GridSpec(3, 2, figure=fig, hspace=0.45, wspace=0.35,
                  left=0.07, right=0.97, top=0.93, bottom=0.07)

    ax_wf   = fig.add_subplot(gs[0, :])   # waterfall (full width)
    ax_bar  = fig.add_subplot(gs[1, 0])   # stacked bar
    ax_tl   = fig.add_subplot(gs[1, 1])   # per-call timeline
    ax_cdf  = fig.add_subplot(gs[2, 0])   # CDF
    ax_box  = fig.add_subplot(gs[2, 1])   # box plot

    fig.suptitle("π₀ ALOHA sim — Inference Latency Deep-Dive",
                 fontsize=18, color="white", fontweight="bold", y=0.97)

    # ── 1. WATERFALL: one sample call per mode ────────────────────────────────
    ax_wf.set_facecolor("#161b22")
    ax_wf.set_title("Inference Call Breakdown (one sample call per mode)",
                    color="white", fontsize=12, pad=8)

    yticks, ylabels = [], []
    for row_idx, (label, timings) in enumerate(mode_timings.items()):
        if not timings:
            continue
        # Pick median-ish call
        idx = len(timings) // 2
        t = timings[idx]
        y = row_idx * 1.4
        yticks.append(y + 0.35)
        ylabels.append(label)

        x = 0
        segments = []
        is_cloud = t.serialize_ms > 0

        if is_cloud:
            segments = [
                (t.preprocess_ms, C_PREP, "Preprocess"),
                (t.serialize_ms,  C_DESER, "Serialize"),
                (t.network_ms,    C_NET,   f"Network ({t.network_ms:.0f} ms)"),
                (t.gpu_ms,        C_GPU,   f"GPU ({t.gpu_ms:.0f} ms)"),
                (t.deserialize_ms,C_DESER, "Deserialize"),
                (t.overhead_ms,   "#555",  "Other"),
            ]
        else:
            segments = [
                (t.preprocess_ms, C_PREP, f"Preprocess ({t.preprocess_ms:.1f} ms)"),
                (t.gpu_ms,        C_GPU,  f"GPU ({t.gpu_ms:.0f} ms)"),
                (t.overhead_ms,   "#555", "Other"),
            ]

        for dur, color, lbl in segments:
            if dur > 0.5:
                ax_wf.barh(y, dur, left=x, height=0.7, color=color, alpha=0.9,
                           edgecolor="none")
                if dur > 15:
                    ax_wf.text(x + dur / 2, y + 0.35, lbl,
                               ha="center", va="center", fontsize=8,
                               color="white", fontweight="bold")
                x += dur

        # Total label
        ax_wf.text(x + 5, y + 0.35, f"  {t.total_ms:.0f} ms total",
                   va="center", fontsize=9, color=mode_colors[label])

    ax_wf.set_yticks(yticks)
    ax_wf.set_yticklabels(ylabels, color="white", fontsize=10)
    ax_wf.set_xlabel("Time (ms)", color="white")
    ax_wf.tick_params(colors="white")
    ax_wf.spines[:].set_color("#444")
    ax_wf.set_xlim(0, None)

    # legend for waterfall
    legend_patches = [
        mpatches.Patch(color=C_PREP,  label="Preprocess"),
        mpatches.Patch(color=C_DESER, label="Serialize / Deserialize"),
        mpatches.Patch(color=C_NET,   label="Network (round-trip)"),
        mpatches.Patch(color=C_GPU,   label="GPU compute"),
        mpatches.Patch(color="#555",  label="Other overhead"),
    ]
    ax_wf.legend(handles=legend_patches, loc="upper right",
                 facecolor="#161b22", edgecolor="#444", labelcolor="white",
                 fontsize=8)

    # ── 2. STACKED BAR: mean breakdown per mode ────────────────────────────────
    ax_bar.set_facecolor("#161b22")
    ax_bar.set_title("Mean Latency Breakdown", color="white", fontsize=11, pad=8)

    bar_labels, bar_data = [], []
    for label, timings in mode_timings.items():
        if not timings:
            continue
        bar_labels.append(label)
        t_arr = [asdict(t) for t in timings]
        bar_data.append({
            "preprocess": np.mean([t["preprocess_ms"] for t in t_arr]),
            "serialize":  np.mean([t["serialize_ms"] for t in t_arr]),
            "network":    np.mean([t["network_ms"] for t in t_arr]),
            "gpu":        np.mean([t["gpu_ms"] for t in t_arr]),
            "deserialize":np.mean([t["deserialize_ms"] for t in t_arr]),
            "overhead":   np.mean([t["overhead_ms"] for t in t_arr]),
        })

    x_pos = np.arange(len(bar_labels))
    components = [
        ("preprocess", C_PREP, "Preprocess"),
        ("serialize",  C_DESER, "Serialize"),
        ("network",    C_NET,  "Network"),
        ("gpu",        C_GPU,  "GPU"),
        ("deserialize",C_DESER, "Deserialize"),
        ("overhead",   "#555", "Other"),
    ]
    bottoms = np.zeros(len(bar_labels))
    for key, color, lbl in components:
        vals = np.array([d[key] for d in bar_data])
        bars = ax_bar.bar(x_pos, vals, bottom=bottoms, color=color, label=lbl,
                          alpha=0.9, edgecolor="none", width=0.55)
        bottoms += vals

    # Total labels on top
    for i, (lbl, total) in enumerate(zip(bar_labels, bottoms)):
        ax_bar.text(i, total + 5, f"{total:.0f} ms", ha="center", va="bottom",
                    color="white", fontsize=9, fontweight="bold")

    ax_bar.set_xticks(x_pos)
    ax_bar.set_xticklabels([l.replace(" ", "\n") for l in bar_labels],
                           color="white", fontsize=8)
    ax_bar.set_ylabel("ms", color="white")
    ax_bar.tick_params(colors="white")
    ax_bar.spines[:].set_color("#444")
    ax_bar.legend(facecolor="#161b22", edgecolor="#444", labelcolor="white",
                  fontsize=7, loc="upper left")

    # ── 3. PER-CALL TIMELINE ──────────────────────────────────────────────────
    ax_tl.set_facecolor("#161b22")
    ax_tl.set_title("Inference Latency per Call (all episodes)", color="white",
                    fontsize=11, pad=8)

    for label, timings in mode_timings.items():
        if not timings:
            continue
        xs = list(range(len(timings)))
        ys = [t.total_ms for t in timings]
        ax_tl.plot(xs, ys, color=mode_colors[label], linewidth=1.5,
                   marker="o", markersize=3, label=label, alpha=0.85)

    ax_tl.set_xlabel("Inference call #", color="white")
    ax_tl.set_ylabel("ms", color="white")
    ax_tl.tick_params(colors="white")
    ax_tl.spines[:].set_color("#444")
    ax_tl.legend(facecolor="#161b22", edgecolor="#444", labelcolor="white", fontsize=8)

    # ── 4. CDF ────────────────────────────────────────────────────────────────
    ax_cdf.set_facecolor("#161b22")
    ax_cdf.set_title("Latency CDF", color="white", fontsize=11, pad=8)

    for label, timings in mode_timings.items():
        if not timings:
            continue
        vals = sorted([t.total_ms for t in timings])
        cdf = np.arange(1, len(vals) + 1) / len(vals)
        ax_cdf.plot(vals, cdf, color=mode_colors[label], linewidth=2, label=label)
        # Mark p50, p95
        p50 = float(np.percentile(vals, 50))
        p95 = float(np.percentile(vals, 95))
        ax_cdf.axvline(p50, color=mode_colors[label], linestyle="--", alpha=0.4, linewidth=1)
        ax_cdf.text(p50 + 3, 0.55, f"p50={p50:.0f}", color=mode_colors[label],
                    fontsize=7, alpha=0.9)

    ax_cdf.set_xlabel("ms", color="white")
    ax_cdf.set_ylabel("Cumulative fraction", color="white")
    ax_cdf.tick_params(colors="white")
    ax_cdf.spines[:].set_color("#444")
    ax_cdf.legend(facecolor="#161b22", edgecolor="#444", labelcolor="white", fontsize=8)
    ax_cdf.set_ylim(0, 1.05)

    # ── 5. BOX PLOT ───────────────────────────────────────────────────────────
    ax_box.set_facecolor("#161b22")
    ax_box.set_title("Latency Distribution", color="white", fontsize=11, pad=8)

    box_data, box_labels, box_colors = [], [], []
    for label, timings in mode_timings.items():
        if not timings:
            continue
        box_data.append([t.total_ms for t in timings])
        box_labels.append(label)
        box_colors.append(mode_colors[label])

    bp = ax_box.boxplot(box_data, patch_artist=True, widths=0.5,
                        medianprops=dict(color="white", linewidth=2),
                        whiskerprops=dict(color="#aaa"),
                        capprops=dict(color="#aaa"),
                        flierprops=dict(marker="x", color="#888", markersize=4))
    for patch, color in zip(bp["boxes"], box_colors):
        patch.set_facecolor(color)
        patch.set_alpha(0.7)

    ax_box.set_xticks(range(1, len(box_labels) + 1))
    ax_box.set_xticklabels([l.replace(" ", "\n") for l in box_labels],
                           color="white", fontsize=8)
    ax_box.set_ylabel("ms", color="white")
    ax_box.tick_params(colors="white")
    ax_box.spines[:].set_color("#444")

    fig.savefig(str(path), dpi=160, bbox_inches="tight", facecolor="#0d1117")
    plt.close(fig)
    print(f"  Saved chart: {path}")


# ── main ──────────────────────────────────────────────────────────────────────

def main():
    if "MUJOCO_GL" not in os.environ:
        os.environ["MUJOCO_GL"] = "egl"
        print("INFO: MUJOCO_GL not set, defaulting to egl")

    parser = argparse.ArgumentParser(description="Record π₀ eval videos + latency charts")
    parser.add_argument("--url", help="Modal server URL (omit for local-only)")
    parser.add_argument("--gpu-label", default="H100",
                        help="GPU label for cloud mode (e.g. H100, A10G) [default: H100]")
    parser.add_argument("--episodes", type=int, default=3,
                        help="Episodes per mode [default: 3]")
    parser.add_argument("--max-steps", type=int, default=400)
    parser.add_argument("--action-horizon", type=int, default=10)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--fps", type=int, default=FPS,
                        help="Video output FPS [default: 30]")
    parser.add_argument("--skip-local", action="store_true",
                        help="Skip local inference (cloud only)")
    parser.add_argument("--transport", choices=["ws", "tcp", "http"], default="ws",
                        help="Cloud transport: ws (WebSocket) | tcp (Modal portal) | http [default: ws]")
    args = parser.parse_args()

    import gymnasium as gym
    import gym_aloha  # noqa

    all_results: dict[str, list[dict]] = {}

    # ── LOCAL ─────────────────────────────────────────────────────────────────
    if not args.skip_local:
        print("\n═══ LOCAL H100 ═══")
        local_policy = TimedLocalPolicy()
        label_local = "LOCAL H100"
        local_episodes = []

        env = gym.make(TASK, obs_type="pixels_agent_pos", render_mode="rgb_array")
        for ep in range(args.episodes):
            seed = args.seed + ep
            print(f"  Episode {ep+1}/{args.episodes} (seed={seed}) …")
            result = run_episode_record(
                env, local_policy, seed, label_local,
                max_steps=args.max_steps, action_horizon=args.action_horizon,
                fps=args.fps,
            )
            status = "✓" if result["success"] else "✗"
            ms = np.median([t.total_ms for t in result["timings"]])
            print(f"    {status}  steps={result['steps']}  infer={ms:.0f}ms median  "
                  f"frames={len(result['frames'])}")
            local_episodes.append(result)
        env.close()

        all_results[label_local] = local_episodes
        print("  Saving local video …")
        all_local_frames = [f for ep in local_episodes for f in ep["frames"]]
        save_video(all_local_frames, VIDEOS_DIR / "local_h100.mp4", fps=args.fps)

    # ── CLOUD ─────────────────────────────────────────────────────────────────
    if args.url or args.transport == "tcp":
        transport_tag = args.transport.upper()
        print(f"\n═══ CLOUD {args.gpu_label} [{transport_tag}] ═══")
        if args.transport == "tcp":
            cloud_policy = TimedCloudPolicyTCP(label=f"Modal {args.gpu_label} TCP")
        elif args.transport == "http":
            cloud_policy = TimedCloudPolicyHTTP(args.url, label=f"Modal {args.gpu_label}")
        else:
            cloud_policy = TimedCloudPolicyWS(args.url, label=f"Modal {args.gpu_label}")
        label_cloud = f"CLOUD {args.gpu_label} {transport_tag}"
        cloud_episodes = []

        env = gym.make(TASK, obs_type="pixels_agent_pos", render_mode="rgb_array")
        try:
            for ep in range(args.episodes):
                seed = args.seed + ep
                print(f"  Episode {ep+1}/{args.episodes} (seed={seed}) …")
                result = run_episode_record(
                    env, cloud_policy, seed, label_cloud,
                    max_steps=args.max_steps, action_horizon=args.action_horizon,
                    fps=args.fps,
                )
                status = "✓" if result["success"] else "✗"
                ms = np.median([t.total_ms for t in result["timings"]])
                print(f"    {status}  steps={result['steps']}  infer={ms:.0f}ms median  "
                      f"frames={len(result['frames'])}")
                cloud_episodes.append(result)
        finally:
            cloud_policy.close()
        env.close()

        all_results[label_cloud] = cloud_episodes
        slug = args.gpu_label.lower().replace(" ", "_")
        transport_slug = args.transport
        print("  Saving cloud video …")
        all_cloud_frames = [f for ep in cloud_episodes for f in ep["frames"]]
        save_video(all_cloud_frames, VIDEOS_DIR / f"cloud_{slug}_{transport_slug}.mp4", fps=args.fps)

        if not args.skip_local and "LOCAL H100" in all_results:
            print("  Saving comparison video …")
            make_comparison_video(
                all_results["LOCAL H100"],
                cloud_episodes,
                VIDEOS_DIR / f"comparison_local_vs_cloud_{slug}_{transport_slug}.mp4",
                fps=args.fps,
            )

    # ── CHARTS ────────────────────────────────────────────────────────────────
    if all_results:
        print("\n═══ CHARTS ═══")
        make_charts(all_results, CHARTS_DIR / "latency_breakdown.png")

        # Print summary
        print("\n  Summary:")
        for label, episodes in all_results.items():
            timings = [t for ep in episodes for t in ep["timings"]]
            sr = sum(1 for ep in episodes if ep["success"]) / len(episodes)
            p50 = np.median([t.total_ms for t in timings])
            p95 = np.percentile([t.total_ms for t in timings], 95)
            print(f"    {label:<20}  SR={sr*100:.0f}%  "
                  f"p50={p50:.0f}ms  p95={p95:.0f}ms")

    print("\nDone.")
    print(f"  Videos: {VIDEOS_DIR}/")
    print(f"  Charts: {CHARTS_DIR}/")


if __name__ == "__main__":
    main()