Yuval Tassa
c232938233
PiperOrigin-RevId: 965914672 Change-Id: Ib349b77055c57453780818dd4d8f48899b7024d5
490 lines
16 KiB
Python
490 lines
16 KiB
Python
# Copyright 2026 DeepMind Technologies Limited
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# https://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""Viewer-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._setup_history()
|
|
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.sim_history = sim.SimHistory()
|
|
self._setup_history()
|
|
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)
|
|
self._setup_history()
|
|
|
|
def _setup_history(self) -> None:
|
|
"""(Re)initialize simulation history recording for the current model."""
|
|
ux.setup_history(
|
|
self.step_control, self.sim_history, self.ux_state, self.model,
|
|
self.data,
|
|
)
|
|
|
|
def reload_model(self) -> None:
|
|
"""Reload the current model from its file."""
|
|
if self.model_path:
|
|
self.load_model_from_file(self.model_path)
|
|
|
|
def align_camera(self, camera: mujoco.MjvCamera) -> None:
|
|
"""Recenter the camera on the model's home camera, else the free camera."""
|
|
cam_id = self.model.vis.global_.cameraid
|
|
if 0 <= cam_id < self.model.ncam:
|
|
self.ux_state.camera_index = ux.set_camera(self.model, camera, cam_id)
|
|
else:
|
|
mujoco.mjv_defaultFreeCamera(self.model, camera)
|
|
|
|
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.
|
|
"""
|
|
if self.ux_state.update_threadpool:
|
|
mujoco.mju_threadpool(self.data, self.ux_state.nthread)
|
|
self.ux_state.update_threadpool = False
|
|
|
|
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')
|
|
# The Simulation panel draws its own collapsible section header.
|
|
ux.simulation_gui(
|
|
self.model,
|
|
self.data,
|
|
self.step_control,
|
|
self.sim_history,
|
|
self.ux_state,
|
|
self.reset_physics,
|
|
self.reload_model,
|
|
lambda: self.align_camera(camera),
|
|
)
|
|
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)
|