add gaia scorer for reward in gym, fix envargs serialization

Author: ollmer

scripts/run_gaia.py

@@ -16,8 +16,8 @@
         benchmark=GaiaBenchmark(split="validation"),
         agent_args=TapeAgentArgs("gaia_agent"),
         comment="Gaia eval",
-        logging_level=logging.DEBUG,
-        logging_level_stdout=logging.DEBUG,
+        logging_level=logging.INFO,
+        logging_level_stdout=logging.INFO,
     )
     print(f"Exp args list len: {len(study.exp_args_list)}")
     study.exp_args_list = study.exp_args_list[:1]

src/agentlab/agents/tapeagent/agent.py

@@ -39,7 +39,7 @@ def __init__(self, agent: Agent, tape: Tape):
         self.tape = tape
 
     def obs_preprocessor(self, obs: Any) -> Any:
-        logger.info(f"Observation: {obs}")
+        logger.info(f"Observations: {type(obs)}")
         return obs
 
     def get_action(self, obs: Observation | list[Observation]) -> tuple[str, TapeAgentInfo]:
@@ -52,10 +52,6 @@ def get_action(self, obs: Observation | list[Observation]) -> tuple[str, TapeAge
         action = None
         while not action:
             for event in self.agent.run(self.tape):
-                if event.final_tape:
-                    logger.info(
-                        f"agent run final tape state: {[type(s).__name__ for s in self.tape]}"
-                    )
                 if not event.step:
                     continue
                 self.tape = self.tape.append(event.step)
@@ -64,7 +60,7 @@ def get_action(self, obs: Observation | list[Observation]) -> tuple[str, TapeAge
                     logger.info(f"Thought: {event.step.llm_view()}")
                 elif isinstance(event.step, Action) and not action:
                     action = event.step
-                    logger.info(f"Action: {action}")
+                    logger.info(f"Action: {action.llm_view()}")
                     # we stop at the first action
                 else:
                     logger.info(f"Other step: {type(event.step)}")

src/agentlab/benchmarks/abstract_env.py

@@ -1,15 +1,12 @@
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 
 import gymnasium as gym
+from dataclasses_json import DataClassJsonMixin
 from pydantic import BaseModel
 
 
-class AbstractEnvArgs(BaseModel, frozen=True):
-    """Easily serialiazable class to store the arguments of an environment"""
-
-    task_seed: int = 0
-    task_name: str = ""
-
+class AbstractEnvArgs(ABC):
     @abstractmethod
     def make_env(self, action_mapping, exp_dir, exp_task_kwargs) -> "AbstractEnv":
         """Create an instance of the environment with the arguments stored in this object.
@@ -25,9 +22,17 @@ def make_env(self, action_mapping, exp_dir, exp_task_kwargs) -> "AbstractEnv":
         """
 
 
+@dataclass
+class SerializableEnvArgs(AbstractEnvArgs, DataClassJsonMixin):
+    """Easily serialiazable class to store the arguments of an environment"""
+
+    task_seed: int = 0
+    task_name: str = ""
+
+
 class AbstractBenchmark(BaseModel):
     name: str
-    env_args_list: list[AbstractEnvArgs]
+    env_args_list: list
 
     def get_version(self) -> int:
         return "1"

src/agentlab/benchmarks/gaia.py

@@ -1,20 +1,22 @@
 import logging
 import os
+import re
 import shutil
+import string
 from pathlib import Path
 from typing import Any, Literal
 
 import datasets
 from pydantic import Field
-from tapeagents.core import Observation, 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
 from tapeagents.tools.code_executor import CodeExecutor
 from tapeagents.tools.media_reader import VideoReader
 from tapeagents.tools.web_search import WebSearch
 
-from agentlab.benchmarks.abstract_env import AbstractBenchmark, AbstractEnvArgs
+from agentlab.benchmarks.abstract_env import AbstractBenchmark, SerializableEnvArgs
 from agentlab.benchmarks.multitool_gym import MultiToolGym
 
 logger = logging.getLogger(__name__)
@@ -38,14 +40,41 @@ def reset(self, seed=None) -> tuple[list[Observation], dict]:
             steps.append(image_obs)
         return steps, {}
 
-    def step(self, action: str) -> tuple[Observation, float, bool, bool, dict]:
-        logger.info(f"env step called with action {type(action)}")
-        return super().step(action)
+    def step(self, action: Action) -> tuple[Observation, float, bool, bool, dict]:
+        logger.info(f"Gym step called with action {type(action)}")
+        observation, reward, terminated, truncated, env_info = super().step(action)
+        logger.info(f"Gym observation: {observation.short_view()}")
+        return observation, reward, terminated, truncated, env_info
 
+    def calculate_reward(self, action: Action) -> float:
+        if isinstance(action, GaiaAnswer):
+            model_answer = action.answer
+            ground_truth = self.task["Final answer"]
+            reward = 1.0 if question_scorer(model_answer, ground_truth) else 0.0
+        else:
+            reward = 0.0
 
-class GaiaGymArgs(AbstractEnvArgs, frozen=True):
+        if reward == 1.0:
+            logger.info(f"Task {self.task['task_id']} solved.")
+        else:
+            logger.info(f"Task {self.task['task_id']} failed.")
+
+        return reward
+
+
+class GaiaGymArgs(SerializableEnvArgs):
     task: dict[str, Any]
-    viewport_chars: int = 64000
+    viewport_chars: int
+    task_seed: int
+    task_name: str
+
+    def __init__(
+        self, task_name: str, task: dict[str, Any], viewport_chars: int = 64000, task_seed: int = 0
+    ):
+        self.task_name = task_name
+        self.task = task
+        self.viewport_chars = viewport_chars
+        self.task_seed = task_seed
 
     def make_env(self, exp_dir: str | Path, action_mapping=None) -> GaiaGym:
         exp_dir = str(exp_dir)
@@ -80,7 +109,7 @@ 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:
-            env_args = GaiaGymArgs(task_name="gaia_" + task["task_id"], task=task)
+            env_args = GaiaGymArgs(task_name="gaia." + task["task_id"], task=task)
             self.env_args_list.append(env_args)
 
 
@@ -143,3 +172,96 @@ class GaiaAnswer(StopStep):
     )
     answer: Any = Field(description="Short final answer")
     long_answer: str = Field(description="Detailed final answer not restricted by format rules")
+
+
+def normalize_number_str(number_str: str) -> float:
+    # we replace these common units and commas to allow
+    # conversion to float
+    for char in ["$", "%", ","]:
+        number_str = number_str.replace(char, "")
+    try:
+        return float(number_str)
+    except ValueError:
+        logger.info(f"String {number_str} cannot be normalized to number str.")
+        return float("inf")
+
+
+def split_string(
+    s: str,
+    char_list: list[str] = [",", ";"],
+) -> list[str]:
+    pattern = f"[{''.join(char_list)}]"
+    return re.split(pattern, s)
+
+
+def question_scorer(
+    model_answer: str,
+    ground_truth: str,
+) -> bool:
+    def is_float(element: any) -> bool:
+        try:
+            float(element)
+            return True
+        except ValueError:
+            return False
+
+    # if gt is a number
+    if is_float(ground_truth):
+        logger.info(f"Evaluating {model_answer} as a number.")
+        normalized_answer = normalize_number_str(model_answer)
+        return normalized_answer == float(ground_truth)
+
+    # if gt is a list
+    elif any(char in ground_truth for char in [",", ";"]):
+        logger.info(f"Evaluating {model_answer} as a comma separated list.")
+        # question with the fish: normalization removes punct
+
+        gt_elems = split_string(ground_truth)
+        ma_elems = split_string(model_answer)
+
+        # check length is the same
+        if len(gt_elems) != len(ma_elems):
+            logger.warning("Answer lists have different lengths, returning False.", UserWarning)
+            return False
+
+        # compare each element as float or str
+        comparisons = []
+        for ma_elem, gt_elem in zip(ma_elems, gt_elems):
+            if is_float(gt_elem):
+                normalized_ma_elem = normalize_number_str(ma_elem)
+                comparisons.append(normalized_ma_elem == float(gt_elem))
+            else:
+                # we do not remove punct since comparisons can include punct
+                comparisons.append(
+                    normalize_str(ma_elem, remove_punct=False)
+                    == normalize_str(gt_elem, remove_punct=False)
+                )
+        return all(comparisons)
+
+    # if gt is a str
+    else:
+        logger.info(f"Evaluating {model_answer} as a string.")
+        return normalize_str(model_answer) == normalize_str(ground_truth)
+
+
+def normalize_str(input_str, remove_punct=True) -> str:
+    """
+    Normalize a string by:
+    - Removing all white spaces
+    - Optionally removing punctuation (if remove_punct is True)
+    - Converting to lowercase
+    Parameters:
+    - input_str: str, the string to normalize
+    - remove_punct: bool, whether to remove punctuation (default: True)
+    Returns:
+    - str, the normalized string
+    """
+    # Remove all white spaces. Required e.g for seagull vs. sea gull
+    no_spaces = re.sub(r"\s", "", input_str)
+
+    # Remove punctuation, if specified.
+    if remove_punct:
+        translator = str.maketrans("", "", string.punctuation)
+        return no_spaces.lower().translate(translator)
+    else:
+        return no_spaces.lower()

src/agentlab/benchmarks/multitool_gym.py

@@ -33,7 +33,7 @@ def step(self, action: Action) -> tuple[Observation, float, bool, bool, dict]:
             observation = self._env.step(action)
         action_exec_stop = time.time()
 
-        reward = self.calculate_reward()
+        reward = self.calculate_reward(action)
 
         truncated = False
 
@@ -45,7 +45,7 @@ def step(self, action: Action) -> tuple[Observation, float, bool, bool, dict]:
         }
         return observation, reward, terminated, truncated, env_info
 
-    def calculate_reward(self) -> float:
+    def calculate_reward(self, action: Action) -> float:
         return 0.0
 
     def close(self):

src/agentlab/experiments/launch_exp.py

@@ -98,7 +98,7 @@ def run_experiments(
         logging.info("All jobs are finished. Calling agent_args.close() on all agents...")
         for exp_args in exp_args_list:
             exp_args.agent_args.close()
-        logging.info("Experiment finished.")
+        logging.info(f"Experiment finished and saved in {study_dir}.")
 
 
 def find_incomplete(study_dir: str | Path, include_errors=True):