From 0dfa4b509a7f8527d5d34b68e371586a93a4ab6f Mon Sep 17 00:00:00 2001 From: Matija Kecman Date: Thu, 2 Jul 2026 09:14:02 -0700 Subject: [PATCH] Refactor Studio customization and message dispatch to use decorator-based handlers This change replaces fixed callback protocols (e.g., ViewerGuiHook, ViewerUpdateHook, SimEventHandler) with a general-purpose, priority-based message and event handling mechanism. Key changes: - Handler decorator and registry: Introduced the `@messages.handler(priority=...)` decorator and `HandlerRegistry` (`handler_registry.py`). Methods marked as handlers are automatically discovered and dispatched by priority (CRITICAL, USER, LIBRARY, INTERNAL) or method resolution order. - Local lifecycle events: Added `ViewerAppInitEvent`, `BuildGuiEvent`, and `UpdateEvent` to `messages.py`. Custom GUI rendering and per-frame update logic can now be implemented as standard event handlers without needing separate interface protocols. - Streamlined launch and app APIs: Replaced individual hook and handler arguments in `launch_passive`, `ViewerApp`, and `ViewerHandle` with unified `viewer_handlers` and `sim_handlers` lists. - Module restructuring: Extracted simulation-side message handling and `ViewerHandle` from `sim_app.py` into a dedicated `viewer_handle.py` module, removing `sim_app.py`. - Sample updates: Migrated existing examples (such as `implot.py`) to use the new handler pattern and lifecycle events. PiperOrigin-RevId: 941729322 Change-Id: I93e7c0edf0a13a8dc854f9e2083451f7c25a1825 --- .../experimental/studio/handler_registry.py | 62 ++++++ .../experimental/studio/launch_passive.py | 82 ++++---- python/mujoco/experimental/studio/messages.py | 84 +++++++- .../experimental/studio/sample/implot.py | 92 ++++---- python/mujoco/experimental/studio/sim_app.py | 170 --------------- .../mujoco/experimental/studio/viewer_app.py | 177 +++++++--------- .../experimental/studio/viewer_handle.py | 198 ++++++++++++++++++ 7 files changed, 507 insertions(+), 358 deletions(-) create mode 100644 python/mujoco/experimental/studio/handler_registry.py delete mode 100644 python/mujoco/experimental/studio/sim_app.py create mode 100644 python/mujoco/experimental/studio/viewer_handle.py diff --git a/python/mujoco/experimental/studio/handler_registry.py b/python/mujoco/experimental/studio/handler_registry.py new file mode 100644 index 00000000..36e8ebc4 --- /dev/null +++ b/python/mujoco/experimental/studio/handler_registry.py @@ -0,0 +1,62 @@ +# Copyright 2026 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Internal handler registration, metadata stamping, and runtime dispatching.""" + +import collections +from typing import Any, Callable +from mujoco.experimental.studio import messages + + +def _discover_handlers( + handler_obj: Any, +) -> list[tuple[messages._HandlerInfo, Callable[..., Any]]]: + """Scan an object for methods marked with handler decorators.""" + handlers = [] + for name in dir(type(handler_obj)): + unbound = getattr(type(handler_obj), name, None) + if unbound is None: + continue + info = getattr(unbound, messages._HANDLER_INFO_ATTR, None) # pylint: disable=protected-access + if info is None: + continue + bound = getattr(handler_obj, name) + handlers.append((info, bound)) + return handlers + + +class HandlerRegistry: + """Runtime registry of discovered handlers from handler instances.""" + + def __init__(self, handlers: list[Any] | None = None) -> None: + self.handlers: dict[ + type[messages.Message], list[tuple[int, Callable[..., Any]]] + ] = collections.defaultdict(list) + for obj in handlers or []: + for info, bound_method in _discover_handlers(obj): + self.handlers[info.message_type].append((info.priority, bound_method)) + + def dispatch(self, message: Any) -> None: + """Dispatches a message to registered handlers in priority order.""" + # Collect all matching handlers for the given message class and all its + # ancestors/superclasses. + message_handlers = [] + for message_type in type(message).__mro__: + message_handlers.extend(self.handlers.get(message_type, [])) + + # Sort by priority; higher values execute first. + message_handlers.sort(key=lambda item: item[0], reverse=True) + for _, handler in message_handlers: + if handler(message): + # Handler returned True to consume the message; stop dispatching. + return diff --git a/python/mujoco/experimental/studio/launch_passive.py b/python/mujoco/experimental/studio/launch_passive.py index 35b2c1c2..d2168da8 100644 --- a/python/mujoco/experimental/studio/launch_passive.py +++ b/python/mujoco/experimental/studio/launch_passive.py @@ -19,22 +19,22 @@ that the calling thread uses to push simulation state into the viewer. import queue import threading +from typing import Any from mujoco.experimental.studio import endpoints from mujoco.experimental.studio import messages -from mujoco.experimental.studio import sim_app from mujoco.experimental.studio import viewer_app +from mujoco.experimental.studio import viewer_handle from mujoco.experimental.studio import viewer_protocol -sa = sim_app -va = viewer_app -vp = viewer_protocol +class _PassiveSnapshotChannel(messages.SnapshotChannel): + """SnapshotChannel for passive mode. -class _PassiveSnapshotChannel: - """SnapshotChannel for passive mode: per-type snapshot by reference, wakes the viewer thread on put.""" + Stores per-type snapshot by reference and wakes the viewer thread when put. + """ - def __init__(self): + def __init__(self) -> None: self._pending_snapshots: dict[ type[messages.Snapshot], messages.Snapshot ] = {} @@ -58,9 +58,9 @@ class _PassiveSnapshotChannel: class _PassiveEventChannel(messages.EventChannel): - """EventChannel for passive mode: a thread-safe queue.Queue of event messages.""" + """EventChannel for passive mode using a thread-safe queue.""" - def __init__(self): + def __init__(self) -> None: self._events: queue.Queue[messages.Event] = queue.Queue() def put(self, event: messages.Event) -> None: @@ -79,15 +79,23 @@ class _PassiveEventChannel(messages.EventChannel): pass -def _run_viewer_target( +def run_viewer_target( config: viewer_protocol.ViewerConfig, viewer_endpoint: endpoints.ViewerEndpoint, - viewer_gui_hook: va.ViewerGuiHook | None = None, - viewer_update_hook: va.ViewerUpdateHook | None = None, - viewer_event_handler: va.ViewerEventHandler | None = None, - viewer_snapshot_handler: va.ViewerSnapshotHandler | None = None, + handlers: list[Any] | None = None, ) -> None: - """Creates the appropriate viewer and runs the viewer loop.""" + """Creates the appropriate viewer and runs the viewer loop. + + Args: + config: Configuration specifying the viewer mode and window settings. + viewer_endpoint: Endpoint for communicating with the simulation side. + handlers: Optional list of viewer-side handler instances, which are classes + with methods decorated with ``@handler``. + + Raises: + ValueError: If the viewer mode requested in config is unknown. + NotImplementedError: If web mode is requested. + """ if config.viewer_mode == viewer_protocol.ViewerMode.NATIVE: from mujoco.experimental.studio import native_viewer # pylint: disable=g-import-not-at-top @@ -98,25 +106,15 @@ def _run_viewer_target( else: raise ValueError(f'Unknown viewer mode: {config.viewer_mode!r}') - va.run_viewer( - viewer, - viewer_endpoint, - viewer_gui_hook=viewer_gui_hook, - viewer_update_hook=viewer_update_hook, - viewer_event_handler=viewer_event_handler, - viewer_snapshot_handler=viewer_snapshot_handler, - ) + viewer_app.run_viewer(viewer, viewer_endpoint, handlers=handlers) def launch_passive( - config: vp.ViewerConfig, + config: viewer_protocol.ViewerConfig, *, - viewer_gui_hook: va.ViewerGuiHook | None = None, - viewer_update_hook: va.ViewerUpdateHook | None = None, - viewer_event_handler: va.ViewerEventHandler | None = None, - viewer_snapshot_handler: va.ViewerSnapshotHandler | None = None, - sim_event_handler: sa.SimEventHandler | None = None, -) -> sa.ViewerHandle: + viewer_handlers: list[Any] | None = None, + sim_handlers: list[Any] | None = None, +) -> viewer_handle.ViewerHandle: """Launches the Studio GUI in a daemon thread without blocking. The viewer runs the full Studio UI (toolbar, options, inspector) on the @@ -125,11 +123,10 @@ def launch_passive( Args: config: Viewer window configuration. - viewer_gui_hook: Optional hook to draw custom ImGui panels. - viewer_update_hook: Optional hook called once per frame before GUI. - viewer_event_handler: Optional handler for sim-to-viewer events. - viewer_snapshot_handler: Optional handler for sim-to-viewer snapshots. - sim_event_handler: Optional handler for viewer-to-sim events. + viewer_handlers: Optional list of viewer-side handler instances, which are + classes with methods decorated with ``@handler``. + sim_handlers: Optional list of sim-side handler instances, which are classes + with methods decorated with ``@handler``. Returns: A ViewerHandle for interacting with the viewer. @@ -142,22 +139,15 @@ def launch_passive( ) thread = threading.Thread( - target=_run_viewer_target, - args=( - config, - viewer_endpoint, - viewer_gui_hook, - viewer_update_hook, - viewer_event_handler, - viewer_snapshot_handler, - ), + target=run_viewer_target, + args=(config, viewer_endpoint, viewer_handlers), daemon=True, ) thread.start() - handle = sa.ViewerHandle( + handle = viewer_handle.ViewerHandle( sim_endpoint, is_alive_fn=thread.is_alive, - sim_event_handler=sim_event_handler, + handlers=sim_handlers, ) return handle diff --git a/python/mujoco/experimental/studio/messages.py b/python/mujoco/experimental/studio/messages.py index 10fd379c..51af9839 100644 --- a/python/mujoco/experimental/studio/messages.py +++ b/python/mujoco/experimental/studio/messages.py @@ -14,8 +14,10 @@ """Messages and Channels for Studio.""" import dataclasses -from typing import Protocol -from typing import runtime_checkable +import enum +import inspect +from typing import Any, Callable, Protocol, runtime_checkable + import mujoco from mujoco.experimental.studio import sim import numpy as np @@ -147,3 +149,81 @@ class StepControlSnapshot(Snapshot): @dataclasses.dataclass(frozen=True) class ExitEvent(Event): """An event requesting to exit.""" + + +# --------------------------------------------------------------------------- +# Message handling decorator. +# --------------------------------------------------------------------------- + +_HANDLER_INFO_ATTR = '_studio_handler_info' + + +@dataclasses.dataclass(frozen=True) +class _HandlerInfo: + """Metadata stamped on a decorated method.""" + + priority: int + message_type: type[Message] + + +class Priority(enum.IntEnum): + """Execution priority levels for message handlers. + + Handlers with higher values execute first. When multiple handlers share the + same priority, their relative order is undefined. + """ + + INTERNAL = 1 # For built-in handlers. + LIBRARY = 10 # For library extensions. + USER = 100 # For user extensions. Default when no priority is specified. + CRITICAL = 1000 # For handlers that must run before everything else. + + +def handler( + fn: Callable[..., Any] | None = None, *, priority: int = Priority.USER +) -> Callable[..., Any]: + """Decorator to mark a class method as a message handler. + + The class method must accept exactly two arguments: self and a message event. + e.g., ``@handler`` or ``@handler(priority=...)``. + + Args: + fn: The method to stamp with handler metadata. + priority: The priority of the handler. + + Returns: + The stamped method. + """ + + def stamp(target: Callable[..., Any]) -> Callable[..., Any]: + fn_name = getattr(target, '__name__', str(target)) + params = list(inspect.signature(target).parameters.values()) + if len(params) != 2: + raise TypeError( + f'{fn_name} must accept exactly two arguments, got {len(params)}' + ) + + message_type = params[1].annotation + if message_type is inspect.Parameter.empty: + raise TypeError( + f'{fn_name}: second parameter must have a type annotation' + ) + + if not ( + isinstance(message_type, type) and issubclass(message_type, Message) + ): + raise TypeError( + f'{fn_name}: second parameter type annotation {message_type} is not' + ' a Message subclass' + ) + + info = _HandlerInfo(priority=priority, message_type=message_type) + setattr(target, _HANDLER_INFO_ATTR, info) + return target + + # Used without parens: e.g., @handler + if fn is not None: + return stamp(fn) + + # Used with parens: e.g., @handler() or @handler(priority=...) + return stamp diff --git a/python/mujoco/experimental/studio/sample/implot.py b/python/mujoco/experimental/studio/sample/implot.py index 6f001b6b..5b9c22dc 100644 --- a/python/mujoco/experimental/studio/sample/implot.py +++ b/python/mujoco/experimental/studio/sample/implot.py @@ -28,16 +28,18 @@ from absl import app as _app from absl import flags as _flags import mujoco from mujoco.experimental.studio import launch_passive -from mujoco.experimental.studio import messages as msg +from mujoco.experimental.studio import messages from mujoco.experimental.studio import parser from mujoco.experimental.studio import sim -from mujoco.experimental.studio import viewer_app as va -from mujoco.experimental.studio import viewer_protocol as vp +from mujoco.experimental.studio import viewer_app +from mujoco.experimental.studio import viewer_protocol import numpy as np from mujoco.experimental.dear_imgui import dear_imgui as imgui from mujoco.experimental.implot import implot +vp = viewer_protocol + _GFX = _flags.DEFINE_enum('gfx', None, vp.GFX_MODES, 'Graphics mode.') _WIDTH = _flags.DEFINE_integer('width', 1200, 'Width of the output image.') _HEIGHT = _flags.DEFINE_integer('height', 800, 'Height of the output image') @@ -106,23 +108,34 @@ def _setup_angle_axis(plot_size: imgui.Vec2) -> None: ) -class PlottingUi(va.ViewerGuiHook): - """Custom GUI component maintaining history buffers for ImPlot charts.""" +class BodyInspectorHandler: + """Handler that draws body-inspection plots using ImGui/ImPlot.""" - def __init__(self): - self.centroid = [np.zeros(3) for _ in range(_N_HISTORY)] - self.euler = [np.zeros(3) for _ in range(_N_HISTORY)] - self.body_id = -1 + def __init__(self) -> None: + self._app: viewer_app.ViewerApp | None = None + self._centroid: list[np.ndarray] = [np.zeros(3) for _ in range(_N_HISTORY)] + self._euler: list[np.ndarray] = [np.zeros(3) for _ in range(_N_HISTORY)] + self._body_id: int = -1 - def build_gui(self, app: va.ViewerApp) -> None: - # Inspect the perturb.select body + @messages.handler + def on_viewer_app_init(self, event: viewer_app.ViewerAppInitEvent) -> None: + """Caches the ViewerApp reference on startup.""" + assert isinstance(event.viewer_app, viewer_app.ViewerApp) + self._app = event.viewer_app + + @messages.handler + def inspect_body(self, _: viewer_app.BuildGuiEvent) -> None: + """Renders the body-inspection charts in ImGui/ImPlot.""" + app = self._app + if app is None: + return if app.viewer.perturb.select > 0: - self.body_id = app.viewer.perturb.select + self._body_id = app.viewer.perturb.select # Display selected body information. - if self.body_id > 0: + if self._body_id > 0: body_name = mujoco.mj_id2name( - app.model, int(mujoco.mjtObj.mjOBJ_BODY), self.body_id + app.model, int(mujoco.mjtObj.mjOBJ_BODY), self._body_id ) io = imgui.GetIO() @@ -137,7 +150,7 @@ class PlottingUi(va.ViewerGuiHook): # ID for the window is constant for all body names. This is needed for # the window to retain its state for all bodies. window_title = ( - f'Inspect Body {body_name or "(???)"!r} ({self.body_id})###Plot' + f'Inspect Body {body_name or "(???)"!r} ({self._body_id})###Plot' ) if imgui.Begin(window_title): avail = imgui.GetContentRegionAvail() @@ -152,10 +165,16 @@ class PlottingUi(va.ViewerGuiHook): plot_flags = _setup_plot_flags(plot_size) if implot.BeginPlot('Centroid vs Time', plot_size, flags=plot_flags): _setup_time_axis(plot_size) - _setup_xpos_axis(self.centroid, plot_size) - implot.PlotLine('x', range(_N_HISTORY), [c[0] for c in self.centroid]) - implot.PlotLine('y', range(_N_HISTORY), [c[1] for c in self.centroid]) - implot.PlotLine('z', range(_N_HISTORY), [c[2] for c in self.centroid]) + _setup_xpos_axis(self._centroid, plot_size) + implot.PlotLine( + 'x', range(_N_HISTORY), [c[0] for c in self._centroid] + ) + implot.PlotLine( + 'y', range(_N_HISTORY), [c[1] for c in self._centroid] + ) + implot.PlotLine( + 'z', range(_N_HISTORY), [c[2] for c in self._centroid] + ) implot.EndPlot() if wide: @@ -164,31 +183,31 @@ class PlottingUi(va.ViewerGuiHook): if implot.BeginPlot('Euler Angle vs Time', plot_size, flags=plot_flags): _setup_time_axis(plot_size) _setup_angle_axis(plot_size) - implot.PlotLine('roll', range(_N_HISTORY), [e[0] for e in self.euler]) implot.PlotLine( - 'pitch', range(_N_HISTORY), [e[1] for e in self.euler] + 'roll', range(_N_HISTORY), [e[0] for e in self._euler] ) - implot.PlotLine('yaw', range(_N_HISTORY), [e[2] for e in self.euler]) + implot.PlotLine( + 'pitch', range(_N_HISTORY), [e[1] for e in self._euler] + ) + implot.PlotLine('yaw', range(_N_HISTORY), [e[2] for e in self._euler]) implot.EndPlot() imgui.End() - # Update plot data - self.centroid.pop(0) - self.euler.pop(0) - if self.body_id > 0 and self.body_id < app.model.nbody: - self.centroid.append(app.data.xpos[self.body_id].copy()) - # Convert quaternion to Euler angles via rotation matrix. - quat = app.data.xquat[self.body_id] + self._centroid.pop(0) + self._euler.pop(0) + if self._body_id > 0 and self._body_id < app.model.nbody: + self._centroid.append(app.data.xpos[self._body_id].copy()) + quat = app.data.xquat[self._body_id] mat = np.zeros(9) mujoco.mju_quat2Mat(mat, quat) # mat is row-major 3x3: R[i,j] = mat[3*i + j]. roll = math.atan2(mat[7], mat[8]) pitch = math.atan2(-mat[6], math.sqrt(mat[7] ** 2 + mat[8] ** 2)) yaw = math.atan2(mat[3], mat[0]) - self.euler.append(np.degrees(np.array([roll, pitch, yaw]))) + self._euler.append(np.degrees(np.array([roll, pitch, yaw]))) else: - self.centroid.append(np.zeros(3)) - self.euler.append(np.zeros(3)) + self._centroid.append(np.zeros(3)) + self._euler.append(np.zeros(3)) def main(argv: list[str]) -> None: @@ -203,9 +222,8 @@ def main(argv: list[str]) -> None: model = data.model title = os.path.basename(sys.argv[0]) - plot_ui = PlottingUi() - config = vp.ViewerConfig( + config = viewer_protocol.ViewerConfig( title=title, width=_WIDTH.value, height=_HEIGHT.value, @@ -213,8 +231,10 @@ def main(argv: list[str]) -> None: viewer_mode=_VIEWER.value, ) - with launch_passive.launch_passive(config, viewer_gui_hook=plot_ui) as handle: - handle.send_to_viewer(msg.ModelEvent(model=model)) + with launch_passive.launch_passive( + config, viewer_handlers=[BodyInspectorHandler()] + ) as handle: + handle.send_to_viewer(messages.ModelEvent(model=model)) step_control = sim.StepControl() while handle.is_running(): diff --git a/python/mujoco/experimental/studio/sim_app.py b/python/mujoco/experimental/studio/sim_app.py deleted file mode 100644 index f4f25acd..00000000 --- a/python/mujoco/experimental/studio/sim_app.py +++ /dev/null @@ -1,170 +0,0 @@ -# Copyright 2026 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Simulation application handle and event hooks.""" - -from typing import Callable, Protocol -import mujoco -from mujoco.experimental.studio import endpoints -from mujoco.experimental.studio import messages -from mujoco.experimental.studio import sim as _sim -import numpy as np - -# ----------------------------------------------------------------------------- -# Custom event handling API. -# Note: See viewer_app.py for hooks to customize the viewer behavior. -# ----------------------------------------------------------------------------- - - -class SimEventHandler(Protocol): - """Invoked once per event received from the viewer.""" - - def handle(self, event: messages.Event) -> bool: - """Return True if event was consumed, False to allow further processing.""" - ... - - -class ViewerHandle: - """A handle held by the simulation to sync the simulation with the viewer.""" - - def __init__( - self, - sim_endpoint: endpoints.SimEndpoint, - *, - sim_event_handler: SimEventHandler | None = None, - is_alive_fn: Callable[[], bool] | None = None, - ): - """Initializes the handle. - - Args: - sim_endpoint: The endpoint to use for communication with the viewer. - sim_event_handler: Optional handler for viewer-to-sim events. - 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. - """ - - self._sim_endpoint = sim_endpoint - self._is_running = True - self._is_alive_fn = is_alive_fn - self._sim_event_handler = sim_event_handler - - def close(self) -> None: - """Signals the viewer to exit and closes the sim endpoint, releasing resources.""" - 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. - self._sim_endpoint.close() - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - self.close() - - def is_running(self) -> bool: - """Returns True while the viewer is open.""" - if self._is_alive_fn is not None and not self._is_alive_fn(): - self.close() - return self._is_running - - def send_to_viewer(self, message: messages.Message) -> None: - self._sim_endpoint.send_to_viewer(message) - - def sync( - self, - model: mujoco.MjModel | None, - data: mujoco.MjData | None, - step_control: _sim.StepControl, - ) -> tuple[mujoco.MjModel | None, mujoco.MjData | None, _sim.StepControl]: - """Syncs the simulation with the viewer and returns the updated sim state. - - This method processes incoming events from the viewer, updates the sim state - accordingly, and sends the current simulation state to the viewer as a - snapshot. - - Args: - model: The current model. - data: The current data. - step_control: The current step control state. - - Returns: - The updated model, data, and step control state. - """ - - integration_sig = int(mujoco.mjtState.mjSTATE_INTEGRATION) - - # Process incoming events from the viewer. - for event in self._sim_endpoint.get_viewer_events(): - if ( - self._sim_event_handler is not None - and self._sim_event_handler.handle(event) - ): - continue - if isinstance(event, messages.ModelEvent): - # A new model was loaded in the viewer (e.g. via file drop). - model = event.model - data = mujoco.MjData(event.model) - mujoco.mj_forward(model, data) - step_control = _sim.StepControl() - elif isinstance(event, messages.PerturbEvent): - if model is not None and data is not None: - state_size = mujoco.mj_stateSize(model, event.state_sig) - if len(event.state) == state_size: - mujoco.mj_setState(model, data, event.state, event.state_sig) - elif isinstance(event, messages.ResetEvent): - if model is not None: - assert data is not None - mujoco.mj_resetData(model, data) - mujoco.mj_forward(model, data) - elif isinstance(event, messages.ExitEvent): - self._is_running = False - - # Process incoming snapshots from the viewer. - for snapshot in self._sim_endpoint.get_viewer_snapshots(): - if isinstance(snapshot, messages.StepControlSnapshot): - step_control.set_pause_state(snapshot.pause_state) - step_control.set_speed(snapshot.speed) - step_control.set_noise_parameters( - snapshot.noise_scale, snapshot.noise_rate - ) - elif isinstance(snapshot, messages.MjOptionSnapshot): - if model is not None: - for field in model.opt._all_fields: # pylint: disable=protected-access - val = getattr(snapshot.opt, field) - try: - getattr(model.opt, field)[:] = val - except (TypeError, AttributeError): - setattr(model.opt, field, val) - - # Send the simulation state to the viewer process as a snapshot. - if model is not None: - assert data is not None - integration_size = mujoco.mj_stateSize(model, integration_sig) - integration_state = np.empty(integration_size, np.float64) - mujoco.mj_getState( - model, - data, - integration_state, - integration_sig, - ) - self._sim_endpoint.send_to_viewer( - messages.StateSnapshot( - state=integration_state, state_sig=integration_sig - ), - ) - - return model, data, step_control diff --git a/python/mujoco/experimental/studio/viewer_app.py b/python/mujoco/experimental/studio/viewer_app.py index 38320e01..0dbd7f5b 100644 --- a/python/mujoco/experimental/studio/viewer_app.py +++ b/python/mujoco/experimental/studio/viewer_app.py @@ -14,10 +14,12 @@ """Viewer component of Studio.""" import copy -from typing import Protocol +import dataclasses +from typing import Any import mujoco from mujoco.experimental.studio import endpoints +from mujoco.experimental.studio import handler_registry from mujoco.experimental.studio import messages from mujoco.experimental.studio import parser from mujoco.experimental.studio import sim @@ -28,25 +30,25 @@ import numpy as np from mujoco.experimental.dear_imgui import dear_imgui as imgui -# ------------------------------------------------------------------------------ -# Viewer application and custom message handlers -# ------------------------------------------------------------------------------ +@dataclasses.dataclass(frozen=True) +class ViewerAppInitEvent(messages.Event): + """Lifecycle event dispatched once when the ViewerApp is initialised. + + Handlers that need access to the ViewerApp should handle this event + and cache the reference. + """ + + viewer_app: 'ViewerApp' -class ViewerEventHandler(Protocol): - """Invoked once per event received from the simulation.""" - - def handle(self, event: messages.Event) -> bool: - """Return True if event was consumed, False to allow further processing.""" - ... +@dataclasses.dataclass(frozen=True) +class BuildGuiEvent(messages.Event): + """Lifecycle event dispatched on every frame on the viewer side to build ImGui elements.""" -class ViewerSnapshotHandler(Protocol): - """Invoked once per snapshot received from the simulation.""" - - def handle(self, snapshot: messages.Snapshot) -> bool: - """Return True if snapshot was consumed, False to allow further processing.""" - ... +@dataclasses.dataclass(frozen=True) +class UpdateEvent(messages.Event): + """Lifecycle event dispatched on every frame on the viewer side before building GUI.""" class ViewerApp: @@ -57,10 +59,16 @@ class ViewerApp: viewer: viewer_protocol.Viewer, endpoint: endpoints.ViewerEndpoint, *, - viewer_event_handler: ViewerEventHandler | None = None, - viewer_snapshot_handler: ViewerSnapshotHandler | None = None, - ): - """Initializes the Studio application.""" + handlers: list[Any] | None = None, + ) -> None: + """Initializes the Studio application. + + Args: + viewer: A viewer display surface conforming to the Viewer protocol. + endpoint: The viewer endpoint for communication with the sim side. + handlers: Optional list of handler instances for viewer-side processing, + which are classes with methods decorated with ``@handler``. + """ self.viewer = viewer self.model: mujoco.MjModel = mujoco.MjSpec().compile() self.data: mujoco.MjData = mujoco.MjData(self.model) @@ -68,8 +76,11 @@ class ViewerApp: self._last_model_id: int | None = id(self.model) self.model_path: str = '' self.endpoint = endpoint - self.viewer_event_handler = viewer_event_handler - self.viewer_snapshot_handler = viewer_snapshot_handler + + # Instantiate handlers from user handlers + framework defaults. + all_handlers: list[Any] = list(handlers or []) + all_handlers.append(self) + self._handlers = handler_registry.HandlerRegistry(all_handlers) self.step_control_state = sim.StepControl() # ONLY for state! self.ux_state = ux.UxState() @@ -82,6 +93,9 @@ class ViewerApp: # pass it around with the camera would be convenient. self._cam_speed = 0.001 + # Dispatch lifecycle event so handlers can cache the app reference. + self._handlers.dispatch(ViewerAppInitEvent(viewer_app=self)) + def close(self) -> None: """Signals the sim to exit and closes the viewer endpoint, releasing resources.""" if not self.should_quit: @@ -112,7 +126,7 @@ class ViewerApp: messages.MjOptionSnapshot(opt=copy.deepcopy(self.model.opt)) ) - def handle_keyboard_events(self): + def handle_keyboard_events(self) -> None: """Handles keyboard events.""" is_freecam_wasd = self.ux_state.camera_index == ux.FREE_CAMERA_IDX @@ -227,25 +241,7 @@ class ViewerApp: # Process incoming events from the simulation. incoming_events = self.endpoint.get_sim_events() for event in incoming_events: - if ( - self.viewer_event_handler is not None - and self.viewer_event_handler.handle(event) - ): - continue - if isinstance(event, messages.ModelEvent): - # A new model was sent (initial load or hot-reload). - # Deep-copy the incoming model to ensure the viewer operates on its own - # isolated copy of the C++ struct. This is required if the sim/viewer - # run in different processes, and means we don't need to lock when - # running in the same process on different threads. - self.model = copy.deepcopy(event.model) - self.data = mujoco.MjData(self.model) - - # Confirm model object is distinct from the incoming event. - assert id(self.model) != id(event.model) - elif isinstance(event, messages.ExitEvent): - self.should_quit = True - self.viewer.close() + self._handlers.dispatch(event) # Detect model change from drop_file or ModelEvent (or external swap). model_changed = False @@ -262,18 +258,8 @@ class ViewerApp: # Process incoming snapshots from the simulation process. incoming_snapshots = self.endpoint.get_sim_snapshots() for snapshot in incoming_snapshots: - if ( - self.viewer_snapshot_handler is not None - and self.viewer_snapshot_handler.handle(snapshot) - ): - continue - if not model_changed and isinstance(snapshot, messages.StateSnapshot): - state_size = mujoco.mj_stateSize(self.model, snapshot.state_sig) - if len(snapshot.state) == state_size: - mujoco.mj_setState( - self.model, self.data, snapshot.state, snapshot.state_sig - ) - mujoco.mj_forward(self.model, self.data) + if not model_changed: + self._handlers.dispatch(snapshot) self.handle_mouse_events() self.handle_keyboard_events() @@ -432,40 +418,43 @@ class ViewerApp: imgui.End() imgui.PopStyleVar(3) + @messages.handler(priority=messages.Priority.INTERNAL) + def _on_model(self, event: messages.ModelEvent) -> bool: + self.model = copy.deepcopy(event.model) + self.data = mujoco.MjData(self.model) + assert id(self.model) != id(event.model) + return True -# ------------------------------------------------------------------------------ -# Viewer application run loop and gui/update customization hooks. -# ------------------------------------------------------------------------------ + @messages.handler(priority=messages.Priority.INTERNAL) + def _on_exit(self, _: messages.ExitEvent) -> bool: + self.should_quit = True + self.viewer.close() + return True + @messages.handler(priority=messages.Priority.INTERNAL) + def _on_state(self, event: messages.StateSnapshot) -> bool: + state_size = mujoco.mj_stateSize(self.model, event.state_sig) + if len(event.state) == state_size: + mujoco.mj_setState( + self.model, self.data, event.state, event.state_sig + ) + mujoco.mj_forward(self.model, self.data) + return True -class ViewerUpdateHook(Protocol): - """Invoked once per viewer frame, after default updates .""" + @messages.handler(priority=messages.Priority.INTERNAL) + def _on_update(self, _: UpdateEvent) -> None: + self.update() - def update(self, app: 'ViewerApp') -> None: - """Update custom state.""" - ... - - -class ViewerGuiHook(Protocol): - """Invoked once per viewer frame after all update() calls. - - By deferring GUI hooks until after all updates we ensure custom UI reflects - the latest application state. - """ - - def build_gui(self, app: 'ViewerApp') -> None: - """Draw custom UI.""" - ... + @messages.handler(priority=messages.Priority.INTERNAL) + def _on_build_gui(self, _: BuildGuiEvent) -> None: + self.build_gui() def run_viewer( viewer: viewer_protocol.Viewer, viewer_endpoint: endpoints.ViewerEndpoint, *, - viewer_gui_hook: ViewerGuiHook | None = None, - viewer_update_hook: ViewerUpdateHook | None = None, - viewer_event_handler: ViewerEventHandler | None = None, - viewer_snapshot_handler: ViewerSnapshotHandler | None = None, + handlers: list[Any] | None = None, ) -> None: """Run the viewer loop with the given viewer and endpoint. @@ -479,36 +468,16 @@ def run_viewer( Args: viewer: A viewer display surface conforming to the Viewer protocol. viewer_endpoint: The viewer endpoint for communication with the sim side. - viewer_gui_hook: Optional hook to draw custom ImGui panels. - viewer_update_hook: Optional hook called once per frame before GUI. - viewer_event_handler: Optional handler for simulation-to-viewer events. - viewer_snapshot_handler: Optional handler for sim-to-viewer snapshots. + handlers: Optional list of handler instances for viewer processing, which + are classes with methods decorated with ``@handler``. """ - app = ViewerApp( - viewer, - viewer_endpoint, - viewer_event_handler=viewer_event_handler, - viewer_snapshot_handler=viewer_snapshot_handler, - ) + # pylint: disable=protected-access + app = ViewerApp(viewer, viewer_endpoint, handlers=handlers) # Viewer main loop. while app.is_running(): - - # Update the viewer state and handle user input. - app.update() - - # Update the viewer state. - if viewer_update_hook is not None: - viewer_update_hook.update(app) - - # Draw the default Studio GUI. - app.build_gui() - - # Allow custom GUI elements to be added. - if viewer_gui_hook is not None: - viewer_gui_hook.build_gui(app) - - # Sync the viewer display surface with the current model and data. + app._handlers.dispatch(UpdateEvent()) + app._handlers.dispatch(BuildGuiEvent()) app.viewer.sync(app.model, app.data) app.close() diff --git a/python/mujoco/experimental/studio/viewer_handle.py b/python/mujoco/experimental/studio/viewer_handle.py new file mode 100644 index 00000000..f87241a6 --- /dev/null +++ b/python/mujoco/experimental/studio/viewer_handle.py @@ -0,0 +1,198 @@ +# Copyright 2026 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Viewer handle and event handlers for the simulation side.""" + +from typing import Any, Callable +import mujoco +from mujoco.experimental.studio import endpoints +from mujoco.experimental.studio import handler_registry +from mujoco.experimental.studio import messages +from mujoco.experimental.studio import sim as _sim +import numpy as np + + +class ViewerHandle: + """A handle for interacting with a running Studio application from the sim.""" + + def __init__( + self, + sim_endpoint: endpoints.SimEndpoint, + *, + handlers: list[Any] | None = None, + is_alive_fn: Callable[[], bool] | None = None, + ) -> None: + """Initializes the ViewerHandle. + + Args: + 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. + """ + + self._sim_endpoint = sim_endpoint + self._is_running = True + self._is_alive_fn = is_alive_fn + self.model: mujoco.MjModel | None = None + self.data: mujoco.MjData | None = None + self.step_control: _sim.StepControl | None = None + + # Instantiate handlers from user handlers + framework defaults. + all_handlers: list[Any] = list(handlers or []) + all_handlers.append(self) + self._handlers = handler_registry.HandlerRegistry(all_handlers) + + def close(self) -> None: + """Signals the viewer to exit and closes the sim endpoint.""" + 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. + self._sim_endpoint.close() + + def __enter__(self) -> 'ViewerHandle': + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: Any, + ) -> None: + self.close() + + def is_running(self) -> bool: + """Returns True while the viewer is open.""" + if self._is_alive_fn is not None and not self._is_alive_fn(): + self.close() + return self._is_running + + def send_to_viewer(self, message: messages.Message) -> None: + """Sends an event or snapshot message to the viewer process. + + Args: + message: The message to send. + """ + self._sim_endpoint.send_to_viewer(message) + + def sync( + self, + model: mujoco.MjModel | None, + data: mujoco.MjData | None, + step_control: _sim.StepControl, + ) -> tuple[mujoco.MjModel | None, mujoco.MjData | None, _sim.StepControl]: + """Syncs the simulation with the viewer and returns the updated sim state. + + This method processes incoming events from the viewer, updates the sim state + accordingly, and sends the current simulation state to the viewer as a + snapshot. + + Args: + model: The current model. + data: The current data. + step_control: The current step control state. + + Returns: + The updated model, data, and step control state. + """ + + self.model, self.data, self.step_control = model, data, step_control + + # Process incoming events from the viewer. + for event in self._sim_endpoint.get_viewer_events(): + self._handlers.dispatch(event) + + # Process incoming snapshots from the viewer. + for snapshot in self._sim_endpoint.get_viewer_snapshots(): + self._handlers.dispatch(snapshot) + + model, data, step_control = self.model, self.data, self.step_control + + # Send the simulation state to the viewer process as a snapshot. + if model is not None: + assert data is not None + integration_sig = int(mujoco.mjtState.mjSTATE_INTEGRATION) + integration_size = mujoco.mj_stateSize(model, integration_sig) + integration_state = np.empty(integration_size, np.float64) + mujoco.mj_getState( + model, + data, + integration_state, + integration_sig, + ) + self._sim_endpoint.send_to_viewer( + messages.StateSnapshot( + state=integration_state, state_sig=integration_sig + ), + ) + + return model, data, step_control + + @messages.handler(priority=messages.Priority.INTERNAL) + def _on_model(self, event: messages.ModelEvent) -> bool: + self.model = event.model + self.data = mujoco.MjData(event.model) + mujoco.mj_forward(self.model, self.data) + self.step_control = _sim.StepControl() + return True + + @messages.handler(priority=messages.Priority.INTERNAL) + def _on_perturb(self, event: messages.PerturbEvent) -> bool: + model = self.model + data = self.data + if model is not None and data is not None: + state_size = mujoco.mj_stateSize(model, event.state_sig) + if len(event.state) == state_size: + mujoco.mj_setState(model, data, event.state, event.state_sig) + return True + + @messages.handler(priority=messages.Priority.INTERNAL) + def _on_reset(self, _: messages.ResetEvent) -> bool: + model = self.model + data = self.data + if model is not None: + assert data is not None + mujoco.mj_resetData(model, data) + mujoco.mj_forward(model, data) + return True + + @messages.handler(priority=messages.Priority.INTERNAL) + def _on_exit(self, _: messages.ExitEvent) -> bool: + self._is_running = False # pylint: disable=protected-access + return True + + @messages.handler(priority=messages.Priority.INTERNAL) + def _on_step_control(self, event: messages.StepControlSnapshot) -> bool: + sc = self.step_control + if sc is not None: + sc.set_pause_state(event.pause_state) + sc.set_speed(event.speed) + sc.set_noise_parameters(event.noise_scale, event.noise_rate) + return True + + @messages.handler(priority=messages.Priority.INTERNAL) + def _on_mjoption(self, event: messages.MjOptionSnapshot) -> bool: + model = self.model + if model is not None: + for field in model.opt._all_fields: # pylint: disable=protected-access + val = getattr(event.opt, field) + try: + getattr(model.opt, field)[:] = val + except (TypeError, AttributeError): + setattr(model.opt, field, val) + return True