Allow custom step function in StepControl::Advance.

The StepControl::Advance method now accepts an optional std::function to be used instead of mj_step. This allows for custom simulation logic to be executed within the stepping loop. The Python bindings for StepControl::advance have been updated to support passing a Python callable as the custom step function.

PiperOrigin-RevId: 925393132
Change-Id: Id2820c4b55cbdfee6b01adc2395e99a990d4e7a3

src/experimental/platform/sim/step_control.cc

@@ -17,6 +17,7 @@
 #include <algorithm>
 #include <chrono>
 #include <cstdlib>
+#include <functional>
 #include <optional>
 #include <ratio>
 
@@ -91,7 +92,8 @@ StepControl::PauseState StepControl::GetPauseState() const {
   return pause_state_;
 }
 
-StepControl::Status StepControl::Advance(mjModel* m, mjData* d) {
+StepControl::Status StepControl::Advance(mjModel* m, mjData* d,
+                                         StepFn step_fn) {
   if (!m) {
     return Status::kOk;
   }
@@ -182,7 +184,11 @@ StepControl::Status StepControl::Advance(mjModel* m, mjData* d) {
 
     mjtNum prev_time = d->time;
     InjectNoise(m, d);
-    mj_step(m, d);
+    if (step_fn) {
+      step_fn(m, d);
+    } else {
+      mj_step(m, d);
+    }
 
     if (mjDISABLED(mjDSBL_AUTORESET)) {
       for (mjtWarning w : kDivergedWarnings) {

src/experimental/platform/sim/step_control.h

@@ -16,6 +16,7 @@
 #define MUJOCO_SRC_EXPERIMENTAL_PLATFORM_SIM_STEP_CONTROL_H_
 
 #include <chrono>
+#include <functional>
 #include <string>
 
 #include <mujoco/mujoco.h>
@@ -24,6 +25,7 @@ namespace mujoco::platform {
 
 using Seconds = std::chrono::duration<double>;
 using Clock = std::chrono::steady_clock;
+using StepFn = std::function<void(mjModel*, mjData*)>;
 
 // State and logic for physics synchronization and stepping.
 class StepControl {
@@ -53,7 +55,7 @@ class StepControl {
       mjWARN_BADQACC, mjWARN_BADQVEL, mjWARN_BADQPOS};
 
   // Steps physics forward, respecting speed settings and refresh budget.
-  Status Advance(mjModel* m, mjData* d);
+  Status Advance(mjModel* m, mjData* d, StepFn step_fn = nullptr);
 
   // Ensures the next call to Advance() will synchronize time and step once.
   void ForceSync();