test gym creation and reset

Author: ollmer

src/agentlab/benchmarks/abstract_env.py

@@ -4,20 +4,6 @@
 from pydantic import BaseModel
 
 
-class AbstractBenchmark(BaseModel):
-    name: str
-    env_args_list: list = None
-
-    def get_version(self) -> int:
-        return "1"
-
-    def prepare_backends(self):
-        pass
-
-    def dependency_graph_over_tasks(self) -> dict[str, list[str]]:
-        return {}
-
-
 class AbstractEnvArgs(BaseModel):
     """Easily serialiazable class to store the arguments of an environment"""
 
@@ -36,6 +22,20 @@ def make_env(self, action_mapping, exp_dir, exp_task_kwargs) -> "AbstractEnv":
         """
 
 
+class AbstractBenchmark(BaseModel):
+    name: str
+    env_args_list: list[AbstractEnvArgs]
+
+    def get_version(self) -> int:
+        return "1"
+
+    def prepare_backends(self):
+        pass
+
+    def dependency_graph_over_tasks(self) -> dict[str, list[str]]:
+        return {}
+
+
 class AbstractEnv(gym.Env, ABC):
     @abstractmethod
     def reset(self, seed: int = None) -> tuple[dict[str, any], dict[str, any]]:

src/agentlab/benchmarks/gaia.py

@@ -5,7 +5,7 @@
 import datasets
 from pydantic import Field
 from tapeagents.core import Observation, StopStep, Thought
-from tapeagents.environment import ContainerExecutor
+from tapeagents.environment import ContainerExecutor, StatefulTool, Tool
 from tapeagents.steps import ImageObservation
 from tapeagents.tools.browser import Browser
 from tapeagents.tools.code_executor import CodeExecutor
@@ -16,29 +16,22 @@
 from agentlab.benchmarks.multitool_gym import MultiToolGym
 
 
-class GaiaBenchmark(AbstractBenchmark):
-    exp_dir: str
-    name: str = "gaia"
-    split: Literal["test", "validation"]
-
-    def model_post_init(self, __context: Any) -> None:
-        self.env_args_list = []
-        dataset = datasets.load_dataset("gaia-benchmark/GAIA", "2023_all")[self.split]
-        for task in dataset:
-            task_dir = os.path.join(self.name, task["task_id"])
-            env_args = GaiaGymArgs(task=task, exp_dir=task_dir)
-            self.env_args_list.append(env_args)
-
-
 class GaiaGym(MultiToolGym):
     task: dict
     exp_dir: str
 
+    def __init__(self, tools: list[Tool | StatefulTool], task: dict, exp_dir: str):
+        super().__init__(tools=tools)
+        self.task = task
+        self.exp_dir = exp_dir
+
     def reset(self) -> tuple[list[Observation], dict]:
         super().reset()
+        print("task:", self.task)
         question = GaiaQuestion.from_task(self.task)
         steps = [question]
         if image_obs := with_image(question):
+            print("image_obs:", image_obs)
             steps.append(image_obs)
         return steps
 
@@ -52,9 +45,9 @@ def make_env(self) -> GaiaGym:
         self.init_code_sandbox()
         tools = [
             WebSearch(),
-            VideoReader(self.exp_dir),
-            Browser(self.exp_dir, viewport_chars=self.viewport_chars),
-            CodeExecutor(self.exp_dir),
+            VideoReader(exp_path=self.exp_dir),
+            Browser(exp_path=self.exp_dir, viewport_chars=self.viewport_chars),
+            CodeExecutor(exp_path=self.exp_dir),
         ]
         env = GaiaGym(tools=tools, task=self.task, exp_dir=self.exp_dir)
         return env
@@ -72,6 +65,21 @@ def init_code_sandbox(self) -> None:
         )
 
 
+class GaiaBenchmark(AbstractBenchmark):
+    exp_dir: str
+    name: str = "gaia"
+    split: Literal["test", "validation"]
+    env_args_list: list[GaiaGymArgs] = None
+
+    def model_post_init(self, __context: Any) -> None:
+        self.env_args_list = []
+        dataset = datasets.load_dataset("gaia-benchmark/GAIA", "2023_all")[self.split]
+        for task in dataset:
+            task_dir = os.path.join(self.name, task["task_id"])
+            env_args = GaiaGymArgs(task=task, exp_dir=task_dir)
+            self.env_args_list.append(env_args)
+
+
 class ExtractedFacts(Thought):
     """
     Thought that contains the list of facts extracted from the document

src/agentlab/benchmarks/multitool_gym.py

@@ -1,15 +1,27 @@
+from typing import Annotated, Union
+
+from pydantic import Field, TypeAdapter
 from tapeagents.core import Action, Observation, Tape
+from tapeagents.environment import ToolCollectionEnvironment
+from tapeagents.tools.base import StatefulTool, Tool
 
 from agentlab.benchmarks.abstract_env import AbstractEnv
 
 EnvTape = Tape[None, Action | Observation]
 
 
 class MultiToolGym(AbstractEnv):
+    def __init__(self, tools: list[Tool | StatefulTool]):
+        self._env = ToolCollectionEnvironment(tools)
+        self._actions = self._env.actions()
+        self._actions_parser: TypeAdapter = TypeAdapter(
+            Annotated[Union[self._actions], Field(discriminator="kind")]
+        )
+
     def reset(self):
         self._env.reset()
 
-    def step(self, action: str):
+    def step(self, action: str) -> tuple[Observation, float, bool, bool, dict]:
         try:
             action_step = self._actions_parser.validate_json(action)
         except Exception:

tests/agents/test_gaia_agent.py

@@ -1,5 +1,9 @@
+import os
+
+from tapeagents.steps import ImageObservation
+
 from agentlab.agents.tapeagent.agent import TapeAgent, TapeAgentArgs
-from agentlab.benchmarks.gaia import GaiaBenchmark
+from agentlab.benchmarks.gaia import GaiaBenchmark, GaiaQuestion
 
 
 def test_agent_creation():
@@ -32,3 +36,24 @@ def test_gaia_bench():
     assert task["Annotator Metadata"]["How long did this take?"] == "1 minute"
     assert task["Annotator Metadata"]["Tools"] == "1. Microsoft Excel"
     assert task["Annotator Metadata"]["Number of tools"] == "1"
+
+
+def test_gaia_gym_reset():
+    exp_dir = "/tmp/"
+    bench = GaiaBenchmark(exp_dir=exp_dir, split="validation")
+
+    args = bench.env_args_list[5]
+    env = args.make_env()
+    steps = env.reset()
+    assert len(steps) == 1
+    assert isinstance(steps[0], GaiaQuestion)
+    assert steps[0].content == args.task["Question"]
+
+    args = bench.env_args_list[20]
+    env = args.make_env()
+    steps = env.reset()
+    assert len(steps) == 2
+    assert isinstance(steps[0], GaiaQuestion)
+    assert steps[0].content == args.task["Question"]
+    assert isinstance(steps[1], ImageObservation)
+    assert os.path.basename(steps[1].image_path) == args.task["file_name"]