Add sim_app.py and viewer_app.py to Studio.
- `sim_app.py` provides a SimEventHook and a Handle class for the simulation process to interact with the viewer, including sending state snapshots and processing viewer events. - `viewer_app.py` implements the main viewer application logic, handling user input, building the Dear ImGui GUI, and managing communication with the simulation process. It also includes hooks for custom viewer behavior. PiperOrigin-RevId: 940994701 Change-Id: I82786872cd885b15454843b3a4d1160ff0b4be2b
This commit is contained in:
Matija Kecman
committed by
Copybara-Service
Copybara-Service
parent
a1b0d2c933
commit
f73e1b6865
@@ -0,0 +1,170 @@
|
||||
# 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.
|
||||
"""Simulation application handle and event hooks."""
|
||||
|
||||
from typing import Callable, Protocol
|
||||
import mujoco
|
||||
from mujoco.experimental.studio import endpoints
|
||||
from mujoco.experimental.studio import messages
|
||||
from mujoco.experimental.studio import sim as _sim
|
||||
import numpy as np
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Custom event handling API.
|
||||
# Note: See viewer_app.py for hooks to customize the viewer behavior.
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
class SimEventHandler(Protocol):
|
||||
"""Invoked once per event received from the viewer."""
|
||||
|
||||
def handle(self, event: messages.Event) -> bool:
|
||||
"""Return True if event was consumed, False to allow further processing."""
|
||||
...
|
||||
|
||||
|
||||
class ViewerHandle:
|
||||
"""A handle held by the simulation to sync the simulation with the viewer."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sim_endpoint: endpoints.SimEndpoint,
|
||||
*,
|
||||
sim_event_handler: SimEventHandler | None = None,
|
||||
is_alive_fn: Callable[[], bool] | None = None,
|
||||
):
|
||||
"""Initializes the handle.
|
||||
|
||||
Args:
|
||||
sim_endpoint: The endpoint to use for communication with the viewer.
|
||||
sim_event_handler: Optional handler for viewer-to-sim events.
|
||||
is_alive_fn: Optional function called to check if the viewer is still
|
||||
alive/responsive. If not provided, the viewer is assumed to be running
|
||||
until `close()` is called.
|
||||
"""
|
||||
|
||||
self._sim_endpoint = sim_endpoint
|
||||
self._is_running = True
|
||||
self._is_alive_fn = is_alive_fn
|
||||
self._sim_event_handler = sim_event_handler
|
||||
|
||||
def close(self) -> None:
|
||||
"""Signals the viewer to exit and closes the sim endpoint, releasing resources."""
|
||||
if self._is_running:
|
||||
self._is_running = False
|
||||
try:
|
||||
self.send_to_viewer(messages.ExitEvent())
|
||||
except Exception: # pylint: disable=broad-exception-caught
|
||||
pass # Ignore exceptions, the viewer may have already closed.
|
||||
self._sim_endpoint.close()
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self.close()
|
||||
|
||||
def is_running(self) -> bool:
|
||||
"""Returns True while the viewer is open."""
|
||||
if self._is_alive_fn is not None and not self._is_alive_fn():
|
||||
self.close()
|
||||
return self._is_running
|
||||
|
||||
def send_to_viewer(self, message: messages.Message) -> None:
|
||||
self._sim_endpoint.send_to_viewer(message)
|
||||
|
||||
def sync(
|
||||
self,
|
||||
model: mujoco.MjModel | None,
|
||||
data: mujoco.MjData | None,
|
||||
step_control: _sim.StepControl,
|
||||
) -> tuple[mujoco.MjModel | None, mujoco.MjData | None, _sim.StepControl]:
|
||||
"""Syncs the simulation with the viewer and returns the updated sim state.
|
||||
|
||||
This method processes incoming events from the viewer, updates the sim state
|
||||
accordingly, and sends the current simulation state to the viewer as a
|
||||
snapshot.
|
||||
|
||||
Args:
|
||||
model: The current model.
|
||||
data: The current data.
|
||||
step_control: The current step control state.
|
||||
|
||||
Returns:
|
||||
The updated model, data, and step control state.
|
||||
"""
|
||||
|
||||
integration_sig = int(mujoco.mjtState.mjSTATE_INTEGRATION)
|
||||
|
||||
# Process incoming events from the viewer.
|
||||
for event in self._sim_endpoint.get_viewer_events():
|
||||
if (
|
||||
self._sim_event_handler is not None
|
||||
and self._sim_event_handler.handle(event)
|
||||
):
|
||||
continue
|
||||
if isinstance(event, messages.ModelEvent):
|
||||
# A new model was loaded in the viewer (e.g. via file drop).
|
||||
model = event.model
|
||||
data = mujoco.MjData(event.model)
|
||||
mujoco.mj_forward(model, data)
|
||||
step_control = _sim.StepControl()
|
||||
elif isinstance(event, messages.PerturbEvent):
|
||||
if model is not None and data is not None:
|
||||
state_size = mujoco.mj_stateSize(model, event.state_sig)
|
||||
if len(event.state) == state_size:
|
||||
mujoco.mj_setState(model, data, event.state, event.state_sig)
|
||||
elif isinstance(event, messages.ResetEvent):
|
||||
if model is not None:
|
||||
assert data is not None
|
||||
mujoco.mj_resetData(model, data)
|
||||
mujoco.mj_forward(model, data)
|
||||
elif isinstance(event, messages.ExitEvent):
|
||||
self._is_running = False
|
||||
|
||||
# Process incoming snapshots from the viewer.
|
||||
for snapshot in self._sim_endpoint.get_viewer_snapshots():
|
||||
if isinstance(snapshot, messages.StepControlSnapshot):
|
||||
step_control.set_pause_state(snapshot.pause_state)
|
||||
step_control.set_speed(snapshot.speed)
|
||||
step_control.set_noise_parameters(
|
||||
snapshot.noise_scale, snapshot.noise_rate
|
||||
)
|
||||
elif isinstance(snapshot, messages.MjOptionSnapshot):
|
||||
if model is not None:
|
||||
for field in model.opt._all_fields: # pylint: disable=protected-access
|
||||
val = getattr(snapshot.opt, field)
|
||||
try:
|
||||
getattr(model.opt, field)[:] = val
|
||||
except (TypeError, AttributeError):
|
||||
setattr(model.opt, field, val)
|
||||
|
||||
# Send the simulation state to the viewer process as a snapshot.
|
||||
if model is not None:
|
||||
assert data is not None
|
||||
integration_size = mujoco.mj_stateSize(model, integration_sig)
|
||||
integration_state = np.empty(integration_size, np.float64)
|
||||
mujoco.mj_getState(
|
||||
model,
|
||||
data,
|
||||
integration_state,
|
||||
integration_sig,
|
||||
)
|
||||
self._sim_endpoint.send_to_viewer(
|
||||
messages.StateSnapshot(
|
||||
state=integration_state, state_sig=integration_sig
|
||||
),
|
||||
)
|
||||
|
||||
return model, data, step_control
|
||||
@@ -0,0 +1,514 @@
|
||||
# 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
|
||||
from typing import Protocol
|
||||
|
||||
import mujoco
|
||||
from mujoco.experimental.studio import endpoints
|
||||
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
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Viewer application and custom message handlers
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ViewerEventHandler(Protocol):
|
||||
"""Invoked once per event received from the simulation."""
|
||||
|
||||
def handle(self, event: messages.Event) -> bool:
|
||||
"""Return True if event was consumed, False to allow further processing."""
|
||||
...
|
||||
|
||||
|
||||
class ViewerSnapshotHandler(Protocol):
|
||||
"""Invoked once per snapshot received from the simulation."""
|
||||
|
||||
def handle(self, snapshot: messages.Snapshot) -> bool:
|
||||
"""Return True if snapshot was consumed, False to allow further processing."""
|
||||
...
|
||||
|
||||
|
||||
class ViewerApp:
|
||||
"""Viewer component of Studio."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
viewer: viewer_protocol.Viewer,
|
||||
endpoint: endpoints.ViewerEndpoint,
|
||||
*,
|
||||
viewer_event_handler: ViewerEventHandler | None = None,
|
||||
viewer_snapshot_handler: ViewerSnapshotHandler | None = None,
|
||||
):
|
||||
"""Initializes the Studio application."""
|
||||
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
|
||||
self.viewer_event_handler = viewer_event_handler
|
||||
self.viewer_snapshot_handler = viewer_snapshot_handler
|
||||
|
||||
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
|
||||
|
||||
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):
|
||||
"""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:
|
||||
if (
|
||||
self.viewer_event_handler is not None
|
||||
and self.viewer_event_handler.handle(event)
|
||||
):
|
||||
continue
|
||||
if isinstance(event, messages.ModelEvent):
|
||||
# A new model was sent (initial load or hot-reload).
|
||||
# Deep-copy the incoming model to ensure the viewer operates on its own
|
||||
# isolated copy of the C++ struct. This is required if the sim/viewer
|
||||
# run in different processes, and means we don't need to lock when
|
||||
# running in the same process on different threads.
|
||||
self.model = copy.deepcopy(event.model)
|
||||
self.data = mujoco.MjData(self.model)
|
||||
|
||||
# Confirm model object is distinct from the incoming event.
|
||||
assert id(self.model) != id(event.model)
|
||||
elif isinstance(event, messages.ExitEvent):
|
||||
self.should_quit = True
|
||||
self.viewer.close()
|
||||
|
||||
# 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 (
|
||||
self.viewer_snapshot_handler is not None
|
||||
and self.viewer_snapshot_handler.handle(snapshot)
|
||||
):
|
||||
continue
|
||||
if not model_changed and isinstance(snapshot, messages.StateSnapshot):
|
||||
state_size = mujoco.mj_stateSize(self.model, snapshot.state_sig)
|
||||
if len(snapshot.state) == state_size:
|
||||
mujoco.mj_setState(
|
||||
self.model, self.data, snapshot.state, snapshot.state_sig
|
||||
)
|
||||
mujoco.mj_forward(self.model, self.data)
|
||||
|
||||
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)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# Viewer application run loop and gui/update customization hooks.
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ViewerUpdateHook(Protocol):
|
||||
"""Invoked once per viewer frame, after default updates ."""
|
||||
|
||||
def update(self, app: 'ViewerApp') -> None:
|
||||
"""Update custom state."""
|
||||
...
|
||||
|
||||
|
||||
class ViewerGuiHook(Protocol):
|
||||
"""Invoked once per viewer frame after all update() calls.
|
||||
|
||||
By deferring GUI hooks until after all updates we ensure custom UI reflects
|
||||
the latest application state.
|
||||
"""
|
||||
|
||||
def build_gui(self, app: 'ViewerApp') -> None:
|
||||
"""Draw custom UI."""
|
||||
...
|
||||
|
||||
|
||||
def run_viewer(
|
||||
viewer: viewer_protocol.Viewer,
|
||||
viewer_endpoint: endpoints.ViewerEndpoint,
|
||||
*,
|
||||
viewer_gui_hook: ViewerGuiHook | None = None,
|
||||
viewer_update_hook: ViewerUpdateHook | None = None,
|
||||
viewer_event_handler: ViewerEventHandler | None = None,
|
||||
viewer_snapshot_handler: ViewerSnapshotHandler | 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.
|
||||
viewer_gui_hook: Optional hook to draw custom ImGui panels.
|
||||
viewer_update_hook: Optional hook called once per frame before GUI.
|
||||
viewer_event_handler: Optional handler for simulation-to-viewer events.
|
||||
viewer_snapshot_handler: Optional handler for sim-to-viewer snapshots.
|
||||
"""
|
||||
app = ViewerApp(
|
||||
viewer,
|
||||
viewer_endpoint,
|
||||
viewer_event_handler=viewer_event_handler,
|
||||
viewer_snapshot_handler=viewer_snapshot_handler,
|
||||
)
|
||||
|
||||
# Viewer main loop.
|
||||
while app.is_running():
|
||||
|
||||
# Update the viewer state and handle user input.
|
||||
app.update()
|
||||
|
||||
# Update the viewer state.
|
||||
if viewer_update_hook is not None:
|
||||
viewer_update_hook.update(app)
|
||||
|
||||
# Draw the default Studio GUI.
|
||||
app.build_gui()
|
||||
|
||||
# Allow custom GUI elements to be added.
|
||||
if viewer_gui_hook is not None:
|
||||
viewer_gui_hook.build_gui(app)
|
||||
|
||||
# Sync the viewer display surface with the current model and data.
|
||||
app.viewer.sync(app.model, app.data)
|
||||
|
||||
app.close()
|
||||
Reference in New Issue
Block a user