chore: final code refactoring for portability and robustness

baseline.py

@@ -1,8 +1,7 @@
 import os
 import random
 import time
-from typing import Tuple
-from openai import OpenAI
+from typing import Any, Tuple
 from env import UIEnv, Action, Observation
 
 VALID_ACTIONS = [
@@ -49,7 +48,7 @@ def heuristic_policy(obs: Observation) -> Action:
     return Action(type="noop")
 
 
-def llm_policy(client: OpenAI, obs: Observation, model_name: str) -> Action:
+def llm_policy(client: Any, obs: Observation, model_name: str) -> Action:
     state_desc = (
         f"Device: {obs.device}\n"
         f"Button Size: {obs.layout.button_size:.2f}\n"
@@ -112,7 +111,7 @@ def llm_policy(client: OpenAI, obs: Observation, model_name: str) -> Action:
     return Action(type="noop")
 
 
-def agent_policy(client: OpenAI, obs: Observation, model_name: str) -> Action:
+def agent_policy(client: Any, obs: Observation, model_name: str) -> Action:
     heuristic_action = heuristic_policy(obs)
     if heuristic_action.type != "noop":
         return heuristic_action
@@ -120,7 +119,7 @@ def agent_policy(client: OpenAI, obs: Observation, model_name: str) -> Action:
         return llm_policy(client, obs, model_name)
 
 
-def run_episode(env: UIEnv, client: OpenAI, model_name: str) -> Tuple[float, bool]:
+def run_episode(env: UIEnv, client: Any, model_name: str) -> Tuple[float, bool]:
     obs = env.reset()
     total_reward = 0.0
     done = False
@@ -144,7 +143,7 @@ def run_episode(env: UIEnv, client: OpenAI, model_name: str) -> Tuple[float, boo
     return total_reward, completed
 
 
-def evaluate_task(task: str, client: OpenAI, model_name: str, n_episodes: int = 1) -> Tuple[float, float, float]:
+def evaluate_task(task: str, client: Any, model_name: str, n_episodes: int = 1) -> Tuple[float, float, float]:
     total_rewards = 0.0
     completions = 0
 
@@ -171,6 +170,7 @@ def main():
     api_key = os.getenv("API_KEY") or os.getenv("HF_TOKEN")
 
     if api_key:
+        from openai import OpenAI
         print("Using Proxy / HF Router API...")
         client = OpenAI(
             base_url=base_url,

env.py

@@ -16,12 +16,15 @@ from pydantic import BaseModel, Field, model_validator
 EPS = 1e-6
 
 def safe_grader(fn):
-    def wrapper(x):
+    def wrapper(x=None):
         try:
             if x is None:
-                val = 0.5
+                val = fn()
             else:
-                val = fn(x)
+                try:
+                    val = fn(x)
+                except TypeError:
+                    val = fn()
 
             if not isinstance(val, (int, float)):
                 val = 0.5
@@ -179,35 +182,6 @@ class UIEnv:
                 print(f"GRADER FAIL [{t.name}]: {e}")
         print(f"VALID TASKS: {valid_tasks}")
 
-        inputs = [None, {}, "test", 0, 1]
-
-        print("=== VALIDATOR SIMULATION ===")
-
-        valid = 0
-
-        for t in self.tasks:
-            task_valid = True
-
-            for inp in inputs:
-                try:
-                    val = t.grader(inp)
-                    print(t.name, inp, val)
-
-                    if not isinstance(val, (int, float)):
-                        task_valid = False
-
-                    if not (0 < val < 1):
-                        task_valid = False
-
-                except Exception as e:
-                    print("CRASH:", t.name, inp, e)
-                    task_valid = False
-
-            if task_valid:
-                valid += 1
-
-        print("VALID TASKS:", valid)
-
         self._layout: Layout = Layout()
         self._device: Literal["mobile", "desktop"] = "desktop"
         self._progress: float = 0.0

inference.py

@@ -1,9 +1,7 @@
 import os
-import json
 import argparse
 from typing import List, Optional
 from env import UIEnv, Observation, Action, clamp_score
-from openai import OpenAI
 
 # Required Environment Variables
 API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
@@ -42,6 +40,7 @@ def run_inference(task_id: str = "easy") -> None:
     # 2. Setup Client
     client = None
     if API_KEY:
+        from openai import OpenAI
         client = OpenAI(
             base_url=API_BASE_URL,
             api_key=API_KEY
@@ -97,4 +96,4 @@ if __name__ == "__main__":
     parser.add_argument("--task", type=str, default=default_task, help="Task difficulty (easy, medium, hard)")
     args = parser.parse_args()
     
-    run_inference(task_id=args.task)
+    run_inference(task_id=args.task)

openenv_core.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from fastapi import FastAPI
+
+
+def create_app(env_cls, action_model, observation_model) -> FastAPI:
+    """Minimal local compatibility layer for this project.
+
+    The upstream package offers a richer OpenEnv integration, but this repo
+    only relies on receiving a FastAPI app instance to extend with routes.
+    """
+    app = FastAPI()
+    app.state.env_cls = env_cls
+    app.state.action_model = action_model
+    app.state.observation_model = observation_model
+    return app
+
+
+def validate() -> None:
+    from validate_local import main
+
+    main()