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

Fixes #1023, #1082

PiperOrigin-RevId: 572303635
Change-Id: Ia8399f816f311b38a7b137f5bd4f359f77a1da1c

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 pert, bool is_passive)
+                  py::object cam, py::object opt,
+                  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
-
-  @property
-  def scn(self):
-    return self._scn
+    self._user_scn = user_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 = (
-        lambda: handle_return.put_nowait(Handle(simulate, scn, cam, opt, pert)))
+    notify_loaded = lambda: handle_return.put_nowait(
+        Handle(simulate, cam, opt, pert, user_scn)
+    )
 
   if run_physics_thread:
     side_thread = threading.Thread(