more info in tape metadata, better tape browser

Author: ollmer

src/agentlab/agents/tapeagent/agent.py

@@ -1,11 +1,13 @@
 import logging
 from dataclasses import dataclass
-from typing import Any, Literal
+from typing import Literal
 
 import bgym
 import hydra
+from pydantic import Field
 from tapeagents.agent import Agent
-from tapeagents.core import Action, Observation, Tape, TapeMetadata, Thought
+from tapeagents.core import Action, Observation, TapeMetadata, Thought
+from tapeagents.core import Tape as BaseTape
 
 from agentlab.agents.agent_args import AgentArgs
 
@@ -23,6 +25,10 @@ class ExtendedMetadata(TapeMetadata):
     other: dict = {}
 
 
+class Tape(BaseTape):
+    metadata: ExtendedMetadata = Field(default_factory=ExtendedMetadata)
+
+
 @dataclass
 class TapeAgentArgs(AgentArgs):
     agent_name: str

src/agentlab/analyze/tapes.py

@@ -4,12 +4,13 @@
 from collections import defaultdict
 from pathlib import Path
 
+import numpy as np
 import yaml
-from tapeagents.core import Step, StepMetadata, Tape
+from tapeagents.core import Step, StepMetadata
 from tapeagents.renderers.camera_ready_renderer import CameraReadyRenderer
 from tapeagents.tape_browser import TapeBrowser
 
-from agentlab.agents.tapeagent.agent import ExtendedMetadata
+from agentlab.agents.tapeagent.agent import ExtendedMetadata, Tape
 
 logger = logging.getLogger(__name__)
 fmt = "%(asctime)s - %(levelname)s - %(name)s:%(lineno)d - %(funcName)s() - %(message)s"
@@ -20,6 +21,10 @@ class WrapperStep(Step):
     content: dict
 
 
+def pretty_yaml(data: dict) -> str:
+    return yaml.dump(data, sort_keys=False, indent=2) if data else ""
+
+
 class TapesRender(CameraReadyRenderer):
 
     @property
@@ -31,36 +36,35 @@ def render_step(self, step: WrapperStep, index: int, **kwargs):
         step_dict = step.content.copy()
         step_dict.pop("metadata", None)
         kind = step_dict.pop("kind", "Step")
+        if kind == "set_next_node":
+            return ""
         # remove empty keys
         step_dict = {k: v for k, v in step_dict.items() if v is not None and v != ""}
         if len(step_dict) == 1:
             content = list(step_dict.values())[0]
         elif kind == "page_observation":
-            content = step_dict["text"]
+            content = step_dict.get("text", pretty_yaml(step_dict))
             if len(content) > 100:
                 summary = content[:100]
                 content = f"<details><summary>{summary}</summary>---<br>{content}</details>"
         elif kind == "python_code_action":
-            content = step_dict["code"]
+            content = step_dict.get("code", pretty_yaml(step_dict))
         elif kind == "code_execution_result":
-            content = yaml.dump(step_dict["result"], sort_keys=False, indent=2)
+            content = pretty_yaml(step_dict.get("result"))
         else:
-            content = yaml.dump(step_dict, sort_keys=False, indent=2) if step_dict else ""
+            content = pretty_yaml(step_dict)
 
-        if kind.endswith("thought"):
+        if step_dict.get("error") or step_dict.get("result", {}).get("exit_code"):
+            class_ = "error"
+        elif kind.endswith("thought"):
             class_ = "thought"
             kind = kind[:-8]
         elif kind.endswith("action"):
             class_ = "action"
             kind = kind[:-7]
         else:
             class_ = "observation"
-        return (
-            f"<div class='basic-renderer-box {class_}'>"
-            f"<h4 class='step-header'>{kind}</h4>"
-            f"<pre class='step-text'>{content}</pre>"
-            f"</div>"
-        )
+        return f"<div class='basic-renderer-box {class_}'><h4 class='step-header'>{kind}</h4><pre class='step-text'>{content}</pre></div>"
 
 
 class TapesBrowser(TapeBrowser):
@@ -89,10 +93,21 @@ def get_context(self, tape: Tape) -> list:
         return []
 
     def get_tape_name(self, i: int, tape: Tape) -> str:
-        return tape[0].content["content"][:32] + "..."
+        errors = [
+            bool(s.content.get("error", False) or s.content.get("result", {}).get("exit_code"))
+            for s in tape.steps
+        ]
+        mark = "✅ " if tape.metadata.reward > 0 else ""
+        if any(errors):
+            mark = "⚠ "
+        if tape.metadata.task.get("file_name"):
+            mark += "📁 "
+        n = f"{tape.metadata.task.get('Level', '')}.{tape.metadata.task.get('number','')}"
+        name = tape[0].content["content"][:32] + "..."
+        return f"{n} {mark}{name}"
 
     def get_exp_label(self, filename: str, tapes: list[Tape]) -> str:
-        acc, n_solved = 0, 0  # calculate_accuracy(tapes)
+        acc, n_solved = self.calculate_accuracy(tapes)
         errors = defaultdict(int)
         prompt_tokens_num = 0
         output_tokens_num = 0
@@ -106,8 +121,10 @@ def get_exp_label(self, filename: str, tapes: list[Tape]) -> str:
             prompt_tokens_num += llm_call.prompt_length_tokens
             output_tokens_num += llm_call.output_length_tokens
             total_cost += llm_call.cost
+        avg_steps = np.mean([len(tape) for tape in tapes])
+        std_steps = np.std([len(tape) for tape in tapes])
         for tape in tapes:
-            if tape.metadata.result in ["", None, "None"]:
+            if not tape.metadata.terminated:
                 no_result += 1
             if tape.metadata.error:
                 errors["fatal"] += 1
@@ -125,9 +142,9 @@ def get_exp_label(self, filename: str, tapes: list[Tape]) -> str:
                 if kind.endswith("action"):
                     actions[kind] += 1
                     last_action = kind
-                if kind == "search_results_observation" and not len(step_dict["serp"]):
+                if kind == "search_results_observation" and not len(step_dict.get("serp")):
                     errors["search_empty"] += 1
-                if kind == "page_observation" and step_dict["error"]:
+                if kind == "page_observation" and step_dict.get("error"):
                     errors["browser"] += 1
                 elif kind == "llm_output_parsing_failure_action":
                     errors["parsing"] += 1
@@ -136,13 +153,15 @@ def get_exp_label(self, filename: str, tapes: list[Tape]) -> str:
                         errors[f"{last_action}"] += 1
                     else:
                         errors["unknown_action_execution_failure"] += 1
-                elif kind == "code_execution_result" and step_dict["result"]["exit_code"]:
-                    errors["code_execution"] += 1
+                elif kind == "code_execution_result":
+                    if step_dict.get("result", {}).get("exit_code"):
+                        errors["code_execution"] += 1
         timers, timer_counts = self.aggregate_timer_times(tapes)
         html = f"<h2>Solved {acc:.2f}%, {n_solved} out of {len(tapes)}</h2>"
         if "all" in filename:
             html += f"Prompt tokens: {prompt_tokens_num}<br>Output tokens: {output_tokens_num}<br>Cost: {total_cost:.2f} USD<h3>Visible</h3>"
         html += f"Prompt tokens: {visible_prompt_tokens_num}<br>Output tokens: {visible_output_tokens_num}<br>Cost: {visible_cost:.2f} USD"
+        html += f"<h2>Steps per tape: {avg_steps:.1f} ± {std_steps:.1f}</h2>"
         if errors:
             errors_str = "<br>".join(f"{k}: {v}" for k, v in errors.items())
             html += f"<h2>No result: {no_result}</h2>"
@@ -158,6 +177,11 @@ def get_exp_label(self, filename: str, tapes: list[Tape]) -> str:
             html += f"<h2>Timings</h2>{timers_str}"
         return html
 
+    def calculate_accuracy(self, tapes: list[Tape]) -> tuple[float, int]:
+        solved = [tape.metadata.reward for tape in tapes]
+        accuracy = 100 * (sum(solved) / len(solved) if solved else 0.0)
+        return accuracy, sum(solved)
+
     def aggregate_timer_times(self, tapes: list[Tape]):
         timer_sums = defaultdict(float)
         timer_counts = defaultdict(int)
@@ -175,7 +199,7 @@ def aggregate_timer_times(self, tapes: list[Tape]):
         return dict(timer_sums), dict(timer_counts)
 
     def load_tapes(self, exp_dir: str) -> list[dict]:
-        tape_dicts = []
+        tapes: list[Tape] = []
         fpath = Path(self.tapes_folder) / exp_dir
         for json_file in fpath.rglob("tape.json"):
             if json_file.stat().st_size == 0:
@@ -189,11 +213,14 @@ def load_tapes(self, exp_dir: str) -> list[dict]:
                         WrapperStep(content=s, metadata=StepMetadata(**s["metadata"]))
                         for s in tape_dict["steps"]
                     ]
-                    tape_dicts.append(tape)
+                    tapes.append(tape)
             except Exception as e:
                 logger.warning(f"Failed to load {json_file}: {e}")
-        logger.info(f"Loaded {len(tape_dicts)} tapes from {exp_dir}")
-        return tape_dicts
+        logger.info(f"Loaded {len(tapes)} tapes from {exp_dir}")
+        return sorted(
+            tapes,
+            key=lambda x: f"{x.metadata.task.get('Level', '')}{x.metadata.task.get('number', 0):03d}",
+        )
 
     def save_annotation(self, step: int, annotation: str, tape_id: int):
         pass

src/agentlab/benchmarks/gaia.py

@@ -77,6 +77,7 @@ def __init__(
 
     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}")
         self.init_code_sandbox(exp_dir)
         tools = [
             WebSearch(),
@@ -90,15 +91,9 @@ def make_env(self, exp_dir: str | Path, action_mapping=None) -> GaiaGym:
     def init_code_sandbox(self, exp_dir: str) -> None:
         code_path = os.path.join(exp_dir, "code")
         os.makedirs(code_path, exist_ok=True)
-        container_name = "gaia_code_sandbox"
+        container_name = f"gaia_code_{self.task['task_id'][:8]}"
         os.environ["COMPUTER_CONTAINER_NAME"] = container_name
-        ContainerExecutor(
-            work_dir=code_path,
-            container_name=container_name,
-            restart_if_exists=False,
-            stop_container=False,
-            no_deps=True,
-        )
+        ContainerExecutor(container_name=container_name, work_dir=code_path, no_deps=True)
 
 
 class GaiaBenchmark(AbstractBenchmark):
@@ -112,9 +107,10 @@ def model_post_init(self, __context: Any) -> None:
         if not self.dataset:
             self.dataset = datasets.load_dataset("gaia-benchmark/GAIA", "2023_all")
         self.env_args_list = []
-        for task in self.dataset[self.split]:
+        for i, task in enumerate(self.dataset[self.split]):
             if self.level != "all" and task["Level"] != self.level:
                 continue
+            task["number"] = i
             env_args = GaiaGymArgs(task_name="gaia." + task["task_id"], task=task)
             self.env_args_list.append(env_args)
         logger.info(f"Loaded {len(self.env_args_list)} tasks from {self.split} split")

src/agentlab/benchmarks/multitool_gym.py

@@ -1,14 +1,13 @@
 import logging
 import time
 
-from tapeagents.core import Action, Observation, StopStep, Tape
+from tapeagents.core import Action, Observation, StopStep
 from tapeagents.environment import ToolCollectionEnvironment
 from tapeagents.tools.base import StatefulTool, Tool
 
 from agentlab.benchmarks.abstract_env import AbstractEnv
 
 logger = logging.getLogger(__name__)
-EnvTape = Tape[None, Action | Observation]
 
 
 class MultiToolGym(AbstractEnv):

src/agentlab/experiments/loop.py

@@ -23,12 +23,17 @@
 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, Tape
+from tapeagents.core import Step, StepMetadata, TapeMetadata
 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, TapeAgent
+from agentlab.agents.tapeagent.agent import (
+    DictObservation,
+    ExtendedMetadata,
+    Tape,
+    TapeAgent,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -314,8 +319,8 @@ def run(self):
                 logger.info("Saving experiment info.")
                 _save_summary_info(episode_info, self.exp_dir, err_msg, stack_trace)
                 if isinstance(agent, TapeAgent):
-                    save_json_tape(agent.final_tape, self.exp_dir, "tape.json")
-                    save_tape_images(agent.final_tape, self.exp_dir / "tape_attachments")
+                    task = getattr(env, "task", {})
+                    save_tape(self.exp_dir, episode_info, task, agent.final_tape)
             except Exception as e:
                 logger.exception(f"Error while saving experiment info: {e}")
             try:
@@ -949,3 +954,12 @@ def as_tape(steps_info: list[StepInfo]) -> Tape:
             )
             steps.append(AssistantStep(content=step_info.action, metadata=step_metadata))
     return Tape(steps=steps)
+
+
+def save_tape(exp_dir: str, episode_info: list[StepInfo], task: dict, tape: Tape):
+    tape.metadata.reward = sum([step.reward for step in episode_info])
+    tape.metadata.truncated = episode_info[-1].truncated
+    tape.metadata.terminated = episode_info[-1].terminated
+    tape.metadata.task = task
+    save_json_tape(tape, exp_dir, "tape.json")
+    save_tape_images(tape, exp_dir / "tape_attachments")