Update app.py

app.py

@@ -268,22 +268,46 @@ def load_traj(repo_id, filename, traj_id):
 
 
 def _extract_latest_obs_value(value):
+    """Return the latest stacked observation only when there is a clear stack axis.
+
+    Important:
+    - [obs_T, C, H, W] or [obs_T, H, W, C] should become the latest frame.
+    - [C, H, W] must NOT be sliced, otherwise an RGB image becomes one
+      grayscale channel.
+    """
     arr = np.asarray(value)
-    if arr.ndim >= 1 and arr.shape[0] in (1, 2, 3, 4):
+
+    # Stacked image observations, e.g. [obs_T, C, H, W] or [obs_T, H, W, C].
+    if arr.ndim == 4 and arr.shape[0] in (1, 2, 3, 4):
+        channel_first = arr.shape[1] in (1, 3, 4)
+        channel_last = arr.shape[-1] in (1, 3, 4)
+        if channel_first or channel_last:
+            return arr[-1]
+
+    # Stacked vector observations, e.g. [obs_T, D]. Keep this for non-image obs.
+    if arr.ndim == 2 and arr.shape[0] in (1, 2):
         return arr[-1]
+
     return arr
 
 
 def _looks_like_image_array(key, value):
-    arr = np.asarray(_extract_latest_obs_value(value))
+    arr = np.asarray(value)
     key_l = str(key).lower()
     key_hint = any(hint in key_l for hint in IMAGE_KEY_HINTS)
 
+    # Remove only a clear stacked-image axis for shape detection.
+    if arr.ndim == 4 and arr.shape[0] in (1, 2, 3, 4):
+        if arr.shape[1] in (1, 3, 4) or arr.shape[-1] in (1, 3, 4):
+            arr = arr[-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)
+    elif arr.ndim == 4:
+        shape_hint = arr.shape[1] in (1, 3, 4) or arr.shape[-1] in (1, 3, 4)
 
     return key_hint or shape_hint
 
@@ -918,6 +942,14 @@ def render_frame(preset_name, custom_repo_id, custom_filename, traj_id, timestep
 
     status_plot, is_valid_start, num_valid_starts = get_cached_status_plot(repo_id, filename, traj_id, timestep, chunk_len)
 
+    image_debug_lines = []
+    for _key in image_keys:
+        if _key in step.get("obs", {}):
+            _arr = np.asarray(step["obs"][_key])
+            image_debug_lines.append(
+                "{} shape={} dtype={}".format(_key, tuple(_arr.shape), _arr.dtype)
+            )
+
     info_lines = [
         "dataset: {} / {}".format(repo_id, filename),
         "detected trajectories: {}".format(n_traj),
@@ -935,6 +967,9 @@ def render_frame(preset_name, custom_repo_id, custom_filename, traj_id, timestep
         "",
         "teacher_action: {}".format(_safe_array_str(step.get("teacher_action", []))),
         "robot_action:   {}".format(_safe_array_str(step.get("robot_action", []))),
+        "",
+        "selected image tensors:",
+        *image_debug_lines,
     ]
 
     if warnings: