From bdb553ff8a5520cdcdb9ff28b16a2db97f207fd8 Mon Sep 17 00:00:00 2001 From: Haroon Qureshi Date: Mon, 26 Jan 2026 01:23:51 -0800 Subject: [PATCH] Add a ModelPlugin that can be used to load and update models. PiperOrigin-RevId: 861077626 Change-Id: I3ef4989cfbfff3632c5915c00ebda1e7a48716e8 --- src/experimental/platform/plugin.cc | 37 +++++++++++++ src/experimental/platform/plugin.h | 36 +++++++++++++ src/experimental/studio/app.cc | 80 +++++++++++++++++++---------- src/experimental/studio/app.h | 4 ++ 4 files changed, 131 insertions(+), 26 deletions(-) diff --git a/src/experimental/platform/plugin.cc b/src/experimental/platform/plugin.cc index f76eba91..13c262c1 100644 --- a/src/experimental/platform/plugin.cc +++ b/src/experimental/platform/plugin.cc @@ -36,10 +36,26 @@ void ForEachGuiPlugin(const std::function& 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::GetSingleton().AppendIfUnique(*plugin); +} + +void ForEachModelPlugin(const std::function& fn) { + auto& table = GlobalTable::GetSingleton(); + for (int i = 0; i < table.count(); ++i) { + const ModelPlugin* plugin = table.GetAtSlot(i); + fn(const_cast(plugin)); + } +} + } // namespace mujoco::platform using mujoco::GlobalTable; using GuiPlugin = mujoco::platform::GuiPlugin; +using ModelPlugin = mujoco::platform::ModelPlugin; template <> const char* GlobalTable::HumanReadableTypeName() { @@ -61,3 +77,24 @@ bool GlobalTable::CopyObject(GuiPlugin& dst, const GuiPlugin& src, Er dst = src; return true; } + +template <> +const char* GlobalTable::HumanReadableTypeName() { + return "model plugin"; +} + +template <> +std::string_view GlobalTable::ObjectKey(const ModelPlugin& plugin) { + return std::string_view(plugin.name); +} + +template <> +bool GlobalTable::ObjectEqual(const ModelPlugin& p1, const ModelPlugin& p2) { + return CaseInsensitiveEqual(p1.name, p2.name); +} + +template <> +bool GlobalTable::CopyObject(ModelPlugin& dst, const ModelPlugin& src, ErrorMessage& err) { + dst = src; + return true; +} diff --git a/src/experimental/platform/plugin.h b/src/experimental/platform/plugin.h index 7207f43a..824b9793 100644 --- a/src/experimental/platform/plugin.h +++ b/src/experimental/platform/plugin.h @@ -16,6 +16,7 @@ #define MUJOCO_SRC_EXPERIMENTAL_PLATFORM_PLUGIN_H_ #include +#include 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& 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& fn); + } // namespace mujoco::platform #endif // MUJOCO_SRC_EXPERIMENTAL_PLATFORM_PLUGIN_H_ diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc index b35314b0..6e20914a 100644 --- a/src/experimental/studio/app.cc +++ b/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( - reinterpret_cast(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(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) { diff --git a/src/experimental/studio/app.h b/src/experimental/studio/app.h index 0d548e48..b1cad2a8 100644 --- a/src/experimental/studio/app.h +++ b/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);