From 35e415be8a7e8fb08185ba506b15ebaf7968c1bb Mon Sep 17 00:00:00 2001 From: Nicolas Rabault Date: Mon, 5 Oct 2026 19:17:05 +0200 Subject: [PATCH 1/2] feat(datasets): replay an episode on the follower arm The episode viewer gets a "Play on robot" button, shown only once a calibrated robot is selected. It sends the episode's recorded actions to the follower at the recording's rate and starts the footage from the first frame at 1x so both can be compared; while the arm moves the button becomes "Stop robot", and it picks a running replay back up when the viewer is opened again. replay.py is an in-process take on lerobot-replay, built like teleoperation: the follower connects synchronously so a failure reaches the user, the actions are read from the episode's data files, and the worker owns the disconnect and the active flag. Recording, teleoperation, inference and replay all drive the follower, so each refuses to start while another runs. --- CLAUDE.md | 3 +- .../components/dataset/EpisodeViewer.test.tsx | 1 + .../src/components/dataset/EpisodeViewer.tsx | 46 ++++-- .../src/components/dataset/ReplayButton.tsx | 142 ++++++++++++++++++ lelab/record.py | 4 +- lelab/replay.py | 141 +++++++++++++++++ lelab/rollout.py | 10 +- lelab/server.py | 22 +++ lelab/teleoperate.py | 4 +- tests/test_replay.py | 123 +++++++++++++++ 10 files changed, 478 insertions(+), 18 deletions(-) create mode 100644 frontend/src/components/dataset/ReplayButton.tsx create mode 100644 lelab/replay.py create mode 100644 tests/test_replay.py diff --git a/CLAUDE.md b/CLAUDE.md index 42253292..544fd7cb 100644 --- a/CLAUDE.md +++ b/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 diff --git a/frontend/src/components/dataset/EpisodeViewer.test.tsx b/frontend/src/components/dataset/EpisodeViewer.test.tsx index faff82e4..c8b3f857 100644 --- a/frontend/src/components/dataset/EpisodeViewer.test.tsx +++ b/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 }); diff --git a/frontend/src/components/dataset/EpisodeViewer.tsx b/frontend/src/components/dataset/EpisodeViewer.tsx index 83f19b7e..2603ccc9 100644 --- a/frontend/src/components/dataset/EpisodeViewer.tsx +++ b/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 = ({ 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 = ({ repoId, detail, onNavigat )} -
- {RATES.map((r) => ( - - ))} +
+ +
+ {RATES.map((r) => ( + + ))} +
diff --git a/frontend/src/components/dataset/ReplayButton.tsx b/frontend/src/components/dataset/ReplayButton.tsx new file mode 100644 index 00000000..c71658b2 --- /dev/null +++ b/frontend/src/components/dataset/ReplayButton.tsx @@ -0,0 +1,142 @@ +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 = ({ 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(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(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]); + + if (!robot?.is_clean) return null; + + const start = async () => { + setPhase("starting"); + try { + const res = await apiRequest(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", + }); + } + }; + + const stop = async () => { + try { + await apiRequest(baseUrl, fetchWithHeaders, "/stop-replay", { method: "POST", action: "Stop" }); + setPhase("idle"); + } catch (e) { + toast({ + title: "Could not stop the robot", + description: e instanceof Error ? e.message : String(e), + variant: "destructive", + }); + } + }; + + if (phase === "playing") { + return ( + + ); + } + + return ( + + ); +}; + +export default ReplayButton; diff --git a/lelab/record.py b/lelab/record.py index 3d71f161..4399d8ee 100644 --- a/lelab/record.py +++ b/lelab/record.py @@ -251,7 +251,7 @@ def handle_start_recording(request: RecordingRequest) -> dict[str, Any]: phase_start_time, \ last_recording_info - from . import rollout as _rollout, teleoperate as _teleoperate + from . import replay as _replay, rollout as _rollout, teleoperate as _teleoperate # Claim the active flag under the lock so two concurrent starts can't both # pass the precondition check. @@ -262,6 +262,8 @@ def handle_start_recording(request: RecordingRequest) -> dict[str, Any]: return {"success": False, "message": "Teleoperation is currently active. Stop it first."} if _rollout.inference_active: return {"success": False, "message": "Inference is currently active. Stop it first."} + if _replay.replay_active: + return {"success": False, "message": "Replay is currently active. Stop it first."} recording_active = True recording_thread = None recording_events = None diff --git a/lelab/replay.py b/lelab/replay.py new file mode 100644 index 00000000..78d19880 --- /dev/null +++ b/lelab/replay.py @@ -0,0 +1,141 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Replay a recorded episode on the follower arm: an in-process take on +lerobot-replay, so the dataset browser can stop it.""" + +import logging +import threading +import time +from typing import Any + +from pydantic import BaseModel + +from lerobot.robots.so_follower import SO101Follower, SO101FollowerConfig +from lerobot.utils.robot_utils import precise_sleep + +from . import episode_media +from .utils.config import setup_follower_calibration_file +from .utils.devices import safe_disconnect_device + +logger = logging.getLogger(__name__) + +# Only the worker (or a failed start) clears this, so the follower stays +# claimed until it is actually released. +replay_active = False +replay_thread: threading.Thread | None = None +replay_error: str | None = None +_stop = threading.Event() +# Guards the start path; the worker owns disconnect so stop() does not race. +_state_lock = threading.Lock() + + +class ReplayRequest(BaseModel): + follower_port: str + follower_config: str + dataset_repo_id: str + episode_index: int + + +def _episode_actions(repo_id: str, episode_index: int) -> tuple[list[dict[str, float]], int]: + """The episode's recorded actions as follower commands, and the dataset fps.""" + dataset_dir = episode_media.resolve_dataset_dir(repo_id) + info = episode_media.read_info(dataset_dir) + names = info["features"]["action"]["names"] + rows = episode_media.load_episode_actions(dataset_dir, episode_index) + return [dict(zip(names, row, strict=True)) for row in rows], info["fps"] + + +def handle_start_replay(request: ReplayRequest) -> dict[str, Any]: + """Connect the follower synchronously, so a failure reaches the caller, + then send the episode's actions at the recording's rate in the background.""" + global replay_active, replay_thread, replay_error + + from . import record as _record, rollout as _rollout, teleoperate as _teleoperate + + with _state_lock: + if replay_active: + return {"success": False, "message": "A replay is already running"} + if _record.recording_active: + return {"success": False, "message": "Recording is currently active. Stop it first."} + if _teleoperate.teleoperation_active: + return {"success": False, "message": "Teleoperation is currently active. Stop it first."} + if _rollout.inference_active: + return {"success": False, "message": "Inference is currently active. Stop it first."} + replay_active = True + replay_error = None + _stop.clear() + + robot = None + try: + actions, fps = _episode_actions(request.dataset_repo_id, request.episode_index) + + follower_id = setup_follower_calibration_file(request.follower_config) + robot = SO101Follower(SO101FollowerConfig(port=request.follower_port, id=follower_id)) + try: + robot.bus.connect() + except Exception as e: + raise RuntimeError( + f"Could not connect to the follower arm on {request.follower_port}. " + "Make sure it's plugged in and powered on, then try again." + ) from e + robot.bus.write_calibration(robot.calibration) + robot.configure() + except Exception as e: + safe_disconnect_device(robot, logger) + replay_active = False + logger.error(f"Failed to start replay: {e}") + return {"success": False, "message": str(e)} + + def replay_worker(): + global replay_active, replay_error + + try: + for action in actions: + if _stop.is_set(): + break + start = time.perf_counter() + robot.send_action(action) + precise_sleep(1 / fps - (time.perf_counter() - start)) + except Exception as e: + logger.error(f"Replay failed: {e}") + replay_error = str(e) + finally: + safe_disconnect_device(robot, logger) + replay_active = False + logger.info("Replay ended") + + replay_thread = threading.Thread(target=replay_worker, name="replay-worker", daemon=True) + replay_thread.start() + + return {"success": True, "message": "Replay started"} + + +def handle_stop_replay() -> dict[str, Any]: + """Signal the worker and wait for it to release the arm.""" + if not replay_active: + return {"success": False, "message": "No replay is running"} + + _stop.set() + worker = replay_thread + if worker is not None and worker.is_alive(): + worker.join(timeout=5.0) + if worker.is_alive(): + logger.warning("Replay worker did not exit within 5s") + + return {"success": True, "message": "Replay stopped"} + + +def handle_replay_status() -> dict[str, Any]: + return {"replay_active": replay_active, "error": replay_error} diff --git a/lelab/rollout.py b/lelab/rollout.py index 09c2a8d5..93d039b3 100644 --- a/lelab/rollout.py +++ b/lelab/rollout.py @@ -301,8 +301,8 @@ def handle_start_inference(request: InferenceRequest) -> dict[str, Any]: global inference_active, _inference_proc, _inference_started_at global _inference_rollout_started_at, _inference_meta - # Mutex with teleop and recording: all three drive the same serial bus. - from . import record as _record, teleoperate as _teleoperate + # Mutex with teleop, recording and replay: they all drive the same serial bus. + from . import record as _record, replay as _replay, teleoperate as _teleoperate with _state_lock: if _teleoperate.teleoperation_active: @@ -317,6 +317,12 @@ def handle_start_inference(request: InferenceRequest) -> dict[str, Any]: "status_code": 409, "message": "Recording is currently active. Stop it first.", } + if _replay.replay_active: + return { + "success": False, + "status_code": 409, + "message": "Replay is currently active. Stop it first.", + } if inference_active: return { "success": False, diff --git a/lelab/server.py b/lelab/server.py index d5fb7372..57943acc 100644 --- a/lelab/server.py +++ b/lelab/server.py @@ -62,6 +62,12 @@ handle_stop_recording, handle_upload_dataset, ) +from .replay import ( + ReplayRequest, + handle_replay_status, + handle_start_replay, + handle_stop_replay, +) from .rollout import ( InferenceRequest, handle_inference_status, @@ -354,6 +360,22 @@ def teleoperation_status(): return handle_teleoperation_status() +@app.post("/start-replay") +def start_replay(request: ReplayRequest): + """Replay a recorded episode on the follower arm""" + return handle_start_replay(request) + + +@app.post("/stop-replay") +def stop_replay(): + return handle_stop_replay() + + +@app.get("/replay-status") +def replay_status(): + return handle_replay_status() + + @app.get("/joint-positions") def get_joint_positions(): """Get current robot joint positions""" diff --git a/lelab/teleoperate.py b/lelab/teleoperate.py index c4d17d04..5a1dd43c 100644 --- a/lelab/teleoperate.py +++ b/lelab/teleoperate.py @@ -118,7 +118,7 @@ def handle_start_teleoperation(request: TeleoperateRequest, websocket_manager=No """ global teleoperation_active, teleoperation_thread, current_robot, current_teleop - from . import record as _record, rollout as _rollout + from . import record as _record, replay as _replay, rollout as _rollout with _state_lock: if teleoperation_active: @@ -127,6 +127,8 @@ def handle_start_teleoperation(request: TeleoperateRequest, websocket_manager=No return {"success": False, "message": "Recording is currently active. Stop it first."} if _rollout.inference_active: return {"success": False, "message": "Inference is currently active. Stop it first."} + if _replay.replay_active: + return {"success": False, "message": "Replay is currently active. Stop it first."} teleoperation_active = True robot = None diff --git a/tests/test_replay.py b/tests/test_replay.py new file mode 100644 index 00000000..6e2ef36f --- /dev/null +++ b/tests/test_replay.py @@ -0,0 +1,123 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for lelab.replay: request schema, mutual exclusion with the other +modes, and the start/stop paths that never move the arm.""" + +from __future__ import annotations + +import pytest + + +def _replay_request(): + from lelab.replay import ReplayRequest + + return ReplayRequest( + follower_port="/dev/follower", + follower_config="follower", + dataset_repo_id="local/ds", + episode_index=0, + ) + + +def test_replay_request_requires_an_episode() -> None: + from pydantic import ValidationError + + from lelab.replay import ReplayRequest + + with pytest.raises(ValidationError): + ReplayRequest(follower_port="/dev/follower", follower_config="follower", dataset_repo_id="local/ds") + + +@pytest.mark.parametrize( + ("module", "flag"), + [ + ("record", "recording_active"), + ("teleoperate", "teleoperation_active"), + ("rollout", "inference_active"), + ], +) +def test_replay_refuses_to_start_while_the_arm_is_busy( + module: str, flag: str, monkeypatch: pytest.MonkeyPatch +) -> None: + import importlib + + import lelab.replay as replay + + monkeypatch.setattr(importlib.import_module(f"lelab.{module}"), flag, True) + + result = replay.handle_start_replay(_replay_request()) + assert result["success"] is False + assert "Stop it first" in result["message"] + assert replay.replay_active is False + + +def test_other_modes_refuse_to_start_during_a_replay(monkeypatch: pytest.MonkeyPatch) -> None: + import lelab.replay as replay + from lelab.record import RecordingRequest, handle_start_recording + from lelab.rollout import InferenceRequest, handle_start_inference + from lelab.teleoperate import TeleoperateRequest, handle_start_teleoperation + + monkeypatch.setattr(replay, "replay_active", True) + ports = {"leader_port": "/dev/leader", "follower_port": "/dev/follower"} + configs = {"leader_config": "leader", "follower_config": "follower"} + + results = [ + handle_start_teleoperation(TeleoperateRequest(**ports, **configs)), + handle_start_recording( + RecordingRequest(**ports, **configs, dataset_repo_id="local/ds", single_task="t") + ), + handle_start_inference( + InferenceRequest(follower_port="/dev/follower", follower_config="follower", policy_ref="p") + ), + ] + for result in results: + assert result["success"] is False + assert "Replay" in result["message"] + + +def test_start_replay_reports_connection_failure(monkeypatch: pytest.MonkeyPatch) -> None: + import lelab.replay as replay + + monkeypatch.setattr(replay, "_episode_actions", lambda repo_id, episode: ([{"gripper.pos": 0.0}], 30)) + monkeypatch.setattr(replay, "setup_follower_calibration_file", lambda config: "follower") + + class _Bus: + def connect(self) -> None: + raise RuntimeError("serial port unavailable") + + followers = [] + + class _Follower: + def __init__(self, config) -> None: + self.bus = _Bus() + self.disconnected = False + followers.append(self) + + def disconnect(self) -> None: + self.disconnected = True + + monkeypatch.setattr(replay, "SO101Follower", _Follower) + + result = replay.handle_start_replay(_replay_request()) + + assert result["success"] is False + assert "/dev/follower" in result["message"] + assert followers[0].disconnected is True + assert replay.replay_active is False + + +def test_stop_replay_when_idle() -> None: + from lelab.replay import handle_stop_replay + + assert handle_stop_replay()["success"] is False From 287a5da627af86a9e7b3f54a6f15e71cff73bdcd Mon Sep 17 00:00:00 2001 From: Nicolas Rabault Date: Thu, 8 Oct 2026 00:31:49 +0200 Subject: [PATCH 2/2] fix(replay): claim the follower under one lock and keep Stop until release Recording, teleoperation, inference and replay each checked the others' flags under their own lock, so two concurrent starts could both claim the follower. They now check and claim under a shared follower_lock. Stopping a replay reports a failure when the worker has not released the arm within 5s, and the button only turns back to Play once the status poll sees the arm released. A running replay keeps its Stop button even when no calibrated robot is selected anymore. ReplayButton gets component tests for restoring a running replay, starting, stopping, a failed poll and a refused start. --- .../components/dataset/ReplayButton.test.tsx | 134 ++++++++++++++++++ .../src/components/dataset/ReplayButton.tsx | 59 ++++---- lelab/record.py | 7 +- lelab/replay.py | 9 +- lelab/rollout.py | 3 +- lelab/teleoperate.py | 6 +- lelab/utils/devices.py | 7 +- tests/test_replay.py | 19 +++ 8 files changed, 201 insertions(+), 43 deletions(-) create mode 100644 frontend/src/components/dataset/ReplayButton.test.tsx diff --git a/frontend/src/components/dataset/ReplayButton.test.tsx b/frontend/src/components/dataset/ReplayButton.test.tsx new file mode 100644 index 00000000..d5d2713e --- /dev/null +++ b/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(); + }); +} + +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(); +}); diff --git a/frontend/src/components/dataset/ReplayButton.tsx b/frontend/src/components/dataset/ReplayButton.tsx index c71658b2..6950b6ff 100644 --- a/frontend/src/components/dataset/ReplayButton.tsx +++ b/frontend/src/components/dataset/ReplayButton.tsx @@ -65,6 +65,38 @@ const ReplayButton: React.FC = ({ repoId, episodeIndex, onSta 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(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 ( + + ); + } + if (!robot?.is_clean) return null; const start = async () => { @@ -93,33 +125,6 @@ const ReplayButton: React.FC = ({ repoId, episodeIndex, onSta } }; - const stop = async () => { - try { - await apiRequest(baseUrl, fetchWithHeaders, "/stop-replay", { method: "POST", action: "Stop" }); - setPhase("idle"); - } catch (e) { - toast({ - title: "Could not stop the robot", - description: e instanceof Error ? e.message : String(e), - variant: "destructive", - }); - } - }; - - if (phase === "playing") { - return ( - - ); - } - return (