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