From 8079ab3d4da4597b83f6884fe361cd93ea0d883b Mon Sep 17 00:00:00 2001 From: Matija Kecman Date: Mon, 6 Jul 2026 20:49:58 -0700 Subject: [PATCH] Refactor Viewer to be a base class owning the communication with the simulation, handler registry, core visualization objects and the render function ViewerApp is simplified, interacting with the endpoint and handlers through the Viewer instance. This change makes it possible to write simulation viewers without using the UI/UX provided by ViewerApp (which should be renamed StudioApp) PiperOrigin-RevId: 943642660 Change-Id: Id8eb3d4e6a27ceda93833367fec4699fa6bc838e --- .../experimental/studio/launch_passive.py | 7 +- python/mujoco/experimental/studio/messages.py | 18 +- .../experimental/studio/native_viewer.py | 87 +++--- .../experimental/studio/sample/implot.py | 10 +- python/mujoco/experimental/studio/studio.py | 53 ---- python/mujoco/experimental/studio/viewer.py | 28 +- .../mujoco/experimental/studio/viewer_app.py | 265 ++++++------------ .../experimental/studio/viewer_protocol.py | 189 +++++++++++-- .../experimental/studio/viewer_utils.py | 49 ++++ 9 files changed, 384 insertions(+), 322 deletions(-) delete mode 100644 python/mujoco/experimental/studio/studio.py create mode 100644 python/mujoco/experimental/studio/viewer_utils.py diff --git a/python/mujoco/experimental/studio/launch_passive.py b/python/mujoco/experimental/studio/launch_passive.py index d2168da8..5ecd17ae 100644 --- a/python/mujoco/experimental/studio/launch_passive.py +++ b/python/mujoco/experimental/studio/launch_passive.py @@ -23,7 +23,6 @@ from typing import Any from mujoco.experimental.studio import endpoints from mujoco.experimental.studio import messages -from mujoco.experimental.studio import viewer_app from mujoco.experimental.studio import viewer_handle from mujoco.experimental.studio import viewer_protocol @@ -100,13 +99,15 @@ def run_viewer_target( if config.viewer_mode == viewer_protocol.ViewerMode.NATIVE: from mujoco.experimental.studio import native_viewer # pylint: disable=g-import-not-at-top - viewer = native_viewer.NativeViewer(config) + viewer = native_viewer.NativeViewer( + config, viewer_endpoint, handlers=handlers + ) elif config.viewer_mode == viewer_protocol.ViewerMode.WEB: raise NotImplementedError('Web viewer not implemented yet') else: raise ValueError(f'Unknown viewer mode: {config.viewer_mode!r}') - viewer_app.run_viewer(viewer, viewer_endpoint, handlers=handlers) + viewer_protocol.run_viewer_loop(viewer) def launch_passive( diff --git a/python/mujoco/experimental/studio/messages.py b/python/mujoco/experimental/studio/messages.py index ea2f1cc2..fbefc0f8 100644 --- a/python/mujoco/experimental/studio/messages.py +++ b/python/mujoco/experimental/studio/messages.py @@ -116,9 +116,25 @@ class ResetEvent(Event): @dataclasses.dataclass(frozen=True) class ModelEvent(Event): - """An event that transports a MuJoCo model.""" + """An event that transports a MuJoCo model. + + Attributes: + model: The compiled MuJoCo model. + path: Optional file path the model was loaded from. + """ model: mujoco.MjModel + path: str = '' + + +@dataclasses.dataclass(frozen=True) +class BuildGuiEvent(Event): + """Lifecycle event dispatched on every frame on the viewer side to build ImGui elements.""" + + +@dataclasses.dataclass(frozen=True) +class UpdateEvent(Event): + """Lifecycle event dispatched on every frame on the viewer side before building GUI.""" @dataclasses.dataclass(frozen=True) diff --git a/python/mujoco/experimental/studio/native_viewer.py b/python/mujoco/experimental/studio/native_viewer.py index 5e1e341a..eefa3b43 100644 --- a/python/mujoco/experimental/studio/native_viewer.py +++ b/python/mujoco/experimental/studio/native_viewer.py @@ -13,27 +13,33 @@ # limitations under the License. """Simulation-agnostic native viewer for MuJoCo models. -This class is simulation-agnostic and as such it does not own the model or data. - 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. """ +from typing import Any + import mujoco +from mujoco.experimental.studio import endpoints from mujoco.experimental.studio import native_viewer_cc as _viewer from mujoco.experimental.studio import ux -from mujoco.experimental.studio import viewer_protocol as vp +from mujoco.experimental.studio import viewer_protocol + from mujoco.experimental.dear_imgui import dear_imgui as imgui -class NativeViewer(vp.Viewer): +class NativeViewer(viewer_protocol.Viewer): """Simulation-agnostic native viewer for MuJoCo models.""" def __init__( self, - config: vp.ViewerConfig, + config: viewer_protocol.ViewerConfig, + endpoint: endpoints.ViewerEndpoint, *, + model: mujoco.MjModel | None = None, + model_path: str = '', + handlers: list[Any] | None = None, camera: mujoco.MjvCamera | None = None, vis_options: mujoco.MjvOption | None = None, perturb: mujoco.MjvPerturb | None = None, @@ -42,46 +48,47 @@ class NativeViewer(vp.Viewer): ) -> None: """Initializes the NativeViewer. - The viewer creates and modifies its own camera, perturbation, and - visualization option objects unless they are provided. - Args: config: Viewer window configuration. + 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. 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. render_flags: Render flags. Internal object is created if None. extra_geoms: List of extra geoms. Internal list is created if None. """ - self.config = config - - # Set members of vp.Viewer. - self.camera = camera or mujoco.MjvCamera() - self.perturb = perturb or mujoco.MjvPerturb() - self.vis_options = vis_options or mujoco.MjvOption() - self.extra_geoms = extra_geoms or [] - if render_flags is not None: - self.render_flags = render_flags - else: - self.render_flags = ux.RenderFlags() - # 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] + super().__init__( + config, + endpoint, + model=model, + model_path=model_path, + handlers=handlers, + camera=camera, + vis_options=vis_options, + perturb=perturb, + render_flags=render_flags, + extra_geoms=extra_geoms, + ) # Create the renderer. self._viewer = _viewer.Viewer( config.title, config.width, config.height, config.gfx or '' ) - # This class does not own the model but we need to know if the model being - # rendered has changed, so we store the unique python object id here so we - # can use it to detect model changes. + # Track the python object id of the model currently loaded in the renderer + # so we can detect when the model changes and re-initialize. self._renderer_model_id = id(None) - self._is_running = True ctx = self._viewer.GetImGuiContext() imgui.SetCurrentContext(ctx) ux.set_imgui_context(ctx) + # Dispatch lifecycle event so handlers can cache the viewer reference. + self.dispatch(viewer_protocol.ViewerInitEvent(viewer=self)) + def _sync_renderer(self, model: mujoco.MjModel) -> None: """Re-initializes the renderer if the model object has changed.""" if id(model) != self._renderer_model_id: @@ -90,26 +97,16 @@ class NativeViewer(vp.Viewer): def is_running(self) -> bool: """Poll for a new frame; returns ``False`` when the window is closed.""" - if not self._is_running: - return False - self._is_running = self._viewer.NewFrame() - return self._is_running + if super().is_running() and not self._viewer.NewFrame(): + self.close() + return super().is_running() - def sync( - self, - model: mujoco.MjModel, - data: mujoco.MjData, - ) -> None: - """Render the scene and present it to the window. - - Args: - model: The MuJoCo model provided by the simulation. - data: The MuJoCo data provided by the simulation. - """ - self._sync_renderer(model) + def sync(self) -> None: + """Render the scene and present it to the window.""" + self._sync_renderer(self.model) self._viewer.Present( - model, - data, + self.model, + self.data, self.perturb, self.camera, self.vis_options, @@ -117,10 +114,6 @@ class NativeViewer(vp.Viewer): self.extra_geoms, ) - def close(self) -> None: - """Close the viewer.""" - self._is_running = False - # TODO(matijak): Remove stop() and rename callers to close(). def stop(self) -> None: """Stop the viewer.""" diff --git a/python/mujoco/experimental/studio/sample/implot.py b/python/mujoco/experimental/studio/sample/implot.py index 5b9c22dc..eb9b87e1 100644 --- a/python/mujoco/experimental/studio/sample/implot.py +++ b/python/mujoco/experimental/studio/sample/implot.py @@ -108,7 +108,7 @@ def _setup_angle_axis(plot_size: imgui.Vec2) -> None: ) -class BodyInspectorHandler: +class BodyInspector: """Handler that draws body-inspection plots using ImGui/ImPlot.""" def __init__(self) -> None: @@ -124,7 +124,7 @@ class BodyInspectorHandler: self._app = event.viewer_app @messages.handler - def inspect_body(self, _: viewer_app.BuildGuiEvent) -> None: + def inspect_body(self, _: messages.BuildGuiEvent) -> None: """Renders the body-inspection charts in ImGui/ImPlot.""" app = self._app if app is None: @@ -215,8 +215,7 @@ def main(argv: list[str]) -> None: print('Usage: implot ') sys.exit(1) - data = parser.parse(argv[1]) - if data is None: + if (data := parser.parse(argv[1])) is None: print(f'Error loading model from {argv[1]!r}') sys.exit(1) model = data.model @@ -232,7 +231,8 @@ def main(argv: list[str]) -> None: ) with launch_passive.launch_passive( - config, viewer_handlers=[BodyInspectorHandler()] + config, + viewer_handlers=[viewer_app.ViewerApp(), BodyInspector()], ) as handle: handle.send_to_viewer(messages.ModelEvent(model=model)) diff --git a/python/mujoco/experimental/studio/studio.py b/python/mujoco/experimental/studio/studio.py deleted file mode 100644 index 8d4b4754..00000000 --- a/python/mujoco/experimental/studio/studio.py +++ /dev/null @@ -1,53 +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. -"""This script runs Studio from Python, visualized in a native viewer.""" - -from absl import app as absl_app -from absl import flags as absl_flags -from mujoco.experimental.studio import native_viewer -from mujoco.experimental.studio import studio_app -from mujoco.experimental.studio import viewer_protocol - -_GFX = absl_flags.DEFINE_enum( - 'gfx', None, viewer_protocol.GFX_MODES, 'Rendering graphics mode.' -) -_WIDTH = absl_flags.DEFINE_integer('width', 1200, 'Width of the output image.') -_HEIGHT = absl_flags.DEFINE_integer('height', 800, 'Height of the output image') - - -def main(argv: list[str]) -> None: - app = studio_app.StudioApp.from_argv(argv) - - # Initialize the viewer. - config = viewer_protocol.ViewerConfig( - width=_WIDTH.value, - height=_HEIGHT.value, - gfx=_GFX.value, - ) - viewer = native_viewer.NativeViewer(config) - - # Main viewer loop. - while viewer.is_running(): - if not app.update_from_viewer(viewer): - break - - app.build_gui(viewer.camera, viewer.vis_options, viewer.render_flags) - - viewer.sync(app.model, app.data) - - viewer.stop() - - -if __name__ == '__main__': - absl_app.run(main) diff --git a/python/mujoco/experimental/studio/viewer.py b/python/mujoco/experimental/studio/viewer.py index 17086154..7d773262 100644 --- a/python/mujoco/experimental/studio/viewer.py +++ b/python/mujoco/experimental/studio/viewer.py @@ -19,6 +19,7 @@ from mujoco.experimental.studio import launch_passive 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 from mujoco.experimental.studio import viewer_protocol vp = viewer_protocol @@ -40,24 +41,23 @@ def main(argv: list[str]) -> None: viewer_mode=_VIEWER.value, ) - # Get the model path, if provided. - model_path = None - if _MJCF_PATH.value is not None: - model_path = _MJCF_PATH.value - elif len(argv) > 1 and not argv[1].startswith('--'): - model_path = argv[1] + # Resolve model path, if provided. + model_path = _MJCF_PATH.value or ( + argv[1] if len(argv) > 1 and not argv[1].startswith('--') else None + ) - with launch_passive.launch_passive(config) as handle: - # Load the model if we have a path. - model, data = None, None - if model_path is not None: - data = parser.parse(model_path) - if data is not None: - model = data.model + # Load model if path was provided. + data, model = None, None + if model_path and (data := parser.parse(model_path)): + model = data.model + with launch_passive.launch_passive( + config, + viewer_handlers=[viewer_app.ViewerApp()], + ) as handle: # Send the model to the viewer, if we have a model. if model is not None: - handle.send_to_viewer(messages.ModelEvent(model=model)) + handle.send_to_viewer(messages.ModelEvent(model=model, path=model_path)) # Run the simulation. step_control = sim.StepControl() diff --git a/python/mujoco/experimental/studio/viewer_app.py b/python/mujoco/experimental/studio/viewer_app.py index 63afdbf7..e294c3ab 100644 --- a/python/mujoco/experimental/studio/viewer_app.py +++ b/python/mujoco/experimental/studio/viewer_app.py @@ -15,18 +15,14 @@ import copy 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 from mujoco.experimental.studio import studio_app_events from mujoco.experimental.studio import ux from mujoco.experimental.studio import viewer_protocol -import numpy as np +from mujoco.experimental.studio import viewer_utils from mujoco.experimental.dear_imgui import dear_imgui as imgui @@ -42,91 +38,63 @@ class ViewerAppInitEvent(messages.Event): viewer_app: 'ViewerApp' -@dataclasses.dataclass(frozen=True) -class BuildGuiEvent(messages.Event): - """Lifecycle event dispatched on every frame on the viewer side to build ImGui elements.""" - - -@dataclasses.dataclass(frozen=True) -class UpdateEvent(messages.Event): - """Lifecycle event dispatched on every frame on the viewer side before building GUI.""" - - class ViewerApp: - """Viewer component of Studio.""" + """ViewerApp wraps a Viewer and adds Studio UI/UX.""" - def __init__( - self, - viewer: viewer_protocol.Viewer, - endpoint: endpoints.ViewerEndpoint, - *, - handlers: list[Any] | None = None, - ) -> None: - """Initializes the Studio application. + @property + def viewer(self) -> viewer_protocol.Viewer: + assert self._viewer is not None + return self._viewer - 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) - mujoco.mj_forward(self.model, self.data) - self._last_model_id: int | None = id(self.model) - self.model_path: str = '' - self.endpoint = endpoint + @viewer.setter + def viewer(self, value: viewer_protocol.Viewer | None) -> None: + self._viewer = value - # 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) + @property + def model(self) -> mujoco.MjModel: + return self.viewer.model - self.step_control_state = sim.StepControl() # ONLY for state! - self.ux_state = ux.UxState() + @model.setter + def model(self, value: mujoco.MjModel) -> None: + self.viewer.model = value + + @property + def model_path(self) -> str: + return self.viewer.model_path + + @model_path.setter + def model_path(self, value: str) -> None: + self.viewer.model_path = value + + @property + def data(self) -> mujoco.MjData: + return self.viewer.data + + @data.setter + def data(self, value: mujoco.MjData) -> None: + self.viewer.data = value + + def __init__(self) -> None: + self._viewer: viewer_protocol.Viewer | None = None self.theme = ux.GuiTheme.LIGHT self.show_stats = False self.show_solver = False - self.should_quit = False self.status = 'Ready' - # TODO(matijak): This should be part of the viewer, also making a struct to - # pass it around with the camera would be convenient. - self._cam_speed = 0.001 + self._reset_app_state() - # Dispatch lifecycle event so handlers can cache the app reference. - self._handlers.dispatch(ViewerAppInitEvent(viewer_app=self)) + @messages.handler(priority=messages.Priority.CRITICAL) + def _on_viewer_init(self, event: viewer_protocol.ViewerInitEvent) -> None: + self.viewer = event.viewer + self.viewer.dispatch(ViewerAppInitEvent(viewer_app=self)) + + def _reset_app_state(self) -> None: + """Resets ViewerApp-specific state (step control, ux).""" + self.step_control_state = sim.StepControl() + self.ux_state = ux.UxState() def close(self) -> None: - """Signals the sim to exit and closes the viewer endpoint, releasing resources.""" - if not self.should_quit: - self.should_quit = True - try: - self.endpoint.send_to_sim(messages.ExitEvent()) - except Exception: # pylint: disable=broad-exception-caught - pass # Ignore exceptions, the sim may have already closed. - self.endpoint.close() self.viewer.close() - def _send_viewer_snapshots(self) -> None: - """Sends per-frame viewer-to-sim snapshots. - - Called once per viewer frame. The snapshot channel coalesces values so the - sim always sees the latest viewer state without accumulating a backlog. - """ - noise_scale, noise_rate = self.step_control_state.get_noise_parameters() - self.endpoint.send_to_sim( - messages.StepControlSnapshot( - pause_state=self.step_control_state.get_pause_state(), - speed=self.step_control_state.get_speed(), - noise_scale=noise_scale, - noise_rate=noise_rate, - ) - ) - self.endpoint.send_to_sim( - messages.MjOptionSnapshot(opt=copy.deepcopy(self.model.opt)) - ) - def handle_keyboard_events(self) -> None: """Handles keyboard events.""" @@ -153,11 +121,11 @@ class ViewerApp: if is_freecam_wasd: handled, cam_speed = ( studio_app_events.handle_freecam_wasd_keyboard_events( - self.model, self.data, self.viewer.camera, self._cam_speed + self.model, self.data, self.viewer.camera, self.viewer.cam_speed ) ) if handled: - self._cam_speed = cam_speed + self.viewer.cam_speed = cam_speed return def handle_camera_tracking_mouse_events(self) -> None: @@ -187,25 +155,17 @@ class ViewerApp: """Reset the physics.""" mujoco.mj_resetData(self.model, self.data) mujoco.mj_forward(self.model, self.data) - self.endpoint.send_to_sim(messages.ResetEvent()) + self.viewer.send_to_sim(messages.ResetEvent()) # Discard any pre-reset snapshots so we don't overwrite the reset state. - self.endpoint.get_sim_snapshots() + self.viewer.get_sim_snapshots() def apply_perturb(self) -> None: """Apply perturbation the model.""" - perturb = self.viewer.perturb - if ( + is_paused = ( self.step_control_state.get_pause_state() - != sim.PauseState.NORMAL_PAUSED - ): - sig = mujoco.mjtState.mjSTATE_XFRC_APPLIED.value - size = mujoco.mj_stateSize(self.model, sig) - zero_state = np.zeros(size, np.float64) - mujoco.mj_setState(self.model, self.data, zero_state, sig) - mujoco.mjv_applyPerturbPose(self.model, self.data, perturb, 0) - mujoco.mjv_applyPerturbForce(self.model, self.data, perturb) - else: - mujoco.mjv_applyPerturbPose(self.model, self.data, perturb, 1) + == sim.PauseState.NORMAL_PAUSED + ) + viewer_utils.apply_perturb(self.viewer, self.model, self.data, is_paused) def reset_physics_gui(self) -> None: """GUI to Reset the physics i.e., the reset button.""" @@ -218,7 +178,7 @@ class ViewerApp: def is_running(self) -> bool: """Returns True if the application should continue running.""" - return self.viewer.is_running() and not self.should_quit + return self.viewer.is_running() def update(self) -> None: """Update the simulation and handle user input. @@ -228,56 +188,44 @@ class ViewerApp: Applies the perturbations and advances the physics. """ drop_file = self.viewer.get_drop_file() - # Handle file drop: load the model, update local state, notify sim. + # Handle file drop: update viewer model/data, reset app state, notify sim. if drop_file: try: data = parser.parse(drop_file) if data is not None: - self.model, self.data = data.model, data - self.model_path = drop_file - self.endpoint.send_to_sim(messages.ModelEvent(model=self.model)) + # Notify all handlers on the sim side + self.viewer.send_to_sim( + messages.ModelEvent(model=data.model, path=drop_file) + ) + # Notify all handlers on the viewer side. + self.viewer.dispatch( + messages.ModelEvent(model=data.model, path=drop_file) + ) + # Discard all snapshots, including any stale state snapshots + self.viewer.get_sim_snapshots() except Exception as ex: # pylint: disable=broad-except print(f'Error loading model from {drop_file!r}: {ex}') - # Process incoming events from the simulation. - incoming_events = self.endpoint.get_sim_events() - for event in incoming_events: - self._handlers.dispatch(event) - - # Detect model change from drop_file or ModelEvent (or external swap). - model_changed = False - if ( - self._last_model_id is not None - and id(self.model) != self._last_model_id - ): - model_changed = True - mujoco.mj_forward(self.model, self.data) - self.step_control_state = sim.StepControl() - self.ux_state = ux.UxState() - self._last_model_id = id(self.model) - - # Process incoming snapshots from the simulation process. - incoming_snapshots = self.endpoint.get_sim_snapshots() - for snapshot in incoming_snapshots: - if not model_changed: - self._handlers.dispatch(snapshot) - + # Handle user input. self.handle_mouse_events() self.handle_keyboard_events() - xfrc_sig: int = int(mujoco.mjtState.mjSTATE_XFRC_APPLIED) - xfrc_size: int = mujoco.mj_stateSize(self.model, xfrc_sig) - xfrc_state: np.ndarray = np.zeros(xfrc_size, np.float64) - # Apply perturbation forces from the viewer. self.apply_perturb() - mujoco.mj_getState(self.model, self.data, xfrc_state, xfrc_sig) - self.endpoint.send_to_sim( - messages.PerturbEvent(state=xfrc_state, state_sig=xfrc_sig) - ) # Send viewer-to-sim snapshots (step control, model options) each frame. - self._send_viewer_snapshots() + noise_scale, noise_rate = self.step_control_state.get_noise_parameters() + self.viewer.send_to_sim( + messages.StepControlSnapshot( + pause_state=self.step_control_state.get_pause_state(), + speed=self.step_control_state.get_speed(), + noise_scale=noise_scale, + noise_rate=noise_rate, + ) + ) + self.viewer.send_to_sim( + messages.MjOptionSnapshot(opt=copy.deepcopy(self.model.opt)) + ) def build_gui(self) -> None: """Emit full Studio UI.""" @@ -419,68 +367,21 @@ class ViewerApp: imgui.End() imgui.PopStyleVar(3) - # CRITICAL priority ensures the model is copied before any other handlers - # are notified of the ModelEvent. @messages.handler(priority=messages.Priority.CRITICAL) 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 False # Do not consume to allow other handlers to recieve the event. + del event # Model/data are owned by the Viewer. + self._reset_app_state() + return False # Do not consume to allow other handlers to receive the event. @messages.handler(priority=messages.Priority.INTERNAL) def _on_exit(self, _: messages.ExitEvent) -> bool: - self.should_quit = True - self.viewer.close() + self.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 - - @messages.handler(priority=messages.Priority.INTERNAL) - def _on_update(self, _: UpdateEvent) -> None: + def _on_update(self, _: messages.UpdateEvent) -> None: self.update() @messages.handler(priority=messages.Priority.INTERNAL) - def _on_build_gui(self, _: BuildGuiEvent) -> None: + def _on_build_gui(self, _: messages.BuildGuiEvent) -> None: self.build_gui() - - -def run_viewer( - viewer: viewer_protocol.Viewer, - viewer_endpoint: endpoints.ViewerEndpoint, - *, - handlers: list[Any] | None = None, -) -> None: - """Run the viewer loop with the given viewer and endpoint. - - This is the common viewer-side entry point used by both passive and - subprocess launchers. - - Runs until the viewer window is closed or the app requests a quit. - On exit, sends an ExitEvent to the sim side so ViewerHandle.sync() detects - the shutdown. - - Args: - viewer: A viewer display surface conforming to the Viewer protocol. - viewer_endpoint: The viewer endpoint for communication with the sim side. - handlers: Optional list of handler instances for viewer processing, which - are classes with methods decorated with ``@handler``. - """ - # pylint: disable=protected-access - app = ViewerApp(viewer, viewer_endpoint, handlers=handlers) - - # Viewer main loop. - while app.is_running(): - app._handlers.dispatch(UpdateEvent()) - app._handlers.dispatch(BuildGuiEvent()) - app.viewer.sync(app.model, app.data) - - app.close() diff --git a/python/mujoco/experimental/studio/viewer_protocol.py b/python/mujoco/experimental/studio/viewer_protocol.py index bff5841d..9dc83cca 100644 --- a/python/mujoco/experimental/studio/viewer_protocol.py +++ b/python/mujoco/experimental/studio/viewer_protocol.py @@ -11,15 +11,17 @@ # 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. -"""Structural protocol defining the common viewer interface. - -ViewerApp uses the protocol for convenience methods that accept any viewer. -""" +"""Base class and configuration for any viewer.""" +import abc +import copy import dataclasses import enum -from typing import Any, Protocol +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 ux import numpy as np @@ -65,6 +67,7 @@ class ViewerConfig: @dataclasses.dataclass class SimToView: """A message sent from the simulation to the viewer.""" + model: mujoco.MjModel | None = None state: np.ndarray | None = None state_sig: int = 0 @@ -74,6 +77,7 @@ class SimToView: @dataclasses.dataclass class ViewToSim: """A message sent from the viewer to the simulation.""" + state: np.ndarray | None = None state_sig: int = 0 reset: bool = False @@ -82,32 +86,183 @@ class ViewToSim: # ----------------------------------------------------------------------------- -# Structural interface for any viewer. +# Base class for any viewer. # ----------------------------------------------------------------------------- -class Viewer(Protocol): - """Structural interface for any viewer.""" +@dataclasses.dataclass(frozen=True) +class ViewerInitEvent(messages.Event): + """Lifecycle event dispatched once when the concrete Viewer is initialized. - camera: mujoco.MjvCamera - perturb: mujoco.MjvPerturb - vis_options: mujoco.MjvOption - render_flags: ux.RenderFlags - extra_geoms: list[mujoco.MjvGeom] + Handlers that need access to the Viewer should handle this event + and cache the reference. + """ - def is_running(self) -> bool: - ... + viewer: 'Viewer' - def sync(self, model: mujoco.MjModel, data: mujoco.MjData) -> None: - ... + +class Viewer(abc.ABC): + """Base class for any viewer. + + Owns the communication endpoint, handler registry and core visualization + objects. The application is rendered by calling ``sync()``. + """ + + def __init__( + self, + config: ViewerConfig, + endpoint: endpoints.ViewerEndpoint, + *, + model: mujoco.MjModel | None = None, + model_path: str = '', + handlers: list[Any] | None = None, + camera: mujoco.MjvCamera | None = None, + vis_options: mujoco.MjvOption | None = None, + perturb: mujoco.MjvPerturb | None = None, + render_flags: ux.RenderFlags | None = None, + extra_geoms: list[mujoco.MjvGeom] | None = None, + ) -> None: + """Initializes the Viewer. + + Args: + config: Viewer window configuration. + endpoint: The viewer endpoint for communication with the sim side. + model: Optional initial MjModel. If None, an empty model is created from + 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. + 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. + render_flags: Render flags. Internal object is created if None. + extra_geoms: List of extra geoms. Internal list is created if None. + """ + self.config = config + self._endpoint = endpoint + self._is_running = True + + # Viewer-owned model and data. + if model is None: + model = mujoco.MjSpec().compile() + self.model: mujoco.MjModel + self.data: mujoco.MjData + self.model_path: str = '' + self.load_model(model, model_path) + + # Visual state. + self.camera = camera or mujoco.MjvCamera() + self.cam_speed = 0.001 + self.perturb = perturb or mujoco.MjvPerturb() + self.vis_options = vis_options or mujoco.MjvOption() + self.extra_geoms = extra_geoms or [] + if render_flags is not None: + self.render_flags = render_flags + else: + self.render_flags = ux.RenderFlags() + # 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) def close(self) -> None: + """Closes the viewer, sends an exit event and shuts down the endpoint.""" + if self._is_running: + self._is_running = False + try: + self.send_to_sim(messages.ExitEvent()) + except Exception: # pylint: disable=broad-exception-caught + pass # Ignore exceptions, the sim may have already closed. + self._endpoint.close() + + def is_running(self) -> bool: + """Returns True while the viewer has not been closed.""" + return self._is_running + + def send_to_sim(self, message: messages.Message) -> None: + """Sends a message to the simulation process.""" + self._endpoint.send_to_sim(message) + + def get_sim_events(self) -> list[messages.Event]: + """Returns all pending events from the simulation.""" + return self._endpoint.get_sim_events() + + def get_sim_snapshots(self) -> list[messages.Snapshot]: + """Returns all pending latest snapshots from the simulation, one per type.""" + return self._endpoint.get_sim_snapshots() + + def dispatch(self, message: messages.Message) -> None: + """Dispatches a message to registered handlers in priority order.""" + self.handlers.dispatch(message) + + def load_model(self, model: mujoco.MjModel, model_path: str = '') -> None: + """Deep-copies a model and creates fresh data for the viewer.""" + self.model_path = model_path + self.model = copy.deepcopy(model) + self.data = mujoco.MjData(self.model) + assert id(self.model) != id(model) + mujoco.mj_forward(self.model, self.data) + + @messages.handler(priority=messages.Priority.CRITICAL) + def _on_model(self, event: messages.ModelEvent) -> bool: + """Deep-copies the incoming model so the Viewer owns its data.""" + self.load_model(event.model, event.path) + self.extra_geoms.clear() + return False # Do not consume; let other handlers see the event. + + @messages.handler(priority=messages.Priority.CRITICAL) + def _on_state(self, event: messages.StateSnapshot) -> bool: + """Applies incoming simulation state to the viewer's model/data.""" + 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 False # Do not consume; let other handlers see the event. + + @abc.abstractmethod + def sync(self) -> None: + """Renders the scene using the viewer's current model and data.""" ... + @abc.abstractmethod def get_drop_file(self) -> str: ... + @abc.abstractmethod def upload_image( self, tex_id: int, img: str | bytes, width: int, height: int, bpp: int ) -> int: ... + + +# ----------------------------------------------------------------------------- +# Standalone viewer loop. +# ----------------------------------------------------------------------------- + + +def run_viewer_loop(viewer: Viewer) -> None: + """Minimal viewer loop: process sim messages, dispatch lifecycle events, sync. + + Runs until the viewer window is closed or an exit event is received. + On exit, closes the viewer (which sends an ExitEvent to the sim side). + + Args: + viewer: A Viewer that owns the endpoint and handler registry. + """ + while viewer.is_running(): + # Process incoming messages. + for event in viewer.get_sim_events(): + viewer.dispatch(event) + for snapshot in viewer.get_sim_snapshots(): + viewer.dispatch(snapshot) + + # Dispatch lifecycle events. + viewer.dispatch(messages.UpdateEvent()) + viewer.dispatch(messages.BuildGuiEvent()) + + # Render the scene. + viewer.sync() + + viewer.close() diff --git a/python/mujoco/experimental/studio/viewer_utils.py b/python/mujoco/experimental/studio/viewer_utils.py new file mode 100644 index 00000000..24d1b012 --- /dev/null +++ b/python/mujoco/experimental/studio/viewer_utils.py @@ -0,0 +1,49 @@ +# 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. +"""Utility functions for Studio viewer implementations.""" + +import mujoco +from mujoco.experimental.studio import messages +from mujoco.experimental.studio import viewer_protocol +import numpy as np + + +def apply_perturb( + viewer: viewer_protocol.Viewer, + model: mujoco.MjModel | None, + data: mujoco.MjData | None, + is_paused: bool = False, +) -> None: + """Apply perturbation to the model and send updated forces to the sim thread.""" + if model is None or data is None: + return + perturb = viewer.perturb + if not is_paused: + sig = int(mujoco.mjtState.mjSTATE_XFRC_APPLIED) + size = mujoco.mj_stateSize(model, sig) + zero_state = np.zeros(size, np.float64) + mujoco.mj_setState(model, data, zero_state, sig) + mujoco.mjv_applyPerturbPose(model, data, perturb, 0) + mujoco.mjv_applyPerturbForce(model, data, perturb) + else: + mujoco.mjv_applyPerturbPose(model, data, perturb, 1) + + xfrc_sig: int = int(mujoco.mjtState.mjSTATE_XFRC_APPLIED) + xfrc_size: int = mujoco.mj_stateSize(model, xfrc_sig) + xfrc_state: np.ndarray = np.zeros(xfrc_size, np.float64) + + mujoco.mj_getState(model, data, xfrc_state, xfrc_sig) + viewer.send_to_sim( + messages.PerturbEvent(state=xfrc_state, state_sig=xfrc_sig) + )