Implement viscous pause in Studio

Viscous pause is like pause but the simulation continues with zero gravity, high viscosity and passive springs disabled. It is useful to move things like cabinets with shelves---just move the shelf and the drawers come along as well.

* Updated the simulation control widgets
* Added Ctrl-space shortcut to toggle play/viscous pause

PiperOrigin-RevId: 869810053
Change-Id: Icf649e1f2631c2a64a5d2f25af9537b6a12bdf2c

src/experimental/platform/gui.cc

@@ -627,7 +627,7 @@ void PhysicsGui(mjModel* model, float min_width) {
     ImGui_InputN("Wind", opt.wind, 3);
     ImGui_InputN("Magnetic", opt.magnetic, 3);
     ImGui_Input("Density", &opt.density, {.min = .1, .max = 1});
-    ImGui_Input("Viscosity", &opt.viscosity, {.min = .1, .max = 1});
+    ImGui_Input("Viscosity", &opt.viscosity, {.min = .1, .max = 10});
     ImGui_Input("Imp Ratio", &opt.impratio, {.min = .1, .max = 1});
     ImGui::TreePop();
   };

src/experimental/platform/imgui_widgets.h

@@ -39,6 +39,7 @@ static constexpr const char ICON_FA_DIAMOND[] = "\xEF\x88\x99";
 static constexpr const char ICON_FA_EJECT[] = "\xEF\x81\x92";
 static constexpr const char ICON_FA_FAST_FORWARD[] = "\xEF\x81\x90";
 static constexpr const char ICON_FA_MOON[] = "\xEF\x86\x86";
+static constexpr const char ICON_FA_MAGIC[] = "\xEF\x83\x90";
 static constexpr const char ICON_FA_PAUSE[] = "\xEF\x81\x8C";
 static constexpr const char ICON_FA_PLAY[] = "\xEF\x81\x8B";
 static constexpr const char ICON_FA_REFRESH[] = "\xEF\x80\xA1";
@@ -262,6 +263,24 @@ inline bool ImGui_IsChordJustPressed(ImGuiKeyChord chord) {
   return ImGui::IsKeyChordPressed(chord, 0);
 }
 
+// Stateful button that displays the given color when active, and shows a
+// semi-transparent hover color (controlled by hover_alpha) when inactive.
+inline bool ImGui_ColorButton(const char* label, bool active, ImColor color,
+                              const ImVec2& size = ImVec2(0, 0),
+                              float hover_alpha = 0.5f) {
+  ScopedStyle style;
+  const ImColor hover(color.Value.x, color.Value.y, color.Value.z,
+                      color.Value.w * hover_alpha);
+  if (active) {
+    style.Color(ImGuiCol_Button, color);
+    style.Color(ImGuiCol_ButtonHovered, color);
+  } else {
+    style.Color(ImGuiCol_ButtonHovered, hover);
+  }
+  style.Color(ImGuiCol_ButtonActive, color);
+  return ImGui::Button(label, size);
+}
+
 // Begin a boxed section with outer borders - use EndBoxSection to close.
 inline bool BeginBoxSection(const char* id, ImGuiTableFlags extra_flags = 0) {
   ImGuiTableFlags flags = ImGuiTableFlags_BordersOuter | extra_flags;

src/experimental/platform/step_control.cc

@@ -13,10 +13,12 @@
 // limitations under the License.
 
 #include "experimental/platform/step_control.h"
+
 #include <algorithm>
 #include <chrono>
 #include <cstdlib>
 #include <ratio>
+
 #include <mujoco/mujoco.h>
 
 namespace mujoco::platform {
@@ -28,17 +30,11 @@ static mjtNum Timer() {
   return Milliseconds(Clock::now() - start).count();
 }
 
-StepControl::StepControl() {
-  mjcb_time = Timer;
-}
+StepControl::StepControl() { mjcb_time = Timer; }
 
-float StepControl::GetSpeedMeasured() const {
-  return speed_measured_;
-}
+float StepControl::GetSpeedMeasured() const { return speed_measured_; }
 
-float StepControl::GetSpeed() const {
-  return speed_;
-}
+float StepControl::GetSpeed() const { return speed_; }
 
 void StepControl::SetSpeed(float speed_percent_real_time) {
   speed_ = std::clamp(speed_percent_real_time, .1f, 100.f);
@@ -59,12 +55,34 @@ void StepControl::SetNoiseParameters(float ctrl_noise_scale,
   ctrl_noise_rate_ = ctrl_noise_rate;
 }
 
-StepControl::Status StepControl::Advance(const mjModel* m, mjData* d) {
+void StepControl::SetPauseState(PauseState state, mjModel* m) {
+  if (pause_state_ == PauseState::kViscousPaused &&
+      state != PauseState::kViscousPaused && m) {
+    mju_copy3(m->opt.gravity, saved_gravity_);
+    m->opt.viscosity = saved_viscosity_;
+    m->opt.disableflags = saved_disableflags_;
+  }
+
+  if (state == PauseState::kViscousPaused &&
+      pause_state_ != PauseState::kViscousPaused && m) {
+    mju_copy3(saved_gravity_, m->opt.gravity);
+    saved_viscosity_ = m->opt.viscosity;
+    saved_disableflags_ = m->opt.disableflags;
+    mju_zero3(m->opt.gravity);
+    m->opt.viscosity = 10;
+    m->opt.disableflags |= mjDSBL_SPRING;
+  }
+
+  pause_state_ = state;
+}
+
+StepControl::Status StepControl::Advance(mjModel* m, mjData* d) {
   if (!m) {
     return Status::kOk;
   }
 
-  if (paused_) {
+
+  if (pause_state_ == PauseState::kNormalPaused) {
     // When we eventually unpause, we need to make sure we sync to immediately
     // and step once. Without this we could step many times before rendering
     // resulting in a noticeable delay before the simulation restarts
@@ -151,7 +169,7 @@ StepControl::Status StepControl::Advance(const mjModel* m, mjData* d) {
       for (mjtWarning w : kDivergedWarnings) {
         if (d->warning[w].number > 0) {
           // Stop stepping if the simulation diverged.
-          paused_ = true;
+          pause_state_ = PauseState::kNormalPaused;
           return Status::kDiverged;
         }
       }
@@ -179,7 +197,7 @@ void StepControl::InjectNoise(const mjModel* m, mjData* d) {
 
   // convert rate and scale to discrete time (Ornstein–Uhlenbeck)
   mjtNum rate = mju_exp(-m->opt.timestep / ctrl_noise_rate_);
-  mjtNum scale = ctrl_noise_std_ * mju_sqrt(1-rate*rate);
+  mjtNum scale = ctrl_noise_std_ * mju_sqrt(1 - rate * rate);
 
   for (int i = 0; i < m->nu; i++) {
     mjtNum bottom = 0;
@@ -187,14 +205,14 @@ void StepControl::InjectNoise(const mjModel* m, mjData* d) {
     mjtNum midpoint = 0;
     mjtNum halfrange = 1;
     if (m->actuator_ctrllimited[i]) {
-      bottom = m->actuator_ctrlrange[2*i];
-      top = m->actuator_ctrlrange[2*i+1];
-      midpoint =  0.5 * (top + bottom);  // target of exponential decay
+      bottom = m->actuator_ctrlrange[2 * i];
+      top = m->actuator_ctrlrange[2 * i + 1];
+      midpoint = 0.5 * (top + bottom);   // target of exponential decay
       halfrange = 0.5 * (top - bottom);  // scales noise
     }
 
     // exponential convergence to midpoint at ctrl_noise_rate
-    d->ctrl[i] = rate * d->ctrl[i] + (1-rate) * midpoint;
+    d->ctrl[i] = rate * d->ctrl[i] + (1 - rate) * midpoint;
 
     // add noise
     d->ctrl[i] += scale * halfrange * mju_standardNormal(nullptr);

src/experimental/platform/step_control.h

@@ -36,6 +36,10 @@ class StepControl {
     // Simulation was not stepped because it is paused.
     kPaused,
 
+    // Simulation is viscously paused (stepping with zero gravity, high
+    // viscosity, and no spring forces).
+    kViscousPaused,
+
     // Simulation diverged with autoreset enabled.
     kAutoReset,
 
@@ -49,7 +53,7 @@ class StepControl {
       mjWARN_BADQACC, mjWARN_BADQVEL, mjWARN_BADQPOS};
 
   // Steps physics forward, respecting speed settings and refresh budget.
-  Status Advance(const mjModel* m, mjData* d);
+  Status Advance(mjModel* m, mjData* d);
 
   // Ensures the next call to Advance() will synchronize time and step once.
   void ForceSync();
@@ -63,13 +67,14 @@ class StepControl {
   void GetNoiseParameters(float& noise_scale, float& noise_rate) const;
   void SetNoiseParameters(float noise_scale, float noise_rate);
 
-  // Returns true if the simulation is paused.
-  bool IsPaused() { return paused_; }
+  enum class PauseState { kUnpaused, kNormalPaused, kViscousPaused };
 
-  // Pauses/unpauses the simulation.
-  void Pause() { paused_ = true; }
-  void Unpause() { paused_ = false; }
-  void TogglePause() { paused_ = !paused_; }
+  // Sets the pause state of the simulation.
+  // m must be non-null for viscous pausing.
+  void SetPauseState(PauseState state, mjModel* m = nullptr);
+
+  // Gets the current pause state of the simulation.
+  PauseState GetPauseState() const { return pause_state_; }
 
   // If the simulation is paused, will perform a single step on the next
   // Advance() call.
@@ -104,8 +109,12 @@ class StepControl {
   // Maximum mis-alignment before re-sync (simulation seconds)
   double sync_misalign_ = .1;
 
-  // Whether or not the simulation is paused.
-  bool paused_ = false;
+  PauseState pause_state_ = PauseState::kUnpaused;
+
+  // Viscous pause state variables
+  mjtNum saved_gravity_[3] = {0};
+  mjtNum saved_viscosity_ = 0;
+  int saved_disableflags_ = 0;
 
   // Perform only a single step on the next call to Advance() if the simulation
   // is paused.