feat: merge local datasets in the browser

frontend/dist/index.html

@@ -10,8 +10,8 @@
     <meta property="og:title" content="LeLab" />
     <meta property="og:description" content="LeRobot but on your browser." />
     <meta property="og:type" content="website" />
-    <script type="module" crossorigin src="/assets/index-BX9JKcmw.js"></script>
-    <link rel="stylesheet" crossorigin href="/assets/index-DJuBJMLM.css">
+    <script type="module" crossorigin src="/assets/index-Dpbx8D4U.js"></script>
+    <link rel="stylesheet" crossorigin href="/assets/index-Dv87rxaM.css">
   </head>
 
   <body>

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

@@ -0,0 +1,41 @@
+// @vitest-environment jsdom
+import React, { act } from "react";
+import { createRoot } from "react-dom/client";
+import { expect, test } from "vitest";
+
+import MergeDatasetsDialog from "./MergeDatasetsDialog";
+
+test("opens with the browsed local dataset retained as the first merge source", async () => {
+  Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true });
+  const host = document.createElement("div");
+  document.body.append(host);
+  const root = createRoot(host);
+
+  try {
+    await act(async () => {
+      root.render(
+        <MergeDatasetsDialog
+          datasets={[
+            { repo_id: "local/browsed", source: "local" },
+            { repo_id: "local/additional", source: "local" },
+          ]}
+          selectedRepoId="local/browsed"
+          open
+          onOpenChange={() => {}}
+          onMerge={async () => {}}
+        />,
+      );
+    });
+
+    expect(document.body.textContent).toContain("local/browsed is included");
+    expect(document.body.textContent).toContain("local/additional");
+    expect(document.body.textContent).toContain("Select at least one additional dataset.");
+    const mergeButton = [...document.body.querySelectorAll("button")].find((button) =>
+      button.textContent?.includes("Merge datasets"),
+    ) as HTMLButtonElement;
+    expect(mergeButton.disabled).toBe(true);
+  } finally {
+    await act(async () => root.unmount());
+    host.remove();
+  }
+});

frontend/src/components/dataset/MergeDatasetsDialog.tsx

@@ -0,0 +1,150 @@
+import { useEffect, useMemo, 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";
+import { DatasetItem } from "@/lib/replayApi";
+
+interface MergeDatasetsDialogProps {
+  datasets: DatasetItem[];
+  selectedRepoId: string;
+  open: boolean;
+  onOpenChange: (open: boolean) => void;
+  onMerge: (sourceRepoIds: string[], outputName: string) => Promise<void>;
+}
+
+export default function MergeDatasetsDialog({
+  datasets,
+  selectedRepoId,
+  open,
+  onOpenChange,
+  onMerge,
+}: MergeDatasetsDialogProps) {
+  const [selected, setSelected] = useState<string[]>([selectedRepoId]);
+  const [outputName, setOutputName] = useState("");
+  const [merging, setMerging] = useState(false);
+  const [error, setError] = useState<string | null>(null);
+
+  const localDatasets = useMemo(
+    () =>
+      datasets.filter(
+        (dataset) =>
+          dataset.repo_id !== selectedRepoId && (dataset.source === "local" || dataset.source === "both"),
+      ),
+    [datasets, selectedRepoId],
+  );
+  const canMerge = selected.length >= 2 && /^[A-Za-z0-9._-]+$/.test(outputName) && !merging;
+
+  const toggleDataset = (repoId: string, checked: boolean) => {
+    setError(null);
+    setSelected((current) => (checked ? [...current, repoId] : current.filter((id) => id !== repoId)));
+  };
+
+  const reset = () => {
+    setSelected([selectedRepoId]);
+    setOutputName("");
+    setMerging(false);
+    setError(null);
+  };
+
+  useEffect(() => {
+    if (open) {
+      setSelected([selectedRepoId]);
+      setOutputName("");
+      setMerging(false);
+      setError(null);
+    }
+  }, [open, selectedRepoId]);
+
+  const handleOpenChange = (nextOpen: boolean) => {
+    if (!merging && !nextOpen) reset();
+    onOpenChange(nextOpen);
+  };
+
+  const handleMerge = async () => {
+    if (!canMerge) return;
+    setMerging(true);
+    try {
+      await onMerge(selected, outputName);
+      reset();
+      onOpenChange(false);
+    } catch (caughtError) {
+      setError(caughtError instanceof Error ? caughtError.message : "Could not merge the selected datasets.");
+    } finally {
+      setMerging(false);
+    }
+  };
+
+  return (
+    <Dialog open={open} onOpenChange={handleOpenChange}>
+      <DialogContent className="border-gray-700 bg-gray-900 text-gray-300 sm:max-w-lg">
+        <DialogHeader>
+          <DialogTitle className="text-white">Merge local datasets</DialogTitle>
+          <DialogDescription>
+            <code>{selectedRepoId}</code> is included. Select at least one more local dataset. Their source files stay
+            unchanged; the merged dataset is saved locally.
+          </DialogDescription>
+        </DialogHeader>
+
+        <div className="space-y-3">
+          <Label className="text-gray-300">Source datasets</Label>
+          <div className="max-h-52 space-y-2 overflow-y-auto rounded-md border border-gray-700 p-3">
+            <p className="text-sm text-gray-200">{selectedRepoId}</p>
+            {localDatasets.map((dataset) => (
+              <label key={dataset.repo_id} className="flex cursor-pointer items-center gap-3 text-sm text-gray-200">
+                <Checkbox
+                  checked={selected.includes(dataset.repo_id)}
+                  onCheckedChange={(checked) => toggleDataset(dataset.repo_id, checked === true)}
+                  disabled={merging}
+                />
+                {dataset.repo_id}
+              </label>
+            ))}
+          </div>
+          <p className="text-xs text-gray-500">Select at least one additional dataset.</p>
+
+          <div className="space-y-2">
+            <Label htmlFor="merged-dataset-name" className="text-gray-300">
+              Merged dataset name
+            </Label>
+            <Input
+              id="merged-dataset-name"
+              value={outputName}
+              onChange={(event) => {
+                setError(null);
+                setOutputName(event.target.value.replace(/[^A-Za-z0-9._-]/g, "_"));
+              }}
+              placeholder="combined_pick_place"
+              disabled={merging}
+              className="border-gray-700 bg-gray-800 text-white"
+            />
+            <p className="text-xs text-gray-500">
+              Saved locally as <code>local/{outputName || "name"}</code>.
+            </p>
+          </div>
+          {error && <p className="text-sm text-red-400">{error}</p>}
+        </div>
+
+        <DialogFooter>
+          <Button variant="outline" onClick={() => handleOpenChange(false)} disabled={merging}>
+            Cancel
+          </Button>
+          <Button onClick={handleMerge} disabled={!canMerge}>
+            {merging && <Loader2 className="mr-2 h-4 w-4 animate-spin" />}
+            {merging ? "Merging…" : "Merge datasets"}
+          </Button>
+        </DialogFooter>
+      </DialogContent>
+    </Dialog>
+  );
+}

frontend/src/lib/replayApi.test.ts

@@ -0,0 +1,21 @@
+import { expect, test, vi } from "vitest";
+
+import { mergeLocalDatasets } from "./replayApi";
+
+test("posts the selected local datasets and output name", async () => {
+  const fetcher = vi.fn().mockResolvedValue(
+    new Response(
+      JSON.stringify({ success: true, repo_id: "local/combined", num_episodes: 4, total_frames: 120 }),
+      { status: 200, headers: { "Content-Type": "application/json" } },
+    ),
+  );
+
+  await expect(mergeLocalDatasets("http://localhost:8000", fetcher, ["local/one", "local/two"], "combined"))
+    .resolves.toEqual({ success: true, repo_id: "local/combined", num_episodes: 4, total_frames: 120 });
+
+  expect(fetcher).toHaveBeenCalledWith("http://localhost:8000/datasets/merge", {
+    method: "POST",
+    signal: undefined,
+    body: JSON.stringify({ source_repo_ids: ["local/one", "local/two"], output_name: "combined" }),
+  });
+});

frontend/src/lib/replayApi.ts

@@ -9,6 +9,13 @@ export interface DatasetItem {
   source: DatasetSource;
 }
 
+export interface MergeDatasetsResult {
+  success: true;
+  repo_id: string;
+  num_episodes: number;
+  total_frames: number;
+}
+
 export async function listDatasets(
   baseUrl: string,
   fetcher: Fetcher,
@@ -35,3 +42,16 @@ export async function listLocalDatasets(
     action: "List local datasets",
   });
 }
+
+export async function mergeLocalDatasets(
+  baseUrl: string,
+  fetcher: Fetcher,
+  sourceRepoIds: string[],
+  outputName: string,
+): Promise<MergeDatasetsResult> {
+  return apiRequest<MergeDatasetsResult>(baseUrl, fetcher, "/datasets/merge", {
+    method: "POST",
+    body: { source_repo_ids: sourceRepoIds, output_name: outputName },
+    action: "Merge datasets",
+  });
+}

frontend/src/pages/EditDataset.tsx

@@ -1,6 +1,6 @@
-import React, { useEffect, useMemo } from "react";
+import React, { useEffect, useMemo, useState } from "react";
 import { useNavigate, useSearchParams } from "react-router-dom";
-import { AlertCircle, ArrowLeft, Database, Loader2, Upload } from "lucide-react";
+import { AlertCircle, ArrowLeft, Database, Loader2, Merge, Upload } from "lucide-react";
 import {
   Select,
   SelectContent,
@@ -11,9 +11,13 @@ import {
 import { Button } from "@/components/ui/button";
 import EpisodeList from "@/components/dataset/EpisodeList";
 import EpisodeViewer from "@/components/dataset/EpisodeViewer";
+import MergeDatasetsDialog from "@/components/dataset/MergeDatasetsDialog";
+import { useApi } from "@/contexts/ApiContext";
 import { useDatasets } from "@/hooks/useDatasets";
 import { useEpisodeDetail, useEpisodes } from "@/hooks/useEpisodes";
+import { useToast } from "@/hooks/use-toast";
 import { formatDuration } from "@/lib/datasetApi";
+import { mergeLocalDatasets } from "@/lib/replayApi";
 
 const Stat: React.FC<{ label: string; value: string }> = ({ label, value }) => (
   <span className="flex items-baseline gap-1.5">
@@ -33,7 +37,10 @@ const EditDataset = () => {
   const navigate = useNavigate();
   const [searchParams, setSearchParams] = useSearchParams();
 
-  const { datasets, loading: datasetsLoading } = useDatasets();
+  const { datasets, loading: datasetsLoading, refresh } = useDatasets();
+  const { baseUrl, fetchWithHeaders } = useApi();
+  const { toast } = useToast();
+  const [mergeOpen, setMergeOpen] = 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"),
@@ -44,6 +51,7 @@ const EditDataset = () => {
   const episodeParam = searchParams.get("episode");
   const episodeIndex = episodeParam === null ? null : Number(episodeParam);
   const selectedSource = datasets.find((d) => d.repo_id === repoId)?.source;
+  const selectedIsLocal = selectedSource === "local" || selectedSource === "both";
 
   const { data, loading, error } = useEpisodes(repoId);
   const {
@@ -62,6 +70,12 @@ const EditDataset = () => {
     return api ? { ...params, api } : params;
   };
   const setDataset = (next: string) => setSearchParams(withApi({ dataset: next }));
+  const handleMerge = async (sourceRepoIds: string[], outputName: string) => {
+    const result = await mergeLocalDatasets(baseUrl, fetchWithHeaders, sourceRepoIds, outputName);
+    refresh();
+    setDataset(result.repo_id);
+    toast({ title: "Datasets merged", description: `${result.num_episodes} episodes are ready to browse.` });
+  };
   const setEpisode = (next: number) => {
     if (!repoId) return;
     setSearchParams(withApi({ dataset: repoId, episode: String(next) }));
@@ -113,6 +127,17 @@ const EditDataset = () => {
                 </SelectContent>
               </Select>
             </div>
+            {repoId && selectedIsLocal && (
+              <Button
+                variant="outline"
+                size="sm"
+                onClick={() => setMergeOpen(true)}
+                className="h-9 gap-1.5 border-gray-800 bg-gray-950 text-xs text-gray-300 hover:bg-gray-900"
+              >
+                <Merge className="h-3.5 w-3.5" />
+                Merge
+              </Button>
+            )}
             {/* Picking a dataset now lands here rather than on /upload, so the
                 upload + delete flow hangs off the page you browse from. */}
             {repoId && (
@@ -206,6 +231,15 @@ const EditDataset = () => {
           </>
         )}
       </div>
+      {repoId && selectedIsLocal && (
+        <MergeDatasetsDialog
+          datasets={datasets}
+          selectedRepoId={repoId}
+          open={mergeOpen}
+          onOpenChange={setMergeOpen}
+          onMerge={handleMerge}
+        />
+      )}
     </div>
   );
 };