Matija Kecman
0dfa4b509a
This change replaces fixed callback protocols (e.g., ViewerGuiHook, ViewerUpdateHook, SimEventHandler) with a general-purpose, priority-based message and event handling mechanism. Key changes: - Handler decorator and registry: Introduced the `@messages.handler(priority=...)` decorator and `HandlerRegistry` (`handler_registry.py`). Methods marked as handlers are automatically discovered and dispatched by priority (CRITICAL, USER, LIBRARY, INTERNAL) or method resolution order. - Local lifecycle events: Added `ViewerAppInitEvent`, `BuildGuiEvent`, and `UpdateEvent` to `messages.py`. Custom GUI rendering and per-frame update logic can now be implemented as standard event handlers without needing separate interface protocols. - Streamlined launch and app APIs: Replaced individual hook and handler arguments in `launch_passive`, `ViewerApp`, and `ViewerHandle` with unified `viewer_handlers` and `sim_handlers` lists. - Module restructuring: Extracted simulation-side message handling and `ViewerHandle` from `sim_app.py` into a dedicated `viewer_handle.py` module, removing `sim_app.py`. - Sample updates: Migrated existing examples (such as `implot.py`) to use the new handler pattern and lifecycle events. PiperOrigin-RevId: 941729322 Change-Id: I93e7c0edf0a13a8dc854f9e2083451f7c25a1825
484 lines
16 KiB
Python
484 lines
16 KiB
Python
# 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.
|
|
"""Viewer component of Studio."""
|
|
|
|
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.dear_imgui import dear_imgui as imgui
|
|
|
|
@dataclasses.dataclass(frozen=True)
|
|
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.
|
|
"""
|
|
|
|
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."""
|
|
|
|
def __init__(
|
|
self,
|
|
viewer: viewer_protocol.Viewer,
|
|
endpoint: endpoints.ViewerEndpoint,
|
|
*,
|
|
handlers: list[Any] | None = None,
|
|
) -> None:
|
|
"""Initializes the Studio application.
|
|
|
|
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
|
|
|
|
# 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)
|
|
|
|
self.step_control_state = sim.StepControl() # ONLY for state!
|
|
self.ux_state = ux.UxState()
|
|
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
|
|
|
|
# Dispatch lifecycle event so handlers can cache the app reference.
|
|
self._handlers.dispatch(ViewerAppInitEvent(viewer_app=self))
|
|
|
|
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."""
|
|
|
|
is_freecam_wasd = self.ux_state.camera_index == ux.FREE_CAMERA_IDX
|
|
if studio_app_events.handle_step_control_keyboard_events(
|
|
self.step_control_state, self.ux_state
|
|
):
|
|
return
|
|
|
|
if studio_app_events.handle_reset_keyboard_events(self.model, self.data):
|
|
self.reset_physics()
|
|
return
|
|
|
|
if studio_app_events.handle_camera_select_keyboard_events(
|
|
self.model, self.viewer.camera, self.ux_state
|
|
):
|
|
return
|
|
|
|
if studio_app_events.handle_vis_options_keyboard_events(
|
|
self.viewer.vis_options, is_freecam_wasd
|
|
):
|
|
return
|
|
|
|
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
|
|
)
|
|
)
|
|
if handled:
|
|
self._cam_speed = cam_speed
|
|
return
|
|
|
|
def handle_camera_tracking_mouse_events(self) -> None:
|
|
"""Handles mouse events for camera tracking."""
|
|
return studio_app_events.handle_camera_tracking_mouse_events(
|
|
self.model,
|
|
self.data,
|
|
self.viewer.camera,
|
|
self.viewer.vis_options,
|
|
self.ux_state,
|
|
)
|
|
|
|
def handle_mouse_events(
|
|
self,
|
|
) -> None:
|
|
"""Handles mouse events."""
|
|
return studio_app_events.handle_mouse_events(
|
|
self.model,
|
|
self.data,
|
|
self.viewer.camera,
|
|
self.viewer.vis_options,
|
|
self.viewer.perturb,
|
|
self.ux_state,
|
|
)
|
|
|
|
def reset_physics(self) -> None:
|
|
"""Reset the physics."""
|
|
mujoco.mj_resetData(self.model, self.data)
|
|
mujoco.mj_forward(self.model, self.data)
|
|
self.endpoint.send_to_sim(messages.ResetEvent())
|
|
# Discard any pre-reset snapshots so we don't overwrite the reset state.
|
|
self.endpoint.get_sim_snapshots()
|
|
|
|
def apply_perturb(self) -> None:
|
|
"""Apply perturbation the model."""
|
|
perturb = self.viewer.perturb
|
|
if (
|
|
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)
|
|
|
|
def reset_physics_gui(self) -> None:
|
|
"""GUI to Reset the physics i.e., the reset button."""
|
|
button_size = imgui.GetFrameHeight()
|
|
square_size = imgui.Vec2(button_size, button_size)
|
|
icon_reset_model = '\uf0e2' # FontAwesome "undo" icon.
|
|
if imgui.Button(icon_reset_model, square_size):
|
|
self.reset_physics()
|
|
imgui.SetItemTooltip('Reset')
|
|
|
|
def is_running(self) -> bool:
|
|
"""Returns True if the application should continue running."""
|
|
return self.viewer.is_running() and not self.should_quit
|
|
|
|
def update(self) -> None:
|
|
"""Update the simulation and handle user input.
|
|
|
|
Handles mouse input to compute perturbations or camera motion.
|
|
Handles keyboard input e.g., for keybindings or camera motion.
|
|
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.
|
|
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))
|
|
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)
|
|
|
|
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()
|
|
|
|
def build_gui(self) -> None:
|
|
"""Emit full Studio UI."""
|
|
ux.setup_theme(self.theme)
|
|
ux.configure_docking_layout()
|
|
|
|
# -- Main menu bar --------------------------------------------------------
|
|
if imgui.BeginMainMenuBar():
|
|
if imgui.BeginMenu('File'):
|
|
if imgui.MenuItem('Quit'):
|
|
self.close()
|
|
imgui.EndMenu()
|
|
if imgui.BeginMenu('Simulation'):
|
|
imgui.EndMenu()
|
|
if imgui.BeginMenu('Charts'):
|
|
if imgui.MenuItem('Solver', '', self.show_solver):
|
|
self.show_solver = not self.show_solver
|
|
if imgui.MenuItem('Stats', '', self.show_stats):
|
|
self.show_stats = not self.show_stats
|
|
imgui.EndMenu()
|
|
if imgui.BeginMenu('Help'):
|
|
if imgui.MenuItem('Stats', '', self.show_stats):
|
|
self.show_stats = not self.show_stats
|
|
imgui.Separator()
|
|
version = f'Version {mujoco.mj_versionString()}'
|
|
imgui.MenuItem(version)
|
|
imgui.EndMenu()
|
|
imgui.EndMainMenuBar()
|
|
|
|
# -- Tool Bar -------------------------------------------------------------
|
|
if imgui.Begin('ToolBar'):
|
|
imgui.PushStyleVar(imgui.StyleVar.CellPadding, imgui.Vec2(0, 0))
|
|
if imgui.BeginTable('##ToolBarTable', 2):
|
|
imgui.TableSetupColumn('', int(imgui.TableColumnFlags.WidthStretch))
|
|
imgui.TableSetupColumn('', int(imgui.TableColumnFlags.WidthFixed))
|
|
|
|
imgui.TableNextColumn()
|
|
self.reset_physics_gui()
|
|
|
|
imgui.SameLine()
|
|
ux.step_control_gui(self.step_control_state, self.ux_state)
|
|
|
|
imgui.TableNextColumn()
|
|
ux.camera_selection_gui(
|
|
self.model,
|
|
self.data,
|
|
self.viewer.camera,
|
|
self.ux_state,
|
|
)
|
|
|
|
imgui.SameLine()
|
|
ux.label_selection_gui(self.viewer.vis_options)
|
|
|
|
imgui.SameLine()
|
|
ux.frame_selection_gui(self.viewer.vis_options)
|
|
|
|
imgui.SameLine()
|
|
changed, self.theme = ux.theme_select_gui(self.theme)
|
|
if changed:
|
|
ux.setup_theme(self.theme)
|
|
|
|
imgui.EndTable()
|
|
imgui.PopStyleVar()
|
|
imgui.End()
|
|
|
|
# -- Left pane: Options ---------------------------------------------------
|
|
node_flags = int(imgui.TreeNodeFlags.SpanAvailWidth) | int(
|
|
imgui.TreeNodeFlags.Framed
|
|
)
|
|
|
|
imgui.Begin('Options')
|
|
if imgui.TreeNodeEx('Physics Settings', node_flags):
|
|
ux.physics_gui(self.model)
|
|
imgui.TreePop()
|
|
if imgui.TreeNodeEx('Rendering Settings', node_flags):
|
|
ux.rendering_gui(
|
|
self.model,
|
|
self.viewer.vis_options,
|
|
self.viewer.render_flags,
|
|
)
|
|
imgui.TreePop()
|
|
if imgui.TreeNodeEx('Visibility Groups', node_flags):
|
|
ux.groups_gui(self.model, self.viewer.vis_options)
|
|
imgui.TreePop()
|
|
if imgui.TreeNodeEx('Visualization', node_flags):
|
|
ux.visualization_gui(
|
|
self.model,
|
|
self.viewer.vis_options,
|
|
self.viewer.camera,
|
|
)
|
|
imgui.TreePop()
|
|
imgui.End()
|
|
|
|
# -- Right pane: Inspector ------------------------------------------------
|
|
imgui.Begin('Inspector')
|
|
if imgui.TreeNodeEx('Noise', node_flags):
|
|
ux.noise_gui(self.step_control_state)
|
|
imgui.TreePop()
|
|
if imgui.TreeNodeEx('Joints', node_flags):
|
|
ux.joints_gui(self.model, self.data, self.viewer.vis_options)
|
|
imgui.TreePop()
|
|
if imgui.TreeNodeEx('Controls', node_flags):
|
|
ux.controls_gui(self.model, self.data, self.viewer.vis_options)
|
|
imgui.TreePop()
|
|
if imgui.TreeNodeEx(
|
|
'Sensors', node_flags | int(imgui.TreeNodeFlags.DefaultOpen)
|
|
):
|
|
ux.sensor_gui(self.model, self.data)
|
|
imgui.TreePop()
|
|
if imgui.TreeNodeEx('Watch', node_flags):
|
|
ux.watch_gui(self.model, self.data, self.ux_state)
|
|
imgui.TreePop()
|
|
if imgui.TreeNodeEx('State', node_flags):
|
|
ux.state_gui(self.model, self.data, self.ux_state)
|
|
imgui.TreePop()
|
|
imgui.End()
|
|
|
|
# -- Floating windows -----------------------------------------------------
|
|
if self.show_solver:
|
|
_, self.show_solver = imgui.Begin('Solver', self.show_solver)
|
|
ux.counts_gui(self.model, self.data)
|
|
ux.convergence_gui(self.model, self.data)
|
|
imgui.End()
|
|
|
|
if self.show_stats:
|
|
_, self.show_stats = imgui.Begin('Stats', self.show_stats)
|
|
paused = (
|
|
self.step_control_state.get_pause_state() != sim.PauseState.UNPAUSED
|
|
)
|
|
ux.stats_gui(self.model, self.data, paused, 0.0)
|
|
imgui.End()
|
|
|
|
# -- Status bar -----------------------------------------------------------
|
|
imgui.PushStyleVar(imgui.StyleVar.CellPadding, imgui.Vec2(0, 0))
|
|
imgui.PushStyleVar(imgui.StyleVar.FramePadding, imgui.Vec2(0, 0))
|
|
imgui.PushStyleVar(imgui.StyleVar.WindowPadding, imgui.Vec2(0, 0))
|
|
if imgui.Begin('StatusBar'):
|
|
imgui.Text(self.status)
|
|
imgui.End()
|
|
imgui.PopStyleVar(3)
|
|
|
|
@messages.handler(priority=messages.Priority.INTERNAL)
|
|
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 True
|
|
|
|
@messages.handler(priority=messages.Priority.INTERNAL)
|
|
def _on_exit(self, _: messages.ExitEvent) -> bool:
|
|
self.should_quit = True
|
|
self.viewer.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:
|
|
self.update()
|
|
|
|
@messages.handler(priority=messages.Priority.INTERNAL)
|
|
def _on_build_gui(self, _: 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()
|