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.
This commit is contained in:
yuanmengqi
2025-06-09 16:35:13 +00:00
parent 168a2694f2
commit 02387f2cee

View File

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