Update app.py

app.py

@@ -1,32 +1,33 @@
 """
-Standalone Hugging Face Space viewer for HDF5 trajectory datasets.
+Standalone Hugging Face Space viewer for TrajectoryBuffer-style HDF5 files.
 
-This version does NOT require your local TrajectoryBuffer class.
-It reads the HDF5 file directly with h5py.
-
-Files needed in the Space:
-  app.py
-  requirements.txt
+Best-practice version:
+- No dependency on your local TrajectoryBuffer Python class.
+- Dataset preset + custom dataset support.
+- Robust HDF5 schema detection for root-level episode_XXXX groups.
+- Auto-detect image keys from each trajectory's observation group.
+- Avoids fragile multiline f-strings in UI status text.
+- Uses slider.release() for timestep rendering to reduce image flicker.
 
 requirements.txt:
-  gradio
-  huggingface_hub
-  h5py
-  numpy
-  pillow
-  matplotlib
+    gradio
+    huggingface_hub
+    h5py
+    numpy
+    pillow
+    matplotlib
 
 Optional:
-  opencv-python-headless
+    opencv-python-headless
 """
 
-import os
 import re
 from functools import lru_cache
 
 import gradio as gr
 import h5py
 import matplotlib
+
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 import numpy as np
@@ -40,8 +41,7 @@ except Exception:
 
 
 # -----------------------------------------------------------------------------
-# Dataset presets.
-# The same Space can visualize multiple HDF5 files by changing repo_id + filename.
+# Dataset presets
 # -----------------------------------------------------------------------------
 DATASET_PRESETS = {
     "Robosuite Square 20260409": {
@@ -67,36 +67,35 @@ DATASET_PRESETS = {
 DEFAULT_PRESET = "Robosuite Square 20260409"
 REPO_TYPE = "dataset"
 DEFAULT_CHUNK_LEN = 16
+
 PREFERRED_IMAGE_KEYS = [
     "image1",
     "image2",
     "agentview_image",
     "robot0_eye_in_hand_image",
+    "front_image",
+    "wrist_image",
 ]
 
+IMAGE_KEY_HINTS = ["rgb", "image", "img", "camera", "cam"]
+
 
 # -----------------------------------------------------------------------------
-# HDF5 helpers
+# Dataset resolution and cache helpers
 # -----------------------------------------------------------------------------
-def _clear_dataset_caches():
-    get_local_hdf5_path.cache_clear()
-    get_trajectory_keys.cache_clear()
-    get_num_trajectories.cache_clear()
-    load_traj.cache_clear()
-
-
 def resolve_dataset(preset_name, custom_repo_id=None, custom_filename=None):
-    """Return (repo_id, filename) from a preset or custom fields."""
     preset_name = preset_name or DEFAULT_PRESET
+
     if preset_name == "Custom":
         repo_id = str(custom_repo_id or "").strip()
         filename = str(custom_filename or "").strip()
         if not repo_id or not filename:
-            raise ValueError("For Custom, provide both repo_id and HDF5 filename/path.")
+            raise ValueError("For Custom mode, provide both repo_id and HDF5 filename/path.")
         return repo_id, filename
 
     if preset_name not in DATASET_PRESETS:
         preset_name = DEFAULT_PRESET
+
     item = DATASET_PRESETS[preset_name]
     return item["repo_id"], item["filename"]
 
@@ -111,75 +110,92 @@ def get_local_hdf5_path(repo_id, filename):
 
 
 def _natural_sort_key(name):
-    m = re.search(r"([0-9]+)$", str(name))
-    return (0, int(m.group(1))) if m else (1, str(name))
+    match = re.search(r"([0-9]+)$", str(name))
+    if match:
+        return 0, int(match.group(1))
+    return 1, str(name)
 
 
 @lru_cache(maxsize=8)
 def get_trajectory_keys(repo_id, filename):
-    """Detect trajectory groups in common HDF5 layouts."""
+    """Return ordered trajectory group paths."""
     path = get_local_hdf5_path(repo_id, filename)
+
     with h5py.File(path, "r") as f:
-        # Your TrajectoryBuffer saves root-level groups:
+        # Your TrajectoryBuffer format:
         #   /episode_0000
         #   /episode_0001
-        #   ...
-        # Some other robotics datasets use /data/demo_0, so keep that fallback.
         root_episode_keys = [
-            k for k in f.keys()
-            if isinstance(f[k], h5py.Group) and str(k).startswith("episode_")
+            key
+            for key in f.keys()
+            if isinstance(f[key], h5py.Group) and str(key).startswith("episode_")
         ]
         if root_episode_keys:
-            group = f
-            prefix = ""
-            group_keys = root_episode_keys
-        elif "data" in f and isinstance(f["data"], h5py.Group):
-            group = f["data"]
-            prefix = "data"
-            group_keys = [k for k in group.keys() if isinstance(group[k], h5py.Group)]
-        else:
-            group = f
-            prefix = ""
-            group_keys = [k for k in group.keys() if isinstance(group[k], h5py.Group)]
-
-    group_keys = sorted(group_keys, key=_natural_sort_key)
-    return tuple(f"{prefix}/{k}" if prefix else k for k in group_keys)
+            return tuple(sorted(root_episode_keys, key=_natural_sort_key))
+
+        # Robomimic-style fallback:
+        #   /data/demo_0
+        #   /data/demo_1
+        if "data" in f and isinstance(f["data"], h5py.Group):
+            data_group = f["data"]
+            keys = [
+                key
+                for key in data_group.keys()
+                if isinstance(data_group[key], h5py.Group)
+            ]
+            return tuple("data/" + key for key in sorted(keys, key=_natural_sort_key))
+
+        # Generic root-level group fallback.
+        keys = [
+            key
+            for key in f.keys()
+            if isinstance(f[key], h5py.Group)
+        ]
+        return tuple(sorted(keys, key=_natural_sort_key))
 
 
 @lru_cache(maxsize=8)
 def get_num_trajectories(repo_id, filename):
-    return max(len(get_trajectory_keys(repo_id, filename)), 1)
+    return len(get_trajectory_keys(repo_id, filename))
 
 
-def inspect_hdf5_tree(preset_name, custom_repo_id, custom_filename, max_lines=160):
-    """Show the HDF5 tree for debugging inside the Space."""
+def inspect_hdf5_tree(preset_name, custom_repo_id, custom_filename, max_lines=180):
     repo_id, filename = resolve_dataset(preset_name, custom_repo_id, custom_filename)
     path = get_local_hdf5_path(repo_id, filename)
+
     lines = []
     with h5py.File(path, "r") as f:
         def visitor(name, obj):
             if len(lines) >= max_lines:
                 return
             if isinstance(obj, h5py.Dataset):
-                lines.append(f"DATASET {name} shape={obj.shape} dtype={obj.dtype}")
+                lines.append(
+                    "DATASET {} shape={} dtype={}".format(name, obj.shape, obj.dtype)
+                )
             elif isinstance(obj, h5py.Group):
-                lines.append(f"GROUP   {name}")
+                lines.append("GROUP   {}".format(name))
+
         f.visititems(visitor)
 
     if len(lines) >= max_lines:
         lines.append("...")
-    return "\n".join(lines) if lines else "No HDF5 contents found."
+
+    if not lines:
+        return "No HDF5 contents found."
+    return chr(10).join(lines)
 
 
-def _read_dataset_value(ds):
-    value = ds[()]
+# -----------------------------------------------------------------------------
+# HDF5 loading helpers
+# -----------------------------------------------------------------------------
+def _read_dataset_value(dataset):
+    value = dataset[()]
     if isinstance(value, bytes):
         return value.decode("utf-8")
     return value
 
 
 def _read_group_recursive(group):
-    """Read a group into nested dictionaries of numpy arrays."""
     out = {}
     for key, obj in group.items():
         if isinstance(obj, h5py.Dataset):
@@ -189,239 +205,266 @@ def _read_group_recursive(group):
     return out
 
 
-def _find_first_existing_key(mapping, candidates):
-    for key in candidates:
+def _find_first_key(mapping, candidate_keys):
+    for key in candidate_keys:
         if key in mapping:
             return key
     return None
 
 
-def _maybe_time_slice(value, t, T):
+def _infer_time_length(data):
+    """Infer trajectory length from common TrajectoryBuffer fields."""
+    for key in ["timesteps", "dones", "robot_actions", "teacher_actions", "actions"]:
+        if key in data:
+            arr = np.asarray(data[key])
+            if arr.ndim >= 1:
+                return int(arr.shape[0])
+
+    obs_group = None
+    if isinstance(data.get("observation"), dict):
+        obs_group = data["observation"]
+    elif isinstance(data.get("obs"), dict):
+        obs_group = data["obs"]
+
+    if obs_group:
+        lengths = []
+        for value in obs_group.values():
+            arr = np.asarray(value)
+            if arr.ndim >= 1:
+                lengths.append(int(arr.shape[0]))
+        if lengths:
+            values, counts = np.unique(lengths, return_counts=True)
+            return int(values[np.argmax(counts)])
+
+    return 1
+
+
+def _slice_time(value, t, T):
     arr = np.asarray(value)
     if arr.ndim >= 1 and arr.shape[0] == T:
         return arr[t]
     return arr
 
 
-def _infer_time_length(data):
-    """Infer T from datasets whose first dimension is time."""
-    candidate_lengths = []
-
-    def collect(obj):
-        if isinstance(obj, dict):
-            for v in obj.values():
-                collect(v)
-        else:
-            arr = np.asarray(obj)
-            if arr.ndim >= 1 and arr.shape[0] > 1:
-                candidate_lengths.append(int(arr.shape[0]))
-
-    collect(data)
-    if not candidate_lengths:
-        return 1
-
-    # The trajectory length should usually be the most common large first dim.
-    values, counts = np.unique(candidate_lengths, return_counts=True)
-    return int(values[np.argmax(counts)])
-
-
 @lru_cache(maxsize=64)
 def load_traj(repo_id, filename, traj_id):
-    """Load one trajectory as a list of step dictionaries.
-
-    Output step format:
-      {
-        "timestep": int,
-        "obs": dict,
-        "teacher_action": np.ndarray,
-        "robot_action": np.ndarray,
-        "no_teacher_action": bool,
-        "no_robot_action": bool,
-      }
-    """
-    path = get_local_hdf5_path(repo_id, filename)
+    """Load one trajectory as list[dict]."""
     traj_keys = get_trajectory_keys(repo_id, filename)
     if not traj_keys:
         return []
 
     traj_id = int(np.clip(int(traj_id), 0, len(traj_keys) - 1))
     traj_key = traj_keys[traj_id]
+    path = get_local_hdf5_path(repo_id, filename)
 
     with h5py.File(path, "r") as f:
-        g = f[traj_key]
-        data = _read_group_recursive(g)
-        attrs = dict(g.attrs)
-
-    # Case A: trajectory group contains step groups: step_0, step_1, ...
-    step_group_keys = [
-        k for k, v in data.items()
-        if isinstance(v, dict) and (str(k).startswith("step") or str(k).isdigit())
-    ]
-    if step_group_keys:
-        traj = []
-        for step_key in sorted(step_group_keys, key=_natural_sort_key):
-            step = data[step_key]
-            obs = step.get("obs", {}) if isinstance(step.get("obs", {}), dict) else {}
-            teacher_action = step.get("teacher_action", step.get("teacher_actions", step.get("action", step.get("actions", np.zeros(1, dtype=np.float32)))))
-            robot_action = step.get("robot_action", step.get("robot_actions", step.get("action", step.get("actions", teacher_action))))
-            traj.append({
-                "timestep": int(step.get("timestep", len(traj))),
-                "obs": obs,
-                "teacher_action": np.asarray(teacher_action),
-                "robot_action": np.asarray(robot_action),
-                "no_teacher_action": bool(np.asarray(step.get("no_teacher_action", step.get("no_teacher_actions", False))).reshape(-1)[0]),
-                "no_robot_action": bool(np.asarray(step.get("no_robot_action", step.get("no_robot_actions", False))).reshape(-1)[0]),
-            })
-        return traj
-
-    # Case B: trajectory group contains array datasets with first dimension T.
-    # Your TrajectoryBuffer layout is:
-    #   /episode_0000/observation/<image_or_state_key>[T,...]
-    #   /episode_0000/robot_actions[T,D]
-    #   /episode_0000/teacher_actions[T,D]
-    #   /episode_0000/no_teacher_actions[T]
-    #   /episode_0000/no_robot_actions[T]
-    #
-    # Keep obs/action aliases for compatibility with other layouts.
+        group = f[traj_key]
+        data = _read_group_recursive(group)
+
     T = _infer_time_length(data)
-    obs_all = {}
-    if isinstance(data.get("observation", {}), dict):
-        obs_all = data.get("observation", {})
-    elif isinstance(data.get("obs", {}), dict):
-        obs_all = data.get("obs", {})
-
-    action_key = _find_first_existing_key(data, ["actions", "action"])
-    teacher_key = _find_first_existing_key(data, ["teacher_actions", "teacher_action"])
-    robot_key = _find_first_existing_key(data, ["robot_actions", "robot_action"])
-    no_teacher_key = _find_first_existing_key(data, ["no_teacher_actions", "no_teacher_action"])
-    no_robot_key = _find_first_existing_key(data, ["no_robot_actions", "no_robot_action"])
+
+    if isinstance(data.get("observation"), dict):
+        obs_all = data["observation"]
+    elif isinstance(data.get("obs"), dict):
+        obs_all = data["obs"]
+    else:
+        obs_all = {}
+
+    action_key = _find_first_key(data, ["actions", "action"])
+    teacher_key = _find_first_key(data, ["teacher_actions", "teacher_action"])
+    robot_key = _find_first_key(data, ["robot_actions", "robot_action"])
+    no_teacher_key = _find_first_key(data, ["no_teacher_actions", "no_teacher_action"])
+    no_robot_key = _find_first_key(data, ["no_robot_actions", "no_robot_action"])
+    done_key = _find_first_key(data, ["dones", "done"])
+    timestep_key = _find_first_key(data, ["timesteps", "timestep"])
+    success_key = _find_first_key(data, ["if_success", "success", "successes"])
 
     traj = []
     for t in range(T):
         obs_t = {}
         for key, value in obs_all.items():
-            obs_t[key] = _maybe_time_slice(value, t, T)
+            obs_t[key] = _slice_time(value, t, T)
 
         default_action = np.zeros(1, dtype=np.float32)
         if action_key is not None:
-            default_action = _maybe_time_slice(data[action_key], t, T)
-
-        teacher_action = _maybe_time_slice(data[teacher_key], t, T) if teacher_key is not None else default_action
-        robot_action = _maybe_time_slice(data[robot_key], t, T) if robot_key is not None else default_action
-        no_teacher = _maybe_time_slice(data[no_teacher_key], t, T) if no_teacher_key is not None else False
-        no_robot = _maybe_time_slice(data[no_robot_key], t, T) if no_robot_key is not None else False
-
-        traj.append({
-            "timestep": t,
-            "obs": obs_t,
-            "teacher_action": np.asarray(teacher_action),
-            "robot_action": np.asarray(robot_action),
-            "no_teacher_action": bool(np.asarray(no_teacher).reshape(-1)[0]),
-            "no_robot_action": bool(np.asarray(no_robot).reshape(-1)[0]),
-        })
+            default_action = _slice_time(data[action_key], t, T)
+
+        teacher_action = default_action
+        if teacher_key is not None:
+            teacher_action = _slice_time(data[teacher_key], t, T)
+
+        robot_action = default_action
+        if robot_key is not None:
+            robot_action = _slice_time(data[robot_key], t, T)
+
+        no_teacher = False
+        if no_teacher_key is not None:
+            no_teacher = _slice_time(data[no_teacher_key], t, T)
+
+        no_robot = False
+        if no_robot_key is not None:
+            no_robot = _slice_time(data[no_robot_key], t, T)
+
+        done = False
+        if done_key is not None:
+            done = _slice_time(data[done_key], t, T)
+
+        timestep = t
+        if timestep_key is not None:
+            timestep_arr = _slice_time(data[timestep_key], t, T)
+            timestep = int(np.asarray(timestep_arr).reshape(-1)[0])
+
+        if_success = False
+        if success_key is not None:
+            if_success = _slice_time(data[success_key], t, T)
+
+        traj.append(
+            {
+                "obs": obs_t,
+                "robot_action": np.asarray(robot_action),
+                "teacher_action": np.asarray(teacher_action),
+                "done": bool(np.asarray(done).reshape(-1)[0]),
+                "timestep": timestep,
+                "no_robot_action": bool(np.asarray(no_robot).reshape(-1)[0]),
+                "no_teacher_action": bool(np.asarray(no_teacher).reshape(-1)[0]),
+                "episode_id": traj_key,
+                "if_success": bool(np.asarray(if_success).reshape(-1)[0]),
+            }
+        )
 
     return traj
 
 
 # -----------------------------------------------------------------------------
-# Visualization helpers
+# Image and plotting helpers
 # -----------------------------------------------------------------------------
 def _extract_latest_obs_value(value):
     arr = np.asarray(value)
-    # Your local script handled stacked recent observations by taking the latest.
-    if arr.ndim >= 1 and arr.shape[0] in (1, 2):
+    # Per-timestep stacked observation commonly has shape [obs_T, C, H, W].
+    if arr.ndim >= 1 and arr.shape[0] in (1, 2, 3, 4):
         return arr[-1]
     return arr
 
 
-def _extract_display_image(value, reverse_channels=False, output_uint8=True):
-    img = _extract_latest_obs_value(value)
-    img = np.asarray(img)
+def _looks_like_image_array(key, value):
+    arr = np.asarray(_extract_latest_obs_value(value))
+    key_l = str(key).lower()
+    key_hint = any(hint in key_l for hint in IMAGE_KEY_HINTS)
 
-    # Your saved image shape can be [obs_T, C, H, W] per timestep.
-    # Take the latest stacked observation, then convert CHW -> HWC.
-    if img.ndim == 4 and img.shape[0] in (1, 2, 3, 4):
-        img = img[-1]
+    shape_hint = False
+    if arr.ndim == 2:
+        shape_hint = True
+    elif arr.ndim == 3:
+        shape_hint = arr.shape[-1] in (1, 3, 4) or arr.shape[0] in (1, 3, 4)
+
+    return key_hint or shape_hint
+
+
+def _float_img_to_uint8(img):
+    arr = img.astype(np.float32)
+    arr_min = float(np.nanmin(arr))
+    arr_max = float(np.nanmax(arr))
+
+    # TrajectoryBuffer saves float images originally in [-1, 1] as uint8.
+    # But for compatibility, handle float [-1, 1], [0, 1], and [0, 255].
+    if arr_min >= -1.01 and arr_max <= 1.01:
+        if arr_min < 0.0:
+            arr = (arr + 1.0) * 0.5
+        arr = np.clip(arr, 0.0, 1.0) * 255.0
+    elif arr_max <= 255.0:
+        arr = np.clip(arr, 0.0, 255.0)
+    else:
+        arr = 255.0 * (arr - arr_min) / max(arr_max - arr_min, 1e-8)
+
+    return np.round(arr).astype(np.uint8)
+
+
+def _extract_display_image(value, reverse_channels=False):
+    img = np.asarray(_extract_latest_obs_value(value))
 
     if img.ndim == 2:
         img = np.repeat(img[..., None], 3, axis=-1)
     elif img.ndim == 3 and img.shape[0] in (1, 3, 4):
         img = np.transpose(img, (1, 2, 0))
 
-    if img.ndim == 3 and img.shape[-1] == 4:
+    if img.ndim == 3 and img.shape[-1] == 1:
+        img = np.repeat(img, 3, axis=-1)
+    elif img.ndim == 3 and img.shape[-1] == 4:
         img = img[..., :3]
 
     if img.ndim != 3:
-        raise ValueError(f"Unsupported image shape: {img.shape}")
+        raise ValueError("Unsupported image shape: {}".format(img.shape))
 
     if img.dtype == np.uint8:
-        img_rgb = img.copy()
+        out = img.copy()
     else:
-        img_rgb = img.astype(np.float32)
-        if np.nanmin(img_rgb) < 0:
-            img_rgb = (img_rgb + 1.0) / 2.0
-        if np.nanmax(img_rgb) > 1.5:
-            img_rgb = img_rgb / 255.0
-        img_rgb = np.clip(img_rgb, 0.0, 1.0)
-        if output_uint8:
-            img_rgb = np.round(img_rgb * 255.0).astype(np.uint8)
+        out = _float_img_to_uint8(img)
 
-    if reverse_channels and img_rgb.shape[-1] == 3:
-        img_rgb = img_rgb[..., ::-1]
+    # Browser display expects RGB. Your current data appears RGB already,
+    # so default reverse_channels=False.
+    if reverse_channels and out.shape[-1] == 3:
+        out = out[..., ::-1]
 
-    return img_rgb
+    return out
 
 
-def _resize_image_for_display(img, display_scale=4):
-    if display_scale is None or float(display_scale) == 1.0:
+def _resize_image_for_display(img, display_scale):
+    scale = float(display_scale)
+    if scale == 1.0:
         return img
 
-    display_scale = float(display_scale)
     h, w = img.shape[:2]
-    new_w = max(1, int(round(w * display_scale)))
-    new_h = max(1, int(round(h * display_scale)))
+    new_size = (max(1, int(round(w * scale))), max(1, int(round(h * scale))))
 
     if cv2 is not None:
-        return cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_NEAREST)
+        return cv2.resize(img, new_size, interpolation=cv2.INTER_NEAREST)
 
     pil_img = Image.fromarray(img)
-    return np.asarray(pil_img.resize((new_w, new_h), resample=Image.Resampling.NEAREST))
+    return np.asarray(pil_img.resize(new_size, resample=Image.Resampling.NEAREST))
 
 
-def _extract_mixed_action_chunk(traj, start_idx, chunk_length=16):
+def _extract_mixed_action_chunk(traj, start_idx, chunk_length):
     chunk = []
     sources = []
     end_idx = min(len(traj), int(start_idx) + int(chunk_length))
+
     for idx in range(int(start_idx), end_idx):
         step = traj[idx]
         use_teacher = not bool(step.get("no_teacher_action", False))
         action = step["teacher_action"] if use_teacher else step["robot_action"]
         chunk.append(np.asarray(action, dtype=np.float32).reshape(-1))
         sources.append("T" if use_teacher else "R")
+
     if not chunk:
         return None, ""
+
     return np.stack(chunk, axis=0), "".join(sources)
 
 
-def _extract_robot_action_chunk(traj, start_idx, chunk_length=16):
+def _extract_robot_action_chunk(traj, start_idx, chunk_length):
     chunk = []
     end_idx = min(len(traj), int(start_idx) + int(chunk_length))
+
     for idx in range(int(start_idx), end_idx):
         step = traj[idx]
         chunk.append(np.asarray(step["robot_action"], dtype=np.float32).reshape(-1))
+
     if not chunk:
         return None
+
     return np.stack(chunk, axis=0)
 
 
-def _safe_array_str(x, precision=3, max_items=24):
-    arr = np.asarray(x).reshape(-1)
+def _safe_array_str(value, precision=3, max_items=24):
+    arr = np.asarray(value).reshape(-1)
     shown = arr[:max_items]
-    suffix = "" if arr.size <= max_items else f" ... +{arr.size - max_items} more"
-    return np.array2string(shown, precision=precision, separator=", ") + suffix
+    text = np.array2string(shown, precision=precision, separator=", ")
+    if arr.size > max_items:
+        text += " ... +{} more".format(arr.size - max_items)
+    return text
 
 
-def _make_action_chunk_plot(mixed_chunk, robot_chunk=None):
+def _make_action_chunk_plot(mixed_chunk, robot_chunk):
     if mixed_chunk is None:
         return None
 
@@ -430,18 +473,24 @@ def _make_action_chunk_plot(mixed_chunk, robot_chunk=None):
         mixed_chunk = mixed_chunk[:, None]
 
     fig, ax = plt.subplots(figsize=(7, 3.2), dpi=140)
-    t = np.arange(mixed_chunk.shape[0])
+    x = np.arange(mixed_chunk.shape[0])
     max_dims = min(mixed_chunk.shape[1], 10)
 
-    for d in range(max_dims):
-        ax.plot(t, mixed_chunk[:, d], label=f"mixed[{d}]")
+    for dim in range(max_dims):
+        ax.plot(x, mixed_chunk[:, dim], label="mixed[{}]".format(dim))
 
     if robot_chunk is not None:
         robot_chunk = np.asarray(robot_chunk, dtype=np.float32)
         if robot_chunk.ndim == 1:
             robot_chunk = robot_chunk[:, None]
-        for d in range(min(robot_chunk.shape[1], max_dims)):
-            ax.plot(t, robot_chunk[:, d], linestyle="--", alpha=0.55, label=f"robot[{d}]")
+        for dim in range(min(robot_chunk.shape[1], max_dims)):
+            ax.plot(
+                x,
+                robot_chunk[:, dim],
+                linestyle="--",
+                alpha=0.55,
+                label="robot[{}]".format(dim),
+            )
 
     ax.set_title("Action chunk")
     ax.set_xlabel("chunk step")
@@ -451,144 +500,193 @@ def _make_action_chunk_plot(mixed_chunk, robot_chunk=None):
     fig.tight_layout()
     fig.canvas.draw()
     rgba = np.asarray(fig.canvas.buffer_rgba())
-    img = rgba[..., :3].copy()
+    image = rgba[..., :3].copy()
     plt.close(fig)
-    return img
+    return image
 
 
+# -----------------------------------------------------------------------------
+# Gradio callbacks
+# -----------------------------------------------------------------------------
 def get_available_image_keys(repo_id, filename, traj_id):
     n_traj = get_num_trajectories(repo_id, filename)
-    traj_id = int(np.clip(int(traj_id), 0, max(n_traj - 1, 0)))
+    if n_traj == 0:
+        return []
+
+    traj_id = int(np.clip(int(traj_id), 0, n_traj - 1))
     traj = load_traj(repo_id, filename, traj_id)
     if not traj:
         return []
 
     obs = traj[0].get("obs", {})
-    image_keys = []
+    keys = []
     for key, value in obs.items():
         try:
-            arr = np.asarray(_extract_latest_obs_value(value))
-            key_l = str(key).lower()
-            key_hint = any(s in key_l for s in ["rgb", "image", "img", "camera", "cam"])
-            looks_like_shape = (
-                arr.ndim == 2
-                or (arr.ndim == 3 and (arr.shape[-1] in (1, 3, 4) or arr.shape[0] in (1, 3, 4)))
-                or (arr.ndim == 4 and arr.shape[0] in (1, 2, 3, 4) and arr.shape[1] in (1, 3, 4))
-            )
-            if key_hint or looks_like_shape:
-                image_keys.append(key)
+            if _looks_like_image_array(key, value):
+                keys.append(key)
         except Exception:
             pass
 
-    ordered = [k for k in PREFERRED_IMAGE_KEYS if k in image_keys]
-    ordered += [k for k in image_keys if k not in ordered]
+    ordered = [key for key in PREFERRED_IMAGE_KEYS if key in keys]
+    ordered.extend([key for key in keys if key not in ordered])
     return ordered
 
 
-# -----------------------------------------------------------------------------
-# Gradio callbacks
-# -----------------------------------------------------------------------------
+def update_custom_visibility(preset_name):
+    visible = preset_name == "Custom"
+    return gr.update(visible=visible), gr.update(visible=visible)
+
+
+def update_after_dataset_change(preset_name, custom_repo_id, custom_filename):
+    repo_id, filename = resolve_dataset(preset_name, custom_repo_id, custom_filename)
+    n_traj = get_num_trajectories(repo_id, filename)
+
+    if n_traj == 0:
+        status = "Loaded `{}` / `{}`".format(repo_id, filename)
+        status = status + chr(10) + "Detected trajectories: 0"
+        return (
+            gr.update(maximum=1, value=0),
+            gr.update(maximum=1, value=0),
+            gr.update(choices=[], value=[]),
+            status,
+        )
+
+    keys = get_available_image_keys(repo_id, filename, 0)
+    traj = load_traj(repo_id, filename, 0)
+
+    status = "Loaded `{}` / `{}`".format(repo_id, filename)
+    status = status + chr(10) + "Detected trajectories: {}".format(n_traj)
+
+    return (
+        gr.update(maximum=max(n_traj - 1, 1), value=0),
+        gr.update(maximum=max(len(traj) - 1, 1), value=0),
+        gr.update(choices=keys, value=keys[:2]),
+        status,
+    )
+
+
 def update_after_traj_change(preset_name, custom_repo_id, custom_filename, traj_id):
     repo_id, filename = resolve_dataset(preset_name, custom_repo_id, custom_filename)
     n_traj = get_num_trajectories(repo_id, filename)
-    traj_id = int(np.clip(int(traj_id), 0, max(n_traj - 1, 0)))
+    if n_traj == 0:
+        return gr.update(maximum=1, value=0), gr.update(choices=[], value=[])
+
+    traj_id = int(np.clip(int(traj_id), 0, n_traj - 1))
     traj = load_traj(repo_id, filename, traj_id)
-    image_keys = get_available_image_keys(repo_id, filename, traj_id)
-    max_step = max(len(traj) - 1, 0)
-    slider_max = max(max_step, 1)  # Gradio requires min < max.
+    keys = get_available_image_keys(repo_id, filename, traj_id)
+
     return (
-        gr.update(maximum=slider_max, value=0),
-        gr.update(choices=image_keys, value=image_keys[:2]),
+        gr.update(maximum=max(len(traj) - 1, 1), value=0),
+        gr.update(choices=keys, value=keys[:2]),
     )
 
 
-def render_frame(preset_name, custom_repo_id, custom_filename, traj_id, timestep, image_keys, chunk_len, display_scale, reverse_channels):
+def render_frame(
+    preset_name,
+    custom_repo_id,
+    custom_filename,
+    traj_id,
+    timestep,
+    image_keys,
+    chunk_len,
+    display_scale,
+    reverse_channels,
+):
     repo_id, filename = resolve_dataset(preset_name, custom_repo_id, custom_filename)
     n_traj = get_num_trajectories(repo_id, filename)
-    traj_id = int(np.clip(int(traj_id), 0, max(n_traj - 1, 0)))
-    traj = load_traj(repo_id, filename, traj_id)
 
+    if n_traj == 0:
+        return [], None, "No trajectory groups found. Open Debug: HDF5 tree."
+
+    traj_id = int(np.clip(int(traj_id), 0, n_traj - 1))
+    traj = load_traj(repo_id, filename, traj_id)
     if not traj:
-        return [], None, "No trajectory could be loaded. Open the HDF5 debug panel to inspect the file layout."
+        return [], None, "Trajectory could not be loaded. Open Debug: HDF5 tree."
 
     timestep = int(np.clip(int(timestep), 0, len(traj) - 1))
     chunk_len = int(chunk_len)
     display_scale = float(display_scale)
-    step = traj[timestep]
-    obs = step.get("obs", {})
 
     if image_keys is None:
         image_keys = []
     if isinstance(image_keys, str):
         image_keys = [image_keys]
 
-    gallery = []
-    errors = []
+    step = traj[timestep]
+    obs = step.get("obs", {})
+
+    gallery_items = []
+    warnings = []
     for key in image_keys:
         if key not in obs:
-            errors.append(f"Missing image key: {key}")
+            warnings.append("Missing image key: {}".format(key))
             continue
         try:
-            img = _extract_display_image(obs[key], reverse_channels=bool(reverse_channels), output_uint8=True)
-            img = _resize_image_for_display(img, display_scale=display_scale)
-            gallery.append((img, key))
+            img = _extract_display_image(
+                obs[key],
+                reverse_channels=bool(reverse_channels),
+            )
+            img = _resize_image_for_display(img, display_scale)
+            gallery_items.append((img, key))
         except Exception as exc:
-            errors.append(f"{key}: {exc}")
+            warnings.append("{}: {}".format(key, exc))
 
-    mixed_chunk, chunk_sources = _extract_mixed_action_chunk(traj, timestep, chunk_length=chunk_len)
-    robot_chunk = _extract_robot_action_chunk(traj, timestep, chunk_length=chunk_len)
+    mixed_chunk, source_mask = _extract_mixed_action_chunk(traj, timestep, chunk_len)
+    robot_chunk = _extract_robot_action_chunk(traj, timestep, chunk_len)
     action_plot = _make_action_chunk_plot(mixed_chunk, robot_chunk)
 
-    teacher_action = step.get("teacher_action", np.zeros(1, dtype=np.float32))
-    robot_action = step.get("robot_action", np.zeros(1, dtype=np.float32))
-    no_teacher = bool(step.get("no_teacher_action", False))
-    no_robot = bool(step.get("no_robot_action", False))
-
     info_lines = [
-        f"detected trajectories: {n_traj}",
-        f"trajectory: {traj_id}",
-        f"timestep: {timestep} / {len(traj) - 1}",
-        f"no_teacher_action: {int(no_teacher)}",
-        f"no_robot_action: {int(no_robot)}",
-        f"chunk_len: {chunk_len}",
-        f"chunk source mask: {chunk_sources}  (T=teacher, R=robot fallback)",
+        "dataset: {} / {}".format(repo_id, filename),
+        "detected trajectories: {}".format(n_traj),
+        "trajectory: {}".format(traj_id),
+        "episode_id: {}".format(step.get("episode_id", "")),
+        "timestep: {} / {}".format(timestep, len(traj) - 1),
+        "saved timestep: {}".format(step.get("timestep", timestep)),
+        "done: {}".format(int(bool(step.get("done", False)))),
+        "if_success: {}".format(int(bool(step.get("if_success", False)))),
+        "no_teacher_action: {}".format(int(bool(step.get("no_teacher_action", False)))),
+        "no_robot_action: {}".format(int(bool(step.get("no_robot_action", False)))),
+        "chunk_len: {}".format(chunk_len),
+        "chunk source mask: {}  (T=teacher, R=robot fallback)".format(source_mask),
         "",
-        f"teacher_action: {_safe_array_str(teacher_action)}",
-        f"robot_action:   {_safe_array_str(robot_action)}",
+        "teacher_action: {}".format(_safe_array_str(step.get("teacher_action", []))),
+        "robot_action:   {}".format(_safe_array_str(step.get("robot_action", []))),
     ]
 
-    if errors:
-        info_lines += ["", "Image warnings:"] + errors
+    if warnings:
+        info_lines.append("")
+        info_lines.append("Image warnings:")
+        info_lines.extend(warnings)
 
-    return gallery, action_plot, "\n".join(info_lines)
+    return gallery_items, action_plot, chr(10).join(info_lines)
 
 
 # -----------------------------------------------------------------------------
 # App
 # -----------------------------------------------------------------------------
 def build_app():
+    repo_id, filename = resolve_dataset(DEFAULT_PRESET)
+
     try:
-        repo_id, filename = resolve_dataset(DEFAULT_PRESET)
         n_traj = get_num_trajectories(repo_id, filename)
-        first_keys = get_available_image_keys(repo_id, filename, 0)
-        startup_error = None
+        first_keys = get_available_image_keys(repo_id, filename, 0) if n_traj else []
+        startup_warning = ""
     except Exception as exc:
-        n_traj = 1
+        n_traj = 0
         first_keys = []
-        startup_error = repr(exc)
+        startup_warning = repr(exc)
+
+    default_status = "Loaded default dataset" + chr(10)
+    default_status += "Detected trajectories: {}".format(n_traj)
 
     with gr.Blocks(title="HDF5 Trajectory Viewer") as demo:
         gr.Markdown(
-            "# HDF5 Trajectory Viewer\n"
-            "Standalone viewer: no local `TrajectoryBuffer` dependency.\n\n"
-            f"Default dataset detected trajectories: **{n_traj}**"
+            "# HDF5 Trajectory Viewer\n\n"
+            "Standalone viewer for TrajectoryBuffer-style HDF5 datasets on Hugging Face."
         )
 
-        if startup_error is not None:
-            gr.Markdown(
-                "⚠️ **Startup warning**\n\n"
-                f"```text\n{startup_error}\n```"
-            )
+        if startup_warning:
+            gr.Markdown("Startup warning: `{}`".format(startup_warning))
 
         with gr.Row():
             preset = gr.Dropdown(
@@ -607,6 +705,13 @@ def build_app():
                 visible=False,
             )
 
+        dataset_status = gr.Textbox(
+            label="Dataset status",
+            lines=2,
+            value=default_status,
+            interactive=False,
+        )
+
         with gr.Row():
             traj_slider = gr.Slider(
                 minimum=0,
@@ -657,37 +762,13 @@ def build_app():
             object_fit="contain",
         )
         action_plot = gr.Image(label="Action chunk plot", type="numpy")
-        info = gr.Textbox(label="Frame info", lines=13)
+        info = gr.Textbox(label="Frame info", lines=16)
 
         with gr.Accordion("Debug: HDF5 tree", open=False):
             inspect_btn = gr.Button("Inspect HDF5 structure")
-            hdf5_tree = gr.Textbox(lines=22, label="HDF5 tree")
-            inspect_btn.click(
-                fn=inspect_hdf5_tree,
-                inputs=[preset, custom_repo_id, custom_filename],
-                outputs=hdf5_tree,
-            )
-
-        def update_custom_visibility(preset_name):
-            visible = preset_name == "Custom"
-            return gr.update(visible=visible), gr.update(visible=visible)
-
-        def update_after_dataset_change(preset_name, custom_repo_id, custom_filename):
-            repo_id, filename = resolve_dataset(preset_name, custom_repo_id, custom_filename)
-            n = get_num_trajectories(repo_id, filename)
-            keys = get_available_image_keys(repo_id, filename, 0)
-            traj = load_traj(repo_id, filename, 0)
-            status_text = "Loaded `{}` / `{}`".format(repo_id, filename)
-            status_text = status_text + chr(10) + "Detected trajectories: {}".format(n)
-            return (
-                gr.update(maximum=max(n - 1, 1), value=0),
-                gr.update(maximum=max(len(traj) - 1, 1), value=0),
-                gr.update(choices=keys, value=keys[:2]),
-                status_text,
-            )
-        dataset_status = gr.Textbox(label="Dataset status", lines=2, value=f"Loaded default dataset
-Detected trajectories: {n_traj}")
+            hdf5_tree = gr.Textbox(lines=24, label="HDF5 tree")
 
+        # Dataset selection and custom-field visibility.
         preset.change(
             fn=update_custom_visibility,
             inputs=preset,
@@ -698,7 +779,17 @@ Detected trajectories: {n_traj}")
             outputs=[traj_slider, timestep_slider, image_keys, dataset_status],
         ).then(
             fn=render_frame,
-            inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
+            inputs=[
+                preset,
+                custom_repo_id,
+                custom_filename,
+                traj_slider,
+                timestep_slider,
+                image_keys,
+                chunk_len,
+                display_scale,
+                reverse_channels,
+            ],
             outputs=[gallery, action_plot, info],
         )
 
@@ -719,39 +810,93 @@ Detected trajectories: {n_traj}")
             outputs=[timestep_slider, image_keys],
         ).then(
             fn=render_frame,
-            inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
+            inputs=[
+                preset,
+                custom_repo_id,
+                custom_filename,
+                traj_slider,
+                timestep_slider,
+                image_keys,
+                chunk_len,
+                display_scale,
+                reverse_channels,
+            ],
             outputs=[gallery, action_plot, info],
         )
 
-        # Use release for the timestep slider so the gallery does not clear/re-render
-        # continuously while the user drags through a trajectory.
+        # Render only after releasing the timestep slider, reducing flicker.
         timestep_slider.release(
             fn=render_frame,
-            inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
+            inputs=[
+                preset,
+                custom_repo_id,
+                custom_filename,
+                traj_slider,
+                timestep_slider,
+                image_keys,
+                chunk_len,
+                display_scale,
+                reverse_channels,
+            ],
             outputs=[gallery, action_plot, info],
         )
 
-        # These controls can re-render immediately because they are changed less often.
         for widget in [image_keys, chunk_len, display_scale, reverse_channels]:
             widget.change(
                 fn=render_frame,
-                inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
+                inputs=[
+                    preset,
+                    custom_repo_id,
+                    custom_filename,
+                    traj_slider,
+                    timestep_slider,
+                    image_keys,
+                    chunk_len,
+                    display_scale,
+                    reverse_channels,
+                ],
                 outputs=[gallery, action_plot, info],
             )
 
         render_btn.click(
             fn=render_frame,
-            inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
+            inputs=[
+                preset,
+                custom_repo_id,
+                custom_filename,
+                traj_slider,
+                timestep_slider,
+                image_keys,
+                chunk_len,
+                display_scale,
+                reverse_channels,
+            ],
             outputs=[gallery, action_plot, info],
         )
 
+        inspect_btn.click(
+            fn=inspect_hdf5_tree,
+            inputs=[preset, custom_repo_id, custom_filename],
+            outputs=hdf5_tree,
+        )
+
         demo.load(
-            fn=update_after_traj_change,
-            inputs=[preset, custom_repo_id, custom_filename, traj_slider],
-            outputs=[timestep_slider, image_keys],
+            fn=update_after_dataset_change,
+            inputs=[preset, custom_repo_id, custom_filename],
+            outputs=[traj_slider, timestep_slider, image_keys, dataset_status],
         ).then(
             fn=render_frame,
-            inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
+            inputs=[
+                preset,
+                custom_repo_id,
+                custom_filename,
+                traj_slider,
+                timestep_slider,
+                image_keys,
+                chunk_len,
+                display_scale,
+                reverse_channels,
+            ],
             outputs=[gallery, action_plot, info],
         )