update (#2)

- update (53214c0902c5bc6a66e7b9dd3ec470450acacc84)


Co-authored-by: iruno <[email protected]>

.gitignore

@@ -4,4 +4,9 @@
 *.gif
 *.bmp
 *.tiff
-*.ico
+*.ico
+*.log
+.gradio/
+__pycache__/
+.env
+.venv/

BrowserGym/browsergym/assistantbench/src/browsergym/assistantbench/task.py

@@ -108,7 +108,7 @@ class AssistantBenchTask(AbstractBrowserTask):
 
     def setup(self, page: Page) -> Tuple[str, dict]:
         logger.info(f"Navigating to start url: {self.start_url}")
-        page.goto(self.start_url, timeout=10000)
+        page.goto(self.start_url, timeout=50000)
         if self.save_predictions and self.output_file:
             # create an empty task entry in the output file (will raise an Exception if the entry is already there)
             add_prediction_to_jsonl(

BrowserGym/browsergym/browsergym.egg-info/PKG-INFO

@@ -0,0 +1,22 @@
+Metadata-Version: 2.4
+Name: browsergym
+Version: 0.13.4
+Summary: BrowserGym: a gym environment for web task automation in the Chromium browser
+Author: Rim Assouel, Léo Boisvert, Massimo Caccia, Alex Drouin, Maxime Gasse, Imene Kerboua, Alex Lacoste, Thibault Le Sellier De Chezelles, Tom Marty, Aman Jaiswal
+License: Apache-2.0
+Classifier: Development Status :: 3 - Alpha
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: License :: OSI Approved :: Apache Software License
+Requires-Python: >3.10
+Description-Content-Type: text/markdown
+Requires-Dist: browsergym-core==0.13.4
+Requires-Dist: browsergym-miniwob==0.13.4
+Requires-Dist: browsergym-webarena==0.13.4
+Requires-Dist: browsergym-visualwebarena==0.13.4
+Requires-Dist: browsergym-assistantbench==0.13.4
+Requires-Dist: browsergym-experiments==0.13.4
+Requires-Dist: browsergym-workarena>=0.4.1
+Requires-Dist: weblinx-browsergym>=0.0.2

BrowserGym/browsergym/browsergym.egg-info/SOURCES.txt

@@ -0,0 +1,6 @@
+pyproject.toml
+browsergym.egg-info/PKG-INFO
+browsergym.egg-info/SOURCES.txt
+browsergym.egg-info/dependency_links.txt
+browsergym.egg-info/requires.txt
+browsergym.egg-info/top_level.txt

BrowserGym/browsergym/browsergym.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

BrowserGym/browsergym/browsergym.egg-info/requires.txt

@@ -0,0 +1,8 @@
+browsergym-core==0.13.4
+browsergym-miniwob==0.13.4
+browsergym-webarena==0.13.4
+browsergym-visualwebarena==0.13.4
+browsergym-assistantbench==0.13.4
+browsergym-experiments==0.13.4
+browsergym-workarena>=0.4.1
+weblinx-browsergym>=0.0.2

BrowserGym/browsergym/browsergym.egg-info/top_level.txt

@@ -0,0 +1 @@
+

BrowserGym/browsergym/core/src/browsergym/core/env.py

@@ -27,6 +27,7 @@ from .observation import (
 )
 from .spaces import AnyBox, AnyDict, Float, Unicode
 from .task import AbstractBrowserTask
+from ..utils.obs import overlay_som, flatten_axtree_to_str
 
 logger = logging.getLogger(__name__)
 
@@ -602,6 +603,8 @@ document.addEventListener("visibilitychange", () => {
         _post_extract(self.page)
 
         # obs is generic to all tasks
+        screenshot_np_array = extract_screenshot(self.page)
+        som_screenshot_np_array = overlay_som(screenshot_np_array, extra_properties)
         obs = {
             "chat_messages": tuple(copy.deepcopy(self.chat.messages)),
             "goal": _try_to_extract_legacy_goal(self.goal_object),  # legacy goal, deprecated
@@ -612,7 +615,7 @@ document.addEventListener("visibilitychange", () => {
             "open_pages_titles": tuple(page.title() for page in self.context.pages),
             "active_page_index": np.asarray([self.context.pages.index(self.page)]),
             "url": self.page.url,  # redundant with "open_pages_urls" and "active_page_index"
-            "screenshot": extract_screenshot(self.page),
+            "som_screenshot": som_screenshot_np_array,
             "dom_object": dom,
             "axtree_object": axtree,
             "extra_element_properties": extra_properties,

BrowserGym/browsergym/core/src/browsergym/core/task.py

@@ -92,7 +92,7 @@ class OpenEndedTask(AbstractBrowserTask):
         self.goal = goal
 
     def setup(self, page: playwright.sync_api.Page) -> tuple[str, dict]:
-        page.goto(self.start_url, timeout=10000)
+        page.goto(self.start_url, timeout=50000)
         return self.goal, {}
 
     def teardown(self) -> None:

BrowserGym/browsergym/visualwebarena/src/browsergym/visualwebarena/task.py

@@ -109,7 +109,7 @@ class GenericVisualWebArenaTask(AbstractBrowserTask):
         # task properties, will be used to set up the browsergym environment
         self.viewport = {"width": 1280, "height": 720}
         self.slow_mo = 1000  # ms
-        self.timeout = 10000  # ms
+        self.timeout = 50000  # ms
 
         self.webarena_instance = VisualWebArenaInstance()
         self.config_file: str = None

BrowserGym/browsergym/webarena/src/browsergym/webarena/task.py

@@ -34,7 +34,7 @@ class GenericWebArenaTask(AbstractBrowserTask):
         # task properties, will be used to set up the browsergym environment
         self.viewport = {"width": 1280, "height": 720}
         self.slow_mo = 1000  # ms
-        self.timeout = 10000  # ms
+        self.timeout = 50000  # ms
 
         self.webarena_instance = WebArenaInstance()
         self.config_file: str = None

Dockerfile

@@ -56,6 +56,11 @@ RUN curl -fsSL https://dl.google.com/linux/linux_signing_key.pub | gpg --dearmor
 # Set up working directory
 WORKDIR /app
 
+COPY templates/ templates/
+COPY browser_agent.py .
+COPY process_run.py .
+COPY utils.py .
+
 # Copy requirements and install Python dependencies
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
@@ -94,6 +99,8 @@ ENV RESOLUTION_HEIGHT=1080
 # COPY supervisord.conf /etc/supervisor/conf.d/supervisord.conf
 
 # EXPOSE 7788 6080 5900
+EXPOSE 7860
+ENV GRADIO_SERVER_NAME="0.0.0.0"
 
 # CMD ["/usr/bin/supervisord", "-c", "/etc/supervisor/conf.d/supervisord.conf"]
-# RUN python3 app.py
+CMD ["python", "app.py"]

agent/checklist.py

@@ -0,0 +1,18 @@
+from .mini_bench.agent import ChecklistGenerationAgent
+
+def generate_checklist(**data):
+    # data: 'intent', 'start_url', 'text_observation'
+    agent_config = {
+        'model_name': 'WPRM/qwen-3b-ar-reward-cot-mtl-checklist-enhanced',
+        'base_url': 'http://165.132.144.84:7701/v1',
+        'api_key': 'empty',
+        'temperature': 0.7,
+        'use_log_probs': True,
+        'use_checklist': True,
+        'use_multimodal': False,
+        'num_generate': 1,
+    }
+    checklist_generation_agent = ChecklistGenerationAgent(agent_config)
+    response_list, cost = checklist_generation_agent.generate_response(data, prompt_type='ours', constraint_str_list=["<think>", "</think>", "<answer>", "</answer>"])
+    response = response_list[0]
+    return response.split("<answer>")[-1].split("</answer>")[0].strip()

agent/mini_bench/agent.py

@@ -0,0 +1,467 @@
+from abc import ABC, abstractmethod
+import time
+import requests
+import json
+import math
+from langsmith import Client
+from langchain_openai import ChatOpenAI
+
+from .prompts import get_messages
+from .prompts.judge_prompt import (
+    JUDGE_OURS_BT_MODELING_PROMPT_TEMPLATE,
+    JUDGE_OURS_BT_MODELING_WO_CHECKLIST_PROMPT_TEMPLATE,
+    JUDGE_OURS_BT_MODELING_MULTIMODAL_PROMPT_TEMPLATE,
+    JUDGE_OURS_BT_MODELING_MULTIMODAL_WO_CHECKLIST_PROMPT_TEMPLATE
+)
+from .prompts.image_utils import image_to_base64_url
+
+MAX_RETRY = 3
+RETRY_SLEEP = 5
+MODEL_COST_MAPPING = {
+    "gpt-4o-mini": {
+        "input_token_cost": 0.15,
+        "output_token_cost": 0.6
+    },
+    "gpt-4o": {
+        "input_token_cost": 2.5,
+        "output_token_cost": 10
+    },
+}
+
+
+class Agent(ABC):
+    @abstractmethod
+    def generate_response(self, inputs: dict) -> str:
+        pass
+
+class BaseAgent(Agent):
+    def __init__(self, agent_config: dict):
+        self.agent_config = agent_config
+        self._setup()
+    
+    def _setup(self):
+        use_log_probs = self.agent_config.get("use_log_probs", False)
+        if use_log_probs:
+            self.llm = ChatOpenAI(
+                model=self.agent_config["model_name"], 
+                base_url=self.agent_config["base_url"], 
+                api_key=self.agent_config["api_key"], 
+                temperature=self.agent_config["temperature"],
+                timeout=300,
+                logprobs=True,
+                top_logprobs=10
+            )
+        else:
+            self.llm = ChatOpenAI(
+                model=self.agent_config["model_name"], 
+                base_url=self.agent_config["base_url"], 
+                api_key=self.agent_config["api_key"], 
+                temperature=self.agent_config["temperature"],
+                timeout=300
+            )
+        self.temperature = self.agent_config["temperature"]
+        self.num_generate = self.agent_config["num_generate"]
+        self.use_checklist = self.agent_config.get("use_checklist", False)
+        self.use_multimodal = self.agent_config.get("use_multimodal", False)
+
+        # setup cost
+        model_cost = MODEL_COST_MAPPING.get(self.agent_config["model_name"], None)
+        if model_cost and "api" in self.agent_config["base_url"]:
+            self.input_token_cost = model_cost["input_token_cost"]
+            self.output_token_cost = model_cost["output_token_cost"]
+        else:
+            self.input_token_cost = 0.0
+            self.output_token_cost = 0.0
+
+    def generate_with_retry(self, model_input, constraint_str_list: list = None):
+        total_input_tokens = 0
+        total_output_tokens = 0
+        if self.temperature == 0:
+            response = self.llm.invoke(model_input)
+            total_input_tokens += response.response_metadata["token_usage"]["prompt_tokens"]
+            total_output_tokens += response.response_metadata["token_usage"]["completion_tokens"]
+        else:
+            for i in range(MAX_RETRY):
+                try:
+                    response = self.llm.invoke(model_input)
+                    total_input_tokens += response.response_metadata["token_usage"]["prompt_tokens"]
+                    total_output_tokens += response.response_metadata["token_usage"]["completion_tokens"]
+                    if constraint_str_list:
+                        pass_constraint_num = 0
+                        for constraint_str in constraint_str_list:
+                            if constraint_str in response.content:
+                                pass_constraint_num += 1
+                        if pass_constraint_num == len(constraint_str_list):
+                            break
+                        else:
+                            print(f"Agent has fomat issue, retry... {i+1}/{MAX_RETRY}")
+                            print(response.content)
+                    else:
+                        break
+                except Exception as e:
+                    print(f"Agent returned an Error: {e}")
+                    response = None
+                    time.sleep(RETRY_SLEEP)
+        
+        cost = self.input_token_cost * total_input_tokens / 1000000 + self.output_token_cost * total_output_tokens / 1000000
+        
+        if response is None:
+            return "", cost
+        else:
+            return response.content, cost
+    
+    def prepare_message(self, model_input: dict, prompt_type: str):
+        message = []
+        return message
+
+    def generate_response(self, model_input: dict, prompt_type: str, constraint_str_list: list = None,):
+        total_cost = 0
+        response_list = []
+        # prepare message
+        message = self.prepare_message(model_input, prompt_type)
+        # print(message)
+
+        # n sampling
+        for i in range(self.num_generate):
+            response, cost = self.generate_with_retry(message, constraint_str_list)
+            response_list.append(response)
+            total_cost += cost
+        
+        return response_list, total_cost
+
+
+class GroundingJudgeAgent(BaseAgent):
+    def __init__(self, agent_config: dict):
+        super().__init__(agent_config)
+        self._setup()
+    
+    def prepare_message(self, model_input: dict, prompt_type):
+        message = get_messages(
+            input_info=model_input,
+            inference_mode="judge_grounding",
+            prompt_type=prompt_type,
+            use_multimodal=self.use_multimodal,
+            text_obs=self.agent_config["text_obs_type"],
+            image_obs=self.agent_config["image_obs_type"]
+        )
+        return message
+
+
+class ProgressJudgeAgent(BaseAgent):
+    def __init__(self, agent_config: dict):
+        super().__init__(agent_config)
+        self._setup()
+    
+    def prepare_message(self, model_input: dict, prompt_type):
+        if self.agent_config["input_type"]=="text_only":
+            use_multimodal = False
+            text_obs = self.agent_config["text_obs_type"]
+            image_obs = None
+        elif self.agent_config["input_type"]=="image_only":
+            use_multimodal = True
+            text_obs = None
+            image_obs = self.agent_config["image_obs_type"]
+        elif self.agent_config["input_type"]=="text_image":
+            use_multimodal = True
+            text_obs = self.agent_config["text_obs_type"]
+            image_obs = self.agent_config["image_obs_type"]
+        else:
+            raise ValueError(f"Invalid input type: {self.agent_config['input_type']}")
+        
+        if self.agent_config["use_in_progress"]:
+            use_in_progress = True
+        else:
+            use_in_progress = False
+
+        message = get_messages(
+            input_info=model_input,
+            inference_mode="judge_progress",
+            prompt_type=prompt_type,
+            use_checklist=self.use_checklist,
+            use_multimodal=use_multimodal,
+            text_obs=text_obs,
+            image_obs=image_obs,
+            use_in_progress=use_in_progress
+        )
+        return message
+    
+    def add_logprob(self, ori_logprob: float, add_logprob: float):
+        if ori_logprob is None:
+            return add_logprob
+        else:
+            ori_prob = math.exp(ori_logprob)
+            add_prob = math.exp(add_logprob)
+            return math.log(ori_prob + add_prob)
+    
+    def get_judge_probs(self, logprobs: list):
+        # target_judge = {
+        #     "yes": [" Yes", "Yes"],
+        #     "no": [" No", "No"],
+        #     "in": [" In", "In"]
+        # }
+        target_judge = {
+            "yes": [
+                " Yes", "ĠYes", "Yes", "ĊYes",
+                "Ġyes", "yes", "Ċyes",
+                "ĠYES", "YES", "ĊYES",
+                "ĠDone", "Done", "ĊDone",
+                "ĠCompleted", "Completed", "ĊCompleted",
+                "ĠCorrect", "Correct", "ĊCorrect"
+            ],
+            "no": [
+                " No", "ĠNo", "No", "ĊNo",
+                "ĠNO", "NO", "ĊNO",
+                "ĠNot", "Not", "ĊNot",
+                "ĠNone", "None", "ĊNone",
+                "ĠNope", "Nope", "ĊNope",
+                "ĠUn", "Un", "ĊUn",
+                "ĠWrong", "Wrong", "ĊWrong"
+            ],
+            "in": [
+                " In", "ĠIn", "In", "ĊIn",
+                "ĠPending", "Pending", "ĊPending",
+                "ĠPart", "Part", "ĊPart",
+                "ĠPartial", "Partial", "ĊPartial",
+                "ĠInProgress", "InProgress", "ĊInProgress"
+            ]
+        }
+        response_str = ""
+        judge_probs_list = []
+        # print(logprobs)
+        for i, log_prob in enumerate(logprobs):
+            # Start to find judge string
+            if "<answer>" in response_str:
+                find_judge_str = None
+                for judge_type in target_judge:
+                    if log_prob["token"] in target_judge[judge_type]:
+                        # print(log_prob)
+                        find_judge_str = judge_type
+                        break
+                if find_judge_str:
+                    # print("find judge str")
+                    token_judge_dict = {
+                        "yes": None,
+                        "no": None,
+                        "in": None
+                    }
+                    if "top_logprobs" in log_prob:
+                        for token_info in log_prob["top_logprobs"]:
+                            for judge_type in target_judge:
+                                for judge_str in target_judge[judge_type]:
+                                    # if judge_str in token_info["token"] and token_info["logprob"] > token_judge_dict[judge_type]:
+                                    #     token_judge_dict[judge_type] = token_info["logprob"]
+                                    if judge_str in token_info["token"]:
+                                        # print(token_info["logprob"])
+                                        token_judge_dict[judge_type] = self.add_logprob(token_judge_dict[judge_type], token_info["logprob"])
+                        # for None case
+                        for judge_type in token_judge_dict:
+                            if token_judge_dict[judge_type] is None:
+                                token_judge_dict[judge_type] = float("-inf")
+                        judge_probs_list.append(token_judge_dict)
+                    else:
+                        # for vllm bugs : no top_logprobs
+                        for judge_type in token_judge_dict:
+                            if judge_type == find_judge_str:
+                                token_judge_dict[judge_type] = log_prob["logprob"]
+                            else:
+                                token_judge_dict[judge_type] = float("-inf")
+                        judge_probs_list.append(token_judge_dict)
+                    # print(token_judge_dict)
+            
+            if "</answer>" in response_str:
+                break
+            
+            response_str += log_prob["token"]
+        # print(response_str.replace("Ġ", " ").replace("Ċ", "\n"))
+        # print(judge_probs_list)
+        if len(judge_probs_list) == 0:
+            return [{
+                "yes": 0.0,
+                "no": 0.0,
+                "in": 0.0
+            }]
+        else:
+            # convert with softmax
+            final_judge_probs_list = []
+            for judge_probs in judge_probs_list:
+                exp_logprobs = [math.exp(x) for x in [judge_probs["yes"], judge_probs["no"], judge_probs["in"]]]
+                sum_exp_logprobs = sum(exp_logprobs)
+                softmax_probs = [x / sum_exp_logprobs for x in exp_logprobs]
+                final_judge_probs_list.append({
+                    "yes": softmax_probs[0], 
+                    "no": softmax_probs[1],
+                    "in": softmax_probs[2]
+                })
+            return final_judge_probs_list
+    
+    def generate_probs(self, model_input: dict, prompt_type: str):
+        total_cost = 0
+        response_list = []
+        # prepare message
+        message = self.prepare_message(model_input, prompt_type)
+        # print(message)
+
+        for i in range(self.num_generate):
+            try:
+                response = self.llm.invoke(message)
+                total_input_tokens = response.response_metadata["token_usage"]["prompt_tokens"]
+                total_output_tokens = response.response_metadata["token_usage"]["completion_tokens"]
+                total_cost = self.input_token_cost * total_input_tokens / 1000000 + self.output_token_cost * total_output_tokens / 1000000  
+                logprobs = response.response_metadata["logprobs"]["content"]
+                response_list.append(
+                    {
+                        "response": response.content,
+                        "judge_probs": self.get_judge_probs(logprobs)
+                    }
+                )
+            except Exception as e:
+                print(f"Error: {e}")
+                # print(response.response_metadata["logprobs"])
+                response_list.append(
+                    {
+                        "response": response.content,
+                        "judge_probs": []
+                    }
+                )
+        return response_list, total_cost
+
+
+class ChecklistGenerationAgent(BaseAgent):
+    def __init__(self, agent_config: dict):
+        super().__init__(agent_config)
+        self._setup()
+    
+    def prepare_message(self, model_input: dict, prompt_type):
+        message = get_messages(
+            input_info=model_input,
+            inference_mode="checklist_generation",
+            prompt_type=prompt_type
+        )
+        return message
+    
+
+class ClassifierRewardAgent(Agent):
+    def __init__(self, url: str, use_checklist: bool = False, use_multimodal: bool = False):
+        self.url = url
+        self.use_checklist = use_checklist
+        self.use_multimodal = use_multimodal
+
+    def _process_multimodal_message(self, prompt: str, image_list: list[str]):
+        multimodal_message = []
+        text_prompt_prefix = prompt.split("<IMAGE_PLACEHOLDER>")[0]
+        text_prompt_suffix = prompt.split("<IMAGE_PLACEHOLDER>")[1]
+        multimodal_message = [
+            {"type": "text", "text": text_prompt_prefix},
+            # {"type": "image_url", "image_url": {"url": image_to_base64_url(image_list[0])}},
+            {"type": "image", "image": image_to_base64_url(image_list[0])},
+            {"type": "text", "text": text_prompt_suffix}
+        ]
+        return multimodal_message
+
+    def _make_query(self, user_prompt_template: dict, model_input: dict | list[dict]):
+        if self.use_multimodal:
+            tmp_user_prompt = user_prompt_template["user"].format(
+                **model_input
+            )
+            user_prompt =  self._process_multimodal_message(tmp_user_prompt, model_input["image_list"])
+        else:
+            user_prompt = user_prompt_template["user"].format(
+                **model_input
+            )
+        assistant_prompt = user_prompt_template["assistant"].format(
+            **model_input
+        )
+        query = [
+            {"role": "user", "content": user_prompt},
+            {"role": "assistant", "content": assistant_prompt}
+        ]
+        return query
+
+    def prepare_message(self, model_input: dict | list[dict], batch: bool = False):
+        if self.use_checklist:
+            if self.use_multimodal:
+                user_prompt_template = JUDGE_OURS_BT_MODELING_MULTIMODAL_PROMPT_TEMPLATE
+            else:
+                user_prompt_template = JUDGE_OURS_BT_MODELING_PROMPT_TEMPLATE
+        else:
+            if self.use_multimodal:
+                user_prompt_template = JUDGE_OURS_BT_MODELING_MULTIMODAL_WO_CHECKLIST_PROMPT_TEMPLATE
+            else:
+                user_prompt_template = JUDGE_OURS_BT_MODELING_WO_CHECKLIST_PROMPT_TEMPLATE
+
+        if self.use_multimodal:
+            if batch:
+                message = [self._make_query(user_prompt_template, input) for input in model_input]
+            else:
+                message = [self._make_query(user_prompt_template, model_input)]
+        else:
+            if batch:
+                message = {
+                    "query": [self._make_query(user_prompt_template, input) for input in model_input],
+                    "promptts": []
+                }
+            else:
+                message = {
+                    "query": self._make_query(user_prompt_template, model_input),
+                    "prompts": []
+                }
+
+        return message
+    
+    def get_rm_scroe(self, message: dict | list):
+        headers = {"Content-Type": "application/json"}
+        
+        try:
+            if self.use_multimodal:
+                response = requests.post(
+                    self.url,
+                    json={"messages": message},
+                    timeout=600
+                )
+            else:
+                response = requests.post(
+                    self.url,
+                    headers=headers,
+                    data=json.dumps(message),
+                    timeout=300
+                )
+            response.raise_for_status()
+            
+            response_json = response.json()
+
+            if "rewards" not in response_json:
+                print(f"Error: 'rewards' key not found in API response: {response_json}")
+                return []
+            
+            if "get_reward" in self.url:
+                # use openrlhf
+                return response_json["rewards"]
+            elif "pooling" in self.url:
+                # use vllm server
+                return response_json["reward"]
+            else:
+                # error
+                raise ValueError(f"Invalid URL: {self.url}")
+
+        except requests.exceptions.Timeout:
+            print(f"Error: Request timed out to {self.url}")
+            return []
+        except requests.exceptions.RequestException as e:
+            print(f"Error during request to {self.url}: {e}")
+            return []
+        except json.JSONDecodeError:
+            print(f"Error: Failed to decode JSON response from {self.url}")
+            return []
+        except KeyError as e:
+             print(f"Error: Missing key {e} in response from {self.url}")
+             return []
+
+
+    def generate_response(self, model_input: dict | list[dict], batch: bool = False):
+        if batch:
+            message = self.prepare_message(model_input, batch=True)
+        else:
+            message = self.prepare_message(model_input)
+        rewards = self.get_rm_scroe(message)
+        
+        return rewards, 0

agent/mini_bench/checklist_eval.py

@@ -0,0 +1,95 @@
+import re
+
+from langchain_openai import ChatOpenAI
+
+from .agent import BaseAgent
+
+SYSTEM_PROMPT = "You are an expert evaluator. Your task is to assess how well a Web Agent’s generated checklist aligns with the reference checklist for a given user instruction."
+
+USER_PROMPT = """# Task Description
+Use the provided task description, evaluation criteria, and both checklists to assign a score from 1 to 5. Justify your rating with a brief explanation that considers both content overlap and logical structure.
+
+## Score Criteria
+- 5: Checklist covers all subgoals, is correct and clearly expressed
+- 4: Minor omissions or phrasing issues but mostly accurate and complete
+- 3: Partially matches, but with noticeable gaps or errors
+- 2: Incomplete or includes incorrect steps
+- 1: Mostly irrelevant, incorrect, or missing the task goal
+
+## User Instruction:
+{intent}
+
+## Reference Checklist:
+{gt_checklist}
+
+## Agent’s Generated Checklist:
+{generated_checklist}
+
+# Output Format
+Your response should be in the following format:
+REASON: [Write 2–4 sentences explaining how well the generated checklist matches the reference. Mention specific matches, omissions, errors, or strengths.]
+SCORE: [1–5]
+"""
+
+
+class ChecklistEvalAgent(BaseAgent):
+    def __init__(self, agent_config: dict):
+        super().__init__(agent_config)
+        self._setup()
+    
+    def prepare_message(self, model_input: dict, prompt_type):
+        message = [
+            {
+                "role": "system",
+                "content": SYSTEM_PROMPT
+            },
+            {
+                "role": "user",
+                "content": USER_PROMPT.format(
+                    intent=model_input["intent"],
+                    gt_checklist=model_input["gt_checklist"],
+                    generated_checklist=model_input["generated_checklist"]
+                )
+            }
+        ]
+        return message
+    
+    def generate_response(self, model_input: dict):
+        total_cost = 0
+        response_list = []
+        # prepare message
+        message = self.prepare_message(model_input)
+
+        # n sampling
+        for _ in range(self.num_generate):
+            response, cost = self.generate_with_retry(message, ["SCORE"])
+            response_list.append(response)
+            total_cost += cost
+
+        return response_list, total_cost
+
+def parsing_score(response: str):
+    score = response.split("SCORE:")[-1].split("\n")[0].strip()
+    match = re.search(r'\d+', score)
+    
+    if match:
+        return int(match.group())
+    else:
+        return None
+
+def average_score(scores: list[int]):
+    if len(scores) == 0:
+        return 0
+    return sum(scores) / len(scores)
+
+def get_score(results: list[dict]):
+    score_list = []
+    for result in results:
+        tmp_scores = [parsing_score(response) for response in result["response"]]
+        scores = [score for score in tmp_scores if score is not None]
+        result["score_list"] = scores
+        final_score = average_score(scores)
+        result["score"] = final_score
+        score_list.append(result)
+
+    return results, score_list

agent/mini_bench/eval_utils.py

@@ -0,0 +1,309 @@
+import re
+import random
+from collections import Counter
+
+from .utils import load_json, save_json, create_html_report
+
+random.seed(42)
+def get_score(response_list: list, indicator: str) -> int:
+    if len(response_list) == 0:
+        return [-100]
+
+    if isinstance(response_list[0], float):
+        return response_list
+    
+    if indicator == "prob":
+        score_list = []
+        for response in response_list:
+            total_score = 0
+            for judge_probs in response:
+                yes_prob = judge_probs.get("yes", 0)
+                in_progress_prob = judge_probs.get("in", 0)
+                total_score += yes_prob + in_progress_prob * 0.5
+            if len(response) > 0:
+                score_list.append(total_score / len(response))
+            else:
+                score_list.append(0)
+        return score_list
+    else:
+        score_list = []
+        for response in response_list:
+            if indicator == "SCORE":
+                if "SCORE" in response:
+                    try:
+                        score_str = response.split("SCORE:")[1].split("\n")[0].strip()
+                    except:
+                        score_str = response.split("SCORE:")[-1].strip()
+                    # find first integer
+                    try:
+                        score = re.search(r'-?\d+', score_str).group()
+                        score_list.append(int(score))
+                    except:
+                        score_list.append(0)
+                else:
+                    try:
+                        score_str = response.split("<answer>")[1].split("</answer>")[0].strip()
+                    except:
+                        score_str = response.split("<answer>")[-1].split("</answer>")[0].strip()
+                    # find "Yes" or "No"
+                    if "Yes" in score_str:
+                        score_list.append(1)
+                    elif "In Progress" in score_str:
+                        score_list.append(0.5)
+                    elif "No" in score_str:
+                        score_list.append(0)
+                    else:
+                        score_list.append(0)
+            elif indicator == "JUDGE":
+                try:
+                    judge_str = response.split("JUDGE:")[1].split("\n")[0].strip()
+                except:
+                    judge_str = response.split("JUDGE:")[-1].strip()
+                if "Yes" in judge_str:
+                    score_list.append(1)
+                elif "No" in judge_str:
+                    score_list.append(0)
+                else:
+                    score_list.append(0)
+            elif indicator == "CHECKLIST EVALUATION":
+                if "<answer>" in response:
+                    try:
+                        checklist_str = response.split("<answer>")[1].split("</answer>")[0].strip()
+                    except:
+                        checklist_str = response.split("<answer>")[-1].split("</answer>")[0].strip()
+                else:
+                    checklist_str = response.split("CHECKLIST EVALUATION:")[-1].strip()
+                
+                count_yes = checklist_str.count("Yes")
+                count_no = checklist_str.count("No")
+                count_in_progress = checklist_str.count("In Progress")
+                try:
+                    total_score = (count_yes + count_in_progress*0.5) / (count_yes + count_no + count_in_progress)
+                except:
+                    total_score = 0
+                score_list.append(total_score)
+            else:
+                raise ValueError(f"Invalid indicator: {indicator}")
+    return score_list
+
+def get_acc_and_mrr(chosen_score, rejected_scores):
+    if len(rejected_scores) == 0:
+        return 0, False
+    
+    same_score_num = rejected_scores.count(chosen_score)
+    all_scores = rejected_scores + [chosen_score]
+    sorted_scores = sorted(all_scores, reverse=True)
+    rank = sorted_scores.index(chosen_score) + 1 + same_score_num  # draw penalty
+    if all(chosen_score > r for r in rejected_scores):
+        accuracy = True
+    else:
+        accuracy = False
+    return 1 / rank, accuracy
+
+def average_score(score_list: list[float]):
+    if len(score_list) == 0:
+        return -100
+    return sum(score_list) / len(score_list)
+
+def self_consistency_score(score_list: list[float]):
+    if len(score_list) == 0:
+        return -100
+    counter = Counter(score_list)
+    return max(counter.values()) / len(score_list)
+
+def get_chosen_rejected_scores(data: dict, agg_func: str):
+    if len(data["chosen"]) == 0:
+        data["chosen"] = [{"score": [-100]}]
+    if len(data["rejected"]) == 0:
+        data["rejected"] = [{"score": [-100]}]
+    if not isinstance(data["chosen"][0], dict):
+        data["chosen"][0]["score"] = [-100]
+    if not isinstance(data["rejected"][0], dict):
+        data["rejected"][0]["score"] = [-100]
+        
+    if agg_func == "average":
+        chosen_score = average_score(data["chosen"][0]["score"])
+        rejected_scores = [average_score(rejected_score["score"]) for rejected_score in data["rejected"]]
+    elif agg_func == "self_consistency":
+        chosen_score = self_consistency_score(data["chosen"][0]["score"])
+        rejected_scores = [self_consistency_score(rejected_score["score"]) for rejected_score in data["rejected"]]
+    else:
+        raise ValueError(f"Invalid agg_func: {agg_func}")
+    return chosen_score, rejected_scores
+
+def get_score_results(results, agg_func):
+    score_dict = {"mrr": [], "accuracy": [], "traj_accuracy": []}
+    task_accuracy = {}
+    for result in results:
+        chosen_score, rejected_scores = get_chosen_rejected_scores(result, agg_func)
+        mrr, accuracy = get_acc_and_mrr(chosen_score, rejected_scores)
+        score_dict["mrr"].append(mrr)
+        score_dict["accuracy"].append(accuracy)
+        if result["task_id"] not in task_accuracy:
+            task_accuracy[result["task_id"]] = []
+        task_accuracy[result["task_id"]].append(accuracy)
+
+    for task_id in task_accuracy:
+        if sum(task_accuracy[task_id]) == len(task_accuracy[task_id]):
+            score_dict["traj_accuracy"].append(True)
+        else:
+            score_dict["traj_accuracy"].append(False)
+
+    return score_dict
+
+def calculate_stats(results, agg_func: str="average"):
+    if len(results) == 0:
+        return {
+            "MRR": 0,
+            "Accuracy": 0,
+            "Traj_Accuracy": 0,
+        }
+    total_score = get_score_results(results, agg_func)
+    stats = {
+        "MRR": sum(total_score["mrr"]) / len(total_score["mrr"]),
+        "Accuracy": sum(total_score["accuracy"]) / len(total_score["accuracy"]),
+        "Traj_Accuracy": sum(total_score["traj_accuracy"]) / len(total_score["traj_accuracy"]),
+    }
+    
+    return stats
+
+def group_by_task(results, split_indicator: str):
+    # sort results by task_id and step_id
+    results.sort(key=lambda x: (x["task_id"], x["step_id"]))
+    # group by task_name
+    grouped_task_dict = {}
+    for result in results:
+        task_name = "task_" + str(result["task_id"]) + "_step_" + str(result["step_id"])
+        if task_name not in grouped_task_dict:
+            grouped_task_dict[task_name] = {
+                "task_id": result["task_id"],
+                "step_id": result["step_id"],
+                "intent": result["intent"],
+                "start_url": result["start_url"],
+                "gt_checklist": result["gt_checklist"],
+                "generated_checklist": result.get("generated_checklist", None)  ,
+                "trajectory": result["trajectory"],
+                "current_url": result["current_url"],
+                "text_observation": result["text_observation"],
+                # "image_list": result["image_list"],
+                "chosen": [],
+                "rejected": [],
+                "source_name": result["source_name"],
+            }
+        
+        response = result["response"] if "response" in result else []
+        type_data = {
+            "thought": result["thought"],
+            "action": result["action"],
+            "response": response,
+            "score": get_score(response, split_indicator) if split_indicator != "prob" else get_score(result["judge_probs"], split_indicator),
+        }
+        if split_indicator == "prob":
+            type_data["judge_probs"] = result["judge_probs"]
+        if result["type"] == "chosen":
+            grouped_task_dict[task_name]["chosen"].append(type_data)
+        elif result["type"] == "rejected":
+            grouped_task_dict[task_name]["rejected"].append(type_data)
+    
+    return list(grouped_task_dict.values())
+
+
+def processing_results(results, evaluation_mode: str, num_generate: int, use_batch: bool=False):
+    if "judge_probs" in results[0]:
+        split_indicator = "prob"
+    else:
+        if evaluation_mode == "judge_with_checklist_generation" or evaluation_mode == "judge_with_gt_checklist":
+            split_indicator = "CHECKLIST EVALUATION" 
+        else:
+            split_indicator = "SCORE"
+
+    # if use_batch is True, make it flattened
+    if use_batch:
+        tmp_results = []
+        for result in results:
+            for d in result:
+                tmp_results.append(d)
+        grouped_results = group_by_task(tmp_results, split_indicator)
+    else:
+        grouped_results = group_by_task(results, split_indicator)
+
+    mind2web_results = []
+    webarena_results = []
+    mind2web_task_results = []
+    mind2web_website_results = []
+    mind2web_domain_results = []
+    
+    for grouped_result in grouped_results:
+        if "mind2web" in grouped_result["source_name"]:
+            mind2web_results.append(grouped_result)
+            if grouped_result["source_name"] == "mind2web_test_task":
+                mind2web_task_results.append(grouped_result)
+            elif grouped_result["source_name"] == "mind2web_test_website":
+                mind2web_website_results.append(grouped_result)
+            elif grouped_result["source_name"] == "mind2web_test_domain":
+                mind2web_domain_results.append(grouped_result)
+        elif "webarena" in grouped_result["source_name"]:
+            webarena_results.append(grouped_result)
+            
+    try:
+        final_stats = {
+            "mind2web": {
+                "MRR": {},
+                "Accuracy": {},
+                "Traj_Accuracy": {},
+            },
+            "webarena": {
+                "MRR": {},
+                "Accuracy": {},
+                "Traj_Accuracy": {},
+            },
+            "mind2web_task": {
+                "MRR": {},
+                "Accuracy": {},
+                "Traj_Accuracy": {},
+            },
+            "mind2web_website": {
+                "MRR": {},
+                "Accuracy": {},
+                "Traj_Accuracy": {},
+            },
+            "mind2web_domain": {
+                "MRR": {},
+                "Accuracy": {},
+                "Traj_Accuracy": {},
+            },
+        }
+        for source_results in [
+            ("mind2web", mind2web_results), 
+            ("webarena", webarena_results),
+            ("mind2web_task", mind2web_task_results),
+            ("mind2web_website", mind2web_website_results),
+            ("mind2web_domain", mind2web_domain_results)
+        ]:
+            average_stats = calculate_stats(source_results[1], "average")
+            self_consistency_stats = calculate_stats(source_results[1], "self_consistency")
+            for metric in average_stats:
+                final_stats[source_results[0]][metric]["Average"] = average_stats[metric]
+            for metric in self_consistency_stats:
+                final_stats[source_results[0]][metric]["Self_Consistency"] = self_consistency_stats[metric]
+        
+        if num_generate == 1:
+            for source_name in final_stats:
+                for metric in final_stats[source_name]:
+                    print(f"{round(100 * final_stats[source_name][metric]['Average'], 2)}", end=", ")
+            print()
+        else:
+            for agg_func in ["Average", "Self_Consistency"]:
+                print(f"{agg_func}")
+                for source_name in final_stats:
+                    for metric in final_stats[source_name]:
+                        print(f"{round(100 * final_stats[source_name][metric][agg_func], 2)}", end=", ")
+                print()
+    except Exception as e:
+        print(e)
+        return grouped_results, None    
+    
+    # add function to convert json format results to html format results
+    # TODO: implement this function
+    # create_html_report(results, "results.html")
+    return grouped_results, final_stats

agent/mini_bench/inference_utils.py

@@ -0,0 +1,87 @@
+import time
+
+from multiprocessing import Process, Manager
+from tqdm import tqdm
+
+
+def worker_main(work_queue, result_queue, process_func, config):
+    while True:
+        item = work_queue.get()
+        if item is None:
+            result_queue.put(None)
+            break
+        try:
+            results, cost = process_func(config, item)
+            result_queue.put((results, cost))
+        except Exception as e:
+            item_info = item.get('idx', item.get('id', 'unknown item'))
+            print(f"Error processing item {item_info}: {e}")
+            result_queue.put(None)
+        finally:
+            work_queue.task_done()
+
+def run_parallel_evaluation(dataset, process_func, config, num_workers, description):
+    """
+    Runs parallel evaluation on the given dataset and returns the results.
+
+    Args:
+        dataset (list or datasets.Dataset): Data to evaluate.
+        process_func (callable): Function to process each data item.
+        config (dict): Configuration for the process_func.
+        num_workers (int): Number of worker processes to use.
+        description (str): Description to display on the tqdm progress bar.
+
+    Returns:
+        tuple: (list of evaluation results, total cost)
+    """
+    manager = Manager()
+    work_queue = manager.Queue()
+    result_queue = manager.Queue()
+
+    # Add data to the work queue
+    dataset_list = list(dataset) if not isinstance(dataset, list) else dataset
+    for data in dataset_list:
+        work_queue.put(data)
+    
+    # Add termination signals for workers
+    for _ in range(num_workers):
+        work_queue.put(None)
+
+    # Start parallel processing
+    processes = []
+    for _ in range(num_workers):
+        p = Process(target=worker_main, args=(work_queue, result_queue, process_func, config))
+        p.start()
+        processes.append(p)
+    
+    # Show progress bar and collect results
+    process_results = []
+    process_cost = 0
+    completed_workers = 0
+
+    with tqdm(total=len(dataset_list), desc=description) as pbar:
+        while completed_workers < num_workers:
+            result_item = result_queue.get()
+            if result_item is None:
+                completed_workers += 1
+            else:
+                results, cost = result_item
+                if results is not None:
+                    process_results.append(results)
+                    process_cost += cost if cost is not None else 0
+                pbar.update(1)
+
+    # Wait for all processes to finish
+    for p in processes:
+        p.join()
+
+    # Collect remaining results
+    while not result_queue.empty():
+        result_item = result_queue.get_nowait()
+        if result_item is not None:
+            results, cost = result_item
+            if results is not None:
+                process_results.append(results)
+                process_cost += cost if cost is not None else 0
+
+    return process_results, process_cost

agent/mini_bench/prompts/__init__.py

@@ -0,0 +1 @@
+from .construct_messages import get_messages