Add prestep/poststep callbacks to StepControl.

PiperOrigin-RevId: 936567801
Change-Id: I62d4d657791fe68101e63c2050adb63929128acf

src/experimental/platform/sim/step_control.cc

@@ -184,11 +184,17 @@ StepControl::Status StepControl::Advance(mjModel* m, mjData* d,
 
     mjtNum prev_time = d->time;
     InjectNoise(m, d);
+    if (pre_step_) {
+      pre_step_(m, d);
+    }
     if (step_fn) {
       step_fn(m, d);
     } else {
       mj_step(m, d);
     }
+    if (post_step_) {
+      post_step_(m, d);
+    }
 
     if (mjDISABLED(mjDSBL_AUTORESET)) {
       for (mjtWarning w : kDivergedWarnings) {
@@ -249,4 +255,12 @@ void StepControl::InjectNoise(const mjModel* m, mjData* d) {
   }
 }
 
+void StepControl::SetPreStepCallback(StepFn step_fn) {
+  pre_step_ = step_fn;
+}
+
+void StepControl::SetPostStepCallback(StepFn step_fn) {
+  post_step_ = step_fn;
+}
+
 }  // namespace mujoco::platform

src/experimental/platform/sim/step_control.h

@@ -69,6 +69,10 @@ class StepControl {
   void GetNoiseParameters(float& noise_scale, float& noise_rate) const;
   void SetNoiseParameters(float noise_scale, float noise_rate);
 
+  // Callbacks that will be invoked before/after each call to mj_step.
+  void SetPreStepCallback(StepFn step_fn);
+  void SetPostStepCallback(StepFn step_fn);
+
   enum class PauseState { kUnpaused, kNormalPaused, kViscousPaused };
 
   // Sets the pause state of the simulation.
@@ -120,6 +124,10 @@ class StepControl {
   // which has the effect of making the constraint solver eventually converge
   // while the simulation is paused.
   bool pause_update_ = false;
+
+  // Callbacks that can be invoked before/after physics is stepped.
+  StepFn pre_step_ = nullptr;
+  StepFn post_step_ = nullptr;
 };
 
 }  // namespace mujoco::platform