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

python/mujoco/experimental/studio/sim.cc

@@ -42,11 +42,21 @@ PYBIND11_MODULE(sim, m) {
       .def(py::init<>())
       .def(
           "advance",
-          [](StepControl& self, mujoco::python::MjModelWrapper& model,
-             mujoco::python::MjDataWrapper& data) {
-            return self.Advance(model.get(), data.get());
+          [](StepControl& self, py::object model_obj, py::object data_obj,
+             py::object step_fn) {
+            auto& model = py::cast<mujoco::python::MjModelWrapper&>(model_obj);
+            auto& data = py::cast<mujoco::python::MjDataWrapper&>(data_obj);
+            if (step_fn.is_none()) {
+              return self.Advance(model.get(), data.get());
+            } else {
+              return self.Advance(
+                  model.get(), data.get(),
+                  [step_fn, model_obj, data_obj](mjModel*, mjData*) {
+                    step_fn(model_obj, data_obj);
+                  });
+            }
           },
-          py::arg("model"), py::arg("data"),
+          py::arg("model"), py::arg("data"), py::arg("step_fn") = py::none(),
           "Step physics forward, respecting speed settings and refresh budget.")
       .def("force_sync", &StepControl::ForceSync,
            "Ensures the next Advance() will synchronize time and step once.")

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();