Replace the page with a Rerun-only viewer

Drops the composed-video, frame-stepper and action-plot tabs. The page is now
an episode list with search, sort and PSNR/MSE, next to an embedded Rerun
viewer showing ground truth and generated side by side on one frame timeline.

Starts empty and assumes no dataset: the repo is typed in, nothing is loaded
until Load is pressed.

Two viewer bugs fixed along the way: ssr_mode must be off, since hydration
re-mounts the wasm canvas, and .change() defaults to trigger_mode
"always_last", which let two runs write into one viewer concurrently.

Pins gradio 6.5.1 to match rerun/gradio-rerun-viewer, the reference Space
known to work with gradio_rerun 0.35.0.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

README.md

@@ -4,7 +4,7 @@ emoji: 🎬
 colorFrom: green
 colorTo: purple
 sdk: gradio
-sdk_version: 6.22.0
+sdk_version: 6.5.1
 python_version: '3.12'
 app_file: app.py
 pinned: false

app.py

@@ -1,5 +1,5 @@
 """
-Side-by-side viewer for forward-dynamics rollouts vs. ground truth.
+Rerun viewer for forward-dynamics rollouts vs. ground truth.
 
 Point it at any HF dataset repo laid out as:
 
@@ -7,20 +7,17 @@ Point it at any HF dataset repo laid out as:
     metrics_summary.json                # optional: per-window MSE / PSNR
     windows/<episode_id>/gt.mp4         # ground-truth window
     windows/<episode_id>/generated.mp4  # model rollout for the same slice
-    windows/<episode_id>/actions.json   # optional: [T, D] driving actions
-    windows/<episode_id>/meta.json      # optional: per-window metadata
 
-Only `windows/<id>/{gt,generated}.mp4` is required. Without a manifest the
-window list is discovered by listing `windows/` one level deep. Videos may
-stack N cameras vertically; N is inferred from the frame aspect ratio.
+Pick an episode and its two videos stream into an embedded Rerun viewer, side
+by side on a shared frame timeline.
 
 Nothing is downloaded at build time and no repo is ever pulled whole. Each
-browser session gets its own scratch directory; window files arrive only when
-that window is opened, at most MAX_WINDOWS_ON_DISK are kept, and the whole
-directory is deleted when the session ends.
+browser session gets its own scratch directory, keeps at most
+MAX_WINDOWS_ON_DISK windows, and is deleted when the session ends. Recordings
+go straight to the viewer as a binary stream, so no .rrd files are written.
 
-DATASET_REPO sets the repo prefilled at startup. Private repos need an
-HF_TOKEN secret with read access.
+The app starts empty and assumes no dataset: type a repo id and press Load.
+Private repos need an HF_TOKEN secret with read access.
 """
 
 from __future__ import annotations
@@ -38,15 +35,12 @@ from pathlib import Path
 
 import cv2
 import gradio as gr
-import imageio_ffmpeg
 import numpy as np
 import rerun as rr
 import rerun.blueprint as rrb
 from gradio_rerun import Rerun
 from huggingface_hub import HfApi, RepoFolder, hf_hub_download
-from matplotlib.figure import Figure
 
-DEFAULT_REPO = os.environ.get("DATASET_REPO", "OneWorld-AI/abc-critical-window-rollout")
 HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or None
 
 BASE_DIR = Path(tempfile.gettempdir()) / "wmviz"
@@ -54,29 +48,14 @@ BASE_DIR.mkdir(parents=True, exist_ok=True)
 
 MAX_WINDOWS_ON_DISK = int(os.environ.get("MAX_WINDOWS_ON_DISK", 12))
 MAX_PAIRS_IN_RAM = 2
-MAX_COMPOSED_IN_RAM = 1
 MAX_DISCOVER = int(os.environ.get("MAX_DISCOVER", 5000))
 SESSION_TTL = int(os.environ.get("SESSION_TTL", 3600))
-
-# Conventional names when a video stacks exactly three cameras.
-THREE_CAM_NAMES = ["Top", "Left wrist", "Right wrist"]
-ALL_CAMS = -1
-NO_CRIT = -1
-
-FONT = cv2.FONT_HERSHEY_SIMPLEX
-HDR, FTR, SEP = 26, 26, 4
-SCALE = 2  # source frames are small; upscale so overlays stay legible
-JPEG_QUALITY = 90  # raw frames make a ~16 MB recording; JPEG brings it to ~1.5 MB
-
-GT_COLOR = (130, 225, 140)
-GEN_COLOR = (120, 180, 255)
-DIFF_COLOR = (215, 215, 215)
-CRIT_COLOR = (255, 95, 95)
+JPEG_QUALITY = 90
 
 SORTS = [
-    "Hardest first (critical PSNR ↑)",
-    "Easiest first (critical PSNR ↓)",
-    "Task A→Z",
+    "Hardest first (critical PSNR up)",
+    "Easiest first (critical PSNR down)",
+    "Task A-Z",
     "Manifest order",
 ]
 
@@ -109,15 +88,10 @@ class Session:
         self.repo = ""
         self.revision: str | None = None
         self.rows: list[dict] = []
-        self.note = ""
         self._on_disk: OrderedDict[str, bool] = OrderedDict()
         self._pairs: OrderedDict[str, tuple[np.ndarray, np.ndarray]] = OrderedDict()
-        self._composed: OrderedDict[tuple, np.ndarray] = OrderedDict()
-        self._meta: dict[str, dict] = {}
         self.touched = time.time()
 
-    # -- fetching ----------------------------------------------------------
-
     def fetch(self, rel: str, optional: bool = False) -> str | None:
         """Download one repo file into this session's directory."""
         self.touched = time.time()
@@ -149,16 +123,12 @@ class Session:
             old, _ = self._on_disk.popitem(last=False)
             shutil.rmtree(self.data / "windows" / old, ignore_errors=True)
             self._pairs.pop(old, None)
-            for key in [k for k in self._composed if k[0] == old]:
-                self._composed.pop(key, None)
             for f in self.render.glob(f"{old}_*"):
                 f.unlink(missing_ok=True)
 
     def disk_windows(self) -> int:
         return len(self._on_disk)
 
-    # -- bounded RAM caches ------------------------------------------------
-
     def pair(self, episode_id: str) -> tuple[np.ndarray, np.ndarray]:
         hit = self._pairs.get(episode_id)
         if hit is not None:
@@ -173,29 +143,6 @@ class Session:
             self._pairs.popitem(last=False)
         return value
 
-    def composed(self, episode_id: str, view: int, show_diff: bool, crit: int) -> np.ndarray:
-        key = (episode_id, view, show_diff, crit)
-        hit = self._composed.get(key)
-        if hit is not None:
-            self._composed.move_to_end(key)
-            return hit
-        value = compose(self.pair(episode_id), view, show_diff, crit)
-        self._composed[key] = value
-        while len(self._composed) > MAX_COMPOSED_IN_RAM:
-            self._composed.popitem(last=False)
-        return value
-
-    def meta(self, episode_id: str) -> dict:
-        if episode_id not in self._meta:
-            path = self.window_file(episode_id, "meta.json", optional=True)
-            try:
-                self._meta[episode_id] = json.load(open(path)) if path else {}
-            except Exception:
-                self._meta[episode_id] = {}
-        return self._meta[episode_id]
-
-    # -- lifecycle ---------------------------------------------------------
-
     def set_repo(self, repo: str, revision: str | None) -> None:
         if (repo, revision) != (self.repo, self.revision):
             self.reset()
@@ -204,8 +151,6 @@ class Session:
     def reset(self) -> None:
         self._on_disk.clear()
         self._pairs.clear()
-        self._composed.clear()
-        self._meta.clear()
         self.rows = []
         for d in (self.data, self.render):
             shutil.rmtree(d, ignore_errors=True)
@@ -213,7 +158,6 @@ class Session:
 
     def close(self) -> None:
         self._pairs.clear()
-        self._composed.clear()
         shutil.rmtree(self.root, ignore_errors=True)
 
 
@@ -249,7 +193,7 @@ def get_session(sid: str | None) -> Session:
 
 
 def release(sid: str | None) -> None:
-    """gr.State delete_callback — fires when the session ends or its TTL expires."""
+    """gr.State delete_callback - fires when the session ends or its TTL expires."""
     with _LOCK:
         s = SESSIONS.pop(sid or "", None)
     if s is not None:
@@ -266,11 +210,8 @@ def _row(order: int, episode_id: str, w: dict, m: dict) -> dict:
         "task": w.get("task") or "?",
         "edge": w.get("edge_type") or "?",
         "arm": w.get("arm") or "?",
-        "amplitude": w.get("amplitude"),
-        "start_frame": w.get("start_frame"),
-        "window": w.get("window"),
         "critical_frame": w.get("critical_frame"),
-        "n_frames": w.get("n_frames"),
+        "start_frame": w.get("start_frame"),
         "full_psnr": m.get("full_psnr"),
         "crit_psnr": m.get("crit_psnr"),
         "full_mse": m.get("full_mse"),
@@ -330,7 +271,7 @@ def build_index(session: Session) -> str:
         eps, truncated = discover_episodes(session)
         if not eps:
             raise FileNotFoundError(
-                "no windows/<episode_id>/ directories found — is this the right repo?"
+                "no windows/<episode_id>/ directories found - is this the right repo?"
             )
         session.rows = [_row(i, ep, {}, by_id.get(ep, {})) for i, ep in enumerate(eps)]
         note = "directory listing (no `windows_manifest.json`)"
@@ -351,18 +292,14 @@ def label_for(row: dict) -> str:
     head = f"{psnr:5.1f} dB" if isinstance(psnr, (int, float)) else "  -- dB"
     tags = " ".join(t for t in (row["edge"], row["arm"]) if t and t != "?")
     parts = [head] + ([tags] if tags else []) + [row["task"], short_id(row["episode_id"])]
-    return " · ".join(parts)
+    return " | ".join(parts)
 
 
-def filter_rows(session: Session, query: str, edge: str, arm: str, sort: str) -> list[dict]:
+def filter_rows(session: Session, query: str, sort: str) -> list[dict]:
     rows = session.rows
     q = (query or "").strip().lower()
     if q:
         rows = [r for r in rows if q in r["task"].lower() or q in r["episode_id"].lower()]
-    if edge != "any":
-        rows = [r for r in rows if r["edge"] == edge]
-    if arm != "any":
-        rows = [r for r in rows if r["arm"] == arm]
 
     def psnr(r, missing):
         v = r["crit_psnr"]
@@ -380,19 +317,8 @@ def filter_rows(session: Session, query: str, edge: str, arm: str, sort: str) ->
 def row_by_id(session: Session, episode_id: str) -> dict:
     for r in session.rows:
         if r["episode_id"] == episode_id:
-            row = dict(r)
-            break
-    else:
-        row = _row(0, episode_id, {}, {})
-    # backfill whatever the manifest did not carry from the window's own meta.json
-    if row["task"] == "?" or row["critical_frame"] is None:
-        meta = session.meta(episode_id)
-        if meta:
-            merged = _row(row["order"], episode_id, meta, {})
-            for k, v in merged.items():
-                if row.get(k) in (None, "?"):
-                    row[k] = v
-    return row
+            return r
+    return _row(0, episode_id, {}, {})
 
 
 # --------------------------------------------------------------------------- video
@@ -413,326 +339,41 @@ def read_video(path: str) -> np.ndarray:
     return np.stack(frames)
 
 
-def n_cameras(frames: np.ndarray) -> int:
-    """How many roughly-square cameras are stacked vertically in the frame."""
-    h, w = frames.shape[1], frames.shape[2]
-    return max(1, min(6, int(round(h / w))))
-
-
-def camera_choices(n: int) -> list[tuple[str, int]]:
-    choices = [(name, i) for i, name in enumerate(camera_names(n))]
-    return ([("All cameras", ALL_CAMS)] + choices) if n > 1 else choices
-
-
-def camera_names(n: int) -> list[str]:
-    return THREE_CAM_NAMES if n == 3 else [f"Camera {i + 1}" for i in range(n)]
-
-
-def selected_cameras(n: int, view: int) -> list[tuple[int, str, str]]:
-    """(band index, entity-path key, display label) for the cameras `view` selects."""
-    names = camera_names(n)
-    idxs = range(n) if view == ALL_CAMS or not 0 <= view < n else [view]
-    return [(i, re.sub(r"[^a-z0-9]+", "_", names[i].lower()), names[i]) for i in idxs]
-
-
-def cam_slices(frames: np.ndarray, view: int) -> list[np.ndarray]:
-    n = n_cameras(frames)
-    h = frames.shape[1] // n
-    rows = range(n) if view == ALL_CAMS or not 0 <= view < n else [view]
-    return [frames[:, i * h : (i + 1) * h] for i in rows]
-
-
-def letterbox_rows(*bands: np.ndarray, thresh: int = 16) -> tuple[int, int]:
-    """Rows to keep after dropping the black bars shared by every given band."""
-    h = bands[0].shape[1]
-    lit = np.zeros(h, bool)
-    for b in bands:
-        lit |= b.reshape(b.shape[0], h, -1).max(axis=(0, 2)) > thresh
-    idx = np.flatnonzero(lit)
-    if idx.size < h * 0.1:  # nothing meaningful to trim
-        return 0, h
-    return int(idx[0]), int(idx[-1]) + 1
-
-
-def crop_view(gt: np.ndarray, gen: np.ndarray, view: int) -> tuple[np.ndarray, np.ndarray]:
-    """Select cameras and trim the letterbox, identically for both videos."""
-    gt_out, gen_out = [], []
-    for g, p in zip(cam_slices(gt, view), cam_slices(gen, view)):
-        lo, hi = letterbox_rows(g, p)
-        gt_out.append(g[:, lo:hi])
-        gen_out.append(p[:, lo:hi])
-    return np.concatenate(gt_out, axis=1), np.concatenate(gen_out, axis=1)
-
-
-def upscale(frames: np.ndarray, scale: int) -> np.ndarray:
-    if scale == 1:
-        return frames
-    return np.stack(
-        [cv2.resize(f, None, fx=scale, fy=scale, interpolation=cv2.INTER_NEAREST) for f in frames]
-    )
-
-
-def panel(frames: np.ndarray, title: str, color: tuple[int, int, int]) -> np.ndarray:
-    t, h, w, _ = frames.shape
-    out = np.zeros((t, HDR + h, w, 3), np.uint8)
-    out[:, HDR:] = frames
-    bar = np.full((HDR, w, 3), 22, np.uint8)
-    cv2.putText(bar, title, (8, HDR - 8), FONT, 0.52, color, 1, cv2.LINE_AA)
-    out[:, :HDR] = bar
-    return out
-
-
-def diff_frames(gt: np.ndarray, gen: np.ndarray, gain: float = 3.0) -> np.ndarray:
-    d = np.abs(gt.astype(np.int16) - gen.astype(np.int16)).mean(axis=-1)
-    d = np.clip(d * gain, 0, 255).astype(np.uint8)
-    return np.stack([cv2.applyColorMap(f, cv2.COLORMAP_INFERNO)[:, :, ::-1] for f in d])
-
-
-def compose(pair: tuple[np.ndarray, np.ndarray], view: int, show_diff: bool, crit_idx: int) -> np.ndarray:
-    """Build the labelled side-by-side strip: [T, H, W, 3] RGB."""
-    gt_c, gen_c = crop_view(pair[0], pair[1], view)
-    gt_v, gen_v = upscale(gt_c, SCALE), upscale(gen_c, SCALE)
-
-    panels = [panel(gt_v, "GROUND TRUTH", GT_COLOR), panel(gen_v, "GENERATED", GEN_COLOR)]
-    if show_diff:
-        panels.append(panel(diff_frames(gt_v, gen_v), "|DIFF| x3", DIFF_COLOR))
-
-    t, ph, _, _ = panels[0].shape
-    gap = np.zeros((t, ph, SEP, 3), np.uint8)
-    stitched = [panels[0]]
-    for p in panels[1:]:
-        stitched += [gap, p]
-    comp = np.concatenate(stitched, axis=2)
-
-    t, h, w, _ = comp.shape
-    out = np.zeros((t, h + FTR, w, 3), np.uint8)
-    out[:, :h] = comp
-    has_crit = 0 <= crit_idx < t
-    tl0, tl1, y = int(w * 0.45), w - 14, FTR // 2
-    for i in range(t):
-        bar = np.full((FTR, w, 3), 22, np.uint8)
-        is_crit = has_crit and i == crit_idx
-        text = f"frame {i:02d}/{t - 1}" + ("   CRITICAL FRAME" if is_crit else "")
-        cv2.putText(bar, text, (8, FTR - 8), FONT, 0.5, CRIT_COLOR if is_crit else (200, 200, 200), 1, cv2.LINE_AA)
-        cv2.line(bar, (tl0, y), (tl1, y), (70, 70, 70), 2, cv2.LINE_AA)
-        if has_crit:
-            cx = int(tl0 + (tl1 - tl0) * crit_idx / max(t - 1, 1))
-            cv2.line(bar, (cx, y - 7), (cx, y + 7), CRIT_COLOR, 2, cv2.LINE_AA)
-        px = int(tl0 + (tl1 - tl0) * i / max(t - 1, 1))
-        cv2.circle(bar, (px, y), 4, (255, 255, 255), -1, cv2.LINE_AA)
-        out[i, h:] = bar
-        if is_crit:
-            cv2.rectangle(out[i], (0, 0), (w - 1, h + FTR - 1), CRIT_COLOR, 3)
-
-    # yuv420p needs even dimensions
-    ph, pw = out.shape[1] % 2, out.shape[2] % 2
-    if ph or pw:
-        out = np.pad(out, ((0, 0), (0, ph), (0, pw), (0, 0)))
-    return out
-
-
-def write_mp4(frames: np.ndarray, path: Path, fps: float) -> str:
-    _, h, w, _ = frames.shape
-    writer = imageio_ffmpeg.write_frames(
-        str(path),
-        size=(w, h),
-        fps=fps,
-        codec="libx264",
-        pix_fmt_in="rgb24",
-        pix_fmt_out="yuv420p",
-        macro_block_size=1,
-        output_params=["-crf", "16", "-preset", "veryfast", "-movflags", "+faststart"],
-    )
-    writer.send(None)
-    for f in frames:
-        writer.send(np.ascontiguousarray(f).tobytes())
-    writer.close()
-    return str(path)
-
-
-def render_video(session: Session, episode_id: str, view: int, show_diff: bool, fps: float, crit: int) -> str:
-    path = session.render / f"{episode_id}_{view}_{int(show_diff)}_{fps:g}.mp4"
-    if not path.exists():
-        write_mp4(session.composed(episode_id, view, show_diff, crit), path, fps)
-    return str(path)
-
-
-# --------------------------------------------------------------------------- actions plot
-
-
-def load_actions(session: Session, episode_id: str) -> np.ndarray | None:
-    path = session.window_file(episode_id, "actions.json", optional=True)
-    if not path:
-        return None
-    try:
-        return np.atleast_2d(np.asarray(json.load(open(path)), dtype=float))
-    except Exception:
-        return None
-
-
-def plot_actions(session: Session, episode_id: str, crit_idx: int, cursor: int | None) -> Figure | None:
-    actions = load_actions(session, episode_id)
-    if actions is None:
-        return None
-
-    t = np.arange(len(actions))
-    dim = actions.shape[1]
-    if dim == 14:  # bimanual: [L arm 6, L grip, R arm 6, R grip]
-        specs = [
-            ("Grippers (1.0 = open, 0.0 = closed)", [(6, "L gripper"), (13, "R gripper")]),
-            ("Left arm joints", [(i, f"L{i}") for i in range(6)]),
-            ("Right arm joints", [(i, f"R{i - 7}") for i in range(7, 13)]),
-        ]
-    else:
-        specs = [(f"Action channels [{len(actions)} x {dim}]", [(i, f"a{i}") for i in range(dim)])]
-
-    fig = Figure(figsize=(7.2, 1.8 * len(specs) + 0.6), dpi=110, layout="constrained")
-    axes = np.atleast_1d(fig.subplots(len(specs), 1, sharex=True))
-    for ax, (title, channels) in zip(axes, specs):
-        for idx, name in channels:
-            ax.plot(t, actions[:, idx], lw=1.4, label=name)
-        ax.set_title(title, fontsize=9, loc="left")
-        if crit_idx >= 0:
-            ax.axvline(crit_idx, color="crimson", ls="--", lw=1.2, zorder=0)
-        if cursor is not None:
-            ax.axvline(cursor, color="0.35", ls=":", lw=1.0, zorder=0)
-        ax.grid(alpha=0.25, lw=0.5)
-        ax.tick_params(labelsize=8)
-        if len(channels) <= 8:
-            ax.legend(fontsize=7, ncol=6, loc="upper right", framealpha=0.7)
-    tail = " (dashed red = critical frame)" if crit_idx >= 0 else ""
-    axes[-1].set_xlabel(f"frame in window{tail}", fontsize=8)
-    return fig
-
-
-# --------------------------------------------------------------------------- rerun recording
-
-
-def action_channels(dim: int) -> list[tuple[str, int]]:
-    """(entity path, column) for each action channel."""
-    if dim == 14:  # bimanual: [L arm 6, L grip, R arm 6, R grip]
-        return (
-            [("actions/gripper/L", 6), ("actions/gripper/R", 13)]
-            + [(f"actions/arm_L/j{i}", i) for i in range(6)]
-            + [(f"actions/arm_R/j{i - 7}", i) for i in range(7, 13)]
-        )
-    return [(f"actions/a{i}", i) for i in range(dim)]
-
-
-def blueprint_for(cams: list[tuple[int, str, str]], show_diff: bool, actions_origin: str | None):
-    """One row per camera (GT | generated | diff), with the actions plot underneath."""
-    rows, shares = [], []
-    for _, key, label in cams:
-        views = [
-            rrb.Spatial2DView(origin=f"gt/{key}", name=f"GT · {label}"),
-            rrb.Spatial2DView(origin=f"generated/{key}", name=f"Generated · {label}"),
-        ]
-        if show_diff:
-            views.append(rrb.Spatial2DView(origin=f"diff/{key}", name=f"|diff| · {label}"))
-        rows.append(rrb.Horizontal(*views))
-        shares.append(3)
-    if actions_origin:
-        rows.append(rrb.TimeSeriesView(origin=actions_origin, name="Driving actions"))
-        shares.append(2)
-    return rrb.Blueprint(rrb.Vertical(*rows, row_shares=shares), collapse_panels=True)
-
-
-def build_rrd(session: Session, episode_id: str, view: int, show_diff: bool, crit: int) -> str:
-    """Log the window to a .rrd: every stream on one shared frame timeline."""
-    path = session.render / f"{episode_id}_{view}_{int(show_diff)}.rrd"
-    if path.exists():
-        return str(path)
-
-    gt, gen = session.pair(episode_id)
-    n = n_cameras(gt)
-    band = gt.shape[1] // n
-    cams = selected_cameras(n, view)
-    actions = load_actions(session, episode_id)
-    channels = action_channels(actions.shape[1]) if actions is not None else []
-    actions_origin = None
-    if actions is not None:
-        actions_origin = "actions/gripper" if actions.shape[1] == 14 else "actions"
-
-    rec = rr.RecordingStream("wmviz", recording_id=path.stem)
-    for i in range(len(gt)):
-        rr.set_time("frame", sequence=i, recording=rec)
-        for c, key, _ in cams:
-            g = gt[i, c * band : (c + 1) * band]
-            p = gen[i, c * band : (c + 1) * band]
-            rr.log(f"gt/{key}", rr.Image(g).compress(jpeg_quality=JPEG_QUALITY), recording=rec)
-            rr.log(f"generated/{key}", rr.Image(p).compress(jpeg_quality=JPEG_QUALITY), recording=rec)
-            if show_diff:
-                d = diff_frames(g[None], p[None])[0]
-                rr.log(f"diff/{key}", rr.Image(d).compress(jpeg_quality=JPEG_QUALITY), recording=rec)
-        if actions is not None and i < len(actions):
-            for entity, col in channels:
-                rr.log(entity, rr.Scalars(float(actions[i, col])), recording=rec)
-        if i == crit:
-            rr.log("critical_frame", rr.TextLog("critical frame"), recording=rec)
-
-    rec.save(str(path), default_blueprint=blueprint_for(cams, show_diff, actions_origin))
-    return str(path)
-
-
 # --------------------------------------------------------------------------- ui glue
 
 
-def metrics_md(session: Session, row: dict) -> str:
+def metrics_md(row: dict) -> str:
     def fmt(v, unit=""):
-        return f"{v:.2f}{unit}" if isinstance(v, (int, float)) else "—"
-
-    out = [f"### {row['task'].replace('_', ' ')}", f"`{row['episode_id']}`", ""]
-    if any(isinstance(row[k], (int, float)) for k in ("full_psnr", "crit_psnr", "full_mse", "crit_mse")):
-        out += [
-            "| | full window | critical frame |",
-            "|---|---|---|",
-            f"| **PSNR** | {fmt(row['full_psnr'], ' dB')} | {fmt(row['crit_psnr'], ' dB')} |",
-            f"| **MSE** | {fmt(row['full_mse'])} | {fmt(row['crit_mse'])} |",
-            "",
-        ]
-    bits = []
-    if row["edge"] != "?":
-        bits.append(f"**{row['edge']}**")
-    if row["arm"] != "?":
-        bits.append(f"arm **{row['arm']}**")
-    if isinstance(row["amplitude"], (int, float)):
-        bits.append(f"amplitude {fmt(row['amplitude'])}")
-    win = row.get("window")
-    if isinstance(win, (list, tuple)) and len(win) == 2:
-        bits.append(f"window frames {win[0]}–{win[1]}")
-    if row["n_frames"]:
-        bits.append(f"of {row['n_frames']}")
+        return f"{v:.2f}{unit}" if isinstance(v, (int, float)) else "-"
+
+    out = [
+        f"**{row['task'].replace('_', ' ')}**",
+        "",
+        f"`{row['episode_id']}`",
+        "",
+        "| | full window | critical frame |",
+        "|---|---|---|",
+        f"| **PSNR** | {fmt(row['full_psnr'], ' dB')} | {fmt(row['crit_psnr'], ' dB')} |",
+        f"| **MSE** | {fmt(row['full_mse'])} | {fmt(row['crit_mse'])} |",
+        "",
+    ]
+    bits = [b for b in (row["edge"], row["arm"]) if b and b != "?"]
     if row["critical_frame"] is not None:
         bits.append(f"critical frame {row['critical_frame']}")
     if bits:
-        out.append(" · ".join(bits))
-    out.append(
-        f"\n<sub>cached this session: {session.disk_windows()}/{MAX_WINDOWS_ON_DISK} windows "
-        "— evicted least-recently-used, all deleted when you close the tab</sub>"
-    )
+        out.append(" | ".join(bits))
     return "\n".join(out)
 
 
-def table_rows(rows: list[dict]) -> list[list]:
-    def num(v):
-        return round(v, 2) if isinstance(v, (int, float)) else None
-
-    return [
-        [short_id(r["episode_id"]), r["task"], r["edge"], r["arm"], num(r["crit_psnr"]), num(r["full_psnr"])]
-        for r in rows
-    ]
-
-
 def critical_index(row: dict) -> int:
-    """Offset of the critical frame inside the window, or NO_CRIT."""
     start, crit = row.get("start_frame"), row.get("critical_frame")
     if start is None or crit is None:
-        return NO_CRIT
+        return -1
     return int(crit) - int(start)
 
 
 def on_load(repo_text, sid):
-    """Resolve a repo and load only its index — no window data is pulled here."""
+    """Resolve a repo and load only its index - no window data is pulled here."""
     sweep()
     session = get_session(sid)
     repo, revision = parse_repo(repo_text)
@@ -744,74 +385,69 @@ def on_load(repo_text, sid):
     except Exception as exc:
         session.rows = []
         gr.Warning(f"Could not load {repo}: {exc}")
-        return session.sid, f"❌ **{repo}** — {exc}"
+        return session.sid, f"**{repo}** failed: {exc}"
     at = f"@{revision}" if revision else ""
-    return session.sid, (
-        f"✅ **{repo}{at}** — {len(session.rows)} windows from {note}. "
-        "Window files download only when you open them."
-    )
+    return session.sid, f"**{repo}{at}** - {len(session.rows)} episodes from {note}"
 
 
-def on_filter(sid, query, edge, arm, sort, current):
+def on_filter(sid, query, sort, current):
     session = get_session(sid)
     if not session.rows:
-        return gr.update(choices=[], value=None), [], []
-    rows = filter_rows(session, query, edge, arm, sort)
+        return gr.update(choices=[], value=None), []
+    rows = filter_rows(session, query, sort)
     choices = [(label_for(r), r["episode_id"]) for r in rows]
     ids = [r["episode_id"] for r in rows]
     value = current if current in ids else (ids[0] if ids else None)
     return (
-        gr.update(choices=choices, value=value, label=f"Window ({len(ids)} of {len(session.rows)})"),
-        table_rows(rows),
+        gr.update(choices=choices, value=value, label=f"Episode ({len(ids)} of {len(session.rows)})"),
         ids,
     )
 
 
-def on_select(sid, episode_id, view, show_diff, fps, frame_idx):
-    session = get_session(sid)
-    if not session.rows or not episode_id:
-        return None, None, "*No window selected.*", None, gr.update(), gr.update()
-    row = row_by_id(session, episode_id)
-    crit = critical_index(row)
-
-    cams = n_cameras(session.pair(episode_id)[0])
-    choices = camera_choices(cams)
-    valid = [v for _, v in choices]
-    view = view if view in valid else choices[0][1]
-    cam_update = gr.update(choices=choices, value=view)
-
-    video = render_video(session, episode_id, view, show_diff, fps, crit)
-    frames = session.composed(episode_id, view, show_diff, crit)
-    last = len(frames) - 1
-    i = int(min(frame_idx, last))
-    plot = plot_actions(session, episode_id, crit, i)
-    slider = gr.update(maximum=last, value=i, label=f"Frame in window (0–{last})")
-    return video, frames[i], metrics_md(session, row), plot, slider, cam_update
-
-
-def on_frame(sid, episode_id, view, show_diff, frame_idx):
+def on_metrics(sid, episode_id):
     session = get_session(sid)
     if not session.rows or not episode_id:
-        return None, None
-    crit = critical_index(row_by_id(session, episode_id))
-    frames = session.composed(episode_id, view, show_diff, crit)
-    i = int(min(frame_idx, len(frames) - 1))
-    return frames[i], plot_actions(session, episode_id, crit, i)
+        return "*No episode selected.*"
+    return metrics_md(row_by_id(session, episode_id))
+
 
+def render_episode(sid, episode_id):
+    """Write the episode to an .rrd: both videos side by side on one frame timeline.
 
-def on_rerun(sid, episode_id, view, show_diff):
+    Handing the viewer a finished file rather than a live binary stream: streaming
+    races the viewer's wasm boot, and chunks that land before it is ready are dropped
+    against a stopped instance.
+    """
     session = get_session(sid)
     if not session.rows or not episode_id:
-        return None
-    crit = critical_index(row_by_id(session, episode_id))
-    return build_rrd(session, episode_id, view, show_diff, crit)
+        # Never hand the viewer a null: that resets and stops the wasm instance, and the
+        # next real value then tries to open a channel on a dead viewer.
+        return gr.skip()
 
+    path = session.render / f"{episode_id}.rrd"
+    if path.exists():
+        return str(path)
 
-def on_rerun_if_active(active, sid, episode_id, view, show_diff):
-    """Only build a recording while the Rerun tab is the one on screen."""
-    if active != "Rerun":
-        return gr.skip()
-    return on_rerun(sid, episode_id, view, show_diff)
+    gt, gen = session.pair(episode_id)
+    crit = critical_index(row_by_id(session, episode_id))
+    rec = rr.RecordingStream("wmviz", recording_id=episode_id)
+    for i in range(len(gt)):
+        rec.set_time("frame", sequence=i)
+        rec.log("ground_truth", rr.Image(gt[i]).compress(jpeg_quality=JPEG_QUALITY))
+        rec.log("generated", rr.Image(gen[i]).compress(jpeg_quality=JPEG_QUALITY))
+        if i == crit:
+            rec.log("critical_frame", rr.TextLog("critical frame"))
+    rec.save(
+        str(path),
+        default_blueprint=rrb.Blueprint(
+            rrb.Horizontal(
+                rrb.Spatial2DView(origin="ground_truth", name="Ground truth"),
+                rrb.Spatial2DView(origin="generated", name="Generated"),
+            ),
+            collapse_panels=True,
+        ),
+    )
+    return str(path)
 
 
 def step(ids, current, delta):
@@ -826,124 +462,65 @@ def step(ids, current, delta):
 
 # --------------------------------------------------------------------------- app
 
-CSS = """
-.window-video video { max-height: 78vh; object-fit: contain; background: #111; }
-.window-still img { max-height: 78vh; object-fit: contain; background: #111; }
-"""
-
 with gr.Blocks(title="Rollout vs. ground truth", fill_width=True, delete_cache=(600, 600)) as demo:
-    gr.Markdown("# Forward-dynamics rollout vs. ground truth")
-
     sid_state = gr.State("", time_to_live=SESSION_TTL, delete_callback=release)
     ids_state = gr.State([])
-    active_tab = gr.State("Playback")
 
     with gr.Row():
-        repo_box = gr.Textbox(
-            DEFAULT_REPO,
-            label="Dataset repo",
-            info="owner/name, optionally owner/name@revision. A repo URL also works.",
-            scale=5,
-            submit_btn=True,
-        )
-        load_btn = gr.Button("Load", variant="primary", scale=1)
-    status = gr.Markdown()
-
-    with gr.Row():
-        with gr.Column(scale=1, min_width=320):
-            query = gr.Textbox(label="Search task or episode", placeholder="e.g. napkin, pillowcase…")
+        with gr.Column(scale=1, min_width=340):
+            gr.Markdown("### Rollout vs. ground truth")
             with gr.Row():
-                edge = gr.Radio(["any", "grasp", "release"], value="any", label="Edge")
-                arm = gr.Radio(["any", "L", "R"], value="any", label="Arm")
+                repo_box = gr.Textbox(
+                    label="Dataset repo",
+                    placeholder="owner/name  (or owner/name@revision, or a dataset URL)",
+                    scale=4,
+                )
+                load_btn = gr.Button("Load", variant="primary", scale=1, min_width=70)
+            status = gr.Markdown("Enter a dataset repo id above and press **Load**.")
+            query = gr.Textbox(label="Search", placeholder="task or episode id")
             sort = gr.Dropdown(SORTS, value=SORTS[0], label="Sort")
-            picker = gr.Dropdown([], label="Window", filterable=True)
+            picker = gr.Dropdown([], label="Episode", filterable=True)
             with gr.Row():
-                prev_btn = gr.Button("← Prev")
-                next_btn = gr.Button("Next →")
-            gr.Markdown("---")
-            view = gr.Radio(camera_choices(3), value=0, label="Camera")
-            show_diff = gr.Checkbox(False, label="Add |GT − generated| difference panel")
-            fps = gr.Slider(2, 30, value=8, step=1, label="Playback fps")
+                prev_btn = gr.Button("Prev")
+                next_btn = gr.Button("Next")
             info = gr.Markdown()
 
-        with gr.Column(scale=2):
-            with gr.Tabs():
-                with gr.Tab("Playback") as tab_playback:
-                    video = gr.Video(
-                        label="GT vs generated",
-                        autoplay=True,
-                        loop=True,
-                        buttons=["download", "fullscreen"],
-                        elem_classes="window-video",
-                    )
-                with gr.Tab("Frame stepper") as tab_stepper:
-                    frame_idx = gr.Slider(0, 32, value=16, step=1, label="Frame in window")
-                    still = gr.Image(label="Frame comparison", elem_classes="window-still")
-                with gr.Tab("Driving actions") as tab_actions:
-                    actions_plot = gr.Plot(label="Driving action chunk")
-                with gr.Tab("Rerun") as tab_rerun:
-                    gr.Markdown(
-                        "Every stream on one shared timeline — scrub once and the cameras and "
-                        "action plots move together. Frames are JPEG-compressed for transfer."
-                    )
-                    rerun_view = Rerun(height=720, streaming=False)
-                with gr.Tab("All windows") as tab_table:
-                    table = gr.Dataframe(
-                        headers=["episode", "task", "edge", "arm", "crit PSNR", "full PSNR"],
-                        datatype=["str", "str", "str", "str", "number", "number"],
-                        interactive=False,
-                        wrap=True,
-                        max_height=650,
-                    )
-
-    filter_inputs = [sid_state, query, edge, arm, sort]
-    filter_outputs = [picker, table, ids_state]
-    view_inputs = [sid_state, picker, view, show_diff, fps, frame_idx]
-    view_outputs = [video, still, info, actions_plot, frame_idx, view]
-    rerun_inputs = [sid_state, picker, view, show_diff]
-
-    def wire_select(event):
-        """Render the window, then refresh the recording only if that tab is open."""
-        return event(on_select, view_inputs, view_outputs).then(
-            on_rerun_if_active, [active_tab] + rerun_inputs, rerun_view
-        )
-
-    def wire_load(event):
-        return wire_select(
-            event(on_load, [repo_box, sid_state], [sid_state, status])
-            .then(on_filter, filter_inputs + [picker], filter_outputs)
-            .then
-        )
-
-    wire_load(load_btn.click)
-    wire_load(repo_box.submit)
-    wire_load(demo.load)
-
-    for control in (query, edge, arm, sort):
-        wire_select(control.change(on_filter, filter_inputs + [picker], filter_outputs).then)
-
-    wire_select(picker.change)
-    for control in (view, show_diff, fps):
-        wire_select(control.change)
-
-    for tab, name in (
-        (tab_playback, "Playback"),
-        (tab_stepper, "Frame stepper"),
-        (tab_actions, "Driving actions"),
-        (tab_table, "All windows"),
-    ):
-        tab.select(lambda name=name: name, None, active_tab)
-    # opening the Rerun tab is what triggers the first build for a window
-    tab_rerun.select(lambda: "Rerun", None, active_tab).then(on_rerun, rerun_inputs, rerun_view)
-    frame_idx.change(on_frame, [sid_state, picker, view, show_diff, frame_idx], [still, actions_plot])
+        with gr.Column(scale=4):
+            # Config matched to rerun/gradio-rerun-viewer: a CSS-string height leaves the
+            # container at zero height while the viewer re-initialises, and it never recovers.
+            viewer = Rerun(
+                streaming=True,
+                height=820,
+                panel_states={"time": "expanded", "blueprint": "hidden", "selection": "hidden"},
+            )
 
+    filter_io = ([sid_state, query, sort, picker], [picker, ids_state])
+
+    # Loading and filtering only repopulate the episode list. Selecting an episode is the
+    # single trigger that streams into the viewer -- if these chains also streamed, setting
+    # the dropdown value here would fire picker.change too and run two generators
+    # concurrently into one viewer, which crashes it.
+    for event in (load_btn.click, repo_box.submit):
+        event(on_load, [repo_box, sid_state], [sid_state, status]).then(on_filter, *filter_io)
+
+    for control in (query, sort):
+        control.change(on_filter, *filter_io)
+
+    picker.change(on_metrics, [sid_state, picker], info)
+    # trigger_mode="once" is load-bearing: .change() defaults to "always_last", which lets a
+    # second run queue up behind the pending one. Two generators streaming into one viewer
+    # stops the first viewer mid-write, which is what left it stuck on "Loading Rerun...".
+    picker.change(
+        render_episode,
+        [sid_state, picker],
+        viewer,
+        concurrency_limit=1,
+        trigger_mode="once",
+    )
     prev_btn.click(lambda i, c: step(i, c, -1), [ids_state, picker], picker)
     next_btn.click(lambda i, c: step(i, c, +1), [ids_state, picker], picker)
 
-    def on_table_select(ids, evt: gr.SelectData):
-        return ids[evt.index[0]] if ids and evt.index[0] < len(ids) else gr.update()
-
-    table.select(on_table_select, ids_state, picker)
-
 if __name__ == "__main__":
-    demo.queue(max_size=16).launch(css=CSS)
+    # ssr_mode must stay off: server-side rendering re-mounts the viewer canvas on
+    # hydration, which kills the wasm instance mid-stream.
+    demo.queue(max_size=16).launch(ssr_mode=False)

requirements.txt

@@ -1,7 +1,5 @@
-gradio==6.22.0
-gradio_rerun>=0.35.0
+gradio==6.5.1
+gradio_rerun==0.35.0
 huggingface_hub>=0.28
 numpy
 opencv-python-headless
-imageio-ffmpeg
-matplotlib