Move pause handling into StepControl.

Also moves perturbation handling up to the App.

Combined, this makes Physics::Update only about advancing the StepControl and resonding to its Status.

PiperOrigin-RevId: 836569506
Change-Id: I148c787f0e1c18baac3e3038ee6ad0e337eb92a8

src/experimental/studio/app.cc

@@ -252,7 +252,16 @@ bool App::Update() {
   // Only update the simulation if a popup window is not open. Note that the
   // simulation itself will only update if it is not paused.
   if (!tmp_.modal_open) {
-    if (physics_->Update(&perturb_)) {
+    // Apply any active perturbations.
+    if (!physics_->GetStepControl().IsPaused() && Model() && Data()) {
+      mju_zero(Data()->xfrc_applied, 6 * Model()->nbody);
+      mjv_applyPerturbPose(Model(), Data(), &perturb_, 0);
+      mjv_applyPerturbForce(Model(), Data(), &perturb_);
+    } else {
+      mjv_applyPerturbPose(Model(), Data(), &perturb_, 1);
+    }
+
+    if (physics_->Update()) {
       profiler_.Update(Model(), Data());
     }
   }
@@ -436,19 +445,19 @@ void App::HandleKeyboardEvents() {
   } else if (ImGui_IsChordJustPressed(ImGuiKey_Equal)) {
     SetSpeedIndex(tmp_.speed_index - 1);
   } else if (ImGui_IsChordJustPressed(ImGuiKey_LeftArrow)) {
-    if (physics_->IsPaused()) {
+    if (physics_->GetStepControl().IsPaused()) {
       physics_->LoadHistory(ui_.scrub_idx - 1);
     }
   } else if (ImGui_IsChordJustPressed(ImGuiKey_RightArrow)) {
-    if (physics_->IsPaused()) {
+    if (physics_->GetStepControl().IsPaused()) {
       if (ui_.scrub_idx == 0) {
-        physics_->RequestSingleStep();
+        physics_->GetStepControl().RequestSingleStep();
       } else {
         physics_->LoadHistory(ui_.scrub_idx + 1);
       }
     }
   } else if (ImGui_IsChordJustPressed(ImGuiKey_Space)) {
-    physics_->TogglePause();
+    physics_->GetStepControl().TogglePause();
   } else if (ImGui_IsChordJustPressed(ImGuiKey_Backspace)) {
     physics_->Reset();
   } else if (ImGui_IsChordJustPressed(ImGuiKey_PageUp)) {
@@ -1054,7 +1063,8 @@ void App::InfoGui() {
   }
   solver_err = mju_log10(mju_max(mjMINVAL, solver_err));
 
-  auto type = physics_->IsPaused() ? mjTIMER_FORWARD : mjTIMER_STEP;
+  auto type =
+      physics_->GetStepControl().IsPaused() ? mjTIMER_FORWARD : mjTIMER_STEP;
   auto cpu =
       Data()->timer[type].duration / mjMAX(1, Data()->timer[type].number);
   auto mempct = 100 * Data()->maxuse_arena / (double)(Data()->narena);
@@ -1110,9 +1120,9 @@ void App::ToolBarGui() {
     ImGui::TableNextColumn();
 
     // Play/pause button.
-    const bool paused = physics_->IsPaused();
+    const bool paused = physics_->GetStepControl().IsPaused();
     if (ImGui::Button(paused ? ICON_PLAY : ICON_PAUSE, ImVec2(144, 32))) {
-      physics_->TogglePause();
+      physics_->GetStepControl().TogglePause();
     }
     ImGui::SetItemTooltip("%s", paused ? "Play" : "Pause");
 
@@ -1209,7 +1219,7 @@ void App::StatusBarGui() {
       ImGui::Text("Not loaded");
     } else if (Model() == nullptr) {
       ImGui::Text("Not loaded");
-    } else if (physics_->IsPaused()) {
+    } else if (physics_->GetStepControl().IsPaused()) {
       ImGui::Text("Paused");
     } else {
       const float desired_realtime = physics_->GetStepControl().GetSpeed();
@@ -1273,7 +1283,7 @@ void App::StatusBarGui() {
     ImGui::SameLine();
     if (ImGui::Button(ICON_NEXT_FRAME)) {
       if (ui_.scrub_idx == 0) {
-        physics_->RequestSingleStep();
+        physics_->GetStepControl().RequestSingleStep();
       } else {
         ui_.scrub_idx = std::min(0, ui_.scrub_idx + 1);
         physics_->LoadHistory(ui_.scrub_idx);
@@ -1326,8 +1336,9 @@ void App::MainMenuGui() {
       ImGui::EndMenu();
     }
     if (ImGui::BeginMenu("Simulation")) {
-      if (ImGui::MenuItem("Pause", "Space", physics_->IsPaused())) {
+      if (ImGui::MenuItem("Pause", "Space",
-        physics_->TogglePause();
+                          physics_->GetStepControl().IsPaused())) {
+        physics_->GetStepControl().TogglePause();
       }
       if (ImGui::MenuItem("Reset", "Backspace")) {
         physics_->Reset();

src/experimental/toolbox/physics.cc

@@ -17,7 +17,6 @@
 #include <algorithm>
 #include <chrono>
 #include <climits>
-#include <cstring>
 #include <functional>
 #include <ratio>
 #include <string>
@@ -62,14 +61,14 @@ bool Physics::ProcessPendingLoad() {
   model_ = LoadMujocoModel(model_file, vfs_);
   if (!model_) {
     error_ = "Error loading model!";
-    paused_ = true;
+    step_control_.Pause();
     model_ = LoadMujocoModel("", vfs_);
   }
 
   data_ = mj_makeData(model_);
   if (!data_) {
     error_ = "Error making data!";
-    paused_ = true;
+    step_control_.Pause();
   }
 
   on_model_loaded_(model_file);
@@ -89,7 +88,7 @@ void Physics::Clear() {
     history_.clear();
     history_cursor_ = 0;
     steps_ = 0;
-    GetStepControl().SetSpeed(100.f);
+    step_control_.SetSpeed(100.f);
 
     error_ = "";
   }
@@ -102,7 +101,7 @@ void Physics::Reset() {
   history_cursor_ = 0;
 }
 
-bool Physics::Update(const mjvPerturb* perturb) {
+bool Physics::Update() {
   ProcessPendingLoad();
 
   if (!model_ || !data_) {
@@ -116,42 +115,17 @@ bool Physics::Update(const mjvPerturb* perturb) {
     }
   }
 
-  if (!IsPaused()) {
+  StepControl::Status status = step_control_.Advance(model_, data_);
-    mju_zero(data_->xfrc_applied, 6 * model_->nbody);
+  if (status == StepControl::Status::kOk) {
-    mjv_applyPerturbPose(model_, data_, perturb, 0);
+    AddToHistory();
-    mjv_applyPerturbForce(model_, data_, perturb);
+  } else if (status == StepControl::Status::kPaused) {
-  } else {
+    // do nothing
-    mjv_applyPerturbPose(model_, data_, perturb, 1);
+  } else if (status == StepControl::Status::kAutoReset) {
-  }
+    Reset();
-
+  } else if (status == StepControl::Status::kDiverged) {
-  if (IsPaused() && !single_step_) {
+    for (mjtWarning w : StepControl::kDivergedWarnings) {
-    // run mj_forward, to update rendering and joint sliders
+      if (data_->warning[w].number > 0) {
-    mj_forward(model_, data_);
+        error_ = mju_warningText(w, data_->warning[w].lastinfo);
-    if (pause_update_) {
-      mju_copy(data_->qacc_warmstart, data_->qacc, model_->nv);
-    }
-
-    // When unpaused 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 (especially for large slowdowns)
-    GetStepControl().ForceSync();
-  } else {
-    if (single_step_) {
-      GetStepControl().ForceSync();
-      single_step_ = false;
-    }
-
-    StepControl::Status status = GetStepControl().Advance(model_, data_);
-    if (status == StepControl::Status::kOk) {
-      AddToHistory();
-    } else if (status == StepControl::Status::kAutoReset) {
-      Reset();
-    } else if (status == StepControl::Status::kDiverged) {
-      for (mjtWarning w : StepControl::kDivergedWarnings) {
-        if (data_->warning[w].number > 0) {
-          paused_ = true;
-          error_ = mju_warningText(w, data_->warning[w].lastinfo);
-        }
       }
     }
   }
@@ -169,10 +143,6 @@ bool Physics::UpdateState(mjtNum* state, unsigned int state_sig) {
   return true;
 }
 
-void Physics::TogglePause() { paused_ = !paused_; }
-
-void Physics::RequestSingleStep() { single_step_ = true; }
-
 void Physics::InitHistory() {
   const int state_size = mj_stateSize(model_, mjSTATE_INTEGRATION);
 
@@ -210,7 +180,7 @@ int Physics::LoadHistory(int offset) {
   }
 
   // Pause simulation when entering history mode.
-  paused_ = true;
+  step_control_.Pause();
 
   // Ensure the offset is within a valid range. It's a negative value since
   // we will be going backwards from the "latest" frame.

src/experimental/toolbox/physics.h

@@ -21,7 +21,6 @@
 #include <string_view>
 #include <vector>
 
-
 #include <mujoco/mujoco.h>
 #include "experimental/toolbox/step_control.h"
 
@@ -53,24 +52,11 @@ class Physics {
   void Reset();
 
   // Advances the state of the simulation.
-  bool Update(const mjvPerturb* perturb);
+  bool Update();
 
   // Sets the state of the simulation.
   bool UpdateState(mjtNum* state, unsigned int state_sig);
 
-  // Renders the state of the simulation.
-  void Render();
-
-  // Returns true if the simulation is paused.
-  bool IsPaused() { return paused_; }
-
-  // Pauses/unpauses the simulation.
-  void TogglePause();
-
-  // If the simulation is paused, will perform a single step on the next
-  // Update() call.
-  void RequestSingleStep();
-
   // Returns the number of steps the simulation has taken.
   int GetStepCount() const { return steps_; }
 
@@ -83,9 +69,6 @@ class Physics {
   // Calling this function will automatically pause the simulation.
   int LoadHistory(int offset);
 
-  // Selects the parent of the currently selected perturb object.
-  void SelectParentPerturb();
-
   // Returns the MuJoCo data structures owned by this Simulation object.
   mjModel* GetModel() { return model_; }
   mjData* GetData() { return data_; }
@@ -106,13 +89,6 @@ class Physics {
 
   OnModelLoadedFn on_model_loaded_;
 
-  // If true and paused, d->qacc_warmstart is set to d->qacc after mj_forward
-  // which has the effect of making the constraint solver eventually converge
-  // while the simulation is paused.
-  bool pause_update_ = false;
-
-  bool paused_ = false;
-  bool single_step_ = false;
   int steps_ = 0;
   std::optional<std::string> pending_load_;
   const mjVFS* vfs_;

src/experimental/toolbox/step_control.cc

@@ -52,6 +52,24 @@ StepControl::Status StepControl::Advance(const mjModel* m, mjData* d) {
     return Status::kOk;
   }
 
+  if (paused_) {
+    // 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
+    // (especially for large slowdowns).
+    force_sync_ = true;
+
+    if (!single_step_) {
+      // Run mj_forward to update rendering and joint sliders.
+      mj_forward(m, d);
+      if (pause_update_) {
+        mju_copy(d->qacc_warmstart, d->qacc, m->nv);
+      }
+      return Status::kPaused;
+    }
+    single_step_ = false;
+  }
+
   const Clock::time_point start_cpu = Clock::now();
   const double slowdown = 100. / std::clamp<double>(speed_, 0.001, 100.);
   double elapsed_cpu = Seconds(start_cpu - sync_cpu_).count();
@@ -119,8 +137,9 @@ StepControl::Status StepControl::Advance(const mjModel* m, mjData* d) {
 
     if (mjDISABLED(mjDSBL_AUTORESET)) {
       for (mjtWarning w : kDivergedWarnings) {
-        // Stop stepping if the simulation diverged.
         if (d->warning[w].number > 0) {
+          // Stop stepping if the simulation diverged.
+          paused_ = true;
           return Status::kDiverged;
         }
       }

src/experimental/toolbox/step_control.h

@@ -31,6 +31,9 @@ class StepControl {
   enum class Status {
     kOk,
 
+    // Simulation was not stepped because it is paused.
+    kPaused,
+
     // Simulation diverged with autoreset enabled.
     kAutoReset,
 
@@ -58,6 +61,18 @@ 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_; }
+
+  // Pauses/unpauses the simulation.
+  void Pause() { paused_ = true; }
+  void Unpause() { paused_ = false; }
+  void TogglePause() { paused_ = !paused_; }
+
+  // If the simulation is paused, will perform a single step on the next
+  // Advance() call.
+  void RequestSingleStep() { single_step_ = true; }
+
  private:
   std::string AdvanceOneStep(const mjModel* m, mjData* d);
 
@@ -86,6 +101,18 @@ class StepControl {
 
   // Maximum mis-alignment before re-sync (simulation seconds)
   double sync_misalign_ = .1;
+
+  // Whether or not the simulation is paused.
+  bool paused_ = false;
+
+  // Perform only a single step on the next call to Advance() if the simulation
+  // is paused.
+  bool single_step_ = false;
+
+  // If true and paused, d->qacc_warmstart is set to d->qacc after mj_forward
+  // which has the effect of making the constraint solver eventually converge
+  // while the simulation is paused.
+  bool pause_update_ = false;
 };
 
 }  // namespace mujoco::toolbox