diff --git a/run_autoglm.py b/run_autoglm.py index 5b3a947..ef514a4 100644 --- a/run_autoglm.py +++ b/run_autoglm.py @@ -164,8 +164,67 @@ class DesktopEnv(DesktopEnvBase): observation['exe_result'] = exe_result return observation, reward, done, info - - def reset(self, *args, **kwargs): + + def reset(self, task_config: Optional[Dict[str, Any]] = None, seed=None, options=None) -> Dict[str, Any]: + # Reset to certain task in OSWorld + logger.info("Resetting environment...") + logger.info("Switching task...") + logger.info("Setting counters...") + self._traj_no += 1 + self._step_no = 0 + self.action_history.clear() + + for attempt in range(MAX_RETRIES): + # Only revert to snapshot if environment has been used (step/setup) + # This optimization is especially important for cloud providers like AWS + # where unnecessary snapshot operations are costly and time-consuming + + if task_config is not None: + # Only consider task proxy requirement if proxy is enabled at system level + task_use_proxy = task_config.get("proxy", False) and self.enable_proxy + if not self.enable_proxy and task_config.get("proxy", False): + logger.info("Task requires proxy but proxy is disabled at system level, ignoring proxy requirement.") + + if task_use_proxy != self.current_use_proxy: + # keep because get_info_from_website depend on this + self.current_use_proxy = task_use_proxy + + if self.is_environment_used: + logger.info("Environment has been used, reverting to snapshot {}...".format(self.snapshot_name)) + self._revert_to_snapshot() + logger.info("Starting emulator...") + self._start_emulator() + logger.info("Emulator started.") + # Reset the usage flag after reverting + self.is_environment_used = False + else: + logger.info("Environment is clean, skipping snapshot revert (provider: {}).".format(self.provider_name)) + + if task_config is not None: + if task_config.get("proxy", False) and self.enable_proxy: + # If using proxy and proxy is enabled, set up the proxy configuration + self.setup_controller._proxy_setup(self.client_password) + self._set_task_info(task_config) + self.setup_controller.reset_cache_dir(self.cache_dir) + logger.info("Setting up environment...") + success = self.setup_controller.setup(self.config, task_config.get("proxy", False) and self.enable_proxy) + if success: + # Mark environment as used when setup is successfully executed + if self.config: # Only mark as used if there were actual setup operations + self.is_environment_used = True + break + else: + logger.error( + "Environment setup failed, retrying (%d/%d)...", + attempt + 1, + MAX_RETRIES, + ) + time.sleep(5) + else: + break + + logger.info("Environment setup complete.") + # Upload tools from autoglm package import mm_agents.autoglm tool_dir = os.path.join(os.path.dirname(mm_agents.autoglm.__file__), 'tools', 'package') @@ -174,14 +233,15 @@ class DesktopEnv(DesktopEnvBase): continue self.setup_controller._upload_file_setup([{ "local_path": os.path.join(tool_dir, file), - "path": os.path.join('/home/user', file) + "path": os.path.join('~', file) }]) # start soffice service for office tools self.setup_controller._launch_setup('soffice --accept="socket,host=localhost,port=2002;urp;" --norestore --nologo --nodefault', shell=True) time.sleep(5) - super().reset(*args, **kwargs) + observation = self._get_obs() + return observation def get_current_apps(self): apps_code = r"""import subprocess;