diff --git a/src/experimental/platform/sim/step_control.cc b/src/experimental/platform/sim/step_control.cc index a41a6bbe..88d33f18 100644 --- a/src/experimental/platform/sim/step_control.cc +++ b/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 diff --git a/src/experimental/platform/sim/step_control.h b/src/experimental/platform/sim/step_control.h index 1717c975..b3bd5b2f 100644 --- a/src/experimental/platform/sim/step_control.h +++ b/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