From 58d6910afaccf8d8e8aeca49c49ea415d6914663 Mon Sep 17 00:00:00 2001 From: Matija Kecman Date: Tue, 30 Jun 2026 14:20:10 -0700 Subject: [PATCH] Refactor StepControl and GUI to return change status. 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 --- python/mujoco/experimental/studio/messages.py | 8 ++- .../experimental/studio/sample/implot.py | 6 +- python/mujoco/experimental/studio/sim.cc | 15 +++- .../mujoco/experimental/studio/studio_app.py | 72 ++++++------------- .../experimental/studio/studio_app_events.py | 38 ++++++---- python/mujoco/experimental/studio/ux.cc | 43 +++++------ src/experimental/platform/sim/step_control.cc | 5 +- src/experimental/platform/ux/gui.cc | 9 +-- src/experimental/platform/ux/gui.h | 6 +- src/experimental/studio/app.cc | 8 +-- 10 files changed, 100 insertions(+), 110 deletions(-) diff --git a/python/mujoco/experimental/studio/messages.py b/python/mujoco/experimental/studio/messages.py index 655fbb85..d36de578 100644 --- a/python/mujoco/experimental/studio/messages.py +++ b/python/mujoco/experimental/studio/messages.py @@ -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) diff --git a/python/mujoco/experimental/studio/sample/implot.py b/python/mujoco/experimental/studio/sample/implot.py index df0a59b2..8188ce00 100644 --- a/python/mujoco/experimental/studio/sample/implot.py +++ b/python/mujoco/experimental/studio/sample/implot.py @@ -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. """ diff --git a/python/mujoco/experimental/studio/sim.cc b/python/mujoco/experimental/studio/sim.cc index dbc4b363..810461a5 100644 --- a/python/mujoco/experimental/studio/sim.cc +++ b/python/mujoco/experimental/studio/sim.cc @@ -14,6 +14,8 @@ // Python bindings for MuJoCo platform simulation components. +#include + #include #include #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."); } diff --git a/python/mujoco/experimental/studio/studio_app.py b/python/mujoco/experimental/studio/studio_app.py index 389480fe..042f4c6e 100644 --- a/python/mujoco/experimental/studio/studio_app.py +++ b/python/mujoco/experimental/studio/studio_app.py @@ -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) diff --git a/python/mujoco/experimental/studio/studio_app_events.py b/python/mujoco/experimental/studio/studio_app_events.py index f3265f89..a60924e3 100644 --- a/python/mujoco/experimental/studio/studio_app_events.py +++ b/python/mujoco/experimental/studio/studio_app_events.py @@ -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): diff --git a/python/mujoco/experimental/studio/ux.cc b/python/mujoco/experimental/studio/ux.cc index 2f5a1dc1..6e5d5a88 100644 --- a/python/mujoco/experimental/studio/ux.cc +++ b/python/mujoco/experimental/studio/ux.cc @@ -14,7 +14,6 @@ // Python bindings for MuJoCo platform UX components. -#include #include #include #include @@ -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", diff --git a/src/experimental/platform/sim/step_control.cc b/src/experimental/platform/sim/step_control.cc index 7f0122da..8492845c 100644 --- a/src/experimental/platform/sim/step_control.cc +++ b/src/experimental/platform/sim/step_control.cc @@ -66,8 +66,11 @@ float StepControl::GetSpeedMeasured() const { return speed_measured_; } float StepControl::GetSpeed() const { return speed_; } void StepControl::SetSpeed(float speed_percent_real_time) { + float prev_speed = speed_; speed_ = std::clamp(speed_percent_real_time, .1f, 100.f); - ForceSync(); + if (speed_ != prev_speed) { + ForceSync(); + } } void StepControl::ForceSync() { force_sync_ = true; } diff --git a/src/experimental/platform/ux/gui.cc b/src/experimental/platform/ux/gui.cc index 3682cdfa..794569ab 100644 --- a/src/experimental/platform/ux/gui.cc +++ b/src/experimental/platform/ux/gui.cc @@ -482,8 +482,7 @@ ImVec4 ConfigureDockingLayout(bool show_toolbar, bool show_status_bar) { return ImVec4(workspace_x, workspace_y, workspace_w, workspace_h); } -void StepControlGui(const mjModel* model, StepControl* step_control, - int& speed_index) { +void StepControlGui(StepControl* step_control, int& speed_index) { platform::ScopedStyle style; bool is_dark = ImGui::GetStyle().Colors[ImGuiCol_WindowBg].x < 0.5f; @@ -1215,13 +1214,15 @@ void GroupsGui(const mjModel* model, mjvOption* vis_options, float min_width) { GroupGui("Skins", vis_options->skingroup); } -void NoiseGui(const mjModel* model, const mjData* data, float& noise_scale, - float& noise_rate) { +void NoiseGui(StepControl* step_control) { + float noise_scale, noise_rate; + step_control->GetNoiseParameters(noise_scale, noise_rate); const float item_width = ImGui::GetWindowWidth() * .6f; ImGui::PushItemWidth(item_width); ImGui::SliderFloat("Noise scale", &noise_scale, 0, 1); ImGui::SliderFloat("Noise rate", &noise_rate, 0, 4); ImGui::PopItemWidth(); + step_control->SetNoiseParameters(noise_scale, noise_rate); } void JointsGui(const mjModel* model, const mjData* data, diff --git a/src/experimental/platform/ux/gui.h b/src/experimental/platform/ux/gui.h index cbb16a6a..15096081 100644 --- a/src/experimental/platform/ux/gui.h +++ b/src/experimental/platform/ux/gui.h @@ -83,8 +83,7 @@ static constexpr std::array kPercentRealTime = { // UX for controlling the simulation stepping. `speed_index` is an index into // kPercentRealTime, an array of available speeds (indices in range [0, 30] map // to real-time percentages in range [100%, 0.1%]). -void StepControlGui(const mjModel* model, StepControl* step_control, - int& speed_index); +void StepControlGui(StepControl* step_control, int& speed_index); // Sets the simulation speed index and updates the StepControl object. void SetSpeedIndex(StepControl* step_control, int& speed_index, @@ -148,8 +147,7 @@ void WatchGui(const mjModel* model, const mjData* data, char* field_name, // UX for controlling noise parameters which can then be applied to the // simulation via StepControl::SetNoiseParameters / StepControl::InjectNoise. -void NoiseGui(const mjModel* model, const mjData* data, float& noise_scale, - float& noise_rate); +void NoiseGui(StepControl* step_control); // UX for the solver convergence chart. void ConvergenceGui(const mjModel* model, mjData* data, diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc index ec9e6379..e886b31f 100644 --- a/src/experimental/studio/app.cc +++ b/src/experimental/studio/app.cc @@ -1509,11 +1509,7 @@ void App::DataInspectorGui() { ImGui::BeginChild("ControlsGui", {0, 0}, child_flags); if (platform::SectionHeader("Controls", node_flags, 0.65f)) { - float noise_scale = 0; - float noise_rate = 0; - step_control_.GetNoiseParameters(noise_scale, noise_rate); - platform::NoiseGui(model(), data(), noise_scale, noise_rate); - step_control_.SetNoiseParameters(noise_scale, noise_rate); + platform::NoiseGui(&step_control_); ImGui::Separator(); platform::ControlsGui(model(), data(), &vis_options_); @@ -1889,7 +1885,7 @@ void App::ToolBarGui() { // Combined (Normal Pause, Viscous Pause, Play) widget and Speed selection. ImGui::SameLine(0, separator_width); - platform::StepControlGui(model(), &step_control_, tmp_.speed_index); + platform::StepControlGui(&step_control_, tmp_.speed_index); ImGui::SameLine(0, separator_width); ImGui::SetNextItemWidth(120);