Improve the logic of env setup; add change wallpaper; add example

desktop_env/envs/desktop_env.py

@@ -26,6 +26,7 @@ class DesktopEnv(gym.Env):
             self,
             path_to_vm: str,
             snapshot_path: str = "base",
+            config: dict = None,
             action_space: str = "pyautogui",
     ):
         # Initialize environment variables
@@ -38,6 +39,7 @@ class DesktopEnv(gym.Env):
         self.host = f"http://{self._get_vm_ip()}:5000"
         self.controller = PythonController(http_server=self.host)
         self.setup_controller = SetupController(http_server=self.host)
+        self.config = config
 
         # mode: human or machine
         assert action_space in ["computer_13", "pyautogui"]
@@ -102,6 +104,9 @@ class DesktopEnv(gym.Env):
         self._start_emulator()
         print("Emulator started.")
 
+        print("Setting up environment...")
+        self.setup_controller.setup(self.config)
+
         observation = self._get_obs()
         return observation
 
@@ -120,9 +125,6 @@ class DesktopEnv(gym.Env):
         done = False  # todo: Define episode termination condition for each example
         info = {}
         return observation, reward, done, info
-    
-    def setup(self, config: dict):
-        self.setup_controller.setup(config) 
 
     def render(self, mode='rgb_array'):
         if mode == 'rgb_array':