TapeAgents export for experiment results

browsergym/experiments/src/browsergym/experiments/loop.py

@@ -657,6 +657,94 @@ def summary_info(self) -> dict:
                 self._summary_info = json.load(f)
         return self._summary_info
 
+    @property
+    def tape(self) -> dict:
+        """
+        TapeAgents (https://github.com/ServiceNow/TapeAgents) framework compatibility.
+        Exports experiment trace in the format of serialized tape.
+        Reuses tape segments if they were already placed in the agent_info during the experiment.
+
+        :returns: dict: serialized tape of the experiment
+        """
+        steps = []
+        for step_info in self.steps_info:
+            if "tape_segment" in step_info.agent_info["extra_info"]:
+                tape_segment = step_info.agent_info["extra_info"]["tape_segment"]
+            else:
+                tape_segment = self._create_tape_segment(step_info)
+            steps += tape_segment
+        metadata = dict(
+            id=str(uuid.uuid4()),
+            author="browsergym_agent",
+            result=self.get_exp_record(),
+        )
+        return dict(steps=steps, metadata=metadata)
+
+    def _create_tape_segment(self, step_info: StepInfo) -> list[dict]:
+        tape_segment = []
+        # extract observation step
+        if step_info.obs is not None:
+            screenshot: str = ""
+            screenshot_som: str = ""
+            obs_dict = step_info.obs.copy()
+            if "screenshot" in obs_dict:
+                screenshot = str(self.exp_dir / f"screenshot_step_{step_info.step}.png")
+                obs_dict.pop("screenshot")
+            if "screenshot_som" in obs_dict:
+                screenshot_som = str(self.exp_dir / f"screenshot_som_step_{step_info.step}.png")
+                obs_dict.pop("screenshot_som")
+            tape_segment.append(
+                dict(
+                    kind="browsergym_observation",
+                    metadata=dict(step=step_info.step),
+                    obs=obs_dict,
+                    screenshot=screenshot,
+                    screenshot_som=screenshot_som,
+                )
+            )
+
+        # extract thought step
+        think = step_info.agent_info.get("think", "")
+        if think:
+            tape_segment.append(
+                dict(kind="browsergym_thought", metadata={"step": step_info.step}, text=think)
+            )
+
+        # extract action step
+        action_str = step_info.action
+        # TODO: this is a naive way to split the arguments, make proper parsing later
+        if "(" in action_str:
+            name, args_str = action_str.split("(", maxsplit=1)
+            args_str = args_str.rstrip(")")
+            arguments = {i: a.strip() for i, a in enumerate(args_str.split(","))}
+        else:
+            name = action_str
+            arguments = {}
+        tape_segment.append(
+            dict(
+                kind="browsergym_action",
+                metadata=dict(
+                    step=step_info.step,
+                    reward=step_info.reward,
+                    raw_reward=step_info.raw_reward,
+                    terminated=step_info.terminated,
+                    truncated=step_info.truncated,
+                    agent_info=step_info.agent_info,
+                    stats=step_info.stats,
+                    task_info=step_info.task_info,
+                ),
+                name=name,
+                arguments=arguments,
+            )
+        )
+        return tape_segment
+
+    def save_tape(self, filename: str = "tape.json"):
+        if os.path.exists(self.exp_dir / filename):
+            raise FileExistsError(f"{filename} already exists in {self.exp_dir}")
+        with open(self.exp_dir / filename, "w") as f:
+            json.dump(self.tape, f, indent=4, ensure_ascii=False)
+
     def get_screenshot(self, step: int, som=False) -> Image:
         key = (step, som)
         if self._screenshots.get(key, None) is None: