Add a ModelPlugin that can be used to load and update models.

PiperOrigin-RevId: 861077626
Change-Id: I3ef4989cfbfff3632c5915c00ebda1e7a48716e8

src/experimental/platform/plugin.cc

@@ -36,10 +36,26 @@ void ForEachGuiPlugin(const std::function<void(GuiPlugin*)>& fn) {
   }
 }
 
+void RegisterModelPlugin(const ModelPlugin* plugin) {
+  if (plugin->name == nullptr || plugin->name[0] == '\0') {
+    mju_error("Plugin name must not be empty or null.");
+  }
+  GlobalTable<ModelPlugin>::GetSingleton().AppendIfUnique(*plugin);
+}
+
+void ForEachModelPlugin(const std::function<void(ModelPlugin*)>& fn) {
+  auto& table = GlobalTable<ModelPlugin>::GetSingleton();
+  for (int i = 0; i < table.count(); ++i) {
+    const ModelPlugin* plugin = table.GetAtSlot(i);
+    fn(const_cast<ModelPlugin*>(plugin));
+  }
+}
+
 }  // namespace mujoco::platform
 
 using mujoco::GlobalTable;
 using GuiPlugin = mujoco::platform::GuiPlugin;
+using ModelPlugin = mujoco::platform::ModelPlugin;
 
 template <>
 const char* GlobalTable<GuiPlugin>::HumanReadableTypeName() {
@@ -61,3 +77,24 @@ bool GlobalTable<GuiPlugin>::CopyObject(GuiPlugin& dst, const GuiPlugin& src, Er
   dst = src;
   return true;
 }
+
+template <>
+const char* GlobalTable<ModelPlugin>::HumanReadableTypeName() {
+  return "model plugin";
+}
+
+template <>
+std::string_view GlobalTable<ModelPlugin>::ObjectKey(const ModelPlugin& plugin) {
+  return std::string_view(plugin.name);
+}
+
+template <>
+bool GlobalTable<ModelPlugin>::ObjectEqual(const ModelPlugin& p1, const ModelPlugin& p2) {
+  return CaseInsensitiveEqual(p1.name, p2.name);
+}
+
+template <>
+bool GlobalTable<ModelPlugin>::CopyObject(ModelPlugin& dst, const ModelPlugin& src, ErrorMessage& err) {
+  dst = src;
+  return true;
+}

src/experimental/platform/plugin.h

@@ -16,6 +16,7 @@
 #define MUJOCO_SRC_EXPERIMENTAL_PLATFORM_PLUGIN_H_
 
 #include <functional>
+#include <mujoco/mujoco.h>
 
 namespace mujoco::platform {
 
@@ -46,6 +47,34 @@ struct GuiPlugin final {
   void* data = nullptr;
 };
 
+struct ModelPlugin final {
+  using GetModelToLoadFn = const char* (*)(ModelPlugin * self, int* size,
+                                           char* content_type,
+                                           int content_type_size,
+                                           char* model_name,
+                                           int model_name_size);
+  using PostModelLoadedFn = void (*)(ModelPlugin* self, const char* model_path);
+  using DoUpdateFn = bool (*)(ModelPlugin* self, mjModel* model, mjData* data);
+
+  // The name of the plugin; must be unique.
+  const char* name = "";
+
+  // Callback for when the plugin wants to load a new model. This function will
+  // return a buffer containing the model data as well as the content type of
+  // the buffer. Returns nullptr if no model needs to be loaded.s
+  GetModelToLoadFn get_model_to_load = nullptr;
+
+  // Callback when a new model is loaded.
+  PostModelLoadedFn post_model_loaded = nullptr;
+
+  // Callback when the physics simulation is updated. Returns true if the
+  // simulation should be stepped.
+  DoUpdateFn do_update = nullptr;
+
+  // Optional data pointer.
+  void* data = nullptr;
+};
+
 // Registers a plugin with a global registry. The plugin must have a
 // case-insensitive unique name.
 void RegisterGuiPlugin(const GuiPlugin* plugin);
@@ -53,6 +82,13 @@ void RegisterGuiPlugin(const GuiPlugin* plugin);
 // Executes the given function for each registered plugin.
 void ForEachGuiPlugin(const std::function<void(GuiPlugin*)>& fn);
 
+// Registers a plugin with a global registry. The plugin must have a
+// case-insensitive unique name.
+void RegisterModelPlugin(const ModelPlugin* plugin);
+
+// Executes the given function for each registered plugin.
+void ForEachModelPlugin(const std::function<void(ModelPlugin*)>& fn);
+
 }  // namespace mujoco::platform
 
 #endif  // MUJOCO_SRC_EXPERIMENTAL_PLATFORM_PLUGIN_H_

src/experimental/studio/app.cc

@@ -156,7 +156,6 @@ void App::ClearModel() {
 
 void App::RequestModelLoad(std::string model_file) {
   pending_load_ = std::move(model_file);
-  model_kind_ = kModelFromFile;
 }
 
 void App::RequestModelReload() {
@@ -277,6 +276,12 @@ void App::InitModel(mjModel* model, mjSpec* spec, mjVFS* vfs,
       SetSpeedIndex(i);
     }
   }
+
+  platform::ForEachModelPlugin([&](platform::ModelPlugin* plugin) {
+    if (plugin->post_model_loaded) {
+      plugin->post_model_loaded(plugin, model_path_.c_str());
+    }
+  });
 }
 
 void App::UpdateFilePaths(const std::string& resolved_path) {
@@ -309,9 +314,16 @@ void App::UpdatePhysics() {
     return;
   }
 
-
   bool stepped = false;
-  {
+  platform::ForEachModelPlugin([&](platform::ModelPlugin* plugin) {
+    if (plugin->do_update) {
+      if (plugin->do_update(plugin, model_, data_)) {
+        stepped = true;
+      }
+    }
+  });
+
+  if (!stepped) {
     if (!step_control_.IsPaused()) {
       mju_zero(data_->xfrc_applied, 6 * model_->nbody);
       mjv_applyPerturbPose(model_, data_, &perturb_, 0);
@@ -370,32 +382,11 @@ void App::LoadHistory(int offset) {
 bool App::Update() {
   const platform::Window::Status status = window_->NewFrame();
 
+  HandleWindowEvents();
   HandleMouseEvents();
   HandleKeyboardEvents();
 
-  // Check to see if a model was dropped on the window.
-  const std::string drop_file = window_->GetDropFile();
-  if (!drop_file.empty()) {
-    RequestModelLoad(drop_file);
-  }
-
-  // Check to see if we need to load a new model.
-  if (pending_load_.has_value()) {
-    std::string load_data = std::move(pending_load_.value());
-    pending_load_.reset();
-    if (model_kind_ == kModelFromBuffer) {
-      // TODO(matijak): need to pass the content type and not assume mjb.
-      LoadModelFromBuffer(
-          std::span<const std::byte>(
-              reinterpret_cast<const std::byte*>(load_data.data()),
-              load_data.size()),
-          "application/mjb", model_name_);
-    } else if (model_kind_ == kModelFromFile) {
-      LoadModelFromFile(load_data);
-    } else {
-      InitEmptyModel();
-    }
-  }
+  ProcessPendingLoads();
 
   // Only update the simulation if a popup window is not open. Note that the
   // simulation itself will only update if it is not paused.
@@ -425,6 +416,43 @@ void App::Render() {
   }
 }
 
+void App::ProcessPendingLoads() {
+  // Check to see if we need to load a new model.
+  if (pending_load_.has_value()) {
+    std::string load_data = std::move(pending_load_.value());
+    pending_load_.reset();
+
+    if (load_data.empty()) {
+      InitEmptyModel();
+    } else {
+      LoadModelFromFile(load_data);
+    }
+  }
+
+  // Check plugins to see if we need to load a new model.
+  platform::ForEachModelPlugin([&](platform::ModelPlugin* plugin) {
+    if (plugin->get_model_to_load) {
+      char model_name[1000] = "";
+      char content_type[1000] = "";
+      int size = 0;
+      const char* buf = plugin->get_model_to_load(
+          plugin, &size, content_type, sizeof(content_type), model_name,
+          sizeof(model_name));
+      if (buf && size) {
+        const std::byte* bytes = reinterpret_cast<const std::byte*>(buf);
+        LoadModelFromBuffer({bytes, bytes + size}, content_type, model_name);
+      }
+    }
+  });
+}
+
+void App::HandleWindowEvents() {
+  const std::string drop_file = window_->GetDropFile();
+  if (!drop_file.empty()) {
+    RequestModelLoad(drop_file);
+  }
+}
+
 void App::HandleMouseEvents() {
   auto& io = ImGui::GetIO();
   if (io.WantCaptureMouse) {

src/experimental/studio/app.h

@@ -168,8 +168,12 @@ class App {
 
   void SetSpeedIndex(int idx);
 
+  void HandleWindowEvents();
   void HandleMouseEvents();
   void HandleKeyboardEvents();
+
+  void ProcessPendingLoads();
+
   void MoveCamera(platform::CameraMotion motion, mjtNum reldx, mjtNum reldy);
 
   void SetupTheme(platform::GuiTheme theme);