Show video output only

app.py

@@ -43,7 +43,7 @@ def simulate_agent(stage_selection):
     try:
         weights_path = hf_hub_download(repo_id=repo_id, filename=filename)
     except Exception as e:
-        return None, f"Weight download failed. Error: {str(e)}"
+        raise gr.Error(f"Weight download failed. Error: {str(e)}")
 
     # Initialize env
     env = gym.make("LunarLander-v3", render_mode="rgb_array")
@@ -58,12 +58,11 @@ def simulate_agent(stage_selection):
         actor.eval()
     except Exception as e:
         env.close()
-        return None, f"Architecture mismatch. Error: {str(e)}"
+        raise gr.Error(f"Architecture mismatch. Error: {str(e)}")
 
     state, _ = env.reset(seed=32)
     done = False
     frames = []
-    total_reward = 0.0
     step_count = 0
 
     while not done and step_count < 600:
@@ -73,15 +72,14 @@ def simulate_agent(stage_selection):
                 frames.append(frame)
         except Exception as e:
             env.close()
-            return None, f"Render failed: {str(e)}"
+            raise gr.Error(f"Render failed: {str(e)}")
             
         state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)
         with torch.no_grad():
             action_logits = actor(state_tensor)
             action = torch.argmax(action_logits, dim=1).item()
             
-        state, reward, terminated, truncated, _ = env.step(action)
-        total_reward += reward
+        state, _, terminated, truncated, _ = env.step(action)
         step_count += 1
         done = terminated or truncated
 
@@ -93,14 +91,9 @@ def simulate_agent(stage_selection):
     try:
         imageio.mimsave(video_filename, frames, fps=fps, codec='libx264', pixelformat='yuv420p')
     except Exception as e:
-        return None, f"Video encoding failed: {str(e)}"
+        raise gr.Error(f"Video encoding failed: {str(e)}")
 
-    logs = (f"Status: Inference complete\n"
-            f"Stage: {stage_selection}\n"
-            f"Total Reward: {total_reward:.2f}\n"
-            f"Steps: {step_count}")
-    
-    return video_filename, logs
+    return video_filename
 
 # 3. Gradio Web UI
 with gr.Blocks(title="Representation over Routing", theme=gr.themes.Base()) as demo:
@@ -122,13 +115,12 @@ with gr.Blocks(title="Representation over Routing", theme=gr.themes.Base()) as d
             run_button = gr.Button("Run Inference", variant="primary")
             
         with gr.Column(scale=2):
-            video_output = gr.Video(label="Environment Render")
-            text_output = gr.Textbox(label="Execution Logs", lines=4)
+            video_output = gr.Video(label="Environment Render", autoplay=True)
 
     run_button.click(
         fn=simulate_agent,
         inputs=[model_dropdown],
-        outputs=[video_output, text_output]
+        outputs=[video_output]
     )
 
 if __name__ == "__main__":