actions whitelist, fixes, support new order of the agent env creation in the loop

Author: ollmer

experiments/run_miniwob.py

@@ -8,6 +8,7 @@
 
 fmt = "%(asctime)s - %(levelname)s - %(name)s:%(lineno)d - %(funcName)s() - %(message)s"
 logging.basicConfig(level=logging.INFO, force=True, format=fmt, handlers=[logging.StreamHandler()])
+logger = logging.getLogger(__name__)
 
 if __name__ == "__main__":
     config = load_config("gaia_l1")

src/agentlab/agents/tapeagent/agent.py

@@ -9,6 +9,7 @@
 from tapeagents.agent import Agent
 from tapeagents.core import Action, Observation, StopStep, TapeMetadata, Thought
 from tapeagents.core import Tape as BaseTape
+from tapeagents.tool_calling import ToolSpec
 
 from agentlab.agents.agent_args import AgentArgs
 
@@ -40,8 +41,12 @@ def load_config(config_name: str) -> DictConfig:
 class TapeAgentArgs(AgentArgs):
     config: DictConfig = None  # type: ignore
 
-    def make_agent(self) -> bgym.Agent:
-        agent: Agent = hydra.utils.instantiate(self.config.agent)
+    def make_agent(self, known_actions: tuple[ToolSpec, ...] | None) -> bgym.Agent:
+        if known_actions is None:
+            agent = hydra.utils.instantiate(self.config.agent)
+        else:
+            tools_description = "\n".join([action.description() for action in known_actions])
+            agent = hydra.utils.instantiate(self.config.agent, known_actions=known_actions, tools_description=tools_description)
         return TapeAgent(agent=agent)
 
 

src/agentlab/backends/browser/base.py

@@ -1,40 +1,44 @@
-from tapeagents.environment import FunctionCall
-from tapeagents.mcp import MCPEnvironment, ToolCallAction
-from tapeagents.tool_calling import as_openai_tool
+from pydantic import BaseModel
+from tapeagents.mcp import MCPEnvironment
+from tapeagents.tool_calling import FunctionCall, ToolCallAction, ToolSpec
 
 
-class BrowserBackend():
-    def run_js(self, js: str):
+class BrowserBackend(BaseModel):
+    def initialize(self) -> None:
         raise NotImplementedError
 
-    def call_tool(self, tool_name: str, arguments: dict) -> str:
+    def run_js(self, js: str):
         raise NotImplementedError
 
-    def tools_description(self) -> str:
+    def goto(self, url: str) -> str:
         raise NotImplementedError
 
-    def tools(self) -> list[dict]:
+    def step(self, action: ToolCallAction) -> str:
         raise NotImplementedError
 
+    def actions(self) -> tuple[ToolSpec]:
+        raise NotImplementedError
 
-class MCPBrowserBackend(BrowserBackend):
-    def __init__(self, config_path: str):
-        self.config_path = config_path
-        self.mcp = MCPEnvironment(config_path=self.config_path)
-        self.mcp.initialize()
 
-    def call_tool(self, tool_name: str, arguments: dict) -> str:
-        action = ToolCallAction(
-            function=FunctionCall(name=tool_name, arguments=arguments)
-        )
-        tool_result = self.mcp.step(action)
-        return tool_result.content.content[0].text
 
+class MCPBrowserBackend(BrowserBackend):
+    config_path: str
+    _mcp = None
+    
+    def initialize(self) -> None:
+        self._mcp = MCPEnvironment(config_path=self.config_path)
+        self._mcp.initialize()
 
-    def tools_description(self) -> str:
-        return self.mcp.tools_description()
+    def step(self, action: ToolCallAction) -> str:
+        return self._call_mcp(action)
 
-    def tools(self) -> list[dict]:
-        actions = self.mcp.actions()
-        tools = [as_openai_tool(a).model_dump() for a in actions]
-        return tools
+    def call_tool(self, tool_name: str, arguments: dict) -> str:
+        return self._call_mcp(ToolCallAction(function=FunctionCall(name=tool_name, arguments=arguments)))
+        
+    def _call_mcp(self, action: ToolCallAction) -> str:
+        tool_result = self._mcp.step(action)
+        texts = [c.text for c in tool_result.content.content]
+        return "\n\n".join(texts)
+
+    def actions(self) -> tuple[ToolSpec]:
+        return self._mcp.actions()

src/agentlab/backends/browser/env.py

@@ -1,43 +1,48 @@
 import logging
 import time
-from typing import Any, Literal
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
 
 from tapeagents.core import Action, Observation, StopStep
+from tapeagents.tool_calling import ToolCallAction, ToolSpec
 
 from agentlab.backends.browser.base import BrowserBackend
-from agentlab.benchmarks.abstract_env import AbstractEnv
+from agentlab.benchmarks.abstract_env import AbstractEnv, AbstractEnvArgs
 from agentlab.benchmarks.miniwob.task import AbstractWebTask
 
 logger = logging.getLogger(__name__)
 
+class GoalObservation(Observation):
+    kind: Literal["goal_observation"] = "goal_observation"
+    goal: str
 
 class PageObservation(Observation):
     kind: Literal["page_observation"] = "page_observation"
     content: str
 
-class BrowserAction(Action):
-    kind: Literal["browser_action"] = "browser_action"
-    name: str
-    arguments: dict[str, Any]
-
 
 class BrowserEnv(AbstractEnv):
     def __init__(self, task_name: str, task: AbstractWebTask, backend: BrowserBackend, seed: int = 0):
         self.task_name = task_name
         self.task = task
         self.seed = seed
-        self.backend = backend
         self._turns = 0
+        self.backend = backend
+        self.backend.initialize()
 
     def reset(self, seed: int):
         self.seed = seed
+        logger.info(f"Open task URL: {self.task.url}")
+        page_content = self.backend.goto(self.task.url)
         setup_js = self.task.get_setup_js()
         if setup_js:
             js_result_str = self.backend.run_js(setup_js)
             logger.info(f"Task reset result: {js_result_str}")
+        return [GoalObservation(goal=js_result_str), PageObservation(content=page_content)], {}
 
-    def step(self, action: BrowserAction) -> tuple[Observation, float, bool, bool, dict]:
-        logger.info(f"BrowserEnv.step() called with action {type(action)}")
+    def step(self, action: ToolCallAction) -> tuple[Observation, float, bool, bool, dict]:
+        logger.info(f"BrowserEnv.step() called with action {action.function.name}")
 
         action_exec_start = time.time()
         finished = isinstance(action, StopStep)
@@ -65,8 +70,8 @@ def step(self, action: BrowserAction) -> tuple[Observation, float, bool, bool, d
         logger.info(f"Action result in observation: {obs_view}")
         return observation, reward, finished, truncated, env_info
 
-    def _step(self, action: Action) -> PageObservation:
-        tool_result = self.backend.call_tool(action.name, action.arguments)
+    def _step(self, action: ToolCallAction) -> PageObservation:
+        tool_result = self.backend.step(action)
         return PageObservation(content=tool_result)
 
     def calculate_reward(self, action: Action, observation: PageObservation) -> float:
@@ -80,3 +85,28 @@ def close(self):
         if teardown_js:
             js_result_str = self.backend.run_js(teardown_js)
             logger.info(f"Task teardown result: {js_result_str}")
+
+    def actions(self) -> list[ToolSpec]:
+        all_actions = self.backend.actions()
+        filtered_actions = self.task.filter_actions(all_actions)
+        logger.info(f"Filtered {len(filtered_actions)} actions out of {len(all_actions)} for task {self.task.dataset}")
+        return filtered_actions
+
+
+@dataclass
+class BrowserEnvArgs(AbstractEnvArgs):
+    task: AbstractWebTask
+    task_seed: int
+    task_name: str
+    backend: BrowserBackend
+
+    def __init__(self, task_name: str, task: AbstractWebTask, backend: BrowserBackend, task_seed: int = 0):
+        self.task_name = task_name
+        self.task = task
+        self.task_seed = task_seed
+        self.backend = backend
+
+    def make_env(self, exp_dir: Path) -> BrowserEnv:
+        env = BrowserEnv(task_name=self.task_name, task=self.task, backend=self.backend, seed=self.task_seed)
+        return env
+

src/agentlab/backends/browser/mcp_playwright.py

@@ -1,14 +1,30 @@
+import logging
+
+from tapeagents.tool_calling import ToolCallAction
+
 from agentlab.backends.browser.base import MCPBrowserBackend
 
+logger = logging.getLogger(__name__)
+
 DEFAULT_CONFIG_PATH = "src/agentlab/backends/browser/mcp_playwright.json"
 
+
 class MCPPlaywright(MCPBrowserBackend):
-    def __init__(self, config_path: str | None = None):
-        super().__init__(config_path or DEFAULT_CONFIG_PATH)
+    config_path: str = DEFAULT_CONFIG_PATH
 
     def run_js(self, js: str):
         raw_response = self.call_tool("browser_evaluate", {"function": js})
         _, half_response = raw_response.split("### Result", maxsplit=1)
         result_str, _ = half_response.split("\n### Ran", maxsplit=1)
         result_str = result_str.strip()
         return result_str
+
+    def step(self, action: ToolCallAction) -> str:
+        tool_result = self._call_mcp(action)
+        logger.info(f"Tool result: {tool_result}")
+        snapshot = self.call_tool("browser_snapshot", {})
+        return snapshot
+
+    def goto(self, url: str) -> str:
+        tool_result = self.call_tool("browser_navigate", {"url": url})
+        return tool_result

src/agentlab/benchmarks/miniwob/benchmark.py

@@ -1,42 +1,18 @@
 import logging
-from dataclasses import dataclass
-from pathlib import Path
 from typing import Any
 
-from pydantic import ConfigDict
-
 from agentlab.backends.browser.base import BrowserBackend
-from agentlab.backends.browser.env import BrowserEnv
-from agentlab.benchmarks.abstract_env import AbstractBenchmark, AbstractEnvArgs
+from agentlab.backends.browser.env import BrowserEnvArgs
+from agentlab.benchmarks.abstract_env import AbstractBenchmark
 from agentlab.benchmarks.miniwob.task import MiniWobTask, get_miniwob_tasks
 
 logger = logging.getLogger(__name__)
 
 
-@dataclass
-class MiniwobArgs(AbstractEnvArgs):
-    task: MiniWobTask
-    task_seed: int
-    task_name: str
-    backend: BrowserBackend
-
-    def __init__(self, task_name: str, task: MiniWobTask, backend: BrowserBackend, task_seed: int = 0):
-        self.task_name = task_name
-        self.task = task
-        self.task_seed = task_seed
-        self.backend = backend
-
-    def make_env(self, exp_dir: Path, action_mapping=None) -> BrowserEnv:
-        env = BrowserEnv(task_name=self.task_name, task=self.task, backend=self.backend, seed=self.task_seed)
-        return env
-
-
 class MiniWobBenchmark(AbstractBenchmark):
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-
     backend: BrowserBackend
     name: str = "miniwob"
-    env_args_list: list[MiniwobArgs] = None  # type: ignore
+    env_args_list: list[BrowserEnvArgs] = None  # type: ignore
     dataset: list[MiniWobTask] = None  # type: ignore
 
     def model_post_init(self, __context: Any) -> None:
@@ -45,6 +21,6 @@ def model_post_init(self, __context: Any) -> None:
             self.dataset = get_miniwob_tasks()
         for task in self.dataset:
             name = f"miniwob.{task.task_id}"
-            env_args = MiniwobArgs(task_name=name, task=task, backend=self.backend)
+            env_args = BrowserEnvArgs(task_name=name, task=task, backend=self.backend)
             self.env_args_list.append(env_args)
         logger.info(f"Loaded {len(self.env_args_list)} miniwob tasks")

src/agentlab/benchmarks/miniwob/task.py

@@ -1,6 +1,6 @@
 
 import os
-from typing import Any
+from typing import Any, ClassVar
 
 from browsergym.miniwob import ALL_MINIWOB_TASKS
 
@@ -16,6 +16,15 @@ class MiniWobTask(AbstractWebTask):
     url: str = None
     remove_human_display: bool = True
     episode_max_time: int = 1000000
+    actions_whitelist: ClassVar[list[str]] = [
+        "browser_press_key",
+        "browser_type",
+        "browser_navigate",
+        "browser_click",
+        "browser_drag",
+        "browser_hover",
+        "browser_select_option",
+    ]
 
     def model_post_init(self, __context: Any):
         self.url = f"{self.base_url}/{self.subdomain}.html"

src/agentlab/benchmarks/web_task.py

@@ -1,11 +1,19 @@
+from typing import ClassVar
 
 from pydantic import BaseModel
+from tapeagents.tool_calling import ToolSpec
 
 
 class AbstractWebTask(BaseModel):
-    name: str
+    dataset: str
+    url: str
     validate_per_step: bool = False
-    
+    actions_whitelist: ClassVar[list[str]] = []
+
+    @classmethod
+    def filter_actions(cls, actions: list[ToolSpec]) -> list[str]:
+        return [action for action in actions if action.function.name in cls.actions_whitelist]
+
     def get_setup_js(self) -> str:
         raise NotImplementedError
 

src/agentlab/experiments/loop.py

@@ -25,6 +25,8 @@
 from PIL import Image
 from tqdm import tqdm
 
+from agentlab.backends.browser.env import BrowserEnvArgs
+
 try:
     from agentlab.agents.tapeagent import TapeAgent, save_tape
 except ImportError:
@@ -414,19 +416,23 @@ def run(self):
         env, step_info, err_msg, stack_trace = None, None, None, None
         try:
             logger.info(f"Running experiment {self.exp_name} in:\n  {self.exp_dir}")
-            agent = self.agent_args.make_agent()
-            if hasattr(agent, "set_task_name"):
-                agent.set_task_name(self.env_args.task_name)
-
-            logger.debug("Agent created.")
-
-            env = self.env_args.make_env(
-                action_mapping=agent.action_set.to_python_code,
-                exp_dir=self.exp_dir,
-                use_raw_page_output=getattr(self.agent_args, "use_raw_page_output", False),
-            )
+            if isinstance(self.env_args, BrowserEnvArgs):
+                env = self.env_args.make_env(exp_dir=self.exp_dir)
+                logger.debug("Environment created.")
+                agent = self.agent_args.make_agent(known_actions=env.actions())
+                logger.debug(f"Agent created with actions: {env.actions()}")
+            else:
+                agent = self.agent_args.make_agent()
+                if hasattr(agent, "set_task_name"):
+                    agent.set_task_name(self.env_args.task_name)
+                logger.debug("Agent created.")
+                env = self.env_args.make_env(
+                    action_mapping=agent.action_set.to_python_code,
+                    exp_dir=self.exp_dir,
+                    use_raw_page_output=getattr(self.agent_args, "use_raw_page_output", False),
+                )
+                logger.debug("Environment created.")
 
-            logger.debug("Environment created.")
             step_info = StepInfo(step=0)
             episode_info = [step_info]
             step_info.from_reset(