        if self.env_args.task_seed is None:
            self.env_args.task_seed = np.random.randint(0, SEED_MAX)

        if self.exp_name is None:
            task_name = self.env_args.task_name
            self.exp_name = f"{self.agent_args.agent_name}_on_{task_name}_{self.env_args.task_seed}"

        # if exp_dir exists, it means it's a re-run, move the old one
        if self.exp_dir is not None:
            _move_old_exp(self.exp_dir)

        self.make_id()

        self.exp_date = datetime.now()
        self._make_dir(exp_root)

        self.exp_dir.mkdir(parents=True, exist_ok=True)
        with open(self.exp_dir / "exp_args.pkl", "wb") as f:
            pickle.dump(self, f)

    def _make_dir(self, exp_root):
        """Create a unique directory for the experiment."""
        date_str = self.exp_date.strftime("%Y-%m-%d_%H-%M-%S")
        exp_str = re.sub(
            r"[\/:*?<>|]", "_", self.exp_name
        )  # sanitize exp_name to be used as a file name (substitute forbidden characters)

        for i in range(1000):
            if i >= 999:  # make sure we don't loop forever
                raise ValueError("Could not find a unique name for the experiment directory.")

            tag = f"_{i}" if i > 0 else ""
            self.exp_dir = Path(exp_root) / f"{date_str}_{exp_str}{tag}"
            if not self.exp_dir.exists():
                break

    # TODO distinguish between agent error and environment or system error. e.g.
    # the parsing error of an action should not be re-run.
    def run(self):
        """Run the experiment and save the results"""

        # start writing logs to run logfile
        self._set_logger()

        # log python environment info
        save_package_versions(self.exp_dir)

        episode_info = []
        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}")
            agent = self.agent_args.make_agent()
            logger.debug(f"Agent created.")

            env = self.env_args.make_env(


run

            pickle.dump(self, f)

    def _make_dir(self, exp_root):
        """Create a unique directory for the experiment."""
        date_str = self.exp_date.strftime("%Y-%m-%d_%H-%M-%S")
        exp_str = re.sub(
            r"[\/:*?<>|]", "_", self.exp_name
        )  # sanitize exp_name to be used as a file name (substitute forbidden characters)

        for i in range(1000):
            if i >= 999:  # make sure we don't loop forever
                raise ValueError("Could not find a unique name for the experiment directory.")

            tag = f"_{i}" if i > 0 else ""
            self.exp_dir = Path(exp_root) / f"{date_str}_{exp_str}{tag}"
            if not self.exp_dir.exists():
                break

    # TODO distinguish between agent error and environment or system error. e.g.
    # the parsing error of an action should not be re-run.
    def run(self):
        """Run the experiment and save the results"""

        # start writing logs to run logfile
        self._set_logger()

        # log python environment info
        save_package_versions(self.exp_dir)

        episode_info = []
        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}")
            agent = self.agent_args.make_agent()
            logger.debug(f"Agent created.")

            env = self.env_args.make_env(
                action_mapping=agent.action_set.to_python_code,
                exp_dir=self.exp_dir,
            )

            logger.debug(f"Environment created.")

            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
            )
            logger.debug(f"Environment reset.")

            while not step_info.is_done:  # set a limit
                logger.debug(f"Starting step {step_info.step}.")
                action = step_info.from_action(agent)
                logger.debug(f"Agent chose action:\n {action}")

                if action is None:
                    # will end the episode after saving the step info.
                    step_info.truncated = True

                step_info.save_step_info(
                    self.exp_dir, save_screenshot=self.save_screenshot, save_som=self.save_som
                )
                logger.debug(f"Step info saved.")

                _send_chat_info(env.unwrapped.chat, action, step_info.agent_info)
                logger.debug(f"Chat info sent.")

                if action is None:
                    logger.debug(f"Agent returned None action. Ending episode.")
                    break

                step_info = StepInfo(step=step_info.step + 1)
                episode_info.append(step_info)

                logger.debug(f"Sending action to environment.")
                step_info.from_step(env, action, obs_preprocessor=agent.obs_preprocessor)
                logger.debug(f"Environment stepped.")

        except Exception as e:
            err_msg = f"Exception uncaught by agent or environment in task {self.env_args.task_name}.\n{type(e).__name__}:\n{e}"
            stack_trace = traceback.format_exc()

            self.err_msg = err_msg
            self.stack_trace = stack_trace

            logger.warning(err_msg + "\n" + stack_trace)
            if _is_debugging() and self.enable_debug:
                logger.warning("Debug mode is enabled. Raising the error.")
                raise

        finally:
            try:
                if step_info is not None:
                    step_info.save_step_info(
                        self.exp_dir, save_screenshot=self.save_screenshot, save_som=self.save_som
                    )
            except Exception as e:
                logger.error(f"Error while saving step info in the finally block: {e}")
            try:
                if (
                    not err_msg
                    and len(episode_info) > 0
                    and not (episode_info[-1].terminated or episode_info[-1].truncated)
                ):
                    e = KeyboardInterrupt("Early termination??")
                    err_msg = f"Exception uncaught by agent or environment in task {self.env_args.task_name}.\n{type(e).__name__}:\n{e}"
                logger.info(f"Saving summary info.")
                self.save_summary_info(episode_info, self.exp_dir, err_msg, stack_trace)
            except Exception as e:
                logger.error(f"Error while saving summary info in the finally block: {e}")
            try:
                if env is not None:
                    env.close()
            except Exception as e:
                logger.error(f"Error while closing the environment in the finally block: {e}")
            try:
                self._unset_logger()  # stop writing logs to run logfile
            except Exception as e:
                logger.error(f"Error while unsetting the logger in the finally block: {e}")

    def _set_logger(self):
        # output logging traces to a log file
        file_handler = logging.FileHandler(self.exp_dir / "experiment.log")
        file_handler.setLevel(self.logging_level)  # same level as console outputs
        formatter = logging.Formatter(
            "%(asctime)s - %(process)d - %(name)s - %(levelname)s - %(message)s"
        )
        file_handler.setFormatter(formatter)
        # output handler
        stream_handler = logging.StreamHandler()
        stream_handler.setLevel(self.logging_level_stdout)
        stream_handler.setFormatter(formatter)
        # setup root logger
        root_logger = logging.getLogger()

        # remove previous stream handlers
        for handler in root_logger.handlers:
            if isinstance(handler, logging.StreamHandler):
                root_logger.removeHandler(handler)


                    not err_msg
                    and len(episode_info) > 0
                    and not (episode_info[-1].terminated or episode_info[-1].truncated)
                ):
                    e = KeyboardInterrupt("Early termination??")
                    err_msg = f"Exception uncaught by agent or environment in task {self.env_args.task_name}.\n{type(e).__name__}:\n{e}"
                logger.info(f"Saving summary info.")
                self.save_summary_info(episode_info, self.exp_dir, err_msg, stack_trace)
            except Exception as e:
                logger.error(f"Error while saving summary info in the finally block: {e}")
            try:
                if env is not None:
                    env.close()
            except Exception as e:
                logger.error(f"Error while closing the environment in the finally block: {e}")
            try:
                self._unset_logger()  # stop writing logs to run logfile
            except Exception as e:
                logger.error(f"Error while unsetting the logger in the finally block: {e}")

    def _set_logger(self):
        # output logging traces to a log file
        file_handler = logging.FileHandler(self.exp_dir / "experiment.log")
        file_handler.setLevel(self.logging_level)  # same level as console outputs
        formatter = logging.Formatter(
            "%(asctime)s - %(process)d - %(name)s - %(levelname)s - %(message)s"
        )
        file_handler.setFormatter(formatter)
        # output handler
        stream_handler = logging.StreamHandler()
        stream_handler.setLevel(self.logging_level_stdout)
        stream_handler.setFormatter(formatter)
        # setup root logger
        root_logger = logging.getLogger()

        # remove previous stream handlers
        for handler in root_logger.handlers:
            if isinstance(handler, logging.StreamHandler):
                root_logger.removeHandler(handler)

        root_logger.setLevel(self.logging_level)
        root_logger.addHandler(file_handler)
        root_logger.addHandler(stream_handler)
        # setup openai logger (don't go below INFO verbosity)
        openai_logger = logging.getLogger("openai._base_client")
        openai_logger.setLevel(max(logging.INFO, self.logging_level))

        self.logging_file_handler = file_handler

    def _unset_logger(self):
        root_logger = logging.getLogger()
        root_logger.removeHandler(self.logging_file_handler)

    def save_summary_info(
        self,
        episode_info: list[StepInfo],
        exp_dir,
        err_msg,
        stack_trace,
    ):
        # bring err from agent_info to the top level
        if err_msg is None:
            err_msg, stack_trace = _extract_err_msg(episode_info)
        else:
            # useful until we get a proper place in agent_xray to view error
            # messages.
            if len(episode_info) == 0:
                episode_info.append(StepInfo())


        stream_handler = logging.StreamHandler()
        stream_handler.setLevel(self.logging_level_stdout)
        stream_handler.setFormatter(formatter)
        # setup root logger
        root_logger = logging.getLogger()

        # remove previous stream handlers
        for handler in root_logger.handlers:
            if isinstance(handler, logging.StreamHandler):
                root_logger.removeHandler(handler)

        root_logger.setLevel(self.logging_level)
        root_logger.addHandler(file_handler)
        root_logger.addHandler(stream_handler)
        # setup openai logger (don't go below INFO verbosity)
        openai_logger = logging.getLogger("openai._base_client")
        openai_logger.setLevel(max(logging.INFO, self.logging_level))

        self.logging_file_handler = file_handler

    def _unset_logger(self):
        root_logger = logging.getLogger()
        root_logger.removeHandler(self.logging_file_handler)

    def save_summary_info(
        self,
        episode_info: list[StepInfo],
        exp_dir,
        err_msg,
        stack_trace,
    ):
        # bring err from agent_info to the top level
        if err_msg is None:
            err_msg, stack_trace = _extract_err_msg(episode_info)
        else:
            # useful until we get a proper place in agent_xray to view error
            # messages.
            if len(episode_info) == 0:
                episode_info.append(StepInfo())
            episode_info[-1].agent_info["err_msg"] = err_msg
            episode_info[-1].agent_info["stack_trace"] = stack_trace

        summary_info = dict(


        root_logger = logging.getLogger()

        # remove previous stream handlers
        for handler in root_logger.handlers:
            if isinstance(handler, logging.StreamHandler):
                root_logger.removeHandler(handler)

        root_logger.setLevel(self.logging_level)
        root_logger.addHandler(file_handler)
        root_logger.addHandler(stream_handler)
        # setup openai logger (don't go below INFO verbosity)
        openai_logger = logging.getLogger("openai._base_client")
        openai_logger.setLevel(max(logging.INFO, self.logging_level))

        self.logging_file_handler = file_handler

    def _unset_logger(self):
        root_logger = logging.getLogger()
        root_logger.removeHandler(self.logging_file_handler)

    def save_summary_info(
        self,
        episode_info: list[StepInfo],
        exp_dir,
        err_msg,
        stack_trace,
    ):
        # bring err from agent_info to the top level
        if err_msg is None:
            err_msg, stack_trace = _extract_err_msg(episode_info)
        else:
            # useful until we get a proper place in agent_xray to view error
            # messages.
            if len(episode_info) == 0:
                episode_info.append(StepInfo())
            episode_info[-1].agent_info["err_msg"] = err_msg
            episode_info[-1].agent_info["stack_trace"] = stack_trace

        summary_info = dict(
            n_steps=len(episode_info) - 1,
            cum_reward=sum([step.reward for step in episode_info]),
            cum_raw_reward=sum([step.raw_reward for step in episode_info if step.raw_reward]),
            err_msg=err_msg,
            stack_trace=stack_trace,
        )
        for key, val in _aggregate_episode_stats(episode_info).items():
            summary_info[f"stats.{key}"] = val

        if len(episode_info) > 0:
            summary_info["terminated"] = episode_info[-1].terminated
            summary_info["truncated"] = episode_info[-1].truncated

        with open(exp_dir / "summary_info.json", "w") as f:
            json.dump(summary_info, f, indent=4)

def _extract_err_msg(episode_info: list[StepInfo]):
    """Extract the last error message from the episode info."""
    errors = [(None, None)]
    for step_info in episode_info:
        if step_info.agent_info is None:
            continue
        err_msg = step_info.agent_info.get("err_msg", None)
        if err_msg is not None:
            errors.append((err_msg, step_info.agent_info.get("stack_trace", None)))

    return errors[-1]

def _aggregate_episode_stats(episode_info: list[StepInfo]):
    """Aggregate StepInfo.stats across episodes.

    It will compute the sum and max of each value in the stats dict.
    These two summaries should cover many use cases. If more are needed, the


class ExpResult:
    """Helper class to load and visualize the results of an experiment.

    attributes are loaded lazily.

    Attributes (lazily loaded):
        exp_args: ExpArgs, the arguments of the experiment.
        steps_info: list[StepInfo], the information of each steps so far
        summary_info: dict, the summary of the experiment.
        screenshots: list[Image], the screenshots of each step.
        screenshots_som: list[Image], the screenshots of each step with set of
            marks inprinted.
        flat_exp_args: dict, the flattened version of exp_args.
        chat_video_path: Path, the path to the chat video. (if record_video=True)
        task_video_path: Path, the path to the task video. (if record_video=True)
        combined_video_path: Path, the path to the combined video. (if video was
            combined)
    """

    def __init__(self, exp_dir) -> None:
        self.exp_dir = Path(exp_dir)
        self._exp_args = None
        self._steps_info = {}
        self._summary_info = None
        self._screenshots = {}
        self._flat_exp_args = None
        self._logs = None

    @property
    def exp_args(self) -> ExpArgs:
        if self._exp_args is None:
            with open(self.exp_dir / "exp_args.pkl", "rb") as f:
                self._exp_args = pickle.load(f)
                # in case experiments were moved
                self._exp_args.exp_dir = self.exp_dir
        return self._exp_args

    def get_step_info(self, step: int) -> StepInfo:
        """Load the step info from the file and return it."""
        if self._steps_info.get(step, None) is None:
            with gzip.open(self.exp_dir / f"step_{step}.pkl.gz", "rb") as f:
                self._steps_info[step] = pickle.load(f)
            if self._steps_info[step].obs:
                if "screenshot" not in self._steps_info[step].obs:
                    try:
                        self._steps_info[step].obs["screenshot"] = np.array(
                            self.get_screenshot(step), dtype=np.uint8


        screenshots: list[Image], the screenshots of each step.
        screenshots_som: list[Image], the screenshots of each step with set of
            marks inprinted.
        flat_exp_args: dict, the flattened version of exp_args.
        chat_video_path: Path, the path to the chat video. (if record_video=True)
        task_video_path: Path, the path to the task video. (if record_video=True)
        combined_video_path: Path, the path to the combined video. (if video was
            combined)
    """

    def __init__(self, exp_dir) -> None:
        self.exp_dir = Path(exp_dir)
        self._exp_args = None
        self._steps_info = {}
        self._summary_info = None
        self._screenshots = {}
        self._flat_exp_args = None
        self._logs = None

    @property
    def exp_args(self) -> ExpArgs:
        if self._exp_args is None:
            with open(self.exp_dir / "exp_args.pkl", "rb") as f:
                self._exp_args = pickle.load(f)
                # in case experiments were moved
                self._exp_args.exp_dir = self.exp_dir
        return self._exp_args

    def get_step_info(self, step: int) -> StepInfo:
        """Load the step info from the file and return it."""
        if self._steps_info.get(step, None) is None:
            with gzip.open(self.exp_dir / f"step_{step}.pkl.gz", "rb") as f:
                self._steps_info[step] = pickle.load(f)
            if self._steps_info[step].obs:
                if "screenshot" not in self._steps_info[step].obs:
                    try:
                        self._steps_info[step].obs["screenshot"] = np.array(
                            self.get_screenshot(step), dtype=np.uint8
                        )
                    except FileNotFoundError:
                        pass
                if "screenshot_som" not in self._steps_info[step].obs:
                    try:
                        self._steps_info[step].obs["screenshot_som"] = np.array(
                            self.get_screenshot(step, som=True), dtype=np.uint8
                        )
                    except FileNotFoundError:


    """

    def __init__(self, exp_dir) -> None:
        self.exp_dir = Path(exp_dir)
        self._exp_args = None
        self._steps_info = {}
        self._summary_info = None
        self._screenshots = {}
        self._flat_exp_args = None
        self._logs = None

    @property
    def exp_args(self) -> ExpArgs:
        if self._exp_args is None:
            with open(self.exp_dir / "exp_args.pkl", "rb") as f:
                self._exp_args = pickle.load(f)
                # in case experiments were moved
                self._exp_args.exp_dir = self.exp_dir
        return self._exp_args

    def get_step_info(self, step: int) -> StepInfo:
        """Load the step info from the file and return it."""
        if self._steps_info.get(step, None) is None:
            with gzip.open(self.exp_dir / f"step_{step}.pkl.gz", "rb") as f:
                self._steps_info[step] = pickle.load(f)
            if self._steps_info[step].obs:
                if "screenshot" not in self._steps_info[step].obs:
                    try:
                        self._steps_info[step].obs["screenshot"] = np.array(
                            self.get_screenshot(step), dtype=np.uint8
                        )
                    except FileNotFoundError:
                        pass
                if "screenshot_som" not in self._steps_info[step].obs:
                    try:
                        self._steps_info[step].obs["screenshot_som"] = np.array(
                            self.get_screenshot(step, som=True), dtype=np.uint8
                        )
                    except FileNotFoundError:
                        pass
                # if goal_object is set to None, it indicates it has been saved into a separate file
                if (
                    "goal_object" in self._steps_info[step].obs
                    and self._steps_info[step].obs["goal_object"] is None
                ):
                    with gzip.open(self.exp_dir / "goal_object.pkl.gz", "rb") as f:
                        goal_object = pickle.load(f)
                        self._steps_info[step].obs["goal_object"] = goal_object

        return self._steps_info[step]

    @property
    def steps_info(self) -> list[StepInfo]:
        step_files = list(self.exp_dir.glob("step_*.pkl.gz"))
        for file in step_files:
            step = int(file.name.split("_")[-1].split(".")[0])
            self.get_step_info(step)

        return [self._steps_info[i] for i in range(len(self._steps_info))]

    @property
    def summary_info(self) -> dict:
        if self._summary_info is None:
            with open(self.exp_dir / "summary_info.json", "r") as f:
                # if length is zero raise file not found error
                if os.fstat(f.fileno()).st_size == 0:
                    raise FileNotFoundError(f"summary_info.json is empty.")
                self._summary_info = json.load(f)
        return self._summary_info


                        pass
                if "screenshot_som" not in self._steps_info[step].obs:
                    try:
                        self._steps_info[step].obs["screenshot_som"] = np.array(
                            self.get_screenshot(step, som=True), dtype=np.uint8
                        )
                    except FileNotFoundError:
                        pass
                # if goal_object is set to None, it indicates it has been saved into a separate file
                if (
                    "goal_object" in self._steps_info[step].obs
                    and self._steps_info[step].obs["goal_object"] is None
                ):
                    with gzip.open(self.exp_dir / "goal_object.pkl.gz", "rb") as f:
                        goal_object = pickle.load(f)
                        self._steps_info[step].obs["goal_object"] = goal_object

        return self._steps_info[step]

    @property
    def steps_info(self) -> list[StepInfo]:
        step_files = list(self.exp_dir.glob("step_*.pkl.gz"))
        for file in step_files:
            step = int(file.name.split("_")[-1].split(".")[0])
            self.get_step_info(step)

        return [self._steps_info[i] for i in range(len(self._steps_info))]

    @property
    def summary_info(self) -> dict:
        if self._summary_info is None:
            with open(self.exp_dir / "summary_info.json", "r") as f:
                # if length is zero raise file not found error
                if os.fstat(f.fileno()).st_size == 0:
                    raise FileNotFoundError(f"summary_info.json is empty.")
                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
        """


                if (
                    "goal_object" in self._steps_info[step].obs
                    and self._steps_info[step].obs["goal_object"] is None
                ):
                    with gzip.open(self.exp_dir / "goal_object.pkl.gz", "rb") as f:
                        goal_object = pickle.load(f)
                        self._steps_info[step].obs["goal_object"] = goal_object

        return self._steps_info[step]

    @property
    def steps_info(self) -> list[StepInfo]:
        step_files = list(self.exp_dir.glob("step_*.pkl.gz"))
        for file in step_files:
            step = int(file.name.split("_")[-1].split(".")[0])
            self.get_step_info(step)

        return [self._steps_info[i] for i in range(len(self._steps_info))]

    @property
    def summary_info(self) -> dict:
        if self._summary_info is None:
            with open(self.exp_dir / "summary_info.json", "r") as f:
                # if length is zero raise file not found error
                if os.fstat(f.fileno()).st_size == 0:
                    raise FileNotFoundError(f"summary_info.json is empty.")
                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=f"browsergym_agent_[{self.exp_args.agent_args.agent_name}]",


tape

    @property
    def steps_info(self) -> list[StepInfo]:
        step_files = list(self.exp_dir.glob("step_*.pkl.gz"))
        for file in step_files:
            step = int(file.name.split("_")[-1].split(".")[0])
            self.get_step_info(step)

        return [self._steps_info[i] for i in range(len(self._steps_info))]

    @property
    def summary_info(self) -> dict:
        if self._summary_info is None:
            with open(self.exp_dir / "summary_info.json", "r") as f:
                # if length is zero raise file not found error
                if os.fstat(f.fileno()).st_size == 0:
                    raise FileNotFoundError(f"summary_info.json is empty.")
                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=f"browsergym_agent_[{self.exp_args.agent_args.agent_name}]",
            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 = copy.deepcopy(step_info.obs)
            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,


        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=f"browsergym_agent_[{self.exp_args.agent_args.agent_name}]",
            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 = copy.deepcopy(step_info.obs)
            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 steps
        function_calls = highlevel_action_parser.parse_string(step_info.action, parse_all=True)
        for name, arguments in function_calls:
            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:
            file_name = f"screenshot_{'som_' if som else ''}step_{step}"
            try:
                with Image.open(self.exp_dir / (file_name + ".png")) as img:
                    self._screenshots[key] = img.copy()
            except FileNotFoundError:
                with Image.open(self.exp_dir / (file_name + ".jpg")) as img:
                    self._screenshots[key] = img.copy()
        return self._screenshots[key]

    def get_screenshots(self, som=False):


        for name, arguments in function_calls:
            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:
            file_name = f"screenshot_{'som_' if som else ''}step_{step}"
            try:
                with Image.open(self.exp_dir / (file_name + ".png")) as img:
                    self._screenshots[key] = img.copy()
            except FileNotFoundError:
                with Image.open(self.exp_dir / (file_name + ".jpg")) as img:
                    self._screenshots[key] = img.copy()
        return self._screenshots[key]

    def get_screenshots(self, som=False):
        files = list(self.exp_dir.glob("screenshot_step_*"))
        max_step = 0
        for file in files:
            step = int(file.name.split("_")[-1].split(".")[0])
            self.get_screenshot(step, som=som)
            max_step = max(max_step, 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:
            file_name = f"screenshot_{'som_' if som else ''}step_{step}"
            try:
                with Image.open(self.exp_dir / (file_name + ".png")) as img:
                    self._screenshots[key] = img.copy()
            except FileNotFoundError:
                with Image.open(self.exp_dir / (file_name + ".jpg")) as img:
                    self._screenshots[key] = img.copy()
        return self._screenshots[key]

    def get_screenshots(self, som=False):
        files = list(self.exp_dir.glob("screenshot_step_*"))
        max_step = 0
        for file in files:
            step = int(file.name.split("_")[-1].split(".")[0])
            self.get_screenshot(step, som=som)
            max_step = max(max_step, step)
        return [self._screenshots.get((i, som), None) for i in range(max_step + 1)]

    @property
    def screenshots(self):
        return self.get_screenshots(som=False)

    @property
    def screenshots_som(self):
        return self.get_screenshots(som=True)

    @property
    def flat_exp_args(self) -> dict:


        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:
            file_name = f"screenshot_{'som_' if som else ''}step_{step}"
            try:
                with Image.open(self.exp_dir / (file_name + ".png")) as img:
                    self._screenshots[key] = img.copy()
            except FileNotFoundError:
                with Image.open(self.exp_dir / (file_name + ".jpg")) as img:
                    self._screenshots[key] = img.copy()
        return self._screenshots[key]

    def get_screenshots(self, som=False):
        files = list(self.exp_dir.glob("screenshot_step_*"))
        max_step = 0
        for file in files:
            step = int(file.name.split("_")[-1].split(".")[0])
            self.get_screenshot(step, som=som)
            max_step = max(max_step, step)
        return [self._screenshots.get((i, som), None) for i in range(max_step + 1)]

    @property
    def screenshots(self):
        return self.get_screenshots(som=False)

    @property
    def screenshots_som(self):
        return self.get_screenshots(som=True)

    @property
    def flat_exp_args(self) -> dict:
        """Return a dict with exp_args flattened."""
        if self._flat_exp_args is None:
            exp_args = asdict(self.exp_args)
            # this will flatten nested dicts
            self._flat_exp_args = _flatten_dict(exp_args)
        return self._flat_exp_args

    def get_exp_record(self) -> dict:
        """Return a dict with exp_args flattened and summary_info."""


screenshots

        if self._screenshots.get(key, None) is None:
            file_name = f"screenshot_{'som_' if som else ''}step_{step}"
            try:
                with Image.open(self.exp_dir / (file_name + ".png")) as img:
                    self._screenshots[key] = img.copy()
            except FileNotFoundError:
                with Image.open(self.exp_dir / (file_name + ".jpg")) as img:
                    self._screenshots[key] = img.copy()
        return self._screenshots[key]

    def get_screenshots(self, som=False):
        files = list(self.exp_dir.glob("screenshot_step_*"))
        max_step = 0
        for file in files:
            step = int(file.name.split("_")[-1].split(".")[0])
            self.get_screenshot(step, som=som)
            max_step = max(max_step, step)
        return [self._screenshots.get((i, som), None) for i in range(max_step + 1)]

    @property
    def screenshots(self):
        return self.get_screenshots(som=False)

    @property
    def screenshots_som(self):
        return self.get_screenshots(som=True)

    @property
    def flat_exp_args(self) -> dict:
        """Return a dict with exp_args flattened."""
        if self._flat_exp_args is None:
            exp_args = asdict(self.exp_args)
            # this will flatten nested dicts
            self._flat_exp_args = _flatten_dict(exp_args)
        return self._flat_exp_args

    def get_exp_record(self) -> dict:
        """Return a dict with exp_args flattened and summary_info."""
        record = {"exp_dir": self.exp_dir}
        try:
            record.update(self.flat_exp_args)
        except FileNotFoundError:


                    self._screenshots[key] = img.copy()
            except FileNotFoundError:
                with Image.open(self.exp_dir / (file_name + ".jpg")) as img:
                    self._screenshots[key] = img.copy()
        return self._screenshots[key]

    def get_screenshots(self, som=False):
        files = list(self.exp_dir.glob("screenshot_step_*"))
        max_step = 0
        for file in files:
            step = int(file.name.split("_")[-1].split(".")[0])
            self.get_screenshot(step, som=som)
            max_step = max(max_step, step)
        return [self._screenshots.get((i, som), None) for i in range(max_step + 1)]

    @property
    def screenshots(self):
        return self.get_screenshots(som=False)

    @property
    def screenshots_som(self):
        return self.get_screenshots(som=True)

    @property
    def flat_exp_args(self) -> dict:
        """Return a dict with exp_args flattened."""
        if self._flat_exp_args is None:
            exp_args = asdict(self.exp_args)
            # this will flatten nested dicts
            self._flat_exp_args = _flatten_dict(exp_args)
        return self._flat_exp_args

    def get_exp_record(self) -> dict:
        """Return a dict with exp_args flattened and summary_info."""
        record = {"exp_dir": self.exp_dir}
        try:
            record.update(self.flat_exp_args)
        except FileNotFoundError:
            pass
        try:
            record.update(self.summary_info)
        except FileNotFoundError:


        return self._screenshots[key]

    def get_screenshots(self, som=False):
        files = list(self.exp_dir.glob("screenshot_step_*"))
        max_step = 0
        for file in files:
            step = int(file.name.split("_")[-1].split(".")[0])
            self.get_screenshot(step, som=som)
            max_step = max(max_step, step)
        return [self._screenshots.get((i, som), None) for i in range(max_step + 1)]

    @property
    def screenshots(self):
        return self.get_screenshots(som=False)

    @property
    def screenshots_som(self):
        return self.get_screenshots(som=True)

    @property
    def flat_exp_args(self) -> dict:
        """Return a dict with exp_args flattened."""
        if self._flat_exp_args is None:
            exp_args = asdict(self.exp_args)
            # this will flatten nested dicts
            self._flat_exp_args = _flatten_dict(exp_args)
        return self._flat_exp_args

    def get_exp_record(self) -> dict:
        """Return a dict with exp_args flattened and summary_info."""
        record = {"exp_dir": self.exp_dir}
        try:
            record.update(self.flat_exp_args)
        except FileNotFoundError:
            pass
        try:
            record.update(self.summary_info)
        except FileNotFoundError:
            pass
        return record

    @property
    def chat_video_path(self) -> Path:
        try:
            return next(self.exp_dir.glob("chat_video/*.webm"))
        except StopIteration:
            raise FileNotFoundError(f"No chat_video found in {self.exp_dir}")


            max_step = max(max_step, step)
        return [self._screenshots.get((i, som), None) for i in range(max_step + 1)]

    @property
    def screenshots(self):
        return self.get_screenshots(som=False)

    @property
    def screenshots_som(self):
        return self.get_screenshots(som=True)

    @property
    def flat_exp_args(self) -> dict:
        """Return a dict with exp_args flattened."""
        if self._flat_exp_args is None:
            exp_args = asdict(self.exp_args)
            # this will flatten nested dicts
            self._flat_exp_args = _flatten_dict(exp_args)
        return self._flat_exp_args

    def get_exp_record(self) -> dict:
        """Return a dict with exp_args flattened and summary_info."""
        record = {"exp_dir": self.exp_dir}
        try:
            record.update(self.flat_exp_args)
        except FileNotFoundError:
            pass
        try:
            record.update(self.summary_info)
        except FileNotFoundError:
            pass
        return record

    @property
    def chat_video_path(self) -> Path:
        try:
            return next(self.exp_dir.glob("chat_video/*.webm"))
        except StopIteration:
            raise FileNotFoundError(f"No chat_video found in {self.exp_dir}")

    @property
    def task_video_path(self) -> Path:
        try:
            return next(self.exp_dir.glob("task_video/*.webm"))
        except StopIteration:
            raise FileNotFoundError(f"No task_video found in {self.exp_dir}")

    @property
    def combined_video_path(self) -> Path:
        return self.exp_dir / "combined_video.mp4"

    @property


        if self._flat_exp_args is None:
            exp_args = asdict(self.exp_args)
            # this will flatten nested dicts
            self._flat_exp_args = _flatten_dict(exp_args)
        return self._flat_exp_args

    def get_exp_record(self) -> dict:
        """Return a dict with exp_args flattened and summary_info."""
        record = {"exp_dir": self.exp_dir}
        try:
            record.update(self.flat_exp_args)
        except FileNotFoundError:
            pass
        try:
            record.update(self.summary_info)
        except FileNotFoundError:
            pass
        return record

    @property
    def chat_video_path(self) -> Path:
        try:
            return next(self.exp_dir.glob("chat_video/*.webm"))
        except StopIteration:
            raise FileNotFoundError(f"No chat_video found in {self.exp_dir}")

    @property
    def task_video_path(self) -> Path:
        try:
            return next(self.exp_dir.glob("task_video/*.webm"))
        except StopIteration:
            raise FileNotFoundError(f"No task_video found in {self.exp_dir}")

    @property
    def combined_video_path(self) -> Path:
        return self.exp_dir / "combined_video.mp4"

    @property
    def logs(self):
        if self._logs is None:
            self._logs = (self.exp_dir / "experiment.log").read_text()
        return self._logs

    @property
    def status(self):


        """Return a dict with exp_args flattened and summary_info."""
        record = {"exp_dir": self.exp_dir}
        try:
            record.update(self.flat_exp_args)
        except FileNotFoundError:
            pass
        try:
            record.update(self.summary_info)
        except FileNotFoundError:
            pass
        return record

    @property
    def chat_video_path(self) -> Path:
        try:
            return next(self.exp_dir.glob("chat_video/*.webm"))
        except StopIteration:
            raise FileNotFoundError(f"No chat_video found in {self.exp_dir}")

    @property
    def task_video_path(self) -> Path:
        try:
            return next(self.exp_dir.glob("task_video/*.webm"))
        except StopIteration:
            raise FileNotFoundError(f"No task_video found in {self.exp_dir}")

    @property
    def combined_video_path(self) -> Path:
        return self.exp_dir / "combined_video.mp4"

    @property
    def logs(self):
        if self._logs is None:
            self._logs = (self.exp_dir / "experiment.log").read_text()
        return self._logs

    @property
    def status(self):
        """Return one of the following status:
        * "done": completed with no error
        * "error": completed with error
        * "incomplete": not completed yet (may be pending or just stalled)
        """
        try:
            summary_info = self.summary_info


            record.update(self.summary_info)
        except FileNotFoundError:
            pass
        return record

    @property
    def chat_video_path(self) -> Path:
        try:
            return next(self.exp_dir.glob("chat_video/*.webm"))
        except StopIteration:
            raise FileNotFoundError(f"No chat_video found in {self.exp_dir}")

    @property
    def task_video_path(self) -> Path:
        try:
            return next(self.exp_dir.glob("task_video/*.webm"))
        except StopIteration:
            raise FileNotFoundError(f"No task_video found in {self.exp_dir}")

    @property
    def combined_video_path(self) -> Path:
        return self.exp_dir / "combined_video.mp4"

    @property
    def logs(self):
        if self._logs is None:
            self._logs = (self.exp_dir / "experiment.log").read_text()
        return self._logs

    @property
    def status(self):
        """Return one of the following status:
        * "done": completed with no error
        * "error": completed with error
        * "incomplete": not completed yet (may be pending or just stalled)
        """
        try:
            summary_info = self.summary_info
        except FileNotFoundError:
            return "incomplete"

        if summary_info.get("err_msg", None) is not None:


logs

    @property
    def chat_video_path(self) -> Path:
        try:
            return next(self.exp_dir.glob("chat_video/*.webm"))
        except StopIteration:
            raise FileNotFoundError(f"No chat_video found in {self.exp_dir}")

    @property
    def task_video_path(self) -> Path:
        try:
            return next(self.exp_dir.glob("task_video/*.webm"))
        except StopIteration:
            raise FileNotFoundError(f"No task_video found in {self.exp_dir}")

    @property
    def combined_video_path(self) -> Path:
        return self.exp_dir / "combined_video.mp4"

    @property
    def logs(self):
        if self._logs is None:
            self._logs = (self.exp_dir / "experiment.log").read_text()
        return self._logs

    @property
    def status(self):
        """Return one of the following status:
        * "done": completed with no error
        * "error": completed with error
        * "incomplete": not completed yet (may be pending or just stalled)
        """
        try:
            summary_info = self.summary_info
        except FileNotFoundError:
            return "incomplete"

        if summary_info.get("err_msg", None) is not None:
            return "error"

        if summary_info.get("terminated", False) or summary_info.get("truncated", False):
            return "done"

        return "incomplete"


status

            raise FileNotFoundError(f"No chat_video found in {self.exp_dir}")

    @property
    def task_video_path(self) -> Path:
        try:
            return next(self.exp_dir.glob("task_video/*.webm"))
        except StopIteration:
            raise FileNotFoundError(f"No task_video found in {self.exp_dir}")

    @property
    def combined_video_path(self) -> Path:
        return self.exp_dir / "combined_video.mp4"

    @property
    def logs(self):
        if self._logs is None:
            self._logs = (self.exp_dir / "experiment.log").read_text()
        return self._logs

    @property
    def status(self):
        """Return one of the following status:
        * "done": completed with no error
        * "error": completed with error
        * "incomplete": not completed yet (may be pending or just stalled)
        """
        try:
            summary_info = self.summary_info
        except FileNotFoundError:
            return "incomplete"

        if summary_info.get("err_msg", None) is not None:
            return "error"

        if summary_info.get("terminated", False) or summary_info.get("truncated", False):
            return "done"

        return "incomplete"

EXP_RESULT_CACHE = {}

def get_exp_result(exp_dir) -> ExpResult:
    """Keep a cache of pre-loaded exp_results for faster loading"""
    exp_dir = str(exp_dir)  # make sure it's not a Path
    exp_result = EXP_RESULT_CACHE.get(exp_dir, None)
    if exp_result is None:
        exp_result = ExpResult(exp_dir)
        EXP_RESULT_CACHE[exp_dir] = exp_result
    return exp_result

def yield_all_exp_results(
    savedir_base: str | Path, progress_fn=tqdm, load_hidden=False, use_cache=True
):
    """Recursively find all experiments from savedir_base folder.


default

    exp_args_paths = []
    for exp_dir in savedir_base:
        exp_args_paths.extend(list(Path(exp_dir).glob("**/exp_args.pkl")))

    if progress_fn is not None:
        exp_args_paths = progress_fn(exp_args_paths, desc="Searching experiments directories.")

    for exp_args_path in exp_args_paths:
        exp_dir = exp_args_path.parent
        if not load_hidden:
            if exp_dir.name.startswith("_") or exp_dir.name.startswith("."):
                continue
        if use_cache:
            yield get_exp_result(exp_dir)
        else:
            yield ExpResult(exp_dir)

class DataclassJSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if is_dataclass(obj):
            return asdict(obj)
        if isinstance(obj, np.integer):
            return int(obj)
        if isinstance(obj, np.floating):
            return float(obj)
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return super().default(obj)

def _move_old_exp(exp_dir):
    """Move the old experiment directory to a new name."""
    exp_dir = Path(exp_dir)
    if exp_dir.exists():
        exp_dir.rename(exp_dir.with_name("_" + exp_dir.name))

def _get_env_name(task_name: str):
    """Register tasks if needed (lazy import) and return environment name."""

    # lazy benchmark import
    if task_name.startswith("miniwob"):
        import browsergym.miniwob
    elif task_name.startswith("workarena"):
        import browsergym.workarena
    elif task_name.startswith("webarena"):
        import browsergym.webarena
        import browsergym.webarenalite


import tiktoken

def count_tokens(text, model="gpt-4"):
    """Count the number of tokens in a text."""

    return len(tiktoken.encoding_for_model(model).encode(text))

def count_messages_token(messages, model="gpt-4"):
    """Count the number of tokens in a list of messages.

    Args:
        messages (list): a list of messages, each message can be a string or a
            list of dicts or an object with a content attribute.
        model (str): the model to use for tokenization.

    Returns:
        int: the number of tokens.
    """
    token_count = 0
    for message in messages:
        if hasattr(message, "content"):
            message = message.content
        elif isinstance(message, dict) and "content" in message:
            message = message["content"]

        if isinstance(message, str):
            token_count += count_tokens(message, model)
        # handles messages with image content
        elif isinstance(message, (list, tuple)):
            for part in message:
                if not isinstance(part, dict):
                    raise ValueError(
                        f"The message is expected to be a list of dicts, but got list of {type(message)}"
                    )
                if part["type"] == "text":
                    token_count += count_tokens(part["text"], model)
        else:
            raise ValueError(
                f"The message is expected to be a string or a list of dicts, but got {type(message)}"
            )
    return token_count


import tiktoken

def count_tokens(text, model="gpt-4"):
    """Count the number of tokens in a text."""

    return len(tiktoken.encoding_for_model(model).encode(text))

def count_messages_token(messages, model="gpt-4"):
    """Count the number of tokens in a list of messages.

    Args:
        messages (list): a list of messages, each message can be a string or a
            list of dicts or an object with a content attribute.
        model (str): the model to use for tokenization.

    Returns:
        int: the number of tokens.
    """
    token_count = 0
    for message in messages:
        if hasattr(message, "content"):
            message = message.content
        elif isinstance(message, dict) and "content" in message:
            message = message["content"]


import fnmatch
import logging
import random
import typing
from dataclasses import dataclass, field
from typing import Literal, Optional

import pandas as pd
from dataclasses_json import DataClassJsonMixin, config

from browsergym.core.action.highlevel import HighLevelActionSet
from browsergym.experiments.loop import EnvArgs

from .metadata.utils import (
    build_env_args_dependency_graphs,
    build_full_task_dependency_graph_from_metadata,
    extract_sparse_task_dependency_graph_from_subset,
    task_list_from_metadata,
)
from .utils import prepare_backend

logger = logging.getLogger(__name__)

@dataclass
class HighLevelActionSetArgs(DataClassJsonMixin):
    subsets: tuple[HighLevelActionSet.ActionSubset] = field(
        metadata=config(
            encoder=lambda x: list(x),
            decoder=lambda x: tuple(x),
        ),
    )
    # custom_actions: list[callable] | None  # non-serializable argument, not supported
    multiaction: bool = False
    strict: bool = False
    retry_with_force: bool = False
    demo_mode: Optional[Literal["off", "default", "all_blue", "only_visible_elements"]] = None

    def __post_init__(self):
        if isinstance(self.subsets, list):
            """Needs to be hashable for AgentLab when uniquely identifying agents."""
            self.subsets = tuple(self.subsets)

    def make_action_set(self):
        return HighLevelActionSet(
            subsets=self.subsets,
            custom_actions=None,
            multiaction=self.multiaction,
            strict=self.strict,
            retry_with_force=self.retry_with_force,
            demo_mode=self.demo_mode,
        )

BenchmarkBackend = Literal[
    "miniwob", "webarena", "visualwebarena", "workarena", "assistantbench", "weblinx"
]

@dataclass
class Benchmark(DataClassJsonMixin):
    name: str
    high_level_action_set_args: HighLevelActionSetArgs
    is_multi_tab: bool
    supports_parallel_seeds: bool
    env_args_list: list[EnvArgs]
    backends: list[BenchmarkBackend]
    task_metadata: Optional[pd.DataFrame] = field(
        default_factory=lambda: None,
        metadata=config(
            encoder=lambda df: df.to_dict(orient="records") if df is not None else None,
            decoder=lambda items: pd.DataFrame(items) if items is not None else None,
        ),
    )

    def __post_init__(self):
        # if no metadata is present, generate a dataframe with single "task_name" column
        if self.task_metadata is None:
            unique_task_names = list(set([env_args.task_name for env_args in self.env_args_list]))
            self.task_metadata = pd.DataFrame(
                [{"task_name": task_name} for task_name in unique_task_names]
            )
        # make sure all tasks in env_args are in the metadata
        metadata_tasks = list(self.task_metadata["task_name"])
        assert all([env_args.task_name in metadata_tasks for env_args in self.env_args_list])
        # check backend values
        for backend in self.backends:
            if backend not in typing.get_args(BenchmarkBackend):
                raise ValueError(
                    f"Unknown Benchmark backend {repr(backend)}. Available backends: {typing.get_args(BenchmarkBackend)}"
                )

    def prepare_backends(self):
        for backend in self.backends:
            logger.info(f"Preparing {backend} backend...")
            prepare_backend(backend)
            logger.info(f"{backend} backend ready")

    def subset_from_split(self, split: Literal["train", "valid", "test"]):
        split_column = "browsergym_split"

        # check for a split column in metadata
        if split_column not in self.task_metadata.columns:
            raise NotImplementedError(
                f"This benchmark does not provide default train/valid/test splits (missing a {repr(split_column)} column in task metadata)"
            )

        # recover the target split
        sub_benchmark = self.subset_from_regexp(split_column, regexp=f"^{split}$")
        sub_benchmark.name = f"{self.name}_{split}"

        # check that the split exists (non-empty task list)
        if not sub_benchmark.env_args_list:
            raise ValueError(f"The default {split} split for this benchmark is empty.")

        return sub_benchmark

    def subset_from_list(
        self,
        task_list: list[str],
        benchmark_name_suffix: Optional[str] = "custom",
        split: Optional[str] = None,
    ):
        """Create a sub-benchmark containing only the specified tasks.

        Args:
            task_list: List of task names to include in the sub-benchmark.
            benchmark_name_suffix: Optional suffix to append to the benchmark name. Defaults to "custom".
            split: Optional split name to append to the benchmark name. Useful for organization.

        Returns:
            Benchmark: A new benchmark instance containing only the specified tasks.

        Raises:
            ValueError: If the resulting task list is empty or if any specified task doesn't exist.
        """
        if not task_list:
            raise ValueError("Task list cannot be empty")

        # Convert task_list to set for more efficient lookups
        task_set = set(task_list)

        # Validate that all requested tasks exist in the original benchmark
        existing_tasks = {env_args.task_name for env_args in self.env_args_list}
        invalid_tasks = task_set - existing_tasks
        if invalid_tasks:
            raise ValueError(f"The following tasks do not exist in the benchmark: {invalid_tasks}")

        name = f"{self.name}_{benchmark_name_suffix}"
        if split:
            name += f"_{split}"

        sub_benchmark = Benchmark(
            name=name,
            high_level_action_set_args=self.high_level_action_set_args,
            is_multi_tab=self.is_multi_tab,
            supports_parallel_seeds=self.supports_parallel_seeds,
            backends=self.backends,
            env_args_list=[
                env_args for env_args in self.env_args_list if env_args.task_name in task_set
            ],
            task_metadata=self.task_metadata,
        )

        # This check is redundant now due to the validation above, but kept for safety
        if not sub_benchmark.env_args_list:
            raise ValueError(
                f"The custom {split if split else ''} split for this benchmark is empty."
            )

        return sub_benchmark

    def subset_from_glob(self, column, glob):
        subset = self.subset_from_regexp(column, regexp=fnmatch.translate(glob))
        subset.name = f"{self.name}[{column}={glob}]"
        return subset

    def subset_from_regexp(self, column, regexp):
        # extract the filtered task_name subset
        task_name_subset = task_list_from_metadata(self.task_metadata, {column: regexp})

        # return the sub benchmark
        return Benchmark(
            name=f"{self.name}[{column}=/{regexp}/]",
            high_level_action_set_args=self.high_level_action_set_args,
            is_multi_tab=self.is_multi_tab,
            supports_parallel_seeds=self.supports_parallel_seeds,
            backends=self.backends,
            env_args_list=[
                env_args
                for env_args in self.env_args_list
                if env_args.task_name in task_name_subset
            ],
            task_metadata=self.task_metadata,
        )

    def subset_from_task_ratio(self, ratio, seed):
        """Get a random subset of the tasks given a ratio and seed."""
        rng = random.Random(seed)
        task_names = list(set([env_args.task_name for env_args in self.env_args_list]))
        rng.shuffle(task_names)
        num_tasks = int(len(task_names) * ratio)
        task_name_subset = task_names[:num_tasks]

        return Benchmark(
            name=f"{self.name}[ratio={ratio}, seed={seed}]",
            high_level_action_set_args=self.high_level_action_set_args,
            is_multi_tab=self.is_multi_tab,
            supports_parallel_seeds=self.supports_parallel_seeds,
            backends=self.backends,
            env_args_list=[
                env_args
                for env_args in self.env_args_list
                if env_args.task_name in task_name_subset
            ],
            task_metadata=self.task_metadata,
        )

    def dependency_graph_over_tasks(self) -> dict[str, list[str]]:
        # recover all unique task_names present in the benchmark
        task_names = list(set([env_args.task_name for env_args in self.env_args_list]))

        # if "depends_on" column is missing, raise a warning and deal with it
        # (we don't want the "depends_on" column to be mandatory)
        if "depends_on" not in self.task_metadata.columns:
            logger.warning(
                f'This benchmark does not provide a dependency graph (missing a "depends_on" column in task metadata). Assuming no task dependencies.'
            )
            zero_dependencies = {task_name: [] for task_name in task_names}
            return zero_dependencies

        # recover the task dependency graph, for tasks in the benchmark only
        task_dependencies = extract_sparse_task_dependency_graph_from_subset(
            task_subset=task_names,
            parents=build_full_task_dependency_graph_from_metadata(
                task_metadata=self.task_metadata
            ),
        )

        return task_dependencies

    def dependency_graphs_over_env_args(self) -> list[dict[str, list[str]]]:
        """
        Returns a list of dependency graphs to be executed sequentially, typically with a full instance reset in-between.
        Ideally, a job scheduler should connect these graphs by injecting a reset task in-between each, which depends on all previous tasks being completed.
        """
        task_dependencies = self.dependency_graph_over_tasks()
        env_args_dependencies = build_env_args_dependency_graphs(
            env_args_list=self.env_args_list,
            task_dependencies=task_dependencies,
            supports_parallel_seeds=self.supports_parallel_seeds,
        )
        return env_args_dependencies


HighLevelActionSetArgs

from typing import Literal, Optional

import pandas as pd
from dataclasses_json import DataClassJsonMixin, config

from browsergym.core.action.highlevel import HighLevelActionSet
from browsergym.experiments.loop import EnvArgs

from .metadata.utils import (
    build_env_args_dependency_graphs,
    build_full_task_dependency_graph_from_metadata,
    extract_sparse_task_dependency_graph_from_subset,
    task_list_from_metadata,
)
from .utils import prepare_backend

logger = logging.getLogger(__name__)

@dataclass
class HighLevelActionSetArgs(DataClassJsonMixin):
    subsets: tuple[HighLevelActionSet.ActionSubset] = field(
        metadata=config(
            encoder=lambda x: list(x),
            decoder=lambda x: tuple(x),
        ),
    )
    # custom_actions: list[callable] | None  # non-serializable argument, not supported
    multiaction: bool = False
    strict: bool = False
    retry_with_force: bool = False
    demo_mode: Optional[Literal["off", "default", "all_blue", "only_visible_elements"]] = None

    def __post_init__(self):
        if isinstance(self.subsets, list):
            """Needs to be hashable for AgentLab when uniquely identifying agents."""
            self.subsets = tuple(self.subsets)

    def make_action_set(self):
        return HighLevelActionSet(
            subsets=self.subsets,
            custom_actions=None,
            multiaction=self.multiaction,
            strict=self.strict,
            retry_with_force=self.retry_with_force,
            demo_mode=self.demo_mode,
        )

BenchmarkBackend = Literal[
    "miniwob", "webarena", "visualwebarena", "workarena", "assistantbench", "weblinx"
]

@dataclass
class Benchmark(DataClassJsonMixin):
    name: str
    high_level_action_set_args: HighLevelActionSetArgs
    is_multi_tab: bool
    supports_parallel_seeds: bool
    env_args_list: list[EnvArgs]
    backends: list[BenchmarkBackend]
    task_metadata: Optional[pd.DataFrame] = field(
        default_factory=lambda: None,
        metadata=config(
            encoder=lambda df: df.to_dict(orient="records") if df is not None else None,
            decoder=lambda items: pd.DataFrame(items) if items is not None else None,


Benchmark

            """Needs to be hashable for AgentLab when uniquely identifying agents."""
            self.subsets = tuple(self.subsets)

    def make_action_set(self):
        return HighLevelActionSet(
            subsets=self.subsets,
            custom_actions=None,
            multiaction=self.multiaction,
            strict=self.strict,
            retry_with_force=self.retry_with_force,
            demo_mode=self.demo_mode,
        )

BenchmarkBackend = Literal[
    "miniwob", "webarena", "visualwebarena", "workarena", "assistantbench", "weblinx"
]

@dataclass
class Benchmark(DataClassJsonMixin):
    name: str
    high_level_action_set_args: HighLevelActionSetArgs
    is_multi_tab: bool
    supports_parallel_seeds: bool
    env_args_list: list[EnvArgs]
    backends: list[BenchmarkBackend]
    task_metadata: Optional[pd.DataFrame] = field(
        default_factory=lambda: None,
        metadata=config(
            encoder=lambda df: df.to_dict(orient="records") if df is not None else None,
            decoder=lambda items: pd.DataFrame(items) if items is not None else None,
        ),
    )

    def __post_init__(self):
        # if no metadata is present, generate a dataframe with single "task_name" column
        if self.task_metadata is None:
            unique_task_names = list(set([env_args.task_name for env_args in self.env_args_list]))
            self.task_metadata = pd.DataFrame(
                [{"task_name": task_name} for task_name in unique_task_names]
            )
        # make sure all tasks in env_args are in the metadata
        metadata_tasks = list(self.task_metadata["task_name"])
        assert all([env_args.task_name in metadata_tasks for env_args in self.env_args_list])
        # check backend values
        for backend in self.backends:
            if backend not in typing.get_args(BenchmarkBackend):
                raise ValueError(
                    f"Unknown Benchmark backend {repr(backend)}. Available backends: {typing.get_args(BenchmarkBackend)}"
                )

    def prepare_backends(self):
        for backend in self.backends:
            logger.info(f"Preparing {backend} backend...")
            prepare_backend(backend)
            logger.info(f"{backend} backend ready")

    def subset_from_split(self, split: Literal["train", "valid", "test"]):
        split_column = "browsergym_split"

        # check for a split column in metadata
        if split_column not in self.task_metadata.columns:
            raise NotImplementedError(
                f"This benchmark does not provide default train/valid/test splits (missing a {repr(split_column)} column in task metadata)"
            )

        # recover the target split
        sub_benchmark = self.subset_from_regexp(split_column, regexp=f"^{split}$")
        sub_benchmark.name = f"{self.name}_{split}"

        # check that the split exists (non-empty task list)
        if not sub_benchmark.env_args_list:
            raise ValueError(f"The default {split} split for this benchmark is empty.")

        return sub_benchmark

    def subset_from_list(
        self,
        task_list: list[str],
        benchmark_name_suffix: Optional[str] = "custom",
        split: Optional[str] = None,
    ):
        """Create a sub-benchmark containing only the specified tasks.

        Args:
            task_list: List of task names to include in the sub-benchmark.
            benchmark_name_suffix: Optional suffix to append to the benchmark name. Defaults to "custom".
            split: Optional split name to append to the benchmark name. Useful for organization.

        Returns:
            Benchmark: A new benchmark instance containing only the specified tasks.

        Raises:
            ValueError: If the resulting task list is empty or if any specified task doesn't exist.
        """
        if not task_list:
            raise ValueError("Task list cannot be empty")

        # Convert task_list to set for more efficient lookups
        task_set = set(task_list)

        # Validate that all requested tasks exist in the original benchmark
        existing_tasks = {env_args.task_name for env_args in self.env_args_list}
        invalid_tasks = task_set - existing_tasks
        if invalid_tasks:
            raise ValueError(f"The following tasks do not exist in the benchmark: {invalid_tasks}")

        name = f"{self.name}_{benchmark_name_suffix}"
        if split:
            name += f"_{split}"

        sub_benchmark = Benchmark(
            name=name,
            high_level_action_set_args=self.high_level_action_set_args,
            is_multi_tab=self.is_multi_tab,
            supports_parallel_seeds=self.supports_parallel_seeds,
            backends=self.backends,
            env_args_list=[
                env_args for env_args in self.env_args_list if env_args.task_name in task_set
            ],
            task_metadata=self.task_metadata,
        )

        # This check is redundant now due to the validation above, but kept for safety
        if not sub_benchmark.env_args_list:
            raise ValueError(
                f"The custom {split if split else ''} split for this benchmark is empty."
            )

        return sub_benchmark

    def subset_from_glob(self, column, glob):
        subset = self.subset_from_regexp(column, regexp=fnmatch.translate(glob))
        subset.name = f"{self.name}[{column}={glob}]"
        return subset

    def subset_from_regexp(self, column, regexp):
        # extract the filtered task_name subset
        task_name_subset = task_list_from_metadata(self.task_metadata, {column: regexp})

        # return the sub benchmark
        return Benchmark(
            name=f"{self.name}[{column}=/{regexp}/]",
            high_level_action_set_args=self.high_level_action_set_args,
            is_multi_tab=self.is_multi_tab,
            supports_parallel_seeds=self.supports_parallel_seeds,
            backends=self.backends,
            env_args_list=[
                env_args
                for env_args in self.env_args_list
                if env_args.task_name in task_name_subset
            ],
            task_metadata=self.task_metadata,
        )

    def subset_from_task_ratio(self, ratio, seed):
        """Get a random subset of the tasks given a ratio and seed."""
        rng = random.Random(seed)
        task_names = list(set([env_args.task_name for env_args in self.env_args_list]))
        rng.shuffle(task_names)
        num_tasks = int(len(task_names) * ratio)
        task_name_subset = task_names[:num_tasks]

        return Benchmark(
            name=f"{self.name}[ratio={ratio}, seed={seed}]",
            high_level_action_set_args=self.high_level_action_set_args,
            is_multi_tab=self.is_multi_tab,
            supports_parallel_seeds=self.supports_parallel_seeds,
            backends=self.backends,
            env_args_list=[
                env_args
                for env_args in self.env_args_list
                if env_args.task_name in task_name_subset
            ],
            task_metadata=self.task_metadata,
        )

    def dependency_graph_over_tasks(self) -> dict[str, list[str]]:
        # recover all unique task_names present in the benchmark
        task_names = list(set([env_args.task_name for env_args in self.env_args_list]))

        # if "depends_on" column is missing, raise a warning and deal with it
        # (we don't want the "depends_on" column to be mandatory)
        if "depends_on" not in self.task_metadata.columns:
            logger.warning(
                f'This benchmark does not provide a dependency graph (missing a "depends_on" column in task metadata). Assuming no task dependencies.'
            )
            zero_dependencies = {task_name: [] for task_name in task_names}
            return zero_dependencies

        # recover the task dependency graph, for tasks in the benchmark only
        task_dependencies = extract_sparse_task_dependency_graph_from_subset(
            task_subset=task_names,
            parents=build_full_task_dependency_graph_from_metadata(
                task_metadata=self.task_metadata
            ),
        )

        return task_dependencies

    def dependency_graphs_over_env_args(self) -> list[dict[str, list[str]]]:
        """
        Returns a list of dependency graphs to be executed sequentially, typically with a full instance reset in-between.
        Ideally, a job scheduler should connect these graphs by injecting a reset task in-between each, which depends on all previous tasks being completed.
        """
        task_dependencies = self.dependency_graph_over_tasks()
        env_args_dependencies = build_env_args_dependency_graphs(
            env_args_list=self.env_args_list,
            task_dependencies=task_dependencies,
            supports_parallel_seeds=self.supports_parallel_seeds,
        )
        return env_args_dependencies


    "miniwob", "webarena", "visualwebarena", "workarena", "assistantbench", "weblinx"
]

@dataclass
class Benchmark(DataClassJsonMixin):
    name: str
    high_level_action_set_args: HighLevelActionSetArgs
    is_multi_tab: bool
    supports_parallel_seeds: bool
    env_args_list: list[EnvArgs]
    backends: list[BenchmarkBackend]
    task_metadata: Optional[pd.DataFrame] = field(
        default_factory=lambda: None,
        metadata=config(
            encoder=lambda df: df.to_dict(orient="records") if df is not None else None,
            decoder=lambda items: pd.DataFrame(items) if items is not None else None,
        ),
    )

    def __post_init__(self):
        # if no metadata is present, generate a dataframe with single "task_name" column
        if self.task_metadata is None:
            unique_task_names = list(set([env_args.task_name for env_args in self.env_args_list]))
            self.task_metadata = pd.DataFrame(
                [{"task_name": task_name} for task_name in unique_task_names]
            )
        # make sure all tasks in env_args are in the metadata
        metadata_tasks = list(self.task_metadata["task_name"])
        assert all([env_args.task_name in metadata_tasks for env_args in self.env_args_list])
        # check backend values
        for backend in self.backends:
            if backend not in typing.get_args(BenchmarkBackend):
                raise ValueError(
                    f"Unknown Benchmark backend {repr(backend)}. Available backends: {typing.get_args(BenchmarkBackend)}"
                )

    def prepare_backends(self):
        for backend in self.backends:
            logger.info(f"Preparing {backend} backend...")
            prepare_backend(backend)
            logger.info(f"{backend} backend ready")

    def subset_from_split(self, split: Literal["train", "valid", "test"]):
        split_column = "browsergym_split"

        # check for a split column in metadata
        if split_column not in self.task_metadata.columns:
            raise NotImplementedError(
                f"This benchmark does not provide default train/valid/test splits (missing a {repr(split_column)} column in task metadata)"
            )

        # recover the target split
        sub_benchmark = self.subset_from_regexp(split_column, regexp=f"^{split}$")
        sub_benchmark.name = f"{self.name}_{split}"


@dataclass
class HighLevelActionSetArgs(DataClassJsonMixin):
    subsets: tuple[HighLevelActionSet.ActionSubset] = field(
        metadata=config(
            encoder=lambda x: list(x),
            decoder=lambda x: tuple(x),
        ),
    )
    # custom_actions: list[callable] | None  # non-serializable argument, not supported
    multiaction: bool = False
    strict: bool = False
    retry_with_force: bool = False
    demo_mode: Optional[Literal["off", "default", "all_blue", "only_visible_elements"]] = None

    def __post_init__(self):
        if isinstance(self.subsets, list):
            """Needs to be hashable for AgentLab when uniquely identifying agents."""
            self.subsets = tuple(self.subsets)

    def make_action_set(self):
        return HighLevelActionSet(
            subsets=self.subsets,
            custom_actions=None,
            multiaction=self.multiaction,
            strict=self.strict,
            retry_with_force=self.retry_with_force,
            demo_mode=self.demo_mode,
        )

BenchmarkBackend = Literal[
    "miniwob", "webarena", "visualwebarena", "workarena", "assistantbench", "weblinx"
]

@dataclass
class Benchmark(DataClassJsonMixin):
    name: str
    high_level_action_set_args: HighLevelActionSetArgs
    is_multi_tab: bool
    supports_parallel_seeds: bool
    env_args_list: list[EnvArgs]
    backends: list[BenchmarkBackend]
    task_metadata: Optional[pd.DataFrame] = field(
        default_factory=lambda: None,
        metadata=config(
            encoder=lambda df: df.to_dict(orient="records") if df is not None else None,
            decoder=lambda items: pd.DataFrame(items) if items is not None else None,


        ),
    )

    def __post_init__(self):
        # if no metadata is present, generate a dataframe with single "task_name" column
        if self.task_metadata is None:
            unique_task_names = list(set([env_args.task_name for env_args in self.env_args_list]))
            self.task_metadata = pd.DataFrame(
                [{"task_name": task_name} for task_name in unique_task_names]
            )
        # make sure all tasks in env_args are in the metadata
        metadata_tasks = list(self.task_metadata["task_name"])
        assert all([env_args.task_name in metadata_tasks for env_args in self.env_args_list])
        # check backend values
        for backend in self.backends:
            if backend not in typing.get_args(BenchmarkBackend):
                raise ValueError(
                    f"Unknown Benchmark backend {repr(backend)}. Available backends: {typing.get_args(BenchmarkBackend)}"
                )

    def prepare_backends(self):
        for backend in self.backends:
            logger.info(f"Preparing {backend} backend...")
            prepare_backend(backend)
            logger.info(f"{backend} backend ready")

    def subset_from_split(self, split: Literal["train", "valid", "test"]):
        split_column = "browsergym_split"

        # check for a split column in metadata
        if split_column not in self.task_metadata.columns:
            raise NotImplementedError(
                f"This benchmark does not provide default train/valid/test splits (missing a {repr(split_column)} column in task metadata)"
            )

        # recover the target split
        sub_benchmark = self.subset_from_regexp(split_column, regexp=f"^{split}$")
        sub_benchmark.name = f"{self.name}_{split}"

        # check that the split exists (non-empty task list)
        if not sub_benchmark.env_args_list:
            raise ValueError(f"The default {split} split for this benchmark is empty.")

        return sub_benchmark


            unique_task_names = list(set([env_args.task_name for env_args in self.env_args_list]))
            self.task_metadata = pd.DataFrame(
                [{"task_name": task_name} for task_name in unique_task_names]
            )
        # make sure all tasks in env_args are in the metadata
        metadata_tasks = list(self.task_metadata["task_name"])
        assert all([env_args.task_name in metadata_tasks for env_args in self.env_args_list])
        # check backend values
        for backend in self.backends:
            if backend not in typing.get_args(BenchmarkBackend):
                raise ValueError(
                    f"Unknown Benchmark backend {repr(backend)}. Available backends: {typing.get_args(BenchmarkBackend)}"
                )

    def prepare_backends(self):
        for backend in self.backends:
            logger.info(f"Preparing {backend} backend...")
            prepare_backend(backend)
            logger.info(f"{backend} backend ready")

    def subset_from_split(self, split: Literal["train", "valid", "test"]):
        split_column = "browsergym_split"

        # check for a split column in metadata
        if split_column not in self.task_metadata.columns:
            raise NotImplementedError(
                f"This benchmark does not provide default train/valid/test splits (missing a {repr(split_column)} column in task metadata)"
            )

        # recover the target split
        sub_benchmark = self.subset_from_regexp(split_column, regexp=f"^{split}$")
        sub_benchmark.name = f"{self.name}_{split}"

        # check that the split exists (non-empty task list)
        if not sub_benchmark.env_args_list:
            raise ValueError(f"The default {split} split for this benchmark is empty.")

        return sub_benchmark

    def subset_from_list(
        self,
        task_list: list[str],
        benchmark_name_suffix: Optional[str] = "custom",
        split: Optional[str] = None,
    ):
        """Create a sub-benchmark containing only the specified tasks.

        Args:
            task_list: List of task names to include in the sub-benchmark.
            benchmark_name_suffix: Optional suffix to append to the benchmark name. Defaults to "custom".
            split: Optional split name to append to the benchmark name. Useful for organization.

        Returns:
            Benchmark: A new benchmark instance containing only the specified tasks.

        Raises:
            ValueError: If the resulting task list is empty or if any specified task doesn't exist.
        """


    def subset_from_split(self, split: Literal["train", "valid", "test"]):
        split_column = "browsergym_split"

        # check for a split column in metadata
        if split_column not in self.task_metadata.columns:
            raise NotImplementedError(
                f"This benchmark does not provide default train/valid/test splits (missing a {repr(split_column)} column in task metadata)"
            )

        # recover the target split
        sub_benchmark = self.subset_from_regexp(split_column, regexp=f"^{split}$")
        sub_benchmark.name = f"{self.name}_{split}"

        # check that the split exists (non-empty task list)
        if not sub_benchmark.env_args_list:
            raise ValueError(f"The default {split} split for this benchmark is empty.")

        return sub_benchmark

    def subset_from_list(
        self,
        task_list: list[str],
        benchmark_name_suffix: Optional[str] = "custom",
        split: Optional[str] = None,
    ):
        """Create a sub-benchmark containing only the specified tasks.

        Args:
            task_list: List of task names to include in the sub-benchmark.
            benchmark_name_suffix: Optional suffix to append to the benchmark name. Defaults to "custom".
            split: Optional split name to append to the benchmark name. Useful for organization.

        Returns:
            Benchmark: A new benchmark instance containing only the specified tasks.

        Raises:
            ValueError: If the resulting task list is empty or if any specified task doesn't exist.
        """
        if not task_list:
            raise ValueError("Task list cannot be empty")

        # Convert task_list to set for more efficient lookups
        task_set = set(task_list)

        # Validate that all requested tasks exist in the original benchmark
        existing_tasks = {env_args.task_name for env_args in self.env_args_list}
        invalid_tasks = task_set - existing_tasks
        if invalid_tasks:
            raise ValueError(f"The following tasks do not exist in the benchmark: {invalid_tasks}")

        name = f"{self.name}_{benchmark_name_suffix}"
        if split:
            name += f"_{split}"

        sub_benchmark = Benchmark(
            name=name,
            high_level_action_set_args=self.high_level_action_set_args,
            is_multi_tab=self.is_multi_tab,
            supports_parallel_seeds=self.supports_parallel_seeds,
            backends=self.backends,
            env_args_list=[
                env_args for env_args in self.env_args_list if env_args.task_name in task_set
            ],
            task_metadata=self.task_metadata,
        )

        # This check is redundant now due to the validation above, but kept for safety
        if not sub_benchmark.env_args_list:
            raise ValueError(
                f"The custom {split if split else ''} split for this benchmark is empty."
            )

        return sub_benchmark

    def subset_from_glob(self, column, glob):
        subset = self.subset_from_regexp(column, regexp=fnmatch.translate(glob))
        subset.name = f"{self.name}[{column}={glob}]"
        return subset

    def subset_from_regexp(self, column, regexp):
        # extract the filtered task_name subset
        task_name_subset = task_list_from_metadata(self.task_metadata, {column: regexp})

        # return the sub benchmark
        return Benchmark(
            name=f"{self.name}[{column}=/{regexp}/]",
            high_level_action_set_args=self.high_level_action_set_args,
            is_multi_tab=self.is_multi_tab,
            supports_parallel_seeds=self.supports_parallel_seeds,
            backends=self.backends,
            env_args_list=[
                env_args
                for env_args in self.env_args_list


        sub_benchmark = Benchmark(
            name=name,
            high_level_action_set_args=self.high_level_action_set_args,
            is_multi_tab=self.is_multi_tab,
            supports_parallel_seeds=self.supports_parallel_seeds,
            backends=self.backends,
            env_args_list=[
                env_args for env_args in self.env_args_list if env_args.task_name in task_set
            ],
            task_metadata=self.task_metadata,
        )

        # This check is redundant now due to the validation above, but kept for safety
        if not sub_benchmark.env_args_list:
            raise ValueError(
                f"The custom {split if split else ''} split for this benchmark is empty."
            )

        return sub_benchmark

    def subset_from_glob(self, column, glob):
        subset = self.subset_from_regexp(column, regexp=fnmatch.translate(glob))
        subset.name = f"{self.name}[{column}={glob}]"
        return subset

    def subset_from_regexp(self, column, regexp):
        # extract the filtered task_name subset
        task_name_subset = task_list_from_metadata(self.task_metadata, {column: regexp})

        # return the sub benchmark
        return Benchmark(
            name=f"{self.name}[{column}=/{regexp}/]",
            high_level_action_set_args=self.high_level_action_set_args,
            is_multi_tab=self.is_multi_tab,
            supports_parallel_seeds=self.supports_parallel_seeds,
            backends=self.backends,
            env_args_list=[
                env_args
                for env_args in self.env_args_list
                if env_args.task_name in task_name_subset
            ],
            task_metadata=self.task_metadata,
        )


            backends=self.backends,
            env_args_list=[
                env_args for env_args in self.env_args_list if env_args.task_name in task_set
            ],
            task_metadata=self.task_metadata,
        )

        # This check is redundant now due to the validation above, but kept for safety
        if not sub_benchmark.env_args_list:
            raise ValueError(
                f"The custom {split if split else ''} split for this benchmark is empty."
            )

        return sub_benchmark

    def subset_from_glob(self, column, glob):
        subset = self.subset_from_regexp(column, regexp=fnmatch.translate(glob))
        subset.name = f"{self.name}[{column}={glob}]"
        return subset

    def subset_from_regexp(self, column, regexp):
        # extract the filtered task_name subset
        task_name_subset = task_list_from_metadata(self.task_metadata, {column: regexp})

        # return the sub benchmark
        return Benchmark(
            name=f"{self.name}[{column}=/{regexp}/]",
            high_level_action_set_args=self.high_level_action_set_args,
            is_multi_tab=self.is_multi_tab,
            supports_parallel_seeds=self.supports_parallel_seeds,
            backends=self.backends,
            env_args_list=[
                env_args
                for env_args in self.env_args_list
                if env_args.task_name in task_name_subset
            ],
            task_metadata=self.task_metadata,
        )

    def subset_from_task_ratio(self, ratio, seed):
        """Get a random subset of the tasks given a ratio and seed."""
        rng = random.Random(seed)
        task_names = list(set([env_args.task_name for env_args in self.env_args_list]))
        rng.shuffle(task_names)
        num_tasks = int(len(task_names) * ratio)
        task_name_subset = task_names[:num_tasks]

        return Benchmark(
            name=f"{self.name}[ratio={ratio}, seed={seed}]",
            high_level_action_set_args=self.high_level_action_set_args,
            is_multi_tab=self.is_multi_tab,
            supports_parallel_seeds=self.supports_parallel_seeds,
            backends=self.backends,
            env_args_list=[
                env_args
                for env_args in self.env_args_list
                if env_args.task_name in task_name_subset
            ],


    def subset_from_regexp(self, column, regexp):
        # extract the filtered task_name subset
        task_name_subset = task_list_from_metadata(self.task_metadata, {column: regexp})

        # return the sub benchmark
        return Benchmark(
            name=f"{self.name}[{column}=/{regexp}/]",
            high_level_action_set_args=self.high_level_action_set_args,
            is_multi_tab=self.is_multi_tab,
            supports_parallel_seeds=self.supports_parallel_seeds,
            backends=self.backends,
            env_args_list=[
                env_args
                for env_args in self.env_args_list
                if env_args.task_name in task_name_subset
            ],
            task_metadata=self.task_metadata,
        )

    def subset_from_task_ratio(self, ratio, seed):
        """Get a random subset of the tasks given a ratio and seed."""
        rng = random.Random(seed)
        task_names = list(set([env_args.task_name for env_args in self.env_args_list]))
        rng.shuffle(task_names)
        num_tasks = int(len(task_names) * ratio)
        task_name_subset = task_names[:num_tasks]

        return Benchmark(
            name=f"{self.name}[ratio={ratio}, seed={seed}]",
            high_level_action_set_args=self.high_level_action_set_args,
            is_multi_tab=self.is_multi_tab,
            supports_parallel_seeds=self.supports_parallel_seeds,
            backends=self.backends,
            env_args_list=[
                env_args
                for env_args in self.env_args_list
                if env_args.task_name in task_name_subset
            ],
            task_metadata=self.task_metadata,
        )

    def dependency_graph_over_tasks(self) -> dict[str, list[str]]:
        # recover all unique task_names present in the benchmark
        task_names = list(set([env_args.task_name for env_args in self.env_args_list]))

        # if "depends_on" column is missing, raise a warning and deal with it
        # (we don't want the "depends_on" column to be mandatory)
        if "depends_on" not in self.task_metadata.columns:
            logger.warning(
                f'This benchmark does not provide a dependency graph (missing a "depends_on" column in task metadata). Assuming no task dependencies.'
            )
            zero_dependencies = {task_name: [] for task_name in task_names}
            return zero_dependencies

        # recover the task dependency graph, for tasks in the benchmark only
        task_dependencies = extract_sparse_task_dependency_graph_from_subset(
            task_subset=task_names,
            parents=build_full_task_dependency_graph_from_metadata(
                task_metadata=self.task_metadata
            ),


        rng = random.Random(seed)
        task_names = list(set([env_args.task_name for env_args in self.env_args_list]))
        rng.shuffle(task_names)
        num_tasks = int(len(task_names) * ratio)
        task_name_subset = task_names[:num_tasks]

        return Benchmark(
            name=f"{self.name}[ratio={ratio}, seed={seed}]",
            high_level_action_set_args=self.high_level_action_set_args,
            is_multi_tab=self.is_multi_tab,
            supports_parallel_seeds=self.supports_parallel_seeds,
            backends=self.backends,
            env_args_list=[
                env_args
                for env_args in self.env_args_list
                if env_args.task_name in task_name_subset
            ],
            task_metadata=self.task_metadata,
        )

    def dependency_graph_over_tasks(self) -> dict[str, list[str]]:
        # recover all unique task_names present in the benchmark
        task_names = list(set([env_args.task_name for env_args in self.env_args_list]))

        # if "depends_on" column is missing, raise a warning and deal with it
        # (we don't want the "depends_on" column to be mandatory)
        if "depends_on" not in self.task_metadata.columns:
            logger.warning(
                f'This benchmark does not provide a dependency graph (missing a "depends_on" column in task metadata). Assuming no task dependencies.'
            )
            zero_dependencies = {task_name: [] for task_name in task_names}
            return zero_dependencies

        # recover the task dependency graph, for tasks in the benchmark only
        task_dependencies = extract_sparse_task_dependency_graph_from_subset(
            task_subset=task_names,
            parents=build_full_task_dependency_graph_from_metadata(
                task_metadata=self.task_metadata
            ),
        )

        return task_dependencies

    def dependency_graphs_over_env_args(self) -> list[dict[str, list[str]]]:
        """
        Returns a list of dependency graphs to be executed sequentially, typically with a full instance reset in-between.
        Ideally, a job scheduler should connect these graphs by injecting a reset task in-between each, which depends on all previous tasks being completed.
        """
        task_dependencies = self.dependency_graph_over_tasks()
        env_args_dependencies = build_env_args_dependency_graphs(
            env_args_list=self.env_args_list,
            task_dependencies=task_dependencies,
            supports_parallel_seeds=self.supports_parallel_seeds,
        )
        return env_args_dependencies