update screenshot command

2023-11-11 02:38:16 +08:00
parent a8aebf5d15
commit 203fae5798
1 changed files with 1 additions and 3 deletions
--- a/desktop_env/envs/desktop_env.py
+++ b/desktop_env/envs/desktop_env.py
@@ -123,9 +123,7 @@ class DesktopEnv(gym.Env):

    def _get_screenshot(self):
        image_path = "./screenshot.png"
-        self.ssh_connection.run("DISPLAY=:0 import -window root screenshot.png")
-        self._execute_command(["scp", "user@192.168.7.128:~/screenshot.png", image_path])
-        self.ssh_connection.run("rm -rf ~/screenshot.png")
+        self._execute_command(["vmrun", "-T", "ws", "-gu", self.username, "-gp", self.password, "captureScreen", self.path_to_vm, image_path])
        return image_path
    
    def _get_obs(self):