Proper termination for the simulate window

Author: pkrack

python/rcs/sim.py

@@ -1,6 +1,6 @@
 import atexit
-import time
 import multiprocessing as mp
+import time
 import uuid
 from copy import deepcopy
 from logging import getLogger
@@ -35,12 +35,13 @@
 FRAME_DURATION = 1.0 / FPS
 
 
-def gui_loop(model_bytes: bytes, close_event: Event, gui_uuid: str):
+def gui_loop(gui_uuid: str, close_event: Event):
+    gui_client = _GuiClient(gui_uuid)
+    model_bytes = gui_client.get_model_bytes()
     with NamedTemporaryFile(mode="wb") as f:
         f.write(model_bytes)
         model = mujoco.MjModel.from_binary_path(f.name)
     data = mujoco.MjData(model)
-    gui_client = _GuiClient(gui_uuid)
     gui_client.set_model_and_data(model._address, data._address)
     mujoco.mj_step(model, data)
     with mujoco.viewer.launch_passive(model, data) as viewer:
@@ -69,28 +70,29 @@ def __init__(self, mjmdl: str | PathLike):
             logger.error(msg)
         self.data = mj.MjData(self.model)
         super().__init__(self.model._address, self.data._address)
-        self._gui_uuid: Optional[str] = None
         self._mp_context = mp.get_context("spawn")
+        self._gui_uuid: Optional[str] = None
         self._gui_client: Optional[GuiClient] = None
-        self._gui_thread: Optional[Thread] = None
+        self._gui_process: Optional[mp.Process] = None
+        self._stop_event: Optional[mp.Event] = None
+
+    def close_gui(self):
+        if self._stop_event is not None:
+            self._stop_event.set()
+        if self._gui_process is not None:
+            self._gui_process.join()
+        self._stop_gui_server()
 
     def open_gui(self):
         if self._gui_uuid is None:
             self._gui_uuid = "rcs_" + str(uuid.uuid4())
             self._start_gui_server(self._gui_uuid)
-            atexit.register(self._stop_gui_server)
         if self._gui_client is None:
-            self._gui_client = _GuiClient(self._gui_uuid)
-            model_bytes = self._gui_client.get_model_bytes()
             ctx = mp.get_context("spawn")
             self._stop_event = ctx.Event()
             self._gui_process = ctx.Process(
                 target=gui_loop,
-                args=(model_bytes, self._stop_event, self._gui_uuid),
+                args=(self._gui_uuid, self._stop_event),
             )
             self._gui_process.start()
-
-    def __del__(self):
-        self._stop_event.set()
-        self._gui_process.join()
-        self._stop_gui_server()
+        atexit.register(self.close_gui)