Merge branch 'main' of https://github.com/ztjhz/DesktopEnv

desktop_env/controllers/python.py

@@ -197,8 +197,10 @@ class PythonController:
             if "text" not in parameters:
                 raise Exception(f"Unknown parameters: {parameters}")
             # deal with special ' and \ characters
-            text = parameters["text"].replace("\\", "\\\\").replace("'", "\\'")
-            self.execute_python_command(f"pyautogui.typewrite('{text}')")
+            # text = parameters["text"].replace("\\", "\\\\").replace("'", "\\'")
+            # self.execute_python_command(f"pyautogui.typewrite('{text}')")
+            text = parameters["text"]
+            self.execute_python_command("pyautogui.typewrite({:})".format(repr(text)))
 
         elif action_type == "PRESS":
             if "key" not in parameters:
@@ -237,6 +239,9 @@ class PythonController:
             keys_para_rep = "', '".join(keys)
             self.execute_python_command(f"pyautogui.hotkey('{keys_para_rep}')")
 
+        elif action_type in ['WAIT', 'FAIL', 'DONE']:
+            pass
+
         else:
             raise Exception(f"Unknown action type: {action_type}")
 

desktop_env/envs/actions.py

@@ -186,5 +186,18 @@ ACTION_SPACE = [
                 "optional": False,
             }
         }
+    },
+    ############################################################################################################
+    {
+        "action_type": "WAIT",
+        "note": "wait until the next action",
+    },
+    {
+        "action_type": "FAIL",
+        "note": "decide the task can not be performed",
+    },
+    {
+        "action_type": "DONE",
+        "note": "decide the task is done",
     }
 ]

desktop_env/envs/desktop_env.py

@@ -1,28 +1,30 @@
 from __future__ import annotations
 
+import logging
 import os
 import subprocess
+import tempfile
 import time
+from typing import Callable, Any, Optional
 # import uuid
 # import platform
 from typing import List, Dict
-from typing import Callable, Any, Optional
-import tempfile
 
 import gymnasium as gym
-# import requests
 
 from desktop_env.controllers.python import PythonController
 from desktop_env.controllers.setup import SetupController
 # from desktop_env.evaluators import eval_funcs
 from desktop_env.evaluators import metrics, getters
 
-import logging
+# import requests
+
 logger = logging.getLogger("desktopenv.env")
 
 Metric = Callable[[Any, Any], float]
 Getter = Callable[[gym.Env, Dict[str, Any]], Any]
 
+
 def _execute_command(command: List[str]) -> None:
     if command[:4] == ["vmrun", "-T", "ws", "start"]:
         p = subprocess.Popen(command)
@@ -84,8 +86,8 @@ class DesktopEnv(gym.Env):
         self.setup_controller = SetupController(vm_ip=self.vm_ip, cache_dir=self.cache_dir)
 
         # 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()
+        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"]
@@ -164,7 +166,7 @@ class DesktopEnv(gym.Env):
             self.evaluator["expected"]["type"])) if "expected" in self.evaluator else None
         self.metric_options: Dict[str, Any] = self.evaluator.get("options", {})
 
-    def reset(self, task_config: Optional[Dict[str, Any]] = None, seed=None, options=None):
+    def reset(self, task_config: Optional[Dict[str, Any]] = None, seed=None, options=None) -> Dict[str, Any]:
         logger.info("Resetting environment...")
 
         logger.info("Switching task...")
@@ -202,11 +204,27 @@ class DesktopEnv(gym.Env):
         time.sleep(5)
         logger.info("Environment setup complete.")
 
-        observation = self._get_obs()
+        observation = {"screenshot": self._get_obs()}
         return observation
 
     def step(self, action, pause=0.5):
         self._step_no += 1
+        self.action_history.append(action)
+
+        reward = 0  # todo: Define reward calculation for each example
+        done = False  # todo: Define episode termination condition for each example
+        info = {}
+
+        # handle the special actions
+        if action in ['WAIT', 'FAIL', 'DONE']:
+            if action == 'WAIT':
+                time.sleep(pause)
+            elif action == 'FAIL':
+                done = True
+                info = {"fail": True}
+            elif action == 'DONE':
+                done = True
+                info = {"done": True}
 
         # fixme: add reminding logic here, decide if the action is valid for the current action_space
         if self.action_space == "computer_13":
@@ -215,19 +233,14 @@ class DesktopEnv(gym.Env):
         elif self.action_space == "pyautogui":
             # the set of all possible python commands insides `pyautogui`
             self.controller.execute_python_command(action)
-        self.action_history.append(action)
 
-        # todo: maybe for the better here we need to add a logic to wait until the rendering is done
-        time.sleep(pause)
         observation = {
             "screenshot": self._get_obs(),
             "accessibility_tree": self.controller.get_accessibility_tree(),
             "terminal": self.controller.get_terminal_output(),
             "instruction": self.instruction
         }
-        reward = 0  # todo: Define reward calculation for each example
-        done = False  # todo: Define episode termination condition for each example
-        info = {}
+
         return observation, reward, done, info
 
     def evaluate(self):