From e77751214b893ab95b24c573008cb2dc7fe1a549 Mon Sep 17 00:00:00 2001
From: Lennard Stoellger <lennard.stoellger@tum.de>
Date: Mon, 27 Jul 2026 14:50:00 +0200
Subject: [PATCH 1/2] refactor(viewer): Make _on_sigint a method and allow
 passing custom cleanup callbacks to the constructor before the server is
 stopped.

---
 src/mjviser/viewer.py | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/src/mjviser/viewer.py b/src/mjviser/viewer.py
index b35fb73..d9cea0f 100644
--- a/src/mjviser/viewer.py
+++ b/src/mjviser/viewer.py
@@ -46,6 +46,7 @@ def __init__(
     reset_fn: Callable[[mujoco.MjModel, mujoco.MjData], None] | None = None,
     num_envs: int = 1,
     server: viser.ViserServer | None = None,
+    cleanup_hook: Callable[[], None] | None = None,
   ) -> None:
     """Create an active viewer.
 
@@ -66,6 +67,7 @@ def __init__(
       num_envs: Number of parallel environments (passed to
         ``ViserMujocoScene``).
       server: Optional viser server. If None, one is created.
+      cleanup_hook: Optional function called on exit to clean up resources.
     """
     self.model = model
     self.data = data
@@ -76,6 +78,7 @@ def __init__(
     self.scene = ViserMujocoScene(self._server, model, num_envs=num_envs)
     self._lock = Lock()
     self.scene.set_refresh_handler(self._refresh_scene_from_gui)
+    self._cleanup_hook = cleanup_hook if cleanup_hook is not None else lambda: None
 
     # Speed.
     self._speed_idx = _SPEEDS.index(1.0)
@@ -86,6 +89,7 @@ def __init__(
     self._paused = False
     self._dirty = True
     self._step_count = 0
+    self._interrupted = False
 
     # Timing.
     self._budget = 0.0
@@ -146,19 +150,17 @@ def _reset(self) -> None:
       mujoco.mj_resetData(self.model, self.data)
       mujoco.mj_forward(self.model, self.data)
 
+  def _on_sigint(self, signum, frame):
+      self._interrupted = True
+      signal.signal(signal.SIGINT, signal.SIG_DFL)
+
   def run(self) -> None:
     """Run the viewer loop until Ctrl+C."""
     self._setup_gui()
 
     prev_handler = signal.getsignal(signal.SIGINT)
-    interrupted = False
-
-    def _on_sigint(signum, frame):
-      nonlocal interrupted
-      interrupted = True
-      signal.signal(signal.SIGINT, signal.SIG_DFL)
 
-    signal.signal(signal.SIGINT, _on_sigint)
+    signal.signal(signal.SIGINT, self._on_sigint)
 
     if self._render_fn is None:
       mujoco.mj_forward(self.model, self.data)
@@ -168,10 +170,11 @@ def _on_sigint(signum, frame):
     self._last_tick = now
     self._stats_last_time = now
     try:
-      while not interrupted:
+      while not self._interrupted:
         self._tick()
         time.sleep(0.001)
     finally:
+      self._cleanup_hook()
       self._server.stop()
       signal.signal(signal.SIGINT, prev_handler)
 

From dd063f6f4ad037572fb4bfe8dd124a797e5f661e Mon Sep 17 00:00:00 2001
From: Kevin Zakka <kevinarmandzakka@gmail.com>
Date: Tue, 28 Jul 2026 21:56:41 -0700
Subject: [PATCH 2/2] Fix indentation in _on_sigint to satisfy ruff format

---
 src/mjviser/viewer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/mjviser/viewer.py b/src/mjviser/viewer.py
index d9cea0f..c91cb70 100644
--- a/src/mjviser/viewer.py
+++ b/src/mjviser/viewer.py
@@ -151,8 +151,8 @@ def _reset(self) -> None:
       mujoco.mj_forward(self.model, self.data)
 
   def _on_sigint(self, signum, frame):
-      self._interrupted = True
-      signal.signal(signal.SIGINT, signal.SIG_DFL)
+    self._interrupted = True
+    signal.signal(signal.SIGINT, signal.SIG_DFL)
 
   def run(self) -> None:
     """Run the viewer loop until Ctrl+C."""