Update run.py and config files

Author: MarcCote

debug_gym/agents/base_agent.py

@@ -371,5 +371,5 @@ def create_agent(
     if agent_args is None:
         raise ValueError("Either agent_args or config must be provided.")
 
-    agent = agent_class(args=agent_args, **agent_kwargs)
+    agent = agent_class(agent_args=agent_args, **agent_kwargs)
     return agent

debug_gym/agents/utils.py

@@ -108,15 +108,6 @@ def load_config():
     with open(args.config_file) as reader:
         config = yaml.safe_load(reader)
 
-    # Parse overriden params.
-    for param in args.params:
-        fqn_key, value = param.split("=")
-        entry_to_change = config
-        keys = fqn_key.split(".")
-        for k in keys[:-1]:
-            entry_to_change = entry_to_change[k]
-        entry_to_change[keys[-1]] = yaml.safe_load(value)
-
     available_agents = [item for item in list(config.keys()) if item != "base"]
 
     if not args.agent:
@@ -130,14 +121,25 @@ def load_config():
     if "base" in config:
         # base config is specified (shared across agents)
         return_config = config["base"]
-        agent_specific_config = config[args.agent]
-        for key in agent_specific_config:
-            # override base config with agent specific config
-            return_config[key] = agent_specific_config[key]
+        # Override base config with agent specific config
+        for key, value in config[args.agent].items():
+            return_config[key] = value
     else:
         # base config is not specified
         return_config = config[args.agent]
 
+    # Parse overriden params.
+    for param in args.params:
+        fqn_key, value = param.split("=")
+        entry_to_change = return_config
+        keys = fqn_key.split(".")
+        for k in keys[:-1]:
+            if k not in entry_to_change:
+                entry_to_change[k] = {}
+
+            entry_to_change = entry_to_change[k]
+        entry_to_change[keys[-1]] = yaml.safe_load(value)
+
     # assume agent type is the key if not specified by the user
     if not return_config.get("agent_type"):
         return_config["agent_type"] = args.agent

debug_gym/gym/envs/__init__.py

@@ -6,12 +6,11 @@
 from debug_gym.gym.envs.swe_bench import SWEBenchEnv
 from debug_gym.gym.envs.swe_bench_debug import SWEBenchDebugEnv
 from debug_gym.gym.envs.swe_smith import SWESmithEnv
+from debug_gym.logger import DebugGymLogger
 
 
 def select_env(env_type: str = None) -> type[RepoEnv]:
     match env_type:
-        case None:
-            return RepoEnv
         case "local":
             return LocalEnv
         case "aider":
@@ -27,4 +26,20 @@ def select_env(env_type: str = None) -> type[RepoEnv]:
         case "r2egym":
             return R2EGymEnv
         case _:
-            raise ValueError(f"Unknown benchmark {env_type}")
+            raise ValueError(f"Unknown environment {env_type}")
+
+
+def load_dataset(config: dict, logger: DebugGymLogger | None = None) -> dict:
+    """Load dataset based on the given config."""
+    if config.get("type") is None:
+        raise ValueError("Dataset config must specify 'type' field.")
+
+    try:
+        env = select_env(config.get("type"))
+    except ValueError as e:
+        raise ValueError(
+            f"Unknown environment type '{config.get('type')}' from dataset's config: {config}"
+        )
+
+    dataset = env.load_dataset(logger=logger, **config)
+    return dataset

debug_gym/gym/envs/aider.py

@@ -138,6 +138,7 @@ def load_dataset(
         problems: str | list[str] | None = None,
         build_image: bool = True,
         logger: object = None,
+        **kwargs,
     ) -> dict:
         if build_image:
             build_docker_image(logger)
@@ -184,4 +185,9 @@ def load_dataset(
 
         problems = utils.filter_problems(dataset, problems)
         dataset = {id: data for id, data in dataset.items() if id in problems}
+
+        # Add env_type to each task_data.
+        for task_data in dataset.values():
+            task_data["env_type"] = "aider"
+
         return dataset

debug_gym/gym/envs/mini_nightmare.py

@@ -152,6 +152,7 @@ def load_dataset(
         problems: str | list[str] | None = None,
         build_image: bool = True,
         logger: object = None,
+        **kwargs,
     ) -> dict:
         if build_image:
             build_docker_image(logger)
@@ -180,4 +181,9 @@ def load_dataset(
 
         problems = utils.filter_problems(dataset, problems)
         dataset = {id: data for id, data in dataset.items() if id in problems}
+
+        # Add env_type to each task_data.
+        for task_data in dataset.values():
+            task_data["env_type"] = "mini_nightmare"
+
         return dataset

debug_gym/gym/envs/r2egym.py

@@ -262,6 +262,7 @@ def load_dataset(
         problems: list | None = None,
         prepull_images: bool = False,
         logger: DebugGymLogger | None = None,
+        **kwargs,
     ) -> dict:
         logger = logger or DebugGymLogger("debug_gym")
         data_path = Path(dataset_id)
@@ -297,9 +298,10 @@ def extract_instance_id(docker_image: str) -> str:
         problems = filter_problems(id2idx, problems, custom_splits, excluded_ids)
         dataset = {problem: ds[id2idx[problem]] for problem in problems}
 
-        # add instance id to each example (name of the image)
+        # Add instance_id (name of the image) and env_type to each task_data.
         for instance_id, task_data in dataset.items():
             task_data["instance_id"] = instance_id
+            task_data["env_type"] = "r2egym"
 
         image_names = set(task_data["docker_image"] for task_data in dataset.values())
         logger.debug(

debug_gym/gym/envs/swe_bench.py

@@ -182,6 +182,7 @@ def load_dataset(
         problems: list | None = None,
         prepull_images: bool = False,
         logger: DebugGymLogger | None = None,
+        **kwargs,
     ) -> dict:
         ds = datasets.load_dataset(dataset_id, revision=dataset_revision)[split]
 
@@ -190,6 +191,10 @@ def load_dataset(
         problems = filter_problems(id2idx, problems)
         dataset = {problem: ds[id2idx[problem]] for problem in problems}
 
+        # Add env_type to each task_data.
+        for task_data in dataset.values():
+            task_data["env_type"] = "swebench"
+
         image_names = set(
             f"sweb.eval.x86_64.{id.replace('__', '_1776_')}" for id in dataset
         )

debug_gym/gym/envs/swe_bench_debug.py

@@ -15,3 +15,13 @@ def eval(self, **kwargs) -> EvalOutput:
         success, output = self.terminal.run(self.entrypoint, timeout=self.run_timeout)
         self.last_eval = EvalOutput(success, output)
         return self.last_eval
+
+    @classmethod
+    def load_dataset(*args, **kwargs) -> dict:
+        dataset = SWEBenchEnv.load_dataset(*args, **kwargs)
+
+        # Add env_type to each task_data.
+        for task_data in dataset.values():
+            task_data["env_type"] = "swebench-debug"
+
+        return dataset

debug_gym/gym/envs/swe_smith.py

@@ -153,6 +153,7 @@ def load_dataset(
         problems: list | None = None,
         prepull_images: bool = False,
         logger: DebugGymLogger | None = None,
+        **kwargs,
     ) -> dict:
         logger = logger or DebugGymLogger("debug_gym")
         data_path = Path(dataset_id)
@@ -182,6 +183,10 @@ def load_dataset(
         problems = filter_problems(id2idx, problems, custom_splits, excluded_ids)
         dataset = {problem: ds[id2idx[problem]] for problem in problems}
 
+        # Add env_type to each task_data.
+        for task_data in dataset.values():
+            task_data["env_type"] = "swesmith"
+
         image_names = set(task_data["image_name"] for task_data in dataset.values())
         logger.debug(
             f"Loaded {len(dataset)} tasks across {len(image_names)} Docker images from {dataset_id}."

scripts/config_aider.yaml

@@ -1,14 +1,17 @@
 base:
     # Environment configs
     output_path: "exps/aider"
-    benchmark: "aider"
-    problems: "all"  # list of problems, e.g., ["wordy"], or "all"
+
     env:
-        type: "aider"
         run_timeout: 20
+
     terminal:
         type: "docker"  # "docker", "kubernetes", or "local"
 
+    dataset:
+        type: "aider"
+        problems: "all"  # list of problems, e.g., ["wordy"], or "all"
+
     # LLM configs
     llm_name: "gpt-4o"