feat(datasets): replay an episode on the follower arm

CLAUDE.md

@@ -37,6 +37,7 @@ Run the Python tests with `pytest` (config in [pyproject.toml](pyproject.toml);
 
 - [record.py](lelab/record.py) — dataset recording (wraps `lerobot.record.record`); patches `lerobot.common.utils.control_utils` keyboard listener so frontend buttons replace arrow-key controls. The worker is a daemon thread, so the server's shutdown hook calls `stop_recording_and_wait()` to let an active session finalize on Ctrl+C.
 - [teleoperate.py](lelab/teleoperate.py) — leader→follower teleoperation (wraps `lerobot.teleoperate`).
+- [replay.py](lelab/replay.py): plays a recorded episode's actions on the follower arm, an in-process take on `lerobot-replay` so the dataset browser can stop it.
 - [calibrate.py](lelab/calibrate.py) — step-by-step web calibration with a `CalibrationManager` singleton and `_step_complete` threading.Event.
 - [train.py](lelab/train.py) — wraps the LeRobot training CLI as a subprocess (psutil for lifecycle, queue for log streaming).
 - [datasets.py](lelab/datasets.py) — dataset listing (local cache + Hub, merged under a `source` field) plus the episode-browsing handlers behind `/dataset-episodes`, `/dataset-episode`, `/dataset-frame`, `/dataset-thumbnails`, `/dataset-motion` and `/dataset-video`.
@@ -45,7 +46,7 @@ Run the Python tests with `pytest` (config in [pyproject.toml](pyproject.toml);
 
 ### State model
 
-Each feature module owns module-level globals (e.g. `recording_active`, `teleoperation_active`, `current_robot`) protected by threads/locks where needed. There's no shared session object — features are mutually exclusive in practice (you can't teleoperate and record simultaneously) but this is **not** enforced in code.
+Each feature module owns module-level globals (e.g. `recording_active`, `teleoperation_active`, `current_robot`) protected by threads/locks where needed. There's no shared session object. Recording, teleoperation, inference and replay all drive the follower, so each start handler checks the others' flags and refuses to run while one is active. Calibration is not part of that check.
 
 ### WebSocket broadcast
 

frontend/src/components/dataset/EpisodeViewer.test.tsx

@@ -10,6 +10,7 @@ vi.mock("@/contexts/ApiContext", () => ({
 }));
 vi.mock("./FilmStrip", () => ({ default: () => null }));
 vi.mock("./MotionTrace", () => ({ default: () => null }));
+vi.mock("./ReplayButton", () => ({ default: () => null }));
 
 test.each([0, 4])("playback loops at file end with an episode offset of %s seconds", async (offset) => {
   Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true });

frontend/src/components/dataset/EpisodeViewer.tsx

@@ -15,6 +15,7 @@ import { useApi } from "@/contexts/ApiContext";
 import { EpisodeDetail, frameUrl, videoUrl } from "@/lib/datasetApi";
 import FilmStrip from "./FilmStrip";
 import MotionTrace from "./MotionTrace";
+import ReplayButton from "./ReplayButton";
 
 const RATES = [0.5, 1, 2] as const;
 type Rate = (typeof RATES)[number];
@@ -160,6 +161,18 @@ const EpisodeViewer: React.FC<EpisodeViewerProps> = ({ repoId, detail, onNavigat
     }
   }, []);
 
+  // The arm replays at the recording's speed; run the footage alongside it.
+  const playFromStart = useCallback(() => {
+    const v = driverRef.current;
+    if (!v) return;
+    setRate(1);
+    seekToFrame(0);
+    void v.play().then(
+      () => setPlaying(true),
+      () => setPlaying(false),
+    );
+  }, [seekToFrame]);
+
   // Keep the element's rate in step with the pills across re-renders and
   // camera swaps.
   useEffect(() => {
@@ -229,19 +242,26 @@ const EpisodeViewer: React.FC<EpisodeViewerProps> = ({ repoId, detail, onNavigat
           </Button>
         )}
 
-        <div className="ml-auto flex rounded-md border border-gray-800 bg-gray-950 p-0.5">
-          {RATES.map((r) => (
-            <button
-              key={r}
-              type="button"
-              onClick={() => setRate(r)}
-              className={`rounded px-2.5 py-1 text-xs font-medium tabular-nums transition ${
-                r === rate ? "bg-gray-800 text-white" : "text-gray-400 hover:text-gray-200"
-              }`}
-            >
-              {r}x
-            </button>
-          ))}
+        <div className="ml-auto flex items-center gap-2">
+          <ReplayButton
+            repoId={repoId}
+            episodeIndex={detail.episode_index}
+            onStart={playFromStart}
+          />
+          <div className="flex rounded-md border border-gray-800 bg-gray-950 p-0.5">
+            {RATES.map((r) => (
+              <button
+                key={r}
+                type="button"
+                onClick={() => setRate(r)}
+                className={`rounded px-2.5 py-1 text-xs font-medium tabular-nums transition ${
+                  r === rate ? "bg-gray-800 text-white" : "text-gray-400 hover:text-gray-200"
+                }`}
+              >
+                {r}x
+              </button>
+            ))}
+          </div>
         </div>
       </div>
 

frontend/src/components/dataset/ReplayButton.test.tsx

@@ -0,0 +1,134 @@
+// @vitest-environment jsdom
+import React, { act } from "react";
+import { createRoot, Root } from "react-dom/client";
+import { afterEach, beforeEach, expect, test, vi } from "vitest";
+import ReplayButton from "./ReplayButton";
+import { RobotRecord } from "@/hooks/useRobots";
+
+const backend = vi.hoisted(() => ({
+  active: false,
+  startOk: true,
+  statusFails: false,
+  requests: [] as { path: string; body?: unknown }[],
+}));
+const robot = vi.hoisted(() => ({ current: null as RobotRecord | null }));
+const toast = vi.hoisted(() => vi.fn());
+const api = vi.hoisted(() => ({ baseUrl: "http://localhost:8000", fetchWithHeaders: () => {} }));
+
+vi.mock("@/contexts/ApiContext", () => ({ useApi: () => api }));
+vi.mock("@/hooks/useRobots", () => ({ useRobots: () => ({ selectedRecord: robot.current }) }));
+vi.mock("@/hooks/use-toast", () => ({ useToast: () => ({ toast }) }));
+vi.mock("@/lib/apiClient", () => ({
+  apiRequest: async (_baseUrl: string, _fetcher: unknown, path: string, opts?: { body?: unknown }) => {
+    backend.requests.push({ path, body: opts?.body });
+    if (path === "/replay-status") {
+      if (backend.statusFails) throw new Error("network");
+      return { replay_active: backend.active, error: null };
+    }
+    if (path === "/start-replay") {
+      backend.active = backend.startOk;
+      return { success: backend.startOk, message: "A replay is already running" };
+    }
+    return { success: true, message: "Replay stopped" };
+  },
+}));
+
+const calibrated: RobotRecord = {
+  name: "arm",
+  leader_port: "/dev/leader",
+  follower_port: "/dev/follower",
+  leader_config: "leader.json",
+  follower_config: "follower.json",
+  cameras: [],
+  is_clean: true,
+};
+
+let host: HTMLDivElement;
+let root: Root;
+const onStart = vi.fn();
+
+const button = (label: string) =>
+  [...host.querySelectorAll("button")].find((b) => b.textContent?.includes(label));
+
+async function mount() {
+  await act(async () => {
+    root.render(<ReplayButton repoId="local/ds" episodeIndex={2} onStart={onStart} />);
+  });
+}
+
+async function poll() {
+  await act(async () => {
+    await vi.advanceTimersByTimeAsync(500);
+  });
+}
+
+beforeEach(() => {
+  Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true });
+  vi.useFakeTimers();
+  Object.assign(backend, { active: false, startOk: true, statusFails: false, requests: [] });
+  robot.current = calibrated;
+  toast.mockClear();
+  onStart.mockClear();
+  host = document.createElement("div");
+  root = createRoot(host);
+});
+
+afterEach(() => {
+  act(() => root.unmount());
+  vi.useRealTimers();
+});
+
+test("hidden until a calibrated robot is selected", async () => {
+  robot.current = null;
+  await mount();
+  expect(host.querySelector("button")).toBeNull();
+});
+
+test("a replay already running stays stoppable without a selected robot", async () => {
+  robot.current = null;
+  backend.active = true;
+  await mount();
+  expect(button("Stop robot")).toBeDefined();
+});
+
+test("starting replays this episode and runs the footage alongside", async () => {
+  await mount();
+  await act(async () => button("Play on robot")!.click());
+  expect(backend.requests.at(-1)).toMatchObject({
+    path: "/start-replay",
+    body: { dataset_repo_id: "local/ds", episode_index: 2, follower_port: "/dev/follower" },
+  });
+  expect(onStart).toHaveBeenCalledOnce();
+  expect(button("Stop robot")).toBeDefined();
+});
+
+test("Stop stays until the status shows the arm released", async () => {
+  backend.active = true;
+  await mount();
+  await act(async () => button("Stop robot")!.click());
+  await poll();
+  expect(button("Stop robot")).toBeDefined();
+
+  backend.active = false;
+  await poll();
+  expect(button("Play on robot")).toBeDefined();
+});
+
+test("a failed status poll keeps the Stop button", async () => {
+  backend.active = true;
+  await mount();
+  backend.statusFails = true;
+  await poll();
+  expect(button("Stop robot")).toBeDefined();
+});
+
+test("a refused start reports it and offers Play again", async () => {
+  backend.startOk = false;
+  await mount();
+  await act(async () => button("Play on robot")!.click());
+  expect(onStart).not.toHaveBeenCalled();
+  expect(toast).toHaveBeenCalledWith(
+    expect.objectContaining({ description: "A replay is already running" }),
+  );
+  expect(button("Play on robot")).toBeDefined();
+});

frontend/src/components/dataset/ReplayButton.tsx

@@ -0,0 +1,147 @@
+import React, { useEffect, useRef, useState } from "react";
+import { Bot, Loader2, Square } from "lucide-react";
+import { Button } from "@/components/ui/button";
+import { useApi } from "@/contexts/ApiContext";
+import { useToast } from "@/hooks/use-toast";
+import { useRobots } from "@/hooks/useRobots";
+import { apiRequest } from "@/lib/apiClient";
+
+interface ReplayResponse {
+  success: boolean;
+  message: string;
+}
+
+interface ReplayStatus {
+  replay_active: boolean;
+  error: string | null;
+}
+
+interface ReplayButtonProps {
+  repoId: string;
+  episodeIndex: number;
+  /** Called once the arm starts moving, so the footage can run alongside it. */
+  onStart: () => void;
+}
+
+/**
+ * Plays the episode's recorded actions on the selected follower arm. Only
+ * shown when that arm is ready to move.
+ */
+const ReplayButton: React.FC<ReplayButtonProps> = ({ repoId, episodeIndex, onStart }) => {
+  const { baseUrl, fetchWithHeaders } = useApi();
+  const { selectedRecord: robot } = useRobots();
+  const { toast } = useToast();
+  const [phase, setPhase] = useState<"idle" | "starting" | "playing">("idle");
+  // Paging to another episode while the arm connects must not drag the
+  // footage back to the one being replayed.
+  const episodeRef = useRef(episodeIndex);
+  episodeRef.current = episodeIndex;
+
+  // A replay keeps running when this button unmounts (another dataset, another
+  // page), so pick it back up to keep it stoppable.
+  useEffect(() => {
+    apiRequest<ReplayStatus>(baseUrl, fetchWithHeaders, "/replay-status")
+      .then((status) => {
+        if (status.replay_active) setPhase("playing");
+      })
+      .catch(() => undefined);
+  }, [baseUrl, fetchWithHeaders]);
+
+  // Poll only while the arm moves, to notice when the episode ends or fails.
+  useEffect(() => {
+    if (phase !== "playing") return;
+    const timer = setInterval(async () => {
+      try {
+        const status = await apiRequest<ReplayStatus>(baseUrl, fetchWithHeaders, "/replay-status");
+        if (status.replay_active) return;
+        setPhase("idle");
+        if (status.error) {
+          toast({ title: "Replay failed", description: status.error, variant: "destructive" });
+        }
+      } catch {
+        // Keep the Stop button: the next poll retries.
+      }
+    }, 500);
+    return () => clearInterval(timer);
+  }, [phase, baseUrl, fetchWithHeaders, toast]);
+
+  // Back to Play only once the status poll sees the arm released.
+  const stop = async () => {
+    try {
+      const res = await apiRequest<ReplayResponse>(baseUrl, fetchWithHeaders, "/stop-replay", {
+        method: "POST",
+        action: "Stop",
+      });
+      if (!res.success) throw new Error(res.message);
+    } catch (e) {
+      toast({
+        title: "Could not stop the robot",
+        description: e instanceof Error ? e.message : String(e),
+        variant: "destructive",
+      });
+    }
+  };
+
+  // A running replay stays stoppable even once the robot is no longer selected.
+  if (phase === "playing") {
+    return (
+      <Button
+        variant="outline"
+        size="sm"
+        onClick={stop}
+        className="h-7 gap-1.5 border-red-900 bg-red-950/40 text-xs text-red-300 hover:bg-red-950 hover:text-red-200"
+      >
+        <Square className="h-3 w-3 fill-current" />
+        Stop robot
+      </Button>
+    );
+  }
+
+  if (!robot?.is_clean) return null;
+
+  const start = async () => {
+    setPhase("starting");
+    try {
+      const res = await apiRequest<ReplayResponse>(baseUrl, fetchWithHeaders, "/start-replay", {
+        method: "POST",
+        body: {
+          follower_port: robot.follower_port,
+          follower_config: robot.follower_config,
+          dataset_repo_id: repoId,
+          episode_index: episodeIndex,
+        },
+        action: "Replay",
+      });
+      if (!res.success) throw new Error(res.message);
+      setPhase("playing");
+      if (episodeRef.current === episodeIndex) onStart();
+    } catch (e) {
+      setPhase("idle");
+      toast({
+        title: "Could not replay this episode",
+        description: e instanceof Error ? e.message : String(e),
+        variant: "destructive",
+      });
+    }
+  };
+
+  return (
+    <Button
+      variant="outline"
+      size="sm"
+      onClick={start}
+      disabled={phase === "starting"}
+      title={`Replay this episode on ${robot.name}`}
+      className="h-7 gap-1.5 border-gray-800 bg-gray-950 text-xs text-gray-300 hover:bg-gray-900"
+    >
+      {phase === "starting" ? (
+        <Loader2 className="h-3.5 w-3.5 animate-spin" />
+      ) : (
+        <Bot className="h-3.5 w-3.5" />
+      )}
+      Play on robot
+    </Button>
+  );
+};
+
+export default ReplayButton;