moving the browsergym.experiment.benchmark module to agentlab

Author: TLSDC

src/agentlab/agents/agent_args.py

@@ -1,6 +1,8 @@
 import bgym
 from bgym import AbstractAgentArgs
 
+from agentlab.experiments.benchmark import Benchmark
+
 
 class AgentArgs(AbstractAgentArgs):
     """Base class for agent arguments for instantiating an agent.
@@ -14,7 +16,7 @@ class MyAgentArgs(AgentArgs):
     Note: for working properly with AgentXRay, the arguments need to be serializable and hasable.
     """
 
-    def set_benchmark(self, benchmark: bgym.Benchmark, demo_mode: bool):
+    def set_benchmark(self, benchmark: Benchmark, demo_mode: bool):
         """Optional method to set benchmark specific flags.
 
         This allows the agent to have minor adjustments based on the benchmark.

src/agentlab/agents/dynamic_prompting.py

@@ -10,13 +10,9 @@
 
 import bgym
 from browsergym.core.action.base import AbstractActionSet
-from browsergym.utils.obs import (
-    flatten_axtree_to_str,
-    flatten_dom_to_str,
-    overlay_som,
-    prune_html,
-)
+from browsergym.utils.obs import flatten_axtree_to_str, flatten_dom_to_str, overlay_som, prune_html
 
+from agentlab.experiments.benchmark import HighLevelActionSetArgs
 from agentlab.llm.llm_utils import (
     BaseMessage,
     ParseError,
@@ -99,7 +95,7 @@ class ObsFlags(Flags):
 
 @dataclass
 class ActionFlags(Flags):
-    action_set: bgym.HighLevelActionSetArgs = None  # should be set by the set_benchmark method
+    action_set: HighLevelActionSetArgs = None  # should be set by the set_benchmark method
     long_description: bool = True
     individual_examples: bool = False
 

src/agentlab/agents/generic_agent/agent_configs.py

@@ -6,6 +6,7 @@
 
 from agentlab.agents import dynamic_prompting as dp
 from agentlab.experiments import args
+from agentlab.experiments.benchmark import HighLevelActionSetArgs
 from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT
 
 from .generic_agent import GenericAgentArgs
@@ -31,7 +32,7 @@
         filter_visible_elements_only=False,
     ),
     action=dp.ActionFlags(
-        action_set=bgym.HighLevelActionSetArgs(
+        action_set=HighLevelActionSetArgs(
             subsets=["bid"],
             multiaction=False,
         ),
@@ -79,7 +80,7 @@
         filter_visible_elements_only=False,
     ),
     action=dp.ActionFlags(
-        action_set=bgym.HighLevelActionSetArgs(
+        action_set=HighLevelActionSetArgs(
             subsets=["bid"],
             multiaction=False,
         ),
@@ -126,7 +127,7 @@
         filter_visible_elements_only=False,
     ),
     action=dp.ActionFlags(
-        action_set=bgym.HighLevelActionSetArgs(
+        action_set=HighLevelActionSetArgs(
             subsets=["bid"],
             multiaction=False,
         ),
@@ -176,7 +177,7 @@
         filter_visible_elements_only=False,
     ),
     action=dp.ActionFlags(
-        action_set=bgym.HighLevelActionSetArgs(
+        action_set=HighLevelActionSetArgs(
             subsets=["bid"],
             multiaction=True,
         ),
@@ -231,7 +232,7 @@
         filter_visible_elements_only=False,
     ),
     action=dp.ActionFlags(
-        action_set=bgym.HighLevelActionSetArgs(
+        action_set=HighLevelActionSetArgs(
             subsets=["bid"],
             multiaction=False,
         ),
@@ -319,7 +320,7 @@
         filter_visible_elements_only=args.Choice([True, False], p=[0.3, 0.7]),
     ),
     action=dp.ActionFlags(
-        action_set=bgym.HighLevelActionSetArgs(
+        action_set=HighLevelActionSetArgs(
             subsets=args.Choice([["bid"], ["bid", "coord"]]),
             multiaction=args.Choice([True, False], p=[0.7, 0.3]),
         ),

src/agentlab/agents/generic_agent/generic_agent.py

@@ -10,19 +10,20 @@
 
 from copy import deepcopy
 from dataclasses import asdict, dataclass
+from functools import partial
 from warnings import warn
 
 import bgym
 from browsergym.experiments.agent import Agent, AgentInfo
 
 from agentlab.agents import dynamic_prompting as dp
 from agentlab.agents.agent_args import AgentArgs
+from agentlab.experiments.benchmark import Benchmark
 from agentlab.llm.chat_api import BaseModelArgs
 from agentlab.llm.llm_utils import Discussion, ParseError, SystemMessage, retry
 from agentlab.llm.tracking import cost_tracker_decorator
 
 from .generic_agent_prompt import GenericPromptFlags, MainPrompt
-from functools import partial
 
 
 @dataclass
@@ -37,7 +38,7 @@ def __post_init__(self):
         except AttributeError:
             pass
 
-    def set_benchmark(self, benchmark: bgym.Benchmark, demo_mode):
+    def set_benchmark(self, benchmark: Benchmark, demo_mode):
         """Override Some flags based on the benchmark."""
         if benchmark.name.startswith("miniwob"):
             self.flags.obs.use_html = True

src/agentlab/agents/generic_agent/reproducibility_agent.py

@@ -23,6 +23,7 @@
 from bs4 import BeautifulSoup
 
 from agentlab.agents.agent_args import AgentArgs
+from agentlab.experiments.benchmark import HighLevelActionSetArgs
 from agentlab.experiments.loop import ExpArgs, ExpResult, yield_all_exp_results
 from agentlab.experiments.study import Study
 from agentlab.llm.chat_api import make_assistant_message
@@ -144,7 +145,7 @@ def _make_backward_compatible(agent_args: GenericAgentArgs):
         if isinstance(action_set, str):
             action_set = action_set.split("+")
 
-        agent_args.flags.action.action_set = bgym.HighLevelActionSetArgs(
+        agent_args.flags.action.action_set = HighLevelActionSetArgs(
             subsets=action_set,
             multiaction=agent_args.flags.action.multi_actions,
         )

src/agentlab/agents/visual_agent/agent_configs.py

@@ -1,9 +1,11 @@
+import bgym
+
+import agentlab.agents.dynamic_prompting as dp
+from agentlab.experiments.benchmark import HighLevelActionSetArgs
 from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT
 
 from .visual_agent import VisualAgentArgs
 from .visual_agent_prompts import PromptFlags
-import agentlab.agents.dynamic_prompting as dp
-import bgym
 
 # the other flags are ignored for this agent.
 DEFAULT_OBS_FLAGS = dp.ObsFlags(
@@ -16,7 +18,7 @@
 )
 
 DEFAULT_ACTION_FLAGS = dp.ActionFlags(
-    action_set=bgym.HighLevelActionSetArgs(subsets=["coord"]),
+    action_set=HighLevelActionSetArgs(subsets=["coord"]),
     long_description=True,
     individual_examples=False,
 )

src/agentlab/agents/visual_agent/visual_agent.py

@@ -15,11 +15,12 @@
 
 from agentlab.agents import dynamic_prompting as dp
 from agentlab.agents.agent_args import AgentArgs
+from agentlab.experiments.benchmark import Benchmark
 from agentlab.llm.chat_api import BaseModelArgs
 from agentlab.llm.llm_utils import Discussion, ParseError, SystemMessage, retry
 from agentlab.llm.tracking import cost_tracker_decorator
 
-from .visual_agent_prompts import PromptFlags, MainPrompt
+from .visual_agent_prompts import MainPrompt, PromptFlags
 
 
 @dataclass
@@ -34,7 +35,7 @@ def __post_init__(self):
         except AttributeError:
             pass
 
-    def set_benchmark(self, benchmark: bgym.Benchmark, demo_mode):
+    def set_benchmark(self, benchmark: Benchmark, demo_mode):
         """Override Some flags based on the benchmark."""
         self.flags.obs.use_tabs = benchmark.is_multi_tab
 

src/agentlab/experiments/benchmark/__init__.py

@@ -0,0 +1,2 @@
+from .base import Benchmark, HighLevelActionSetArgs
+from .configs import DEFAULT_BENCHMARKS