adding resample benchmark objects

src/agentlab/experiments/custom_benchmark.py

@@ -0,0 +1,135 @@
+from abc import abstractmethod
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional
+
+import benchmarks
+import numpy as np
+import pandas as pd
+from bgym import Benchmark, EnvArgs, HighLevelActionSetArgs
+from browsergym.experiments.benchmark.base import BenchmarkBackend
+from browsergym.experiments.benchmark.utils import make_env_args_list_from_repeat_tasks
+from dataclasses_json import DataClassJsonMixin, config
+from torch import threshold
+
+from agentlab.analyze.inspect_results import load_result_df
+from agentlab.experiments.study import Study
+
+
+@dataclass
+class ResampleBenchmark(Benchmark):
+    exp_dir: Path = None
+    name: str = None
+    high_level_action_set_args: HighLevelActionSetArgs = None
+    is_multi_tab: bool = None
+    supports_parallel_seeds: bool = None
+    env_args_list: list[EnvArgs] = None
+    backends: list[BenchmarkBackend] = None
+    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):
+        assert self.exp_dir is not None
+        study = Study.load(self.exp_dir)
+        benchmark = study.benchmark
+
+        self.name = f"resample-{benchmark.name}"
+        self.high_level_action_set_args = benchmark.high_level_action_set_args
+        self.is_multi_tab = benchmark.is_multi_tab
+        self.supports_parallel_seeds = benchmark.supports_parallel_seeds
+        self.backends = benchmark.backends
+        # we discard the task_metadata to create new ones in post_init
+
+        values = self.evaluate(study, benchmark.env_args_list)
+        selected_env_args = self.select(values, benchmark.env_args_list)
+
+        if len(selected_env_args) == 0:
+            raise ValueError("No env_args selected, lower restrictions")
+
+        self.env_args_list = selected_env_args
+
+        super().__post_init__()
+
+    @abstractmethod
+    def evaluate(self, study, env_args_list):
+        pass
+
+    @abstractmethod
+    def select(self, values, env_args_list):
+        pass
+
+
+@dataclass
+class AllTasksBenchmark(ResampleBenchmark):
+    def evaluate(self, study, env_args_list):
+        return [0] * len(env_args_list)
+
+    def select(self, values, env_args_list):
+        return env_args_list
+
+
+@dataclass
+class HighVarianceBenchmark(ResampleBenchmark):
+    threshold: float = 0.2
+
+    def evaluate(self, study: Study, env_args_list):
+        result_df = load_result_df(study.dir)
+        return dict(result_df.groupby("env.task_name")["cum_reward"].std())
+
+    def select(self, values, env_args_list):
+        selected_env_args = []
+        for env_args in env_args_list:
+            if values[env_args.task_name] > self.threshold:
+                selected_env_args.append(env_args)
+        return selected_env_args
+
+
+@dataclass
+class StochasticHighVarianceBenchmark(ResampleBenchmark):
+    regulation_threshold: float = 0.1
+    total_seeds = 600
+    min_seeds = 2
+    random_seed = 42
+
+    def evaluate(self, study: Study, env_args_list):
+        result_df = load_result_df(study.dir)
+        var = result_df.groupby("env.task_name")["cum_reward"].var()
+        probs = dict((var + self.regulation_threshold) / (var + self.regulation_threshold).sum())
+        return probs
+
+    def select(self, values, env_args_list: list[EnvArgs]):
+        selected_env_args = []
+        max_steps = env_args_list[0].max_steps
+        for task_name, p in values.items():
+            # ceil to avoid missing seeds
+            n_seeds = np.random.RandomState(self.random_seed).poisson(p * self.total_seeds)
+            n_seeds = max(n_seeds, self.min_seeds)
+            for seed in np.random.RandomState(self.random_seed).randint(0, 2**32, n_seeds):
+                selected_env_args.append(
+                    EnvArgs(
+                        task_name=task_name,
+                        task_seed=int(seed),
+                        max_steps=max_steps,
+                        headless=True,
+                        record_video=False,
+                        wait_for_user_message=False,
+                        viewport=None,
+                        slow_mo=None,
+                        storage_state=None,
+                        task_kwargs=None,
+                    )
+                )
+        return selected_env_args
+
+
+if __name__ == "__main__":
+    exp_dir = Path(
+        "/Users/t.lesellierdechezell/agentlab_results/2025-03-04_14-43-48_genericagent-gpt-4o-mini-2024-07-18-on-miniwob"
+    )
+    benchmark = StochasticHighVarianceBenchmark(exp_dir=exp_dir)
+    print(benchmark.env_args_list)