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:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user