From 9620c75cf237d6aaf72952b5cf9fd6e54bcd10f2 Mon Sep 17 00:00:00 2001
From: Jing Hua <tohjinghua123@gmail.com>
Date: Thu, 2 Nov 2023 12:09:36 +0800
Subject: [PATCH] finish reset step

---
 desktop_env/envs/desktop_env.py | 26 ++++++++++++++++++++++++--
 main.py                         |  2 +-
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/desktop_env/envs/desktop_env.py b/desktop_env/envs/desktop_env.py
index bda6ed8..bcb67d3 100644
--- a/desktop_env/envs/desktop_env.py
+++ b/desktop_env/envs/desktop_env.py
@@ -1,6 +1,7 @@
 from enum import Enum
 import subprocess
 from fabric import Connection
+import time
 
 import gymnasium as gym
 from gymnasium import spaces
@@ -51,10 +52,26 @@ class DesktopEnv(gym.Env):
         # Additional setup
         self.metadata = {'render.modes': ['rgb_array']}
         self._start_emulator()
+        self._wait_for_emulator_load()
 
     def _start_emulator(self):
         self._execute_command(["vmrun", "start", self.path_to_vm])
 
+    def _wait_for_emulator_load(self):
+        while True:
+            try:
+                output = subprocess.check_output(f"vmrun -T ws list", shell=True, stderr=subprocess.STDOUT)
+                output = output.decode()
+                if self.path_to_vm.lstrip("~/") in output:
+                    print("VM is running.")
+                    return
+                else:
+                    print("Waiting for VM to start...")
+                    time.sleep(5)
+            except subprocess.CalledProcessError as e:
+                print(f"Error executing command: {e.output.decode().strip()}")
+                return
+
     def _execute_command(self, command: list[str]) -> None:
         process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         stdout, stderr = process.communicate()
@@ -70,7 +87,7 @@ class DesktopEnv(gym.Env):
         return result.stdout.strip()
 
     def _save_state(self):
-        self._execute_command(["vmrun", "snapshot", self.snapshot_path])
+        self._execute_command(["vmrun", "-T", "ws" "snapshot", self.path_to_vm, self.snapshot_path])
 
     def _click(self, click: MouseClick):
         self._execute_xdotool_command(f"click {click.value}")
@@ -103,7 +120,12 @@ class DesktopEnv(gym.Env):
             return np.array(img)
 
     def reset(self):
-        self._execute_command(["vmrun", "revertToSnapshot", self.snapshot_path])
+        input()
+        self._execute_command(["vmrun", "-T", "ws", "revertToSnapshot", self.path_to_vm, self.snapshot_path])
+        input()
+        self._start_emulator()
+        input()
+        self._wait_for_emulator_load()
         observation = self._get_obs()
 
         return observation
diff --git a/main.py b/main.py
index 2e3ca27..0b20ce1 100644
--- a/main.py
+++ b/main.py
@@ -38,7 +38,7 @@ def human_agent():
     """
     Runs the Gym environment with human input.
     """
-    env = DesktopEnv(path_to_vm="~/vmware/Ubuntu 64-bit/Ubuntu 64-bit.vmx",
+    env = DesktopEnv(path_to_vm="/home/yuri/vmware/Ubuntu 64-bit/Ubuntu 64-bit.vmx",
                  username="user",
                  password="password",
                  host="192.168.7.128")