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]
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()