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
This commit is contained in:
Matija Kecman
committed by
Copybara-Service
Copybara-Service
parent
ab1af24911
commit
f1c8d3a58f
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
+9
-9
@@ -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))
|
||||
|
||||
@@ -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))
|
||||
|
||||
|
||||
@@ -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))
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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'
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user