Add option to only sync the state to improve sync performance in passive mode.

By default this option is off, which preserves the current behavior of syncing all of `mjModel` and `mjData`.

PiperOrigin-RevId: 776498017
Change-Id: I02127c9397efbabebb89b0e0d139b00b6e665d35

doc/changelog.rst

@@ -15,6 +15,8 @@ General
 ^^^^^^^
 - Added support for setting the initial camera in the viewer using
   :ref:`visual/global/cameraid<visual-global-cameraid>`.
+- Added support to only sync the state in the Python :ref:`passive viewer<PyViewerPassive>`'s ``Sync`` method, this is
+  useful to improve performance. The default behavior is unchanged and copies the entire model and data.
 
 Bug fixes
 ^^^^^^^^^

doc/python.rst

@@ -50,25 +50,17 @@ Interactive viewer
 
 An interactive GUI viewer is provided as part of the Python package in the ``mujoco.viewer`` module. It is based on the
 same codebase as the :ref:`simulate<saSimulate>` application that ships with the MuJoCo binary releases. Three distinct
-use cases are supported:
+use cases are supported: :ref:`managed viewer<PyViewerManaged>`, :ref:`standalone app<PyViewerApp>`, and :ref:`passive
-
+viewer<PyViewerPassive>`.
-.. _PyViewerApp:
-
-Standalone app
---------------
-
-- ``python -m mujoco.viewer`` launches an empty visualization session, where a model can be loaded by drag-and-drop.
-- ``python -m mujoco.viewer --mjcf=/path/to/some/mjcf.xml`` launches a visualization session for the specified
-  model file.
 
 .. _PyViewerManaged:
 
 Managed viewer
 --------------
 
-Called from a Python program/script, through the function ``viewer.launch``. This function *blocks user code* to
+The ``viewer.launch`` function launches the interactive viewer and *blocks user code* which is useful to support precise
-support precise timing of the physics loop. This mode should be used if user code is implemented as
+timing of the physics loop. This mode should be used if user code is implemented as :ref:`engine
-:ref:`engine plugins<exPlugin>` or :ref:`physics callbacks<glPhysics>`, and is called by MuJoCo during :ref:`mj_step`.
+plugins<exPlugin>` or :ref:`physics callbacks<glPhysics>`, and is called by MuJoCo during :ref:`mj_step`.
 
 - ``viewer.launch()`` launches an empty visualization session, where a model can be loaded by drag-and-drop.
 - ``viewer.launch(model)`` launches a visualization session for the given ``mjModel`` where the visualizer
@@ -76,14 +68,26 @@ support precise timing of the physics loop. This mode should be used if user cod
 - ``viewer.launch(model, data)`` is the same as above, except that the visualizer operates directly on the given
   ``mjData`` instance -- upon exit the ``data`` object will have been modified.
 
+.. _PyViewerApp:
+
+Standalone app
+--------------
+
+The ``mujoco.viewer`` Python package uses the ``if __name__ == '__main__'`` mechanism to allow the :ref:`managed
+viewer<PyViewerManaged>` to be called directly from the command line as a standalone app:
+
+- ``python -m mujoco.viewer`` launches an empty visualization session, where a model can be loaded by drag-and-drop.
+- ``python -m mujoco.viewer --mjcf=/path/to/some/mjcf.xml`` launches a visualization session for the specified
+  model file.
+
 .. _PyViewerPassive:
 
 Passive viewer
 --------------
 
-By calling ``viewer.launch_passive(model, data)``. This function *does not block*, allowing user code to continue
+The ``viewer.launch_passive`` function launches the interactive viewer in a way which *does not block*, allowing user
-execution. In this mode, the user's script is responsible for timing and advancing the physics state, and mouse-drag
+code to continue execution. In this mode, the user's script is responsible for timing and advancing the physics state,
-perturbations will not work unless the user explicitly synchronizes incoming events.
+and mouse-drag perturbations will not work unless the user explicitly synchronizes incoming events.
 
 .. warning::
   On MacOS, ``launch_passive`` requires that the user script is executed via a special ``mjpython`` launcher.
@@ -103,10 +107,14 @@ attributes:
   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(state_only=False)``: synchronizes between the user's ``mjModel``, ``mjData`` and the GUI. In order to allow
-  ``sync``. In order to allow user scripts to make arbitrary modifications to ``mjModel`` and ``mjData`` without
+  user scripts to make arbitrary modifications to ``mjModel`` and ``mjData`` without needing to hold the viewer lock,
-  needing to hold the viewer lock, the passive viewer does not access or modify these structs outside of ``sync``
+  the passive viewer does not access or modify these structs outside of ``sync`` calls. If the ``state_only`` argument
-  calls.
+  is ``True``, instead of syncing everything, only the ``mjData`` fields corresponding to
+  :ref:`mjSTATE_INTEGRATION<mjtState>` are synced, followed by a call to :ref:`mj_forward`. The latter option is much
+  faster, but would not pick up arbitrary changes as in the default case. Changes made via the GUI are picked up in
+  either case but changing e.g., ``mjModel.geom_rgba`` via code will be picked up when ``state_only=False`` but not when
+  ``state_only=True``.
 
   User scripts must call ``sync`` in order for the viewer to reflect physics state changes. The ``sync`` function
   also transfers user inputs from the GUI back into ``mjOption`` (inside ``mjModel``) and ``mjData``, including
@@ -1060,5 +1068,5 @@ non-exhaustive list of specific mujoco-py features:
    This is the one context in which the MuJoCo library (and therefore also ``mujoco``) is stateful: it holds a copy in
    memory of the last XML that was compiled, which is used in :ref:`mujoco.mj_saveLastXML(fname) <mj_saveLastXML>`. Note
    that mujoco-py’s implementation has a convenient extra feature, whereby the pose (as determined by ``sim.data``’s
-   state) is transformed to a keyframe that’s added to the model before saving.  This extra feature is not currently
+   state) is transformed to a keyframe that’s added to the model before saving. This extra feature is not currently
    available in ``mujoco``.

python/mujoco/simulate.cc

@@ -336,6 +336,7 @@ PYBIND11_MODULE(_simulate, pymodule) {
            CallIfNotNull(&mujoco::Simulate::LoadMessageClear),
            py::call_guard<py::gil_scoped_release>())
       .def("sync", CallIfNotNull(&mujoco::Simulate::Sync),
+           py::arg("state_only") = false,
            py::call_guard<py::gil_scoped_release>())
       .def("add_to_history", CallIfNotNull(&mujoco::Simulate::AddToHistory),
            py::call_guard<py::gil_scoped_release>())

python/mujoco/viewer.py

@@ -230,10 +230,10 @@ class Handle:
       return sim.lock()
     return contextlib.nullcontext()
 
-  def sync(self):
+  def sync(self, state_only: bool = False):
     sim = self._get_sim()
     if sim is not None:
-      sim.sync()  # locks internally
+      sim.sync(state_only)  # locks internally
 
   def update_hfield(self, hfieldid: int):
     sim = self._get_sim()

simulate/simulate.cc

@@ -1912,7 +1912,7 @@ Simulate::Simulate(std::unique_ptr<PlatformUIAdapter> platform_ui,
 //------------------------- Synchronize render and physics threads ---------------------------------
 
 // operations which require holding the mutex, prevents racing with physics thread
-void Simulate::Sync() {
+void Simulate::Sync(bool state_only) {
   MutexLock lock(this->mtx);
 
   if (!m_) {
@@ -2159,8 +2159,17 @@ void Simulate::Sync() {
   if (!is_passive_) {
     mjv_updateScene(m_, d_, &this->opt, &this->pert, &this->cam, mjCAT_ALL, &this->scn);
   } else {
-    mjv_copyModel(m_passive_, m_);
+    if (state_only) {
-    mjv_copyData(d_passive_, m_passive_, d_);
+      int state_size = mj_stateSize(m_, mjSTATE_INTEGRATION);
+      mjtNum* state = new mjtNum[state_size];
+      mj_getState(m_, d_, state, mjSTATE_INTEGRATION);
+      mj_setState(m_passive_, d_passive_, state, mjSTATE_INTEGRATION);
+      mj_forward(m_passive_, d_passive_);
+      delete[] state;
+    } else {
+      mjv_copyModel(m_passive_, m_);
+      mjv_copyData(d_passive_, m_passive_, d_);
+    }
 
     // append geoms from user_scn to scratch space
     if (user_scn) {

simulate/simulate.h

@@ -54,9 +54,10 @@ class Simulate {
       std::unique_ptr<PlatformUIAdapter> platform_ui_adapter,
       mjvCamera* cam, mjvOption* opt, mjvPerturb* pert, bool is_passive);
 
-  // Synchronize mjModel and mjData state with UI inputs, and update
+  // Synchronize state with UI inputs, and update visualization.  If state_only
-  // visualization.
+  // is false mjData and mjModel will be updated, otherwise only the subset of
-  void Sync();
+  // mjData corresponding to mjSTATE_INTEGRATION will be synced.
+  void Sync(bool state_only = false);
 
   void UpdateHField(int hfieldid);
   void UpdateMesh(int meshid);