Add plugin for adding mjvGeoms to the scene.

PiperOrigin-RevId: 936709560
Change-Id: I69a01fc50e2d4b013d260d524515719c434bbf5f

src/experimental/platform/hal/renderer.cc

@@ -14,6 +14,7 @@
 
 #include "experimental/platform/hal/renderer.h"
 
+#include <algorithm>
 #include <chrono>
 #include <cstddef>
 #include <functional>
@@ -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<std::byte> pixels) {
+                      std::span<std::byte> pixels,
+                      std::span<mjvGeom> 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<int>(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) {

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<std::byte> pixels = {});
+              int height, std::span<std::byte> pixels = {},
+              std::span<mjvGeom> extra_geoms = {});
 
   // Populates the given output buffer with RGB888 pixel data. The size of the
   // output buffer must be at least width * height * 3.

src/experimental/platform/ux/plugin.cc

@@ -20,11 +20,6 @@
 #include <mujoco/mujoco.h>
 #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 <typename T>
@@ -71,7 +66,8 @@ void ForEachPlugin(const std::function<void(T*)>& fn) {
   template void ForEachPlugin<PLUGIN>(const std::function<void(PLUGIN*)>& 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");

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,

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<platform::ScenePlugin>([&](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_);

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_;