Allow user-defined visualization geoms to be added to the passive viewer.

Fixes #1023, #1082

PiperOrigin-RevId: 572303635
Change-Id: Ia8399f816f311b38a7b137f5bd4f359f77a1da1c

simulate/main.cc

@@ -474,9 +474,6 @@ int main(int argc, char** argv) {
   // scan for libraries in the plugin directory to load additional plugins
   scanPluginLibraries();
 
-  mjvScene scn;
-  mjv_defaultScene(&scn);
-
   mjvCamera cam;
   mjv_defaultCamera(&cam);
 
@@ -489,7 +486,7 @@ int main(int argc, char** argv) {
   // simulate object encapsulates the UI
   auto sim = std::make_unique<mj::Simulate>(
       std::make_unique<mj::GlfwAdapter>(),
-      &scn, &cam, &opt, &pert, /* is_passive = */ false
+      &cam, &opt, &pert, /* is_passive = */ false
   );
 
   const char* filename = nullptr;

simulate/simulate.cc

@@ -1737,16 +1737,15 @@ namespace mujoco {
 namespace mju = ::mujoco::sample_util;
 
 Simulate::Simulate(std::unique_ptr<PlatformUIAdapter> platform_ui,
-                   mjvScene* scn, mjvCamera* cam,
-                   mjvOption* opt, mjvPerturb* pert,
+                   mjvCamera* cam, mjvOption* opt, mjvPerturb* pert,
                    bool is_passive)
     : is_passive_(is_passive),
-      scn(*scn),
       cam(*cam),
       opt(*opt),
       pert(*pert),
       platform_ui(std::move(platform_ui)),
       uistate(this->platform_ui->state()) {
+  mjv_defaultScene(&scn);
   mjv_defaultSceneState(&scnstate_);
 }
 
@@ -2011,6 +2010,23 @@ void Simulate::Sync() {
     mjv_updateScene(m_, d_, &this->opt, &this->pert, &this->cam, mjCAT_ALL, &this->scn);
   } else {
     mjv_updateSceneState(m_, d_, &this->opt, &scnstate_);
+
+    // append geoms from user_scn to scnstate_ scratch space
+    if (user_scn) {
+      int ngeom = user_scn->ngeom;
+      int maxgeom = scnstate_.scratch.maxgeom - scnstate_.scratch.ngeom;
+      if (ngeom > maxgeom) {
+        mj_warning(d_, mjWARN_VGEOMFULL, scnstate_.scratch.maxgeom);
+        ngeom = maxgeom;
+      }
+      if (ngeom > 0) {
+        std::memcpy(scnstate_.scratch.geoms + scnstate_.scratch.ngeom,
+                    user_scn->geoms,
+                    sizeof(mjvGeom) * ngeom);
+        scnstate_.scratch.ngeom += ngeom;
+      }
+    }
+
     mjopt_prev_ = scnstate_.model.opt;
     warn_vgeomfull_prev_ = scnstate_.data.warning[mjWARN_VGEOMFULL].number;
   }
@@ -2172,9 +2188,8 @@ void Simulate::LoadOnRenderThread() {
   }
 
   // re-create scene and context
-  if (!this->is_passive_) {
-    mjv_makeScene(this->m_, &this->scn, kMaxGeom);
-  } else {
+  mjv_makeScene(this->m_, &this->scn, kMaxGeom);
+  if (this->is_passive_) {
     mjopt_prev_ = m_->opt;
     opt_prev_ = opt;
     cam_prev_ = cam;

simulate/simulate.h

@@ -51,8 +51,7 @@ class Simulate {
   // create object and initialize the simulate ui
   Simulate(
       std::unique_ptr<PlatformUIAdapter> platform_ui_adapter,
-      mjvScene* scn, mjvCamera* cam,
-      mjvOption* opt, mjvPerturb* pert, bool is_passive);
+      mjvCamera* cam, mjvOption* opt, mjvPerturb* pert, bool is_passive);
 
   // Synchronize mjModel and mjData state with UI inputs, and update
   // visualization.
@@ -233,7 +232,7 @@ class Simulate {
   int camera = 0;
 
   // abstract visualization
-  mjvScene& scn;
+  mjvScene scn;
   mjvCamera& cam;
   mjvOption& opt;
   mjvPerturb& pert;
@@ -243,6 +242,9 @@ class Simulate {
   mjvFigure figsize = {};
   mjvFigure figsensor = {};
 
+  // additional user-defined visualization geoms (used in passive mode)
+  mjvScene* user_scn = nullptr;
+
   // OpenGL rendering and UI
   int refresh_rate = 60;
   int window_pos[2] = {0};