diff --git a/python/mujoco/experimental/studio/launch_passive.py b/python/mujoco/experimental/studio/launch_passive.py index 5ecd17ae..215c940f 100644 --- a/python/mujoco/experimental/studio/launch_passive.py +++ b/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 diff --git a/python/mujoco/experimental/studio/native_viewer.py b/python/mujoco/experimental/studio/native_viewer.py index dceaba8f..f7f444d5 100644 --- a/python/mujoco/experimental/studio/native_viewer.py +++ b/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: - """Poll for a new frame; returns ``False`` when the window is closed.""" - if super().is_running() and not self._viewer.NewFrame(): - self.close() - return super().is_running() + def prepare_next_frame(self) -> bool: + """Advances to the next frame; returns False when the window is closed.""" + if not self._viewer.NewFrame(): + self._is_running = False + return False + return True def sync(self) -> None: """Render the scene and present it to the window.""" diff --git a/python/mujoco/experimental/studio/viewer_handle.py b/python/mujoco/experimental/studio/viewer_handle.py index f87241a6..ef910c0e 100644 --- a/python/mujoco/experimental/studio/viewer_handle.py +++ b/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 - alive/responsive. If not provided, the viewer is assumed to be running - until ``close()`` is called. + is_alive_fn: Optional liveness check; without one the viewer is assumed to + be running 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': diff --git a/python/mujoco/experimental/studio/viewer_protocol.py b/python/mujoco/experimental/studio/viewer_protocol.py index 9dc83cca..e6193d90 100644 --- a/python/mujoco/experimental/studio/viewer_protocol.py +++ b/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: - self._is_running = False - try: - self.send_to_sim(messages.ExitEvent()) - except Exception: # pylint: disable=broad-exception-caught - pass # Ignore exceptions, the sim may have already closed. - self._endpoint.close() + if self._closed: + return + self._closed = True + self._is_running = False + try: + self.send_to_sim(messages.ExitEvent()) + 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(): - # Process incoming messages. + while True: + # 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. - viewer.dispatch(messages.UpdateEvent()) - viewer.dispatch(messages.BuildGuiEvent()) + # Skip rendering when prepare_next_frame returned no active frame. + # e.g., no browser is connected to the web viewer. + if frame: + # Dispatch lifecycle events. + viewer.dispatch(messages.UpdateEvent()) + viewer.dispatch(messages.BuildGuiEvent()) - # Render the scene. - viewer.sync() + # Render the scene. + viewer.sync() viewer.close()