Add a Rerun viewer tab

Logs the window to an .rrd and renders it with gradio_rerun: each camera's
GT/generated/diff as image streams and every action channel as a scalar
series, all on one shared frame timeline. This is the thing the composed mp4
could not do — the action plots were a static PNG that could not be scrubbed
against the video.

Frames are JPEG-compressed into the recording (~1.5 MB rather than ~16 MB for
a 3-camera window). The recording is built only when the tab is opened, and
is evicted along with its window like the composed mp4.

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

README.md

@@ -29,6 +29,10 @@ URL. `DATASET_REPO` sets what is prefilled at startup.
   frame when the dataset marks one.
 - **Driving actions** — the action chunk fed to the model. A 14-D chunk is split into grippers /
   left-arm / right-arm; any other width is plotted as raw channels.
+- **Rerun** — the same window in a [Rerun](https://rerun.io) viewer: every camera and every
+  action channel on one shared frame timeline, so scrubbing moves all of them together. One row
+  per camera (GT | generated | optional diff) over a time-series plot of the actions, with the
+  critical frame marked. Built only when you open the tab, and only for the cameras selected.
 - **All windows** — sortable table of every window; click a row to load it.
 
 Filter by task, gripper edge and arm; sort by critical-frame PSNR to jump to the hardest or
@@ -69,6 +73,13 @@ too large to fit on the Space.
   `SESSION_TTL` and directories orphaned by an earlier process, so nothing accumulates if a tab
   dies without a clean shutdown.
 
+Rerun recordings follow the same rule: built on first view of the tab, cached beside the
+composed mp4, and evicted with their window. Frames go into the recording as JPEG (quality 90) —
+a raw recording of one 3-camera window is ~16 MB against ~1.5 MB compressed, which is the
+difference between usable and not over a browser connection. The trade-off is that extreme zoom
+in the Rerun viewer shows JPEG artifacts; the **Frame stepper** tab renders from the raw decoded
+frames if you need a pixel-exact look.
+
 The sidebar shows live occupancy (`cached this session: n/N windows`).
 
 ## Configuration

app.py

@@ -40,6 +40,9 @@ 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
 
@@ -63,6 +66,7 @@ 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)
@@ -147,7 +151,7 @@ class Session:
             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}_*.mp4"):
+            for f in self.render.glob(f"{old}_*"):
                 f.unlink(missing_ok=True)
 
     def disk_windows(self) -> int:
@@ -416,11 +420,21 @@ def n_cameras(frames: np.ndarray) -> int:
 
 
 def camera_choices(n: int) -> list[tuple[str, int]]:
-    names = THREE_CAM_NAMES if n == 3 else [f"Camera {i + 1}" for i in range(n)]
-    choices = [(name, i) for i, name in enumerate(names)]
+    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
@@ -546,15 +560,21 @@ def render_video(session: Session, episode_id: str, view: int, show_diff: bool,
 # --------------------------------------------------------------------------- actions plot
 
 
-def plot_actions(session: Session, episode_id: str, crit_idx: int, cursor: int | None) -> Figure | None:
+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:
-        actions = np.atleast_2d(np.asarray(json.load(open(path)), dtype=float))
+        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]
@@ -585,6 +605,75 @@ def plot_actions(session: Session, episode_id: str, crit_idx: int, cursor: int |
     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
 
 
@@ -710,6 +799,21 @@ def on_frame(sid, episode_id, view, show_diff, frame_idx):
     return frames[i], plot_actions(session, episode_id, crit, i)
 
 
+def on_rerun(sid, episode_id, view, show_diff):
+    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)
+
+
+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)
+
+
 def step(ids, current, delta):
     if not ids:
         return None
@@ -732,6 +836,7 @@ with gr.Blocks(title="Rollout vs. ground truth", fill_width=True, delete_cache=(
 
     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(
@@ -763,7 +868,7 @@ with gr.Blocks(title="Rollout vs. ground truth", fill_width=True, delete_cache=(
 
         with gr.Column(scale=2):
             with gr.Tabs():
-                with gr.Tab("Playback"):
+                with gr.Tab("Playback") as tab_playback:
                     video = gr.Video(
                         label="GT vs generated",
                         autoplay=True,
@@ -771,12 +876,18 @@ with gr.Blocks(title="Rollout vs. ground truth", fill_width=True, delete_cache=(
                         buttons=["download", "fullscreen"],
                         elem_classes="window-video",
                     )
-                with gr.Tab("Frame stepper"):
+                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"):
+                with gr.Tab("Driving actions") as tab_actions:
                     actions_plot = gr.Plot(label="Driving action chunk")
-                with gr.Tab("All windows"):
+                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"],
@@ -789,12 +900,19 @@ with gr.Blocks(title="Rollout vs. ground truth", fill_width=True, delete_cache=(
     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 (
+        return wire_select(
             event(on_load, [repo_box, sid_state], [sid_state, status])
             .then(on_filter, filter_inputs + [picker], filter_outputs)
-            .then(on_select, view_inputs, view_outputs)
+            .then
         )
 
     wire_load(load_btn.click)
@@ -802,13 +920,21 @@ with gr.Blocks(title="Rollout vs. ground truth", fill_width=True, delete_cache=(
     wire_load(demo.load)
 
     for control in (query, edge, arm, sort):
-        control.change(on_filter, filter_inputs + [picker], filter_outputs).then(
-            on_select, view_inputs, view_outputs
-        )
+        wire_select(control.change(on_filter, filter_inputs + [picker], filter_outputs).then)
 
-    picker.change(on_select, view_inputs, view_outputs)
+    wire_select(picker.change)
     for control in (view, show_diff, fps):
-        control.change(on_select, view_inputs, view_outputs)
+        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])
 
     prev_btn.click(lambda i, c: step(i, c, -1), [ids_state, picker], picker)

requirements.txt

@@ -1,4 +1,5 @@
 gradio==6.22.0
+gradio_rerun>=0.35.0
 huggingface_hub>=0.28
 numpy
 opencv-python-headless