From b60eb2a93342660341ba8849ff703c662464bc93 Mon Sep 17 00:00:00 2001
From: Timothyxxx <384084775@qq.com>
Date: Thu, 18 Jan 2024 01:43:57 +0800
Subject: [PATCH] VM resolution adjust support

---
 desktop_env/envs/desktop_env.py | 63 ++++++++++++++++++++++++++++++---
 1 file changed, 59 insertions(+), 4 deletions(-)

diff --git a/desktop_env/envs/desktop_env.py b/desktop_env/envs/desktop_env.py
index 786ed72..94e3b0d 100644
--- a/desktop_env/envs/desktop_env.py
+++ b/desktop_env/envs/desktop_env.py
@@ -5,7 +5,7 @@ import os
 import subprocess
 import tempfile
 import time
-from typing import Callable, Any, Optional
+from typing import Callable, Any, Optional, Tuple
 # import uuid
 # import platform
 from typing import List, Dict
@@ -48,7 +48,8 @@ class DesktopEnv(gym.Env):
             action_space: str = "computer_13",
             task_config: Dict[str, Any] = None,
             tmp_dir: str = "tmp",
-            cache_dir: str = "cache"
+            cache_dir: str = "cache",
+            screen_size: Tuple[int] = (1920, 1080)
     ):
         """
         Args:
@@ -73,6 +74,7 @@ class DesktopEnv(gym.Env):
         self.path_to_vm = os.path.abspath(os.path.expandvars(os.path.expanduser(path_to_vm)))
         self.tmp_dir_base: str = tmp_dir
         self.cache_dir_base: str = cache_dir
+        self.vm_screen_size = screen_size
 
         # task-aware stuffs
         # todo: handling the logic of snapshot directory
@@ -80,6 +82,7 @@ class DesktopEnv(gym.Env):
 
         # Initialize emulator and controller
         logger.info("Initializing...")
+        self._config_screen_size()
         self._start_emulator()
         self.vm_ip = self._get_vm_ip()
         self.controller = PythonController(vm_ip=self.vm_ip)
@@ -87,7 +90,6 @@ class DesktopEnv(gym.Env):
 
         # Meta info of the VM, move to the reset() function
         self.vm_platform: str = ""  # self.controller.get_vm_platform()
-        self.vm_screen_size = None  # self.controller.get_vm_screen_size()
 
         # mode: human or machine
         assert action_space in ["computer_13", "pyautogui"]
@@ -101,6 +103,56 @@ class DesktopEnv(gym.Env):
         self._step_no: int = 0
         self.action_history: List[Dict[str, any]] = []
 
+    def _config_screen_size(self):
+        def calculate_vram_size(width, height, bits_per_pixel=32):
+            """
+            Calculate VRAM size for given width, height, and color depth.
+            Color depth defaults to 32 bits per pixel.
+            """
+            bytes_per_pixel = bits_per_pixel // 8
+            vram_size = width * height * bytes_per_pixel
+            return vram_size
+        if not os.path.isfile(self.path_to_vm):
+            logger.warning(f"The specified vmx file does not exist: {self.path_to_vm}")
+            return False
+
+        width, height = self.vm_screen_size
+        vramSize = calculate_vram_size(width, height)
+
+        try:
+            with open(self.path_to_vm, 'r') as file:
+                lines = file.readlines()
+
+            new_lines = []
+            for line in lines:
+                if "svga.autodetect" in line:
+                    continue
+                elif "svga.vramSize" in line:
+                    continue
+                elif "displayWidth" in line:
+                    continue
+                elif "displayHeight" in line:
+                    continue
+                else:
+                    new_lines.append(line)
+
+            # Append new settings for screen size and VRAM.
+            new_lines.append(f'svga.autodetect = "TRUE"\n')
+            new_lines.append(f'svga.vramSize = "{vramSize}"\n')
+            new_lines.append(f'displayWidth = "{width}"\n')
+            new_lines.append(f'displayHeight = "{height}"\n')
+
+            with open(self.path_to_vm, 'w') as file:
+                file.writelines(new_lines)
+            logger.info(f"Screen size for {self.path_to_vm} set to {width}x{height} with VRAM size {vramSize} bytes")
+            return True
+        except IOError as e:
+            logger.error(f"An IOError occurred: {e}")
+            return False
+        except Exception as e:
+            logger.error(f"An error occurred: {e}")
+            return False
+
     def _start_emulator(self):
         while True:
             try:
@@ -119,7 +171,7 @@ class DesktopEnv(gym.Env):
                 logger.error(f"Error executing command: {e.output.decode().strip()}")
 
     def _get_vm_ip(self):
-        max_retries = 10
+        max_retries = 20
         logger.info("Getting IP Address...")
         for _ in range(max_retries):
             try:
@@ -190,6 +242,8 @@ class DesktopEnv(gym.Env):
         _execute_command(["vmrun", "-T", "ws", "revertToSnapshot", self.path_to_vm, self.snapshot_path])
         time.sleep(5)
 
+        self._config_screen_size()
+        print(self.vm_screen_size)
         logger.info("Starting emulator...")
         self._start_emulator()
         logger.info("Emulator started.")
@@ -197,6 +251,7 @@ class DesktopEnv(gym.Env):
         logger.info("Get meta info of the VM...")
         self.vm_platform = self.controller.get_vm_platform()
         self.vm_screen_size = self.controller.get_vm_screen_size()
+        print(self.vm_screen_size)
 
         logger.info("Setting up environment...")
         self.setup_controller.setup(self.config)