fixes

Author: ollmer

experiments/run_miniwob.py

@@ -1,12 +1,16 @@
 import logging
 import os
 
+from bgym import DEFAULT_BENCHMARKS
 from dotenv import load_dotenv
 
+from agentlab.agents.generic_agent.agent_configs import FLAGS_GPT_4o
+from agentlab.agents.generic_agent.generic_agent import GenericAgentArgs
 from agentlab.agents.tapeagent.agent import TapeAgentArgs, load_config
 from agentlab.backends.browser.mcp_playwright import MCPPlaywright
 from agentlab.benchmarks.miniwob import MiniWobBenchmark
 from agentlab.experiments.study import make_study
+from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT
 
 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()])
@@ -15,15 +19,26 @@
 
 if __name__ == "__main__":
     config = load_config("miniwob")
+
+    # benchmark = DEFAULT_BENCHMARKS["miniwob"]()
+    agent_args = GenericAgentArgs(
+        chat_model_args=CHAT_MODEL_ARGS_DICT["anthropic/claude-sonnet-4-20250514"],
+        flags=FLAGS_GPT_4o,
+    )
+
+    benchmark = MiniWobBenchmark(backend=MCPPlaywright())
+    # agent_args =TapeAgentArgs(agent_name=config.name, config=config)
+
+
     study = make_study(
-        benchmark=MiniWobBenchmark(backend=MCPPlaywright()),
-        agent_args=TapeAgentArgs(agent_name=config.name, config=config),
+        benchmark=benchmark,
+        agent_args=agent_args,
         comment=config.comment,
         logging_level=logging.INFO,
         logging_level_stdout=logging.INFO,
     )
     if os.environ.get("AGENTLAB_DEBUG"):
-        study.exp_args_list = study.exp_args_list[:1]
+        study.exp_args_list = study.exp_args_list[1:2]
         study.run(n_jobs=1, n_relaunch=1, parallel_backend="sequential")
     else:
         study.run(n_jobs=config.n_jobs, n_relaunch=1, parallel_backend=config.parallel_backend)

src/agentlab/actions.py

@@ -1,15 +1,27 @@
+import json
+import logging
+
 from bgym import AbstractActionSet
 from tapeagents.tool_calling import FunctionCall, ToolCallAction, ToolSpec
 
 from agentlab.llm.llm_utils import parse_html_tags_raise
 
+logger = logging.getLogger(__name__)
+
 
 class ToolsActionSet(AbstractActionSet):
-    def __init__(self, actions:list[ToolSpec]):
+    multiaction: bool = False
+    strict: bool = False
+
+    def __init__(self, actions: list[ToolSpec]):
         self.actions = actions
 
     def describe(self, with_long_description: bool = True, with_examples: bool = True) -> str:
-        tools_description = "\n".join([action.description() for action in self.actions])
+        descs = []
+        for action in self.actions:
+            desc = f"## {action.description()}.\n Schema: {action.model_dump_json(indent=2)}"
+            descs.append(desc)
+        tools_description = "\n".join(descs)
         return tools_description
 
     def example_action(self, abstract: bool) -> str:
@@ -35,13 +47,26 @@ def example_action(self, abstract: bool) -> str:
 }
 </action>
 """
+
     @classmethod
     def parse_action(cls, llm_output: str) -> ToolCallAction:
-        content_dict, valid, retry_message = parse_html_tags_raise(llm_output, keys=["action"])
-        if not valid or "action" not in content_dict:
-            raise ValueError(f"Invalid action: llm_output: {llm_output}, retry_message: {retry_message}")
-        action_str = content_dict["action"]
-        return ToolCallAction(function=FunctionCall(name=action_str["name"], arguments=action_str["arguments"]))
+        logger.info(f"Parsing action: {llm_output}")
+        if "<action>" in llm_output:
+            content_dict, valid, retry_message = parse_html_tags_raise(llm_output, keys=["action"])
+            if not valid or "action" not in content_dict:
+                raise ValueError(
+                    f"Invalid action: llm_output: {llm_output}, retry_message: {retry_message}"
+                )
+            action_str = content_dict["action"]
+        else:
+            action_str = llm_output
+        try:
+            action_dict = json.loads(action_str)
+        except json.JSONDecodeError:
+            raise ValueError(f"Failed to parse action: {action_str}")
+        return ToolCallAction(
+            function=FunctionCall(name=action_dict["name"], arguments=action_dict["arguments"])
+        )
 
     def to_python_code(self, action) -> str:
-        return action.model_dump_json(indent=2)
+        return action

src/agentlab/agents/generic_agent/generic_agent.py

@@ -16,7 +16,9 @@
 import bgym
 from bgym import Benchmark
 from browsergym.experiments.agent import Agent, AgentInfo
+from tapeagents.tool_calling import ToolSpec
 
+from agentlab.actions import ToolsActionSet
 from agentlab.agents import dynamic_prompting as dp
 from agentlab.agents.agent_args import AgentArgs
 from agentlab.llm.chat_api import BaseModelArgs
@@ -65,9 +67,12 @@ def prepare(self):
     def close(self):
         return self.chat_model_args.close_server()
 
-    def make_agent(self):
+    def make_agent(self, actions: list[ToolSpec] | None = None):
         return GenericAgent(
-            chat_model_args=self.chat_model_args, flags=self.flags, max_retry=self.max_retry
+            chat_model_args=self.chat_model_args,
+            flags=self.flags,
+            max_retry=self.max_retry,
+            actions=actions,
         )
 
 
@@ -78,15 +83,21 @@ def __init__(
         chat_model_args: BaseModelArgs,
         flags: GenericPromptFlags,
         max_retry: int = 4,
+        actions: list[ToolSpec] | None = None,
     ):
 
         self.chat_llm = chat_model_args.make_model()
         self.chat_model_args = chat_model_args
         self.max_retry = max_retry
 
         self.flags = flags
-        self.action_set = self.flags.action.action_set.make_action_set()
-        self._obs_preprocessor = dp.make_obs_preprocessor(flags.obs)
+        if actions is not None:
+            self.action_set = ToolsActionSet(actions=actions)
+            self.flags.action.action_set = self.action_set
+            self._obs_preprocessor = lambda obs: obs
+        else:
+            self.action_set = self.flags.action.action_set.make_action_set()
+            self._obs_preprocessor = dp.make_obs_preprocessor(flags.obs)
 
         self._check_flag_constancy()
         self.reset(seed=None)

src/agentlab/backends/browser/base.py

@@ -13,18 +13,23 @@ def run_js(self, js: str):
     def goto(self, url: str) -> str:
         raise NotImplementedError
 
+    def page_snapshot(self) -> str:
+        raise NotImplementedError
+
     def step(self, action: ToolCallAction) -> str:
         raise NotImplementedError
 
     def actions(self) -> tuple[ToolSpec]:
         raise NotImplementedError
 
+    def close(self) -> None:
+        raise NotImplementedError
 
 
 class MCPBrowserBackend(BrowserBackend):
     config_path: str
     _mcp = None
-    
+
     def initialize(self) -> None:
         self._mcp = MCPEnvironment(config_path=self.config_path)
         self._mcp.initialize()
@@ -33,12 +38,17 @@ def step(self, action: ToolCallAction) -> str:
         return self._call_mcp(action)
 
     def call_tool(self, tool_name: str, arguments: dict) -> str:
-        return self._call_mcp(ToolCallAction(function=FunctionCall(name=tool_name, arguments=arguments)))
-        
+        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()
+        return self._mcp.actions()
+
+    def close(self) -> None:
+        self._mcp.close()

src/agentlab/backends/browser/env.py

@@ -1,3 +1,4 @@
+import json
 import logging
 import time
 from dataclasses import dataclass
@@ -10,38 +11,60 @@
 from agentlab.actions import ToolsActionSet
 from agentlab.backends.browser.base import BrowserBackend
 from agentlab.benchmarks.abstract_env import AbstractEnv, AbstractEnvArgs
-from agentlab.benchmarks.miniwob.task import AbstractWebTask
+from agentlab.benchmarks.web_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 BrowserEnv(AbstractEnv):
-    def __init__(self, task_name: str, task: AbstractWebTask, backend: BrowserBackend, seed: int = 0):
+    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._turns = 0
         self.max_turns = task.max_turns
         self.backend = backend
         self.backend.initialize()
+        self.goal = ""
 
     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)
+        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)], {}
+            js_out = self.backend.run_js(setup_js)
+            out_dict = json.loads(js_out)
+            logger.info(f"Task setup result: {out_dict}")
+            goal = out_dict["goal"]
+            done = out_dict["done"]
+            task_start_time = out_dict["task_start_time"]
+            logger.info(f"Task start time: {task_start_time}")
+            if done:
+                raise ValueError("Task is already done")
+            self.goal = goal
+            logger.info(f"Task goal: {self.goal}")
+        page_content = self.backend.page_snapshot()
+        logger.info(f"Initial obs: {page_content}")
+        return {
+            "goal_object": [{"type": "text", "text": self.goal}],
+            "pruned_html": page_content,
+            "axtree_txt": "",
+            "last_action_error": "",
+            "focused_element_bid": "none",
+        }, {}
 
     def step(self, action: ToolCallAction | str) -> tuple[Observation, float, bool, bool, dict]:
         if isinstance(action, str):
@@ -51,49 +74,67 @@ def step(self, action: ToolCallAction | str) -> tuple[Observation, float, bool,
         action_exec_start = time.time()
         finished = isinstance(action, StopStep)
         if finished:
-            observation = Observation()  # empty observation
+            observation = {
+                "goal_object": [{"type": "text", "text": self.goal}],
+                "pruned_html": "Task finished",
+                "axtree_txt": "",
+                "last_action_error": "",
+                "focused_element_bid": "none",
+            }
         else:
             observation = self._step(action)
         action_exec_stop = time.time()
         self._turns += 1
+        logger.info(f"Obs:\n{observation['pruned_html']}")
 
         truncated = self._turns >= self.max_turns
 
         if self.task.validate_per_step or finished or truncated:
-            reward = self.calculate_reward(action, observation)
+            reward, other = self.calculate_reward(action, observation)
+            if other.get("done", False):
+                finished = True
         else:
             reward = 0.0
+            other = {}
 
         env_info = {
-            "step_metadata": observation.metadata,
             "action_exec_start": action_exec_start,
             "action_exec_stop": action_exec_stop,
             "action_exec_timeout": 0.0,
-        }
+        } | other
         obs_view = observation.short_view() if isinstance(observation, Observation) else observation
         logger.info(f"Action result in observation: {obs_view}")
         return observation, reward, finished, truncated, env_info
 
-    def _step(self, action: ToolCallAction) -> PageObservation:
+    def _step(self, action: ToolCallAction) -> dict:
         tool_result = self.backend.step(action)
-        return PageObservation(content=tool_result)
+        return {
+            "goal_object": [{"type": "text", "text": self.goal}],
+            "pruned_html": tool_result,
+            "axtree_txt": "",
+            "last_action_error": "",
+            "focused_element_bid": "none",
+        }
 
-    def calculate_reward(self, action: Action, observation: PageObservation) -> float:
+    def calculate_reward(self, action: Action, observation: PageObservation) -> tuple[float, dict]:
         validate_js = self.task.get_step_validate_js()
         validate_result = self.backend.run_js(validate_js)
         reward, other = self.task.parse_validation_result(validate_result)
-        return reward
+        return reward, other
 
     def close(self):
         teardown_js = self.task.get_teardown_js()
         if teardown_js:
             js_result_str = self.backend.run_js(teardown_js)
             logger.info(f"Task teardown result: {js_result_str}")
+        self.backend.close()
 
     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}")
+        logger.info(
+            f"Filtered {len(filtered_actions)} actions out of {len(all_actions)} for task {self.task.dataset}"
+        )
         return filtered_actions
 
 
@@ -104,13 +145,16 @@ class BrowserEnvArgs(AbstractEnvArgs):
     task_name: str
     backend: BrowserBackend
 
-    def __init__(self, task_name: str, task: AbstractWebTask, backend: BrowserBackend, task_seed: int = 0):
+    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)
+        env = BrowserEnv(
+            task_name=self.task_name, task=self.task, backend=self.backend, seed=self.task_seed
+        )
         return env
-

src/agentlab/benchmarks/miniwob/__init__.py

@@ -1,4 +1,4 @@
 from .benchmark import MiniWobBenchmark
 from .task import MiniWobTask
 
-__all__ = ["MiniWobBenchmark", "MiniWobTask"]
+__all__ = ["MiniWobBenchmark", "MiniWobTask"]