lzy/sci-gui-agent-benchmark
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

finish reset step

Browse Source
This commit is contained in:
Jing Hua
2023-11-02 12:09:36 +08:00
parent fd3787ee28
commit 9620c75cf2
2 changed files with 25 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
desktop_env/envs/desktop_env.py
View File

@@ -1,6 +1,7 @@
from enum import Enum from enum import Enum
import subprocess import subprocess
from fabric import Connection from fabric import Connection
import time
import gymnasium as gym import gymnasium as gym
from gymnasium import spaces from gymnasium import spaces
@@ -51,10 +52,26 @@ class DesktopEnv(gym.Env):
# Additional setup # Additional setup
self.metadata = {'render.modes': ['rgb_array']} self.metadata = {'render.modes': ['rgb_array']}
self._start_emulator() self._start_emulator()
self._wait_for_emulator_load()
def _start_emulator(self): def _start_emulator(self):
self._execute_command(["vmrun", "start", self.path_to_vm]) 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: def _execute_command(self, command: list[str]) -> None:
process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE) process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stdout, stderr = process.communicate() stdout, stderr = process.communicate()
@@ -70,7 +87,7 @@ class DesktopEnv(gym.Env):
return result.stdout.strip() return result.stdout.strip()
def _save_state(self): 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): def _click(self, click: MouseClick):
self._execute_xdotool_command(f"click {click.value}") self._execute_xdotool_command(f"click {click.value}")
@@ -103,7 +120,12 @@ class DesktopEnv(gym.Env):
return np.array(img) return np.array(img)
def reset(self): 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() observation = self._get_obs()
return observation return observation

2
main.py
View File

@@ -38,7 +38,7 @@ def human_agent():
""" """
Runs the Gym environment with human input. 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", username="user",
password="password", password="password",
host="192.168.7.128") host="192.168.7.128")