From 09f7154e574cef81a03137ab1cacee171a7f5ad7 Mon Sep 17 00:00:00 2001 From: Matija Kecman Date: Fri, 27 Jun 2025 02:48:58 -0700 Subject: [PATCH] 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 | 2 ++ doc/python.rst | 50 +++++++++++++++++++++++---------------- python/mujoco/simulate.cc | 1 + python/mujoco/viewer.py | 4 ++-- simulate/simulate.cc | 15 +++++++++--- simulate/simulate.h | 7 +++--- 6 files changed, 50 insertions(+), 29 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index ecc4a0da..2f8e4375 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -15,6 +15,8 @@ General ^^^^^^^ - Added support for setting the initial camera in the viewer using :ref:`visual/global/cameraid`. +- Added support to only sync the state in the Python :ref:`passive viewer`'s ``Sync`` method, this is + useful to improve performance. The default behavior is unchanged and copies the entire model and data. Bug fixes ^^^^^^^^^ diff --git a/doc/python.rst b/doc/python.rst index c6490edc..2d22de9d 100644 --- a/doc/python.rst +++ b/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` application that ships with the MuJoCo binary releases. Three distinct -use cases are supported: - -.. _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. +use cases are supported: :ref:`managed viewer`, :ref:`standalone app`, and :ref:`passive +viewer`. .. _PyViewerManaged: Managed viewer -------------- -Called from a Python program/script, through the function ``viewer.launch``. This function *blocks user code* to -support precise timing of the physics loop. This mode should be used if user code is implemented as -:ref:`engine plugins` or :ref:`physics callbacks`, and is called by MuJoCo during :ref:`mj_step`. +The ``viewer.launch`` function launches the interactive viewer and *blocks user code* which is useful to support precise +timing of the physics loop. This mode should be used if user code is implemented as :ref:`engine +plugins` or :ref:`physics callbacks`, 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` 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 -execution. In this mode, the user's script is responsible for timing and advancing the physics state, and mouse-drag -perturbations will not work unless the user explicitly synchronizes incoming events. +The ``viewer.launch_passive`` function launches the interactive viewer in a way which *does not block*, allowing user +code to continue execution. In this mode, the user's script is responsible for timing and advancing the physics state, +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``. In order to allow user scripts to make arbitrary modifications to ``mjModel`` and ``mjData`` without - needing to hold the viewer lock, the passive viewer does not access or modify these structs outside of ``sync`` - calls. +- ``sync(state_only=False)``: synchronizes between the user's ``mjModel``, ``mjData`` and the GUI. In order to allow + user scripts to make arbitrary modifications to ``mjModel`` and ``mjData`` without needing to hold the viewer lock, + the passive viewer does not access or modify these structs outside of ``sync`` calls. If the ``state_only`` argument + is ``True``, instead of syncing everything, only the ``mjData`` fields corresponding to + :ref:`mjSTATE_INTEGRATION` 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) `. 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``. diff --git a/python/mujoco/simulate.cc b/python/mujoco/simulate.cc index a62d1b6b..f44bd9bd 100644 --- a/python/mujoco/simulate.cc +++ b/python/mujoco/simulate.cc @@ -336,6 +336,7 @@ PYBIND11_MODULE(_simulate, pymodule) { CallIfNotNull(&mujoco::Simulate::LoadMessageClear), py::call_guard()) .def("sync", CallIfNotNull(&mujoco::Simulate::Sync), + py::arg("state_only") = false, py::call_guard()) .def("add_to_history", CallIfNotNull(&mujoco::Simulate::AddToHistory), py::call_guard()) diff --git a/python/mujoco/viewer.py b/python/mujoco/viewer.py index c1de94ec..7140773d 100644 --- a/python/mujoco/viewer.py +++ b/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() diff --git a/simulate/simulate.cc b/simulate/simulate.cc index b692302e..17ac9e82 100644 --- a/simulate/simulate.cc +++ b/simulate/simulate.cc @@ -1912,7 +1912,7 @@ Simulate::Simulate(std::unique_ptr 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_); - mjv_copyData(d_passive_, m_passive_, d_); + if (state_only) { + 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) { diff --git a/simulate/simulate.h b/simulate/simulate.h index 03c0a0b1..a949615e 100644 --- a/simulate/simulate.h +++ b/simulate/simulate.h @@ -54,9 +54,10 @@ class Simulate { std::unique_ptr platform_ui_adapter, mjvCamera* cam, mjvOption* opt, mjvPerturb* pert, bool is_passive); - // Synchronize mjModel and mjData state with UI inputs, and update - // visualization. - void Sync(); + // Synchronize state with UI inputs, and update visualization. If state_only + // is false mjData and mjModel will be updated, otherwise only the subset of + // mjData corresponding to mjSTATE_INTEGRATION will be synced. + void Sync(bool state_only = false); void UpdateHField(int hfieldid); void UpdateMesh(int meshid);