From 5637f7432744676e0282c3fa0a66fbd165870b98 Mon Sep 17 00:00:00 2001 From: Haroon Qureshi Date: Tue, 23 Jun 2026 09:16:05 -0700 Subject: [PATCH] Add plugin for adding mjvGeoms to the scene. PiperOrigin-RevId: 936709560 Change-Id: I69a01fc50e2d4b013d260d524515719c434bbf5f --- src/experimental/platform/hal/renderer.cc | 9 ++++++++- src/experimental/platform/hal/renderer.h | 3 ++- src/experimental/platform/ux/plugin.cc | 14 +++++--------- src/experimental/platform/ux/plugin.h | 15 +++++++++++++++ src/experimental/studio/app.cc | 17 ++++++++++++++++- src/experimental/studio/app.h | 2 ++ 6 files changed, 48 insertions(+), 12 deletions(-) diff --git a/src/experimental/platform/hal/renderer.cc b/src/experimental/platform/hal/renderer.cc index 81d983b6..080da346 100644 --- a/src/experimental/platform/hal/renderer.cc +++ b/src/experimental/platform/hal/renderer.cc @@ -14,6 +14,7 @@ #include "experimental/platform/hal/renderer.h" +#include #include #include #include @@ -136,7 +137,8 @@ void Renderer::Deinit() { void Renderer::Render(const mjModel* model, mjData* data, const mjvPerturb* perturb, mjvCamera* camera, const mjvOption* vis_option, int width, int height, - std::span pixels) { + std::span pixels, + std::span extra_geoms) { if (!initialized_) { return; } @@ -157,6 +159,11 @@ void Renderer::Render(const mjModel* model, mjData* data, } mjv_updateScene(model, data, vis_option, perturb, camera, mjCAT_ALL, &scene_); + const int nextra_geoms = + std::min(extra_geoms.size(), scene_.maxgeom - scene_.ngeom); + for (int i = 0; i < nextra_geoms; ++i) { + scene_.geoms[scene_.ngeom++] = extra_geoms[i]; + } const bool render_to_texture = !pixels.empty(); if (render_to_texture) { diff --git a/src/experimental/platform/hal/renderer.h b/src/experimental/platform/hal/renderer.h index f94a7570..4d1e1db8 100644 --- a/src/experimental/platform/hal/renderer.h +++ b/src/experimental/platform/hal/renderer.h @@ -73,7 +73,8 @@ class Renderer { // otherwise renders to the `native_window` provided at construction. void Render(const mjModel* model, mjData* data, const mjvPerturb* perturb, mjvCamera* camera, const mjvOption* vis_option, int width, - int height, std::span pixels = {}); + int height, std::span pixels = {}, + std::span extra_geoms = {}); // Populates the given output buffer with RGB888 pixel data. The size of the // output buffer must be at least width * height * 3. diff --git a/src/experimental/platform/ux/plugin.cc b/src/experimental/platform/ux/plugin.cc index 058e9bd5..60bf0c7e 100644 --- a/src/experimental/platform/ux/plugin.cc +++ b/src/experimental/platform/ux/plugin.cc @@ -20,11 +20,6 @@ #include #include "engine/engine_global_table.h" -using GuiPlugin = mujoco::platform::GuiPlugin; -using ModelPlugin = mujoco::platform::ModelPlugin; -using KeyHandlerPlugin = mujoco::platform::KeyHandlerPlugin; -using SpecEditorPlugin = mujoco::platform::SpecEditorPlugin; - namespace mujoco::platform { template @@ -71,7 +66,8 @@ void ForEachPlugin(const std::function& fn) { template void ForEachPlugin(const std::function& fn); \ } -MUJOCO_SPECIALIZE_PLUGIN(GuiPlugin, "gui plugin"); -MUJOCO_SPECIALIZE_PLUGIN(ModelPlugin, "model plugin"); -MUJOCO_SPECIALIZE_PLUGIN(KeyHandlerPlugin, "key handler plugin"); -MUJOCO_SPECIALIZE_PLUGIN(SpecEditorPlugin, "spec editor plugin"); +MUJOCO_SPECIALIZE_PLUGIN(mujoco::platform::GuiPlugin, "gui plugin"); +MUJOCO_SPECIALIZE_PLUGIN(mujoco::platform::ModelPlugin, "model plugin"); +MUJOCO_SPECIALIZE_PLUGIN(mujoco::platform::ScenePlugin, "scene plugin"); +MUJOCO_SPECIALIZE_PLUGIN(mujoco::platform::KeyHandlerPlugin, "key handler plugin"); +MUJOCO_SPECIALIZE_PLUGIN(mujoco::platform::SpecEditorPlugin, "spec editor plugin"); diff --git a/src/experimental/platform/ux/plugin.h b/src/experimental/platform/ux/plugin.h index 79a9bf72..42b32197 100644 --- a/src/experimental/platform/ux/plugin.h +++ b/src/experimental/platform/ux/plugin.h @@ -53,6 +53,21 @@ struct GuiPlugin final { void* data = nullptr; }; +// Plugin for adding additional scene geoms. +struct ScenePlugin final { + using EnhanceSceneFn = void (*)(ScenePlugin* self, const mjModel* model, + mjData* data, mjvScene* scene); + + // The name of the plugin; must be unique. + const char* name = ""; + + // This function will add additional mjvGeom objects to the mjvScene. + EnhanceSceneFn enhance_scene = nullptr; + + // Optional data pointer. + void* data = nullptr; +}; + // Plugin for loading and updating models. struct ModelPlugin final { using GetModelToLoadFn = const char* (*)(ModelPlugin* self, int* size, diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc index 7b5855e6..0d411b52 100644 --- a/src/experimental/studio/app.cc +++ b/src/experimental/studio/app.cc @@ -95,6 +95,8 @@ App::App(Config config) mjv_defaultPerturb(&perturb_); mjv_defaultCamera(&camera_); mjv_defaultOption(&vis_options_); + std::memset(&plugin_scene_, 0, sizeof(mjvScene)); + mjv_makeScene(nullptr, &plugin_scene_, 2000); profiler_.Clear(); @@ -104,6 +106,10 @@ App::App(Config config) [this](const mjModel* m, mjData* d) { PostStep(m, d); }); } +App::~App() { + mjv_freeScene(&plugin_scene_); +} + void App::SwitchGraphicsMode(int width, int height, platform::GraphicsMode mode) { renderer_.reset(); @@ -389,8 +395,17 @@ void App::Render() { } else { pixels_.clear(); } + + plugin_scene_.ngeom = 0; + platform::ForEachPlugin([&](auto* plugin) { + if (plugin->enhance_scene) { + plugin->enhance_scene(plugin, model(), data(), &plugin_scene_); + } + }); + renderer_->Render(model(), data(), &perturb_, &camera_, &vis_options_, - width * scale, height * scale, pixels_); + width * scale, height * scale, pixels_, + {plugin_scene_.geoms, (size_t)plugin_scene_.ngeom}); window_->EndFrame(); window_->Present(pixels_); diff --git a/src/experimental/studio/app.h b/src/experimental/studio/app.h index d70fb739..cd1d50b4 100644 --- a/src/experimental/studio/app.h +++ b/src/experimental/studio/app.h @@ -64,6 +64,7 @@ class App { }; explicit App(Config config); + ~App(); // Loads an empty mjModel. void InitEmptyModel(); @@ -263,6 +264,7 @@ class App { mjvCamera camera_; mjvPerturb perturb_; mjvOption vis_options_; + mjvScene plugin_scene_; UiState ui_; UiTempState tmp_;