Refactor viewer loop and shutdown handling to support conditional rendering and clean exits.

This change introduces a get_frame method to the viewer protocol, replacing the previous is_running check. This allows the viewer loop to skip rendering and lifecycle events when no active frame is ready (for example, when no browser is connected to the web viewer).

Additionally, this change:
- Adds a shutdown hook to ViewerHandle to wait for the viewer to cleanly release resources before exiting.
- Implements an ExitEvent handler in the viewer protocol to stop the run loop.
- Adds an http_port configuration option for the web viewer.
- Prevents double-closing in the viewer.

PiperOrigin-RevId: 954603868
Change-Id: Iaba7d131dc94191de51150464514dd10fa7834b0

python/mujoco/experimental/studio/launch_passive.py

@@ -149,6 +149,7 @@ def launch_passive(
   handle = viewer_handle.ViewerHandle(
       sim_endpoint,
       is_alive_fn=thread.is_alive,
+      shutdown_fn=thread.join,
       handlers=sim_handlers,
   )
   return handle

python/mujoco/experimental/studio/native_viewer.py

@@ -99,11 +99,12 @@ class NativeViewer(viewer_protocol.Viewer):
       self._viewer.InitRenderer(model)
       self._renderer_model_id = id(model)
 
-  def is_running(self) -> bool:
+  def prepare_next_frame(self) -> bool:
-    """Poll for a new frame; returns ``False`` when the window is closed."""
+    """Advances to the next frame; returns False when the window is closed."""
-    if super().is_running() and not self._viewer.NewFrame():
+    if not self._viewer.NewFrame():
-      self.close()
+      self._is_running = False
-    return super().is_running()
+      return False
+    return True
 
   def sync(self) -> None:
     """Render the scene and present it to the window."""

python/mujoco/experimental/studio/viewer_handle.py

@@ -21,6 +21,15 @@ from mujoco.experimental.studio import messages
 from mujoco.experimental.studio import sim as _sim
 import numpy as np
 
+# Launcher-owned liveness check: returns True while the viewer is still alive.
+# This enables the handle to notice a viewer died without sending ExitEvent.
+IsAliveFn = Callable[[], bool]
+
+# Launcher-owned shutdown function: waits up to the given timeout (seconds) for
+# the viewer to finish after close() has sent the ExitEvent. Waiting lets the
+# viewer release its resources before the interpreter tears itself down.
+ShutdownFn = Callable[[float], None]
+
 
 class ViewerHandle:
   """A handle for interacting with a running Studio application from the sim."""
@@ -30,7 +39,8 @@ class ViewerHandle:
       sim_endpoint: endpoints.SimEndpoint,
       *,
       handlers: list[Any] | None = None,
-      is_alive_fn: Callable[[], bool] | None = None,
+      is_alive_fn: IsAliveFn | None = None,
+      shutdown_fn: ShutdownFn | None = None,
   ) -> None:
     """Initializes the ViewerHandle.
 
@@ -38,14 +48,15 @@ class ViewerHandle:
       sim_endpoint: The endpoint to use for communication with the viewer.
       handlers: Optional list of handler instances for sim-side processing,
         which are classes with methods decorated with ``@handler``.
-      is_alive_fn: Optional function called to check if the viewer is still
+      is_alive_fn: Optional liveness check; without one the viewer is assumed to
-        alive/responsive. If not provided, the viewer is assumed to be running
+        be running until ``close()`` is called.
-        until ``close()`` is called.
+      shutdown_fn: Optional launcher-owned shutdown hook, called by ``close()``.
     """
 
     self._sim_endpoint = sim_endpoint
     self._is_running = True
     self._is_alive_fn = is_alive_fn
+    self._shutdown_fn = shutdown_fn
     self.model: mujoco.MjModel | None = None
     self.data: mujoco.MjData | None = None
     self.step_control: _sim.StepControl | None = None
@@ -56,13 +67,15 @@ class ViewerHandle:
     self._handlers = handler_registry.HandlerRegistry(all_handlers)
 
   def close(self) -> None:
-    """Signals the viewer to exit and closes the sim endpoint."""
+    """Signals the viewer to exit and waits for it to shut down."""
     if self._is_running:
       self._is_running = False
       try:
         self.send_to_viewer(messages.ExitEvent())
       except Exception:  # pylint: disable=broad-exception-caught
         pass  # Ignore exceptions, the viewer may have already closed.
+      if self._shutdown_fn is not None:
+        self._shutdown_fn(5.0)
     self._sim_endpoint.close()
 
   def __enter__(self) -> 'ViewerHandle':

python/mujoco/experimental/studio/viewer_protocol.py

@@ -58,6 +58,9 @@ class ViewerConfig:
   height: int = 800
   gfx: str = ''
   viewer_mode: ViewerMode = ViewerMode.NATIVE
+  # Web viewer only: public port. 0 picks the first free port starting at 8080,
+  # so several viewers can run side by side.
+  http_port: int = 0
 
 
 # Legacy message types kept for backward compatibility.
@@ -141,6 +144,7 @@ class Viewer(abc.ABC):
     self.config = config
     self._endpoint = endpoint
     self._is_running = True
+    self._closed = False
 
     # Viewer-owned model and data.
     if model is None:
@@ -169,13 +173,21 @@ class Viewer(abc.ABC):
 
   def close(self) -> None:
     """Closes the viewer, sends an exit event and shuts down the endpoint."""
-    if self._is_running:
+    if self._closed:
-      self._is_running = False
+      return
-      try:
+    self._closed = True
-        self.send_to_sim(messages.ExitEvent())
+    self._is_running = False
-      except Exception:  # pylint: disable=broad-exception-caught
+    try:
-        pass  # Ignore exceptions, the sim may have already closed.
+      self.send_to_sim(messages.ExitEvent())
-      self._endpoint.close()
+    except Exception:  # pylint: disable=broad-exception-caught
+      pass  # Ignore exceptions, the sim may have already closed.
+    self._endpoint.close()
+
+  @messages.handler(priority=messages.Priority.CRITICAL)
+  def _on_exit(self, _: messages.ExitEvent) -> bool:
+    """Stops the viewer loop when the sim side requests an exit."""
+    self._is_running = False
+    return False  # Do not consume; app handlers may want cleanup too.
 
   def is_running(self) -> bool:
     """Returns True while the viewer has not been closed."""
@@ -221,6 +233,11 @@ class Viewer(abc.ABC):
       mujoco.mj_forward(self.model, self.data)
     return False  # Do not consume; let other handlers see the event.
 
+  @abc.abstractmethod
+  def prepare_next_frame(self) -> bool:
+    """Advances to the next frame; returns whether one is ready to render."""
+    ...
+
   @abc.abstractmethod
   def sync(self) -> None:
     """Renders the scene using the viewer's current model and data."""
@@ -251,18 +268,30 @@ def run_viewer_loop(viewer: Viewer) -> None:
   Args:
     viewer: A Viewer that owns the endpoint and handler registry.
   """
-  while viewer.is_running():
+  while True:
-    # Process incoming messages.
+    # Get the next frame; this also gets the frame's mouse/keyboard events.
+    frame = viewer.prepare_next_frame()
+
+    # Process incoming simulation events.
     for event in viewer.get_sim_events():
       viewer.dispatch(event)
+
+    # Stop the loop if the viewer is not running (endpoint will be closed).
+    if not viewer.is_running():
+      break
+
+    # Process incoming simulation snapshots.
     for snapshot in viewer.get_sim_snapshots():
       viewer.dispatch(snapshot)
 
-    # Dispatch lifecycle events.
+    # Skip rendering when prepare_next_frame returned no active frame.
-    viewer.dispatch(messages.UpdateEvent())
+    # e.g., no browser is connected to the web viewer.
-    viewer.dispatch(messages.BuildGuiEvent())
+    if frame:
+      # Dispatch lifecycle events.
+      viewer.dispatch(messages.UpdateEvent())
+      viewer.dispatch(messages.BuildGuiEvent())
 
-    # Render the scene.
+      # Render the scene.
-    viewer.sync()
+      viewer.sync()
 
   viewer.close()