huggingface · nobullryder · Jun 25, 2026
diff --git a/frontend/dist/assets/index-BYEP-1Zk.js → frontend/dist/assets/index-CuYRVmdE.js b/frontend/dist/assets/index-BYEP-1Zk.js → frontend/dist/assets/index-CuYRVmdE.js
diff --git a/frontend/dist/assets/index-Da-rRmcX.css b/frontend/dist/assets/index-Da-rRmcX.css
diff --git a/frontend/dist/assets/index-DmcGvCXo.css b/frontend/dist/assets/index-DmcGvCXo.css
diff --git a/frontend/dist/index.html b/frontend/dist/index.html
@@ -1,20 +1,20 @@
-<!DOCTYPE html>
-<html lang="en">
-  <head>
-    <meta charset="UTF-8" />
-    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-    <title>LeLab</title>
-    <link rel="icon" type="image/x-icon" href="/favicon.ico" />
-    <meta name="description" content="LeRobot but on your browser." />
-
-    <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-BYEP-1Zk.js"></script>
-    <link rel="stylesheet" crossorigin href="/assets/index-DmcGvCXo.css">
-  </head>
-
-  <body>
-    <div id="root"></div>
-  </body>
-</html>
+<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <title>LeLab</title>
+    <link rel="icon" type="image/x-icon" href="/favicon.ico" />
+    <meta name="description" content="LeRobot but on your browser." />
+
+    <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-CuYRVmdE.js"></script>
+    <link rel="stylesheet" crossorigin href="/assets/index-Da-rRmcX.css">
+  </head>
+
+  <body>
+    <div id="root"></div>
+  </body>
+</html>

diff --git a/frontend/src/components/training/PolicyExtraDialog.tsx b/frontend/src/components/training/PolicyExtraDialog.tsx
@@ -0,0 +1,81 @@
+import React from "react";
+import {
+  Dialog,
+  DialogContent,
+  DialogDescription,
+  DialogHeader,
+  DialogTitle,
+} from "@/components/ui/dialog";
+import { useInstallExtra } from "@/hooks/useInstallExtra";
+import {
+  InstallProgress,
+  InstallTitleIcon,
+  RestartInstructions,
+  installTitle,
+} from "./InstallProgress";
+
+interface Props {
+  open: boolean;
+  onOpenChange: (open: boolean) => void;
+  policyType: string;
+  packageName: string; // the probed module, e.g. "transformers"
+  installTarget: string; // e.g. "lerobot[smolvla]"
+  installHint: string; // e.g. "pip install 'lerobot[smolvla]'"
+}
+
+// Some policies (smolvla, pi0, pi0_fast, diffusion) need an optional LeRobot
+// extra. This catches the missing package before training starts and offers a
+// one-click install, instead of the run dying with a buried ImportError.
+const PolicyExtraDialog: React.FC<Props> = ({
+  open,
+  onOpenChange,
+  policyType,
+  packageName,
+  installTarget,
+  installHint,
+}) => {
+  const install = useInstallExtra(`system/policy-extra/${policyType}`, open);
+  const title = `${policyType.toUpperCase()} needs an extra package`;
+
+  return (
+    <Dialog open={open} onOpenChange={onOpenChange}>
+      <DialogContent className="bg-slate-800 border-slate-700 text-white max-w-2xl">
+        <DialogHeader>
+          <DialogTitle className="flex items-center gap-3 text-white">
+            <InstallTitleIcon state={install.state} />
+            {installTitle(install.state, title)}
+          </DialogTitle>
+          <DialogDescription className="sr-only">
+            Install {installTarget} to train {policyType}.
+          </DialogDescription>
+        </DialogHeader>
+
+        <div className="space-y-4">
+          <InstallProgress
+            state={install.state}
+            error={install.error}
+            logs={install.logs}
+            logBoxRef={install.logBoxRef}
+            onInstall={install.handleInstall}
+            onRetry={install.handleRetry}
+            installHint={installHint}
+            packageName={installTarget}
+            idleTitle={title}
+            idleDescription={
+              <>
+                Training a <span className="font-semibold">{policyType}</span> policy needs the{" "}
+                <code className="px-1 py-0.5 rounded bg-slate-900 text-sky-300">{packageName}</code>{" "}
+                package (installed via{" "}
+                <code className="px-1 py-0.5 rounded bg-slate-900 text-sky-300">{installTarget}</code>),
+                which isn't in this environment yet. Install it to train this policy.
+              </>
+            }
+            doneDescription={<RestartInstructions purpose={`${policyType} training`} />}
+          />
+        </div>
+      </DialogContent>
+    </Dialog>
+  );
+};
+
+export default PolicyExtraDialog;
diff --git a/frontend/src/pages/Training.tsx b/frontend/src/pages/Training.tsx
@@ -10,6 +10,7 @@ import ConfigurationTab from "@/components/training/ConfigurationTab";
 import MonitoringStats from "@/components/training/monitoring/MonitoringStats";
 import TrainingLogs from "@/components/training/monitoring/TrainingLogs";
 import TrainingExtraGate from "@/components/training/TrainingExtraGate";
+import PolicyExtraDialog from "@/components/training/PolicyExtraDialog";
 import HfAuthBanner from "@/components/landing/HfAuthBanner";
 
 import { Button } from "@/components/ui/button";
@@ -131,6 +132,12 @@ const ConfigurationMode: React.FC = () => {
   const [trainingExtraInstallHint, setTrainingExtraInstallHint] = useState<string>("pip install accelerate");
   const [localJobRunning, setLocalJobRunning] = useState<boolean>(false);
   const [isStarting, setIsStarting] = useState(false);
+  const [policyExtra, setPolicyExtra] = useState<{
+    policyType: string;
+    packageName: string;
+    installTarget: string;
+    installHint: string;
+  } | null>(null);
   const [authenticated, setAuthenticated] = useState<boolean>(false);
   const [flavors, setFlavors] = useState<RunnerFlavor[]>([]);
   const [hardwareLoading, setHardwareLoading] = useState(true);
@@ -191,6 +198,31 @@ const ConfigurationMode: React.FC = () => {
       toast({ title: "Error", description: "Dataset repository ID is required", variant: "destructive" });
       return;
     }
+
+    // Pre-flight: smolvla/pi0/diffusion need an optional package. Catch it here
+    // with a one-click installer instead of a buried ImportError after the job
+    // has already started.
+    try {
+      const r = await fetchWithHeaders(
+        `${baseUrl}/system/policy-extra/${trainingConfig.policy_type}`,
+      );
+      if (r.ok) {
+        const extra = await r.json();
+        if (extra.needs_extra && !extra.available) {
+          setPolicyExtra({
+            policyType: trainingConfig.policy_type,
+            packageName: extra.package,
+            installTarget: extra.install_target,
+            installHint: extra.install_hint,
+          });
+          return;
+        }
+      }
+    } catch {
+      // Check failed (offline / older backend) — fall through and let the job
+      // report any problem itself.
+    }
+
     setIsStarting(true);
     try {
       const job = await startTrainingJob(baseUrl, fetchWithHeaders, configToRequest(trainingConfig));
@@ -306,6 +338,17 @@ const ConfigurationMode: React.FC = () => {
           })()}
         </div>
       </div>
+
+      {policyExtra && (
+        <PolicyExtraDialog
+          open={!!policyExtra}
+          onOpenChange={(o) => !o && setPolicyExtra(null)}
+          policyType={policyExtra.policyType}
+          packageName={policyExtra.packageName}
+          installTarget={policyExtra.installTarget}
+          installHint={policyExtra.installHint}
+        />
+      )}
     </div>
   );
 };

diff --git a/lelab/server.py b/lelab/server.py
@@ -101,8 +101,11 @@
 from .utils.hf_auth import cached_whoami, handle_hf_auth_status, handle_hf_login, shared_hf_api
 from .utils.system import (
     handle_get_cuda_status,
+    handle_get_policy_extra,
     handle_get_training_extra,
     handle_get_wandb_extra,
+    handle_install_policy_extra,
+    handle_install_policy_extra_status,
     handle_install_training_extra,
     handle_install_training_extra_status,
     handle_install_wandb_extra,
@@ -750,6 +753,25 @@ def install_wandb_extra_status():
     return handle_install_wandb_extra_status()
 
 
+@app.get("/system/policy-extra/{policy_type}")
+def get_policy_extra(policy_type: str):
+    """Whether the optional LeRobot extra a policy needs (e.g. transformers for
+    smolvla/pi0, diffusers for diffusion) is importable. Core policies report available."""
+    return handle_get_policy_extra(policy_type)
+
+
+@app.post("/system/policy-extra/{policy_type}/install")
+def install_policy_extra(policy_type: str):
+    """Spawn `pip install lerobot[<extra>]` for the policy's extra in the background."""
+    return handle_install_policy_extra(policy_type)
+
+
+@app.get("/system/policy-extra/{policy_type}/install-status")
+def install_policy_extra_status(policy_type: str):
+    """Return the policy extra's install state plus any pending log lines (drained on read)."""
+    return handle_install_policy_extra_status(policy_type)
+
+
 @app.get("/system/update-check")
 def update_check():
     """Report whether a newer LeLab commit exists on GitHub (cached, silent on failure)."""

diff --git a/lelab/utils/system.py b/lelab/utils/system.py
@@ -186,6 +186,85 @@ def handle_install_wandb_extra_status() -> dict[str, Any]:
     return wandb_install_manager.get_status()
 
 
+# --------------------------------------------------------------------------- #
+# Policy extras
+# --------------------------------------------------------------------------- #
+# Some LeRobot policies import an optional extra at construction time; training
+# (or inference) otherwise dies with a buried ImportError once the subprocess is
+# already running. Map each such policy to the module we probe and the
+# ``pip install lerobot[extra]`` target. Policies not listed (act, vqbet, tdmpc,
+# sac, reward_classifier) need nothing extra.
+POLICY_EXTRAS: dict[str, tuple[str, str]] = {
+    # policy_type: (probe_module, install_target)
+    "smolvla": ("transformers", "lerobot[smolvla]"),
+    "pi0": ("transformers", "lerobot[pi]"),
+    "pi0_fast": ("transformers", "lerobot[pi]"),
+    "diffusion": ("diffusers", "lerobot[diffusion]"),
+}
+
+# One install manager per install target (lerobot[smolvla] / lerobot[pi] / …),
+# created lazily so pi0 and pi0_fast share the lerobot[pi] install.
+_policy_install_managers: dict[str, InstallManager] = {}
+
+
+def _policy_install_manager(policy_type: str) -> InstallManager | None:
+    spec = POLICY_EXTRAS.get(policy_type)
+    if spec is None:
+        return None
+    target = spec[1]
+    mgr = _policy_install_managers.get(target)
+    if mgr is None:
+        mgr = InstallManager(target)
+        _policy_install_managers[target] = mgr
+    return mgr
+
+
+def handle_get_policy_extra(policy_type: str) -> dict[str, Any]:
+    """Whether the optional extra a policy needs is importable right now.
+
+    Probed live (not cached at import) so a restart after installing is picked
+    up. Policies that need nothing report ``available`` so the UI never blocks
+    them.
+    """
+    spec = POLICY_EXTRAS.get(policy_type)
+    if spec is None:
+        return {
+            "policy_type": policy_type,
+            "needs_extra": False,
+            "available": True,
+            "package": "",
+            "install_target": "",
+            "install_hint": "",
+        }
+    probe, target = spec
+    try:
+        available = importlib.util.find_spec(probe) is not None
+    except (ImportError, ValueError):
+        available = False
+    return {
+        "policy_type": policy_type,
+        "needs_extra": True,
+        "available": available,
+        "package": probe,
+        "install_target": target,
+        "install_hint": f"pip install '{target}'",
+    }
+
+
+def handle_install_policy_extra(policy_type: str) -> dict[str, Any]:
+    mgr = _policy_install_manager(policy_type)
+    if mgr is None:
+        return {"started": False, "message": f"'{policy_type}' needs no extra package."}
+    return mgr.start()
+
+
+def handle_install_policy_extra_status(policy_type: str) -> dict[str, Any]:
+    mgr = _policy_install_manager(policy_type)
+    if mgr is None:
+        return {"state": "done", "error": None, "logs": []}
+    return mgr.get_status()
+
+
 # Detect the common Windows/LeLab mismatch where an NVIDIA GPU is visible to the
 # OS, but the active PyTorch build cannot use CUDA. Do not auto-install torch.
 

diff --git a/tests/test_utils_system.py b/tests/test_utils_system.py
@@ -127,3 +127,55 @@ def test_cuda_status_endpoint_returns_expected_shape(client) -> None:
     assert response.status_code == 200
     body = response.json()
     assert set(body) >= {"gpu_present", "cuda_available", "mismatch", "install_hint", "docs_url"}
+
+
+def test_policy_extra_maps_policies_to_install_targets() -> None:
+    """smolvla/pi0/pi0_fast/diffusion map to the right probe module + lerobot[extra]."""
+    from lelab.utils.system import handle_get_policy_extra
+
+    smol = handle_get_policy_extra("smolvla")
+    assert smol["needs_extra"] is True
+    assert smol["package"] == "transformers"
+    assert smol["install_target"] == "lerobot[smolvla]"
+    assert "lerobot[smolvla]" in smol["install_hint"]
+
+    # pi0 and pi0_fast share the lerobot[pi] extra; diffusion uses diffusers.
+    assert handle_get_policy_extra("pi0")["install_target"] == "lerobot[pi]"
+    assert handle_get_policy_extra("pi0_fast")["install_target"] == "lerobot[pi]"
+    assert handle_get_policy_extra("diffusion")["package"] == "diffusers"
+    assert handle_get_policy_extra("diffusion")["install_target"] == "lerobot[diffusion]"
+
+
+def test_policy_extra_core_policy_needs_nothing() -> None:
+    from lelab.utils.system import handle_get_policy_extra
+
+    act = handle_get_policy_extra("act")
+    assert act["needs_extra"] is False
+    assert act["available"] is True
+    assert act["install_target"] == ""
+
+
+def test_policy_extra_available_reflects_find_spec(monkeypatch) -> None:
+    import importlib.util
+
+    from lelab.utils import system
+
+    monkeypatch.setattr(importlib.util, "find_spec", lambda name: object())
+    assert system.handle_get_policy_extra("smolvla")["available"] is True
+    monkeypatch.setattr(importlib.util, "find_spec", lambda name: None)
+    assert system.handle_get_policy_extra("smolvla")["available"] is False
+
+
+def test_policy_extra_install_is_noop_for_core_policy() -> None:
+    from lelab.utils.system import handle_install_policy_extra, handle_install_policy_extra_status
+
+    assert handle_install_policy_extra("act")["started"] is False
+    assert handle_install_policy_extra_status("act")["state"] == "done"
+
+
+def test_policy_extra_route_known_and_core(client) -> None:
+    smol = client.get("/system/policy-extra/smolvla").json()
+    assert smol["needs_extra"] is True
+    assert smol["install_target"] == "lerobot[smolvla]"
+    core = client.get("/system/policy-extra/act").json()
+    assert core["needs_extra"] is False