diff --git a/doc/changelog.rst b/doc/changelog.rst index 798cfde7..e0255298 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -131,6 +131,8 @@ Python bindings 23. Fixed `#870 `__ where calling ``update_scene`` with an invalid camera name used the default camera. +24. Added ``user_scn`` to the :ref:`passive viewer` handle, which allows users to add custom + visualization geoms (`#1023 `__). Simulate ^^^^^^^^ @@ -139,18 +141,18 @@ Simulate :align: right :width: 240px -24. Added **state history** mechanism to :ref:`simulate` and the managed +25. Added **state history** mechanism to :ref:`simulate` and the managed :ref:`Python viewer`. 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 `__ by `Levi Burner `__. Bug fixes ^^^^^^^^^ -26. Fixed a bug that was causing :ref:`geom margin` to be ignored during the construction of +27. Fixed a bug that was causing :ref:`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. diff --git a/doc/includes/references.h b/doc/includes/references.h index dd9a261f..09e29799 100644 --- a/doc/includes/references.h +++ b/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 { diff --git a/doc/python.rst b/doc/python.rst index cc1b2c39..8bb75072 100644 --- a/doc/python.rst +++ b/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`, - :ref:`mjvOption`, and :ref:`mjvPerturb` structs, respectively. +- ``cam``, ``opt``, and ``pert`` properties: correspond to :ref:`mjvCamera`, :ref:`mjvOption`, and :ref:`mjvPerturb` + 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``, - ``cam``, ``opt``, and ``pert`` properties of the viewer handle. + state. These include the ``mjModel`` and ``mjData`` instance passed to ``launch_passive``, and also the ``cam``, + ``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 diff --git a/include/mujoco/mjvisualize.h b/include/mujoco/mjvisualize.h index 2433db5a..ac1c2aed 100644 --- a/include/mujoco/mjvisualize.h +++ b/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 { diff --git a/introspect/structs.py b/introspect/structs.py index de0edd9b..787bca1a 100644 --- a/introspect/structs.py +++ b/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', diff --git a/python/mujoco/simulate.cc b/python/mujoco/simulate.cc index 0f24026e..820c73b8 100644 --- a/python/mujoco/simulate.cc +++ b/python/mujoco/simulate.cc @@ -69,19 +69,23 @@ class UIAdapterWithPyCallback : public Adapter { class SimulateWrapper { public: SimulateWrapper(std::unique_ptr 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().get(), + std::move(platform_ui_adapter), cam.cast().get(), opt.cast().get(), pert.cast().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().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()) .def("load", &SimulateWrapper::Load) - .def("load_message_clear", CallIfNotNull(&mujoco::Simulate::LoadMessageClear), + .def("load_message_clear", + CallIfNotNull(&mujoco::Simulate::LoadMessageClear), py::call_guard()) .def("sync", CallIfNotNull(&mujoco::Simulate::Sync), py::call_guard()) diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc index df00498e..85af8ad2 100644 --- a/python/mujoco/structs.cc +++ b/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), diff --git a/python/mujoco/viewer.py b/python/mujoco/viewer.py index 2a0be3b9..6492d41c 100644 --- a/python/mujoco/viewer.py +++ b/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( diff --git a/simulate/main.cc b/simulate/main.cc index 8ce2c4ae..57763646 100644 --- a/simulate/main.cc +++ b/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( std::make_unique(), - &scn, &cam, &opt, &pert, /* is_passive = */ false + &cam, &opt, &pert, /* is_passive = */ false ); const char* filename = nullptr; diff --git a/simulate/simulate.cc b/simulate/simulate.cc index d25362e3..2a4d98d5 100644 --- a/simulate/simulate.cc +++ b/simulate/simulate.cc @@ -1737,16 +1737,15 @@ namespace mujoco { namespace mju = ::mujoco::sample_util; Simulate::Simulate(std::unique_ptr 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; diff --git a/simulate/simulate.h b/simulate/simulate.h index 286002b8..d50a6677 100644 --- a/simulate/simulate.h +++ b/simulate/simulate.h @@ -51,8 +51,7 @@ class Simulate { // create object and initialize the simulate ui Simulate( std::unique_ptr 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}; diff --git a/src/engine/engine_vis_state.c b/src/engine/engine_vis_state.c index f2d0eb25..90fb58ea 100644 --- a/src/engine/engine_vis_state.c +++ b/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); } } } diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 2debba82..7469600f 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/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 {}