chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor Viewer to be a base class owning the communication with the simulation, handler registry, core visualization objects and the render function

Browse Source 
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
This commit is contained in:
Matija Kecman
2026-07-06 20:49:58 -07:00
committed by Copybara-Service
parent ae58855c15
commit 8079ab3d4d
9 changed files with 384 additions and 322 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/experimental/studio/viewer_app.py
+83 -182
View File
@@ -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()