Sync from huggingface/lerobot-dataset-visualizer via hub-sync

src/app/[org]/[dataset]/[episode]/fetch-data.ts

@@ -22,7 +22,13 @@ import {
   depthEncodingFromFeature,
 } from "@/utils/colormaps";
 import type { VideoInfo } from "@/types";
-import { LeRobotDataset, parseInfo } from "@huggingface/lerobot";
+import {
+  LeRobotDataset,
+  HttpError,
+  fetchRange,
+  formatPathTemplate,
+  parseInfo,
+} from "@huggingface/lerobot";
 import type { LeRobotEpisode, LeRobotFrames } from "@huggingface/lerobot";
 import { authHeaders, getAuthToken } from "@/utils/auth";
 
@@ -590,6 +596,109 @@ export async function getEpisodeData(
   }
 }
 
+const v2ContiguousIndexes = new Map<string, Promise<boolean>>();
+/// The package's first 8 KB read of episodes.jsonl already holds the first few dozen episodes, so
+/// probing for those would only add a round trip.
+const V2_PROBE_FROM_EPISODE = 32;
+
+/**
+ * Whether line N of `meta/episodes.jsonl` is episode N, judged from its first and last lines.
+ *
+ * The package finds line N by re-downloading ever longer prefixes of the file (8 KB × 4^k) one
+ * after another and gives up past 4 MB, so the upper episodes of a large v2 dataset could not open.
+ * An index running 0..total-1 lets the episode be located from info.json instead; indexes that
+ * start elsewhere (10, 20, 30, ... in behavior-1k) keep the package's lookup.
+ */
+function isV2IndexContiguous(
+  repoId: string,
+  totalEpisodes: number,
+): Promise<boolean> {
+  const url = leRobotDataset(repoId).fileUrl("meta/episodes.jsonl");
+  let contiguous = v2ContiguousIndexes.get(url);
+  if (!contiguous) {
+    const read = (start: number, end?: number) =>
+      fetchRange(url, start, end, {
+        additionalFetchHeaders: authHeaders(),
+      }).then(({ bytes }) => new TextDecoder().decode(bytes));
+    const headRead = read(0, 4095);
+    /// The Hub answers a suffix longer than the file with a 416; the head then holds all of it.
+    const tailRead = read(-4096).catch((error) => {
+      if (error instanceof HttpError && error.status === 416) return headRead;
+      throw error;
+    });
+    contiguous = Promise.all([headRead, tailRead]).then(
+      ([head, tail]) => {
+        try {
+          const first = JSON.parse(head.split("\n")[0]);
+          const last = JSON.parse(tail.trimEnd().split("\n").pop() ?? "");
+          return (
+            first.episode_index === 0 &&
+            last.episode_index === totalEpisodes - 1
+          );
+        } catch {
+          /// A line longer than the 4 KB window: the answer would not change on a retry.
+          return false;
+        }
+      },
+      () => {
+        v2ContiguousIndexes.delete(url);
+        return false;
+      },
+    );
+    v2ContiguousIndexes.set(url, contiguous);
+  }
+  return contiguous;
+}
+
+/**
+ * Enough of the package's episode for getEpisodeDataV2, built from info.json alone. Only
+ * `videos[].url` and `data` are filled in: `length`, `durationSec`, `tasks` and `toSec` stay empty,
+ * since that path takes its duration from the chart timestamps.
+ */
+async function v2EpisodeFromTemplates(
+  repoId: string,
+  episodeId: number,
+): Promise<LeRobotEpisode> {
+  const dataset = leRobotDataset(repoId);
+  const info = await dataset.info();
+  /// The package's lookup keeps answering for ids that name no episode ("not found" past the end).
+  if (
+    !Number.isInteger(episodeId) ||
+    episodeId < 0 ||
+    episodeId >= info.totalEpisodes
+  ) {
+    return loadLeRobotEpisode(repoId, episodeId);
+  }
+  const { videoPath } = info;
+  const vars = {
+    episode_chunk: Math.floor(episodeId / info.chunksSize),
+    episode_index: episodeId,
+  };
+  return {
+    index: episodeId,
+    length: 0,
+    durationSec: 0,
+    tasks: [],
+    videos:
+      videoPath === undefined
+        ? []
+        : info.cameras.map((camera) => ({
+            cameraKey: camera.key,
+            url: dataset.fileUrl(
+              formatPathTemplate(videoPath, { ...vars, video_key: camera.key }),
+            ),
+            fromSec: 0,
+            toSec: 0,
+          })),
+    data: {
+      url: dataset.fileUrl(formatPathTemplate(info.dataPath, vars)),
+      /// The file holds this one episode, and hyparquet stops at its last row.
+      fromRow: 0,
+      toRow: Number.MAX_SAFE_INTEGER,
+    },
+  };
+}
+
 // Legacy v2.x data loading
 async function getEpisodeDataV2(
   repoId: string,
@@ -620,7 +729,19 @@ async function getEpisodeDataV2(
           .map((x) => parseInt(x.trim(), 10))
           .filter((x) => !isNaN(x));
 
-  const leRobotEpisode = await loadLeRobotEpisode(repoId, episodeId);
+  /// Both ways of locating the episode need the package's info.json; fetch it alongside the probe.
+  const probe =
+    episodeId >= V2_PROBE_FROM_EPISODE ||
+    v2ContiguousIndexes.has(
+      leRobotDataset(repoId).fileUrl("meta/episodes.jsonl"),
+    );
+  const [contiguous] = await Promise.all([
+    probe && isV2IndexContiguous(repoId, info.total_episodes),
+    leRobotDataset(repoId).info(),
+  ]);
+  const leRobotEpisode = contiguous
+    ? await v2EpisodeFromTemplates(repoId, episodeId)
+    : await loadLeRobotEpisode(repoId, episodeId);
   const videosInfo =
     info.video_path !== null ? toVideosInfo(leRobotEpisode, info, false) : [];
 

src/components/urdf-viewer.tsx

@@ -57,6 +57,104 @@ function loadStlGeometry(url: string): Promise<THREE.BufferGeometry> {
   return loading;
 }
 
+/// One download and parse per DAE file per session: the OpenArm URDF references finger.dae 4 times
+/// and each arm link twice, and the viewer remounts on every visit to the 3D Replay tab.
+const daeSceneCache = new Map<string, Promise<THREE.Object3D>>();
+
+function loadDaeScene(url: string): Promise<THREE.Object3D> {
+  let loading = daeSceneCache.get(url);
+  if (!loading) {
+    loading = new ColladaLoader()
+      .loadAsync(url)
+      .then((collada) => {
+        // DAE (Collada) files — ColladaLoader yields whatever the .dae author
+        // baked in: flat MeshPhongMaterial/MeshBasicMaterial colors plus, in
+        // OpenArm's case, ~23 per-file PointLight/SpotLight nodes. The stray
+        // lights caused the scene to look pure-white everywhere regardless of
+        // our own lighting, and the flat materials looked cartoonish. We strip
+        // both and rebuild every mesh with a MeshStandardMaterial bucketed into
+        // one of three archetypes (carbon-black, brushed metal, off-white paint)
+        // based on the original base-color lightness.
+        if (url.includes("openarm")) {
+          const strayLights: THREE.Object3D[] = [];
+          collada.scene.traverse((child) => {
+            if (
+              (child as THREE.Light).isLight &&
+              !(child instanceof THREE.AmbientLight)
+            ) {
+              strayLights.push(child);
+            }
+          });
+          for (const l of strayLights) l.parent?.remove(l);
+
+          collada.scene.traverse((child) => {
+            if (!(child instanceof THREE.Mesh) || !child.material) return;
+
+            const originals = Array.isArray(child.material)
+              ? child.material
+              : [child.material];
+
+            const rebuilt = originals.map((orig) => {
+              const srcColor =
+                (orig as THREE.MeshStandardMaterial).color ??
+                new THREE.Color("#c0c4cc");
+              const hsl = { h: 0, s: 0, l: 0 };
+              srcColor.getHSL(hsl);
+
+              // Archetype classification by original lightness.
+              let color: THREE.Color;
+              let metalness: number;
+              let roughness: number;
+              let envMapIntensity: number;
+              if (hsl.l < 0.3) {
+                // Carbon / anodised structural parts
+                color = new THREE.Color().setHSL(hsl.h, 0.02, 0.09);
+                metalness = 0.15;
+                roughness = 0.75;
+                envMapIntensity = 0.6;
+              } else if (hsl.l < 0.7) {
+                // Brushed metal joint collars / accents
+                color = new THREE.Color().setHSL(hsl.h, 0.04, 0.42);
+                metalness = 0.75;
+                roughness = 0.35;
+                envMapIntensity = 1.1;
+              } else {
+                // Off-white painted plates
+                color = new THREE.Color().setHSL(hsl.h, 0.03, 0.6);
+                metalness = 0.1;
+                roughness = 0.5;
+                envMapIntensity = 0.9;
+              }
+
+              const mat = new THREE.MeshStandardMaterial({
+                color,
+                metalness,
+                roughness,
+                envMapIntensity,
+                side: THREE.DoubleSide,
+              });
+              orig.dispose?.();
+              return mat;
+            });
+
+            child.material = Array.isArray(child.material)
+              ? rebuilt
+              : rebuilt[0];
+            child.castShadow = true;
+            child.receiveShadow = true;
+          });
+        }
+        return collada.scene;
+      })
+      .catch((err) => {
+        daeSceneCache.delete(url);
+        throw err;
+      });
+    daeSceneCache.set(url, loading);
+  }
+  return loading;
+}
+
 // URDFs + meshes are hosted in the Hub bucket at
 // https://huggingface.co/buckets/lerobot/robot-urdfs. URDFLoader resolves
 // relative mesh paths against the URDF's own URL, so the bucket layout
@@ -410,93 +508,14 @@ function RobotScene({
       return group;
     };
     loader.loadMeshCb = (url, mgr, onLoad) => {
-      // DAE (Collada) files — ColladaLoader yields whatever the .dae author
-      // baked in: flat MeshPhongMaterial/MeshBasicMaterial colors plus, in
-      // OpenArm's case, ~23 per-file PointLight/SpotLight nodes. The stray
-      // lights caused the scene to look pure-white everywhere regardless of
-      // our own lighting, and the flat materials looked cartoonish. We strip
-      // both and rebuild every mesh with a MeshStandardMaterial bucketed into
-      // one of three archetypes (carbon-black, brushed metal, off-white paint)
-      // based on the original base-color lightness.
       if (url.endsWith(".dae")) {
-        const colladaLoader = new ColladaLoader(mgr);
-        colladaLoader.load(
-          url,
-          (collada) => {
-            if (isOpenArm) {
-              const strayLights: THREE.Object3D[] = [];
-              collada.scene.traverse((child) => {
-                if (
-                  (child as THREE.Light).isLight &&
-                  !(child instanceof THREE.AmbientLight)
-                ) {
-                  strayLights.push(child);
-                }
-              });
-              for (const l of strayLights) l.parent?.remove(l);
-
-              collada.scene.traverse((child) => {
-                if (!(child instanceof THREE.Mesh) || !child.material) return;
-
-                const originals = Array.isArray(child.material)
-                  ? child.material
-                  : [child.material];
-
-                const rebuilt = originals.map((orig) => {
-                  const srcColor =
-                    (orig as THREE.MeshStandardMaterial).color ??
-                    new THREE.Color("#c0c4cc");
-                  const hsl = { h: 0, s: 0, l: 0 };
-                  srcColor.getHSL(hsl);
-
-                  // Archetype classification by original lightness.
-                  let color: THREE.Color;
-                  let metalness: number;
-                  let roughness: number;
-                  let envMapIntensity: number;
-                  if (hsl.l < 0.3) {
-                    // Carbon / anodised structural parts
-                    color = new THREE.Color().setHSL(hsl.h, 0.02, 0.09);
-                    metalness = 0.15;
-                    roughness = 0.75;
-                    envMapIntensity = 0.6;
-                  } else if (hsl.l < 0.7) {
-                    // Brushed metal joint collars / accents
-                    color = new THREE.Color().setHSL(hsl.h, 0.04, 0.42);
-                    metalness = 0.75;
-                    roughness = 0.35;
-                    envMapIntensity = 1.1;
-                  } else {
-                    // Off-white painted plates
-                    color = new THREE.Color().setHSL(hsl.h, 0.03, 0.6);
-                    metalness = 0.1;
-                    roughness = 0.5;
-                    envMapIntensity = 0.9;
-                  }
-
-                  const mat = new THREE.MeshStandardMaterial({
-                    color,
-                    metalness,
-                    roughness,
-                    envMapIntensity,
-                    side: THREE.DoubleSide,
-                  });
-                  orig.dispose?.();
-                  return mat;
-                });
-
-                child.material = Array.isArray(child.material)
-                  ? rebuilt
-                  : rebuilt[0];
-                child.castShadow = true;
-                child.receiveShadow = true;
-              });
-            }
-            onLoad(collada.scene);
-          },
-          undefined,
-          (err) => onLoad(new THREE.Object3D(), err as Error),
-        );
+        /// URDFLoader re-parents and re-orients what it is handed, so it gets a clone (sharing
+        /// geometry and materials) and the cached scene stays intact.
+        mgr.itemStart(url);
+        loadDaeScene(url)
+          .then((daeScene) => onLoad(daeScene.clone()))
+          .catch((err) => onLoad(new THREE.Object3D(), err as Error))
+          .finally(() => mgr.itemEnd(url));
         return;
       }
       // STL files — apply final PBR materials directly here. We used to do a