Update app.py

app.py

@@ -1,9 +1,4 @@
-import gradio as gr
-import pybullet as p
-from contactvla.env import BoxPushEnv
-from contactvla.mppi import MPPI
-from contactvla.llm_feedback import LLMFeedback
-import numpy as np
+import cv2, numpy as np, imageio, tempfile, os
 
 def run_demo(api_key):
     cid = p.connect(p.DIRECT)
@@ -11,21 +6,37 @@ def run_demo(api_key):
     mppi = MPPI(env)
     llm = LLMFeedback(api_key)
 
+    cam_w, cam_h = 320, 240
+    view = p.computeViewMatrix([0.5, -2, 1], [0.5, 0, 0.5], [0, 0, 1])
+    proj = p.computeProjectionMatrixFOV(60, cam_w / cam_h, 0.1, 10)
+    frames = []
+
     env.reset()
     history = []
-    for step in range(10):
+
+    for step in range(20):
         u = mppi.compute_control()
         box_pos, ee_pos = env.step(u)
         cost = env.state_cost()
-        history.append({"step":step,"box_pos":box_pos.tolist(),"cost":cost})
+        history.append({"step": step, "box_pos": box_pos.tolist(), "cost": cost})
+
+        # render frame off-screen
+        _, _, rgb, _, _ = p.getCameraImage(cam_w, cam_h, view, proj, renderer=p.ER_TINY_RENDERER)
+        frame = np.array(rgb, dtype=np.uint8).reshape(cam_h, cam_w, 4)[:, :, :3]
+        frames.append(frame)
+
     code, expl = llm.ask_state_cost_fn("Flip the box using the wall", history, env)
     env.update_state_cost(code)
-    return f"LLM updated cost:\n{code}\n\nExplanation:\n{expl}"
+
+    tmp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
+    imageio.mimsave(tmp_path, frames, fps=10)
+
+    return tmp_path, f"LLM updated cost:\n{code}\n\nExplanation:\n{expl}"
 
 gr.Interface(
     fn=run_demo,
-    inputs=gr.Textbox(label="Enter your OpenAI API Key", type="password"),
-    outputs="text",
+    inputs=gr.Textbox(label="Enter your OpenAI API Key (without quotation marks)", type="password"),
+    outputs=["video", "text"],  
     title="ContactVLA: LLM-guided Box Pushing",
-    description="Enter your OpenAI API key to run an LLM-driven robot control demo."
+    description="Watch the box being pushed and flipped using an LLM-driven controller."
 ).launch()