chore: import upstream snapshot with attribution

osworld_setup/s3/OSWorld.md

@@ -0,0 +1,54 @@
+# Deplying Agent S3 in OSWorld
+
+# Step 1: Set up Agent S3
+
+Follow the [README.md](https://github.com/simular-ai/Agent-S/blob/main/README.md) to set up Agent S3.
+
+# Step 2: Copying Over Run Files
+
+If you haven't already, please follow the [OSWorld environment setup](https://github.com/xlang-ai/OSWorld/blob/main/README.md). We've provided the relevant OSWorld run files for evaluation in this `osworld_setup` folder. Please copy this over to your OSWorld folder. `run_local.py` is for if you want to run locally on VMWare and `run.py` and `lib_run_single.py` are for if you want to run on AWS. All run commands in order are provided in the `run.sh`. Copy over the files in `osworld_setup/s3/bbon` as well. 
+
+# Step 3: Switch the AMI 
+
+Switch image AMI for the AWS provider in `desktop_env/providers/aws/manager.py` is set to `"ami-0b505e9d0d99ba88c"`.
+
+# Step 4: Generating Facts
+
+After completing your OSWorld runs and having result directories, run `generate_facts.py` to generate fact captions for screenshot pairs:
+
+```bash
+python osworld_setup/s3/bbon/generate_facts.py \
+  --results-dirs \
+    results1/pyautogui/screenshot/gpt-5-2025-08-07 \
+    results2/pyautogui/screenshot/gpt-5-2025-08-07 \
+  --model "gpt-5-2025-08-07" \
+  --engine-type "openai" \
+  --temperature 1.0
+```
+
+This will populate your result directories with `fact_captions.jsonl` files containing behavioral descriptions of screenshot differences.
+
+# Step 5: Run the Judge
+
+Finally, run `run_judge.py` to evaluate the trajectories using the generated fact captions:
+
+```bash
+python osworld_setup/s3/bbon/run_judge.py \
+  --results-dirs \
+    results1/pyautogui/screenshot/gpt-5-2025-08-07 \
+    results2/pyautogui/screenshot/gpt-5-2025-08-07 \
+  --output-dir "judge_results" \
+  --examples-path "evaluation_examples/examples" \
+  --model "gpt-5-2025-08-07" \
+  --engine-type "openai" \
+  --temperature 1.0
+```
+
+This will:
+- Compare trajectories across different result directories
+- Use the facts to judge which trajectory performs better
+- Generate evaluation results
+- Save results to the specified output directory
+
+The judge will create files like `BoN2.json`, `BoN3.json`, etc., showing the performance comparison as you add more trajectories.
+

osworld_setup/s3/bbon/generate_facts.py

@@ -0,0 +1,241 @@
+import os
+import json
+import asyncio
+import argparse
+from typing import List, Optional
+from dotenv import load_dotenv
+
+from gui_agents.s3.bbon.behavior_narrator import BehaviorNarrator
+from utils import get_new_tasks_classification
+
+load_dotenv()
+
+
+async def generate_single_fact_caption(
+    task_dir: str,
+    screenshot_files: List[str],
+    i: int,
+    judge: BehaviorNarrator,
+    trajectory_lines: List[str],
+):
+    """Generate a single fact caption for a screenshot pair."""
+    before_file = os.path.join(task_dir, screenshot_files[i])
+    after_file = os.path.join(task_dir, screenshot_files[i + 1])
+
+    # Load action from trajectory data if available
+    pyautogui_action = None
+    if i < len(trajectory_lines):
+        try:
+            data = json.loads(trajectory_lines[i])
+            pyautogui_action = data.get("exec_code")
+        except:
+            pass
+
+    if pyautogui_action is None:
+        raise ValueError(f"No pyautogui action found for step {i+1}")
+
+    # Read image bytes
+    try:
+        with open(before_file, "rb") as f:
+            before_bytes = f.read()
+        with open(after_file, "rb") as f:
+            after_bytes = f.read()
+    except Exception as e:
+        raise Exception(f"Error reading images: {e}")
+
+    # Generate fact caption using behavior narrator
+    result = await asyncio.to_thread(
+        judge.judge,
+        screenshot_num=i + 1,
+        before_img_bytes=before_bytes,
+        after_img_bytes=after_bytes,
+        pyautogui_action=pyautogui_action,
+    )
+    result["screenshot_num"] = i + 1
+
+    return result
+
+
+async def generate_fact_captions_parallel(
+    task_dir: str,
+    judge: BehaviorNarrator,
+    step_semaphore: Optional[asyncio.Semaphore] = None,
+):
+    """Generate fact captions for a task directory when they don't exist (parallelized version)."""
+    print(f"Generating fact captions for {task_dir}...")
+
+    # Find all screenshot files
+    screenshot_files = []
+    for filename in os.listdir(task_dir):
+        if filename.startswith("step_") and filename.endswith(".png"):
+            screenshot_files.append(filename)
+
+    # Sort by step number
+    def extract_step_num(filename):
+        try:
+            return int(filename.split("_")[1].split(".")[0])
+        except:
+            return 0
+
+    screenshot_files.sort(key=extract_step_num)
+
+    if len(screenshot_files) < 2:
+        print(f"Not enough screenshots to generate fact captions in {task_dir}")
+        return []
+
+    # Load trajectory data once
+    trajectory_lines = []
+    trajectory_file = os.path.join(task_dir, "traj.jsonl")
+    if os.path.exists(trajectory_file):
+        try:
+            with open(trajectory_file, "r") as f:
+                trajectory_lines = f.readlines()
+        except:
+            pass
+
+    # Use shared semaphore to limit concurrent judge calls
+    if step_semaphore is None:
+        step_semaphore = asyncio.Semaphore(5)  # Default limit
+
+    async def bounded_task(task_func, *args, **kwargs):
+        async with step_semaphore:
+            return await task_func(*args, **kwargs)
+
+    try:
+        # Create bounded tasks for parallel execution
+        bounded_tasks = [
+            bounded_task(
+                generate_single_fact_caption,
+                task_dir,
+                screenshot_files,
+                i,
+                judge,
+                trajectory_lines,
+            )
+            for i in range(len(screenshot_files) - 1)
+        ]
+        results = await asyncio.gather(*bounded_tasks, return_exceptions=True)
+    except Exception as e:
+        print(f"Error in parallel execution: {e}")
+        return []
+
+    # Process results and save to file
+    fact_captions = []
+    successful_results = []
+    fact_captions_file = os.path.join(task_dir, "fact_captions.jsonl")
+
+    for i, result in enumerate(results):
+        if isinstance(result, Exception):
+            print(f"Error generating fact caption for step {i+1}: {result}")
+            continue
+        successful_results.append(result)
+        fact_caption = f"Fact Caption from Screenshot {result['screenshot_num']}: {result['fact_answer']}"
+        fact_captions.append(fact_caption)
+
+    # Save all results to file at once
+    if successful_results:
+        with open(fact_captions_file, "w") as f:
+            for result in successful_results:
+                f.write(json.dumps(result) + "\n")
+
+    print(f"Generated {len(fact_captions)} fact captions for {task_dir}")
+    return fact_captions
+
+
+async def main(engine_params: dict, results_dirs: List[str]):
+    """Main function to generate fact captions for multiple task directories.
+
+    Args:
+        engine_params: Engine parameters for BehaviorNarrator
+        results_dirs: List of results directories to analyze for task classification
+    """
+    # Get task IDs automatically using get_new_tasks_classification
+    tasks_classification = get_new_tasks_classification(results_dirs)
+    task_ids = tasks_classification["variance"]
+
+    print(f"Found {len(task_ids)} variance tasks to process")
+    judge = BehaviorNarrator(engine_params=engine_params)
+
+    # Get concurrency settings from environment
+    per_step = int(os.getenv("DIFFCAP_PER_STEP_CONCURRENCY", "100"))
+    per_taskdir = int(os.getenv("DIFFCAP_PER_TASKDIR_CONCURRENCY", "4"))
+
+    # Build list of task directories to process
+    task_dirs = []
+    for task_id in task_ids:
+        domain, example_id = task_id.split("/")
+
+        # Check each results directory for this task
+        for results_dir in results_dirs:
+            task_dir = os.path.join(results_dir, domain, example_id)
+
+            try:
+                if "fact_captions.jsonl" in os.listdir(task_dir):
+                    print(f"Fact captions already exist for {task_dir}")
+                    continue
+            except FileNotFoundError:
+                continue
+
+            task_dirs.append(task_dir)
+
+    if not task_dirs:
+        print("No new task directories to process.")
+        return
+
+    print(f"Scheduling {len(task_dirs)} task directories...")
+
+    # Set up semaphores for concurrency control
+    shared_step_semaphore = asyncio.Semaphore(per_step)
+    taskdir_semaphore = asyncio.Semaphore(per_taskdir)
+
+    async def run_one(task_dir):
+        async with taskdir_semaphore:
+            print(f"Processing {task_dir}")
+            return await generate_fact_captions_parallel(
+                task_dir, judge, step_semaphore=shared_step_semaphore
+            )
+
+    # Execute all tasks in parallel
+    results = await asyncio.gather(
+        *[run_one(d) for d in task_dirs], return_exceptions=True
+    )
+
+    # Report results
+    failures = sum(1 for r in results if isinstance(r, Exception))
+    if failures:
+        print(
+            f"Completed with {failures} failures out of {len(task_dirs)} task directories."
+        )
+    else:
+        print("Completed all task directories successfully.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Generate fact captions for OSWorld task directories"
+    )
+    parser.add_argument(
+        "--results-dirs",
+        nargs="+",
+        required=True,
+        help="List of results directories to analyze for task classification",
+    )
+    parser.add_argument(
+        "--model", default="gpt-5-2025-08-07", help="Model to use for generation"
+    )
+    parser.add_argument("--engine-type", default="openai", help="Engine type")
+    parser.add_argument(
+        "--temperature", type=float, default=1.0, help="Temperature for generation"
+    )
+
+    args = parser.parse_args()
+
+    # Engine parameters
+    engine_params = {
+        "model": args.model,
+        "engine_type": args.engine_type,
+        "temperature": args.temperature,
+    }
+
+    print(f"Results directories: {args.results_dirs}")
+    asyncio.run(main(engine_params, args.results_dirs))

osworld_setup/s3/bbon/run_judge.py

@@ -0,0 +1,278 @@
+import json
+import os
+import asyncio
+import argparse
+import concurrent.futures
+from typing import List, Tuple, Optional
+from dotenv import load_dotenv
+from tqdm.asyncio import tqdm_asyncio
+
+load_dotenv()
+
+from utils import (
+    get_new_tasks_classification,
+    evaluate_comparative_results,
+    load_task_instruction,
+    load_facts,
+)
+from gui_agents.s3.bbon.comparative_judge import ComparativeJudge
+
+
+def run_judge(
+    task: str, task_instruction: str, result_dirs: List[str], judge: ComparativeJudge
+) -> Tuple[str, str, Optional[str]]:
+    """
+    Fact captions + initial/final screenshots judging.
+    Pipeline: load trajectories → load existing fact captions → include initial/final screenshots → judge.
+    """
+    # 1. Use provided task instruction
+    # task_instruction is now a direct input parameter
+
+    # 2. Load fact captions for all trajectories
+    all_fact_captions = []
+    for result_dir in result_dirs:
+        task_dir = os.path.join(result_dir, task.split("/")[0], task.split("/")[1])
+        fact_captions = load_facts(task_dir)
+        all_fact_captions.append(fact_captions)
+
+    # 3. Use the new Judge class method
+    return judge.judge(task_instruction, task, result_dirs, all_fact_captions)
+
+
+def evaluate_trajectories(
+    task: str, task_instruction: str, result_dirs: List[str], judge: ComparativeJudge
+) -> Tuple[str, str, dict]:
+    """Wrapper that runs fact-only MCQ judge and returns results."""
+    answer, thoughts, selected_trajectory = run_judge(
+        task, task_instruction, result_dirs, judge
+    )
+
+    record = {
+        "selected_trajectory": selected_trajectory,
+        "answer": answer,
+        "thoughts": thoughts,
+    }
+
+    print(f"✅ Added task {task} (MCQ fact-only)")
+    return answer, thoughts, record
+
+
+asyncio.get_event_loop().set_default_executor(
+    concurrent.futures.ThreadPoolExecutor(max_workers=100)
+)
+
+
+async def run_async(
+    task: str, task_instruction: str, result_dirs: List[str], judge: ComparativeJudge
+):
+    """Async wrapper for fact-only MCQ evaluation."""
+    return await asyncio.to_thread(
+        evaluate_trajectories,
+        task=task,
+        task_instruction=task_instruction,
+        result_dirs=result_dirs,
+        judge=judge,
+    )
+
+
+async def evaluate_and_save(
+    result_dirs: List[str],
+    output_file_path: str,
+    examples_path: str,
+    engine_params: dict,
+):
+    """Main evaluation function that processes tasks and saves results."""
+    res = get_new_tasks_classification(results_dirs=result_dirs)
+    for key in res:
+        print(f"{key}: {res[key]}")
+    optimal, minimum, expected_value = (
+        res["optimal"],
+        res["minimum"],
+        res["expected_value"],
+    )
+    print(f"optimal score: {optimal}, minimum score: {minimum}")
+
+    variance = res["variance"]
+
+    judge = ComparativeJudge(engine_params=engine_params)
+
+    # Load existing results
+    if os.path.exists(output_file_path):
+        with open(output_file_path, "r", encoding="utf-8") as f:
+            try:
+                data = json.load(f)
+                if not isinstance(data, dict):
+                    data = {}
+            except json.JSONDecodeError:
+                data = {}
+    else:
+        data = {}
+
+    # Prepare async tasks only for tasks not yet in data
+    tasks = []
+    task_names = []
+    for task in variance:
+        if str(task) in data:
+            print(f"⚠️ Task {task} already exists in results — skipping.")
+            continue
+
+        # Load task instruction from examples path
+        task_instruction = load_task_instruction(task, examples_path)
+        if task_instruction is None:
+            print(f"⚠️ No task instruction found for {task}, skipping...")
+            continue
+
+        tasks.append(run_async(task, task_instruction, result_dirs, judge))
+        task_names.append(task)
+
+    # Run only new tasks
+    results = await tqdm_asyncio.gather(*tasks)
+    # Merge into existing results
+    for task, (ans, thoughts, record) in zip(task_names, results):
+        data[str(task)] = record
+
+    os.makedirs(os.path.dirname(output_file_path), exist_ok=True)
+    with open(output_file_path, "w") as f:
+        json.dump(data, f, indent=2)
+
+    res = evaluate_comparative_results(result_dirs, json_path=output_file_path)
+    gain, maximum_gain = res
+    data["score"] = {
+        "optimal": optimal,
+        "minimum": minimum,
+        "expected_value": expected_value,
+        "res": res,
+        "actual score": minimum + gain,
+    }
+    os.makedirs(os.path.dirname(output_file_path), exist_ok=True)
+    with open(output_file_path, "w") as f:
+        json.dump(data, f, indent=2)
+
+    return results
+
+
+async def run_experiment(
+    shuffled_runs: List[str],
+    output_dir: str,
+    examples_path: str,
+    engine_params: dict,
+    start_round: int = 2,
+    max_rounds: int = None,
+):
+    """
+    Run fact-only experiments progressively: start_round vs start_round+1, etc.
+    """
+    if max_rounds is None:
+        max_rounds = len(shuffled_runs)
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    for i in range(start_round, max_rounds + 1):  # start at start_round (default 2)
+        test_dirs = shuffled_runs[:i]
+        output_file_path = os.path.join(output_dir, f"BoN{i}.json")
+
+        print(f"Running fact-only experiment with {i} dirs → {output_file_path}")
+        await evaluate_and_save(
+            test_dirs, output_file_path, examples_path, engine_params
+        )
+
+
+async def main(
+    shuffled_runs: List[str] = None,
+    output_dir: str = None,
+    examples_path: str = None,
+    engine_params: dict = None,
+    start_round: int = 2,
+    max_rounds: int = None,
+):
+    """Main function to run fact-only judge experiments.
+
+    Args:
+        shuffled_runs: List of result directory paths to compare
+        output_dir: Directory to save results
+        examples_path: Path to examples directory containing task instructions
+        engine_params: Engine parameters for the judge
+        start_round: Starting round number (default: 2)
+        max_rounds: Maximum number of rounds to run (default: len(shuffled_runs))
+    """
+    if shuffled_runs is None:
+        print("Error: shuffled_runs must be provided")
+        return
+
+    if output_dir is None:
+        print("Error: output_dir must be provided")
+        return
+
+    if examples_path is None:
+        print("Error: examples_path must be provided")
+        return
+
+    if engine_params is None:
+        print("Error: engine_params must be provided")
+        return
+
+    await run_experiment(
+        shuffled_runs, output_dir, examples_path, engine_params, start_round, max_rounds
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Run fact-only judge experiments on OSWorld task directories"
+    )
+    parser.add_argument(
+        "--results-dirs",
+        nargs="+",
+        required=True,
+        help="List of results directories to analyze",
+    )
+    parser.add_argument("--output-dir", required=True, help="Directory to save results")
+    parser.add_argument(
+        "--examples-path",
+        required=True,
+        help="Path to examples directory containing task instructions",
+    )
+    parser.add_argument(
+        "--start-round", type=int, default=2, help="Starting round number (default: 2)"
+    )
+    parser.add_argument(
+        "--max-rounds",
+        type=int,
+        default=None,
+        help="Maximum number of rounds to run (default: len(results_dirs))",
+    )
+    parser.add_argument(
+        "--model", default="gpt-5-2025-08-07", help="Model to use for judging"
+    )
+    parser.add_argument("--engine-type", default="openai", help="Engine type")
+    parser.add_argument(
+        "--temperature", type=float, default=1.0, help="Temperature for generation"
+    )
+
+    args = parser.parse_args()
+
+    # Engine parameters
+    engine_params = {
+        "model": args.model,
+        "engine_type": args.engine_type,
+        "temperature": args.temperature,
+    }
+
+    print(f"Results directories: {args.results_dirs}")
+    print(f"Output directory: {args.output_dir}")
+    print(f"Examples path: {args.examples_path}")
+    print(f"Start round: {args.start_round}")
+    print(f"Max rounds: {args.max_rounds}")
+    print(f"Engine params: {engine_params}")
+
+    # Run fact-only evaluation
+    asyncio.run(
+        main(
+            shuffled_runs=args.results_dirs,
+            output_dir=args.output_dir,
+            examples_path=args.examples_path,
+            engine_params=engine_params,
+            start_round=args.start_round,
+            max_rounds=args.max_rounds,
+        )
+    )

osworld_setup/s3/bbon/utils.py

@@ -0,0 +1,301 @@
+import logging
+import os
+import re
+import json
+from PIL import Image
+from typing import Optional, List
+import base64
+
+
+def image_to_openai_message_format(
+    image_path: str, caption: str = None
+) -> Optional[dict]:
+    """Convert an image file to OpenAI message format."""
+    if not os.path.exists(image_path):
+        print(f"Image file not found: {image_path}")
+        return None
+
+    try:
+        with open(image_path, "rb") as f:
+            image_bytes = f.read()
+
+        if not image_bytes:
+            print(f"Empty image file: {image_path}")
+            return None
+
+        base64_image = base64.b64encode(image_bytes).decode("utf-8")
+
+        if not base64_image:
+            print(f"Failed to encode image to base64: {image_path}")
+            return None
+
+        content = []
+        if caption:
+            content.append({"type": "text", "text": caption})
+
+        content.append(
+            {
+                "type": "image_url",
+                "image_url": {"url": f"data:image/png;base64,{base64_image}"},
+            }
+        )
+
+        return {"role": "user", "content": content}
+
+    except Exception as e:
+        print(f"Error processing image {image_path}: {e}")
+        return None
+
+
+def load_facts(task_dir: str) -> List[str]:
+    """Load existing facts from facts.jsonl file."""
+    fact_captions_file = os.path.join(task_dir, "fact_captions.jsonl")
+
+    if not os.path.exists(fact_captions_file):
+        print(f"fact_captions.jsonl not found at {fact_captions_file}")
+        return []
+
+    fact_captions = []
+    with open(fact_captions_file, "r") as f:
+        for line in f:
+            if line.strip():
+                data = json.loads(line)
+                if "fact_answer" in data:
+                    fact_captions.append(data["fact_answer"])
+
+    return fact_captions
+
+
+def load_task_instruction(task: str, examples_path: str) -> Optional[str]:
+    """
+    Load task instruction from examples path.
+
+    Args:
+        task: Task ID in format "domain/example_id"
+        examples_path: Path to the examples directory (e.g., "/home/ubuntu/Simular/OSWorld/evaluation_examples/examples")
+
+    Returns:
+        Task instruction string or None if not found
+    """
+    domain, example_id = task.split("/", 1)
+
+    # Construct path to the JSON file
+    json_file_path = os.path.join(examples_path, domain, f"{example_id}.json")
+
+    if not os.path.exists(json_file_path):
+        logging.warning(f"Example file not found: {json_file_path}")
+        return None
+
+    try:
+        with open(json_file_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+
+        # Extract instruction from the JSON
+        if "instruction" in data:
+            instruction = data["instruction"]
+            if instruction and instruction.strip():
+                return instruction.strip()
+
+        logging.warning(f"No 'instruction' key found in {json_file_path}")
+        return None
+
+    except Exception as e:
+        logging.warning(f"Error reading example file {json_file_path}: {e}")
+        return None
+
+
+def get_final_screenshot_file(result_dir: str) -> str:
+    """
+    Finds the screenshot file with the largest valid step index in the given directory.
+    Works with filenames like step_0.png, step_1_20250.png, step-2.png, etc.
+    Only considers .png files (case-insensitive).
+    If the highest index file is invalid/corrupted, it tries the next lower index.
+    Returns None if no valid matching files are found.
+    """
+    # First, collect all valid step files with their indices
+    step_files = {}
+    pattern = re.compile(r"step[_\-]?(\d+)", re.IGNORECASE)
+
+    for fname in os.listdir(result_dir):
+        if not fname.lower().endswith(".png"):
+            continue
+        match = pattern.match(fname)
+        if match:
+            idx = int(match.group(1))
+            step_files[idx] = fname
+    if not step_files:
+        return None
+    # Sort indices in descending order (highest first)
+    sorted_indices = sorted(step_files.keys(), reverse=True)
+    # Try each file from highest to lowest index
+    for idx in sorted_indices:
+        fname = step_files[idx]
+        file_path = os.path.join(result_dir, fname)
+        # Check if file exists and is valid
+        if os.path.exists(file_path) and is_valid_image(file_path):
+            return fname
+        else:
+            print(
+                f"Invalid or corrupted image at step {idx}: {fname}, trying previous step..."
+            )
+    return None
+
+
+def is_valid_image(file_path: str) -> bool:
+    """
+    Check if an image file is valid by trying to open it with PIL.
+    Also checks if file is not empty.
+    """
+    try:
+        # Check file size first (quick check)
+        if os.path.getsize(file_path) == 0:
+            return False
+
+        # Try to open and verify the image
+        with Image.open(file_path) as img:
+            img.verify()  # This will raise an exception if image is corrupted
+            return True
+    except Exception as e:
+        print(f"Image validation failed for {file_path}: {e}")
+        return False
+
+
+def get_new_tasks_classification(results_dirs: [str]):
+    # Step 1: collect domain/task_ids for each trajectory
+    tasks_per_dir = []
+    for results_dir in results_dirs:
+        domain_tasks = set()
+        for domain in os.listdir(results_dir):
+            domain_dir = os.path.join(results_dir, domain)
+            if not os.path.isdir(domain_dir):
+                continue
+            for task_id in os.listdir(domain_dir):
+                task_dir = os.path.join(domain_dir, task_id)
+                if os.path.isdir(task_dir):
+                    domain_tasks.add(f"{domain}/{task_id}")
+        tasks_per_dir.append(domain_tasks)
+
+    # Step 2: find tasks common to all trajectories
+    common_tasks = set.intersection(*tasks_per_dir)
+
+    constant_tasks = []
+    variance_tasks = []
+    constant_tasks_scores = []
+    optimal_sum = 0.0
+    expected_value = 0.0
+
+    # Step 3: evaluate each common task
+    for domain_task in sorted(common_tasks):
+        domain, task_id = domain_task.split("/", 1)
+        results = []
+        for results_dir in results_dirs:
+            task_dir = os.path.join(results_dir, domain, task_id)
+            result_file = os.path.join(task_dir, "result.txt")
+            if os.path.isfile(result_file):
+                with open(result_file, "r") as f:
+                    try:
+                        val = float(f.read().strip())
+                        results.append(val)
+                    except ValueError:
+                        continue
+
+        if not results:  # skip if no valid results
+            logging.warning(f"No valid results for {domain_task}")
+            continue
+
+        # classification
+        if all(r == results[0] for r in results):
+            constant_tasks.append(domain_task)
+            constant_tasks_scores.append(results[0])
+        else:
+            variance_tasks.append(domain_task)
+
+        # accumulate min/optimal
+        # minimum_sum += min(results) #We incorrectly also counted the minimum sum of variance tasks, we should not do this
+        optimal_sum += max(results)
+        expected_value += sum(results) / len(results)
+
+    return {
+        "constant": constant_tasks,  # We dont evaluate constant tasks
+        "variance": variance_tasks,  # We evaluate variance tasks
+        "minimum": sum(
+            constant_tasks_scores
+        ),  # sum of constant tasks scores (easy + hard)
+        "optimal": optimal_sum,  # If we get the best score, we get the optimal score
+        "expected_value": expected_value,  # If we get the average score across all tasks for all trajectories, we get the expected value
+    }
+
+
+def check_selected_trajectory(results_dirs: [str], selected_trajectory: str, task: str):
+    """
+    results_dirs: list of directories in format results_dir/<domain>/<task_id>
+    selected_trajectory: the path of the selected trajectory
+    task: string in format "<domain>/<task_id>"
+
+    Returns (selected_val, optimal_val)
+    """
+    domain, task_id = task.split("/")
+    all_results = []
+
+    if not any(
+        os.path.commonpath([os.path.abspath(selected_trajectory), os.path.abspath(rd)])
+        == os.path.abspath(rd)
+        for rd in results_dirs
+    ):
+        return None, None
+
+    for rd in results_dirs:
+        result_file = os.path.join(rd, domain, task_id, "result.txt")
+        if os.path.isfile(result_file):
+            try:
+                all_results.append(float(open(result_file).read().strip()))
+            except ValueError:
+                pass
+
+    selected_file = os.path.join(selected_trajectory, domain, task_id, "result.txt")
+    if not os.path.isfile(selected_file):
+        return None, max(all_results) if all_results else None
+
+    try:
+        selected_val = float(open(selected_file).read().strip())
+    except ValueError:
+        return None, max(all_results) if all_results else None
+
+    optimal_val = max(all_results) if all_results else selected_val
+    return selected_val, optimal_val
+
+
+def evaluate_comparative_results(results_dirs: [str], json_path: str = None):
+    """
+    Opens comparative_judge_results.json (default) or a given path,
+    evaluates each task, and returns results.
+
+    Args:
+        results_dirs: list of result directories
+        json_path: optional path to comparative_judge_results.json
+
+    Returns:
+        dict mapping task -> {"selected_val": float or None, "optimal_val": float or None}
+    """
+    judge_score = 0
+    optimal_score = 0
+    if json_path is None:
+        json_path = "comparative_judge_results.json"
+
+    with open(json_path, "r") as f:
+        data = json.load(f)
+
+    results = {}
+    for task, info in data.items():
+        selected_trajectory = info.get("selected_trajectory")
+        if selected_trajectory:
+            selected_val, optimal_val = check_selected_trajectory(
+                results_dirs, selected_trajectory, task
+            )
+            if selected_val is not None and optimal_val is not None:
+                print(
+                    f"task: {task}, selected_val: {selected_val}, optimal_val: {optimal_val}"
+                )
+                judge_score += selected_val
+                optimal_score += optimal_val
+    return judge_score, optimal_score

osworld_setup/s3/lib_run_single.py

@@ -0,0 +1,90 @@
+import datetime
+import json
+import logging
+import os
+import time
+from typing import *
+from wrapt_timeout_decorator import *
+
+logger = logging.getLogger("desktopenv.experiment")
+
+
+def run_single_example(
+    agent, env, example, max_steps, instruction, args, example_result_dir, scores
+):
+    runtime_logger = setup_logger(example, example_result_dir)
+    try:
+        agent.reset(runtime_logger)
+    except Exception as e:
+        agent.reset()
+
+    env.reset(task_config=example)
+    time.sleep(60)  # Wait for the environment to be ready
+    obs = env._get_obs()  # Get the initial observation
+
+    with open(os.path.join(example_result_dir, f"step_0.png"), "wb") as _f:
+        _f.write(obs["screenshot"])
+
+    with open(
+        os.path.join(example_result_dir, "instruction.txt"), "w", encoding="utf-8"
+    ) as f:
+        f.write(instruction)
+
+    done = False
+    step_idx = 0
+    # env.controller.start_recording()
+    while not done and step_idx < max_steps:
+        response, actions = agent.predict(instruction, obs)
+        for action in actions:
+            action_timestamp = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
+            logger.info("Step %d: %s", step_idx + 1, action)
+            obs, reward, done, info = env.step(action, args.sleep_after_execution)
+
+            logger.info("Reward: %.2f", reward)
+            logger.info("Done: %s", done)
+            # Save screenshot and trajectory information
+            with open(
+                os.path.join(
+                    example_result_dir, f"step_{step_idx + 1}_{action_timestamp}.png"
+                ),
+                "wb",
+            ) as _f:
+                _f.write(obs["screenshot"])
+
+            response.update(
+                {
+                    "step_num": step_idx + 1,
+                    "action_timestamp": action_timestamp,
+                    "action": action,
+                    "reward": reward,
+                    "done": done,
+                    "info": info,
+                    "screenshot_file": f"step_{step_idx + 1}_{action_timestamp}.png",
+                }
+            )
+            with open(
+                os.path.join(example_result_dir, "traj.jsonl"), "a", encoding="utf-8"
+            ) as f:
+                f.write(json.dumps(response, ensure_ascii=False))
+                f.write("\n")
+            if done:
+                logger.info("The episode is done.")
+                break
+        step_idx += 1
+    result = env.evaluate()
+    logger.info("Result: %.2f", result)
+    scores.append(result)
+    with open(
+        os.path.join(example_result_dir, "result.txt"), "w", encoding="utf-8"
+    ) as f:
+        f.write(f"{result}\n")
+    # env.controller.end_recording(os.path.join(example_result_dir, "recording.mp4"))
+
+
+def setup_logger(example, example_result_dir):
+    runtime_logger = logging.getLogger(f"desktopenv.example.{example['id']}")
+    runtime_logger.setLevel(logging.DEBUG)
+    runtime_logger.addHandler(
+        logging.FileHandler(os.path.join(example_result_dir, "runtime.log"))
+    )
+    return runtime_logger

osworld_setup/s3/run.py

@@ -0,0 +1,578 @@
+"""OSWorld's run.py with AgentS2."""
+
+"""Script to run end-to-end evaluation on the benchmark.
+Utils and basic architecture credit to https://github.com/web-arena-x/webarena/blob/main/run.py.
+"""
+
+import argparse
+import datetime
+import json
+import logging
+import os
+import sys
+import signal
+import time
+from multiprocessing import Process, Manager, current_process, Queue
+
+
+import lib_run_single
+from desktop_env.desktop_env import DesktopEnv
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+#  Logger Configs {{{ #
+logger = logging.getLogger()
+logger.setLevel(logging.DEBUG)
+
+datetime_str: str = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
+
+stdout_handler = logging.StreamHandler(sys.stdout)
+
+stdout_handler.setLevel(logging.INFO)
+
+formatter = logging.Formatter(
+    fmt="\x1b[1;33m[%(asctime)s \x1b[31m%(levelname)s \x1b[32m%(module)s/%(lineno)d-%(processName)s\x1b[1;33m] \x1b[0m%(message)s"
+)
+
+stdout_handler.setFormatter(formatter)
+
+stdout_handler.addFilter(logging.Filter("desktopenv"))
+
+logger.addHandler(stdout_handler)
+#  }}} Logger Configs #
+
+logger = logging.getLogger("desktopenv.experiment")
+
+
+# Global variables for signal handling
+active_environments = []
+processes = []
+is_terminating = False
+
+
+def distribute_tasks(test_all_meta: dict) -> list:
+    all_tasks = []
+    for domain, examples in test_all_meta.items():
+        for example_id in examples:
+            all_tasks.append((domain, example_id))
+    return all_tasks
+
+
+def process_signal_handler(signum, frame, env_idx):
+    logger.info(f"Process {env_idx + 1} received signal {signum}. Shutting down...")
+    local_vars = frame.f_locals
+    active_environments = local_vars.get("active_environments", [])
+    for env in active_environments:
+        if env is not None:
+            try:
+                logger.info(f"Process {env_idx + 1} closing environment...")
+                env.close()
+                logger.info(f"Process {env_idx + 1} environment closed successfully")
+            except Exception as e:
+                logger.error(f"Process {env_idx + 1} error closing environment: {e}")
+    logger.info(f"Process {env_idx + 1} shutdown complete. Exiting.")
+    sys.exit(0)
+
+
+def run_env_tasks(
+    task_queue: Queue,
+    args: argparse.Namespace,
+    shared_scores: list,
+    engine_params,
+    engine_params_for_grounding,
+):
+    active_environments = []
+    env = None
+    try:
+        # Use IMAGE_ID_MAP for AWS provider to get snapshot_name
+        snapshot_name = None
+        region = getattr(args, "region", None)
+        if args.provider_name == "aws" and region is not None:
+            try:
+                from desktop_env.providers.aws.manager import IMAGE_ID_MAP
+
+                screen_size = (args.screen_width, args.screen_height)
+                snapshot_name = IMAGE_ID_MAP[region].get(
+                    screen_size, IMAGE_ID_MAP[region][(1920, 1080)]
+                )
+            except Exception as e:
+                logger.error(f"Failed to get snapshot_name from IMAGE_ID_MAP: {e}")
+                snapshot_name = None
+        from gui_agents.s3.agents.agent_s import AgentS3
+        from gui_agents.s3.agents.grounding import OSWorldACI
+
+        env = DesktopEnv(
+            path_to_vm=args.path_to_vm,
+            action_space=args.action_space,
+            provider_name=args.provider_name,
+            region=region,
+            snapshot_name=snapshot_name,
+            screen_size=(args.screen_width, args.screen_height),
+            headless=args.headless,
+            os_type="Ubuntu",
+            require_a11y_tree=args.observation_type
+            in ["a11y_tree", "screenshot_a11y_tree", "som"],
+            enable_proxy=True,
+            client_password=getattr(args, "client_password", ""),
+        )
+        grounding_agent = OSWorldACI(
+            env=env,
+            platform="linux",
+            engine_params_for_generation=engine_params,
+            engine_params_for_grounding=engine_params_for_grounding,
+            width=args.screen_width,
+            height=args.screen_height,
+        )
+        agent = AgentS3(
+            engine_params,
+            grounding_agent,
+            platform="linux",
+        )
+
+        active_environments.append(env)
+        logger.info(f"Process {current_process().name} started.")
+        while True:
+            try:
+                item = task_queue.get(timeout=5)
+            except Exception:
+                break
+            domain, example_id = item
+            try:
+                config_file = os.path.join(
+                    args.test_config_base_dir, f"examples/{domain}/{example_id}.json"
+                )
+                with open(config_file, "r", encoding="utf-8") as f:
+                    example = json.load(f)
+                instruction = example["instruction"]
+                example_result_dir = os.path.join(
+                    args.result_dir,
+                    args.action_space,
+                    args.observation_type,
+                    args.model,
+                    domain,
+                    example_id,
+                )
+                os.makedirs(example_result_dir, exist_ok=True)
+                logger.info(f"[{current_process().name}][Domain]: {domain}")
+                logger.info(f"[{current_process().name}][Example ID]: {example_id}")
+                logger.info(f"[{current_process().name}][Instruction]: {instruction}")
+                try:
+                    lib_run_single.run_single_example(
+                        agent,
+                        env,
+                        example,
+                        args.max_steps,
+                        instruction,
+                        args,
+                        example_result_dir,
+                        shared_scores,
+                    )
+                except Exception as e:
+                    import traceback
+
+                    logger.error(
+                        f"Exception in {current_process().name} {domain}/{example_id}: {e}"
+                    )
+                    logger.error(traceback.format_exc())
+                    try:
+                        env.controller.end_recording(
+                            os.path.join(example_result_dir, "recording.mp4")
+                        )
+                    except Exception as rec_e:
+                        logger.error(f"Failed to end recording: {rec_e}")
+                    with open(os.path.join(example_result_dir, "traj.jsonl"), "a") as f:
+                        f.write(json.dumps({"Error": f"{domain}/{example_id} - {e}"}))
+                        f.write("\n")
+            except Exception as e:
+                logger.error(f"Task-level error in {current_process().name}: {e}")
+                import traceback
+
+                logger.error(traceback.format_exc())
+    except Exception as e:
+        logger.error(f"Process-level error in {current_process().name}: {e}")
+        import traceback
+
+        logger.error(traceback.format_exc())
+    finally:
+        logger.info(f"{current_process().name} cleaning up environment...")
+        try:
+            if env:
+                env.close()
+                logger.info(f"{current_process().name} environment closed successfully")
+        except Exception as e:
+            logger.error(
+                f"{current_process().name} error during environment cleanup: {e}"
+            )
+
+
+def signal_handler(signum, frame):
+    global is_terminating, active_environments, processes
+    if is_terminating:
+        return
+    is_terminating = True
+    logger.info(f"Received signal {signum}. Gracefully shutting down...")
+    for env in active_environments:
+        try:
+            logger.info(f"Closing environment...")
+            env.close()
+            logger.info(f"Environment closed successfully")
+        except Exception as e:
+            logger.error(f"Error closing environment: {e}")
+    for p in processes:
+        if p.is_alive():
+            try:
+                logger.info(f"Sending termination signal to process {p.name}...")
+                p.terminate()
+            except Exception as e:
+                logger.error(f"Error sending termination signal to process: {e}")
+    time.sleep(1)
+    for p in processes:
+        if p.is_alive():
+            try:
+                logger.info(f"Forcefully terminating process {p.name}...")
+                import signal as sig
+
+                os.kill(p.pid, sig.SIGKILL)
+            except Exception as e:
+                logger.error(f"Error forcefully terminating process: {e}")
+    logger.info("Shutdown complete. Exiting.")
+    sys.exit(0)
+
+
+def config() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run end-to-end evaluation on the benchmark"
+    )
+
+    # environment config
+    parser.add_argument("--path_to_vm", type=str, default=None)
+    parser.add_argument(
+        "--provider_name",
+        type=str,
+        default="vmware",
+        help="Virtualization provider (vmware, docker, aws, azure, gcp, virtualbox)",
+    )
+    parser.add_argument(
+        "--headless", action="store_true", help="Run in headless machine"
+    )
+    parser.add_argument(
+        "--action_space", type=str, default="pyautogui", help="Action type"
+    )
+    parser.add_argument(
+        "--observation_type",
+        choices=["screenshot", "a11y_tree", "screenshot_a11y_tree", "som"],
+        default="screenshot",
+        help="Observation type",
+    )
+    parser.add_argument(
+        "--num_envs",
+        type=int,
+        default=1,
+        help="Number of environments to run in parallel",
+    )
+    parser.add_argument("--screen_width", type=int, default=1920)
+    parser.add_argument("--screen_height", type=int, default=1080)
+    parser.add_argument("--sleep_after_execution", type=float, default=1.0)
+    parser.add_argument("--max_steps", type=int, default=15)
+
+    parser.add_argument("--domain", type=str, default="all")
+    parser.add_argument(
+        "--test_all_meta_path", type=str, default="evaluation_examples/test_all.json"
+    )
+    parser.add_argument(
+        "--test_config_base_dir", type=str, default="evaluation_examples"
+    )
+    parser.add_argument("--result_dir", type=str, default="./results")
+
+    parser.add_argument(
+        "--region", type=str, default="us-east-1", help="AWS region for the VM"
+    )
+    parser.add_argument(
+        "--client_password", type=str, default="", help="Client password"
+    )
+
+    # agent config
+    parser.add_argument("--max_trajectory_length", type=int, default=8)
+
+    # lm config
+    parser.add_argument("--model_provider", type=str, default="openai")
+    parser.add_argument("--model", type=str, default="gpt-4o")
+    parser.add_argument(
+        "--model_url",
+        type=str,
+        default="",
+        help="The URL of the main generation model API.",
+    )
+    parser.add_argument(
+        "--model_api_key",
+        type=str,
+        default="",
+        help="The API key of the main generation model.",
+    )
+    parser.add_argument(
+        "--model_temperature",
+        type=float,
+        default=None,
+        help="Temperature to fix the generation model at (e.g. o3 can only be run with 1.0)",
+    )
+
+    # grounding model config
+    parser.add_argument(
+        "--ground_provider",
+        type=str,
+        required=True,
+        help="The provider for the grounding model",
+    )
+    parser.add_argument(
+        "--ground_url", type=str, required=True, help="The URL of the grounding model"
+    )
+    parser.add_argument(
+        "--ground_api_key",
+        type=str,
+        default="",
+        help="The API key of the grounding model.",
+    )
+    parser.add_argument(
+        "--ground_model",
+        type=str,
+        required=True,
+        help="The model name for the grounding model",
+    )
+    parser.add_argument(
+        "--grounding_width",
+        type=int,
+        required=True,
+        help="Width of screenshot image after processor rescaling",
+    )
+    parser.add_argument(
+        "--grounding_height",
+        type=int,
+        required=True,
+        help="Height of screenshot image after processor rescaling",
+    )
+
+    args = parser.parse_args()
+
+    return args
+
+
+def test(args: argparse.Namespace, test_all_meta: dict) -> None:
+    global processes
+    logger.info("Args: %s", args)
+    all_tasks = distribute_tasks(test_all_meta)
+    logger.info(f"Total tasks: {len(all_tasks)}")
+
+    engine_params = {
+        "engine_type": args.model_provider,
+        "model": args.model,
+        "base_url": getattr(args, "model_url", ""),
+        "api_key": getattr(args, "model_api_key", ""),
+        "temperature": getattr(args, "model_temperature", None),
+    }
+    engine_params_for_grounding = {
+        "engine_type": args.ground_provider,
+        "model": args.ground_model,
+        "base_url": getattr(args, "ground_url", ""),
+        "api_key": getattr(args, "ground_api_key", ""),
+        "grounding_width": args.grounding_width,
+        "grounding_height": args.grounding_height,
+    }
+
+    with Manager() as manager:
+        shared_scores = manager.list()
+        task_queue = manager.Queue()
+        for item in all_tasks:
+            task_queue.put(item)
+        num_envs = args.num_envs
+        processes = []
+        for i in range(num_envs):
+            p = Process(
+                target=run_env_tasks,
+                args=(
+                    task_queue,
+                    args,
+                    shared_scores,
+                    engine_params,
+                    engine_params_for_grounding,
+                ),
+                name=f"EnvProcess-{i+1}",
+            )
+            p.daemon = True
+            p.start()
+            processes.append(p)
+            logger.info(f"Started process {p.name} with PID {p.pid}")
+        try:
+            while True:
+                alive_count = 0
+                for idx, p in enumerate(processes):
+                    if not p.is_alive():
+                        logger.warning(f"Process {p.name} died, restarting...")
+                        new_p = Process(
+                            target=run_env_tasks,
+                            args=(
+                                task_queue,
+                                args,
+                                shared_scores,
+                                engine_params,
+                                engine_params_for_grounding,
+                            ),
+                            name=f"EnvProcess-Restart-{idx+1}",
+                        )
+                        new_p.daemon = True
+                        new_p.start()
+                        processes[idx] = new_p
+                        logger.info(
+                            f"Restarted process {new_p.name} with PID {new_p.pid}"
+                        )
+                    else:
+                        alive_count += 1
+                if task_queue.empty():
+                    logger.info("All tasks finished.")
+                    break
+                if alive_count == 0:
+                    logger.error("All processes died, exiting.")
+                    break
+                time.sleep(5)
+            for p in processes:
+                p.join()
+        except KeyboardInterrupt:
+            logger.info(
+                "Main process received KeyboardInterrupt. Initiating graceful shutdown..."
+            )
+            raise
+        except Exception as e:
+            logger.error(
+                f"Unexpected error while waiting for processes: {e}", exc_info=True
+            )
+            for p in processes:
+                if p.is_alive():
+                    try:
+                        logger.info(f"Terminating process {p.name} due to error...")
+                        p.terminate()
+                    except Exception as term_e:
+                        logger.error(f"Error terminating process {p.name}: {term_e}")
+            raise
+        scores = list(shared_scores)
+    logger.info(f"Average score: {sum(scores) / len(scores) if scores else 0}")
+
+
+def get_unfinished(
+    action_space, use_model, observation_type, result_dir, total_file_json
+):
+    target_dir = os.path.join(result_dir, action_space, observation_type, use_model)
+
+    if not os.path.exists(target_dir):
+        return total_file_json
+
+    finished = {}
+    for domain in os.listdir(target_dir):
+        finished[domain] = []
+        domain_path = os.path.join(target_dir, domain)
+        if os.path.isdir(domain_path):
+            for example_id in os.listdir(domain_path):
+                if example_id == "onboard":
+                    continue
+                example_path = os.path.join(domain_path, example_id)
+                if os.path.isdir(example_path):
+                    if "result.txt" not in os.listdir(example_path):
+                        # empty all files under example_id
+                        for file in os.listdir(example_path):
+                            os.remove(os.path.join(example_path, file))
+                    else:
+                        finished[domain].append(example_id)
+
+    if not finished:
+        return total_file_json
+
+    for domain, examples in finished.items():
+        if domain in total_file_json:
+            total_file_json[domain] = [
+                x for x in total_file_json[domain] if x not in examples
+            ]
+
+    return total_file_json
+
+
+def get_result(action_space, use_model, observation_type, result_dir, total_file_json):
+    target_dir = os.path.join(result_dir, action_space, observation_type, use_model)
+    if not os.path.exists(target_dir):
+        print("New experiment, no result yet.")
+        return None
+
+    all_result = []
+
+    for domain in os.listdir(target_dir):
+        domain_path = os.path.join(target_dir, domain)
+        if os.path.isdir(domain_path):
+            for example_id in os.listdir(domain_path):
+                example_path = os.path.join(domain_path, example_id)
+                if os.path.isdir(example_path):
+                    if "result.txt" in os.listdir(example_path):
+                        # empty all files under example_id
+                        try:
+                            all_result.append(
+                                float(
+                                    open(
+                                        os.path.join(example_path, "result.txt"), "r"
+                                    ).read()
+                                )
+                            )
+                        except:
+                            all_result.append(0.0)
+
+    if not all_result:
+        print("New experiment, no result yet.")
+        return None
+    else:
+        print("Current Success Rate:", sum(all_result) / len(all_result) * 100, "%")
+        return all_result
+
+
+if __name__ == "__main__":
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
+    ####### The complete version of the list of examples #######
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+    args = config()
+
+    # save args to json in result_dir/action_space/observation_type/model/args.json
+    path_to_args = os.path.join(
+        args.result_dir,
+        args.action_space,
+        args.observation_type,
+        args.model,
+        "args.json",
+    )
+    os.makedirs(os.path.dirname(path_to_args), exist_ok=True)
+    with open(path_to_args, "w", encoding="utf-8") as f:
+        json.dump(vars(args), f, indent=4)
+
+    with open(args.test_all_meta_path, "r", encoding="utf-8") as f:
+        test_all_meta = json.load(f)
+
+    if args.domain != "all":
+        test_all_meta = {args.domain: test_all_meta[args.domain]}
+
+    test_file_list = get_unfinished(
+        args.action_space,
+        args.model,
+        args.observation_type,
+        args.result_dir,
+        test_all_meta,
+    )
+    left_info = ""
+    for domain in test_file_list:
+        left_info += f"{domain}: {len(test_file_list[domain])}\n"
+    logger.info(f"Left tasks:\n{left_info}")
+
+    get_result(
+        args.action_space,
+        args.model,
+        args.observation_type,
+        args.result_dir,
+        test_all_meta,
+    )
+    test(args, test_file_list)

osworld_setup/s3/run.sh

@@ -0,0 +1,54 @@
+# Step 1: Complete 2 or more rollouts on either AWS or locally
+python run.py \
+  --provider_name "aws" \
+  --headless \
+  --num_envs 10 \
+  --max_steps 100 \
+  --domain "all" \
+  --test_all_meta_path evaluation_examples/test_nogdrive.json \
+  --result_dir "results" \
+  --region "us-east-1" \
+  --model_provider "openai" \
+  --model "gpt-5-2025-08-07" \
+  --model_temperature 1.0 \
+  --ground_provider "huggingface" \
+  --ground_url "<YOUR_HUGGINGFACE_ENDPOINT_URL>/v1" \
+  --grounding_width 1920 \
+  --grounding_height 1080 \
+  --sleep_after_execution 3
+
+python run_local.py \
+  --path_to_vm "/Users/user/OSWorld/vmware_vm_data/Ubuntu0/Ubuntu0.vmx" \
+  --provider_name "vmware" \
+  --headless \
+  --max_steps 100 \
+  --domain "all" \
+  --test_all_meta_path evaluation_examples/test_nogdrive.json \
+  --result_dir "results" \
+  --model_provider "openai" \
+  --model "gpt-5-2025-08-07" \
+  --model_temperature 1.0 \
+  --ground_provider "huggingface" \
+  --ground_url "<YOUR_HUGGINGFACE_ENDPOINT_URL>/v1" \
+  --grounding_width 1920 \
+  --grounding_height 1080
+
+# Step 2: Generate Facts
+python generate_facts.py \
+  --results-dirs \
+    results1/pyautogui/screenshot/gpt-5-2025-08-07 \
+    results2/pyautogui/screenshot/gpt-5-2025-08-07 \
+  --model "gpt-5-2025-08-07" \
+  --engine-type "openai" \
+  --temperature 1.0
+
+# Step 3: Run the Judge. Make sure the order of the results-dirs is the same as the order above.
+python run_judge.py \
+  --results-dirs \
+    results1/pyautogui/screenshot/gpt-5-2025-08-07 \
+    results2/pyautogui/screenshot/gpt-5-2025-08-07 \
+  --output-dir "judge_results" \
+  --examples-path "evaluation_examples/examples" \
+  --model "gpt-5-2025-08-07" \
+  --engine-type "openai" \
+  --temperature 1.0

osworld_setup/s3/run_local.py

@@ -0,0 +1,417 @@
+"""Script to run end-to-end evaluation on the benchmark.
+Utils and basic architecture credit to https://github.com/web-arena-x/webarena/blob/main/run.py.
+"""
+
+import argparse
+import datetime
+import json
+import logging
+import os
+import sys
+
+from tqdm import tqdm
+
+import lib_run_single
+from desktop_env.desktop_env import DesktopEnv
+from gui_agents.s3.agents.agent_s import AgentS3
+from gui_agents.s3.agents.grounding import OSWorldACI
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# Almost deprecated since it's not multi-env, use run_multienv_*.py instead
+
+#  Logger Configs {{{ #
+logger = logging.getLogger()
+logger.setLevel(logging.DEBUG)
+
+datetime_str: str = datetime.datetime.now().strftime("%Y%m%d@%H%M%S")
+
+file_handler = logging.FileHandler(
+    os.path.join("logs", "normal-{:}.log".format(datetime_str)), encoding="utf-8"
+)
+debug_handler = logging.FileHandler(
+    os.path.join("logs", "debug-{:}.log".format(datetime_str)), encoding="utf-8"
+)
+stdout_handler = logging.StreamHandler(sys.stdout)
+sdebug_handler = logging.FileHandler(
+    os.path.join("logs", "sdebug-{:}.log".format(datetime_str)), encoding="utf-8"
+)
+
+file_handler.setLevel(logging.INFO)
+debug_handler.setLevel(logging.DEBUG)
+stdout_handler.setLevel(logging.INFO)
+sdebug_handler.setLevel(logging.DEBUG)
+
+formatter = logging.Formatter(
+    fmt="\x1b[1;33m[%(asctime)s \x1b[31m%(levelname)s \x1b[32m%(module)s/%(lineno)d-%(processName)s\x1b[1;33m] \x1b[0m%(message)s"
+)
+file_handler.setFormatter(formatter)
+debug_handler.setFormatter(formatter)
+stdout_handler.setFormatter(formatter)
+sdebug_handler.setFormatter(formatter)
+
+stdout_handler.addFilter(logging.Filter("desktopenv"))
+sdebug_handler.addFilter(logging.Filter("desktopenv"))
+
+logger.addHandler(file_handler)
+logger.addHandler(debug_handler)
+logger.addHandler(stdout_handler)
+logger.addHandler(sdebug_handler)
+#  }}} Logger Configs #
+
+logger = logging.getLogger("desktopenv.experiment")
+
+
+def config() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run end-to-end evaluation on the benchmark"
+    )
+
+    # environment config
+    parser.add_argument("--path_to_vm", type=str, default=None)
+    parser.add_argument(
+        "--provider_name",
+        type=str,
+        default="vmware",
+        help="Virtualization provider (vmware, docker, aws, azure, gcp, virtualbox)",
+    )
+    parser.add_argument(
+        "--headless", action="store_true", help="Run in headless machine"
+    )
+    parser.add_argument(
+        "--action_space", type=str, default="pyautogui", help="Action type"
+    )
+    parser.add_argument(
+        "--observation_type",
+        choices=["screenshot", "a11y_tree", "screenshot_a11y_tree", "som"],
+        default="screenshot",
+        help="Observation type",
+    )
+    parser.add_argument("--screen_width", type=int, default=1920)
+    parser.add_argument("--screen_height", type=int, default=1080)
+    parser.add_argument("--sleep_after_execution", type=float, default=3.0)
+    parser.add_argument("--max_steps", type=int, default=15)
+
+    # agent config
+    parser.add_argument("--max_trajectory_length", type=int, default=3)
+    parser.add_argument(
+        "--test_config_base_dir", type=str, default="evaluation_examples"
+    )
+
+    # lm config
+    parser.add_argument("--model", type=str, default="gpt-4o")
+    parser.add_argument("--temperature", type=float, default=1.0)
+
+    # AgentS2 specific config
+    parser.add_argument("--model_provider", type=str, default="openai")
+    parser.add_argument(
+        "--model_url",
+        type=str,
+        default="",
+        help="The URL of the main generation model API.",
+    )
+    parser.add_argument(
+        "--model_api_key",
+        type=str,
+        default="",
+        help="The API key of the main generation model.",
+    )
+    parser.add_argument(
+        "--model_temperature",
+        type=float,
+        default=None,
+        help="Temperature to fix the generation model at (e.g. o3 can only be run with 1.0)",
+    )
+
+    # grounding model config
+    parser.add_argument(
+        "--ground_provider",
+        type=str,
+        required=True,
+        help="The provider for the grounding model",
+    )
+    parser.add_argument(
+        "--ground_url", type=str, required=True, help="The URL of the grounding model"
+    )
+    parser.add_argument(
+        "--ground_api_key",
+        type=str,
+        default="",
+        help="The API key of the grounding model.",
+    )
+    parser.add_argument(
+        "--ground_model",
+        type=str,
+        required=True,
+        help="The model name for the grounding model",
+    )
+    parser.add_argument(
+        "--grounding_width",
+        type=int,
+        required=True,
+        help="Width of screenshot image after processor rescaling",
+    )
+    parser.add_argument(
+        "--grounding_height",
+        type=int,
+        required=True,
+        help="Height of screenshot image after processor rescaling",
+    )
+
+    # example config
+    parser.add_argument("--domain", type=str, default="all")
+    parser.add_argument(
+        "--test_all_meta_path", type=str, default="evaluation_examples/test_all.json"
+    )
+
+    # logging related
+    parser.add_argument("--result_dir", type=str, default="./results")
+    args = parser.parse_args()
+
+    return args
+
+
+def test(args: argparse.Namespace, test_all_meta: dict) -> None:
+    scores = []
+    max_steps = args.max_steps
+
+    # log args
+    logger.info("Args: %s", args)
+    # set wandb project
+    cfg_args = {
+        "path_to_vm": args.path_to_vm,
+        "provider_name": args.provider_name,
+        "headless": args.headless,
+        "action_space": args.action_space,
+        "observation_type": args.observation_type,
+        "screen_width": args.screen_width,
+        "screen_height": args.screen_height,
+        "sleep_after_execution": args.sleep_after_execution,
+        "max_steps": args.max_steps,
+        "max_trajectory_length": args.max_trajectory_length,
+        "model": args.model,
+        "temperature": args.temperature,
+        "result_dir": args.result_dir,
+    }
+
+    # AgentS2 configuration
+    engine_params = {
+        "engine_type": args.model_provider,
+        "model": args.model,
+        "base_url": getattr(args, "model_url", ""),
+        "api_key": getattr(args, "model_api_key", ""),
+        "temperature": getattr(args, "model_temperature", None),
+    }
+    engine_params_for_grounding = {
+        "engine_type": args.ground_provider,
+        "model": args.ground_model,
+        "base_url": getattr(args, "ground_url", ""),
+        "api_key": getattr(args, "ground_api_key", ""),
+        "grounding_width": args.grounding_width,
+        "grounding_height": args.grounding_height,
+    }
+
+    env = DesktopEnv(
+        provider_name=args.provider_name,
+        path_to_vm=args.path_to_vm,
+        action_space=args.action_space,
+        screen_size=(args.screen_width, args.screen_height),
+        headless=args.headless,
+        os_type="Ubuntu",
+        require_a11y_tree=args.observation_type
+        in ["a11y_tree", "screenshot_a11y_tree", "som"],
+        enable_proxy=True,
+    )
+
+    grounding_agent = OSWorldACI(
+        env=env,
+        platform="linux",
+        engine_params_for_generation=engine_params,
+        engine_params_for_grounding=engine_params_for_grounding,
+        width=args.screen_width,
+        height=args.screen_height,
+    )
+    agent = AgentS3(
+        engine_params,
+        grounding_agent,
+        platform="linux",
+    )
+
+    for domain in tqdm(test_all_meta, desc="Domain"):
+        for example_id in tqdm(test_all_meta[domain], desc="Example", leave=False):
+            config_file = os.path.join(
+                args.test_config_base_dir, f"examples/{domain}/{example_id}.json"
+            )
+            with open(config_file, "r", encoding="utf-8") as f:
+                example = json.load(f)
+
+            logger.info(f"[Domain]: {domain}")
+            logger.info(f"[Example ID]: {example_id}")
+
+            instruction = example["instruction"]
+
+            logger.info(f"[Instruction]: {instruction}")
+            # wandb each example config settings
+            cfg_args["instruction"] = instruction
+            cfg_args["start_time"] = datetime.datetime.now().strftime(
+                "%Y:%m:%d-%H:%M:%S"
+            )
+            # run.config.update(cfg_args)
+
+            example_result_dir = os.path.join(
+                args.result_dir,
+                args.action_space,
+                args.observation_type,
+                args.model,
+                domain,
+                example_id,
+            )
+            os.makedirs(example_result_dir, exist_ok=True)
+            # example start running
+            try:
+                lib_run_single.run_single_example(
+                    agent,
+                    env,
+                    example,
+                    max_steps,
+                    instruction,
+                    args,
+                    example_result_dir,
+                    scores,
+                )
+            except Exception as e:
+                logger.error(f"Exception in {domain}/{example_id}: {e}")
+                # Only attempt to end recording if controller exists (not Docker provider)
+                if hasattr(env, "controller") and env.controller is not None:
+                    env.controller.end_recording(
+                        os.path.join(example_result_dir, "recording.mp4")
+                    )
+                with open(os.path.join(example_result_dir, "traj.jsonl"), "a") as f:
+                    f.write(
+                        json.dumps(
+                            {"Error": f"Time limit exceeded in {domain}/{example_id}"}
+                        )
+                    )
+                    f.write("\n")
+
+    env.close()
+    logger.info(f"Average score: {sum(scores) / len(scores)}")
+
+
+def get_unfinished(
+    action_space, use_model, observation_type, result_dir, total_file_json
+):
+    target_dir = os.path.join(result_dir, action_space, observation_type, use_model)
+
+    if not os.path.exists(target_dir):
+        return total_file_json
+
+    finished = {}
+    for domain in os.listdir(target_dir):
+        finished[domain] = []
+        domain_path = os.path.join(target_dir, domain)
+        if os.path.isdir(domain_path):
+            for example_id in os.listdir(domain_path):
+                if example_id == "onboard":
+                    continue
+                example_path = os.path.join(domain_path, example_id)
+                if os.path.isdir(example_path):
+                    if "result.txt" not in os.listdir(example_path):
+                        # empty all files under example_id
+                        for file in os.listdir(example_path):
+                            os.remove(os.path.join(example_path, file))
+                    else:
+                        finished[domain].append(example_id)
+
+    if not finished:
+        return total_file_json
+
+    for domain, examples in finished.items():
+        if domain in total_file_json:
+            total_file_json[domain] = [
+                x for x in total_file_json[domain] if x not in examples
+            ]
+
+    return total_file_json
+
+
+def get_result(action_space, use_model, observation_type, result_dir, total_file_json):
+    target_dir = os.path.join(result_dir, action_space, observation_type, use_model)
+    if not os.path.exists(target_dir):
+        print("New experiment, no result yet.")
+        return None
+
+    all_result = []
+
+    for domain in os.listdir(target_dir):
+        domain_path = os.path.join(target_dir, domain)
+        if os.path.isdir(domain_path):
+            for example_id in os.listdir(domain_path):
+                example_path = os.path.join(domain_path, example_id)
+                if os.path.isdir(example_path):
+                    if "result.txt" in os.listdir(example_path):
+                        # empty all files under example_id
+                        try:
+                            all_result.append(
+                                float(
+                                    open(
+                                        os.path.join(example_path, "result.txt"), "r"
+                                    ).read()
+                                )
+                            )
+                        except:
+                            all_result.append(0.0)
+
+    if not all_result:
+        print("New experiment, no result yet.")
+        return None
+    else:
+        print("Current Success Rate:", sum(all_result) / len(all_result) * 100, "%")
+        return all_result
+
+
+if __name__ == "__main__":
+    ####### The complete version of the list of examples #######
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+    args = config()
+
+    # save args to json in result_dir/action_space/observation_type/model/args.json
+    path_to_args = os.path.join(
+        args.result_dir,
+        args.action_space,
+        args.observation_type,
+        args.model,
+        "args.json",
+    )
+    os.makedirs(os.path.dirname(path_to_args), exist_ok=True)
+    with open(path_to_args, "w", encoding="utf-8") as f:
+        json.dump(vars(args), f, indent=4)
+
+    with open(args.test_all_meta_path, "r", encoding="utf-8") as f:
+        test_all_meta = json.load(f)
+
+    if args.domain != "all":
+        test_all_meta = {args.domain: test_all_meta[args.domain]}
+
+    test_file_list = get_unfinished(
+        args.action_space,
+        args.model,
+        args.observation_type,
+        args.result_dir,
+        test_all_meta,
+    )
+    left_info = ""
+    for domain in test_file_list:
+        left_info += f"{domain}: {len(test_file_list[domain])}\n"
+    logger.info(f"Left tasks:\n{left_info}")
+
+    get_result(
+        args.action_space,
+        args.model,
+        args.observation_type,
+        args.result_dir,
+        test_all_meta,
+    )
+    test(args, test_file_list)