Kolaps27/UI-layout-optimizer

1import os2import argparse3from typing import List, Optional4from env import UIEnv, Observation, Action, clamp_score5 6# Required Environment Variables7API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")8MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")9API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")10IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")11BENCHMARK = os.getenv("BENCHMARK", "ui_layout_optimizer")12 13def log_start(task: str, env: str, model: str) -> None:14    print(f"[START] task={task} env={env} model={model}", flush=True)15 16def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:17    error_val = error if error else "null"18    done_val = str(done).lower()19    print(20        f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",21        flush=True,22    )23 24def log_end(success: bool, steps: int, rewards: List[float]) -> None:25    rewards_str = ",".join(f"{r:.2f}" for r in rewards)26    success_val = str(success).lower()27    print(f"[END] success={success_val} steps={steps} rewards={rewards_str}", flush=True)28 29def run_inference(task_id: str = "easy") -> None:30    """31    Standard OpenEnv inference entry point.32    Evaluates agent performance on concrete objectives.33    """34    log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)35 36    # 1. Setup Environment37    env = UIEnv(seed=42, task=task_id)38    obs = env.reset()39    40    # 2. Setup Client41    client = None42    if API_KEY:43        from openai import OpenAI44        client = OpenAI(45            base_url=API_BASE_URL,46            api_key=API_KEY47        )48 49    done = False50    step_count = 051    total_reward = 0.052    completed = False53    rewards: List[float] = []54 55    while not done:56        step_count += 157 58        # 3. Perform Inference Step59        if client:60            from baseline import agent_policy61            action = agent_policy(client, obs, MODEL_NAME)62        else:63            # Fallback to heuristic if no API key is provided64            from baseline import heuristic_policy65            action = heuristic_policy(obs)66        67        # Format action for logging68        action_str = action.type69        if action.value is not None:70            action_str += f"({action.value})"71 72        # 4. Step Environment73        obs, reward, done, info = env.step(action)74        75        rewards.append(reward)76        total_reward += reward77        78        error = info.get("error") # Can be parsed dynamically if environment fails on step internally79        80        log_step(step=step_count, action=action_str, reward=reward, done=done, error=error)81 82        if info.get("outcome") == "complete" or info.get("completed") is True:83            completed = True84 85    # 5. Retrieve agent performance score from evaluator86    task_obj = env.task_dict[task_id]87    score = task_obj.grader()88 89    # Enforce strict (0,1) bound90    score = clamp_score(score)91    log_end(success=completed, steps=step_count, rewards=rewards)92 93if __name__ == "__main__":94    parser = argparse.ArgumentParser(description="Run UIEnv Inference")95    default_task = os.getenv("TASK", os.getenv("MY_ENV_TASK", "easy"))96    parser.add_argument("--task", type=str, default=default_task, help="Task difficulty (easy, medium, hard)")97    args = parser.parse_args()98    99    run_inference(task_id=args.task)100