From f1c8d3a58fbad33c21d5915f9a1fdb037b93820c Mon Sep 17 00:00:00 2001 From: Matija Kecman Date: Mon, 10 Aug 2026 07:49:47 -0700 Subject: [PATCH] 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 --- .../experimental/studio/launch_passive.py | 38 +++++++++---------- .../experimental/studio/native_viewer.py | 11 +++--- ...handler_registry.py => plugin_registry.py} | 18 ++++----- .../experimental/studio/sample/ghost.py | 6 ++- .../experimental/studio/sample/implot.py | 2 +- python/mujoco/experimental/studio/viewer.py | 2 +- .../mujoco/experimental/studio/viewer_app.py | 20 ++++++++-- .../experimental/studio/viewer_handle.py | 22 +++++------ .../experimental/studio/viewer_protocol.py | 20 +++++----- .../mujoco/experimental/studio/web_viewer.py | 13 ++++--- 10 files changed, 84 insertions(+), 68 deletions(-) rename python/mujoco/experimental/studio/{handler_registry.py => plugin_registry.py} (82%) diff --git a/python/mujoco/experimental/studio/launch_passive.py b/python/mujoco/experimental/studio/launch_passive.py index 1c6d05c4..a891baa9 100644 --- a/python/mujoco/experimental/studio/launch_passive.py +++ b/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, - handlers: list[Any] | None = None, + endpoint: endpoints.ViewerEndpoint, + 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. - handlers: Optional list of viewer-side handler instances, which are classes + endpoint: Endpoint for communicating with the simulation side. + 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( - config, viewer_endpoint, handlers=handlers - ) + viewer = native_viewer.NativeViewer(config, endpoint, plugins=plugins) 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, - sim_handlers: list[Any] | None = None, + viewer_plugins: 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 - rendering thread. The caller keeps running and pushes state via - ``handle.sync()``. + The viewer runs on the rendering thread and renders the scene along with any + GUI built by registered plugins (such as ``ViewerApp``). The caller keeps + 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 diff --git a/python/mujoco/experimental/studio/native_viewer.py b/python/mujoco/experimental/studio/native_viewer.py index f7f444d5..b0c27db7 100644 --- a/python/mujoco/experimental/studio/native_viewer.py +++ b/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 -separating the viewer and simulation. See the sample/ folder for examples of how -to use these classes. +See the documentation for viewer_app.py for more details on the architecture +separating the viewer and simulation. """ 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, diff --git a/python/mujoco/experimental/studio/handler_registry.py b/python/mujoco/experimental/studio/plugin_registry.py similarity index 82% rename from python/mujoco/experimental/studio/handler_registry.py rename to python/mujoco/experimental/studio/plugin_registry.py index 36e8ebc4..3eeba21b 100644 --- a/python/mujoco/experimental/studio/handler_registry.py +++ b/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)): - unbound = getattr(type(handler_obj), name, None) + for name in dir(type(plugin)): + 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: - """Runtime registry of discovered handlers from handler instances.""" +class PluginRegistry: + """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)) diff --git a/python/mujoco/experimental/studio/sample/ghost.py b/python/mujoco/experimental/studio/sample/ghost.py index f2da0955..3937d938 100644 --- a/python/mujoco/experimental/studio/sample/ghost.py +++ b/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)) diff --git a/python/mujoco/experimental/studio/sample/implot.py b/python/mujoco/experimental/studio/sample/implot.py index 7ff57b43..7b74f20c 100644 --- a/python/mujoco/experimental/studio/sample/implot.py +++ b/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)) diff --git a/python/mujoco/experimental/studio/viewer.py b/python/mujoco/experimental/studio/viewer.py index 5ca3f365..c3983e21 100644 --- a/python/mujoco/experimental/studio/viewer.py +++ b/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: diff --git a/python/mujoco/experimental/studio/viewer_app.py b/python/mujoco/experimental/studio/viewer_app.py index d65c307d..8a0396a9 100644 --- a/python/mujoco/experimental/studio/viewer_app.py +++ b/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 - and cache the reference. + Plugins that need access to the ViewerApp should handle this event and cache + the reference. """ viewer_app: 'ViewerApp' diff --git a/python/mujoco/experimental/studio/viewer_handle.py b/python/mujoco/experimental/studio/viewer_handle.py index ef910c0e..79a439a9 100644 --- a/python/mujoco/experimental/studio/viewer_handle.py +++ b/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, - which are classes with methods decorated with ``@handler``. + plugins: Optional list of plugin instances for sim-side processing, which + 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. - all_handlers: list[Any] = list(handlers or []) - all_handlers.append(self) - self._handlers = handler_registry.HandlerRegistry(all_handlers) + # Instantiate handlers from user plugins + framework defaults. + all_plugins: list[Any] = list(plugins or []) + all_plugins.append(self) + 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 diff --git a/python/mujoco/experimental/studio/viewer_protocol.py b/python/mujoco/experimental/studio/viewer_protocol.py index 9c7c4864..8c133003 100644 --- a/python/mujoco/experimental/studio/viewer_protocol.py +++ b/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. - all_handlers = [self] + list(handlers or []) - self.handlers = handler_registry.HandlerRegistry(all_handlers) + # Plugin infrastructure. + all_plugins = [self] + list(plugins or []) + 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. diff --git a/python/mujoco/experimental/studio/web_viewer.py b/python/mujoco/experimental/studio/web_viewer.py index ac47c995..304b50c7 100644 --- a/python/mujoco/experimental/studio/web_viewer.py +++ b/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