feat: OpenAI support and OpenEnv compliance fixes

Dockerfile

@@ -10,11 +10,11 @@ COPY . /app
 # Install any needed packages specified in requirements.txt
 RUN pip install --no-cache-dir -r requirements.txt
 
-# Make port 80 available to the world outside this container
-EXPOSE 80
+# Expose port (HF Spaces uses 7860)
+EXPOSE 7860
 
 # Environment variable for the HF token (can be overridden at runtime)
 ENV HF_TOKEN=""
 
-# Run baseline.py when the container launches
-CMD ["python", "baseline.py"]
+# Run the FastAPI server
+CMD ["uvicorn", "backend.main:app", "--host", "0.0.0.0", "--port", "7860"]

backend/main.py

@@ -106,7 +106,13 @@ async def reset_env():
     global current_obs, episode_done
     current_obs = env.reset()
     episode_done = False
-    return {"observation": obs_to_dict(current_obs), "done": False}
+    obs_data = obs_to_dict(current_obs)
+    return {
+        "observation": obs_data, 
+        "state": obs_data,  # OpenEnv compliance
+        "done": False,
+        "info": {"step_count": 0, "progress": 0.0}
+    }
 
 
 @app.post("/step")
@@ -126,8 +132,10 @@ async def step_env(req: StepRequest):
     current_obs = obs
     episode_done = done
 
+    obs_data = obs_to_dict(obs)
     return {
-        "observation": obs_to_dict(obs),
+        "observation": obs_data,
+        "state": obs_data, # OpenEnv compliance
         "reward": round(reward, 4),
         "done": done,
         "info": {

baseline.py

@@ -49,7 +49,7 @@ def heuristic_policy(obs: Observation) -> Action:
     return Action(type="noop")
 
 
-def llm_policy(client: OpenAI, obs: Observation) -> Action:
+def llm_policy(client: OpenAI, obs: Observation, model_name: str) -> Action:
     state_desc = (
         f"Device: {obs.device}\n"
         f"Button Size: {obs.layout.button_size:.2f}\n"
@@ -71,7 +71,7 @@ def llm_policy(client: OpenAI, obs: Observation) -> Action:
     for attempt in range(max_retries + 1):
         try:
             response = client.chat.completions.create(
-                model="katanemo/Arch-Router-1.5B",
+                model=model_name,
                 messages=[
                     {"role": "system", "content": "You are a UI optimization agent."},
                     {"role": "user", "content": prompt},
@@ -112,15 +112,15 @@ def llm_policy(client: OpenAI, obs: Observation) -> Action:
     return Action(type="noop")
 
 
-def agent_policy(client: OpenAI, obs: Observation) -> Action:
+def agent_policy(client: OpenAI, obs: Observation, model_name: str) -> Action:
     heuristic_action = heuristic_policy(obs)
     if heuristic_action.type != "noop":
         return heuristic_action
     else:
-        return llm_policy(client, obs)
+        return llm_policy(client, obs, model_name)
 
 
-def run_episode(env: UIEnv, client: OpenAI) -> Tuple[float, bool]:
+def run_episode(env: UIEnv, client: OpenAI, model_name: str) -> Tuple[float, bool]:
     obs = env.reset()
     total_reward = 0.0
     done = False
@@ -128,7 +128,7 @@ def run_episode(env: UIEnv, client: OpenAI) -> Tuple[float, bool]:
     steps = 0
 
     while not done and steps < MAX_STEPS:
-        action = agent_policy(client, obs)
+        action = agent_policy(client, obs, model_name)
         obs, reward, done, info = env.step(action)
         total_reward += reward
         steps += 1
@@ -136,7 +136,7 @@ def run_episode(env: UIEnv, client: OpenAI) -> Tuple[float, bool]:
         if info.get("outcome") == "complete":
             completed = True
 
-        time.sleep(5)
+        time.sleep(1) # Reduced for HF readiness
 
         if DEBUG:
             print(f"    step={steps}  action={action.type}  reward={reward:+.3f}  outcome={info.get('outcome')}")
@@ -144,14 +144,14 @@ def run_episode(env: UIEnv, client: OpenAI) -> Tuple[float, bool]:
     return total_reward, completed
 
 
-def evaluate_task(task: str, client: OpenAI, n_episodes: int = 1) -> Tuple[float, float, float]:
+def evaluate_task(task: str, client: OpenAI, model_name: str, n_episodes: int = 1) -> Tuple[float, float, float]:
     total_rewards = 0.0
     completions = 0
 
     for ep in range(n_episodes):
         env = load_env(task)
 
-        reward, completed = run_episode(env, client)
+        reward, completed = run_episode(env, client, model_name)
         total_rewards += reward
         if completed:
             completions += 1
@@ -167,24 +167,35 @@ def evaluate_task(task: str, client: OpenAI, n_episodes: int = 1) -> Tuple[float
 
 
 def main():
+    openai_key = os.getenv("OPENAI_API_KEY")
     hf_token = os.getenv("HF_TOKEN")
-    if not hf_token:
-        print("Error: HF_TOKEN environment variable not set.")
+    
+    if openai_key:
+        print("Using OpenAI API...")
+        client = OpenAI(api_key=openai_key)
+        model_name = os.getenv("OPENAI_MODEL", "gpt-4o-mini")
+        provider_name = f"OpenAI ({model_name})"
+    elif hf_token:
+        print("Using Hugging Face Router...")
+        client = OpenAI(
+            base_url="https://router.huggingface.co/v1",
+            api_key=hf_token
+        )
+        model_name = "katanemo/Arch-Router-1.5B"
+        provider_name = f"HF Router ({model_name})"
+    else:
+        print("Error: Neither OPENAI_API_KEY nor HF_TOKEN environment variable set.")
         return
 
-    client = OpenAI(
-        base_url="https://router.huggingface.co/v1",
-        api_key=os.getenv("HF_TOKEN")
-    )
     tasks = ["easy", "medium", "hard"]
 
     print("=" * 50)
-    print("  UIEnv Baseline Evaluation (Hugging Face Router)")
+    print(f"  UIEnv Baseline Evaluation ({provider_name})")
     print("=" * 50)
 
     for task in tasks:
         print(f"\n> Evaluating task: {task}...")
-        avg_reward, completion_rate, score = evaluate_task(task, client)
+        avg_reward, completion_rate, score = evaluate_task(task, client, model_name)
         print(f"\nTask: {task}")
         print(f"  Avg Reward:      {avg_reward:.4f}")
         print(f"  Completion Rate: {completion_rate:.4f}")

requirements.txt

@@ -1,3 +1,6 @@
 openai
 pydantic
 numpy
+fastapi
+uvicorn
+python-multipart