AWS Enhancement (#48)

* Fix the path registration after reset

* image id

* Handle lock stuff

* Debug

* Debug

* Update

* Update

* Update

* Update

* Update

* Reorder the reset steps

* Reorder the reset steps

* Reorder the reset steps

* Finish and clean

---------

Co-authored-by: XinyuanWangCS <xywang626@gmail.com>

desktop_env/desktop_env.py

@@ -52,6 +52,7 @@ class DesktopEnv(gym.Env):
             require_terminal (bool): whether to require terminal output
         """
         # Initialize VM manager and vitualization provider
+        self.region = region
         self.manager, self.provider = create_vm_manager_and_provider(provider_name, region)
 
         # Initialize environment variables
@@ -95,7 +96,13 @@ class DesktopEnv(gym.Env):
     def _revert_to_snapshot(self):
         # Revert to certain snapshot of the virtual machine, and refresh the path to vm and ip of vm
         # due to the fact it could be changed when implemented by cloud services
-        self.path_to_vm = self.provider.revert_to_snapshot(self.path_to_vm, self.snapshot_name)
+        path_to_vm = self.provider.revert_to_snapshot(self.path_to_vm, self.snapshot_name)
+        if path_to_vm and not path_to_vm == self.path_to_vm:
+            # path_to_vm has to be a new path
+            self.manager.delete_vm(self.path_to_vm, self.region)
+            self.manager.add_vm(path_to_vm, self.region)
+            self.manager.occupy_vm(path_to_vm, os.getpid(), self.region)
+            self.path_to_vm = path_to_vm
 
     def _save_state(self, snapshot_name=None):
         # Save the current virtual machine state to a certain snapshot name