Update app.py

app.py

@@ -2,25 +2,29 @@ import gradio as gr
 import numpy as np, imageio, tempfile, os
 import pybullet as p
 import openai
-from openai import AuthenticationError 
+from openai import AuthenticationError
 from contactvla.env import BoxPushEnv
 from contactvla.mppi import MPPI
 from contactvla.llm_feedback import LLMFeedback
 
+
 def run_demo(api_key):
+    """Main demo entrypoint for Hugging Face Space"""
 
-    for i in range(p.getNumConnections()):
-        try:
-            p.disconnect(i)
-        except:
-            pass
+    # --- 1. PyBullet cleanup: disconnect any previous physics servers ---
+    try:
+        p.disconnect()  # closes any active simulation session
+    except Exception:
+        pass
 
+    # --- 2. Validate OpenAI API key ---
     api_key = api_key.strip()
     if not api_key:
         return None, "⚠️ Please enter a valid OpenAI API key."
 
     try:
         openai.api_key = api_key
+        # Quick ping test for key validation
         openai.chat.completions.create(
             model="gpt-4o-mini",
             messages=[{"role": "user", "content": "ping"}],
@@ -30,7 +34,8 @@ def run_demo(api_key):
         return None, "⚠️ Invalid OpenAI API key. Please check and try again."
     except Exception as e:
         return None, f"⚠️ Could not reach OpenAI API: {e}"
-    
+
+    # --- 3. Initialize PyBullet env ---
     cid = p.connect(p.DIRECT)
     env = BoxPushEnv(cid)
     mppi = MPPI(env)
@@ -44,29 +49,55 @@ def run_demo(api_key):
     env.reset()
     history = []
 
-    for step in range(20):
+    # --- 4. Run MPPI controller ---
+    for step in range(30):
         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(),
+            "ee_pos": ee_pos.tolist(),
+            "cost": cost
+        })
 
-        # render frame off-screen
+        # Off-screen rendering (TinyRenderer)
         _, _, 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)
+    # --- 5. Query LLM for improved cost function ---
+    try:
+        code, expl = llm.ask_state_cost_fn("Flip the box using the wall", history, env)
+        env.update_state_cost(code)
+    except AuthenticationError:
+        return None, "⚠️ Your API key was rejected during LLM call. Please try again."
+    except Exception as e:
+        return None, f"⚠️ Error while querying LLM: {e}"
 
+    # --- 6. Render final video ---
     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}"
+    # --- 7. Disconnect PyBullet to clean up ---
+    p.disconnect()
+
+    # --- 8. Return output ---
+    return tmp_path, f"✅ LLM updated cost:\n{code}\n\nExplanation:\n{expl}"
+
 
+# --- Gradio UI ---
 gr.Interface(
     fn=run_demo,
-    inputs=gr.Textbox(label="Enter your OpenAI API Key (without quotation marks)", type="password"),
-    outputs=["video", "text"],  
-    title="ContactVLA: LLM-guided Box Pushing",
-    description="Watch the box being pushed and flipped using an LLM-driven controller."
+    inputs=gr.Textbox(
+        label="Enter your OpenAI API Key (without quotation marks)",
+        type="password",
+        placeholder="sk-..."
+    ),
+    outputs=["video", "text"],
+    title="ContactVLA: LLM-guided Box Flipping",
+    description=(
+        "This demo uses a PyBullet simulation of a Panda robot pushing a box against a wall. "
+        "After several steps, an LLM rewrites the cost function guiding the control policy."
+    ),
 ).launch()