Matija Kecman
58d6910afa
StepControl methods SetSpeed, SetNoiseParameters, and SetPauseState now return a boolean indicating whether the state was actually changed. The GUI functions NoiseGui and StepControlGui also return a boolean reflecting if any parameters were modified. Noise parameters are now directly managed by StepControl rather than being stored in UxState. The StepControlEvent message has been expanded to include all step control parameters. Keyboard event handling has been slightly refactored to better align with event-based updates. PiperOrigin-RevId: 940656342 Change-Id: I3701623a2fdcfc4c084b5db0ae19048669040de4
451 lines
15 KiB
Python
451 lines
15 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-agnostic Python implementation of Studio.
|
|
|
|
Architecture:
|
|
StudioApp owns the simulation state (model, data) and the UI logic.
|
|
Viewer classes (e.g., NativeViewer) own the window (if required), renderer,
|
|
camera, and visualization options. The viewer never stores references to
|
|
model or data. Instead, the caller passes them each frame via
|
|
viewer.sync(model, data). This ensures the viewer always renders the current
|
|
model, even if StudioApp.load_model_from_file() swaps it.
|
|
|
|
The class can be used to implement the full Studio application in Python. By
|
|
using the more granular member functions it can also build simple apps that only
|
|
use a subset of the Studio UI. This configuration is fully dynamic, there is
|
|
nothing to configure in advance, you can change your app by changing the
|
|
functions that get called each frame. This class is also viewer-agnostic and as
|
|
such does not own camera, vis_options or perturb objects (these are provided by
|
|
the viewer).
|
|
|
|
See the sample/ folder for usage examples.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import typing
|
|
|
|
import mujoco
|
|
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
|
|
|
|
# Type alias for a custom physics step function.
|
|
StepFn = typing.Callable[[mujoco.MjModel, mujoco.MjData], None]
|
|
|
|
|
|
def load_model_from_file(
|
|
model_path: str,
|
|
) -> tuple[mujoco.MjModel, mujoco.MjData] | None:
|
|
"""Loads a model and data from a file path."""
|
|
try:
|
|
data = parser.parse(model_path)
|
|
return data.model, data
|
|
except Exception as ex: # pylint: disable=broad-except
|
|
print(f'Error loading model from {model_path!r}: {ex}')
|
|
return None
|
|
|
|
|
|
class StudioApp:
|
|
"""Viewer-agnostic Python implementation of Studio."""
|
|
|
|
@classmethod
|
|
def from_argv(cls, argv: list[str]) -> 'StudioApp':
|
|
"""Constructs a StudioApp by parsing a model path from command-line args."""
|
|
if len(argv) < 2:
|
|
model = mujoco.MjSpec().compile()
|
|
data = mujoco.MjData(model)
|
|
app = cls(model, data)
|
|
app.step_control.set_pause_state(sim.PauseState.NORMAL_PAUSED)
|
|
return app
|
|
|
|
model_path = argv[1]
|
|
|
|
res = load_model_from_file(model_path)
|
|
if res is None:
|
|
sys.exit(-1)
|
|
model, data = res
|
|
|
|
app = cls(model, data)
|
|
app.model_path = model_path
|
|
return app
|
|
|
|
def load_model_from_file(
|
|
self, model_path: str
|
|
) -> tuple[mujoco.MjModel, mujoco.MjData] | None:
|
|
"""Loads a new model from a file, replacing the current model and data."""
|
|
res = load_model_from_file(model_path)
|
|
if res is None:
|
|
self.status = f'Error loading model from {model_path!r}'
|
|
return None
|
|
|
|
model, data = res
|
|
self.model = model
|
|
self.data = data
|
|
self.model_path = model_path
|
|
self.step_control = sim.StepControl()
|
|
self.ux_state = ux.UxState()
|
|
self.status = f'Loaded: {os.path.basename(model_path)!r}'
|
|
|
|
return model, data
|
|
|
|
def __init__(
|
|
self,
|
|
model: mujoco.MjModel,
|
|
data: mujoco.MjData,
|
|
):
|
|
"""Initializes the Studio application."""
|
|
self.model = model
|
|
self.data = data
|
|
self.model_path = ''
|
|
|
|
self.step_control = sim.StepControl()
|
|
self.ux_state = ux.UxState()
|
|
self.theme = ux.GuiTheme.LIGHT
|
|
self.show_info = 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 handle_vis_options_keyboard_events(
|
|
self,
|
|
vis_options: mujoco.MjvOption,
|
|
is_freecam_wasd: bool,
|
|
) -> bool:
|
|
"""Toggles visualization flags based on keyboard shortcuts.
|
|
|
|
Args:
|
|
vis_options: The visualization options to modify.
|
|
is_freecam_wasd: If True, keys Q/E/A/D are reserved for camera movement
|
|
and will not toggle visualization flags.
|
|
|
|
Returns:
|
|
True if a key was handled, False otherwise.
|
|
"""
|
|
return studio_app_events.handle_vis_options_keyboard_events(
|
|
vis_options, is_freecam_wasd
|
|
)
|
|
|
|
def handle_step_control_keyboard_events(self) -> bool:
|
|
"""Handles keyboard shortcuts for simulation stepping control.
|
|
|
|
Returns:
|
|
True if a key was handled, False otherwise.
|
|
"""
|
|
return studio_app_events.handle_step_control_keyboard_events(
|
|
self.step_control, self.ux_state
|
|
)
|
|
|
|
def handle_freecam_wasd_keyboard_events(
|
|
self,
|
|
camera: mujoco.MjvCamera,
|
|
) -> bool:
|
|
"""Handles keyboard shortcuts for free camera movement."""
|
|
handled, self._cam_speed = (
|
|
studio_app_events.handle_freecam_wasd_keyboard_events(
|
|
self.model, self.data, camera, self._cam_speed
|
|
)
|
|
)
|
|
return handled
|
|
|
|
def handle_keyboard_events(
|
|
self,
|
|
camera: mujoco.MjvCamera,
|
|
vis_options: mujoco.MjvOption,
|
|
) -> bool:
|
|
"""Handle keyboard events according to Studio's bindings."""
|
|
handled, self._cam_speed = studio_app_events.handle_keyboard_events(
|
|
self.model,
|
|
self.data,
|
|
camera,
|
|
vis_options,
|
|
self.step_control,
|
|
self.ux_state,
|
|
self._cam_speed,
|
|
)
|
|
return handled
|
|
|
|
def handle_camera_tracking_mouse_events(
|
|
self,
|
|
camera: mujoco.MjvCamera,
|
|
vis_options: mujoco.MjvOption,
|
|
) -> None:
|
|
"""Handles mouse events for camera tracking."""
|
|
studio_app_events.handle_camera_tracking_mouse_events(
|
|
self.model, self.data, camera, vis_options, self.ux_state
|
|
)
|
|
|
|
def handle_mouse_events(
|
|
self,
|
|
camera: mujoco.MjvCamera,
|
|
vis_options: mujoco.MjvOption,
|
|
perturb: mujoco.MjvPerturb,
|
|
) -> None:
|
|
"""Handles mouse events."""
|
|
studio_app_events.handle_mouse_events(
|
|
self.model, self.data, camera, vis_options, 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)
|
|
|
|
def apply_perturb(self, perturb: mujoco.MjvPerturb) -> None:
|
|
"""Apply perturbation the model."""
|
|
if self.step_control.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 update_physics(
|
|
self,
|
|
perturb: mujoco.MjvPerturb,
|
|
*,
|
|
step_fn: StepFn | None = None,
|
|
) -> None:
|
|
"""Applies the perturbations and advances the physics.
|
|
|
|
Args:
|
|
perturb: The MuJoCo perturbation object.
|
|
step_fn: Optional custom physics step function. When provided, it is
|
|
called instead of ``step_control.advance``. The function receives
|
|
``(model, data)`` and should step the simulation in-place.
|
|
"""
|
|
self.apply_perturb(perturb)
|
|
|
|
if step_fn is not None:
|
|
step_fn(self.model, self.data)
|
|
else:
|
|
advance_status = self.step_control.advance(self.model, self.data)
|
|
if advance_status == sim.StepStatus.AUTO_RESET:
|
|
self.reset_physics()
|
|
|
|
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 (called by update())."""
|
|
return not self.should_quit
|
|
|
|
def update(
|
|
self,
|
|
camera: mujoco.MjvCamera,
|
|
vis_options: mujoco.MjvOption,
|
|
perturb: mujoco.MjvPerturb,
|
|
*,
|
|
drop_file: str = '',
|
|
step_fn: StepFn | None = None,
|
|
) -> bool:
|
|
"""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.
|
|
The argument objects are provided by the viewer.
|
|
|
|
Args:
|
|
camera: The MuJoCo camera object.
|
|
vis_options: The MuJoCo visualization options.
|
|
perturb: The MuJoCo perturbation object.
|
|
drop_file: Path of a file dropped into the viewer window. If non-empty the
|
|
current model is replaced with the dropped file.
|
|
step_fn: Optional custom physics step function. When provided, it is
|
|
called instead of ``step_control.advance``.
|
|
|
|
Returns:
|
|
Whether the application should continue running, this is a
|
|
convenience to allow this function to be used in a while loop.
|
|
"""
|
|
if drop_file:
|
|
self.load_model_from_file(drop_file)
|
|
|
|
self.handle_mouse_events(camera, vis_options, perturb)
|
|
self.handle_keyboard_events(camera, vis_options)
|
|
self.update_physics(perturb, step_fn=step_fn)
|
|
return self.is_running()
|
|
|
|
def update_from_viewer(
|
|
self,
|
|
viewer: viewer_protocol.Viewer,
|
|
*,
|
|
step_fn: StepFn | None = None,
|
|
) -> bool:
|
|
"""Convenience wrapper around update() that unpacks viewer attributes.
|
|
|
|
Args:
|
|
viewer: A viewer conforming to the Viewer protocol.
|
|
step_fn: Optional custom physics step function. When provided, it is
|
|
called instead of ``step_control.advance``.
|
|
|
|
Returns:
|
|
Whether the application should continue running.
|
|
"""
|
|
return self.update(
|
|
viewer.camera,
|
|
viewer.vis_options,
|
|
viewer.perturb,
|
|
drop_file=viewer.get_drop_file(),
|
|
step_fn=step_fn,
|
|
)
|
|
|
|
def build_gui(
|
|
self,
|
|
camera: mujoco.MjvCamera,
|
|
vis_options: mujoco.MjvOption,
|
|
render_flags: ux.RenderFlags,
|
|
) -> 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.should_quit = True
|
|
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('Info', '', self.show_info):
|
|
self.show_info = not self.show_info
|
|
imgui.EndMenu()
|
|
if imgui.BeginMenu('Help'):
|
|
if imgui.MenuItem('Info', '', self.show_info):
|
|
self.show_info = not self.show_info
|
|
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, self.ux_state)
|
|
|
|
imgui.TableNextColumn()
|
|
ux.camera_selection_gui(self.model, self.data, camera, self.ux_state)
|
|
|
|
imgui.SameLine()
|
|
ux.label_selection_gui(vis_options)
|
|
|
|
imgui.SameLine()
|
|
ux.frame_selection_gui(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, vis_options, render_flags)
|
|
imgui.TreePop()
|
|
if imgui.TreeNodeEx('Visibility Groups', node_flags):
|
|
ux.groups_gui(self.model, vis_options)
|
|
imgui.TreePop()
|
|
if imgui.TreeNodeEx('Visualization', node_flags):
|
|
ux.visualization_gui(self.model, vis_options, camera)
|
|
imgui.TreePop()
|
|
imgui.End()
|
|
|
|
# -- Right pane: Inspector ------------------------------------------------
|
|
imgui.Begin('Inspector')
|
|
if imgui.TreeNodeEx('Noise', node_flags):
|
|
ux.noise_gui(self.step_control)
|
|
imgui.TreePop()
|
|
if imgui.TreeNodeEx('Joints', node_flags):
|
|
ux.joints_gui(self.model, self.data, vis_options)
|
|
imgui.TreePop()
|
|
if imgui.TreeNodeEx('Controls', node_flags):
|
|
ux.controls_gui(self.model, self.data, 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_info:
|
|
_, self.show_info = imgui.Begin('Info', self.show_info)
|
|
paused = self.step_control.get_pause_state() != sim.PauseState.UNPAUSED
|
|
ux.info_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)
|