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
This commit is contained in:
Matija Kecman
committed by
Copybara-Service
Copybara-Service
parent
ae58855c15
commit
8079ab3d4d
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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 <model_path.xml>')
|
||||
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))
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
Reference in New Issue
Block a user