Make episode visualizer nicer (colors, formatting) add detection for still epsiodes and fix urdf visualizer

next.config.ts

@@ -2,13 +2,13 @@ import type { NextConfig } from "next";
 import packageJson from './package.json';
 
 const nextConfig: NextConfig = {
-
   typescript: {
     ignoreBuildErrors: true,
   },
   eslint: {
     ignoreDuringBuilds: true,
   },
+  transpilePackages: ["three"],
   generateBuildId: () => packageJson.version,
 };
 

package-lock.json

@@ -17,7 +17,8 @@
         "react-dom": "^19.0.0",
         "react-icons": "^5.5.0",
         "recharts": "^2.15.3",
-        "three": "^0.182.0"
+        "three": "^0.182.0",
+        "urdf-loader": "^0.12.6"
       },
       "devDependencies": {
         "@eslint/eslintrc": "^3",
@@ -6849,6 +6850,15 @@
         "@unrs/resolver-binding-win32-x64-msvc": "1.11.1"
       }
     },
+    "node_modules/urdf-loader": {
+      "version": "0.12.6",
+      "resolved": "https://registry.npmjs.org/urdf-loader/-/urdf-loader-0.12.6.tgz",
+      "integrity": "sha512-EwpgOCPe6Tep2+MXoo/r13keHaKQXMcM+4s9+jX0NRxNS/QSNuP5JPdk5AIgWEoEB43AkEj9Vk+Nr53NkXgSbA==",
+      "license": "Apache-2.0",
+      "peerDependencies": {
+        "three": ">=0.152.0"
+      }
+    },
     "node_modules/uri-js": {
       "version": "4.4.1",
       "resolved": "https://registry.npmjs.org/uri-js/-/uri-js-4.4.1.tgz",

package.json

@@ -19,7 +19,8 @@
     "react-dom": "^19.0.0",
     "react-icons": "^5.5.0",
     "recharts": "^2.15.3",
-    "three": "^0.182.0"
+    "three": "^0.182.0",
+    "urdf-loader": "^0.12.6"
   },
   "devDependencies": {
     "@eslint/eslintrc": "^3",

src/app/[org]/[dataset]/[episode]/actions.ts

@@ -5,8 +5,10 @@ import type { DatasetMetadata } from "@/utils/parquetUtils";
 import {
   loadAllEpisodeLengthsV3,
   loadAllEpisodeFrameInfo,
+  loadCrossEpisodeActionVariance,
   type EpisodeLengthStats,
   type EpisodeFramesData,
+  type CrossEpisodeVarianceData,
 } from "./fetch-data";
 
 export async function fetchEpisodeLengthStats(
@@ -28,3 +30,12 @@ export async function fetchEpisodeFrames(
   return loadAllEpisodeFrameInfo(repoId, version, info as unknown as DatasetMetadata);
 }
 
+export async function fetchCrossEpisodeVariance(
+  org: string,
+  dataset: string,
+): Promise<CrossEpisodeVarianceData | null> {
+  const repoId = `${org}/${dataset}`;
+  const { version, info } = await getDatasetVersionAndInfo(repoId);
+  return loadCrossEpisodeActionVariance(repoId, version, info as unknown as DatasetMetadata, info.fps);
+}
+

src/app/[org]/[dataset]/[episode]/episode-viewer.tsx

@@ -19,12 +19,14 @@ import {
   type ColumnMinMax,
   type EpisodeLengthStats,
   type EpisodeFramesData,
+  type CrossEpisodeVarianceData,
 } from "./fetch-data";
-import { fetchEpisodeLengthStats, fetchEpisodeFrames } from "./actions";
+import { fetchEpisodeLengthStats, fetchEpisodeFrames, fetchCrossEpisodeVariance } from "./actions";
 
 const URDFViewer = lazy(() => import("@/components/urdf-viewer"));
+const ActionInsightsPanel = lazy(() => import("@/components/action-insights-panel"));
 
-type ActiveTab = "episodes" | "statistics" | "frames" | "urdf";
+type ActiveTab = "episodes" | "statistics" | "frames" | "insights" | "urdf";
 
 export default function EpisodeViewer({
   data,
@@ -87,6 +89,9 @@ function EpisodeViewerInner({ data, org, dataset }: { data: EpisodeData; org?: s
   const [episodeFramesData, setEpisodeFramesData] = useState<EpisodeFramesData | null>(null);
   const [framesLoading, setFramesLoading] = useState(false);
   const framesLoadedRef = useRef(false);
+  const [crossEpData, setCrossEpData] = useState<CrossEpisodeVarianceData | null>(null);
+  const [insightsLoading, setInsightsLoading] = useState(false);
+  const insightsLoadedRef = useRef(false);
 
   const loadStats = () => {
     if (statsLoadedRef.current) return;
@@ -113,10 +118,21 @@ function EpisodeViewerInner({ data, org, dataset }: { data: EpisodeData; org?: s
       .finally(() => setFramesLoading(false));
   };
 
+  const loadInsights = () => {
+    if (insightsLoadedRef.current || !org || !dataset) return;
+    insightsLoadedRef.current = true;
+    setInsightsLoading(true);
+    fetchCrossEpisodeVariance(org, dataset)
+      .then(setCrossEpData)
+      .catch((err) => console.error("[cross-ep] Failed:", err))
+      .finally(() => setInsightsLoading(false));
+  };
+
   const handleTabChange = (tab: ActiveTab) => {
     setActiveTab(tab);
     if (tab === "statistics") loadStats();
     if (tab === "frames") loadFrames();
+    if (tab === "insights") loadInsights();
   };
 
   // Use context for time sync
@@ -288,6 +304,19 @@ function EpisodeViewerInner({ data, org, dataset }: { data: EpisodeData; org?: s
             <span className="absolute bottom-0 left-0 right-0 h-0.5 bg-orange-500" />
           )}
         </button>
+        <button
+          className={`px-6 py-2.5 text-sm font-medium transition-colors relative ${
+            activeTab === "insights"
+              ? "text-orange-400"
+              : "text-slate-400 hover:text-slate-200"
+          }`}
+          onClick={() => handleTabChange("insights")}
+        >
+          Action Insights
+          {activeTab === "insights" && (
+            <span className="absolute bottom-0 left-0 right-0 h-0.5 bg-orange-500" />
+          )}
+        </button>
         {isSO101Robot(datasetInfo.robot_type) && (
           <button
             className={`px-6 py-2.5 text-sm font-medium transition-colors relative ${
@@ -405,6 +434,17 @@ function EpisodeViewerInner({ data, org, dataset }: { data: EpisodeData; org?: s
             <OverviewPanel data={episodeFramesData} loading={framesLoading} />
           )}
 
+          {activeTab === "insights" && (
+            <Suspense fallback={<Loading />}>
+              <ActionInsightsPanel
+                flatChartData={data.flatChartData}
+                fps={datasetInfo.fps}
+                crossEpisodeData={crossEpData}
+                crossEpisodeLoading={insightsLoading}
+              />
+            </Suspense>
+          )}
+
           {activeTab === "urdf" && (
             <Suspense fallback={<Loading />}>
               <URDFViewer data={data} />

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

@@ -1337,6 +1337,309 @@ export async function loadAllEpisodeFrameInfo(
   return { cameras, framesByCamera };
 }
 
+// ─── Cross-episode action variance ──────────────────────────────
+
+export type LowMovementEpisode = { episodeIndex: number; totalMovement: number };
+
+export type AggVelocityStat = {
+  name: string;
+  std: number;
+  maxAbs: number;
+  bins: number[];
+  lo: number;
+  hi: number;
+};
+
+export type AggAutocorrelation = {
+  chartData: Record<string, number>[];
+  suggestedChunk: number | null;
+  shortKeys: string[];
+};
+
+export type CrossEpisodeVarianceData = {
+  actionNames: string[];
+  timeBins: number[];
+  variance: number[][];
+  numEpisodes: number;
+  lowMovementEpisodes: LowMovementEpisode[];
+  aggVelocity: AggVelocityStat[];
+  aggAutocorrelation: AggAutocorrelation | null;
+};
+
+export async function loadCrossEpisodeActionVariance(
+  repoId: string,
+  version: string,
+  info: DatasetMetadata,
+  fps: number,
+  maxEpisodes = 500,
+  numTimeBins = 50,
+): Promise<CrossEpisodeVarianceData | null> {
+  const actionEntry = Object.entries(info.features)
+    .find(([key, f]) => key === "action" && f.shape.length === 1);
+  if (!actionEntry) {
+    console.warn("[cross-ep] No action feature found. Available features:", Object.entries(info.features).map(([k, f]) => `${k}(${f.dtype}, shape=${JSON.stringify(f.shape)})`).join(", "));
+    return null;
+  }
+
+  const [actionKey, actionMeta] = actionEntry;
+  const actionDim = actionMeta.shape[0];
+
+  let names: unknown = actionMeta.names;
+  while (typeof names === "object" && names !== null && !Array.isArray(names)) {
+    names = Object.values(names)[0];
+  }
+  const actionNames = Array.isArray(names)
+    ? (names as string[]).map(n => `${actionKey}${SERIES_NAME_DELIMITER}${n}`)
+    : Array.from({ length: actionDim }, (_, i) => `${actionKey}${SERIES_NAME_DELIMITER}${i}`);
+
+  // Collect episode metadata
+  type EpMeta = { index: number; chunkIdx: number; fileIdx: number; from: number; to: number };
+  const allEps: EpMeta[] = [];
+
+  if (version === "v3.0") {
+    let fileIndex = 0;
+    while (true) {
+      const path = `meta/episodes/chunk-000/file-${fileIndex.toString().padStart(3, "0")}.parquet`;
+      try {
+        const buf = await fetchParquetFile(buildVersionedUrl(repoId, version, path));
+        const rows = await readParquetAsObjects(buf, []);
+        if (rows.length === 0 && fileIndex > 0) break;
+        for (const row of rows) {
+          const parsed = parseEpisodeRowSimple(row);
+          allEps.push({
+            index: parsed.episode_index,
+            chunkIdx: parsed.data_chunk_index,
+            fileIdx: parsed.data_file_index,
+            from: parsed.dataset_from_index,
+            to: parsed.dataset_to_index,
+          });
+        }
+        fileIndex++;
+      } catch { break; }
+    }
+  } else {
+    for (let i = 0; i < info.total_episodes; i++) {
+      allEps.push({ index: i, chunkIdx: 0, fileIdx: 0, from: 0, to: 0 });
+    }
+  }
+
+  if (allEps.length < 2) {
+    console.warn(`[cross-ep] Only ${allEps.length} episode(s) found in metadata, need ≥2`);
+    return null;
+  }
+  console.log(`[cross-ep] Found ${allEps.length} episodes in metadata, sampling up to ${maxEpisodes}`);
+
+  // Sample episodes evenly
+  const sampled = allEps.length <= maxEpisodes
+    ? allEps
+    : Array.from({ length: maxEpisodes }, (_, i) =>
+        allEps[Math.round((i * (allEps.length - 1)) / (maxEpisodes - 1))]
+      );
+
+  // Load action data per episode, tracking episode index alongside
+  const episodeActions: { index: number; actions: number[][] }[] = [];
+
+  if (version === "v3.0") {
+    const byFile = new Map<string, EpMeta[]>();
+    for (const ep of sampled) {
+      const key = `${ep.chunkIdx}-${ep.fileIdx}`;
+      if (!byFile.has(key)) byFile.set(key, []);
+      byFile.get(key)!.push(ep);
+    }
+
+    for (const [, eps] of byFile) {
+      const ep0 = eps[0];
+      const dataPath = `data/chunk-${ep0.chunkIdx.toString().padStart(3, "0")}/file-${ep0.fileIdx.toString().padStart(3, "0")}.parquet`;
+      try {
+        const buf = await fetchParquetFile(buildVersionedUrl(repoId, version, dataPath));
+        const rows = await readParquetAsObjects(buf, []);
+        const fileStart = rows.length > 0 && rows[0].index !== undefined ? Number(rows[0].index) : 0;
+
+        for (const ep of eps) {
+          const localFrom = Math.max(0, ep.from - fileStart);
+          const localTo = Math.min(rows.length, ep.to - fileStart);
+          const actions: number[][] = [];
+          for (let r = localFrom; r < localTo; r++) {
+            const raw = rows[r]?.[actionKey];
+            if (Array.isArray(raw)) actions.push(raw.map(Number));
+          }
+          if (actions.length > 0) episodeActions.push({ index: ep.index, actions });
+        }
+      } catch { /* skip file */ }
+    }
+  } else {
+    const chunkSize = info.chunks_size || 1000;
+    for (const ep of sampled) {
+      const chunk = Math.floor(ep.index / chunkSize);
+      const dataPath = formatStringWithVars(info.data_path, {
+        episode_chunk: chunk.toString().padStart(3, "0"),
+        episode_index: ep.index.toString().padStart(6, "0"),
+      });
+      try {
+        const buf = await fetchParquetFile(buildVersionedUrl(repoId, version, dataPath));
+        const rows = await readParquetAsObjects(buf, []);
+        const actions: number[][] = [];
+        for (const row of rows) {
+          const raw = row[actionKey];
+          if (Array.isArray(raw)) {
+            actions.push(raw.map(Number));
+          } else {
+            const vec: number[] = [];
+            for (let d = 0; d < actionDim; d++) {
+              const v = row[`${actionKey}.${d}`] ?? row[d];
+              vec.push(typeof v === "number" ? v : Number(v) || 0);
+            }
+            actions.push(vec);
+          }
+        }
+        if (actions.length > 0) episodeActions.push({ index: ep.index, actions });
+      } catch { /* skip */ }
+    }
+  }
+
+  if (episodeActions.length < 2) {
+    console.warn(`[cross-ep] Only ${episodeActions.length} episode(s) had loadable action data out of ${sampled.length} sampled`);
+    return null;
+  }
+  console.log(`[cross-ep] Loaded action data for ${episodeActions.length}/${sampled.length} episodes`);
+
+  // Resample each episode to numTimeBins and compute variance
+  const timeBins = Array.from({ length: numTimeBins }, (_, i) => i / (numTimeBins - 1));
+  const sums = Array.from({ length: numTimeBins }, () => new Float64Array(actionDim));
+  const sumsSq = Array.from({ length: numTimeBins }, () => new Float64Array(actionDim));
+  const counts = new Uint32Array(numTimeBins);
+
+  for (const { actions: epActions } of episodeActions) {
+    const T = epActions.length;
+    for (let b = 0; b < numTimeBins; b++) {
+      const srcIdx = Math.min(Math.round(timeBins[b] * (T - 1)), T - 1);
+      const row = epActions[srcIdx];
+      for (let d = 0; d < actionDim; d++) {
+        const v = row[d] ?? 0;
+        sums[b][d] += v;
+        sumsSq[b][d] += v * v;
+      }
+      counts[b]++;
+    }
+  }
+
+  const variance: number[][] = [];
+  for (let b = 0; b < numTimeBins; b++) {
+    const row: number[] = [];
+    const n = counts[b];
+    for (let d = 0; d < actionDim; d++) {
+      if (n < 2) { row.push(0); continue; }
+      const mean = sums[b][d] / n;
+      row.push(sumsSq[b][d] / n - mean * mean);
+    }
+    variance.push(row);
+  }
+
+  // Per-episode average movement per frame: mean L2 norm of frame-to-frame action deltas
+  const movementScores: LowMovementEpisode[] = episodeActions.map(({ index, actions: ep }) => {
+    if (ep.length < 2) return { episodeIndex: index, totalMovement: 0 };
+    let total = 0;
+    for (let t = 1; t < ep.length; t++) {
+      let sumSq = 0;
+      for (let d = 0; d < actionDim; d++) {
+        const delta = (ep[t][d] ?? 0) - (ep[t - 1][d] ?? 0);
+        sumSq += delta * delta;
+      }
+      total += Math.sqrt(sumSq);
+    }
+    const avgPerFrame = total / (ep.length - 1);
+    return { episodeIndex: index, totalMovement: Math.round(avgPerFrame * 10000) / 10000 };
+  });
+
+  movementScores.sort((a, b) => a.totalMovement - b.totalMovement);
+  const lowMovementEpisodes = movementScores.slice(0, 10);
+
+  // Aggregated velocity stats: pool deltas from all episodes
+  const shortName = (k: string) => { const p = k.split(SERIES_NAME_DELIMITER); return p.length > 1 ? p[p.length - 1] : k; };
+
+  const aggVelocity: AggVelocityStat[] = (() => {
+    const binCount = 30;
+    return Array.from({ length: actionDim }, (_, d) => {
+      const deltas: number[] = [];
+      for (const { actions: ep } of episodeActions) {
+        for (let t = 1; t < ep.length; t++) {
+          deltas.push((ep[t][d] ?? 0) - (ep[t - 1][d] ?? 0));
+        }
+      }
+      if (deltas.length === 0) return { name: shortName(actionNames[d]), std: 0, maxAbs: 0, bins: [], lo: 0, hi: 0 };
+      let sum = 0, maxAbs = 0, lo = Infinity, hi = -Infinity;
+      for (const v of deltas) { sum += v; const a = Math.abs(v); if (a > maxAbs) maxAbs = a; if (v < lo) lo = v; if (v > hi) hi = v; }
+      const mean = sum / deltas.length;
+      let varSum = 0; for (const v of deltas) varSum += (v - mean) ** 2;
+      const std = Math.sqrt(varSum / deltas.length);
+      const range = hi - lo || 1;
+      const binW = range / binCount;
+      const bins = new Array(binCount).fill(0);
+      for (const v of deltas) { let b = Math.floor((v - lo) / binW); if (b >= binCount) b = binCount - 1; bins[b]++; }
+      return { name: shortName(actionNames[d]), std, maxAbs, bins, lo, hi };
+    });
+  })();
+
+  // Aggregated autocorrelation: average per-episode ACFs
+  const aggAutocorrelation: AggAutocorrelation | null = (() => {
+    const maxLag = Math.min(100, Math.floor(
+      episodeActions.reduce((min, e) => Math.min(min, e.actions.length), Infinity) / 2
+    ));
+    if (maxLag < 2) return null;
+
+    const avgAcf: number[][] = Array.from({ length: actionDim }, () => new Array(maxLag).fill(0));
+    let epCount = 0;
+
+    for (const { actions: ep } of episodeActions) {
+      if (ep.length < maxLag * 2) continue;
+      epCount++;
+      for (let d = 0; d < actionDim; d++) {
+        const vals = ep.map(row => row[d] ?? 0);
+        const n = vals.length;
+        const m = vals.reduce((a, b) => a + b, 0) / n;
+        const centered = vals.map(v => v - m);
+        const vari = centered.reduce((a, v) => a + v * v, 0);
+        if (vari === 0) continue;
+        for (let lag = 1; lag <= maxLag; lag++) {
+          let s = 0;
+          for (let t = 0; t < n - lag; t++) s += centered[t] * centered[t + lag];
+          avgAcf[d][lag - 1] += s / vari;
+        }
+      }
+    }
+
+    if (epCount === 0) return null;
+    for (let d = 0; d < actionDim; d++) for (let l = 0; l < maxLag; l++) avgAcf[d][l] /= epCount;
+
+    const shortKeys = actionNames.map(shortName);
+    const chartData = Array.from({ length: maxLag }, (_, lag) => {
+      const row: Record<string, number> = { lag: lag + 1, time: (lag + 1) / fps };
+      shortKeys.forEach((k, d) => { row[k] = avgAcf[d][lag]; });
+      return row;
+    });
+
+    // Suggested chunk: median lag where ACF drops below 0.5
+    const lags = avgAcf.map(acf => { const i = acf.findIndex(v => v < 0.5); return i >= 0 ? i + 1 : null; }).filter(Boolean) as number[];
+    const suggestedChunk = lags.length > 0 ? lags.sort((a, b) => a - b)[Math.floor(lags.length / 2)] : null;
+
+    return { chartData, suggestedChunk, shortKeys };
+  })();
+
+  return { actionNames, timeBins, variance, numEpisodes: episodeActions.length, lowMovementEpisodes, aggVelocity, aggAutocorrelation };
+}
+
+// Load only flatChartData for a specific episode (used by URDF viewer episode switching)
+export async function loadEpisodeFlatChartData(
+  repoId: string,
+  version: string,
+  info: DatasetMetadata,
+  episodeId: number,
+): Promise<Record<string, number>[]> {
+  const episodeMetadata = await loadEpisodeMetadataV3Simple(repoId, version, episodeId);
+  const { flatChartData } = await loadEpisodeDataV3(repoId, version, info, episodeMetadata);
+  return flatChartData;
+}
+
 // Safe wrapper for UI error display
 export async function getEpisodeDataSafe(
   org: string,

src/app/page.tsx

@@ -130,14 +130,6 @@ function HomeInner() {
         <h1 className="text-4xl md:text-5xl font-bold mb-6 drop-shadow-lg">
           LeRobot Dataset Visualizer
         </h1>
-        <a
-          href="https://x.com/RemiCadene/status/1825455895561859185"
-          target="_blank"
-          rel="noopener noreferrer"
-          className="text-sky-400 font-medium text-lg underline mb-8 inline-block hover:text-sky-300 transition-colors"
-        >
-          create & train your own robots
-        </a>
         <form onSubmit={handleGo} className="flex gap-2 justify-center mt-6">
           <input
             ref={inputRef}

src/components/action-insights-panel.tsx

@@ -0,0 +1,545 @@
+"use client";
+
+import React, { useMemo } from "react";
+import {
+  LineChart,
+  Line,
+  XAxis,
+  YAxis,
+  CartesianGrid,
+  ResponsiveContainer,
+  Tooltip,
+} from "recharts";
+import type { CrossEpisodeVarianceData, LowMovementEpisode, AggVelocityStat, AggAutocorrelation } from "@/app/[org]/[dataset]/[episode]/fetch-data";
+
+const DELIMITER = " | ";
+const COLORS = [
+  "#f97316", "#3b82f6", "#22c55e", "#ef4444", "#a855f7",
+  "#eab308", "#06b6d4", "#ec4899", "#14b8a6", "#f59e0b",
+  "#6366f1", "#84cc16",
+];
+
+function shortName(key: string): string {
+  const parts = key.split(DELIMITER);
+  return parts.length > 1 ? parts[parts.length - 1] : key;
+}
+
+function getActionKeys(row: Record<string, number>): string[] {
+  return Object.keys(row)
+    .filter(k => k.startsWith("action") && k !== "timestamp")
+    .sort();
+}
+
+// ─── Autocorrelation ─────────────────────────────────────────────
+
+function computeAutocorrelation(values: number[], maxLag: number): number[] {
+  const n = values.length;
+  const mean = values.reduce((a, b) => a + b, 0) / n;
+  const centered = values.map(v => v - mean);
+  const variance = centered.reduce((a, v) => a + v * v, 0);
+  if (variance === 0) return Array(maxLag).fill(0);
+
+  const result: number[] = [];
+  for (let lag = 1; lag <= maxLag; lag++) {
+    let sum = 0;
+    for (let t = 0; t < n - lag; t++) sum += centered[t] * centered[t + lag];
+    result.push(sum / variance);
+  }
+  return result;
+}
+
+function findDecorrelationLag(acf: number[], threshold = 0.5): number | null {
+  const idx = acf.findIndex(v => v < threshold);
+  return idx >= 0 ? idx + 1 : null;
+}
+
+function AutocorrelationSection({ data, fps, agg, numEpisodes }: { data: Record<string, number>[]; fps: number; agg?: AggAutocorrelation | null; numEpisodes?: number }) {
+  const actionKeys = useMemo(() => (data.length > 0 ? getActionKeys(data[0]) : []), [data]);
+  const maxLag = useMemo(() => Math.min(Math.floor(data.length / 2), 100), [data]);
+
+  const fallback = useMemo(() => {
+    if (agg) return null;
+    if (actionKeys.length === 0 || maxLag < 2) return { chartData: [], suggestedChunk: null, shortKeys: [] as string[] };
+
+    const acfs = actionKeys.map(key => {
+      const values = data.map(row => row[key] ?? 0);
+      return computeAutocorrelation(values, maxLag);
+    });
+
+    const rows = Array.from({ length: maxLag }, (_, lag) => {
+      const row: Record<string, number> = { lag: lag + 1, time: (lag + 1) / fps };
+      actionKeys.forEach((key, ki) => { row[shortName(key)] = acfs[ki][lag]; });
+      return row;
+    });
+
+    const lags = acfs.map(acf => findDecorrelationLag(acf, 0.5)).filter(Boolean) as number[];
+    const suggested = lags.length > 0 ? lags.sort((a, b) => a - b)[Math.floor(lags.length / 2)] : null;
+
+    return { chartData: rows, suggestedChunk: suggested, shortKeys: actionKeys.map(shortName) };
+  }, [data, actionKeys, maxLag, fps, agg]);
+
+  const { chartData, suggestedChunk, shortKeys } = agg ?? fallback ?? { chartData: [], suggestedChunk: null, shortKeys: [] };
+  const numEpisodesLabel = agg ? ` (${numEpisodes} episodes sampled)` : " (current episode)";
+
+  if (shortKeys.length === 0) return <p className="text-slate-500 italic">No action columns found.</p>;
+
+  return (
+    <div className="bg-slate-800/60 rounded-lg p-5 border border-slate-700 space-y-4">
+      <div>
+        <h3 className="text-sm font-semibold text-slate-200">Action Autocorrelation<span className="text-xs text-slate-500 ml-2 font-normal">{numEpisodesLabel}</span></h3>
+        <p className="text-xs text-slate-400 mt-1">
+          Shows how correlated each action dimension is with itself over increasing time lags.
+          Where autocorrelation drops below 0.5 suggests a <span className="text-orange-400 font-medium">natural action chunk boundary</span> — actions
+          beyond this lag are essentially independent, so executing them open-loop offers diminishing returns.
+          <br />
+          <span className="text-slate-500">
+            Grounded in the theoretical result that chunk length should scale logarithmically with system stability constants
+            (Zhang et al., 2025 — arXiv:2507.09061, Theorem 1).
+          </span>
+        </p>
+      </div>
+
+      {suggestedChunk && (
+        <div className="flex items-center gap-3 bg-orange-500/10 border border-orange-500/30 rounded-md px-4 py-2.5">
+          <span className="text-orange-400 font-bold text-lg tabular-nums">{suggestedChunk}</span>
+          <div>
+            <p className="text-sm text-orange-300 font-medium">
+              Suggested chunk length: {suggestedChunk} steps ({(suggestedChunk / fps).toFixed(2)}s)
+            </p>
+            <p className="text-xs text-slate-400">Median lag where autocorrelation drops below 0.5 across action dimensions</p>
+          </div>
+        </div>
+      )}
+
+      <div className="h-64">
+        <ResponsiveContainer width="100%" height="100%">
+          <LineChart data={chartData} margin={{ top: 8, right: 16, left: 0, bottom: 16 }}>
+            <CartesianGrid strokeDasharray="3 3" stroke="#334155" />
+            <XAxis
+              dataKey="lag"
+              stroke="#94a3b8"
+              label={{ value: "Lag (steps)", position: "insideBottom", offset: -8, fill: "#94a3b8", fontSize: 11 }}
+            />
+            <YAxis stroke="#94a3b8" domain={[-0.2, 1]} />
+            <Tooltip
+              contentStyle={{ background: "#1e293b", border: "1px solid #475569", borderRadius: 6 }}
+              labelFormatter={(v) => `Lag ${v} (${(Number(v) / fps).toFixed(2)}s)`}
+              formatter={(v: number) => v.toFixed(3)}
+            />
+            <Line
+              dataKey={() => 0.5}
+              stroke="#64748b"
+              strokeDasharray="6 4"
+              dot={false}
+              name="0.5 threshold"
+              legendType="none"
+              isAnimationActive={false}
+            />
+            {shortKeys.map((name, i) => (
+              <Line
+                key={name}
+                dataKey={name}
+                stroke={COLORS[i % COLORS.length]}
+                dot={false}
+                strokeWidth={1.5}
+                legendType="none"
+                isAnimationActive={false}
+              />
+            ))}
+          </LineChart>
+        </ResponsiveContainer>
+      </div>
+
+      {/* Custom legend */}
+      <div className="flex flex-wrap gap-x-4 gap-y-1 px-1">
+        {shortKeys.map((name, i) => (
+          <div key={name} className="flex items-center gap-1.5">
+            <span className="w-3 h-[3px] rounded-full shrink-0" style={{ background: COLORS[i % COLORS.length] }} />
+            <span className="text-[11px] text-slate-400">{name}</span>
+          </div>
+        ))}
+      </div>
+    </div>
+  );
+}
+
+// ─── Action Velocity ─────────────────────────────────────────────
+
+function ActionVelocitySection({ data, agg, numEpisodes }: { data: Record<string, number>[]; agg?: AggVelocityStat[]; numEpisodes?: number }) {
+  const actionKeys = useMemo(() => (data.length > 0 ? getActionKeys(data[0]) : []), [data]);
+
+  const fallbackStats = useMemo(() => {
+    if (agg && agg.length > 0) return null;
+    if (actionKeys.length === 0 || data.length < 2) return [];
+
+    return actionKeys.map(key => {
+      const values = data.map(row => row[key] ?? 0);
+      const deltas = values.slice(1).map((v, i) => v - values[i]);
+      const mean = deltas.reduce((a, b) => a + b, 0) / deltas.length;
+      const std = Math.sqrt(deltas.reduce((a, d) => a + (d - mean) ** 2, 0) / deltas.length);
+      const maxAbs = Math.max(...deltas.map(Math.abs));
+      const binCount = 30;
+      const lo = Math.min(...deltas);
+      const hi = Math.max(...deltas);
+      const range = hi - lo || 1;
+      const binW = range / binCount;
+      const bins: number[] = new Array(binCount).fill(0);
+      for (const d of deltas) { let b = Math.floor((d - lo) / binW); if (b >= binCount) b = binCount - 1; bins[b]++; }
+      return { name: shortName(key), std, maxAbs, bins, lo, hi };
+    });
+  }, [data, actionKeys, agg]);
+
+  const stats = (agg && agg.length > 0) ? agg : fallbackStats ?? [];
+  const isAgg = agg && agg.length > 0;
+
+  if (stats.length === 0) return <p className="text-slate-500 italic">No action data for velocity analysis.</p>;
+
+  const maxBinCount = Math.max(...stats.flatMap(s => s.bins));
+  const maxStd = Math.max(...stats.map(s => s.std));
+
+  const insight = useMemo(() => {
+    const smooth = stats.filter(s => s.std / maxStd < 0.4);
+    const moderate = stats.filter(s => s.std / maxStd >= 0.4 && s.std / maxStd < 0.7);
+    const jerky = stats.filter(s => s.std / maxStd >= 0.7);
+    const isGripper = (n: string) => /grip/i.test(n);
+    const jerkyNonGripper = jerky.filter(s => !isGripper(s.name));
+    const jerkyGripper = jerky.filter(s => isGripper(s.name));
+    const smoothRatio = smooth.length / stats.length;
+
+    let verdict: { label: string; color: string };
+    if (smoothRatio >= 0.6 && jerkyNonGripper.length === 0)
+      verdict = { label: "Smooth", color: "text-green-400" };
+    else if (jerkyNonGripper.length <= 2 && smoothRatio >= 0.3)
+      verdict = { label: "Moderate", color: "text-yellow-400" };
+    else
+      verdict = { label: "Jerky", color: "text-red-400" };
+
+    const lines: string[] = [];
+    if (smooth.length > 0)
+      lines.push(`${smooth.length} smooth (${smooth.map(s => s.name).join(", ")})`);
+    if (moderate.length > 0)
+      lines.push(`${moderate.length} moderate (${moderate.map(s => s.name).join(", ")})`);
+    if (jerkyNonGripper.length > 0)
+      lines.push(`${jerkyNonGripper.length} jerky (${jerkyNonGripper.map(s => s.name).join(", ")})`);
+    if (jerkyGripper.length > 0)
+      lines.push(`${jerkyGripper.length} gripper${jerkyGripper.length > 1 ? "s" : ""} jerky — expected for binary open/close`);
+
+    let tip: string;
+    if (verdict.label === "Smooth")
+      tip = "Actions are consistent — longer action chunks should work well.";
+    else if (verdict.label === "Moderate")
+      tip = "Some dimensions show abrupt changes. Consider moderate chunk sizes.";
+    else
+      tip = "Many dimensions are jerky. Use shorter action chunks and consider filtering outlier episodes.";
+
+    return { verdict, lines, tip };
+  }, [stats, maxStd]);
+
+  return (
+    <div className="bg-slate-800/60 rounded-lg p-5 border border-slate-700 space-y-4">
+      <div>
+        <h3 className="text-sm font-semibold text-slate-200">Action Velocity (Δa) — Smoothness Proxy<span className="text-xs text-slate-500 ml-2 font-normal">{isAgg ? `(${numEpisodes} episodes sampled)` : "(current episode)"}</span></h3>
+        <p className="text-xs text-slate-400 mt-1">
+          Shows the distribution of frame-to-frame action changes (Δa = a<sub>t+1</sub> − a<sub>t</sub>) for each dimension.
+          A <span className="text-green-400">tight distribution around zero</span> means smooth, predictable control — the system
+          is likely stable and benefits from longer action chunks.
+          <span className="text-red-400"> Fat tails or high std</span> indicate jerky demonstrations, suggesting shorter chunks
+          and potentially beneficial noise injection.
+          <br />
+          <span className="text-slate-500">
+            Relates to the Lipschitz constant L<sub>π</sub> and smoothness C<sub>π</sub> in Zhang et al. (2025), which govern
+            compounding error bounds (Assumptions 3.1, 4.1).
+          </span>
+        </p>
+      </div>
+
+      {/* Per-dimension mini histograms + stats */}
+      <div className="grid gap-2" style={{ gridTemplateColumns: "repeat(auto-fill, minmax(180px, 1fr))" }}>
+        {stats.map((s, si) => {
+          const barH = 28;
+          return (
+            <div key={s.name} className="bg-slate-900/50 rounded-md px-2.5 py-2 space-y-1">
+              <p className="text-[11px] font-medium text-slate-200 truncate" title={s.name}>{s.name}</p>
+              <div className="flex gap-2 text-[9px] text-slate-400 tabular-nums">
+                <span>σ={s.std.toFixed(4)}</span>
+                <span>|Δ|<sub>max</sub>={s.maxAbs.toFixed(4)}</span>
+              </div>
+              <svg width="100%" viewBox={`0 0 ${s.bins.length} ${barH}`} preserveAspectRatio="none" className="h-7 rounded" aria-label={`Δa distribution for ${s.name}`}>
+                {[...s.bins].map((count, bi) => {
+                  const h = maxBinCount > 0 ? (count / maxBinCount) * barH : 0;
+                  return <rect key={bi} x={bi} y={barH - h} width={0.85} height={h} fill={COLORS[si % COLORS.length]} opacity={0.7} />;
+                })}
+              </svg>
+              <div className="h-1 w-full bg-slate-700 rounded-full overflow-hidden">
+                <div
+                  className="h-full rounded-full"
+                  style={{
+                    width: `${Math.min(100, (s.std / maxStd) * 100)}%`,
+                    background: s.std / maxStd < 0.4 ? "#22c55e" : s.std / maxStd < 0.7 ? "#eab308" : "#ef4444",
+                  }}
+                />
+              </div>
+            </div>
+          );
+        })}
+      </div>
+
+      <div className="bg-slate-900/60 rounded-md px-4 py-3 border border-slate-700/60 space-y-1.5">
+        <p className="text-sm font-medium text-slate-200">
+          Overall: <span className={insight.verdict.color}>{insight.verdict.label}</span>
+        </p>
+        <ul className="text-xs text-slate-400 space-y-0.5 list-disc list-inside">
+          {insight.lines.map((l, i) => <li key={i}>{l}</li>)}
+        </ul>
+        <p className="text-xs text-slate-500 pt-1">{insight.tip}</p>
+      </div>
+    </div>
+  );
+}
+
+// ─── Cross-Episode Variance Heatmap ──────────────────────────────
+
+function VarianceHeatmap({ data, loading }: { data: CrossEpisodeVarianceData | null; loading: boolean }) {
+  if (loading) {
+    return (
+      <div className="bg-slate-800/60 rounded-lg p-5 border border-slate-700">
+        <h3 className="text-sm font-semibold text-slate-200 mb-2">Cross-Episode Action Variance</h3>
+        <div className="flex items-center gap-2 text-slate-400 text-sm py-8 justify-center">
+          <svg className="animate-spin h-4 w-4" viewBox="0 0 24 24" fill="none">
+            <circle className="opacity-25" cx="12" cy="12" r="10" stroke="currentColor" strokeWidth="4" />
+            <path className="opacity-75" fill="currentColor" d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4z" />
+          </svg>
+          Loading cross-episode data (sampled up to 500 episodes)…
+        </div>
+      </div>
+    );
+  }
+
+  if (!data) {
+    return (
+      <div className="bg-slate-800/60 rounded-lg p-5 border border-slate-700">
+        <h3 className="text-sm font-semibold text-slate-200 mb-2">Cross-Episode Action Variance</h3>
+        <p className="text-slate-500 italic text-sm">Not enough episodes or no action data to compute variance.</p>
+      </div>
+    );
+  }
+
+  const { actionNames, timeBins, variance, numEpisodes } = data;
+  const numDims = actionNames.length;
+  const numBins = timeBins.length;
+
+  // Find global max variance for color scale
+  const maxVar = Math.max(...variance.flat(), 1e-10);
+
+  // Heatmap dimensions
+  const cellW = Math.max(6, Math.min(14, Math.floor(560 / numBins)));
+  const cellH = Math.max(20, Math.min(36, Math.floor(300 / numDims)));
+  const labelW = 100;
+  const svgW = labelW + numBins * cellW + 60; // 60 for color bar
+  const svgH = numDims * cellH + 40; // 40 for x-axis label
+
+  function varColor(v: number): string {
+    const t = Math.sqrt(v / maxVar); // sqrt for better visual spread
+    // Dark blue → teal → orange
+    const r = Math.round(t * 249);
+    const g = Math.round(t < 0.5 ? 80 + t * 200 : 180 - (t - 0.5) * 200);
+    const b = Math.round((1 - t) * 200 + 30);
+    return `rgb(${r},${g},${b})`;
+  }
+
+  return (
+    <div className="bg-slate-800/60 rounded-lg p-5 border border-slate-700 space-y-4">
+      <div>
+        <h3 className="text-sm font-semibold text-slate-200">
+          Cross-Episode Action Variance
+          <span className="text-xs text-slate-500 ml-2 font-normal">({numEpisodes} episodes sampled)</span>
+        </h3>
+        <p className="text-xs text-slate-400 mt-1">
+          Shows how much each action dimension varies across episodes at each point in time (normalized 0–100%).
+          <span className="text-orange-400"> High-variance regions</span> indicate multi-modal or inconsistent demonstrations —
+          generative policies (diffusion, flow-matching) and action chunking help here by modeling multiple modes.
+          <span className="text-blue-400"> Low-variance regions</span> indicate consistent behavior across demonstrations.
+          <br />
+          <span className="text-slate-500">
+            Relates to the &quot;coverage&quot; discussion in Zhang et al. (2025) — regions with low variance may lack the
+            exploratory coverage needed to prevent compounding errors (Section 4).
+          </span>
+        </p>
+      </div>
+
+      <div className="overflow-x-auto">
+        <svg width={svgW} height={svgH} className="block">
+          {/* Heatmap cells */}
+          {variance.map((row, bi) =>
+            row.map((v, di) => (
+              <rect
+                key={`${bi}-${di}`}
+                x={labelW + bi * cellW}
+                y={di * cellH}
+                width={cellW}
+                height={cellH}
+                fill={varColor(v)}
+                stroke="#1e293b"
+                strokeWidth={0.5}
+              >
+                <title>{`${shortName(actionNames[di])} @ ${(timeBins[bi] * 100).toFixed(0)}%: var=${v.toFixed(5)}`}</title>
+              </rect>
+            ))
+          )}
+
+          {/* Y-axis: action names */}
+          {actionNames.map((name, di) => (
+            <text
+              key={di}
+              x={labelW - 4}
+              y={di * cellH + cellH / 2}
+              textAnchor="end"
+              dominantBaseline="central"
+              className="fill-slate-400"
+              fontSize={Math.min(11, cellH - 4)}
+            >
+              {shortName(name)}
+            </text>
+          ))}
+
+          {/* X-axis labels */}
+          {[0, 0.25, 0.5, 0.75, 1].map(frac => {
+            const binIdx = Math.round(frac * (numBins - 1));
+            return (
+              <text
+                key={frac}
+                x={labelW + binIdx * cellW + cellW / 2}
+                y={numDims * cellH + 14}
+                textAnchor="middle"
+                className="fill-slate-400"
+                fontSize={9}
+              >
+                {(frac * 100).toFixed(0)}%
+              </text>
+            );
+          })}
+          <text
+            x={labelW + (numBins * cellW) / 2}
+            y={numDims * cellH + 30}
+            textAnchor="middle"
+            className="fill-slate-500"
+            fontSize={10}
+          >
+            Episode progress
+          </text>
+
+          {/* Color bar */}
+          {Array.from({ length: 10 }, (_, i) => {
+            const t = i / 9;
+            const barX = labelW + numBins * cellW + 16;
+            const barH = (numDims * cellH) / 10;
+            return (
+              <rect
+                key={i}
+                x={barX}
+                y={(9 - i) * barH}
+                width={12}
+                height={barH}
+                fill={varColor(t * maxVar)}
+              />
+            );
+          })}
+          <text
+            x={labelW + numBins * cellW + 34}
+            y={10}
+            className="fill-slate-500"
+            fontSize={8}
+            dominantBaseline="central"
+          >
+            high
+          </text>
+          <text
+            x={labelW + numBins * cellW + 34}
+            y={numDims * cellH - 4}
+            className="fill-slate-500"
+            fontSize={8}
+            dominantBaseline="central"
+          >
+            low
+          </text>
+        </svg>
+      </div>
+    </div>
+  );
+}
+
+// ─── Low-Movement Episodes ──────────────────────────────────────
+
+function LowMovementSection({ episodes }: { episodes: LowMovementEpisode[] }) {
+  if (episodes.length === 0) return null;
+
+  const maxMovement = Math.max(...episodes.map(e => e.totalMovement), 1e-10);
+
+  return (
+    <div className="bg-slate-800/60 rounded-lg p-5 border border-slate-700 space-y-3">
+      <div>
+        <h3 className="text-sm font-semibold text-slate-200">Lowest-Movement Episodes</h3>
+        <p className="text-xs text-slate-400 mt-1">
+          Episodes with the lowest average action change per frame (mean ‖Δa<sub>t</sub>‖). Very low values may indicate the robot
+          was <span className="text-yellow-400">standing still</span> or the episode was recorded incorrectly.
+        </p>
+      </div>
+      <div className="grid gap-2" style={{ gridTemplateColumns: "repeat(auto-fill, minmax(220px, 1fr))" }}>
+        {episodes.map(ep => (
+          <div key={ep.episodeIndex} className="bg-slate-900/50 rounded-md px-3 py-2 flex items-center gap-3">
+            <span className="text-xs text-slate-300 font-medium shrink-0">ep {ep.episodeIndex}</span>
+            <div className="flex-1 min-w-0">
+              <div className="h-1.5 bg-slate-700 rounded-full overflow-hidden">
+                <div
+                  className="h-full rounded-full"
+                  style={{
+                    width: `${Math.max(2, (ep.totalMovement / maxMovement) * 100)}%`,
+                    background: ep.totalMovement / maxMovement < 0.15 ? "#ef4444" : ep.totalMovement / maxMovement < 0.4 ? "#eab308" : "#22c55e",
+                  }}
+                />
+              </div>
+            </div>
+            <span className="text-[10px] text-slate-500 tabular-nums shrink-0">{ep.totalMovement.toFixed(2)}</span>
+          </div>
+        ))}
+      </div>
+    </div>
+  );
+}
+
+// ─── Main Panel ──────────────────────────────────────────────────
+
+interface ActionInsightsPanelProps {
+  flatChartData: Record<string, number>[];
+  fps: number;
+  crossEpisodeData: CrossEpisodeVarianceData | null;
+  crossEpisodeLoading: boolean;
+}
+
+const ActionInsightsPanel: React.FC<ActionInsightsPanelProps> = ({
+  flatChartData,
+  fps,
+  crossEpisodeData,
+  crossEpisodeLoading,
+}) => {
+  return (
+    <div className="max-w-5xl mx-auto py-6 space-y-8">
+      <div>
+        <h2 className="text-xl font-bold text-slate-100">Action Insights</h2>
+        <p className="text-sm text-slate-400 mt-1">
+          Data-driven analysis to guide action chunking, data quality assessment, and training configuration.
+        </p>
+      </div>
+
+      <AutocorrelationSection data={flatChartData} fps={fps} agg={crossEpisodeData?.aggAutocorrelation} numEpisodes={crossEpisodeData?.numEpisodes} />
+      <ActionVelocitySection data={flatChartData} agg={crossEpisodeData?.aggVelocity} numEpisodes={crossEpisodeData?.numEpisodes} />
+      <VarianceHeatmap data={crossEpisodeData} loading={crossEpisodeLoading} />
+      {crossEpisodeData?.lowMovementEpisodes && (
+        <LowMovementSection episodes={crossEpisodeData.lowMovementEpisodes} />
+      )}
+    </div>
+  );
+};
+
+export default ActionInsightsPanel;
+

src/components/data-recharts.tsx

@@ -10,6 +10,7 @@ import {
   CartesianGrid,
   ResponsiveContainer,
   Tooltip,
+  ReferenceLine,
 } from "recharts";
 
 type ChartRow = Record<string, number | Record<string, number>>;
@@ -21,9 +22,14 @@ type DataGraphProps = {
 
 import React, { useMemo } from "react";
 
-// Use the same delimiter as the data processing
 const SERIES_NAME_DELIMITER = " | ";
 
+const CHART_COLORS = [
+  "#f97316", "#3b82f6", "#22c55e", "#ef4444", "#a855f7",
+  "#eab308", "#06b6d4", "#ec4899", "#14b8a6", "#f59e0b",
+  "#6366f1", "#84cc16",
+];
+
 export const DataRecharts = React.memo(
   ({ data, onChartsReady }: DataGraphProps) => {
     // Shared hoveredTime for all graphs
@@ -112,11 +118,10 @@ const SingleDataGraph = React.memo(
       }
     });
 
-    // Assign a color per group (and for singles)
     const allGroups = [...Object.keys(groups), ...singles];
     const groupColorMap: Record<string, string> = {};
     allGroups.forEach((group, idx) => {
-      groupColorMap[group] = `hsl(${idx * (360 / allGroups.length)}, 100%, 50%)`;
+      groupColorMap[group] = CHART_COLORS[idx % CHART_COLORS.length];
     });
 
     // Find the closest data point to the current time for highlighting
@@ -160,11 +165,10 @@ const SingleDataGraph = React.memo(
         }
       });
 
-      // Assign a color per group (and for singles)
       const allGroups = [...Object.keys(groups), ...singles];
       const groupColorMap: Record<string, string> = {};
       allGroups.forEach((group, idx) => {
-        groupColorMap[group] = `hsl(${idx * (360 / allGroups.length)}, 100%, 50%)`;
+        groupColorMap[group] = CHART_COLORS[idx % CHART_COLORS.length];
       });
 
       const isGroupChecked = (group: string) => groups[group].every(k => visibleKeys.includes(k));
@@ -187,58 +191,59 @@ const SingleDataGraph = React.memo(
       };
 
       return (
-        <div className="grid grid-cols-[repeat(auto-fit,250px)] gap-4 mx-8">
-          {/* Grouped keys */}
+        <div className="flex flex-wrap gap-x-5 gap-y-2 px-1 pt-2">
           {Object.entries(groups).map(([group, children]) => {
             const color = groupColorMap[group];
             return (
-              <div key={group} className="mb-2">
-                <label className="flex gap-2 cursor-pointer select-none font-semibold">
+              <div key={group}>
+                <label className="flex items-center gap-1.5 cursor-pointer select-none">
                   <input
                     type="checkbox"
                     checked={isGroupChecked(group)}
                     ref={el => { if (el) el.indeterminate = isGroupIndeterminate(group); }}
                     onChange={() => handleGroupCheckboxChange(group)}
-                    className="size-3.5 mt-1"
+                    className="size-3"
                     style={{ accentColor: color }}
                   />
-                  <span className="text-sm w-40 text-white">{group}</span>
+                  <span className="text-[11px] font-semibold text-slate-200">{group}</span>
                 </label>
-                <div className="pl-7 flex flex-col gap-1 mt-1">
-                  {children.map((key) => (
-                    <label key={key} className="flex gap-2 cursor-pointer select-none">
-                      <input
-                        type="checkbox"
-                        checked={visibleKeys.includes(key)}
-                        onChange={() => handleCheckboxChange(key)}
-                        className="size-3.5 mt-1"
-                        style={{ accentColor: color }}
-                      />
-                      <span className={`text-xs break-all w-36 ${visibleKeys.includes(key) ? "text-white" : "text-gray-400"}`}>{key.slice(group.length + 1)}</span>
-                      <span className={`text-xs font-mono ml-auto ${visibleKeys.includes(key) ? "text-orange-300" : "text-gray-500"}`}>
-                        {typeof currentData[key] === "number" ? currentData[key].toFixed(2) : "--"}
-                      </span>
-                    </label>
-                  ))}
+                <div className="pl-5 flex flex-col gap-0.5 mt-0.5">
+                  {children.map((key) => {
+                    const label = key.split(SERIES_NAME_DELIMITER).pop() ?? key;
+                    return (
+                      <label key={key} className="flex items-center gap-1.5 cursor-pointer select-none">
+                        <input
+                          type="checkbox"
+                          checked={visibleKeys.includes(key)}
+                          onChange={() => handleCheckboxChange(key)}
+                          className="size-2.5"
+                          style={{ accentColor: color }}
+                        />
+                        <span className={`text-[10px] ${visibleKeys.includes(key) ? "text-slate-300" : "text-slate-500"}`}>{label}</span>
+                        <span className={`text-[10px] font-mono tabular-nums ml-1 ${visibleKeys.includes(key) ? "text-orange-300/80" : "text-slate-600"}`}>
+                          {typeof currentData[key] === "number" ? currentData[key].toFixed(2) : "–"}
+                        </span>
+                      </label>
+                    );
+                  })}
                 </div>
               </div>
             );
           })}
-          {/* Singles (non-grouped) */}
           {singles.map((key) => {
             const color = groupColorMap[key];
             return (
-              <label key={key} className="flex gap-2 cursor-pointer select-none">
+              <label key={key} className="flex items-center gap-1.5 cursor-pointer select-none">
                 <input
                   type="checkbox"
                   checked={visibleKeys.includes(key)}
                   onChange={() => handleCheckboxChange(key)}
-                  className="size-3.5 mt-1"
+                  className="size-3"
                   style={{ accentColor: color }}
                 />
-                <span className={`text-sm break-all w-40 ${visibleKeys.includes(key) ? "text-white" : "text-gray-400"}`}>{key}</span>
-                <span className={`text-sm font-mono ml-auto ${visibleKeys.includes(key) ? "text-orange-300" : "text-gray-500"}`}>
-                  {typeof currentData[key] === "number" ? currentData[key].toFixed(2) : "--"}
+                <span className={`text-[11px] ${visibleKeys.includes(key) ? "text-slate-200" : "text-slate-500"}`}>{key}</span>
+                <span className={`text-[10px] font-mono tabular-nums ml-1 ${visibleKeys.includes(key) ? "text-orange-300/80" : "text-slate-600"}`}>
+                  {typeof currentData[key] === "number" ? currentData[key].toFixed(2) : "–"}
                 </span>
               </label>
             );
@@ -247,14 +252,30 @@ const SingleDataGraph = React.memo(
       );
     };
 
+    // Derive chart title from the grouped feature names
+    const chartTitle = useMemo(() => {
+      const featureNames = Object.keys(groups);
+      if (featureNames.length > 0) {
+        const suffixes = featureNames.map(g => {
+          const parts = g.split(SERIES_NAME_DELIMITER);
+          return parts[parts.length - 1];
+        });
+        return suffixes.join(", ");
+      }
+      return singles.join(", ");
+    }, [groups, singles]);
+
     return (
-      <div className="w-full">
-        <div className="w-full h-80" onMouseLeave={handleMouseLeave}>
+      <div className="w-full bg-slate-800/40 rounded-lg border border-slate-700/50 p-3">
+        {chartTitle && (
+          <p className="text-xs font-medium text-slate-300 mb-1 px-1 truncate" title={chartTitle}>{chartTitle}</p>
+        )}
+        <div className="w-full h-72" onMouseLeave={handleMouseLeave}>
           <ResponsiveContainer width="100%" height="100%">
             <LineChart
               data={chartData}
               syncId="episode-sync"
-              margin={{ top: 24, right: 16, left: 0, bottom: 16 }}
+              margin={{ top: 12, right: 12, left: -8, bottom: 8 }}
               onClick={handleClick}
               onMouseMove={(state) => {
                 const payload = state?.activePayload?.[0]?.payload as { timestamp?: number } | undefined;
@@ -262,33 +283,24 @@ const SingleDataGraph = React.memo(
               }}
               onMouseLeave={handleMouseLeave}
             >
-              <CartesianGrid strokeDasharray="3 3" stroke="#444" />
+              <CartesianGrid strokeDasharray="3 3" stroke="#334155" strokeOpacity={0.6} />
               <XAxis
                 dataKey="timestamp"
-                label={{
-                  value: "time",
-                  position: "insideBottomLeft",
-                  fill: "#cbd5e1",
-                }}
                 domain={[
                   chartData.at(0)?.timestamp ?? 0,
                   chartData.at(-1)?.timestamp ?? 0,
                 ]}
-                ticks={useMemo(
-                  () =>
-                    Array.from(
-                      new Set(chartData.map((d) => Math.ceil(d.timestamp))),
-                    ),
-                  [chartData],
-                )}
-                stroke="#cbd5e1"
-                minTickGap={20} // Increased for fewer ticks
+                tickFormatter={(v: number) => `${v.toFixed(1)}s`}
+                stroke="#64748b"
+                tick={{ fontSize: 10, fill: "#94a3b8" }}
+                minTickGap={30}
                 allowDataOverflow={true}
               />
               <YAxis
                 domain={["auto", "auto"]}
-                stroke="#cbd5e1"
-                interval={0}
+                stroke="#64748b"
+                tick={{ fontSize: 10, fill: "#94a3b8" }}
+                width={45}
                 allowDataOverflow={true}
               />
 
@@ -301,9 +313,14 @@ const SingleDataGraph = React.memo(
                 }
               />
 
-              {/* Render lines for visible dataKeys only */}
+              <ReferenceLine
+                x={currentTime}
+                stroke="#f97316"
+                strokeWidth={1.5}
+                strokeOpacity={0.7}
+              />
+
               {dataKeys.map((key) => {
-                // Use group color for all keys in a group
                 const group = key.includes(SERIES_NAME_DELIMITER) ? key.split(SERIES_NAME_DELIMITER)[0] : key;
                 const color = groupColorMap[group];
                 let strokeDasharray: string | undefined = undefined;

src/components/urdf-viewer.tsx

@@ -1,141 +1,24 @@
 "use client";
 
-import React, { useState, useEffect, useRef, useMemo, useCallback, Suspense } from "react";
-import { Canvas, useFrame, useLoader } from "@react-three/fiber";
-import { OrbitControls, Grid, Environment } from "@react-three/drei";
+import React, { useState, useEffect, useRef, useMemo, useCallback } from "react";
+import { Canvas, useThree, useFrame } from "@react-three/fiber";
+import { OrbitControls, Grid, Html } from "@react-three/drei";
 import * as THREE from "three";
-import { STLLoader } from "three/addons/loaders/STLLoader.js";
-import {
-  SO101_JOINTS,
-  SO101_LINKS,
-  MATERIAL_COLORS,
-  autoMatchJoints,
-  type JointDef,
-  type MeshDef,
-} from "@/lib/so101-robot";
+import URDFLoader from "urdf-loader";
+import type { URDFRobot } from "urdf-loader";
+import { STLLoader } from "three/examples/jsm/loaders/STLLoader.js";
 import type { EpisodeData } from "@/app/[org]/[dataset]/[episode]/fetch-data";
 
 const SERIES_DELIM = " | ";
+const SCALE = 10;
 
-// ─── STL Mesh component ───
-function STLMesh({ mesh }: { mesh: MeshDef }) {
-  const geometry = useLoader(STLLoader, mesh.file);
-  const color = MATERIAL_COLORS[mesh.material];
-  return (
-    <mesh
-      geometry={geometry}
-      position={mesh.origin.xyz}
-      rotation={new THREE.Euler(...mesh.origin.rpy, "XYZ")}
-    >
-      <meshStandardMaterial
-        color={color}
-        metalness={mesh.material === "motor" ? 0.7 : 0.1}
-        roughness={mesh.material === "motor" ? 0.3 : 0.6}
-      />
-    </mesh>
-  );
-}
-
-// ─── Link visual: renders all meshes for a link ───
-function LinkVisual({ linkIndex }: { linkIndex: number }) {
-  const link = SO101_LINKS[linkIndex];
-  if (!link) return null;
-  return (
-    <>
-      {link.meshes.map((mesh, i) => (
-        <STLMesh key={i} mesh={mesh} />
-      ))}
-    </>
-  );
+function getUrdfUrl(robotType: string | null): string {
+  const lower = (robotType ?? "").toLowerCase();
+  if (lower.includes("so100") && !lower.includes("so101")) return "/urdf/so101/so100.urdf";
+  return "/urdf/so101/so101_new_calib.urdf";
 }
 
-// ─── Joint group: applies origin transform + joint rotation ───
-function JointGroup({
-  joint,
-  angle,
-  linkIndex,
-  children,
-}: {
-  joint: JointDef;
-  angle: number;
-  linkIndex: number;
-  children?: React.ReactNode;
-}) {
-  const rotRef = useRef<THREE.Group>(null);
-
-  useEffect(() => {
-    if (rotRef.current) {
-      rotRef.current.quaternion.setFromAxisAngle(new THREE.Vector3(...joint.axis), angle);
-    }
-  }, [angle, joint.axis]);
-
-  return (
-    <group position={joint.origin.xyz} rotation={new THREE.Euler(...joint.origin.rpy, "XYZ")}>
-      <group ref={rotRef}>
-        <LinkVisual linkIndex={linkIndex} />
-        {children}
-      </group>
-    </group>
-  );
-}
-
-// ─── Full robot arm ───
-function RobotArm({ angles }: { angles: Record<string, number> }) {
-  return (
-    <group>
-      {/* Base link (no parent joint) */}
-      <LinkVisual linkIndex={0} />
-
-      {/* shoulder_pan → shoulder_link (1) */}
-      <JointGroup joint={SO101_JOINTS[0]} angle={angles.shoulder_pan ?? 0} linkIndex={1}>
-        {/* shoulder_lift → upper_arm_link (2) */}
-        <JointGroup joint={SO101_JOINTS[1]} angle={angles.shoulder_lift ?? 0} linkIndex={2}>
-          {/* elbow_flex → lower_arm_link (3) */}
-          <JointGroup joint={SO101_JOINTS[2]} angle={angles.elbow_flex ?? 0} linkIndex={3}>
-            {/* wrist_flex → wrist_link (4) */}
-            <JointGroup joint={SO101_JOINTS[3]} angle={angles.wrist_flex ?? 0} linkIndex={4}>
-              {/* wrist_roll → gripper_link (5) */}
-              <JointGroup joint={SO101_JOINTS[4]} angle={angles.wrist_roll ?? 0} linkIndex={5}>
-                {/* gripper → moving_jaw (6) */}
-                <JointGroup joint={SO101_JOINTS[5]} angle={angles.gripper ?? 0} linkIndex={6} />
-              </JointGroup>
-            </JointGroup>
-          </JointGroup>
-        </JointGroup>
-      </JointGroup>
-    </group>
-  );
-}
-
-// ─── Playback driver (advances frame inside Canvas render loop) ───
-function PlaybackDriver({
-  playing,
-  fps,
-  totalFrames,
-  frameRef,
-}: {
-  playing: boolean;
-  fps: number;
-  totalFrames: number;
-  frameRef: React.MutableRefObject<number>;
-}) {
-  const elapsed = useRef(0);
-  useFrame((_, delta) => {
-    if (!playing) {
-      elapsed.current = 0;
-      return;
-    }
-    elapsed.current += delta;
-    const frameDelta = Math.floor(elapsed.current * fps);
-    if (frameDelta > 0) {
-      elapsed.current -= frameDelta / fps;
-      frameRef.current = (frameRef.current + frameDelta) % totalFrames;
-    }
-  });
-  return null;
-}
-
-// ─── Detect raw servo values (0-4096) vs radians ───
+// Detect raw servo values (0-4096) vs radians
 function detectAndConvert(values: number[]): number[] {
   if (values.length === 0) return values;
   const max = Math.max(...values.map(Math.abs));
@@ -143,7 +26,7 @@ function detectAndConvert(values: number[]): number[] {
   return values;
 }
 
-// ─── Group columns by feature prefix ───
+// Group flat chart columns by feature prefix
 function groupColumnsByPrefix(keys: string[]): Record<string, string[]> {
   const groups: Record<string, string[]> = {};
   for (const key of keys) {
@@ -156,6 +39,144 @@ function groupColumnsByPrefix(keys: string[]): Record<string, string[]> {
   return groups;
 }
 
+// Auto-match dataset columns to URDF joint names
+function autoMatchJoints(urdfJointNames: string[], columnKeys: string[]): Record<string, string> {
+  const mapping: Record<string, string> = {};
+  for (const jointName of urdfJointNames) {
+    const lower = jointName.toLowerCase();
+    const exactMatch = columnKeys.find((k) => {
+      const suffix = (k.split(SERIES_DELIM).pop()?.trim() ?? k).toLowerCase();
+      return suffix === lower;
+    });
+    if (exactMatch) { mapping[jointName] = exactMatch; continue; }
+    const fuzzy = columnKeys.find((k) => k.toLowerCase().includes(lower));
+    if (fuzzy) mapping[jointName] = fuzzy;
+  }
+  return mapping;
+}
+
+// ─── Robot scene (imperative, inside Canvas) ───
+function RobotScene({
+  urdfUrl,
+  jointValues,
+  onJointsLoaded,
+}: {
+  urdfUrl: string;
+  jointValues: Record<string, number>;
+  onJointsLoaded: (names: string[]) => void;
+}) {
+  const { scene } = useThree();
+  const robotRef = useRef<URDFRobot | null>(null);
+  const [loading, setLoading] = useState(true);
+  const [error, setError] = useState<string | null>(null);
+
+  useEffect(() => {
+    setLoading(true);
+    setError(null);
+
+    const manager = new THREE.LoadingManager();
+    const loader = new URDFLoader(manager);
+
+    loader.loadMeshCb = (url, mgr, onLoad) => {
+      const stlLoader = new STLLoader(mgr);
+      stlLoader.load(
+        url,
+        (geometry) => {
+          const isMotor = url.includes("sts3215");
+          const material = new THREE.MeshStandardMaterial({
+            color: isMotor ? "#1a1a1a" : "#FFD700",
+            metalness: isMotor ? 0.7 : 0.1,
+            roughness: isMotor ? 0.3 : 0.6,
+          });
+          onLoad(new THREE.Mesh(geometry, material));
+        },
+        undefined,
+        (err) => onLoad(new THREE.Object3D(), err as Error),
+      );
+    };
+
+    loader.load(
+      urdfUrl,
+      (robot) => {
+        robotRef.current = robot;
+        robot.rotateOnAxis(new THREE.Vector3(1, 0, 0), -Math.PI / 2);
+        robot.traverse((c) => { c.castShadow = true; });
+        robot.updateMatrixWorld(true);
+        robot.scale.set(SCALE, SCALE, SCALE);
+        scene.add(robot);
+
+        const revolute = Object.values(robot.joints)
+          .filter((j) => j.jointType === "revolute" || j.jointType === "continuous")
+          .map((j) => j.name);
+        onJointsLoaded(revolute);
+        setLoading(false);
+      },
+      undefined,
+      (err) => {
+        console.error("Error loading URDF:", err);
+        setError(String(err));
+        setLoading(false);
+      },
+    );
+
+    return () => {
+      if (robotRef.current) {
+        scene.remove(robotRef.current);
+        robotRef.current = null;
+      }
+    };
+  }, [urdfUrl, scene, onJointsLoaded]);
+
+  useFrame(() => {
+    if (!robotRef.current) return;
+    for (const [name, value] of Object.entries(jointValues)) {
+      robotRef.current.setJointValue(name, value);
+    }
+  });
+
+  if (loading) return <Html center><span className="text-white text-lg">Loading robot…</span></Html>;
+  if (error) return <Html center><span className="text-red-400">Failed to load URDF</span></Html>;
+  return null;
+}
+
+// ─── Playback ticker (inside Canvas) ───
+function PlaybackDriver({
+  playing, fps, totalFrames, frameRef, setFrame,
+}: {
+  playing: boolean; fps: number; totalFrames: number;
+  frameRef: React.MutableRefObject<number>;
+  setFrame: React.Dispatch<React.SetStateAction<number>>;
+}) {
+  const elapsed = useRef(0);
+  const last = useRef(0);
+
+  useEffect(() => {
+    if (!playing) return;
+    let raf: number;
+    const tick = () => {
+      raf = requestAnimationFrame(tick);
+      const now = performance.now();
+      const dt = (now - last.current) / 1000;
+      last.current = now;
+      if (dt > 0 && dt < 0.5) {
+        elapsed.current += dt;
+        const frameDelta = Math.floor(elapsed.current * fps);
+        if (frameDelta > 0) {
+          elapsed.current -= frameDelta / fps;
+          frameRef.current = (frameRef.current + frameDelta) % totalFrames;
+          setFrame(frameRef.current);
+        }
+      }
+    };
+    last.current = performance.now();
+    elapsed.current = 0;
+    raf = requestAnimationFrame(tick);
+    return () => cancelAnimationFrame(raf);
+  }, [playing, fps, totalFrames, frameRef, setFrame]);
+
+  return null;
+}
+
 // ═══════════════════════════════════════
 // ─── Main URDF Viewer ───
 // ═══════════════════════════════════════
@@ -163,7 +184,13 @@ export default function URDFViewer({ data }: { data: EpisodeData }) {
   const { flatChartData, datasetInfo } = data;
   const totalFrames = flatChartData.length;
   const fps = datasetInfo.fps || 30;
+  const urdfUrl = useMemo(() => getUrdfUrl(datasetInfo.robot_type), [datasetInfo.robot_type]);
 
+  // URDF joint names (set after robot loads)
+  const [urdfJointNames, setUrdfJointNames] = useState<string[]>([]);
+  const onJointsLoaded = useCallback((names: string[]) => setUrdfJointNames(names), []);
+
+  // Feature group selection
   const columnGroups = useMemo(() => {
     if (totalFrames === 0) return {};
     return groupColumnsByPrefix(Object.keys(flatChartData[0]));
@@ -183,47 +210,46 @@ export default function URDFViewer({ data }: { data: EpisodeData }) {
   useEffect(() => setSelectedGroup(defaultGroup), [defaultGroup]);
 
   const selectedColumns = columnGroups[selectedGroup] ?? [];
-  const autoMapping = useMemo(() => autoMatchJoints(selectedColumns), [selectedColumns]);
+
+  // Joint mapping (re-compute when URDF joints or selected columns change)
+  const autoMapping = useMemo(
+    () => autoMatchJoints(urdfJointNames, selectedColumns),
+    [urdfJointNames, selectedColumns],
+  );
   const [mapping, setMapping] = useState<Record<string, string>>(autoMapping);
   useEffect(() => setMapping(autoMapping), [autoMapping]);
 
+  // Playback
   const [frame, setFrame] = useState(0);
   const [playing, setPlaying] = useState(false);
   const frameRef = useRef(0);
 
-  useEffect(() => {
-    if (!playing) return;
-    const interval = setInterval(() => setFrame(frameRef.current), 33);
-    return () => clearInterval(interval);
-  }, [playing]);
-
   const handleFrameChange = useCallback((e: React.ChangeEvent<HTMLInputElement>) => {
     const f = parseInt(e.target.value);
     setFrame(f);
     frameRef.current = f;
   }, []);
 
-  const jointAngles = useMemo(() => {
-    if (totalFrames === 0) return {};
+  // Compute joint values for current frame
+  const jointValues = useMemo(() => {
+    if (totalFrames === 0 || urdfJointNames.length === 0) return {};
     const row = flatChartData[Math.min(frame, totalFrames - 1)];
     const rawValues: number[] = [];
-    const jointNames: string[] = [];
+    const names: string[] = [];
 
-    for (const joint of SO101_JOINTS) {
-      const col = mapping[joint.name];
+    for (const jn of urdfJointNames) {
+      const col = mapping[jn];
       if (col && typeof row[col] === "number") {
         rawValues.push(row[col]);
-        jointNames.push(joint.name);
+        names.push(jn);
       }
     }
 
     const converted = detectAndConvert(rawValues);
-    const angles: Record<string, number> = {};
-    jointNames.forEach((name, i) => {
-      angles[name] = converted[i];
-    });
-    return angles;
-  }, [flatChartData, frame, mapping, totalFrames]);
+    const values: Record<string, number> = {};
+    names.forEach((n, i) => { values[n] = converted[i]; });
+    return values;
+  }, [flatChartData, frame, mapping, totalFrames, urdfJointNames]);
 
   const currentTime = totalFrames > 0 ? (frame / fps).toFixed(2) : "0.00";
   const totalTime = (totalFrames / fps).toFixed(2);
@@ -236,66 +262,46 @@ export default function URDFViewer({ data }: { data: EpisodeData }) {
     <div className="flex-1 flex flex-col overflow-hidden">
       {/* 3D Viewport */}
       <div className="flex-1 min-h-0 bg-slate-950 rounded-lg overflow-hidden border border-slate-700">
-        <Canvas camera={{ position: [0.35, 0.25, 0.3], fov: 45, near: 0.001, far: 10 }}>
+        <Canvas camera={{ position: [0.3 * SCALE, 0.25 * SCALE, 0.3 * SCALE], fov: 45, near: 0.01, far: 100 }}>
           <ambientLight intensity={0.5} />
-          <directionalLight position={[3, 5, 4]} intensity={1.2} castShadow />
+          <directionalLight position={[3, 5, 4]} intensity={1.2} />
           <directionalLight position={[-2, 3, -2]} intensity={0.4} />
           <hemisphereLight args={["#b1e1ff", "#444444", 0.4]} />
-          <Suspense fallback={null}>
-            <RobotArm angles={jointAngles} />
-          </Suspense>
+          <RobotScene urdfUrl={urdfUrl} jointValues={jointValues} onJointsLoaded={onJointsLoaded} />
           <Grid
-            args={[1, 1]}
-            cellSize={0.02}
+            args={[10, 10]}
+            cellSize={0.2}
             cellThickness={0.5}
             cellColor="#334155"
-            sectionSize={0.1}
+            sectionSize={1}
             sectionThickness={1}
             sectionColor="#475569"
-            fadeDistance={1}
+            fadeDistance={10}
             position={[0, 0, 0]}
           />
-          <OrbitControls target={[0, 0.1, 0]} />
-          <PlaybackDriver playing={playing} fps={fps} totalFrames={totalFrames} frameRef={frameRef} />
+          <OrbitControls target={[0, 0.8, 0]} />
+          <PlaybackDriver playing={playing} fps={fps} totalFrames={totalFrames} frameRef={frameRef} setFrame={setFrame} />
         </Canvas>
       </div>
 
-      {/* Controls Panel */}
+      {/* Controls */}
       <div className="bg-slate-800/90 border-t border-slate-700 p-3 space-y-3 shrink-0">
-        {/* Playback bar */}
+        {/* Timeline */}
         <div className="flex items-center gap-3">
           <button
-            onClick={() => {
-              setPlaying(!playing);
-              if (!playing) frameRef.current = frame;
-            }}
+            onClick={() => { setPlaying(!playing); if (!playing) frameRef.current = frame; }}
             className="w-8 h-8 flex items-center justify-center rounded bg-orange-600 hover:bg-orange-500 text-white transition-colors shrink-0"
           >
             {playing ? (
-              <svg width="12" height="14" viewBox="0 0 12 14">
-                <rect x="1" y="1" width="3" height="12" fill="white" />
-                <rect x="8" y="1" width="3" height="12" fill="white" />
-              </svg>
+              <svg width="12" height="14" viewBox="0 0 12 14"><rect x="1" y="1" width="3" height="12" fill="white" /><rect x="8" y="1" width="3" height="12" fill="white" /></svg>
             ) : (
-              <svg width="12" height="14" viewBox="0 0 12 14">
-                <polygon points="2,1 11,7 2,13" fill="white" />
-              </svg>
+              <svg width="12" height="14" viewBox="0 0 12 14"><polygon points="2,1 11,7 2,13" fill="white" /></svg>
             )}
           </button>
-          <input
-            type="range"
-            min={0}
-            max={Math.max(totalFrames - 1, 0)}
-            value={frame}
-            onChange={handleFrameChange}
-            className="flex-1 h-1.5 accent-orange-500 cursor-pointer"
-          />
-          <span className="text-xs text-slate-400 tabular-nums w-28 text-right shrink-0">
-            {currentTime}s / {totalTime}s
-          </span>
-          <span className="text-xs text-slate-500 tabular-nums w-20 text-right shrink-0">
-            F {frame}/{totalFrames - 1}
-          </span>
+          <input type="range" min={0} max={Math.max(totalFrames - 1, 0)} value={frame}
+            onChange={handleFrameChange} className="flex-1 h-1.5 accent-orange-500 cursor-pointer" />
+          <span className="text-xs text-slate-400 tabular-nums w-28 text-right shrink-0">{currentTime}s / {totalTime}s</span>
+          <span className="text-xs text-slate-500 tabular-nums w-20 text-right shrink-0">F {frame}/{totalFrames - 1}</span>
         </div>
 
         {/* Data source + joint mapping */}
@@ -304,17 +310,10 @@ export default function URDFViewer({ data }: { data: EpisodeData }) {
             <label className="text-xs text-slate-400">Data source</label>
             <div className="flex gap-1 flex-wrap">
               {groupNames.map((name) => (
-                <button
-                  key={name}
-                  onClick={() => setSelectedGroup(name)}
+                <button key={name} onClick={() => setSelectedGroup(name)}
                   className={`px-2 py-1 text-xs rounded transition-colors ${
-                    selectedGroup === name
-                      ? "bg-orange-600 text-white"
-                      : "bg-slate-700 text-slate-300 hover:bg-slate-600"
-                  }`}
-                >
-                  {name}
-                </button>
+                    selectedGroup === name ? "bg-orange-600 text-white" : "bg-slate-700 text-slate-300 hover:bg-slate-600"
+                  }`}>{name}</button>
               ))}
             </div>
           </div>
@@ -330,29 +329,23 @@ export default function URDFViewer({ data }: { data: EpisodeData }) {
                 </tr>
               </thead>
               <tbody>
-                {SO101_JOINTS.map((joint) => (
-                  <tr key={joint.name} className="border-t border-slate-700/50">
-                    <td className="px-1 py-0.5 text-slate-300 font-mono">{joint.name}</td>
+                {urdfJointNames.map((jointName) => (
+                  <tr key={jointName} className="border-t border-slate-700/50">
+                    <td className="px-1 py-0.5 text-slate-300 font-mono">{jointName}</td>
                     <td className="px-1 text-slate-600">→</td>
                     <td className="px-1 py-0.5">
-                      <select
-                        value={mapping[joint.name] ?? ""}
-                        onChange={(e) => setMapping((m) => ({ ...m, [joint.name]: e.target.value }))}
-                        className="bg-slate-900 text-slate-200 text-xs rounded px-1 py-0.5 border border-slate-600 w-full max-w-[200px]"
-                      >
+                      <select value={mapping[jointName] ?? ""}
+                        onChange={(e) => setMapping((m) => ({ ...m, [jointName]: e.target.value }))}
+                        className="bg-slate-900 text-slate-200 text-xs rounded px-1 py-0.5 border border-slate-600 w-full max-w-[200px]">
                         <option value="">-- unmapped --</option>
                         {selectedColumns.map((col) => {
                           const label = col.split(SERIES_DELIM).pop() ?? col;
-                          return (
-                            <option key={col} value={col}>
-                              {label}
-                            </option>
-                          );
+                          return <option key={col} value={col}>{label}</option>;
                         })}
                       </select>
                     </td>
                     <td className="px-1 py-0.5 text-right tabular-nums text-slate-400 font-mono">
-                      {jointAngles[joint.name] !== undefined ? jointAngles[joint.name].toFixed(3) : "—"}
+                      {jointValues[jointName] !== undefined ? jointValues[jointName].toFixed(3) : "—"}
                     </td>
                   </tr>
                 ))}

src/lib/so101-robot.ts

@@ -1,154 +1,5 @@
-export type JointDef = {
-  name: string;
-  origin: { xyz: [number, number, number]; rpy: [number, number, number] };
-  axis: [number, number, number];
-  limits: [number, number];
-};
-
-export type MeshDef = {
-  file: string;
-  origin: { xyz: [number, number, number]; rpy: [number, number, number] };
-  material: "3d_printed" | "motor";
-};
-
-export type LinkDef = {
-  name: string;
-  meshes: MeshDef[];
-};
-
-const ASSET_BASE = "/urdf/so101/assets";
-const P = Math.PI;
-
-// ─── Visual meshes per link (from URDF) ───
-export const SO101_LINKS: LinkDef[] = [
-  {
-    name: "base_link",
-    meshes: [
-      { file: `${ASSET_BASE}/base_motor_holder_so101_v1.stl`, origin: { xyz: [-0.00636471, -9.94414e-05, -0.0024], rpy: [P / 2, 0, P / 2] }, material: "3d_printed" },
-      { file: `${ASSET_BASE}/base_so101_v2.stl`, origin: { xyz: [-0.00636471, 0, -0.0024], rpy: [P / 2, 0, P / 2] }, material: "3d_printed" },
-      { file: `${ASSET_BASE}/sts3215_03a_v1.stl`, origin: { xyz: [0.0263353, 0, 0.0437], rpy: [0, 0, 0] }, material: "motor" },
-      { file: `${ASSET_BASE}/waveshare_mounting_plate_so101_v2.stl`, origin: { xyz: [-0.0309827, -0.000199441, 0.0474], rpy: [P / 2, 0, P / 2] }, material: "3d_printed" },
-    ],
-  },
-  {
-    name: "shoulder_link",
-    meshes: [
-      { file: `${ASSET_BASE}/sts3215_03a_v1.stl`, origin: { xyz: [-0.0303992, 0.000422241, -0.0417], rpy: [P / 2, P / 2, 0] }, material: "motor" },
-      { file: `${ASSET_BASE}/motor_holder_so101_base_v1.stl`, origin: { xyz: [-0.0675992, -0.000177759, 0.0158499], rpy: [P / 2, -P / 2, 0] }, material: "3d_printed" },
-      { file: `${ASSET_BASE}/rotation_pitch_so101_v1.stl`, origin: { xyz: [0.0122008, 2.22413e-05, 0.0464], rpy: [-P / 2, 0, 0] }, material: "3d_printed" },
-    ],
-  },
-  {
-    name: "upper_arm_link",
-    meshes: [
-      { file: `${ASSET_BASE}/sts3215_03a_v1.stl`, origin: { xyz: [-0.11257, -0.0155, 0.0187], rpy: [-P, 0, -P / 2] }, material: "motor" },
-      { file: `${ASSET_BASE}/upper_arm_so101_v1.stl`, origin: { xyz: [-0.065085, 0.012, 0.0182], rpy: [P, 0, 0] }, material: "3d_printed" },
-    ],
-  },
-  {
-    name: "lower_arm_link",
-    meshes: [
-      { file: `${ASSET_BASE}/under_arm_so101_v1.stl`, origin: { xyz: [-0.0648499, -0.032, 0.0182], rpy: [P, 0, 0] }, material: "3d_printed" },
-      { file: `${ASSET_BASE}/motor_holder_so101_wrist_v1.stl`, origin: { xyz: [-0.0648499, -0.032, 0.018], rpy: [-P, 0, 0] }, material: "3d_printed" },
-      { file: `${ASSET_BASE}/sts3215_03a_v1.stl`, origin: { xyz: [-0.1224, 0.0052, 0.0187], rpy: [-P, 0, -P] }, material: "motor" },
-    ],
-  },
-  {
-    name: "wrist_link",
-    meshes: [
-      { file: `${ASSET_BASE}/sts3215_03a_no_horn_v1.stl`, origin: { xyz: [0, -0.0424, 0.0306], rpy: [P / 2, P / 2, 0] }, material: "motor" },
-      { file: `${ASSET_BASE}/wrist_roll_pitch_so101_v2.stl`, origin: { xyz: [0, -0.028, 0.0181], rpy: [-P / 2, -P / 2, 0] }, material: "3d_printed" },
-    ],
-  },
-  {
-    name: "gripper_link",
-    meshes: [
-      { file: `${ASSET_BASE}/sts3215_03a_v1.stl`, origin: { xyz: [0.0077, 0.0001, -0.0234], rpy: [-P / 2, 0, 0] }, material: "motor" },
-      { file: `${ASSET_BASE}/wrist_roll_follower_so101_v1.stl`, origin: { xyz: [0, -0.000218214, 0.000949706], rpy: [-P, 0, 0] }, material: "3d_printed" },
-    ],
-  },
-  {
-    name: "moving_jaw_link",
-    meshes: [
-      { file: `${ASSET_BASE}/moving_jaw_so101_v1.stl`, origin: { xyz: [0, 0, 0.0189], rpy: [0, 0, 0] }, material: "3d_printed" },
-    ],
-  },
-];
-
-// Kinematic chain: each joint connects a parent link to a child link
-// Index in SO101_LINKS: base=0, shoulder=1, upper_arm=2, lower_arm=3, wrist=4, gripper=5, jaw=6
-export const SO101_JOINTS: JointDef[] = [
-  {
-    name: "shoulder_pan",
-    origin: { xyz: [0.0388353, -8.97657e-09, 0.0624], rpy: [P, 4.18253e-17, -P] },
-    axis: [0, 0, 1],
-    limits: [-1.91986, 1.91986],
-  },
-  {
-    name: "shoulder_lift",
-    origin: { xyz: [-0.0303992, -0.0182778, -0.0542], rpy: [-P / 2, -P / 2, 0] },
-    axis: [0, 0, 1],
-    limits: [-1.74533, 1.74533],
-  },
-  {
-    name: "elbow_flex",
-    origin: { xyz: [-0.11257, -0.028, 1.73763e-16], rpy: [0, 0, P / 2] },
-    axis: [0, 0, 1],
-    limits: [-1.69, 1.69],
-  },
-  {
-    name: "wrist_flex",
-    origin: { xyz: [-0.1349, 0.0052, 0], rpy: [0, 0, -P / 2] },
-    axis: [0, 0, 1],
-    limits: [-1.65806, 1.65806],
-  },
-  {
-    name: "wrist_roll",
-    origin: { xyz: [0, -0.0611, 0.0181], rpy: [P / 2, 0.0486795, P] },
-    axis: [0, 0, 1],
-    limits: [-2.74385, 2.84121],
-  },
-  {
-    name: "gripper",
-    origin: { xyz: [0.0202, 0.0188, -0.0234], rpy: [P / 2, 0, 0] },
-    axis: [0, 0, 1],
-    limits: [-0.174533, 1.74533],
-  },
-];
-
-export const MATERIAL_COLORS = {
-  "3d_printed": "#FFD700",
-  motor: "#1a1a1a",
-} as const;
-
 export function isSO101Robot(robotType: string | null): boolean {
   if (!robotType) return false;
   const lower = robotType.toLowerCase();
   return lower.includes("so100") || lower.includes("so101") || lower === "so_follower";
 }
-
-// Collect all unique STL file paths for preloading
-export function getAllSTLPaths(): string[] {
-  const paths = new Set<string>();
-  for (const link of SO101_LINKS) {
-    for (const mesh of link.meshes) {
-      paths.add(mesh.file);
-    }
-  }
-  return [...paths];
-}
-
-// Auto-match dataset columns to URDF joint names
-export function autoMatchJoints(columnKeys: string[]): Record<string, string> {
-  const mapping: Record<string, string> = {};
-  for (const joint of SO101_JOINTS) {
-    const exactMatch = columnKeys.find((k) => {
-      const suffix = k.split(" | ").pop()?.trim() ?? k;
-      return suffix === joint.name;
-    });
-    if (exactMatch) { mapping[joint.name] = exactMatch; continue; }
-    const fuzzy = columnKeys.find((k) => k.toLowerCase().includes(joint.name));
-    if (fuzzy) mapping[joint.name] = fuzzy;
-  }
-  return mapping;
-}