diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index cc3cf52c..4994dc57 100644 --- a/frontend/src/App.tsx +++ b/frontend/src/App.tsx @@ -11,6 +11,7 @@ import Recording from "@/pages/Recording"; import Training from "@/pages/Training"; import Inference from "@/pages/Inference"; import EditDataset from "@/pages/EditDataset"; +import RerecordEpisode from "@/pages/RerecordEpisode"; import Upload from "@/pages/Upload"; import NotFound from "@/pages/NotFound"; @@ -46,6 +47,7 @@ function App() { } /> } /> } /> + } /> } /> diff --git a/frontend/src/lib/jobsApi.ts b/frontend/src/lib/jobsApi.ts index 088b727b..a5289439 100644 --- a/frontend/src/lib/jobsApi.ts +++ b/frontend/src/lib/jobsApi.ts @@ -200,6 +200,28 @@ export async function stopJob( }); } +export async function pauseJob( + baseUrl: string, + fetcher: Fetcher, + id: string, +): Promise { + return apiRequest(baseUrl, fetcher, `/jobs/${id}/pause`, { + method: "POST", + action: "Pause job", + }); +} + +export async function resumeJob( + baseUrl: string, + fetcher: Fetcher, + id: string, +): Promise { + return apiRequest(baseUrl, fetcher, `/jobs/${id}/resume`, { + method: "POST", + action: "Resume job", + }); +} + export async function deleteJob( baseUrl: string, fetcher: Fetcher, diff --git a/frontend/src/pages/Landing.tsx b/frontend/src/pages/Landing.tsx index 33832569..8d834114 100644 --- a/frontend/src/pages/Landing.tsx +++ b/frontend/src/pages/Landing.tsx @@ -1,6 +1,6 @@ import React, { useState, useRef, useEffect } from "react"; import { useNavigate } from "react-router-dom"; -import { ChevronsUpDown } from "lucide-react"; +import { ChevronsUpDown, Pencil, RefreshCw } from "lucide-react"; import { Button } from "@/components/ui/button"; import { useToast } from "@/hooks/use-toast"; import LandingTopBar from "@/components/landing/LandingTopBar"; @@ -266,6 +266,27 @@ const Landing = () => { +
+

+ Edit Datasets +

+ + +

Create a model diff --git a/frontend/src/pages/Recording.tsx b/frontend/src/pages/Recording.tsx index 0520130e..9caebfa9 100644 --- a/frontend/src/pages/Recording.tsx +++ b/frontend/src/pages/Recording.tsx @@ -14,6 +14,7 @@ import { RotateCcw, Square, SkipForward, + Pause, Play, Volume2, VolumeX, @@ -54,10 +55,11 @@ interface RecordingConfig { streaming_encoding: boolean; } -type Phase = "preparing" | "recording" | "resetting" | "completed"; +type Phase = "preparing" | "recording" | "resetting" | "paused" | "completed"; interface BackendStatus { recording_active: boolean; + recording_paused: boolean; current_phase: string; current_episode?: number; total_episodes?: number; @@ -73,6 +75,8 @@ interface BackendStatus { stop_recording: boolean; exit_early: boolean; rerecord_episode: boolean; + pause_recording: boolean; + resume_recording: boolean; }; } @@ -243,23 +247,6 @@ const Recording = () => { } if (!status.recording_active && status.session_ended) { - // A failure can land after episodes were already saved, so surface - // the reason either way and only go home when nothing survived it. - if (status.current_phase === "error") { - const saved = status.saved_episodes || 0; - toast({ - title: saved > 0 ? "Recording Interrupted" : "Recording Failed", - description: - saved > 0 - ? `${saved} episode(s) were saved before the session failed: ${status.error || "unknown error"}` - : status.error || "The recording session failed to start.", - variant: "destructive", - }); - if (saved === 0) { - navigate("/"); - return; - } - } const datasetInfo = { dataset_repo_id: status.dataset_repo_id || recordingConfig.dataset_repo_id, @@ -278,7 +265,7 @@ const Recording = () => { pollStatus(); const statusInterval = setInterval(pollStatus, 1000); return () => clearInterval(statusInterval); - }, [recordingSessionStarted, recordingConfig, navigate, baseUrl, fetchWithHeaders, toast]); + }, [recordingSessionStarted, recordingConfig, navigate, baseUrl, fetchWithHeaders]); const formatTime = (seconds: number): string => { const mins = Math.floor(seconds / 60); @@ -392,6 +379,26 @@ const Recording = () => { } }, [backendStatus, baseUrl, fetchWithHeaders, toast]); + const handlePauseRecording = useCallback(async () => { + if (!backendStatus?.available_controls.pause_recording) return; + try { + await fetchWithHeaders(`${baseUrl}/pause-recording`, { method: "POST" }); + toast({ title: "Recording paused", description: "Current episode saved. Click Resume to continue." }); + } catch { + toast({ title: "Error", description: "Failed to pause recording.", variant: "destructive" }); + } + }, [backendStatus, baseUrl, fetchWithHeaders, toast]); + + const handleResumeRecording = useCallback(async () => { + if (!backendStatus?.available_controls.resume_recording) return; + try { + await fetchWithHeaders(`${baseUrl}/resume-recording`, { method: "POST" }); + toast({ title: "Recording resumed" }); + } catch { + toast({ title: "Error", description: "Failed to resume recording.", variant: "destructive" }); + } + }, [backendStatus, baseUrl, fetchWithHeaders, toast]); + const handleStopRecording = useCallback(async () => { if (!backendStatus?.available_controls.stop_recording) return; try { @@ -510,6 +517,7 @@ const Recording = () => { if (currentPhase === "recording") return `RECORDING EPISODE ${currentEpisode}`; if (currentPhase === "resetting") return "RESET — GET READY"; if (currentPhase === "preparing") return "PREPARING SESSION"; + if (currentPhase === "paused") return "⏸ PAUSED — CLICK RESUME TO CONTINUE"; return "SESSION COMPLETE"; }; @@ -518,6 +526,8 @@ const Recording = () => { ? { dot: "bg-red-500", pill: "bg-red-500/15 text-red-300", timer: "text-green-400", bar: "bg-green-500", button: "bg-green-500 hover:bg-green-600" } : currentPhase === "resetting" ? { dot: "bg-orange-500", pill: "bg-orange-500/15 text-orange-300", timer: "text-orange-400", bar: "bg-orange-500", button: "bg-orange-500 hover:bg-orange-600" } + : currentPhase === "paused" + ? { dot: "bg-amber-400", pill: "bg-amber-400/15 text-amber-300", timer: "text-amber-400", bar: "bg-amber-400", button: "bg-amber-500 hover:bg-amber-600" } : { dot: "bg-gray-500", pill: "bg-gray-500/15 text-gray-300", timer: "text-gray-400", bar: "bg-gray-500", button: "bg-gray-500" }; const primaryLabel = @@ -560,6 +570,30 @@ const Recording = () => { > {muted ? : } + {backendStatus.available_controls.pause_recording && ( + + )} + {backendStatus.available_controls.resume_recording && ( + + )}

); }; - -export default Recording; diff --git a/frontend/src/pages/RerecordEpisode.tsx b/frontend/src/pages/RerecordEpisode.tsx new file mode 100644 index 00000000..678b08b5 --- /dev/null +++ b/frontend/src/pages/RerecordEpisode.tsx @@ -0,0 +1,771 @@ +import React, { useCallback, useEffect, useRef, useState } from "react"; +import { useNavigate } from "react-router-dom"; +import { + ArrowLeft, + Check, + Loader2, + Play, + Plus, + RefreshCw, + Trash2, + Video, + X, +} from "lucide-react"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { NumberInput } from "@/components/ui/number-input"; +import { useApi } from "@/contexts/ApiContext"; +import { useToast } from "@/hooks/use-toast"; +import { useRobots } from "@/hooks/useRobots"; +import CameraConfiguration, { + CameraConfig, +} from "@/components/recording/CameraConfiguration"; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "@/components/ui/select"; +import { + Dialog, + DialogContent, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; + +// ─── Types ──────────────────────────────────────────────────────────────────── + +interface DatasetMeta { + repo_id: string; + num_episodes: number | null; + num_frames: number | null; + loadable: boolean; +} + +interface EpisodeInfo { + episode_index: number; + length: number; + task: string; +} + +interface CameraVideoInfo { + key: string; + rel_path: string; + from_timestamp: number; + to_timestamp: number; +} + +// ─── Replay modal ───────────────────────────────────────────────────────────── + +function EpisodeReplayModal({ + repoId, + episodeIndex, + baseUrl, + fetchWithHeaders, + onClose, +}: { + repoId: string; + episodeIndex: number; + baseUrl: string; + fetchWithHeaders: typeof fetch; + onClose: () => void; +}) { + const [cameras, setCameras] = useState([]); + const [loading, setLoading] = useState(true); + const [error, setError] = useState(null); + + useEffect(() => { + setLoading(true); + setError(null); + fetchWithHeaders( + `${baseUrl}/dataset-episode-video-info?repo_id=${encodeURIComponent(repoId)}&episode_index=${episodeIndex}` + ) + .then((r) => r.json()) + .then((data) => { + if (data.success) setCameras(data.cameras ?? []); + else setError(data.message ?? "Failed to load video info"); + }) + .catch((e) => setError(String(e))) + .finally(() => setLoading(false)); + }, [repoId, episodeIndex, baseUrl, fetchWithHeaders]); + + return ( + { if (!open) onClose(); }}> + + + + + Episode #{episodeIndex} + + + + {loading ? ( +
+ +
+ ) : error ? ( +
{error}
+ ) : cameras.length === 0 ? ( +
No video found for this episode.
+ ) : ( +
1 ? "grid-cols-2" : "grid-cols-1"}`}> + {cameras.map((cam) => { + const shortKey = cam.key.split(".").pop() ?? cam.key; + const src = `${baseUrl}/dataset-video-file?repo_id=${encodeURIComponent(repoId)}&path=${encodeURIComponent(cam.rel_path)}#t=${cam.from_timestamp},${cam.to_timestamp}`; + return ( +
+
+ ); + })} +
+ )} + + + ); +} + +// ─── Sub-components ────────────────────────────────────────────────────────── + +function EpisodeCard({ + ep, + selected, + onToggle, + onPlay, +}: { + ep: EpisodeInfo; + selected: boolean; + onToggle: () => void; + onPlay: () => void; +}) { + return ( +
+
+ +
+ {ep.length} frames + +
+
+ {ep.task && ( +

{ep.task}

+ )} +
+ ); +} + +// ─── Main Page ──────────────────────────────────────────────────────────────── + +const RerecordEpisode: React.FC = () => { + const navigate = useNavigate(); + const { baseUrl, fetchWithHeaders } = useApi(); + const { toast } = useToast(); + const { selectedRecord, availableNames, records, selectRobot, isLoading: robotsLoading } = useRobots(); + + // ── Dataset selection ── + const [datasets, setDatasets] = useState([]); + const [datasetsLoading, setDatasetsLoading] = useState(false); + const [selectedDataset, setSelectedDataset] = useState(null); + + // ── Episode list ── + const [episodes, setEpisodes] = useState([]); + const [episodesLoading, setEpisodesLoading] = useState(false); + const [selectedEpisodes, setSelectedEpisodes] = useState>(new Set()); + const [replayEpisodeIndex, setReplayEpisodeIndex] = useState(null); + + // ── Recording config ── + const [singleTask, setSingleTask] = useState(""); + const [numNewEpisodes, setNumNewEpisodes] = useState(1); + const [episodeTimeS, setEpisodeTimeS] = useState(60); + const [resetTimeS, setResetTimeS] = useState(15); + const [cameras, setCameras] = useState([]); + + // ── Action state ── + const [working, setWorking] = useState(false); + const releaseStreamsRef = useRef<(() => void) | null>(null); + + // Load datasets on mount + useEffect(() => { + setDatasetsLoading(true); + fetchWithHeaders(`${baseUrl}/edit/datasets`) + .then((r) => r.json()) + .then((data: DatasetMeta[]) => + setDatasets(data.filter((d) => d.loadable && (d.num_episodes ?? 0) > 0)) + ) + .catch((e) => console.error("Failed to load datasets:", e)) + .finally(() => setDatasetsLoading(false)); + }, [baseUrl, fetchWithHeaders]); + + // Load episode list when dataset is selected + const loadEpisodes = useCallback( + (repoId: string) => { + setEpisodesLoading(true); + setEpisodes([]); + setSelectedEpisodes(new Set()); + fetchWithHeaders( + `${baseUrl}/dataset-episodes?repo_id=${encodeURIComponent(repoId)}` + ) + .then((r) => r.json()) + .then((data) => { + if (data.success) setEpisodes(data.episodes ?? []); + else + toast({ + title: "Could not load episodes", + description: data.message, + variant: "destructive", + }); + }) + .catch((e) => + toast({ + title: "Error", + description: String(e), + variant: "destructive", + }) + ) + .finally(() => setEpisodesLoading(false)); + }, + [baseUrl, fetchWithHeaders, toast] + ); + + const handleSelectDataset = useCallback( + (repoId: string) => { + setSelectedDataset(repoId); + loadEpisodes(repoId); + // Pre-fill task from dataset episodes if available + }, + [loadEpisodes] + ); + + // Sync cameras from robot profile when robot changes + useEffect(() => { + setCameras(selectedRecord ? [...(selectedRecord.cameras ?? [])] : []); + }, [selectedRecord]); + + // Sync task from first episode + useEffect(() => { + if (episodes.length > 0 && !singleTask) { + const firstTask = episodes[0].task; + if (firstTask) setSingleTask(firstTask); + } + }, [episodes, singleTask]); + + const toggleEpisode = useCallback((idx: number) => { + setSelectedEpisodes((prev) => { + const next = new Set(prev); + if (next.has(idx)) next.delete(idx); + else next.add(idx); + return next; + }); + }, []); + + const selectAll = useCallback(() => { + setSelectedEpisodes(new Set(episodes.map((e) => e.episode_index))); + }, [episodes]); + + const clearAll = useCallback(() => { + setSelectedEpisodes(new Set()); + }, []); + + // Total episodes that will be recorded = selected (to replace) + extra new ones + const totalToRecord = selectedEpisodes.size > 0 ? selectedEpisodes.size + (numNewEpisodes - selectedEpisodes.size > 0 ? numNewEpisodes - selectedEpisodes.size : 0) : numNewEpisodes; + + const canStart = + !!selectedDataset && + !!selectedRecord && + selectedRecord.is_clean && + !!singleTask.trim() && + !working; + + const handleStart = useCallback(async () => { + if (!selectedDataset || !selectedRecord) return; + + setWorking(true); + + try { + // Step 1: Delete selected episodes in-place (if any) + if (selectedEpisodes.size > 0) { + const resp = await fetchWithHeaders( + `${baseUrl}/edit/delete-episodes-inplace`, + { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + dataset_repo_id: selectedDataset, + episode_indices: Array.from(selectedEpisodes), + }), + } + ); + const data = await resp.json(); + if (!data.success) { + toast({ + title: "Delete failed", + description: data.message, + variant: "destructive", + }); + setWorking(false); + return; + } + toast({ + title: "Episodes deleted", + description: data.message, + }); + } + + // Step 2: Build recording config and navigate to Recording page + const cameraDict = cameras.reduce( + (acc, cam) => { + acc[cam.name] = { + type: cam.type, + camera_index: cam.camera_index, + width: cam.width, + height: cam.height, + fps: cam.fps, + ...(cam.fourcc ? { fourcc: cam.fourcc } : {}), + ...(cam.backend ? { backend: cam.backend } : {}), + }; + return acc; + }, + {} as Record + ); + + const nToRecord = + selectedEpisodes.size > 0 + ? Math.max(selectedEpisodes.size, numNewEpisodes) + : numNewEpisodes; + + const recordingConfig = { + leader_port: selectedRecord.leader_port, + follower_port: selectedRecord.follower_port, + leader_config: selectedRecord.leader_config, + follower_config: selectedRecord.follower_config, + dataset_repo_id: selectedDataset, + single_task: singleTask.trim(), + num_episodes: nToRecord, + episode_time_s: episodeTimeS, + reset_time_s: resetTimeS, + fps: 30, + video: true, + push_to_hub: false, + resume: true, + streaming_encoding: true, + cameras: cameraDict, + }; + + // Release any camera preview streams before navigation + if (releaseStreamsRef.current) releaseStreamsRef.current(); + + navigate("/recording", { state: { recordingConfig } }); + } catch (e) { + toast({ title: "Error", description: String(e), variant: "destructive" }); + setWorking(false); + } + }, [ + selectedDataset, + selectedRecord, + selectedEpisodes, + cameras, + numNewEpisodes, + singleTask, + episodeTimeS, + resetTimeS, + baseUrl, + fetchWithHeaders, + navigate, + toast, + ]); + + const shortId = (r: string) => + r.includes("/") ? r.split("/")[1] : r; + + return ( +
+
+ {/* Header */} +
+ +
+

Add / Re-record Episodes

+

+ Delete bad demonstrations and record replacements, or append new + episodes to an existing dataset. +

+
+
+ +
+ {/* Left column: dataset + episode selection */} +
+ {/* Dataset picker */} +
+

+

+ {datasetsLoading ? ( +
+ + Loading datasets… +
+ ) : datasets.length === 0 ? ( +

No local datasets found.

+ ) : ( + + )} + {selectedDataset && ( +

{selectedDataset}

+ )} +
+ + {/* Episode list */} + {selectedDataset && ( +
+
+

+ + Episodes to Delete + {selectedEpisodes.size > 0 && ( + + {selectedEpisodes.size} selected + + )} +

+
+ + + +
+
+ + {episodesLoading ? ( +
+ + Loading episodes… +
+ ) : episodes.length === 0 ? ( +

No episodes found.

+ ) : ( +
+ {episodes.map((ep) => ( + toggleEpisode(ep.episode_index)} + onPlay={() => setReplayEpisodeIndex(ep.episode_index)} + /> + ))} +
+ )} + + {selectedEpisodes.size > 0 && ( +

+ + {selectedEpisodes.size} episode(s) will be permanently + deleted before recording starts. +

+ )} +
+ )} +
+ + {/* Right column: recording config */} +
+ {/* Robot selection */} +
+

Robot

+ {robotsLoading ? ( +
+ + Loading robots… +
+ ) : availableNames.length === 0 ? ( +

+ No robots configured. Set one up on the Landing page first. +

+ ) : ( + + )} + {selectedRecord && !selectedRecord.is_clean && ( +

+ This robot needs calibration before recording. +

+ )} +
+ + {/* Recording parameters */} +
+

Recording Parameters

+ +
+ + setSingleTask(e.target.value)} + placeholder="e.g., pick and place the rubber object" + className="bg-gray-800 border-gray-700 text-white" + /> +
+ +
+ + 0 ? selectedEpisodes.size : 1} + value={numNewEpisodes} + onChange={(v) => { + if (v !== undefined) setNumNewEpisodes(v); + }} + className="bg-gray-800 border-gray-700 text-white" + /> + {selectedEpisodes.size > 0 && numNewEpisodes < selectedEpisodes.size && ( +

+ Will record {selectedEpisodes.size} episodes to replace the + deleted ones. +

+ )} +
+ +
+
+ + { + if (v !== undefined) setEpisodeTimeS(v); + }} + className="bg-gray-800 border-gray-700 text-white" + /> +
+
+ + { + if (v !== undefined) setResetTimeS(v); + }} + className="bg-gray-800 border-gray-700 text-white" + /> +
+
+
+ + {/* Cameras */} +
+ +

+ Cameras are taken from the selected robot profile. +

+
+ + {/* Action */} +
+
+ {selectedEpisodes.size > 0 ? ( + + Will delete{" "} + + {selectedEpisodes.size} + {" "} + episode(s), then record{" "} + + {Math.max(selectedEpisodes.size, numNewEpisodes)} + {" "} + new episode(s). + + ) : ( + + Will append{" "} + + {numNewEpisodes} + {" "} + new episode(s) to the selected dataset. + + )} +
+ + {!selectedRecord && ( +

+ Select a configured robot to enable recording. +

+ )} + {!selectedDataset && ( +

+ Select a dataset to start. +

+ )} +
+
+
+
+ + {replayEpisodeIndex !== null && selectedDataset && ( + setReplayEpisodeIndex(null)} + /> + )} +
+ ); +}; + +export default RerecordEpisode; diff --git a/frontend/src/pages/Training.tsx b/frontend/src/pages/Training.tsx index a361d042..950f92f7 100644 --- a/frontend/src/pages/Training.tsx +++ b/frontend/src/pages/Training.tsx @@ -15,7 +15,7 @@ import HfAuthBanner from "@/components/landing/HfAuthBanner"; import { Button } from "@/components/ui/button"; import { Tooltip, TooltipTrigger, TooltipContent } from "@/components/ui/tooltip"; -import { Loader2, Play, Square, Trash2, ArrowLeft } from "lucide-react"; +import { Loader2, Pause, Play, Square, Trash2, ArrowLeft } from "lucide-react"; import { DatasetItem, listDatasets } from "@/lib/replayApi"; import { @@ -27,6 +27,8 @@ import { listJobs, startTrainingJob, stopJob, + pauseJob, + resumeJob, deleteJob, listRunnerHardware, RunnerFlavor, @@ -50,7 +52,7 @@ function jobToStatus(job: JobRecord | null, isStarting: boolean): TrainingStatus }; } return { - training_active: job.state === "running", + training_active: job.state === "running" || job.state === "paused", current_step: job.metrics.current_step, total_steps: job.metrics.total_steps, current_loss: job.metrics.current_loss ?? undefined, @@ -58,9 +60,9 @@ function jobToStatus(job: JobRecord | null, isStarting: boolean): TrainingStatus grad_norm: job.metrics.grad_norm ?? undefined, eta_seconds: job.metrics.eta_seconds ?? undefined, available_controls: { - stop_training: job.state === "running", - pause_training: false, - resume_training: false, + stop_training: job.state === "running" || job.state === "paused", + pause_training: job.state === "running", + resume_training: job.state === "paused", }, }; } @@ -421,7 +423,7 @@ const MonitoringMode: React.FC<{ jobId: string }> = ({ jobId }) => { tick(); const id = setInterval(() => { if (cancelled) return; - if (jobStateRef.current && jobStateRef.current !== "running") return; + if (jobStateRef.current && jobStateRef.current !== "running" && jobStateRef.current !== "paused") return; tick(); }, 5000); return () => { @@ -456,7 +458,7 @@ const MonitoringMode: React.FC<{ jobId: string }> = ({ jobId }) => { tick(); const id = setInterval(() => { if (cancelled) return; - if (jobStateRef.current && jobStateRef.current !== "running") return; + if (jobStateRef.current && jobStateRef.current !== "running" && jobStateRef.current !== "paused") return; tick(); }, POLL_INTERVAL_MS); return () => { @@ -502,6 +504,36 @@ const MonitoringMode: React.FC<{ jobId: string }> = ({ jobId }) => { } }; + const handlePause = async () => { + if (!job) return; + try { + const next = await pauseJob(baseUrl, fetchWithHeaders, job.id); + setJob(next); + toast({ title: "Training paused" }); + } catch (e) { + toast({ + title: "Pause failed", + description: e instanceof Error ? e.message : String(e), + variant: "destructive", + }); + } + }; + + const handleResume = async () => { + if (!job) return; + try { + const next = await resumeJob(baseUrl, fetchWithHeaders, job.id); + setJob(next); + toast({ title: "Training resumed" }); + } catch (e) { + toast({ + title: "Resume failed", + description: e instanceof Error ? e.message : String(e), + variant: "destructive", + }); + } + }; + const handleDelete = async () => { if (!job) return; if (!window.confirm("Delete this run? This wipes the output directory.")) return; @@ -542,6 +574,8 @@ const MonitoringMode: React.FC<{ jobId: string }> = ({ jobId }) => { } const isRunning = job.state === "running"; + const isPaused = job.state === "paused"; + const isActive = isRunning || isPaused; return (
@@ -590,10 +624,21 @@ const MonitoringMode: React.FC<{ jobId: string }> = ({ jobId }) => {

- {isRunning ? ( - + {isActive ? ( +
+ {isPaused ? ( + + ) : ( + + )} + +
) : (