Merge branch 'zdy'

desktop_env/envs/desktop_env.py

@@ -17,10 +17,12 @@ from desktop_env.controllers.setup import SetupController
 # from desktop_env.evaluators import eval_funcs
 from desktop_env.evaluators import metrics, getters
 
+import logging
+logger = logging.getLogger("desktopenv.env")
+
 Metric = Callable[[Any, Any], float]
 Getter = Callable[[gym.Env, Dict[str, Any]], Any]
 
-
 def _execute_command(command: List[str]) -> None:
     if command[:4] == ["vmrun", "-T", "ws", "start"]:
         p = subprocess.Popen(command)
@@ -68,22 +70,11 @@ class DesktopEnv(gym.Env):
         self.cache_dir_base: str = cache_dir
 
         # task-aware stuffs
-        self.snapshot_path = task_config["snapshot"]  # todo: handling the logic of snapshot directory
-        self.task_id: str = task_config["id"]
-        self.cache_dir: str = os.path.join(self.cache_dir_base, self.task_id)
-        os.makedirs(self.cache_dir, exist_ok=True)
-        self.instruction = task_config["instruction"]
-        self.config = task_config["config"]
-
-        self.evaluator = task_config["evaluator"]
-        self.metric: Metric = getattr(metrics, self.evaluator["func"])
-        self.result_getter: Getter = getattr(getters, "get_{:}".format(self.evaluator["result"]["type"]))
-        self.expected_getter: Getter = getattr(getters, "get_{:}".format(
-            self.evaluator["expected"]["type"])) if "expected" in self.evaluator else None
-        self.metric_options: Dict[str, Any] = self.evaluator.get("options", {})
+        # todo: handling the logic of snapshot directory
+        self._set_task_info(task_config)
 
         # Initialize emulator and controller
-        print("Initializing...")
+        logger.info("Initializing...")
         self._start_emulator()
         self.vm_ip = self._get_vm_ip()
         self.host = f"http://{self.vm_ip}:5000"
@@ -110,26 +101,26 @@ class DesktopEnv(gym.Env):
                 output: List[str] = output.splitlines()
                 # if self.path_to_vm.lstrip("~/") in output:
                 if self.path_to_vm in output:
-                    print("VM is running.")
+                    logger.info("VM is running.")
                     break
                 else:
-                    print("Starting VM...")
+                    logger.info("Starting VM...")
                     _execute_command(["vmrun", "-T", "ws", "start", self.path_to_vm])
                     time.sleep(3)
             except subprocess.CalledProcessError as e:
-                print(f"Error executing command: {e.output.decode().strip()}")
+                logger.error(f"Error executing command: {e.output.decode().strip()}")
 
     def _get_vm_ip(self):
         max_retries = 10
-        print("Getting IP Address...")
+        logger.info("Getting IP Address...")
         for _ in range(max_retries):
             try:
                 output = _execute_command(["vmrun", "-T", "ws", "getGuestIPAddress", self.path_to_vm]).strip()
-                print(f"IP address: {output}")
+                logger.info(f"IP address: {output}")
                 return output
             except:
                 time.sleep(5)
-                print("Retrying...")
+                logger.info("Retrying...")
         raise Exception("Failed to get VM IP address!")
 
     def _save_state(self):
@@ -152,52 +143,54 @@ class DesktopEnv(gym.Env):
         screenshot_image_path = self._get_screenshot()
         return screenshot_image_path
 
+    def _set_task_info(self, task_config: Dict[str, Any]):
+        self.snapshot_path = task_config["snapshot"]
+        self.task_id: str = task_config["id"]
+        self.cache_dir: str = os.path.join(self.cache_dir_base, self.task_id)
+        os.makedirs(self.cache_dir, exist_ok=True)
+        self.instruction = task_config["instruction"]
+        self.config = task_config["config"]
+
+        self.evaluator = task_config["evaluator"]
+        self.metric: Metric = getattr(metrics, self.evaluator["func"])
+        self.result_getter: Getter = getattr(getters, "get_{:}".format(self.evaluator["result"]["type"]))
+        self.expected_getter: Getter = getattr(getters, "get_{:}".format(
+            self.evaluator["expected"]["type"])) if "expected" in self.evaluator else None
+        self.metric_options: Dict[str, Any] = self.evaluator.get("options", {})
+
     def reset(self, task_config: Optional[Dict[str, Any]] = None, seed=None, options=None):
-        print("Resetting environment...")
+        logger.info("Resetting environment...")
 
-        print("Switching task...")
+        logger.info("Switching task...")
         if task_config is not None:
-            self.snapshot_path = task_config["snapshot"]
-            self.task_id = task_config["id"]
-            self.cache_dir = os.path.join(self.cache_dir_base, self.task_id)
-            os.makedirs(self.cache_dir, exist_ok=True)
-            self.instruction = task_config["instruction"]
-            self.config = task_config["config"]
-
-            self.evaluator = task_config["evaluator"]
-            self.metric: Metric = getattr(metrics, self.evaluator["func"])
-            self.result_getter: Getter = getattr(getters, "get_{:}".format(self.evaluator["result"]["type"]))
-            self.expected_getter: Getter = getattr(getters, "get_{:}".format(
-                self.evaluator["expected"]["type"])) if "expected" in self.evaluator else None
-            self.metric_options = self.evaluator.get("options", {})
-
+            self._set_task_info(task_config)
             self.setup_controller.reset_cache_dir(self.cache_dir)
 
-        print("Setting counters...")
+        logger.info("Setting counters...")
         self._traj_no += 1
         self._step_no = 0
         self.action_history.clear()
 
-        print("Setup new temp dir...")
+        logger.info("Setup new temp dir...")
         self.tmp_dir = tempfile.mkdtemp(
             prefix="{:d}.{:}.".format(self._traj_no, self.task_id),
             dir=self.tmp_dir_base
         )
         os.makedirs(os.path.join(self.tmp_dir, "screenshots"))
 
-        print("Reverting to snapshot to {}...".format(self.snapshot_path))
+        logger.info("Reverting to snapshot to {}...".format(self.snapshot_path))
         _execute_command(["vmrun", "-T", "ws", "revertToSnapshot", self.path_to_vm, self.snapshot_path])
         time.sleep(5)
 
-        print("Starting emulator...")
+        logger.info("Starting emulator...")
         self._start_emulator()
-        print("Emulator started.")
+        logger.info("Emulator started.")
 
-        print("Setting up environment...")
+        logger.info("Setting up environment...")
         self.setup_controller.setup(self.config)
 
         time.sleep(5)
-        print("Environment setup complete.")
+        logger.info("Environment setup complete.")
 
         observation = self._get_obs()
         return observation
@@ -229,6 +222,9 @@ class DesktopEnv(gym.Env):
         """
         Evaluate whether the task is successfully completed.
         """
+
+        self.setup_controller.setup(self.evaluator["postconfig"])
+
         result_state = self.result_getter(self, self.evaluator["result"])
         expected_state = self.expected_getter(self, self.evaluator["expected"]) if "expected" in self.evaluator \
             else None