Add PreStep and PostStep plugin callbacks.

PiperOrigin-RevId: 936665884
Change-Id: Ie04562b02bfd61a62f9de77b3cef0fff534bf547

src/experimental/platform/ux/plugin.h

@@ -62,6 +62,10 @@ struct ModelPlugin final {
                                            int model_name_size);
   using PostModelLoadedFn = void (*)(ModelPlugin* self, const char* model_path);
   using DoUpdateFn = bool (*)(ModelPlugin* self, mjModel* model, mjData* data);
+  using PreStepFn = void (*)(ModelPlugin* self, const mjModel* model,
+                             mjData* data);
+  using PostStepFn = void (*)(ModelPlugin* self, const mjModel* model,
+                              mjData* data);
 
   // The name of the plugin; must be unique.
   const char* name = "";
@@ -79,6 +83,14 @@ struct ModelPlugin final {
   // simulation should be stepped.
   DoUpdateFn do_update = nullptr;
 
+  // Callback immediately before mj_step is called on the model. This may be
+  // called multiple times per update/frame.
+  PreStepFn pre_step = nullptr;
+
+  // Callback immediately after mj_step is called on the model. This may be
+  // called multiple times per update/frame.
+  PostStepFn post_step = nullptr;
+
   // Optional data pointer.
   void* data = nullptr;
 };

src/experimental/studio/app.cc

@@ -97,6 +97,11 @@ App::App(Config config)
   mjv_defaultOption(&vis_options_);
 
   profiler_.Clear();
+
+  step_control_.SetPreStepCallback(
+      [this](const mjModel* m, mjData* d) { PreStep(m, d); });
+  step_control_.SetPostStepCallback(
+      [this](const mjModel* m, mjData* d) { PostStep(m, d); });
 }
 
 void App::SwitchGraphicsMode(int width, int height,
@@ -329,6 +334,22 @@ void App::UpdatePhysics() {
   }
 }
 
+void App::PreStep(const mjModel* m, mjData* d) {
+  platform::ForEachPlugin<platform::ModelPlugin>([&](auto* plugin) {
+    if (plugin->pre_step) {
+      plugin->pre_step(plugin, m, d);
+    }
+  });
+}
+
+void App::PostStep(const mjModel* m, mjData* d) {
+  platform::ForEachPlugin<platform::ModelPlugin>([&](auto* plugin) {
+    if (plugin->post_step) {
+      plugin->post_step(plugin, m, d);
+    }
+  });
+}
+
 void App::LoadHistory(int offset) {
   std::span<mjtNum> state = sim_history_.SetIndex(offset);
   if (!state.empty()) {

src/experimental/studio/app.h

@@ -198,6 +198,8 @@ class App {
 
   void ResetPhysics();
   void UpdatePhysics();
+  void PreStep(const mjModel* m, mjData* d);
+  void PostStep(const mjModel* m, mjData* d);
 
   void LoadSettings();
   void SaveSettings();