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

Refactor StepControl and GUI to return change status.

Browse Source 
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
This commit is contained in:
Matija Kecman
2026-06-30 14:20:10 -07:00
committed by Copybara-Service
parent f4b2f76def
commit 58d6910afa
10 changed files with 100 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/experimental/studio/messages.py
+5 -3
View File
@@ -13,7 +13,6 @@
# limitations under the License.
"""Messages and Channels for Studio."""
import dataclasses
from typing import Protocol
from typing import runtime_checkable
@@ -129,10 +128,13 @@ class PerturbEvent(Event):
@dataclasses.dataclass(frozen=True)
class SetPauseStateEvent(Event):
"""An event requesting to set the pause state."""
class StepControlEvent(Event):
"""An event carrying the full step control state from the viewer to the sim."""
pause_state: sim.PauseState
speed: float
noise_scale: float
noise_rate: float
@dataclasses.dataclass(frozen=True)
python/mujoco/experimental/studio/sample/implot.py
+3 -3
View File
@@ -13,9 +13,9 @@
# limitations under the License.
"""Example to run studio in the native viewer with responsive ImPlot UI.
This script runs a Studio viewer in-process and adds an 'Inspect Body' window
using ImGui and ImPlot bindings to visualize selected body data. The example
demonstrates how responsive UI layout rules are easily implemented.
This script runs a Studio viewer and adds an 'Inspect Body' window using ImGui
and ImPlot bindings to visualize selected body data. The example demonstrates
how responsive UI layout rules are easily implemented.
Provide an MJCF model file via the first command-line argument to launch.
"""
python/mujoco/experimental/studio/sim.cc
+14 -1
View File
@@ -14,6 +14,8 @@
// Python bindings for MuJoCo platform simulation components.
#include <tuple>
#include <mujoco/mujoco.h>
#include <mujoco/experimental/platform/sim/step_control.h>
#include "structs.h"
@@ -70,5 +72,16 @@ PYBIND11_MODULE(sim, m) {
.def("get_pause_state", &StepControl::GetPauseState,
"Returns the current pause state.")
.def("request_single_step", &StepControl::RequestSingleStep,
"Request a single step if paused.");
"Request a single step if paused.")
.def(
"get_noise_parameters",
[](const StepControl& self) {
float noise_scale, noise_rate;
self.GetNoiseParameters(noise_scale, noise_rate);
return std::make_tuple(noise_scale, noise_rate);
},
"Returns the noise parameters.")
.def("set_noise_parameters", &StepControl::SetNoiseParameters,
py::arg("noise_scale"), py::arg("noise_rate"),
"Sets the noise parameters.");
}
python/mujoco/experimental/studio/studio_app.py
+22 -50
View File
@@ -39,7 +39,7 @@ import typing
import mujoco
from mujoco.experimental.studio import parser
from mujoco.experimental.studio import sim
from mujoco.experimental.studio import studio_app_events as events
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
@@ -141,10 +141,7 @@ class StudioApp:
Returns:
True if a key was handled, False otherwise.
"""
if imgui.GetIO().WantCaptureKeyboard:
return False
return events.handle_vis_options_keyboard_events(
return studio_app_events.handle_vis_options_keyboard_events(
vis_options, is_freecam_wasd
)
@@ -154,11 +151,8 @@ class StudioApp:
Returns:
True if a key was handled, False otherwise.
"""
if imgui.GetIO().WantCaptureKeyboard:
return False
return events.handle_step_control_keyboard_events(
self.model, self.data, self.step_control, self.ux_state
return studio_app_events.handle_step_control_keyboard_events(
self.step_control, self.ux_state
)
def handle_freecam_wasd_keyboard_events(
@@ -166,11 +160,10 @@ class StudioApp:
camera: mujoco.MjvCamera,
) -> bool:
"""Handles keyboard shortcuts for free camera movement."""
if imgui.GetIO().WantCaptureKeyboard:
return False
handled, self._cam_speed = events.handle_freecam_wasd_keyboard_events(
self.model, self.data, camera, self._cam_speed
handled, self._cam_speed = (
studio_app_events.handle_freecam_wasd_keyboard_events(
self.model, self.data, camera, self._cam_speed
)
)
return handled
@@ -180,31 +173,16 @@ class StudioApp:
vis_options: mujoco.MjvOption,
) -> bool:
"""Handle keyboard events according to Studio's bindings."""
if imgui.GetIO().WantCaptureKeyboard:
return False
is_freecam_wasd = self.ux_state.camera_index == ux.FREE_CAMERA_IDX
if events.handle_step_control_keyboard_events(
self.model, self.data, self.step_control, self.ux_state
):
return True
if events.handle_camera_select_keyboard_events(
self.model, camera, self.ux_state
):
return True
if events.handle_vis_options_keyboard_events(vis_options, is_freecam_wasd):
return True
if is_freecam_wasd:
handled, self._cam_speed = events.handle_freecam_wasd_keyboard_events(
self.model, self.data, camera, self._cam_speed
)
if handled:
return True
return False
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,
@@ -212,10 +190,7 @@ class StudioApp:
vis_options: mujoco.MjvOption,
) -> None:
"""Handles mouse events for camera tracking."""
if imgui.GetIO().WantCaptureMouse:
return
events.handle_camera_tracking_mouse_events(
studio_app_events.handle_camera_tracking_mouse_events(
self.model, self.data, camera, vis_options, self.ux_state
)
@@ -226,10 +201,7 @@ class StudioApp:
perturb: mujoco.MjvPerturb,
) -> None:
"""Handles mouse events."""
if imgui.GetIO().WantCaptureMouse:
return
events.handle_mouse_events(
studio_app_events.handle_mouse_events(
self.model, self.data, camera, vis_options, perturb, self.ux_state
)
@@ -391,7 +363,7 @@ class StudioApp:
self.reset_physics_gui()
imgui.SameLine()
ux.step_control_gui(self.model, self.step_control, self.ux_state)
ux.step_control_gui(self.step_control, self.ux_state)
imgui.TableNextColumn()
ux.camera_selection_gui(self.model, self.data, camera, self.ux_state)
@@ -434,7 +406,7 @@ class StudioApp:
# -- Right pane: Inspector ------------------------------------------------
imgui.Begin('Inspector')
if imgui.TreeNodeEx('Noise', node_flags):
ux.noise_gui(self.model, self.data, self.ux_state)
ux.noise_gui(self.step_control)
imgui.TreePop()
if imgui.TreeNodeEx('Joints', node_flags):
ux.joints_gui(self.model, self.data, vis_options)
python/mujoco/experimental/studio/studio_app_events.py
+25 -13
View File
@@ -140,16 +140,12 @@ def handle_vis_options_keyboard_events(
def handle_step_control_keyboard_events(
model: mujoco.MjModel,
data: mujoco.MjData,
step_control: sim.StepControl,
ux_state: ux.UxState,
) -> bool:
"""Handles keyboard shortcuts for simulation stepping control.
Args:
model: The MuJoCo model.
data: The MuJoCo data.
step_control: The simulation step control object.
ux_state: The UX state object.
@@ -174,17 +170,30 @@ def handle_step_control_keyboard_events(
else:
step_control.set_pause_state(sim.PauseState.UNPAUSED)
return True
elif pressed(imgui.Key.Backspace):
elif pressed(imgui.Key.Minus):
ux_state.speed_index = ux.set_speed_index(
step_control, ux_state.speed_index, ux_state.speed_index + 1
)
return True
elif pressed(imgui.Key.Equal):
ux_state.speed_index = ux.set_speed_index(
step_control, ux_state.speed_index, ux_state.speed_index - 1
)
return True
return False
def handle_reset_keyboard_events(
model: mujoco.MjModel, data: mujoco.MjData
) -> bool:
"""Handles keyboard shortcuts for simulation reset."""
if imgui.IsKeyChordPressed(imgui.Key.Backspace):
if model is not None and data is not None:
mujoco.mj_resetData(model, data)
mujoco.mj_forward(model, data)
return True
elif pressed(imgui.Key.Minus):
ux.set_speed_index(step_control, ux_state, ux_state.speed_index + 1)
return True
elif pressed(imgui.Key.Equal):
ux.set_speed_index(step_control, ux_state, ux_state.speed_index - 1)
return True
return True
return False
@@ -362,7 +371,10 @@ def handle_keyboard_events(
return False, cam_speed
is_freecam_wasd = ux_state.camera_index == ux.FREE_CAMERA_IDX
if handle_step_control_keyboard_events(model, data, step_control, ux_state):
if handle_step_control_keyboard_events(step_control, ux_state):
return True, cam_speed
if handle_reset_keyboard_events(model, data):
return True, cam_speed
if handle_camera_select_keyboard_events(model, camera, ux_state):
python/mujoco/experimental/studio/ux.cc
+18 -25
View File
@@ -14,7 +14,6 @@
// Python bindings for MuJoCo platform UX components.
#include <algorithm>
#include <array>
#include <string>
#include <tuple>
@@ -44,10 +43,6 @@ struct UxState {
char watch_field_name[256] = {0};
int watch_field_index = 0;
// Read/edited by noise_gui
float noise_scale = 0.0f;
float noise_rate = 0.0f;
// Read/edited by camera_selection_gui
int camera_index = mujoco::platform::kTumbleCameraIdx;
};
@@ -75,8 +70,6 @@ PYBIND11_MODULE(ux, m) {
.def_readwrite("state", &UxState::state)
.def_readwrite("state_sig", &UxState::state_sig)
.def_readwrite("watch_field_index", &UxState::watch_field_index)
.def_readwrite("noise_scale", &UxState::noise_scale)
.def_readwrite("noise_rate", &UxState::noise_rate)
.def_readwrite("camera_index", &UxState::camera_index)
.def_property(
"watch_field_name",
@@ -115,13 +108,11 @@ PYBIND11_MODULE(ux, m) {
m.def(
"step_control_gui",
[](const mujoco::python::MjModelWrapper& model,
mujoco::platform::StepControl* step_control, UxState& ux_state) {
[](mujoco::platform::StepControl* step_control, UxState& ux_state) {
py::gil_scoped_release no_gil;
mujoco::platform::StepControlGui(model.get(), step_control,
ux_state.speed_index);
mujoco::platform::StepControlGui(step_control, ux_state.speed_index);
},
py::arg("model"), py::arg("step_control"), py::arg("ux_state"),
py::arg("step_control"), py::arg("ux_state"),
"Render the simulation stepping control GUI. Modifies "
"ux_state.speed_index.");
@@ -170,19 +161,25 @@ PYBIND11_MODULE(ux, m) {
[](const mujoco::python::MjModelWrapper& model,
mujoco::python::MjvCameraWrapper& camera, int request_idx) {
py::gil_scoped_release no_gil;
return mujoco::platform::SetCamera(model.get(), camera.get(), request_idx);
return mujoco::platform::SetCamera(model.get(), camera.get(),
request_idx);
},
py::arg("model"), py::arg("camera"), py::arg("request_idx"),
"Set the camera index and update the camera object.");
m.def(
"set_speed_index",
[](mujoco::platform::StepControl* step_control, UxState& ux_state, int idx) {
py::gil_scoped_release no_gil;
mujoco::platform::SetSpeedIndex(step_control, ux_state.speed_index, idx);
[](mujoco::platform::StepControl* step_control, int speed_index,
int request_idx) {
{
py::gil_scoped_release no_gil;
mujoco::platform::SetSpeedIndex(step_control, speed_index,
request_idx);
}
return speed_index;
},
py::arg("step_control"), py::arg("ux_state"), py::arg("idx"),
"Set the simulation speed index.");
py::arg("step_control"), py::arg("speed_index"), py::arg("request_idx"),
"Set the simulation speed index. Returns new_speed_index.");
m.def(
"physics_gui",
@@ -294,15 +291,11 @@ PYBIND11_MODULE(ux, m) {
m.def(
"noise_gui",
[](const mujoco::python::MjModelWrapper& model,
mujoco::python::MjDataWrapper& data, UxState& ux_state) {
[](mujoco::platform::StepControl* step_control) {
py::gil_scoped_release no_gil;
mujoco::platform::NoiseGui(model.get(), data.get(),
ux_state.noise_scale, ux_state.noise_rate);
mujoco::platform::NoiseGui(step_control);
},
py::arg("model"), py::arg("data"), py::arg("ux_state"),
"Render the noise UI. Modifies ux_state.noise_scale and "
"ux_state.noise_rate.");
py::arg("step_control"), "Render the noise UI.");
m.def(
"convergence_gui",