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);