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

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

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

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

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():

python/mujoco/experimental/studio/sim_app.py

@@ -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

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()

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