Fix GIMP bug; Speedup the environment, when there is not a11y tree needed, we can do no controller.get

desktop_env/envs/desktop_env.py

@@ -58,7 +58,8 @@ class DesktopEnv(gym.Env):
             tmp_dir: str = "tmp",
             cache_dir: str = "cache",
             screen_size: Tuple[int] = (1920, 1080),
-            headless: bool = False
+            headless: bool = False,
+            require_a11y_tree: bool = True,
     ):
         """
         Args:
@@ -77,6 +78,7 @@ class DesktopEnv(gym.Env):
         self.cache_dir_base: str = cache_dir
         self.vm_screen_size = screen_size  # todo: add the logic to get the screen size from the VM
         self.headless = headless
+        self.require_a11y_tree = require_a11y_tree
 
         os.makedirs(self.tmp_dir_base, exist_ok=True)
 
@@ -248,7 +250,7 @@ class DesktopEnv(gym.Env):
 
         observation = {
             "screenshot": self._get_obs(),
-            "accessibility_tree": self.controller.get_accessibility_tree(),
+            "accessibility_tree": self.controller.get_accessibility_tree() if self.require_a11y_tree else None,
         }
         return observation
 
@@ -284,7 +286,7 @@ class DesktopEnv(gym.Env):
 
         observation = {
             "screenshot": self._get_obs(),
-            "accessibility_tree": self.controller.get_accessibility_tree(),
+            "accessibility_tree": self.controller.get_accessibility_tree() if self.require_a11y_tree else None,
             # "terminal": self.controller.get_terminal_output(),
             "instruction": self.instruction
         }

desktop_env/evaluators/metrics/__init__.py

@@ -77,6 +77,7 @@ from .general import (
     literal_match
 )
 from .gimp import (
+    check_structure_sim_resized,
     check_brightness_decrease_and_structure_sim,
     check_contrast_increase_and_structure_sim,
     check_saturation_increase_and_structure_sim,