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