Pre-render frames into MP4 video for smooth playback

Replace frame-by-frame generator with pre-rendered video using imageio.
Swap gr.Image for gr.Video with autoplay and loop.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

app.py

@@ -6,10 +6,12 @@ Loads precomputed rollout cache and auto-plays animation on sample selection.
 import io
 import os
 import pickle
+import tempfile
 import time
 from pathlib import Path
 
 import gradio as gr
+import imageio
 import matplotlib
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
@@ -185,7 +187,7 @@ def _format_params_html(params):
 # Callbacks — auto-play animation as a generator
 # ---------------------------------------------------------------------------
 def on_random_sample(channel_name):
-    """Pick a random trajectory and auto-play the full animation."""
+    """Pick a random trajectory and pre-render a smooth video."""
     traj_idx = int(np.random.randint(0, NUM_TRAJS))
     entry = CACHE[traj_idx]
     num_steps = entry["num_steps"]
@@ -199,19 +201,22 @@ def on_random_sample(channel_name):
         f"&nbsp;|&nbsp; Steps: <b>{num_steps}</b></div>"
     )
 
-    # Yield each frame for auto-play
+    # Pre-render all frames
+    frames = []
     for step in range(num_steps):
-        fig = render_comparison_frame(traj_idx, step, channel_name)
-        time_str = ""
-        if entry["time_values"] is not None and step < len(entry["time_values"]):
-            time_str = f"t = {entry['time_values'][step]:.3e} &nbsp;|&nbsp; "
-        step_info = (
-            f"<div class='metrics-bar'>"
-            f"{time_str}"
-            f"Step <b>{step+1}/{num_steps}</b></div>"
-        )
-        yield fig, params_html, step_info
-        time.sleep(0.1)
+        img = render_comparison_frame(traj_idx, step, channel_name)
+        frames.append(np.array(img))
+
+    # Write frames to an MP4 video
+    video_path = tempfile.mktemp(suffix=".mp4")
+    fps = 10
+    writer = imageio.get_writer(video_path, fps=fps, codec="libx264",
+                                output_params=["-pix_fmt", "yuv420p"])
+    for frame in frames:
+        writer.append_data(frame)
+    writer.close()
+
+    return video_path, params_html, info
 
 
 # ---------------------------------------------------------------------------
@@ -343,7 +348,7 @@ def main():
         metrics_bar = gr.HTML(
             value="<div class='metrics-bar'>Click <b>Random Sample</b> to begin</div>"
         )
-        plot_output = gr.Image(label="Comparison")
+        plot_output = gr.Video(label="Comparison", autoplay=True, loop=True)
 
         random_btn.click(
             fn=on_random_sample,

requirements.txt

@@ -2,3 +2,4 @@ numpy
 matplotlib
 Pillow
 huggingface_hub
+imageio[ffmpeg]