File size: 2,894 Bytes
a7aef48
664c7f1
a7aef48
0570cb1
77dbc3b
6d7ddff
 
 
36f09aa
77dbc3b
36f09aa
77dbc3b
60a3d23
77dbc3b
 
 
 
60a3d23
0e0da75
 
 
 
 
 
77dbc3b
0e0da75
 
 
 
 
 
 
 
 
77dbc3b
36f09aa
 
 
 
 
1d0da0e
 
 
 
 
36f09aa
 
1d0da0e
ac5c2f7
36f09aa
 
 
77dbc3b
 
 
 
 
 
1d0da0e
 
 
 
 
77dbc3b
 
 
 
 
 
 
1d0da0e
 
 
 
77dbc3b
 
 
 
36f09aa
 
 
77dbc3b
 
 
 
 
 
6d7ddff
77dbc3b
 
 
 
36f09aa
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
import gradio as gr
import numpy as np, imageio, tempfile, os
import pybullet as p
import openai
from openai import AuthenticationError
from coral.env import BoxPushEnv
from coral.mppi import MPPI
from coral.llm_feedback import LLMFeedback


def run_demo(api_key):
    """Main demo entrypoint for Hugging Face Space"""

    try:
        p.disconnect()  # closes any active simulation session
    except Exception:
        pass

    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"}],
            max_tokens=1
        )
    except AuthenticationError:
        return None, "⚠️ Invalid OpenAI API key. Please check and try again."
    except Exception as e:
        return None, f"⚠️ Could not reach OpenAI API: {e}"

    cid = p.connect(p.DIRECT)
    env = BoxPushEnv(cid)
    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(35):
        u = mppi.compute_control()
        box_pos, ee_pos = env.step(u)
        cost = env.state_cost()
        history.append({
            "step": step,
            "box_pos": box_pos.tolist(),
            "ee_pos": ee_pos.tolist(),
            "cost": cost
        })

        _, _, 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)

    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}"

    tmp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
    imageio.mimsave(tmp_path, frames, fps=10)

    p.disconnect()

    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 (without quotation marks)",
        type="password",
        placeholder="sk-..."
    ),
    outputs=["video", "text"],
    title="CoRAL: 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()