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

@@ -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):
+class StepControlEvent(Event):
-  """An event requesting to set the pause state."""
+  """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

@@ -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
+This script runs a Studio viewer and adds an 'Inspect Body' window using ImGui
-using ImGui and ImPlot bindings to visualize selected body data. The example
+and ImPlot bindings to visualize selected body data. The example demonstrates
-demonstrates how responsive UI layout rules are easily implemented.
+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,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

@@ -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 studio_app_events.handle_vis_options_keyboard_events(
-      return False
-
-    return 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 studio_app_events.handle_step_control_keyboard_events(
-      return False
+        self.step_control, self.ux_state
-
-    return events.handle_step_control_keyboard_events(
-        self.model, self.data, 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:
+    handled, self._cam_speed = (
-      return False
+        studio_app_events.handle_freecam_wasd_keyboard_events(
-
+            self.model, self.data, camera, self._cam_speed
-    handled, self._cam_speed = 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:
+    handled, self._cam_speed = studio_app_events.handle_keyboard_events(
-      return False
+        self.model,
-
+        self.data,
-    is_freecam_wasd = self.ux_state.camera_index == ux.FREE_CAMERA_IDX
+        camera,
-    if events.handle_step_control_keyboard_events(
+        vis_options,
-        self.model, self.data, self.step_control, self.ux_state
+        self.step_control,
-    ):
+        self.ux_state,
-      return True
+        self._cam_speed,
-
+    )
-    if events.handle_camera_select_keyboard_events(
+    return handled
-        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
 
   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:
+    studio_app_events.handle_camera_tracking_mouse_events(
-      return
-
-    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:
+    studio_app_events.handle_mouse_events(
-      return
-
-    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

@@ -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
+      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 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

@@ -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,
+        mujoco::platform::StepControlGui(step_control, ux_state.speed_index);
-                                         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) {
+      [](mujoco::platform::StepControl* step_control, int speed_index,
-        py::gil_scoped_release no_gil;
+         int request_idx) {
-        mujoco::platform::SetSpeedIndex(step_control, ux_state.speed_index, 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"),
+      py::arg("step_control"), py::arg("speed_index"), py::arg("request_idx"),
-      "Set the simulation speed index.");
+      "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::platform::StepControl* step_control) {
-         mujoco::python::MjDataWrapper& data, UxState& ux_state) {
         py::gil_scoped_release no_gil;
-        mujoco::platform::NoiseGui(model.get(), data.get(),
+        mujoco::platform::NoiseGui(step_control);
-                                   ux_state.noise_scale, ux_state.noise_rate);
       },
-      py::arg("model"), py::arg("data"), py::arg("ux_state"),
+      py::arg("step_control"), "Render the noise UI.");
-      "Render the noise UI. Modifies ux_state.noise_scale and "
-      "ux_state.noise_rate.");
 
   m.def(
       "convergence_gui",

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

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,
+void StepControlGui(StepControl* step_control, int& speed_index) {
-                    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,
+void NoiseGui(StepControl* step_control) {
-              float& noise_rate) {
+  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,

src/experimental/platform/ux/gui.h

@@ -83,8 +83,7 @@ static constexpr std::array<const char*, 31> 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,
+void StepControlGui(StepControl* step_control, int& speed_index);
-                    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,
+void NoiseGui(StepControl* step_control);
-              float& noise_rate);
 
 // UX for the solver convergence chart.
 void ConvergenceGui(const mjModel* model, mjData* data,

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;
+    platform::NoiseGui(&step_control_);
-    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);
     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);