Add Support for QWEN models from API (QWEN-max, etc.); Improve on the robustness of getting observation

desktop_env/envs/desktop_env.py

@@ -138,24 +138,9 @@ class DesktopEnv(gym.Env):
     def _save_state(self):
         _execute_command(["vmrun", "-T", "ws" "snapshot", self.path_to_vm, self.snapshot_name])
 
-    def _get_screenshot(self):
-        screenshot = None
-        # Get the screenshot and save to the image_path
-        max_retries = 20
-        for _ in range(max_retries):
-            screenshot = self.controller.get_screenshot()
-            if screenshot is not None:
-                break
-            time.sleep(1)
-
-        if screenshot is None:
-            logger.error("Failed to get screenshot!")
-
-        return screenshot
-
     def _get_obs(self):
         return {
-            "screenshot": self._get_screenshot(),
+            "screenshot": self.controller.get_screenshot(),
             "accessibility_tree": self.controller.get_accessibility_tree() if self.require_a11y_tree else None,
             "terminal": self.controller.get_terminal_output() if self.require_terminal else None,
             "instruction": self.instruction