diff --git a/desktop_env/envs/desktop_env.py b/desktop_env/envs/desktop_env.py
index a08c919..8b51c5b 100644
--- 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):