Renamed the notion of handler classes to plugin classes in the Studio Python API

This change refactors the Python API to use the term "plugins" instead of "handlers" for classes containing decorated handler methods. The term handler is still used for an annotated method of a plugin class that handles a specific message type. Also improved some documentation.

PiperOrigin-RevId: 962150099
Change-Id: I34a8cc410cd784b088605490cfa3baccea5c9e71

python/mujoco/experimental/studio/launch_passive.py

@@ -11,9 +11,9 @@
 # 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.
-"""Non-blocking Studio launcher.
+"""Non-blocking viewer launcher.
 
-Launches the full Studio GUI in a daemon thread and returns a ``ViewerHandle``
+Launches the viewer in a daemon thread and returns a ``ViewerHandle``
 that the calling thread uses to push simulation state into the viewer.
 """
 
@@ -80,27 +80,25 @@ class _PassiveEventChannel(messages.EventChannel):
 
 def run_viewer_target(
     config: viewer_protocol.ViewerConfig,
-    viewer_endpoint: endpoints.ViewerEndpoint,
+    endpoint: endpoints.ViewerEndpoint,
-    handlers: list[Any] | None = None,
+    plugins: list[Any] | None = None,
 ) -> None:
   """Creates the appropriate viewer and runs the viewer loop.
 
   Args:
     config: Configuration specifying the viewer window settings.
-    viewer_endpoint: Endpoint for communicating with the simulation side.
+    endpoint: Endpoint for communicating with the simulation side.
-    handlers: Optional list of viewer-side handler instances, which are classes
+    plugins: Optional list of viewer-side plugin instances, which are classes
       with methods decorated with ``@handler``.
   """
   if config.gfx in ('web', 'webgl'):  # In future we may add 'webgpu' here too.
     from mujoco.experimental.studio import web_viewer  # pylint: disable=g-import-not-at-top
 
-    viewer = web_viewer.WebViewer(config, viewer_endpoint, handlers=handlers)
+    viewer = web_viewer.WebViewer(config, endpoint, plugins=plugins)
   else:
     from mujoco.experimental.studio import native_viewer  # pylint: disable=g-import-not-at-top
 
-    viewer = native_viewer.NativeViewer(
+    viewer = native_viewer.NativeViewer(config, endpoint, plugins=plugins)
-        config, viewer_endpoint, handlers=handlers
-    )
 
   viewer_protocol.run_viewer_loop(viewer)
 
@@ -108,20 +106,20 @@ def run_viewer_target(
 def launch_passive(
     config: viewer_protocol.ViewerConfig,
     *,
-    viewer_handlers: list[Any] | None = None,
+    viewer_plugins: list[Any] | None = None,
-    sim_handlers: list[Any] | None = None,
+    sim_plugins: list[Any] | None = None,
 ) -> viewer_handle.ViewerHandle:
-  """Launches the Studio GUI in a daemon thread without blocking.
+  """Launches the viewer in a daemon thread without blocking.
 
-  The viewer runs the full Studio UI (toolbar, options, inspector) on the
+  The viewer runs on the rendering thread and renders the scene along with any
-  rendering thread.  The caller keeps running and pushes state via
+  GUI built by registered plugins (such as ``ViewerApp``). The caller keeps
-  ``handle.sync()``.
+  running and pushes state via ``handle.sync()``.
 
   Args:
     config: Viewer window configuration.
-    viewer_handlers: Optional list of viewer-side handler instances, which are
+    viewer_plugins: Optional list of viewer-side plugin instances, which are
       classes with methods decorated with ``@handler``.
-    sim_handlers: Optional list of sim-side handler instances, which are classes
+    sim_plugins: Optional list of sim-side plugin instances, which are classes
       with methods decorated with ``@handler``.
 
   Returns:
@@ -136,7 +134,7 @@ def launch_passive(
 
   thread = threading.Thread(
       target=run_viewer_target,
-      args=(config, viewer_endpoint, viewer_handlers),
+      args=(config, viewer_endpoint, viewer_plugins),
       daemon=True,
   )
   thread.start()
@@ -145,6 +143,6 @@ def launch_passive(
       sim_endpoint,
       is_alive_fn=thread.is_alive,
       shutdown_fn=thread.join,
-      handlers=sim_handlers,
+      plugins=sim_plugins,
   )
   return handle

python/mujoco/experimental/studio/native_viewer.py

@@ -13,9 +13,8 @@
 # limitations under the License.
 """Simulation-agnostic native viewer for MuJoCo models.
 
-See the documentation for studio_app.py for more details on the architecture
+See the documentation for viewer_app.py for more details on the architecture
-separating the viewer and simulation. See the sample/ folder for examples of how
+separating the viewer and simulation.
-to use these classes.
 """
 
 from typing import Any
@@ -40,7 +39,7 @@ class NativeViewer(viewer_protocol.Viewer):
       *,
       model: mujoco.MjModel | None = None,
       model_path: str = '',
-      handlers: list[Any] | None = None,
+      plugins: list[Any] | None = None,
       camera: mujoco.MjvCamera | None = None,
       vis_options: mujoco.MjvOption | None = None,
       perturb: mujoco.MjvPerturb | None = None,
@@ -54,7 +53,7 @@ class NativeViewer(viewer_protocol.Viewer):
       endpoint: The viewer endpoint for communication with the sim side.
       model: Optional initial MjModel. Forwarded to the base Viewer.
       model_path: Optional path to the model file.
-      handlers: Optional list of handler instances.
+      plugins: Optional list of plugin instances.
       camera: Camera parameters. Internal object is created if None.
       vis_options: Visualization options. Internal object is created if None.
       perturb: Perturbation parameters. Internal object is created if None.
@@ -66,7 +65,7 @@ class NativeViewer(viewer_protocol.Viewer):
         endpoint,
         model=model,
         model_path=model_path,
-        handlers=handlers,
+        plugins=plugins,
         camera=camera,
         vis_options=vis_options,
         perturb=perturb,

python/mujoco/experimental/studio/plugin_registry.py

@@ -11,7 +11,7 @@
 # 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."""
+"""Plugin registry, handler discovery, and runtime dispatching."""
 
 import collections
 from typing import Any, Callable
@@ -19,30 +19,30 @@ from mujoco.experimental.studio import messages
 
 
 def _discover_handlers(
-    handler_obj: Any,
+    plugin: Any,
 ) -> list[tuple[messages._HandlerInfo, Callable[..., Any]]]:
   """Scan an object for methods marked with handler decorators."""
   handlers = []
-  for name in dir(type(handler_obj)):
+  for name in dir(type(plugin)):
-    unbound = getattr(type(handler_obj), name, None)
+    unbound = getattr(type(plugin), 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)
+    bound = getattr(plugin, name)
     handlers.append((info, bound))
   return handlers
 
 
-class HandlerRegistry:
+class PluginRegistry:
-  """Runtime registry of discovered handlers from handler instances."""
+  """Runtime registry of discovered handlers from plugin instances."""
 
-  def __init__(self, handlers: list[Any] | None = None) -> None:
+  def __init__(self, plugins: 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 obj in plugins or []:
       for info, bound_method in _discover_handlers(obj):
         self.handlers[info.message_type].append((info.priority, bound_method))
 

python/mujoco/experimental/studio/sample/ghost.py

@@ -162,7 +162,9 @@ def main(argv: list[str]) -> None:
       argv[1] if len(argv) > 1 and not argv[1].startswith('--') else None
   )
   if not model_path:
-    raise _app.UsageError('Please provide a model path argument or --model flag.')
+    raise _app.UsageError(
+        'Please provide a model path argument or --model flag.'
+    )
 
   data = None
   try:
@@ -185,7 +187,7 @@ def main(argv: list[str]) -> None:
 
   with launch_passive.launch_passive(
       config,
-      viewer_handlers=[ghost_renderer],
+      viewer_plugins=[ghost_renderer],
   ) as handle:
     handle.send_to_viewer(messages.ModelEvent(model=model))
 

python/mujoco/experimental/studio/sample/implot.py

@@ -238,7 +238,7 @@ def main(argv: list[str]) -> None:
 
   with launch_passive.launch_passive(
       config,
-      viewer_handlers=[viewer_app.ViewerApp(), BodyInspector()],
+      viewer_plugins=[viewer_app.ViewerApp(), BodyInspector()],
   ) as handle:
     handle.send_to_viewer(messages.ModelEvent(model=model))
 

python/mujoco/experimental/studio/viewer.py

@@ -55,7 +55,7 @@ def main(argv: list[str]) -> None:
 
   with launch_passive.launch_passive(
       config,
-      viewer_handlers=[viewer_app.ViewerApp()],
+      viewer_plugins=[viewer_app.ViewerApp()],
   ) as handle:
     # Send the model to the viewer, if we have a model.
     if model is not None:

python/mujoco/experimental/studio/viewer_app.py

@@ -11,7 +11,21 @@
 # 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 component of Studio."""
+"""Python implementation of the default Studio Viewer application UI.
+
+The class can be used as a plugin to run the full Studio Viewer application in
+Python.
+
+Architecture:
+  ViewerApp provides the full default UI/UX as a viewer-side plugin. This class
+  is viewer-agnostic and as such does not own camera, vis_options, or perturb
+  objects (these are provided by the viewer).
+
+  Viewer classes (NativeViewer and WebViewer) own the renderer, camera,
+  vis_options, and local deep-copied model/data used for rendering. The sim side
+  owns the physical simulation and pushes state snapshots to the viewer via
+  ViewerHandle.sync().
+"""
 
 import copy
 import dataclasses
@@ -32,8 +46,8 @@ from mujoco.experimental.dear_imgui import dear_imgui as imgui
 class ViewerAppInitEvent(messages.Event):
   """Lifecycle event dispatched once when the ViewerApp is initialised.
 
-  Handlers that need access to the ViewerApp should handle this event
+  Plugins that need access to the ViewerApp should handle this event and cache
-  and cache the reference.
+  the reference.
   """
 
   viewer_app: 'ViewerApp'

python/mujoco/experimental/studio/viewer_handle.py

@@ -16,8 +16,8 @@
 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 plugin_registry
 from mujoco.experimental.studio import sim as _sim
 import numpy as np
 
@@ -32,13 +32,13 @@ ShutdownFn = Callable[[float], None]
 
 
 class ViewerHandle:
-  """A handle for interacting with a running Studio application from the sim."""
+  """A handle for interacting with a running viewer application from the sim."""
 
   def __init__(
       self,
       sim_endpoint: endpoints.SimEndpoint,
       *,
-      handlers: list[Any] | None = None,
+      plugins: list[Any] | None = None,
       is_alive_fn: IsAliveFn | None = None,
       shutdown_fn: ShutdownFn | None = None,
   ) -> None:
@@ -46,8 +46,8 @@ class ViewerHandle:
 
     Args:
       sim_endpoint: The endpoint to use for communication with the viewer.
-      handlers: Optional list of handler instances for sim-side processing,
+      plugins: Optional list of plugin instances for sim-side processing, which
-        which are classes with methods decorated with ``@handler``.
+        are classes with methods decorated with ``@handler``.
       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()``.
@@ -61,10 +61,10 @@ class ViewerHandle:
     self.data: mujoco.MjData | None = None
     self.step_control: _sim.StepControl | None = None
 
-    # Instantiate handlers from user handlers + framework defaults.
+    # Instantiate handlers from user plugins + framework defaults.
-    all_handlers: list[Any] = list(handlers or [])
+    all_plugins: list[Any] = list(plugins or [])
-    all_handlers.append(self)
+    all_plugins.append(self)
-    self._handlers = handler_registry.HandlerRegistry(all_handlers)
+    self._plugins = plugin_registry.PluginRegistry(all_plugins)
 
   def close(self) -> None:
     """Signals the viewer to exit and waits for it to shut down."""
@@ -128,11 +128,11 @@ class ViewerHandle:
 
     # Process incoming events from the viewer.
     for event in self._sim_endpoint.get_viewer_events():
-      self._handlers.dispatch(event)
+      self._plugins.dispatch(event)
 
     # Process incoming snapshots from the viewer.
     for snapshot in self._sim_endpoint.get_viewer_snapshots():
-      self._handlers.dispatch(snapshot)
+      self._plugins.dispatch(snapshot)
 
     model, data, step_control = self.model, self.data, self.step_control
 

python/mujoco/experimental/studio/viewer_protocol.py

@@ -20,8 +20,8 @@ import enum
 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 plugin_registry
 from mujoco.experimental.studio import ux
 import numpy as np
 
@@ -88,7 +88,7 @@ class ViewToSim:
 class ViewerInitEvent(messages.Event):
   """Lifecycle event dispatched once when the concrete Viewer is initialized.
 
-  Handlers that need access to the Viewer should handle this event
+  Plugins that need access to the Viewer should handle this event
   and cache the reference.
   """
 
@@ -98,7 +98,7 @@ class ViewerInitEvent(messages.Event):
 class Viewer(abc.ABC):
   """Base class for any viewer.
 
-  Owns the communication endpoint, handler registry and core visualization
+  Owns the communication endpoint, plugin registry and core visualization
   objects. The application is rendered by calling ``sync()``.
   """
 
@@ -109,7 +109,7 @@ class Viewer(abc.ABC):
       *,
       model: mujoco.MjModel | None = None,
       model_path: str = '',
-      handlers: list[Any] | None = None,
+      plugins: list[Any] | None = None,
       camera: mujoco.MjvCamera | None = None,
       vis_options: mujoco.MjvOption | None = None,
       perturb: mujoco.MjvPerturb | None = None,
@@ -125,7 +125,7 @@ class Viewer(abc.ABC):
         an empty MjSpec. The Viewer deep-copies this model and creates its own
         MjData.
       model_path: Optional path to the model file.
-      handlers: Optional list of handler instances for viewer-side processing.
+      plugins: Optional list of plugin instances for viewer-side processing.
       camera: Camera parameters. Internal object is created if None.
       vis_options: Visualization options. Internal object is created if None.
       perturb: Perturbation parameters. Internal object is created if None.
@@ -158,9 +158,9 @@ class Viewer(abc.ABC):
       # Initted to match mujoco/src/engine/engine_vis_init.c
       self.render_flags.flags = [1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 1]
 
-    # Handler infrastructure.
+    # Plugin infrastructure.
-    all_handlers = [self] + list(handlers or [])
+    all_plugins = [self] + list(plugins or [])
-    self.handlers = handler_registry.HandlerRegistry(all_handlers)
+    self.plugins = plugin_registry.PluginRegistry(all_plugins)
 
   def close(self) -> None:
     """Closes the viewer, sends an exit event and shuts down the endpoint."""
@@ -198,7 +198,7 @@ class Viewer(abc.ABC):
 
   def dispatch(self, message: messages.Message) -> None:
     """Dispatches a message to registered handlers in priority order."""
-    self.handlers.dispatch(message)
+    self.plugins.dispatch(message)
 
   def load_model(self, model: mujoco.MjModel, model_path: str = '') -> None:
     """Deep-copies a model and creates fresh data for the viewer."""
@@ -257,7 +257,7 @@ def run_viewer_loop(viewer: Viewer) -> None:
   On exit, closes the viewer (which sends an ExitEvent to the sim side).
 
   Args:
-    viewer: A Viewer that owns the endpoint and handler registry.
+    viewer: A Viewer that owns the endpoint and plugin registry.
   """
   while True:
     # Get the next frame; this also gets the frame's mouse/keyboard events.

python/mujoco/experimental/studio/web_viewer.py

@@ -13,12 +13,15 @@
 # limitations under the License.
 """Simulation-agnostic web viewer for MuJoCo models.
 
+See the documentation for viewer_app.py for more details on the architecture
+separating the viewer and simulation.
+
 The WebViewer streams UI and simulation state to a browser:
 
   * The ImGui UI is built into a headless ImGui context and streamed to the
     browser with the NetImgui protocol through a WebSocket-to-TCP proxy.
     Input captured in the browser flows back over the same connection and is
-    injected into the headless context, so all viewer-side handlers work
+    injected into the headless context, so all viewer-side plugins work
     unmodified.
   * Physics state and render function state are streamed to the browser over
     a WebSocket with latest-wins semantics. Note that Message types are only
@@ -149,7 +152,7 @@ class WebViewer(viewer_protocol.Viewer):
       *,
       model: mujoco.MjModel | None = None,
       model_path: str = '',
-      handlers: list[Any] | None = None,
+      plugins: list[Any] | None = None,
       camera: mujoco.MjvCamera | None = None,
       vis_options: mujoco.MjvOption | None = None,
       perturb: mujoco.MjvPerturb | None = None,
@@ -166,7 +169,7 @@ class WebViewer(viewer_protocol.Viewer):
       endpoint: The viewer endpoint for communication with the sim side.
       model: Optional initial MjModel. Forwarded to the base Viewer.
       model_path: Optional path to the model file.
-      handlers: Optional list of handler instances.
+      plugins: Optional list of plugin instances.
       camera: Camera parameters. Internal object is created if None.
       vis_options: Visualization options. Internal object is created if None.
       perturb: Perturbation parameters. Internal object is created if None.
@@ -187,7 +190,7 @@ class WebViewer(viewer_protocol.Viewer):
         endpoint,
         model=model,
         model_path=model_path,
-        handlers=handlers,
+        plugins=plugins,
         camera=camera,
         vis_options=vis_options,
         perturb=perturb,
@@ -294,7 +297,7 @@ class WebViewer(viewer_protocol.Viewer):
   def _on_model(self, event: messages.ModelEvent) -> bool:
     """Loads the new model, then restarts the servers to serve it.
 
-    The handler registry discovers handlers by name, so this override replaces
+    The plugin registry discovers handlers by name, so this override replaces
     the base Viewer's _on_model and must call it explicitly to load the model
     before the servers serialize it. The browser reconnects to the new state
     server, notices the changed model identity in the payload, and reloads