put video recording under flag, lint

Author: ollmer

experiments/run_osworld.py

@@ -29,9 +29,9 @@ def main():
 
     if os.environ.get("AGENTLAB_DEBUG"):
         task_ids = get_task_ids()
-        study.exp_args_list = [exp_args for exp_args in study.exp_args_list if exp_args.env_args.task["id"] in task_ids]
+        study.exp_args_list = [exp_args for exp_args in study.exp_args_list if exp_args.env_args.task["id"] in task_ids]  # type: ignore
         print(f"Debug on {len(study.exp_args_list)} experiments")
-        study.run(n_jobs=2, n_relaunch=1, parallel_backend="ray")
+        study.run(n_jobs=4, n_relaunch=1, parallel_backend="ray")
     else:
         study.run(n_jobs=n_jobs, n_relaunch=1, parallel_backend="ray")
 

src/agentlab/agents/tool_use_agent/tool_use_agent.py

@@ -384,6 +384,7 @@ def set_benchmark(self, benchmark: AgentLabBenchmark | BgymBenchmark, demo_mode:
         if benchmark_name == "osworld":
             self.config.obs.use_osworld_obs_preprocessor = True
 
+
 class ToolUseAgent(bgym.Agent):
     def __init__(
         self,
@@ -598,5 +599,7 @@ def get_action(self, obs: Any) -> float:
         multiaction=False,  # whether to use multi-action or not
         action_subsets=("coord",),
     ),
-    action_set=OSWorldActionSet("computer_13"),  # or "pyautogui" #TODO: agent config should only be some primitive types.
-)
+    action_set=OSWorldActionSet(
+        "computer_13"
+    ),  # or "pyautogui" #TODO: agent config should only be some primitive types.
+)

src/agentlab/benchmarks/osworld.py

@@ -343,6 +343,7 @@ def __init__(
         enable_proxy: bool,
         max_steps: int,
         exp_dir: Path,
+        record_video: bool = True,
     ):
         self.task = task
         self.env_info = {
@@ -375,13 +376,18 @@ def __init__(
         self._step_count = 0
         self.max_steps = max_steps
         self.exp_dir = exp_dir
+        self.record_video = record_video
 
     def reset(self, seed: int | None = None) -> tuple[dict[str, Any], dict[str, Any]]:
         self.env.reset(task_config=self.task, seed=seed)
         logging.info(f"Start solving task: {self.task['instruction']}")
-        time.sleep(60) # Wait for the environment to be ready, as in https://github.com/xlang-ai/OSWorld/blob/main/lib_run_single.py#L15
-        raw_obs = self.env._get_obs() # Get the initial observation
-        self.env.controller.start_recording()
+        time.sleep(
+            60
+        )  # Wait for the environment to be ready, as in https://github.com/xlang-ai/OSWorld/blob/main/lib_run_single.py#L15
+        raw_obs = self.env._get_obs()  # Get the initial observation
+        if self.record_video:
+            self.env.controller.start_recording()
+            logging.info("Started recording the environment video")
         obs = self.to_agentlab_observation(raw_obs)
         self._step_count = 0
         return obs, self.env_info
@@ -520,9 +526,10 @@ def parse_agentlab_action_str_to_func_args(action: str):
         return None, None, None
 
     def close(self):
-        video_name = str(self.exp_dir / "recording.mp4")
-        self.env.controller.end_recording(video_name)
-        logger.info(f"Recorded video saved to {video_name}")
+        if self.record_video:
+            video_name = str(self.exp_dir / "recording.mp4")
+            self.env.controller.end_recording(video_name)
+            logger.info(f"Recorded video saved to {video_name}")
         return self.env.close()
 
 
@@ -671,10 +678,8 @@ def fix_settings_file_path_in_config(self, task) -> str:
         osworld_repo = os.getenv("OSWORLD_REPO", "OSWorld")
         updated_task = deepcopy(task)  # Avoid modifying the original task
         for config in updated_task["config"]:
-                if config.get("parameters", False) and config["parameters"].get(
-                    "settings_file", False
-                ):
-                    config["parameters"]["settings_file"] = os.path.join(
-                        osworld_repo, config["parameters"]["settings_file"]
-                    )
+            if config.get("parameters", False) and config["parameters"].get("settings_file", False):
+                config["parameters"]["settings_file"] = os.path.join(
+                    osworld_repo, config["parameters"]["settings_file"]
+                )
         return updated_task

src/agentlab/benchmarks/osworld_axtree_preprocessing.py

@@ -25,6 +25,7 @@ def collect_leaf_nodes(node, leaf_nodes):
     collect_leaf_nodes(root, leaf_nodes)
     return leaf_nodes
 
+
 attributes_ns_ubuntu = "https://accessibility.windows.example.org/ns/attributes"
 attributes_ns_windows = "https://accessibility.windows.example.org/ns/attributes"
 state_ns_ubuntu = "https://accessibility.ubuntu.example.org/ns/state"
@@ -337,4 +338,3 @@ def trim_accessibility_tree(linearized_accessibility_tree, max_tokens):
         linearized_accessibility_tree = enc.decode(tokens[:max_tokens])
         linearized_accessibility_tree += "[...]\n"
     return linearized_accessibility_tree
-