From 0e6ac4c4b1f5776737e7fa9006d2a2cefd6cad00 Mon Sep 17 00:00:00 2001 From: "ravindu.somawansa" Date: Tue, 20 Oct 2026 09:34:06 +0200 Subject: [PATCH] feat(datasets): edit an episode's task in the dataset browser --- .../src/components/dataset/EpisodeList.tsx | 57 +++++--- .../dataset/EpisodeTaskDialog.test.tsx | 61 ++++++++ .../components/dataset/EpisodeTaskDialog.tsx | 130 +++++++++++++++++ frontend/src/hooks/useEpisodes.ts | 9 +- frontend/src/lib/replayApi.ts | 22 +++ frontend/src/pages/EditDataset.tsx | 31 +++- lelab/datasets.py | 60 +++++++- lelab/server.py | 22 ++- tests/test_datasets.py | 133 ++++++++++++++++++ 9 files changed, 497 insertions(+), 28 deletions(-) create mode 100644 frontend/src/components/dataset/EpisodeTaskDialog.test.tsx create mode 100644 frontend/src/components/dataset/EpisodeTaskDialog.tsx diff --git a/frontend/src/components/dataset/EpisodeList.tsx b/frontend/src/components/dataset/EpisodeList.tsx index 120a24aca..36186b60a 100644 --- a/frontend/src/components/dataset/EpisodeList.tsx +++ b/frontend/src/components/dataset/EpisodeList.tsx @@ -1,13 +1,16 @@ import React, { useEffect, useRef } from "react"; +import { Pencil } from "lucide-react"; import { EpisodeSummary, formatDuration } from "@/lib/datasetApi"; interface EpisodeListProps { episodes: EpisodeSummary[]; selected: number | null; onSelect: (episodeIndex: number) => void; + /** Shows an edit button on the selected episode's task when set. */ + onEditTask?: (episodeIndex: number) => void; } -const EpisodeList: React.FC = ({ episodes, selected, onSelect }) => { +const EpisodeList: React.FC = ({ episodes, selected, onSelect, onEditTask }) => { const activeRef = useRef(null); // Follow the selection. Paging with Previous/Next, or opening a link straight @@ -39,29 +42,43 @@ const EpisodeList: React.FC = ({ episodes, selected, onSelect {episodes.map((ep) => { const active = ep.episode_index === selected; return ( - + + {ep.episode_index} + + + {ep.tasks.length > 0 ? ep.tasks.join(", ") : "—"} + + + {formatDuration(ep.duration_s)} + + + {active && onEditTask && ( + + )} + ); })} diff --git a/frontend/src/components/dataset/EpisodeTaskDialog.test.tsx b/frontend/src/components/dataset/EpisodeTaskDialog.test.tsx new file mode 100644 index 000000000..71c0ec3a6 --- /dev/null +++ b/frontend/src/components/dataset/EpisodeTaskDialog.test.tsx @@ -0,0 +1,61 @@ +// @vitest-environment jsdom +import React, { act } from "react"; +import { createRoot } from "react-dom/client"; +import { expect, test } from "vitest"; + +import EpisodeTaskDialog from "./EpisodeTaskDialog"; + +test("saves the task for the browsed episode or for every episode", async () => { + Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true }); + const host = document.createElement("div"); + document.body.append(host); + const root = createRoot(host); + const saved: [string, number | null][] = []; + + try { + await act(async () => { + root.render( + {}} + onSave={async (task, episodeIndex) => { + saved.push([task, episodeIndex]); + }} + />, + ); + }); + + const input = document.getElementById("episode-task") as HTMLInputElement; + const saveButton = () => + [...document.body.querySelectorAll("button")].find( + (b) => b.textContent === "Save task", + ) as HTMLButtonElement; + const type = async (value: string) => { + const setter = Object.getOwnPropertyDescriptor(HTMLInputElement.prototype, "value")!.set!; + await act(async () => { + setter.call(input, value); + input.dispatchEvent(new Event("input", { bubbles: true })); + }); + }; + + expect(input.value).toBe("pick the cube"); + + await type(" "); + expect(saveButton().disabled).toBe(true); + + await type(" pick the bowl "); + await act(async () => saveButton().click()); + expect(saved).toEqual([["pick the bowl", 3]]); + + await act(async () => (document.getElementById("episode-task-all") as HTMLButtonElement).click()); + await act(async () => saveButton().click()); + expect(saved.at(-1)).toEqual(["pick the bowl", null]); + } finally { + await act(async () => root.unmount()); + host.remove(); + } +}); diff --git a/frontend/src/components/dataset/EpisodeTaskDialog.tsx b/frontend/src/components/dataset/EpisodeTaskDialog.tsx new file mode 100644 index 000000000..bc7d3bf02 --- /dev/null +++ b/frontend/src/components/dataset/EpisodeTaskDialog.tsx @@ -0,0 +1,130 @@ +import { useEffect, useState } from "react"; +import { Loader2 } from "lucide-react"; + +import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; + +interface EpisodeTaskDialogProps { + repoId: string; + episodeIndex: number; + currentTask: string; + totalEpisodes: number; + open: boolean; + onOpenChange: (open: boolean) => void; + /** `episodeIndex` is null when the task applies to every episode. */ + onSave: (task: string, episodeIndex: number | null) => Promise; +} + +export default function EpisodeTaskDialog({ + repoId, + episodeIndex, + currentTask, + totalEpisodes, + open, + onOpenChange, + onSave, +}: EpisodeTaskDialogProps) { + const [task, setTask] = useState(""); + const [allEpisodes, setAllEpisodes] = useState(false); + const [saving, setSaving] = useState(false); + const [error, setError] = useState(null); + + useEffect(() => { + if (open) { + setTask(currentTask); + setAllEpisodes(false); + setSaving(false); + setError(null); + } + }, [open, currentTask]); + + const canSave = task.trim() !== "" && !saving; + + const handleOpenChange = (nextOpen: boolean) => { + if (saving) return; + onOpenChange(nextOpen); + }; + + const handleSave = async () => { + if (!canSave) return; + setSaving(true); + try { + await onSave(task.trim(), allEpisodes ? null : episodeIndex); + onOpenChange(false); + } catch (caughtError) { + setError(caughtError instanceof Error ? caughtError.message : "Could not save the task."); + } finally { + setSaving(false); + } + }; + + return ( + + + + Edit task + + The task is the instruction a policy is trained on, saved with each episode of{" "} + {repoId}. + + + +
+
+ + { + setError(null); + setTask(event.target.value); + }} + onKeyDown={(event) => { + if (event.key === "Enter") void handleSave(); + }} + disabled={saving} + className="border-gray-700 bg-gray-800 text-white" + /> +
+ {totalEpisodes > 1 && ( +
+ setAllEpisodes(checked === true)} + disabled={saving} + className="border-gray-500" + /> + +
+ )} + {error &&

{error}

} +
+ + + + + + + + ); +} diff --git a/frontend/src/hooks/useEpisodes.ts b/frontend/src/hooks/useEpisodes.ts index 4158ab4d6..791c6dddf 100644 --- a/frontend/src/hooks/useEpisodes.ts +++ b/frontend/src/hooks/useEpisodes.ts @@ -1,4 +1,4 @@ -import { useEffect, useState } from "react"; +import { useCallback, useEffect, useState } from "react"; import { useApi } from "@/contexts/ApiContext"; import { EpisodeDetail, @@ -13,6 +13,7 @@ export const useEpisodes = (repoId: string | null) => { const [data, setData] = useState(null); const [loading, setLoading] = useState(false); const [error, setError] = useState(null); + const [version, setVersion] = useState(0); useEffect(() => { setData(null); @@ -44,11 +45,13 @@ export const useEpisodes = (repoId: string | null) => { if (!controller.signal.aborted) setLoading(false); }); return () => controller.abort(); - }, [baseUrl, fetchWithHeaders, repoId]); + }, [baseUrl, fetchWithHeaders, repoId, version]); + + const refresh = useCallback(() => setVersion((v) => v + 1), []); // Effects run after render: do not let the new dataset auto-select an // episode from the old list before this effect clears it. - return { data: data?.repo_id === repoId ? data : null, loading, error }; + return { data: data?.repo_id === repoId ? data : null, loading, error, refresh }; }; /** One episode's detail. Either arg null → idle. */ diff --git a/frontend/src/lib/replayApi.ts b/frontend/src/lib/replayApi.ts index 21c591ec2..576600364 100644 --- a/frontend/src/lib/replayApi.ts +++ b/frontend/src/lib/replayApi.ts @@ -55,3 +55,25 @@ export async function mergeLocalDatasets( action: "Merge datasets", }); } + +export interface EpisodeTaskResult { + success: true; + repo_id: string; + task: string; + updated_episodes: number; +} + +/** Set the task of one episode, or of every episode when `episodeIndex` is null. */ +export async function setEpisodeTask( + baseUrl: string, + fetcher: Fetcher, + repoId: string, + task: string, + episodeIndex: number | null, +): Promise { + return apiRequest(baseUrl, fetcher, "/datasets/episode-task", { + method: "POST", + body: { repo_id: repoId, task, episode_index: episodeIndex }, + action: "Edit task", + }); +} diff --git a/frontend/src/pages/EditDataset.tsx b/frontend/src/pages/EditDataset.tsx index 53ec628a1..f2a6cd559 100644 --- a/frontend/src/pages/EditDataset.tsx +++ b/frontend/src/pages/EditDataset.tsx @@ -15,6 +15,7 @@ import { import { Button } from "@/components/ui/button"; import DeleteDatasetDialog from "@/components/dataset/DeleteDatasetDialog"; import EpisodeList from "@/components/dataset/EpisodeList"; +import EpisodeTaskDialog from "@/components/dataset/EpisodeTaskDialog"; import EpisodeViewer from "@/components/dataset/EpisodeViewer"; import MergeDatasetsDialog from "@/components/dataset/MergeDatasetsDialog"; import DatasetPicker from "@/components/landing/DatasetPicker"; @@ -27,7 +28,7 @@ import { useRobots } from "@/hooks/useRobots"; import { useToast } from "@/hooks/use-toast"; import { formatDuration } from "@/lib/datasetApi"; import { openHubViewer } from "@/lib/hubViewer"; -import { DatasetItem, mergeLocalDatasets } from "@/lib/replayApi"; +import { DatasetItem, mergeLocalDatasets, setEpisodeTask } from "@/lib/replayApi"; // The robot_type LeRobot stores for the SO-101 follower. const FOLLOWER_ROBOT_TYPE = "so_follower"; @@ -55,6 +56,7 @@ const EditDataset = () => { const { toast } = useToast(); const [mergeOpen, setMergeOpen] = useState(false); const [deleteOpen, setDeleteOpen] = useState(false); + const [taskOpen, setTaskOpen] = useState(false); // A Hub-only dataset has no videos on disk to decode. const localDatasets = useMemo( () => datasets.filter((d) => d.source === "local" || d.source === "both"), @@ -67,7 +69,7 @@ const EditDataset = () => { const selectedSource = datasets.find((d) => d.repo_id === repoId)?.source; const selectedIsLocal = selectedSource === "local" || selectedSource === "both"; - const { data, loading, error } = useEpisodes(repoId); + const { data, loading, error, refresh: refreshEpisodes } = useEpisodes(repoId); const { detail, loading: detailLoading, @@ -110,6 +112,18 @@ const EditDataset = () => { setDataset(result.repo_id); toast({ title: "Datasets merged", description: `${result.num_episodes} episodes are ready to browse.` }); }; + const handleSaveTask = async (task: string, targetEpisode: number | null) => { + if (!repoId) return; + const result = await setEpisodeTask(baseUrl, fetchWithHeaders, repoId, task, targetEpisode); + refreshEpisodes(); + toast({ + title: "Task saved", + description: + targetEpisode === null + ? `All ${result.updated_episodes} episodes now have the task "${result.task}".` + : `Episode ${targetEpisode} now has the task "${result.task}".`, + }); + }; const setEpisode = (next: number) => { if (!repoId) return; setSearchParams(withApi({ dataset: repoId, episode: String(next) })); @@ -284,6 +298,7 @@ const EditDataset = () => { episodes={data.episodes} selected={episodeIndex} onSelect={setEpisode} + onEditTask={selectedIsLocal ? () => setTaskOpen(true) : undefined} />
@@ -317,6 +332,18 @@ const EditDataset = () => { /> )} + {repoId && selectedIsLocal && data?.success && episodeIndex !== null && ( + ep.episode_index === episodeIndex)?.tasks[0] ?? ""} + totalEpisodes={data.total_episodes} + open={taskOpen} + onOpenChange={setTaskOpen} + onSave={handleSaveTask} + /> + )} + {repoId && selectedIsLocal && ( } +def handle_set_episode_task(repo_id: str, task: str, episode_index: int | None = None) -> dict[str, Any]: + """Set the task of one episode of a local dataset, or of every episode when `episode_index` is None. + + LeRobot's `modify_tasks` rewrites the dataset in place, so it runs on a + staged copy that is swapped in afterwards: a failure leaves the source as it + was. Only `meta/` and `data/` are rewritten, so everything else (videos, + images) is hard-linked into the copy where the filesystem allows it. + """ + task = task.strip() + if not task: + raise ValueError("The task cannot be empty") + + try: + episode_media.ensure_episode_index(repo_id) + except episode_media.UnreadableDatasetError as exc: + raise ValueError(str(exc)) from exc + source_dir = episode_media.resolve_dataset_dir(repo_id) + + total = episode_media.read_info(source_dir).get("total_episodes", 0) + if episode_index is not None and not 0 <= episode_index < total: + raise ValueError(f"Episode {episode_index} is not in {repo_id}") + + def copy_or_link(src: str, dst: str) -> None: + if Path(src).relative_to(source_dir).parts[0] in ("meta", "data"): + shutil.copy2(src, dst) + else: + try: + os.link(src, dst) + except OSError: # e.g. a filesystem without hard links + shutil.copy2(src, dst) + + staging = Path(tempfile.mkdtemp(prefix=f".{source_dir.name}.task-", dir=source_dir.parent)) + edited_dir = staging / "edited" + try: + shutil.copytree(source_dir, edited_dir, copy_function=copy_or_link) + dataset = LeRobotDataset(repo_id=repo_id, root=edited_dir) + if episode_index is None: + modify_tasks(dataset, new_task=task) + else: + modify_tasks(dataset, episode_tasks={episode_index: task}) + source_dir.rename(staging / "old") + try: + edited_dir.rename(source_dir) + except OSError: + (staging / "old").rename(source_dir) + raise + finally: + shutil.rmtree(staging, ignore_errors=True) + + return { + "success": True, + "repo_id": repo_id, + "task": task, + "updated_episodes": total if episode_index is None else 1, + } + + # ── episode browsing ──────────────────────────────────────────────────────── # # Everything below reads a *local* dataset only. Browsing a Hub-only dataset diff --git a/lelab/server.py b/lelab/server.py index d5fb7372c..5caa57d51 100644 --- a/lelab/server.py +++ b/lelab/server.py @@ -24,13 +24,13 @@ import threading import time from pathlib import Path -from typing import Any, Literal +from typing import Annotated, Any, Literal from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import FileResponse, JSONResponse, StreamingResponse from fastapi.staticfiles import StaticFiles -from pydantic import BaseModel +from pydantic import BaseModel, Field from starlette.datastructures import Headers from starlette.exceptions import HTTPException as StarletteHTTPException from starlette.responses import Response @@ -148,6 +148,13 @@ class MergeDatasetsRequest(BaseModel): output_name: str +class EpisodeTaskRequest(BaseModel): + repo_id: str + task: Annotated[str, Field(min_length=1)] + # None sets the task of every episode. + episode_index: Annotated[int, Field(ge=0)] | None = None + + # Cache for HF Jobs hardware flavors (5-minute TTL) _flavors_cache: dict = {"data": None, "fetched_at": 0.0} _FLAVOR_CACHE_TTL_SECONDS = 300.0 @@ -434,6 +441,17 @@ def merge_local_datasets(request: MergeDatasetsRequest): raise HTTPException(status_code=400, detail=str(exc)) from exc +@app.post("/datasets/episode-task") +def set_episode_task(request: EpisodeTaskRequest): + """Set the task of one episode, or of every episode, of a local dataset.""" + if _record.recording_active: + raise HTTPException(status_code=409, detail="Stop the recording before editing a dataset") + try: + return dataset_browser.handle_set_episode_task(request.repo_id, request.task, request.episode_index) + except (DatasetNotFoundError, ValueError) as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + @app.get("/dataset-episodes") def dataset_episodes(repo_id: str): """List the episodes of a local dataset.""" diff --git a/tests/test_datasets.py b/tests/test_datasets.py index da9dd6ce5..9319f3f1b 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -262,6 +262,139 @@ def test_merge_route_returns_result_and_maps_validation_error( assert response.json()["detail"] == "bad input" +def _record_dataset(repo_id: str, episodes: int) -> Path: + """Write a finished LeRobot dataset whose episode i has every value set to i.""" + import numpy as np + + from lerobot.datasets import LeRobotDataset + + features = {"action": {"dtype": "float32", "shape": (1,), "names": ["a"]}} + dataset = LeRobotDataset.create(repo_id, fps=10, features=features, use_videos=False) + for ep in range(episodes): + for _ in range(3): + dataset.add_frame({"action": np.full(1, ep, dtype=np.float32), "task": f"task {ep}"}) + dataset.save_episode() + dataset.finalize() + return dataset.root + + +def _episode_tasks(repo_id: str) -> list[str]: + from lelab import episode_media + from lerobot.datasets import LeRobotDataset + + dataset = LeRobotDataset(repo_id=repo_id, root=episode_media.resolve_dataset_dir(repo_id)) + return [dataset.meta.episodes[i]["tasks"][0] for i in range(dataset.meta.total_episodes)] + + +def test_set_episode_task_relabels_one_episode(tmp_lerobot_home: Path) -> None: + from lelab.datasets import handle_set_episode_task + + root = _record_dataset("alice/demo", episodes=3) + # Stands in for the videos, which the edit must link rather than copy. + (root / "videos").mkdir() + (root / "videos" / "clip.mp4").write_bytes(b"frames") + inode = (root / "videos" / "clip.mp4").stat().st_ino + + result = handle_set_episode_task("alice/demo", " pick the bowl ", episode_index=1) + + assert result == { + "success": True, + "repo_id": "alice/demo", + "task": "pick the bowl", + "updated_episodes": 1, + } + assert _episode_tasks("alice/demo") == ["task 0", "pick the bowl", "task 2"] + assert (root / "videos" / "clip.mp4").stat().st_ino == inode + assert not list((tmp_lerobot_home / "alice").glob(".demo.task-*")) + + +def test_set_episode_task_relabels_every_episode(tmp_lerobot_home: Path) -> None: + from lelab.datasets import handle_set_episode_task + + _record_dataset("alice/demo", episodes=3) + + result = handle_set_episode_task("alice/demo", "pick the bowl") + + assert result["updated_episodes"] == 3 + assert _episode_tasks("alice/demo") == ["pick the bowl"] * 3 + + +@pytest.mark.parametrize( + ("task", "episode_index", "message"), + [(" ", 0, "cannot be empty"), ("pick", 3, "Episode 3 is not in alice/demo")], +) +def test_set_episode_task_rejects_invalid_request( + tmp_lerobot_home: Path, task: str, episode_index: int, message: str +) -> None: + from lelab.datasets import handle_set_episode_task + + _record_dataset("alice/demo", episodes=3) + with pytest.raises(ValueError, match=message): + handle_set_episode_task("alice/demo", task, episode_index) + assert _episode_tasks("alice/demo") == ["task 0", "task 1", "task 2"] + + +def test_set_episode_task_leaves_source_untouched_on_failure( + tmp_lerobot_home: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from lelab import datasets as datasets_mod + + _record_dataset("alice/demo", episodes=2) + real_modify = datasets_mod.modify_tasks + + def failing_modify(dataset, **kwargs): + real_modify(dataset, **kwargs) + raise RuntimeError("disk full") + + monkeypatch.setattr(datasets_mod, "modify_tasks", failing_modify) + + with pytest.raises(RuntimeError, match="disk full"): + datasets_mod.handle_set_episode_task("alice/demo", "pick the bowl") + + assert _episode_tasks("alice/demo") == ["task 0", "task 1"] + assert not list((tmp_lerobot_home / "alice").glob(".demo.task-*")) + + +def test_episode_task_route_refuses_while_recording( + client: TestClient, monkeypatch: pytest.MonkeyPatch +) -> None: + from lelab import datasets as datasets_mod, record + + monkeypatch.setattr(record, "recording_active", True) + monkeypatch.setattr( + datasets_mod, "handle_set_episode_task", lambda *_: pytest.fail("must not edit while recording") + ) + + response = client.post("/datasets/episode-task", json={"repo_id": "alice/demo", "task": "pick"}) + assert response.status_code == 409 + + +@pytest.mark.parametrize("body", [{"task": ""}, {"task": "pick", "episode_index": -1}]) +def test_episode_task_route_rejects_bad_request_before_editing( + client: TestClient, monkeypatch: pytest.MonkeyPatch, body: dict +) -> None: + from lelab import datasets as datasets_mod + + monkeypatch.setattr(datasets_mod, "handle_set_episode_task", lambda *_: pytest.fail("must not edit")) + response = client.post("/datasets/episode-task", json={"repo_id": "alice/demo", **body}) + assert response.status_code == 422 + + +def test_episode_task_route_maps_validation_error( + client: TestClient, monkeypatch: pytest.MonkeyPatch +) -> None: + from lelab import datasets as datasets_mod + + monkeypatch.setattr( + datasets_mod, "handle_set_episode_task", lambda *_: (_ for _ in ()).throw(ValueError("bad input")) + ) + response = client.post( + "/datasets/episode-task", json={"repo_id": "alice/demo", "task": "pick", "episode_index": 0} + ) + assert response.status_code == 400 + assert response.json()["detail"] == "bad input" + + # ── episode browsing ──────────────────────────────────────────────────────── # # The handlers are covered end-to-end through the routes, so the query-param