workarena bench, reuse bgym task inside

Author: ollmer

src/agentlab/backends/browser/env.py

@@ -3,8 +3,6 @@
 from dataclasses import dataclass
 from pathlib import Path
 
-from browsergym.core.task import AbstractBrowserTask
-
 from agentlab.actions import ToolCall, ToolsActionSet, ToolSpec
 from agentlab.backends.browser.base import BrowserBackend
 from agentlab.benchmarks.abstract_env import AbstractEnv, AbstractEnvArgs
@@ -27,7 +25,7 @@ def final_step():
 
 class BrowserEnv(AbstractEnv):
     def __init__(
-        self, task_name: str, task: AbstractWebTask | AbstractBrowserTask, backend: BrowserBackend, seed: int = 0
+        self, task_name: str, task: AbstractWebTask, backend: BrowserBackend, seed: int = 0
     ):
         self.task_name = task_name
         self.task = task
@@ -36,20 +34,12 @@ def __init__(
         self.backend = backend
         self.backend.initialize()
         self.goal = ""
-        if isinstance(self.task, AbstractBrowserTask) and not self.backend.has_pw_page:
-            raise ValueError(
-                "Legacy task requires a backend with direct playwright page access."
-            )
 
     def reset(self, seed: int):
         self.seed = seed
-        if isinstance(self.task, AbstractBrowserTask):
-            self.goal, task_info = self.task.setup(page=self.backend.page)
-            obs = self._get_obs()
-        else:
-            self.goal, task_info = self.task.setup(backend=self.backend) 
-            obs = self._get_obs()
-            obs = self.task.obs_postprocess(obs)
+        self.goal, task_info = self.task.setup(backend=self.backend)
+        obs = self._get_obs()
+        obs = self.task.obs_postprocess(obs)
         return obs, task_info
 
     def _get_obs(self) -> dict:
@@ -86,21 +76,15 @@ def step(self, action: ToolCall | str) -> tuple[dict, float, bool, bool, dict]:
 
         observation = self.obs_postprocess(observation)
 
-        if isinstance(self.task, AbstractBrowserTask):
-            reward, done, _, info = self.task.validate(page=self.backend.page, chat_messages=[])
-        elif self.task.validate_per_step or done or truncated:
-            reward, info = self.task.validate()
-            if info.get("done", False):
-                done = True
-        else:
-            reward = 0.0
-            info = {}
+        reward, info = self.task.validate()
+        if info.get("done", False):
+            done = True
 
         env_info = {
             **info,
             "action_exec_start": action_exec_start,
             "action_exec_stop": action_exec_stop,
-            "action_exec_timeout": 0.0
+            "action_exec_timeout": 0.0,
         }
         logger.info(f"Action result in observation: {observation}")
         return observation, reward, done, truncated, env_info

src/agentlab/backends/browser/playwright.py

@@ -1,4 +1,5 @@
 import logging
+import time
 from io import BytesIO
 from typing import Any, Callable
 
@@ -82,6 +83,10 @@ def browser_mouse_click_xy(self, x: int, y: int):
         """Click at a given x, y coordinate using the mouse."""
         self._page.mouse.click(x, y, delay=100)
 
+    def browser_wait(self, seconds: int):
+        """Wait for a given number of seconds, up to 10 seconds."""
+        time.sleep(min(seconds, 10))
+
     def evaluate_js(self, js: str):
         js_result = self._page.evaluate(js)
         logger.info(f"JS result: {js_result}")

src/agentlab/benchmarks/workarena/__init__.py

@@ -0,0 +1,4 @@
+from .benchmark import WorkArenaBenchmark
+from .task import WorkarenaTask
+
+__all__ = ["WorkArenaBenchmark", "WorkarenaTask"]

src/agentlab/benchmarks/workarena/benchmark.py

@@ -0,0 +1,56 @@
+import logging
+from typing import Any
+
+from browsergym.workarena import get_all_tasks_agents
+from browsergym.workarena.instance import SNowInstance
+from pydantic import ConfigDict
+from ray.cloudpickle import instance
+
+from agentlab.actions import ToolsActionSet
+from agentlab.backends.browser.base import BrowserBackend
+from agentlab.backends.browser.env import BrowserEnvArgs
+from agentlab.benchmarks.abstract_env import AbstractBenchmark
+
+from .task import WorkarenaTask
+
+logger = logging.getLogger(__name__)
+
+
+class WorkArenaBenchmark(AbstractBenchmark):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    backend_cls: type[BrowserBackend]
+    name: str = "workarena"
+    level: str = "l1"
+    env_args_list: list[BrowserEnvArgs] = None  # type: ignore
+    dataset: list[WorkarenaTask] = None  # type: ignore
+    is_multi_tab: bool = False
+    high_level_action_set_args: ToolsActionSet = None  # type: ignore
+    _snow_instance: SNowInstance = None  # type: ignore
+
+    def model_post_init(self, __context: Any) -> None:
+        self.name = f"workarena_{self.level}_{self.backend_cls.__name__.lower()}"
+        self._snow_instance = SNowInstance()
+        self.env_args_list = []
+        if self.dataset is None:
+            task_seed_tuples = get_all_tasks_agents(filter=self.level)
+            self.dataset = self.load_tasks(task_seed_tuples, self.level)
+        for task in self.dataset:
+            env_args = BrowserEnvArgs(task=task, backend_cls=self.backend_cls)
+            self.env_args_list.append(env_args)
+        logger.info(f"Loaded {len(self.env_args_list)} workarena tasks")
+
+    def load_tasks(self, task_seed_tuples: list[tuple[type, int]], level: str) -> list[WorkarenaTask]:
+        tasks = []
+        
+        for task_cls, seed in task_seed_tuples:
+            task = WorkarenaTask(
+                url="",
+                task_id=task_cls.get_task_id(),
+                instance=self._snow_instance,
+                task_cls=task_cls,
+                level=level,
+                seed=seed,
+            )
+            tasks.append(task)
+        return tasks

src/agentlab/benchmarks/workarena/task.py

@@ -0,0 +1,57 @@
+import logging
+from typing import ClassVar
+
+from browsergym.utils.obs import prune_html
+from browsergym.workarena.instance import SNowInstance
+from browsergym.workarena.tasks.base import AbstractServiceNowTask
+from pydantic import ConfigDict
+
+from agentlab.backends.browser import BrowserBackend
+from agentlab.benchmarks.web_task import AbstractWebTask
+
+logger = logging.getLogger(__name__)
+
+
+class WorkarenaTask(AbstractWebTask):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    dataset: str = "workarena"
+    level: str
+    task_cls: type[AbstractServiceNowTask]
+    seed: int
+    instance: SNowInstance
+    _task_obj: AbstractServiceNowTask = None  # type: ignore
+    actions_whitelist: ClassVar[list[str]] = [
+        "browser_press_key",
+        "browser_type",
+        "browser_click",
+        "browser_drag",
+        "browser_hover",
+        "browser_select_option",
+        "browser_mouse_click_xy",
+        "browser_wait",
+    ]
+
+    def setup(self, backend: BrowserBackend) -> tuple[str, dict]:
+        if not backend.has_pw_page:
+            raise ValueError("Workarena task requires a backend with playwright page access.")
+        self._backend = backend
+        self._task_obj = self.task_cls(instance=self.instance, seed=self.seed) # type: ignore
+        self.url = self._task_obj.start_url
+        goal, info = self._task_obj.setup(backend.page)
+        logger.info(f"Current backend page URL: {backend.page.url}")
+        # backend.goto(self.url)
+        return goal, info
+
+    def teardown(self) -> None:
+        self._task_obj.teardown()
+
+    def validate(self) -> tuple[float, dict]:
+        reward, done, _, info = self._task_obj.validate(page=self._backend.page, chat_messages=[])
+        info["done"] = done
+        return reward, info
+
+    def obs_postprocess(self, obs: dict) -> dict:
+        html = obs.pop("html", "")
+        obs["pruned_html"] = prune_html(html)
+        return obs