Update app.py

app.py

@@ -445,13 +445,110 @@ def get_cached_gallery_items(repo_id, filename, traj_id, timestep, image_keys_tu
     return gallery_items, tuple(warnings)
 
 
+def _compute_valid_start_indices(traj, min_seq_len):
+    """Match the original local script's valid-start heuristic.
+
+    A timestep is valid when the following min_seq_len steps all have
+    no_teacher_action == False.
+    """
+    total_steps = len(traj)
+    min_seq_len = int(max(1, min_seq_len))
+    no_teacher = np.asarray(
+        [int(bool(step.get("no_teacher_action", False))) for step in traj],
+        dtype=np.int32,
+    )
+
+    valid_indices = []
+    max_start = total_steps - min_seq_len + 1
+    for t in range(max(0, max_start)):
+        if int(np.sum(no_teacher[t:t + min_seq_len])) == 0:
+            valid_indices.append(t)
+
+    return no_teacher, valid_indices
+
+
+def _make_trajectory_status_plot(traj, timestep, min_seq_len):
+    """Render the same high-level status figure as the local matplotlib tool.
+
+    Shows:
+      - orange no_teacher_action step plot
+      - green triangles for algorithmic valid start points
+      - black vertical cursor at current timestep
+    """
+    total_steps = len(traj)
+    if total_steps == 0:
+        return None, False, 0
+
+    timestep = int(np.clip(int(timestep), 0, total_steps - 1))
+    timesteps = np.asarray(
+        [int(np.asarray(step.get("timestep", idx)).reshape(-1)[0]) for idx, step in enumerate(traj)],
+        dtype=np.int32,
+    )
+    no_teacher, valid_indices = _compute_valid_start_indices(traj, min_seq_len)
+    is_valid_start = timestep in set(valid_indices)
+
+    fig, ax = plt.subplots(figsize=(9, 2.4), dpi=140)
+
+    ax.step(
+        np.arange(total_steps),
+        no_teacher,
+        where="post",
+        label="no_teacher_action",
+        color="orange",
+    )
+
+    if valid_indices:
+        ax.scatter(
+            valid_indices,
+            [-0.15] * len(valid_indices),
+            color="green",
+            marker="^",
+            s=18,
+            label="Valid Start (len >= {})".format(int(min_seq_len)),
+        )
+
+    ax.axvline(timestep, color="black", linestyle="-", alpha=0.85, linewidth=1.5)
+    ax.set_xlim(0, max(total_steps - 1, 1))
+    ax.set_ylim(-0.38, 1.1)
+    ax.set_ylabel("Flag")
+    ax.set_xlabel("Timestep index")
+    ax.set_yticks([0, 1])
+    ax.set_yticklabels(["False", "True"])
+    ax.grid(True, axis="x", alpha=0.2)
+
+    title = "no_teacher_action | step {} / {}".format(timestep, total_steps - 1)
+    if is_valid_start:
+        title += " | VALID START"
+    ax.set_title(title)
+    ax.legend(loc="upper right", fontsize=8)
+
+    # Add saved timestep annotation if the stored timestep is not the same as index.
+    saved_timestep = int(timesteps[timestep]) if len(timesteps) else timestep
+    if saved_timestep != timestep:
+        ax.text(
+            0.01,
+            0.04,
+            "saved timestep: {}".format(saved_timestep),
+            transform=ax.transAxes,
+            fontsize=8,
+            va="bottom",
+            ha="left",
+        )
+
+    fig.tight_layout()
+    fig.canvas.draw()
+    rgba = np.asarray(fig.canvas.buffer_rgba())
+    image = rgba[..., :3].copy()
+    plt.close(fig)
+
+    return image, bool(is_valid_start), len(valid_indices)
+
+
 @lru_cache(maxsize=8192)
-def get_cached_action_plot(repo_id, filename, traj_id, timestep, chunk_len):
+def get_cached_status_plot(repo_id, filename, traj_id, timestep, min_seq_len):
     traj = load_traj(repo_id, filename, int(traj_id))
     timestep = int(np.clip(int(timestep), 0, len(traj) - 1))
-    mixed_chunk, source_mask = _extract_mixed_action_chunk(traj, timestep, int(chunk_len))
-    robot_chunk = _extract_robot_action_chunk(traj, timestep, int(chunk_len))
-    return _make_action_chunk_plot(mixed_chunk, robot_chunk), source_mask
+    return _make_trajectory_status_plot(traj, timestep, int(min_seq_len))
 
 
 def preload_current_trajectory(preset_name, custom_repo_id, custom_filename, traj_id, image_keys, chunk_len, display_scale, reverse_channels):
@@ -474,7 +571,7 @@ def preload_current_trajectory(preset_name, custom_repo_id, custom_filename, tra
     total = len(traj)
     for t in range(total):
         get_cached_gallery_items(repo_id, filename, traj_id, t, image_keys_tuple, float(display_scale), bool(reverse_channels))
-        get_cached_action_plot(repo_id, filename, traj_id, t, int(chunk_len))
+        get_cached_status_plot(repo_id, filename, traj_id, t, int(chunk_len))
 
     status = "Preloaded trajectory {}".format(traj_id)
     status += "\nFrames cached: {}".format(total)
@@ -675,7 +772,7 @@ def render_frame(preset_name, custom_repo_id, custom_filename, traj_id, timestep
     )
     warnings = list(warnings_tuple)
 
-    action_plot, source_mask = get_cached_action_plot(repo_id, filename, traj_id, timestep, chunk_len)
+    status_plot, is_valid_start, num_valid_starts = get_cached_status_plot(repo_id, filename, traj_id, timestep, chunk_len)
 
     info_lines = [
         "dataset: {} / {}".format(repo_id, filename),
@@ -688,8 +785,9 @@ def render_frame(preset_name, custom_repo_id, custom_filename, traj_id, timestep
         "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),
+        "valid-window length: {}".format(chunk_len),
+        "valid_start: {}".format(int(bool(is_valid_start))),
+        "num_valid_starts: {}".format(num_valid_starts),
         "",
         "teacher_action: {}".format(_safe_array_str(step.get("teacher_action", []))),
         "robot_action:   {}".format(_safe_array_str(step.get("robot_action", []))),
@@ -700,7 +798,7 @@ def render_frame(preset_name, custom_repo_id, custom_filename, traj_id, timestep
         info_lines.append("Image warnings:")
         info_lines.extend(warnings)
 
-    return gallery_items, action_plot, "\n".join(info_lines)
+    return gallery_items, status_plot, "\n".join(info_lines)
 
 
 def build_app():
@@ -720,7 +818,8 @@ def build_app():
     with gr.Blocks(title="HDF5 Trajectory Viewer") as demo:
         gr.Markdown(
             "# HDF5 Trajectory Viewer\n\n"
-            "Standalone viewer for TrajectoryBuffer-style HDF5 datasets on Hugging Face."
+            "Standalone viewer for TrajectoryBuffer-style HDF5 datasets on Hugging Face.\n\n"
+            "The status plot matches the local labeling view: orange `no_teacher_action`, green valid-start markers, and a black timestep cursor."
         )
 
         if startup_warning:
@@ -743,7 +842,7 @@ def build_app():
 
         with gr.Row():
             image_keys = gr.CheckboxGroup(choices=first_keys, value=first_keys[:2], label="Image keys")
-            chunk_len = gr.Slider(minimum=1, maximum=64, value=DEFAULT_CHUNK_LEN, step=1, label="Action chunk length")
+            chunk_len = gr.Slider(minimum=1, maximum=64, value=DEFAULT_CHUNK_LEN, step=1, label="Valid-window length")
             display_scale = gr.Slider(minimum=1, maximum=10, value=4, step=1, label="Image display scale")
             reverse_channels = gr.Checkbox(value=False, label="Reverse channels BGR↔RGB")
 
@@ -757,7 +856,7 @@ def build_app():
         trajectory_video = gr.Video(label="Trajectory video: smooth browser-side playback")
 
         gallery = gr.Gallery(label="Observation images", columns=2, height="auto", object_fit="contain")
-        action_plot = gr.Image(label="Action chunk plot", type="numpy")
+        status_plot = gr.Image(label="Trajectory status: no_teacher_action and valid starts", type="numpy")
         info = gr.Textbox(label="Frame info", lines=16)
 
         with gr.Accordion("Debug: HDF5 tree", open=False):
@@ -775,7 +874,7 @@ def build_app():
         ).then(
             fn=render_frame,
             inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
-            outputs=[gallery, action_plot, info],
+            outputs=[gallery, status_plot, info],
         )
 
         custom_repo_id.submit(
@@ -796,26 +895,26 @@ def build_app():
         ).then(
             fn=render_frame,
             inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
-            outputs=[gallery, action_plot, info],
+            outputs=[gallery, status_plot, info],
         )
 
         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],
-            outputs=[gallery, action_plot, info],
+            outputs=[gallery, status_plot, info],
         )
 
         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],
-                outputs=[gallery, action_plot, info],
+                outputs=[gallery, status_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],
-            outputs=[gallery, action_plot, info],
+            outputs=[gallery, status_plot, info],
         )
 
         preload_btn.click(
@@ -843,7 +942,7 @@ def build_app():
         ).then(
             fn=render_frame,
             inputs=[preset, custom_repo_id, custom_filename, traj_slider, timestep_slider, image_keys, chunk_len, display_scale, reverse_channels],
-            outputs=[gallery, action_plot, info],
+            outputs=[gallery, status_plot, info],
         )
 
     return demo