remaining fixes, eval now matched with the old tapeagents evals

Author: ollmer

src/agentlab/agents/tapeagent/agent.py

@@ -26,12 +26,12 @@ class ExtendedMetadata(TapeMetadata):
 
 
 class Tape(BaseTape):
-    metadata: ExtendedMetadata = Field(default_factory=ExtendedMetadata)
+    metadata: ExtendedMetadata = Field(default_factory=ExtendedMetadata)  # type: ignore
 
 
 @dataclass
 class TapeAgentArgs(AgentArgs):
-    agent_name: str
+    agent_name: str = "tape_agent"
 
     def make_agent(self) -> bgym.Agent:
         with hydra.initialize(config_path="conf", version_base="1.1"):

src/agentlab/benchmarks/gaia.py

@@ -1,17 +1,17 @@
+import fcntl
 import logging
 import os
 import re
 import shutil
 import string
 from dataclasses import dataclass
-from math import exp
 from pathlib import Path
 from typing import Any, Literal
 
 import datasets
 from pdf2image import convert_from_path
 from pydantic import Field
-from tapeagents.core import Action, Observation, Step, StopStep, Thought
+from tapeagents.core import Action, Observation, StopStep, Thought
 from tapeagents.environment import ContainerExecutor, StatefulTool, Tool
 from tapeagents.steps import ImageObservation
 from tapeagents.tools.browser import Browser
@@ -78,7 +78,7 @@ def make_env(self, exp_dir: str | Path, action_mapping=None) -> GaiaGym:
         exp_dir = str(exp_dir)
         logger.info(f"Init gaia env with directory {exp_dir}")
         os.environ["TAPEAGENTS_SQLITE_DB"] = os.path.join(exp_dir, "tapedata.sqlite")
-        self.init_code_sandbox(exp_dir)
+        init_code_sandbox(exp_dir)
         tools = [
             WebSearch(),
             VideoReader(exp_path=exp_dir),
@@ -88,34 +88,40 @@ def make_env(self, exp_dir: str | Path, action_mapping=None) -> GaiaGym:
         env = GaiaGym(tools=tools, task=self.task, exp_dir=exp_dir)
         return env
 
-    def init_code_sandbox(self, exp_dir: str) -> None:
-        # Use a common code directory for all tasks in the experiment, which is mounted in the container
-        root_exp_dir = Path(exp_dir).parent
-        code_path = os.path.join(root_exp_dir, "shared_code")
-        os.makedirs(code_path, exist_ok=True)
 
-        container_name = "gaia_code_shared"
-        os.environ["COMPUTER_CONTAINER_NAME"] = container_name
+def init_code_sandbox(exp_dir: str) -> None:
+    # Use a common code directory for all tasks in the experiment, which is mounted in the container
+    root_exp_dir = Path(exp_dir).parent
+    code_path = os.path.join(root_exp_dir, "shared_code")
+    os.makedirs(code_path, exist_ok=True)
 
-        # symlink task code to the shared code directory
-        task_code_path = os.path.join(exp_dir, "code")
-        if not os.path.exists(task_code_path):
-            os.symlink(code_path, task_code_path)
+    container_name = "gaia_code_shared"
+    os.environ["COMPUTER_CONTAINER_NAME"] = container_name
 
+    # symlink task code to the shared code directory
+    task_code_path = os.path.join(exp_dir, "code")
+    if not os.path.exists(task_code_path):
+        os.symlink(code_path, task_code_path)
+
+    try:
         ContainerExecutor(container_name=container_name, work_dir=code_path, no_deps=True)
+    except Exception as e:
+        logger.warning(f"Failed to initialize container executor: {e}")
 
 
 class GaiaBenchmark(AbstractBenchmark):
     name: str = "gaia"
     split: Literal["test", "validation"]
     level: Literal["1", "2", "3", "all"] = "all"
-    env_args_list: list[GaiaGymArgs] = None
-    dataset: dict = Field(default_factory=dict)
+    env_args_list: list[GaiaGymArgs] = None  # type: ignore
+    dataset: dict = None  # type: ignore
 
     def model_post_init(self, __context: Any) -> None:
         if not self.dataset:
             self.dataset = datasets.load_dataset(
-                "gaia-benchmark/GAIA", "2023_all", trust_remote_code=True
+                path="gaia-benchmark/GAIA",
+                name="2023_all",
+                trust_remote_code=True,
             )  # type: ignore
         self.env_args_list = []
         number = 0
@@ -134,7 +140,7 @@ class ExtractedFacts(Thought):
     Thought that contains the list of facts extracted from the document
     """
 
-    kind: Literal["extracted_facts_thought"] = "extracted_facts_thought"
+    kind: Literal["extracted_facts_thought"] = "extracted_facts_thought"  # type: ignore
     extracted_facts: list[str] | dict[str, Any] | str = Field(
         description="facts extracted from the observation"
     )

src/agentlab/benchmarks/multitool_gym.py

@@ -11,12 +11,15 @@
 
 
 class MultiToolGym(AbstractEnv):
-    def __init__(self, tools: list[Tool | StatefulTool]):
+    def __init__(self, tools: list[Tool | StatefulTool], max_turns: int = 50):
         self._env = ToolCollectionEnvironment(tools)
         self._actions = self._env.actions()
+        self.max_turns = max_turns
+        self._turns = 0
 
     def reset(self):
         self._env.reset()
+        self._turns = 0
 
     def step(self, action: Action) -> tuple[Observation, float, bool, bool, dict]:
         logger.info(f"Gym {self.__class__.__name__} step called with action {type(action)}")
@@ -28,11 +31,13 @@ def step(self, action: Action) -> tuple[Observation, float, bool, bool, dict]:
             observation = Observation()  # empty observation
         else:
             observation = self._env.step(action)
+            terminated = isinstance(observation, StopStep)
         action_exec_stop = time.time()
+        self._turns += 1
 
         reward = self.calculate_reward(action)
 
-        truncated = False
+        truncated = self._turns >= self.max_turns
 
         env_info = {
             "step_metadata": observation.metadata,

src/agentlab/experiments/loop.py

@@ -23,17 +23,12 @@
 from browsergym.experiments.utils import count_messages_token, count_tokens
 from dataclasses_json import DataClassJsonMixin
 from PIL import Image
-from tapeagents.core import Step, StepMetadata, TapeMetadata
+from tapeagents.core import Step, StepMetadata
 from tapeagents.dialog_tape import AssistantStep, AssistantThought
 from tapeagents.io import save_json_tape, save_tape_images
 from tqdm import tqdm
 
-from agentlab.agents.tapeagent.agent import (
-    DictObservation,
-    ExtendedMetadata,
-    Tape,
-    TapeAgent,
-)
+from agentlab.agents.tapeagent.agent import DictObservation, Tape, TapeAgent
 
 logger = logging.getLogger(__name__)
 
@@ -237,9 +232,10 @@ def run(self):
         self._set_logger()
 
         # log python environment info
-        save_package_versions(self.exp_dir)
+        save_package_versions(Path(self.exp_dir))
 
         episode_info = []
+        agent = None
         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}")
@@ -255,7 +251,7 @@ def run(self):
             step_info = StepInfo(step=0)
             episode_info = [step_info]
             step_info.from_reset(
-                env, seed=self.env_args.task_seed, obs_preprocessor=agent.obs_preprocessor
+                env, seed=self.env_args.task_seed or 0, obs_preprocessor=agent.obs_preprocessor
             )
             logger.debug("Environment reset.")