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

Fixes #1023, #1082

PiperOrigin-RevId: 572303635
Change-Id: Ia8399f816f311b38a7b137f5bd4f359f77a1da1c

doc/changelog.rst

@@ -131,6 +131,8 @@ Python bindings
 
 23. Fixed `#870 <https://github.com/google-deepmind/mujoco/issues/870>`__ where calling ``update_scene`` with an invalid
     camera name used the default camera.
+24. Added ``user_scn`` to the :ref:`passive viewer<PyViewerPassive>` handle, which allows users to add custom
+    visualization geoms (`#1023 <https://github.com/google-deepmind/mujoco/issues/870>`__).
 
 Simulate
 ^^^^^^^^
@@ -139,18 +141,18 @@ Simulate
    :align: right
    :width: 240px
 
-24. Added **state history** mechanism to :ref:`simulate<saSimulate>` and the managed
+25. Added **state history** mechanism to :ref:`simulate<saSimulate>` and the managed
     :ref:`Python viewer<PyViewerManaged>`. State history can be viewed by scrubbing the History slider and (more
     precisely) with the left and right arrow keys. See screen capture:
 
-25. The ``LOADING...`` label is now shown correctly.
+26. The ``LOADING...`` label is now shown correctly.
     `Contribution <https://github.com/google-deepmind/mujoco/pull/1070>`__ by
     `Levi Burner <https://github.com/aftersomemath>`__.
 
 Bug fixes
 ^^^^^^^^^
 
-26. Fixed a bug that was causing :ref:`geom margin<body-geom-margin>` to be ignored during the construction of
+27. Fixed a bug that was causing :ref:`geom margin<body-geom-margin>` to be ignored during the construction of
     midphase collision trees.
 
 27. Fixed a bug that was generating incorrect values in ``efc_diagApprox`` for weld equality constraints.

doc/includes/references.h

@@ -1989,7 +1989,7 @@ struct mjvSceneState_ {
   int nbuffer;                     // size of the buffer in bytes
   void* buffer;                    // heap-allocated memory for all arrays in this struct
   int maxgeom;                     // maximum number of mjvGeom supported by this state object
-  mjvScene plugincache;            // scratch space for vis geoms inserted by plugins
+  mjvScene scratch;                // scratch space for vis geoms inserted by the user and plugins
 
   // fields in mjModel that are necessary to re-render a scene
   struct {

doc/python.rst

@@ -96,13 +96,13 @@ perturbations will not work unless the user explicitly synchronizes incoming eve
 The ``launch_passive`` function returns a handle which can be used to interact with the viewer. It has the following
 attributes:
 
-- ``scn``, ``cam``, ``opt``, and ``pert`` properties: correspond to :ref:`mjvScene`, :ref:`mjvCamera`,
+- ``cam``, ``opt``, and ``pert`` properties: correspond to :ref:`mjvCamera`, :ref:`mjvOption`, and :ref:`mjvPerturb`
-  :ref:`mjvOption`, and :ref:`mjvPerturb` structs, respectively.
+  structs, respectively.
 
 - ``lock()``: provides a mutex lock for the viewer as a context manager. Since the viewer operates its own
   thread, user code must ensure that it is holding the viewer lock before modifying any physics or visualization
-  state. These include the ``mjModel`` and ``mjData`` instance passed to ``launch_passive``, and also the ``scn``,
+  state. These include the ``mjModel`` and ``mjData`` instance passed to ``launch_passive``, and also the ``cam``,
-  ``cam``, ``opt``, and ``pert`` properties of the viewer handle.
+  ``opt``, and ``pert`` properties of the viewer handle.
 
 - ``sync()``: synchronizes state between ``mjModel``, ``mjData``, and GUI user inputs since the previous call to
   ``sync``. In order to allow user scripts to make arbitrary modifications to ``mjModel`` and ``mjData`` without
@@ -124,6 +124,37 @@ attributes:
 - ``is_running()``: returns ``True`` if the viewer window is running and ``False`` if it is closed.
   This method can be safely called without locking.
 
+- ``user_scn``: an :ref:`mjvScene` object that allows users to add custom visualization geoms to the rendered scene.
+  This is separate from the ``mjvScene`` that the viewer uses internally to render the final scene, and is entirely
+  under the user's control. User scripts can call e.g. :ref:`mjv_initGeom` or :ref:`mjv_makeConnector` to add
+  visualization geoms to ``user_scn``, and upon the next call to ``sync()``, the viewer will incorporate
+  these geoms to future rendered images. For example:
+
+  .. code-block:: python
+
+    with mujoco.viewer.launch_passive(m, d, key_callback=key_callback) as viewer:
+      while viewer.is_running():
+        ...
+        # Step the physics.
+        mujoco.mj_step(m, d)
+
+        # Add a 3x3x3 grid of variously colored spheres to the middle of the scene.
+        viewer.user_scn.ngeom = 0
+        i = 0
+        for x, y, z in itertools.product(*((range(-1, 2),) * 3)):
+          mujoco.mjv_initGeom(
+              viewer.user_scn.geoms[i],
+              type=mujoco.mjtGeom.mjGEOM_SPHERE,
+              size=[0.02, 0, 0],
+              pos=0.1*np.array([x, y, z]),
+              mat=np.eye(3).flatten(),
+              rgba=0.5*np.array([x + 1, y + 1, z + 1, 2])
+          )
+          i += 1
+        viewer.user_scn.ngeom = i
+        viewer.sync()
+        ...
+
 The viewer handle can also be used as a context manager which calls ``close()`` automatically upon exit. A minimal
 example of a user script that uses ``launch_passive`` might look like the following. (Note that example is a simple
 illustrative example that does **not** necessarily keep the physics ticking at the correct wallclock rate.)
@@ -183,7 +214,6 @@ pause or resume the run loop when the spacebar is pressed.
         viewer.sync()
       ...
 
-
 .. _PyUsage:
 
 Basic usage

include/mujoco/mjvisualize.h

@@ -375,7 +375,7 @@ struct mjvSceneState_ {
   int nbuffer;                     // size of the buffer in bytes
   void* buffer;                    // heap-allocated memory for all arrays in this struct
   int maxgeom;                     // maximum number of mjvGeom supported by this state object
-  mjvScene plugincache;            // scratch space for vis geoms inserted by plugins
+  mjvScene scratch;                // scratch space for vis geoms inserted by the user and plugins
 
   // fields in mjModel that are necessary to re-render a scene
   struct {

introspect/structs.py

@@ -5476,9 +5476,9 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  doc='maximum number of mjvGeom supported by this state object',
              ),
              StructFieldDecl(
-                 name='plugincache',
+                 name='scratch',
                  type=ValueType(name='mjvScene'),
-                 doc='scratch space for vis geoms inserted by plugins',
+                 doc='scratch space for vis geoms inserted by the user and plugins',  # pylint: disable=line-too-long
              ),
              StructFieldDecl(
                  name='model',

python/mujoco/simulate.cc

@@ -69,19 +69,23 @@ class UIAdapterWithPyCallback : public Adapter {
 class SimulateWrapper {
  public:
   SimulateWrapper(std::unique_ptr<PlatformUIAdapter> platform_ui_adapter,
-                  py::object scn, py::object cam, py::object opt,
+                  py::object cam, py::object opt,
-                  py::object pert, bool is_passive)
+                  py::object pert, py::object user_scn, bool is_passive)
       : simulate_(new mujoco::Simulate(
-            std::move(platform_ui_adapter), scn.cast<MjvSceneWrapper&>().get(),
+            std::move(platform_ui_adapter),
             cam.cast<MjvCameraWrapper&>().get(),
             opt.cast<MjvOptionWrapper&>().get(),
             pert.cast<MjvPerturbWrapper&>().get(), is_passive)),
         m_(py::none()),
         d_(py::none()),
-        scn_(scn),
         cam_(cam),
         opt_(opt),
-        pert_(pert) {}
+        pert_(pert),
+        user_scn_(user_scn) {
+    if (!user_scn.is_none()) {
+      simulate_->user_scn = user_scn_.cast<MjvSceneWrapper&>().get();
+    }
+  }
 
   ~SimulateWrapper() { Destroy(); }
 
@@ -127,10 +131,10 @@ class SimulateWrapper {
   // simulate object.
   py::object m_;
   py::object d_;
-  py::object scn_;
   py::object cam_;
   py::object opt_;
   py::object pert_;
+  py::object user_scn_;
 
   mjModel* m_raw_ = nullptr;
   mjData* d_raw_ = nullptr;
@@ -207,7 +211,8 @@ PYBIND11_MODULE(_simulate, pymodule) {
       .def("load_message", CallIfNotNull(&mujoco::Simulate::LoadMessage),
            py::call_guard<py::gil_scoped_release>())
       .def("load", &SimulateWrapper::Load)
-      .def("load_message_clear", CallIfNotNull(&mujoco::Simulate::LoadMessageClear),
+      .def("load_message_clear",
+           CallIfNotNull(&mujoco::Simulate::LoadMessageClear),
            py::call_guard<py::gil_scoped_release>())
       .def("sync", CallIfNotNull(&mujoco::Simulate::Sync),
            py::call_guard<py::gil_scoped_release>())

python/mujoco/structs.cc

@@ -1050,7 +1050,9 @@ MjvGeomWrapper::MjWrapper()
         static_assert(sizeof(ptr_->mat) == sizeof(ptr_->mat[0])*9);
         return InitPyArray(std::array{3, 3}, ptr_->mat, owner_);
       }()),
-      X(rgba) {}
+      X(rgba) {
+  mjv_initGeom(ptr_, mjGEOM_NONE, nullptr, nullptr, nullptr, nullptr);
+}
 
 MjvGeomWrapper::MjWrapper(raw::MjvGeom* ptr, py::handle owner)
     : WrapperBase(ptr, owner),

python/mujoco/viewer.py

@@ -67,20 +67,16 @@ class Handle:
   def __init__(
       self,
       sim: _Simulate,
-      scn: mujoco.MjvScene,
       cam: mujoco.MjvCamera,
       opt: mujoco.MjvOption,
       pert: mujoco.MjvPerturb,
+      user_scn: Optional[mujoco.MjvScene],
   ):
     self._sim = weakref.ref(sim)
-    self._scn = scn
     self._cam = cam
     self._opt = opt
     self._pert = pert
-
+    self._user_scn = user_scn
-  @property
-  def scn(self):
-    return self._scn
 
   @property
   def cam(self):
@@ -94,6 +90,10 @@ class Handle:
   def perturb(self):
     return self._pert
 
+  @property
+  def user_scn(self):
+    return self._user_scn
+
   def close(self):
     sim = self._sim()
     if sim is not None:
@@ -340,14 +340,16 @@ def _launch_internal(
 
     loader = _loader
 
-  if model and not run_physics_thread:
-    scn = mujoco.MjvScene(model, _Simulate.MAX_GEOM)
-  else:
-    scn = mujoco.MjvScene()
   cam = mujoco.MjvCamera()
   opt = mujoco.MjvOption()
   pert = mujoco.MjvPerturb()
-  simulate = _Simulate(scn, cam, opt, pert, run_physics_thread, key_callback)
+  if model and not run_physics_thread:
+    user_scn = mujoco.MjvScene(model, _Simulate.MAX_GEOM)
+  else:
+    user_scn = None
+  simulate = _Simulate(
+      cam, opt, pert, user_scn, run_physics_thread, key_callback
+  )
 
   # Initialize GLFW if not using mjpython.
   if _MJPYTHON is None:
@@ -357,8 +359,9 @@ def _launch_internal(
 
   notify_loaded = None
   if handle_return:
-    notify_loaded = (
+    notify_loaded = lambda: handle_return.put_nowait(
-        lambda: handle_return.put_nowait(Handle(simulate, scn, cam, opt, pert)))
+        Handle(simulate, cam, opt, pert, user_scn)
+    )
 
   if run_physics_thread:
     side_thread = threading.Thread(

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,
+                   mjvCamera* cam, mjvOption* opt, mjvPerturb* pert,
-                   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);
-    mjv_makeScene(this->m_, &this->scn, kMaxGeom);
+  if (this->is_passive_) {
-  } else {
     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,
+      mjvCamera* cam, mjvOption* opt, mjvPerturb* pert, bool is_passive);
-      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};

src/engine/engine_vis_state.c

@@ -48,19 +48,19 @@ static inline size_t roundUpToCacheLine(size_t n) {
 // set default scene
 void mjv_defaultSceneState(mjvSceneState* scnstate) {
   memset(scnstate, 0, sizeof(mjvSceneState));
-  mjv_defaultScene(&scnstate->plugincache);
+  mjv_defaultScene(&scnstate->scratch);
 }
 
 
 
 // allocate and init scene state
 void mjv_makeSceneState(const mjModel* m, const mjData* d, mjvSceneState* scnstate, int maxgeom) {
-  mjv_freeScene(&scnstate->plugincache);
+  mjv_freeScene(&scnstate->scratch);
   mju_free(scnstate->buffer);
 
 #ifdef MEMORY_SANITIZER
   __msan_allocated_memory(scnstate, sizeof(mjvSceneState));
-  mjv_defaultScene(&scnstate->plugincache);
+  mjv_defaultScene(&scnstate->scratch);
 #endif
 
   scnstate->nbuffer = 0;
@@ -144,14 +144,14 @@ void mjv_makeSceneState(const mjModel* m, const mjData* d, mjvSceneState* scnsta
     mjERROR("mjvSceneState buffer is not fully used");
   }
 
-  mjv_makeScene(m, &scnstate->plugincache, maxgeom);
+  mjv_makeScene(m, &scnstate->scratch, maxgeom);
 }
 
 
 
 // free scene state
 void mjv_freeSceneState(mjvSceneState* scnstate) {
-  mjv_freeScene(&scnstate->plugincache);
+  mjv_freeScene(&scnstate->scratch);
   mju_free(scnstate->buffer);
   mjv_defaultSceneState(scnstate);
 }
@@ -229,14 +229,14 @@ int mjv_updateSceneFromState(const mjvSceneState* scnstate, const mjvOption* opt
   int warning_start = d.warning[mjWARN_VGEOMFULL].number;
 
   // copy mjvGeoms added by plugins
-  int nplugingeom = scnstate->plugincache.ngeom;
+  int nplugingeom = scnstate->scratch.ngeom;
   if (nplugingeom > scn->maxgeom) {
     mj_warning(&d, mjWARN_VGEOMFULL, scn->maxgeom);
     scn->ngeom = scn->maxgeom;
   } else {
     scn->ngeom = nplugingeom;
   }
-  memcpy(scn->geoms, scnstate->plugincache.geoms, sizeof(mjvGeom) * scn->ngeom);
+  memcpy(scn->geoms, scnstate->scratch.geoms, sizeof(mjvGeom) * scn->ngeom);
 
   // add all categories
   mjv_addGeoms(&m, &d, opt, pert, catmask, scn);
@@ -272,7 +272,7 @@ void mjv_updateSceneState(const mjModel* m, mjData* d, const mjvOption* opt,
 #undef X
 
   // Update plugin visualization cache.
-  scnstate->plugincache.ngeom = 0;
+  scnstate->scratch.ngeom = 0;
   if (m->nplugin) {
     const int nslot = mjp_pluginCount();
     // iterate over plugins, call visualize if defined
@@ -283,7 +283,7 @@ void mjv_updateSceneState(const mjModel* m, mjData* d, const mjvOption* opt,
         mjERROR("invalid plugin slot: %d", slot);
       }
       if (plugin->visualize) {
-        plugin->visualize(m, d, opt, &scnstate->plugincache, i);
+        plugin->visualize(m, d, opt, &scnstate->scratch, i);
       }
     }
   }

unity/Runtime/Bindings/MjBindings.cs

@@ -6148,7 +6148,7 @@ public unsafe struct mjvSceneState_ {
   public int nbuffer;
   public void* buffer;
   public int maxgeom;
-  public mjvScene_ plugincache;
+  public mjvScene_ scratch;
   public model model;
   public data data;
 }public struct mjuiItem_ {}public struct mjfItemEnable {}