Make StepControlGui a reusable component and refactor/simplify the implementation

PiperOrigin-RevId: 900045404
Change-Id: I9976ece042710ddce0a7c5d99a6ac85a0ef06d2f

src/experimental/platform/sim/step_control.cc

@@ -83,10 +83,14 @@ void StepControl::SetNoiseParameters(float ctrl_noise_scale,
   ctrl_noise_rate_ = ctrl_noise_rate;
 }
 
-void StepControl::SetPauseState(PauseState state, mjModel* m) {
+void StepControl::SetPauseState(PauseState state) {
   pause_state_ = state;
 }
 
+StepControl::PauseState StepControl::GetPauseState() const {
+  return pause_state_;
+}
+
 StepControl::Status StepControl::Advance(mjModel* m, mjData* d) {
   if (!m) {
     return Status::kOk;
@@ -184,7 +188,7 @@ StepControl::Status StepControl::Advance(mjModel* m, mjData* d) {
       for (mjtWarning w : kDivergedWarnings) {
         if (d->warning[w].number > 0) {
           // Stop stepping if the simulation diverged.
-          pause_state_ = PauseState::kNormalPaused;
+          SetPauseState(PauseState::kNormalPaused);
           return Status::kDiverged;
         }
       }

src/experimental/platform/sim/step_control.h

@@ -70,11 +70,10 @@ class StepControl {
   enum class PauseState { kUnpaused, kNormalPaused, kViscousPaused };
 
   // Sets the pause state of the simulation.
-  // m must be non-null for viscous pausing.
-  void SetPauseState(PauseState state, mjModel* m = nullptr);
+  void SetPauseState(PauseState state);
 
   // Gets the current pause state of the simulation.
-  PauseState GetPauseState() const { return pause_state_; }
+  PauseState GetPauseState() const;
 
   // If the simulation is paused, will perform a single step on the next
   // Advance() call.

src/experimental/platform/ux/gui.cc

@@ -15,10 +15,12 @@
 #include "experimental/platform/ux/gui.h"
 
 #include <algorithm>
+#include <array>
 #include <cmath>
 #include <cstdio>
 #include <limits>
 #include <string>
+#include <string_view>
 #include <vector>
 
 #include <imgui.h>
@@ -26,10 +28,23 @@
 #include <implot.h>
 #include <mujoco/mujoco.h>
 #include "experimental/platform/helpers.h"
+#include "experimental/platform/sim/step_control.h"
 #include "experimental/platform/ux/imgui_widgets.h"
 #include "experimental/platform/ux/interaction.h"
 
 namespace mujoco::platform {
+namespace {
+struct SpeedStatus {
+  bool misaligned;
+  float measured;
+};
+
+static SpeedStatus IsSpeedMisaligned(const StepControl& step_control) {
+  const float desired = step_control.GetSpeed();
+  const float measured = step_control.GetSpeedMeasured();
+  return {std::abs(measured - desired) > 0.1f * desired, measured};
+}
+}  // namespace
 
 static ImVec2 GetFlexElementSize(int num_cols) {
   const float width = (ImGui::GetContentRegionAvail().x / num_cols) -
@@ -317,6 +332,74 @@ ImVec4 ConfigureDockingLayout() {
   return ImVec4(workspace_x, workspace_y, workspace_w, workspace_h);
 }
 
+void StepControlGui(const mjModel* model, StepControl* step_control,
+                    int& speed_index) {
+  platform::ScopedStyle style;
+  style.Var(ImGuiStyleVar_FrameRounding, 2.f);
+
+  const ImColor yellow(255, 215, 0, 255);
+  const ImColor green(40, 180, 40, 255);
+  const float scale = ImGui::GetWindowDpiScale();
+  ImVec2 button_size(48.f * scale, 32.f * scale);
+
+  auto make_button = [&](const char* icon, StepControl::PauseState target_state,
+                         ImColor color, const char* tooltip = "",
+                         float hover_alpha = 1.f) {
+    bool active = step_control->GetPauseState() == target_state;
+    if (ImGui_ColorButton(icon, active, color, button_size, hover_alpha)) {
+      step_control->SetPauseState(target_state);
+    }
+    if (!std::string_view(tooltip).empty()) {
+      ImGui::SetItemTooltip("%s", tooltip);
+    }
+  };
+
+  make_button(ICON_FA_PAUSE, StepControl::PauseState::kNormalPaused, yellow,
+              "Pause");
+  ImGui::SameLine(0.f, 0.f);
+  make_button(ICON_FA_MAGIC, StepControl::PauseState::kViscousPaused, yellow,
+              "Viscous Pause");
+  ImGui::SameLine(0.f, 0.f);
+  make_button(ICON_FA_PLAY, StepControl::PauseState::kUnpaused, green, "", .6f);
+
+  // Speed selection.
+  ImGui::SameLine();
+  const float pad_y = (button_size.y - ImGui::GetFontSize()) * .5f;
+  ImGui::PushStyleVar(ImGuiStyleVar_FramePadding,
+                      ImVec2(ImGui::GetStyle().FramePadding.x + 5.f, pad_y));
+
+  const auto [misaligned, measured] = IsSpeedMisaligned(*step_control);
+  char speed_preview[64];
+  if (misaligned) {
+    snprintf(speed_preview, sizeof(speed_preview), "%s%s (%-4.1f%%)",
+             ICON_FA_TACHOMETER, kPercentRealTime[speed_index], measured);
+  } else {
+    snprintf(speed_preview, sizeof(speed_preview), "%s%s", ICON_FA_TACHOMETER,
+             kPercentRealTime[speed_index]);
+  }
+
+  ImGui::SetNextItemWidth(ImGui::CalcTextSize(speed_preview).x +
+                          ImGui::GetStyle().FramePadding.x * 2.f);
+  if (ImGui::BeginCombo("##Speed", speed_preview,
+                        ImGuiComboFlags_NoArrowButton)) {
+    for (int n = 0; n < kPercentRealTime.size(); n++) {
+      if (ImGui::Selectable(kPercentRealTime[n], (speed_index == n))) {
+        speed_index = std::clamp<int>(n, 0, kPercentRealTime.size() - 1);
+        float speed = std::stof(kPercentRealTime[speed_index]);
+        step_control->SetSpeed(speed);
+      }
+    }
+    ImGui::EndCombo();
+  }
+
+  ImGui::PopStyleVar();
+  if (misaligned) {
+    ImGui::SetItemTooltip("%s", "Desired Speed (Measured Speed)");
+  } else {
+    ImGui::SetItemTooltip("%s", "Desired Speed");
+  }
+}
+
 bool ThemeSelectGui(GuiTheme* theme) {
   static constexpr const char* ICON_DARKMODE = ICON_FA_CIRCLE;
   static constexpr const char* ICON_LIGHTMODE = ICON_FA_CIRCLE_O;

src/experimental/platform/ux/gui.h

@@ -23,10 +23,12 @@
 // by the caller. In most cases, this is already stored in mjModel, mjData,
 // mjvOption, etc. But, some functions take additional arguments as needed.
 
+#include <array>
 #include <vector>
 
 #include <imgui.h>
 #include <mujoco/mujoco.h>
+#include "experimental/platform/sim/step_control.h"
 
 namespace mujoco::platform {
 
@@ -62,6 +64,22 @@ void SetupTheme(GuiTheme theme);
 // be used to place additional elements (e.g. floating charts).
 ImVec4 ConfigureDockingLayout();
 
+// logarithmically spaced real-time slow-down coefficients (percent)
+// clang-format off
+static constexpr std::array<const char*, 31> kPercentRealTime = {
+"100.0 ", " 80.0 ", " 66.0 ", " 50.0 ", " 40.0 ", " 33.0 ", " 25.0 ", " 20.0 ", " 16.0 ", " 13.0 ",
+" 10.0 ", "  8.0 ", "  6.6 ", "  5.0 ", "  4.0 ", "  3.3 ", "  2.5 ", "  2.0 ", "  1.6 ", "  1.3 ",
+"  1.0 ", "  0.8 ", "  0.7 ", "  0.5 ", "  0.4 ", "  0.33", "  0.25", "  0.2 ", "  0.16", "  0.13",
+"  0.1 ",
+};
+// clang-format on
+
+// 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);
+
 // UX for selecting the GUI theme.
 bool ThemeSelectGui(GuiTheme* theme);