From 02387f2ceebbb7433b4160900131c02009208b95 Mon Sep 17 00:00:00 2001
From: yuanmengqi <yuanmengqi@mail.ustc.edu.cn>
Date: Mon, 9 Jun 2025 16:35:13 +0000
Subject: [PATCH] feat: update DesktopEnv to support VMware provider and add
 proxy configuration

- Changed default provider name from "aws" to "vmware".
- Introduced `enable_proxy` parameter to control proxy support.
- Enhanced retry logic in the `reset` method to use a constant for maximum retries.
- Updated proxy handling to respect the new `enable_proxy` setting.
---
 desktop_env/desktop_env.py | 33 +++++++++++++++++++++++----------
 1 file changed, 23 insertions(+), 10 deletions(-)

diff --git a/desktop_env/desktop_env.py b/desktop_env/desktop_env.py
index 8aff6a8..58a178d 100644
--- a/desktop_env/desktop_env.py
+++ b/desktop_env/desktop_env.py
@@ -18,15 +18,15 @@ logger = logging.getLogger("desktopenv.env")
 Metric = Callable[[Any, Any], float]
 Getter = Callable[[gym.Env, Dict[str, Any]], Any]
 
+MAX_RETRIES = 5
 
 class DesktopEnv(gym.Env):
     """
     DesktopEnv with OpenAI Gym interface. It provides a desktop environment for setting and evaluating desktop automation tasks.
     """
-    #TODO:provider_name: str = "vmware",
     def __init__(
             self,
-            provider_name: str = "aws",
+            provider_name: str = "vmware",
             region: str = None,
             path_to_vm: str = None,
             snapshot_name: str = "init_state",
@@ -37,6 +37,7 @@ class DesktopEnv(gym.Env):
             require_a11y_tree: bool = True,
             require_terminal: bool = False,
             os_type: str = "Ubuntu",
+            enable_proxy: bool = False,
     ):
         """
         Args:
@@ -51,10 +52,13 @@ class DesktopEnv(gym.Env):
             headless (bool): whether to run the VM in headless mode
             require_a11y_tree (bool): whether to require accessibility tree
             require_terminal (bool): whether to require terminal output
+            os_type (str): operating system type, default to "Ubuntu"
+            enable_proxy (bool): whether to enable proxy support, default to False
         """
         # Initialize VM manager and vitualization provider
         self.region = region
         self.provider_name = provider_name
+        self.enable_proxy = enable_proxy  # Store proxy enablement setting
 
         # Default TODO:
         self.server_port = 5000
@@ -145,7 +149,7 @@ class DesktopEnv(gym.Env):
         self.provider.stop_emulator(self.path_to_vm)
 
     def reset(self, task_config: Optional[Dict[str, Any]] = None, seed=None, options=None) -> Dict[str, Any]:
-        retry = 0
+        
         # Reset to certain task in OSWorld
         logger.info("Resetting environment...")
         logger.info("Switching task...")
@@ -153,11 +157,15 @@ class DesktopEnv(gym.Env):
         self._traj_no += 1
         self._step_no = 0
         self.action_history.clear()
-        while retry < 5:
 
+        for attempt in range(MAX_RETRIES):
             # Check and handle proxy requirement changes BEFORE starting emulator
             if task_config is not None:
-                task_use_proxy = task_config.get("proxy", False)
+                # 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:
                     logger.info(f"Task proxy requirement changed: {self.current_use_proxy} -> {task_use_proxy}")
                     
@@ -196,14 +204,19 @@ class DesktopEnv(gym.Env):
                 if success:
                     break
                 else:
-                    logger.error("Environment setup failed, retrying...")
-                    retry += 1
+                    logger.error(
+                        "Environment setup failed, retrying (%d/%d)...",
+                        attempt + 1,
+                        MAX_RETRIES,
+                    )
                     time.sleep(5)
-                    
+            else:
+                break
+            
         logger.info("Environment setup complete.")
     
-        if task_config.get("proxy", False):
-            # If using proxy, set up the proxy configuration
+        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()
 
         observation = self._get_obs()