Merge branch 'google-deepmind:main' into usd-fixes

doc/APIreference/APItypes.rst

@@ -1393,6 +1393,8 @@ Alternative orientation specifiers.
 
 .. _ArrayHandles:
 
+.. _mjBuffer:
+
 .. _mjString:
 
 .. _mjStringVec:
@@ -1417,6 +1419,7 @@ C handles for C++ strings and vector types. When using from C, use the provided
 
    #ifdef __cplusplus
      // C++: defined to be compatible with corresponding std types
+     using mjBuffer      = std::vector<std::byte>;
      using mjString      = std::string;
      using mjStringVec   = std::vector<std::string>;
      using mjIntVec      = std::vector<int>;
@@ -1426,6 +1429,7 @@ C handles for C++ strings and vector types. When using from C, use the provided
      using mjDoubleVec   = std::vector<double>;
    #else
      // C: opaque types
+     typedef void mjBuffer;
      typedef void mjString;
      typedef void mjStringVec;
      typedef void mjIntVec;

doc/APIreference/functions.rst

@@ -2589,7 +2589,7 @@ mjd_transitionFD
 
 .. mujoco-include:: mjd_transitionFD
 
-Finite-differenced discrete-time transition matrices.
+Compute finite-differenced discrete-time transition matrices.
 
 Letting :math:`x, u` denote the current :ref:`state<gePhysicsState>` and :ref:`control<geInput>`
 vector in an mjData instance, and letting :math:`y, s` denote the next state and sensor
@@ -2610,12 +2610,26 @@ These matrices and their dimensions are:
 - All outputs are optional (can be NULL).
 - ``eps`` is the finite-differencing epsilon.
 - ``flg_centered`` denotes whether to use forward (0) or centered (1) differences.
-- Accuracy can be somewhat improved if solver :ref:`iterations<option-iterations>` are set to a
-  fixed (small) value and solver :ref:`tolerance<option-tolerance>` is set to 0. This insures that
-  all calls to the solver will perform exactly the same number of iterations.
+- The Runge-Kutta integrator (:ref:`mjINT_RK4<mjtIntegrator>`) is not supported.
 
-.. attention::
-   - The Runge-Kutta 4th-order integrator (``mjINT_RK4``) is not supported.
+.. admonition:: Improving speed and accuracy
+   :class: tip
+
+   warmstart
+     If warm-starts are not :ref:`disabled<option-flag-warmstart>`, the warm-start accelerations
+     ``mjData.qacc_warmstart`` which are present at call-time are loaded at the start of every relevant pipeline call,
+     to preserve determinism. If solver computations are an expensive part of the simulation, the following trick can
+     lead to significant speed-ups: First call :ref:`mj_forward` to let the solver converge, then reduce :ref:`solver
+     iterations<option-iterations>` significantly, then call :ref:`mjd_transitionFD`, finally, restore the original
+     value of :ref:`iterations<option-iterations>`. Because we are already near the solution, few iteration are required
+     to find the new minimum. This is especially true for the :ref:`Newton<option-solver>` solver, where the required
+     number of iteration for convergence near the minimum can be as low as 1.
+
+   tolerance
+      Accuracy can be improved if solver :ref:`tolerance<option-tolerance>` is set to 0. This means that all calls to
+      the solver will perform exactly the same number of iterations, preventing numerical errors due to early
+      termination. Of course, this means that :ref:`solver iterations<option-iterations>` should be small, to not tread
+      water at the minimum. This method and the one described above can and should be combined.
 
 .. _mjd_inverseFD:
 

doc/APIreference/functions_override.rst

@@ -567,7 +567,7 @@ outputs of derivative functions are the trailing rather than leading arguments.
 
 .. _mjd_transitionFD:
 
-Finite-differenced discrete-time transition matrices.
+Compute finite-differenced discrete-time transition matrices.
 
 Letting :math:`x, u` denote the current :ref:`state<gePhysicsState>` and :ref:`control<geInput>`
 vector in an mjData instance, and letting :math:`y, s` denote the next state and sensor
@@ -588,12 +588,27 @@ These matrices and their dimensions are:
 - All outputs are optional (can be NULL).
 - ``eps`` is the finite-differencing epsilon.
 - ``flg_centered`` denotes whether to use forward (0) or centered (1) differences.
-- Accuracy can be somewhat improved if solver :ref:`iterations<option-iterations>` are set to a
-  fixed (small) value and solver :ref:`tolerance<option-tolerance>` is set to 0. This insures that
-  all calls to the solver will perform exactly the same number of iterations.
+- The Runge-Kutta integrator (:ref:`mjINT_RK4<mjtIntegrator>`) is not supported.
+
+.. admonition:: Improving speed and accuracy
+   :class: tip
+
+   warmstart
+     If warm-starts are not :ref:`disabled<option-flag-warmstart>`, the warm-start accelerations
+     ``mjData.qacc_warmstart`` which are present at call-time are loaded at the start of every relevant pipeline call,
+     to preserve determinism. If solver computations are an expensive part of the simulation, the following trick can
+     lead to significant speed-ups: First call :ref:`mj_forward` to let the solver converge, then reduce :ref:`solver
+     iterations<option-iterations>` significantly, then call :ref:`mjd_transitionFD`, finally, restore the original
+     value of :ref:`iterations<option-iterations>`. Because we are already near the solution, few iteration are required
+     to find the new minimum. This is especially true for the :ref:`Newton<option-solver>` solver, where the required
+     number of iteration for convergence near the minimum can be as low as 1.
+
+   tolerance
+      Accuracy can be improved if solver :ref:`tolerance<option-tolerance>` is set to 0. This means that all calls to
+      the solver will perform exactly the same number of iterations, preventing numerical errors due to early
+      termination. Of course, this means that :ref:`solver iterations<option-iterations>` should be small, to not tread
+      water at the minimum. This method and the one described above can and should be combined.
 
-.. attention::
-   - The Runge-Kutta 4th-order integrator (``mjINT_RK4``) is not supported.
 
 .. _mjd_inverseFD:
 

doc/XMLreference.rst

@@ -175,6 +175,12 @@ replicating 200 times, suffixes will be ``000, 001, ...`` etc). All referencing
 and namespaced appropriately. Detailed examples of models using replicate can be found in the
 `model/replicate/ <https://github.com/google-deepmind/mujoco/tree/main/model/replicate>`__ directory.
 
+There is a caveat concerning :ref:`keyframes<keyframe>` when using replicate. Since :ref:`mjs_attachFrame` is used to
+self-attach multiple times the enclosed kinematic tree, if this tree contains further :ref:`attach<body-attach>`
+elements, keyframes will not be replicated nor namespaced by :ref:`replicate<replicate>`, but they will be attached and
+namespaced once by the innermost call of :ref:`mjs_attachFrame` or :ref:`mjs_attachBody`. See the limitations discussed
+in :ref:`attach<body-attach>`.
+
 .. _replicate-count:
 
 :at:`count`: :at-val:`int, required`
@@ -1106,7 +1112,8 @@ to convert to one of the other supported formats.
 
 .. _legacy-msh-docs:
 
-MSH file format
+.. collapse:: Legacy MSH file format
+
    The binary MSH file starts with 4 integers specifying the number of vertex positions (nvertex), vertex normals
    (nnormal), vertex texture coordinates (ntexcoord), and vertex indices making up the faces (nface), followed by the
    numeric data. nvertex must be at least 4. nnormal and ntexcoord can be zero (in which case the corresponding data is
@@ -1190,7 +1197,8 @@ The full list of processing steps applied by the compiler to each mesh is as fol
    normals. If sharp edges are encountered, the renderer uses the face normals to preserve the visual information about
    the edge, unless :ref:`smoothnormal<asset-mesh-smoothnormal>` is true.
    Note that normals cannot be provided with STL meshes;
-#. Scale, translate and rotate the vertices and normals, re-normalize the normals in case of scaling;
+#. Scale, translate and rotate the vertices and normals, re-normalize the normals in case of scaling. Save these
+   transformations in ``mjModel.mesh_{pos, quat, scale}``.
 #. Construct the convex hull if specified;
 #. Find the centroid of all triangle faces, and construct the union-of-pyramids representation. Triangles whose area is
    too small (below the :ref:`mjMINVAL <glNumeric>` value of 1E-14) result in compile error;
@@ -3806,7 +3814,7 @@ all attachments will appear in the saved XML file.
    - An entire model cannot be attached (i.e. including all elements, referenced or not).
    - All assets from the child model will be copied in, whether they are referenced or not.
    - Self-attach or circular references are not checked for and will lead to infinite loops.
-   - :ref:`Keyframes<keyframe>` are not yet supported. When attaching, all keyframes will be deleted.
+   - :ref:`Keyframes<keyframe>` are attached once, so they are not replicated in nested attachments.
 
 .. _body-attach-model:
 

doc/changelog.rst

@@ -7,8 +7,29 @@ Upcoming version (not yet released)
 
 General
 ^^^^^^^
-1. Add :ref:`mjSpec` option for creating a texture from a buffer.
+1. Added :ref:`mjSpec` option for creating a texture from a buffer.
 2. :ref:`shellinertia <body-geom-shellinertia>` is now supported by all geom types.
+3. Added support for :ref:`attaching<meAttachment>` keyframes.
+
+MJX
+^^^
+4. Added ``efc_pos`` to ``mjx.Data`` (:github:issue:`1388`).
+5. Added position-dependent sensors: ``MAGNETOMETER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``,
+   ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``, ``CLOCK``.
+6. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device.
+7. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``.
+
+Bug fixes
+^^^^^^^^^
+8. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`,
+   contribution by :github:user:`michael-ahn`).
+9. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit
+   integrators, wrong derivatives would be computed.
+
+Python bindings
+^^^^^^^^^^^^^^^
+10. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`).
+
 
 Version 3.2.2 (Aug 8, 2024)
 ---------------------------

doc/computation/index.rst

@@ -1722,20 +1722,23 @@ The top-level function :ref:`mj_inverse` invokes the following sequence of compu
 Derivatives
 -----------
 
-MuJoCo's entire computational pipline including its constraint solver are analytically differentiable. Writing
-efficient implementations of these derivatives is a long term goal of the development team. Analytic derivatives of the
-smooth dynamics (excluding constraints) with respect to velocity are already computed and enable the two
+MuJoCo's entire computational pipline including its constraint solver are analytically differentiable in principle.
+Writing efficient implementations of these derivatives is a long term goal of the development team. Analytic derivatives
+of the smooth dynamics (excluding constraints) with respect to velocity are already computed and enable the two
 :ref:`implicit integrators<geIntegration>`.
 
+Note that the default value of the :ref:`solver impedance<CSolverImpedance>` is such that contacts are *not*
+differentiable by default, and needs to be :ref:`set to 0<solimp0>` in order for contact-force onset to be smooth.
+
 Two functions are currently available which use efficient finite-differencing in order to compute dynamics Jacobians:
 
 :ref:`mjd_transitionFD`:
   Computes state-transition and control-transition Jacobians for the discrete-time forward dynamics (:ref:`mj_step`).
-  See :ref:`API documentation<mjd_transitionFD>`.
+  See :ref:`documentation<mjd_transitionFD>`.
 
 :ref:`mjd_inverseFD`:
-  Computes Jacobians for the continuous-time inverse dynamics (:ref:`mj_inverse`).
-  See :ref:`API documentation<mjd_inverseFD>`.
+  Computes Jacobians for the continuous or discrete-time inverse dynamics (:ref:`mj_inverse`).
+  See :ref:`documentation<mjd_inverseFD>`.
 
 These derivatives are made efficient by exploiting MuJoCo's configurable computation pipeline so that quantities are not
 recomputed when not required. For example when differencing with respect to controls, quantities which depend only on

doc/modeling.rst

@@ -347,8 +347,15 @@ of the function :math:`d(r)` is determined by the element-specific parameter vec
    For friction loss or friction dimensions of elliptic cones, the violation :math:`r` is identically zero, so
    only :math:`d(0)` affects these constraints, all other :at:`solimp` values are ignored.
 
-   .. tip::
-      For completely smooth dynamics, limits and contacts should have :math:`d_0=0`.
+   .. _solimp0:
+
+   .. admonition:: Smoothness and differentiability
+      :class: tip
+
+      For completely smooth (differentiable) dynamics, limits and contacts should have :math:`d_0=0` (``solimp[0]=0``).
+      Specifically for contacts, the :ref:`mixing rules<solmixing>` of geom-associated solver parameters should be kept
+      in mind. See also discussion of derivatives in the :ref:`Computation chapter<derivatives>` and in the
+      :ref:`mjd_transitionFD` documentation.
 
 .. _CSolverReference:
 
@@ -485,11 +492,14 @@ are as follows:
 **margin**, **gap**
    The maximum of the two geom margins (or gaps respectively) is used. The geom priority is ignored here, because the
    margin and gap are distance properties and a one-sided specification makes little sense.
+
+.. _solmixing:
+
 **solref**, **solimp**
-   If one of the two geoms has higher priority, its solref and solimp parameters are used. If both geoms have the same
-   priority, the weighted average is used. The weights are proportional to the solmix attributes, i.e., weight1 =
-   solmix1 / (solmix1 + solmix2) and similarly for weight2. There is one important exception to this weighted averaging
-   rule. If solref for either geom is non-positive, i.e., it relies on the direct format,
+   If one of the two geoms has higher :ref:`priority<body-geom-priority>`, its solref and solimp parameters are used. If
+   both geoms have the same priority, the weighted average is used. The weights are proportional to the solmix
+   attributes, i.e., weight1 = solmix1 / (solmix1 + solmix2) and similarly for weight2. There is one important exception
+   to this weighted averaging rule. If solref for either geom is non-positive, i.e., it relies on the direct format,
    then the element-wise minimum is used regardless of solmix. This is because averaging solref parameters in different
    formats would be meaningless.
 

doc/python.rst

@@ -653,7 +653,7 @@ msh2obj.py
 
 The `msh2obj.py <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/msh2obj.py>`__ script converts the
 :ref:`legacy .msh format<legacy-msh-docs>` for surface meshes (different from the possibly-volumetric
-:ref:`gmsh format<gmsh-file-docs>` also using .msh), to OBJ files. The legacy format is depricated and will be removed
+:ref:`gmsh format<gmsh-file-docs>` also using .msh), to OBJ files. The legacy format is deprecated and will be removed
 in a future release. Please convert all legacy files to OBJ.
 
 

mjx/mujoco/mjx/_src/constraint.py

@@ -529,7 +529,7 @@ def make_constraint(m: Model, d: Data) -> Data:
   if not efcs:
     z = jp.empty(0)
     d = d.replace(efc_J=jp.empty((0, m.nv)))
-    d = d.replace(efc_D=z, efc_aref=z, efc_frictionloss=z)
+    d = d.replace(efc_D=z, efc_aref=z, efc_frictionloss=z, efc_pos=z)
     return d
 
   efc = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *efcs)
@@ -539,10 +539,10 @@ def make_constraint(m: Model, d: Data) -> Data:
     k, b, imp = _kbi(m, efc.solref, efc.solimp, efc.pos_imp)
     r = jp.maximum(efc.invweight * (1 - imp) / imp, mujoco.mjMINVAL)
     aref = -b * (efc.J @ d.qvel) - k * imp * efc.pos_aref
-    return aref, r
+    return aref, r, efc.pos_aref
 
-  aref, r = fn(efc)
-  d = d.replace(efc_J=efc.J, efc_D=1 / r, efc_aref=aref)
+  aref, r, pos = fn(efc)
+  d = d.replace(efc_J=efc.J, efc_D=1 / r, efc_aref=aref, efc_pos=pos)
   d = d.replace(efc_frictionloss=jp.zeros_like(r))
 
   return d

mjx/mujoco/mjx/_src/constraint_test.py

@@ -65,6 +65,8 @@ class ConstraintTest(parameterized.TestCase):
       _assert_eq(d.efc_aref, dx.efc_aref[order][:d.nefc], 'efc_aref')
       _assert_eq(0, dx.efc_aref[order][d.nefc:], 'efc_aref')
       _assert_eq(d.efc_D, dx.efc_D[order][:d.nefc], 'efc_D')
+      _assert_eq(d.efc_pos, dx.efc_pos[order][:d.nefc], 'efc_pos')
+
 
   def test_disable_refsafe(self):
     m = test_util.load_test_file('constraints.xml')

mjx/mujoco/mjx/_src/io.py

@@ -70,15 +70,25 @@ def _make_statistic(s: mujoco.MjStatistic) -> types.Statistic:
 
 
 def put_model(
-    m: mujoco.MjModel, device=None, _check_unsupported=True
+    m: mujoco.MjModel, device=None, _full_compat: bool = False  # pylint: disable=invalid-name
 ) -> types.Model:
-  """Puts mujoco.MjModel onto a device, resulting in mjx.Model."""
+  """Puts mujoco.MjModel onto a device, resulting in mjx.Model.
+
+  Args:
+    m: the model to put onto device
+    device: which device to use - if unspecified picks the default device
+    _full_compat: put all MjModel fields onto device irrespective of MJX support
+        This is an experimental feature.  Avoid using it for now.
+
+  Returns:
+    an mjx.Model placed on device
+  """
 
   mesh_geomid = set()
   for g1, g2, ip in collision_driver.geom_pairs(m):
     t1, t2 = m.geom_type[[g1, g2]]
     # check collision function exists for type pair
-    if _check_unsupported and not collision_driver.has_collision_fn(t1, t2):
+    if not collision_driver.has_collision_fn(t1, t2) and not _full_compat:
       t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2)
       raise NotImplementedError(f'({t1}, {t2}) collisions not implemented.')
     # margin/gap not supported for meshes and height fields
@@ -88,7 +98,7 @@ def put_model(
         margin = m.pair_margin[ip]
       else:
         margin = m.geom_margin[g1] + m.geom_margin[g2]
-      if _check_unsupported and margin.any():
+      if margin.any() and not _full_compat:
         t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2)
         raise NotImplementedError(f'({t1}, {t2}) margin/gap not implemented.')
     for t, g in [(t1, g1), (t2, g2)]:
@@ -106,16 +116,19 @@ def put_model(
       (m.wrap_type, types.WrapType, mujoco.mjtWrap),
   ):
     missing = set(enum_field) - set(enum_type)
-    if _check_unsupported and missing:
+    if missing and not _full_compat:
       raise NotImplementedError(
           f'{[mj_type(m) for m in missing]} not supported'
       )
 
-  if _check_unsupported and not np.allclose(m.dof_frictionloss, 0):
+  if not np.allclose(m.dof_frictionloss, 0) and not _full_compat:
     raise NotImplementedError('dof_frictionloss is not implemented.')
 
-  mjx_only = {'mesh_convex', 'geom_rbound_hfield'}
-  mj_field_names = {f.name for f in types.Model.fields()} - mjx_only
+  mj_field_names = {
+      f.name
+      for f in types.Model.fields()
+      if f.metadata.get('restricted_to') != 'mjx'
+  }
   fields = {f: getattr(m, f) for f in mj_field_names}
   fields['geom_rbound_hfield'] = fields['geom_rbound']
   fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3))
@@ -135,29 +148,156 @@ def put_model(
   return jax.device_put(model, device=device)
 
 
-def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
-  """Allocate and initialize Data."""
+def make_data(
+    m: Union[types.Model, mujoco.MjModel],
+    device=None,
+    _full_compat: bool = False,  # pylint: disable=invalid-name
+) -> types.Data:
+  """Allocate and initialize Data.
+
+  Args:
+    m: the model to use
+    device: which device to use - if unspecified picks the default device
+    _full_compat: create all MjData fields on device irrespective of MJX support
+        This is an experimental feature.  Avoid using it for now.
+        If using this flag, also use _full_compat for put_model.
+
+  Returns:
+    an initialized mjx.Data placed on device
+  """
   dim = collision_driver.make_condim(m)
   efc_type = constraint.make_efc_type(m, dim)
   efc_address = constraint.make_efc_address(m, dim, efc_type)
   ne, nf, nl, nc = constraint.counts(efc_type)
   ncon, nefc = dim.size, ne + nf + nl + nc
 
-  contact = types.Contact(
-      dist=jp.zeros((ncon,), dtype=float),
-      pos=jp.zeros((ncon, 3), dtype=float),
-      frame=jp.zeros((ncon, 3, 3), dtype=float),
-      includemargin=jp.zeros((ncon,), dtype=float),
-      friction=jp.zeros((ncon, 5), dtype=float),
-      solref=jp.zeros((ncon, mujoco.mjNREF), dtype=float),
-      solreffriction=jp.zeros((ncon, mujoco.mjNREF), dtype=float),
-      solimp=jp.zeros((ncon, mujoco.mjNIMP), dtype=float),
-      dim=dim,
-      geom1=jp.full((ncon,), -1, dtype=int),
-      geom2=jp.full((ncon,), -1, dtype=int),
-      geom=jp.full((ncon, 2), -1, dtype=int),
-      efc_address=efc_address,
-  )
+  with jax.default_device(device):
+    contact = types.Contact(
+        dist=jp.zeros((ncon,), dtype=float),
+        pos=jp.zeros((ncon, 3), dtype=float),
+        frame=jp.zeros((ncon, 3, 3), dtype=float),
+        includemargin=jp.zeros((ncon,), dtype=float),
+        friction=jp.zeros((ncon, 5), dtype=float),
+        solref=jp.zeros((ncon, mujoco.mjNREF), dtype=float),
+        solreffriction=jp.zeros((ncon, mujoco.mjNREF), dtype=float),
+        solimp=jp.zeros((ncon, mujoco.mjNIMP), dtype=float),
+        dim=dim,
+        geom1=jp.full((ncon,), -1, dtype=jp.int32),
+        geom2=jp.full((ncon,), -1, dtype=jp.int32),
+        geom=jp.full((ncon, 2), -1, dtype=jp.int32),
+        efc_address=efc_address,
+    )
+
+    zero_fields = {
+        'solver_niter': (int,),
+        'time': (float,),
+        'qvel': (m.nv, float),
+        'act': (m.na, float),
+        'qacc_warmstart': (m.nv, float),
+        'ctrl': (m.nu, float),
+        'qfrc_applied': (m.nv, float),
+        'xfrc_applied': (m.nbody, 6, float),
+        'eq_active': (m.neq, jp.uint8),
+        'mocap_pos': (m.nmocap, 3, float),
+        'mocap_quat': (m.nmocap, 4, float),
+        'qacc': (m.nv, float),
+        'act_dot': (m.na, float),
+        'userdata': (m.nuserdata, float),
+        'sensordata': (m.nsensordata, float),
+        'xpos': (m.nbody, 3, float),
+        'xquat': (m.nbody, 4, float),
+        'xmat': (m.nbody, 3, 3, float),
+        'xipos': (m.nbody, 3, float),
+        'ximat': (m.nbody, 3, 3, float),
+        'xanchor': (m.njnt, 3, float),
+        'xaxis': (m.njnt, 3, float),
+        'geom_xpos': (m.ngeom, 3, float),
+        'geom_xmat': (m.ngeom, 3, 3, float),
+        'site_xpos': (m.nsite, 3, float),
+        'site_xmat': (m.nsite, 3, 3, float),
+        'cam_xpos': (m.ncam, 3, float),
+        'cam_xmat': (m.ncam, 3, 3, float),
+        'light_xpos': (m.nlight, 3, float),
+        'light_xdir': (m.nlight, 3, float),
+        'subtree_com': (m.nbody, 3, float),
+        'cdof': (m.nv, 6, float),
+        'cinert': (m.nbody, 10, float),
+        'flexvert_xpos': (m.nflexvert, 3, float),
+        'flexelem_aabb': (m.nflexelem, 6, float),
+        'flexedge_J_rownnz': (m.nflexedge, jp.int32),
+        'flexedge_J_rowadr': (m.nflexedge, jp.int32),
+        'flexedge_J_colind': (m.nflexedge, m.nv, jp.int32),
+        'flexedge_J': (m.nflexedge, m.nv, float),
+        'flexedge_length': (m.nflexedge, float),
+        'ten_wrapadr': (m.ntendon, jp.int32),
+        'ten_wrapnum': (m.ntendon, jp.int32),
+        'ten_J_rownnz': (m.ntendon, jp.int32),
+        'ten_J_rowadr': (m.ntendon, jp.int32),
+        'ten_J_colind': (m.ntendon, m.nv, jp.int32),
+        'ten_J': (m.ntendon, m.nv, float),
+        'ten_length': (m.ntendon, float),
+        'wrap_obj': (m.nwrap, 2, jp.int32),
+        'wrap_xpos': (m.nwrap, 6, float),
+        'actuator_length': (m.nu, float),
+        'actuator_moment': (m.nu, m.nv, float),
+        'crb': (m.nbody, 10, float),
+        'qM': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float),
+        'qLD': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float),
+        'qLDiagInv': (m.nM, float) if support.is_sparse(m) else (0, float),
+        'qLDiagSqrtInv': (m.nv, float),
+        'bvh_aabb_dyn': (m.nbvhdynamic, 6, float),
+        'bvh_active': (m.nbvh, jp.uint8),
+        'flexedge_velocity': (m.nflexedge, float),
+        'ten_velocity': (m.ntendon, float),
+        'actuator_velocity': (m.nu, float),
+        'cvel': (m.nbody, 6, float),
+        'cdof_dot': (m.nv, 6, float),
+        'qfrc_bias': (m.nv, float),
+        'qfrc_spring': (m.nv, float),
+        'qfrc_damper': (m.nv, float),
+        'qfrc_gravcomp': (m.nv, float),
+        'qfrc_fluid': (m.nv, float),
+        'qfrc_passive': (m.nv, float),
+        'subtree_linvel': (m.nbody, 3, float),
+        'subtree_angmom': (m.nbody, 3, float),
+        'qH': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float),
+        'qHDiagInv': (m.nv, float),
+        'D_rownnz': (m.nv, jp.int32),
+        'D_rowadr': (m.nv, jp.int32),
+        'D_colind': (m.nD, jp.int32),
+        'B_rownnz': (m.nbody, jp.int32),
+        'B_rowadr': (m.nbody, jp.int32),
+        'B_colind': (m.nB, jp.int32),
+        'qDeriv': (m.nD, float),
+        'qLU': (m.nD, float),
+        'actuator_force': (m.nu, float),
+        'qfrc_actuator': (m.nv, float),
+        'qfrc_smooth': (m.nv, float),
+        'qacc_smooth': (m.nv, float),
+        'qfrc_constraint': (m.nv, float),
+        'qfrc_inverse': (m.nv, float),
+        'cacc': (m.nbody, 6, float),
+        'cfrc_int': (m.nbody, 6, float),
+        'cfrc_ext': (m.nbody, 6, float),
+        'efc_J': (nefc, m.nv, float),
+        'efc_pos': (nefc, float),
+        'efc_frictionloss': (nefc, float),
+        'efc_D': (nefc, float),
+        'efc_aref': (nefc, float),
+        'efc_force': (nefc, float),
+        '_qM_sparse': (m.nM, float),
+        '_qLD_sparse': (m.nM, float),
+        '_qLDiagInv_sparse': (m.nv, float),
+    }
+
+    if not _full_compat:
+      for f in types.Data.fields():
+        if f.metadata.get('restricted_to') in ('mujoco', 'mjx'):
+          zero_fields[f.name] = (0, zero_fields[f.name][-1])
+
+    zero_fields = {
+        k: jp.zeros(v[:-1], dtype=v[-1]) for k, v in zero_fields.items()
+    }
 
   d = types.Data(
       ne=ne,
@@ -165,118 +305,10 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
       nl=nl,
       nefc=nefc,
       ncon=ncon,
-      solver_niter=jp.zeros((), dtype=int),
-      time=jp.zeros((), dtype=float),
       qpos=jp.array(m.qpos0),
-      qvel=jp.zeros((m.nv,), dtype=float),
-      act=jp.zeros((m.na,), dtype=float),
-      qacc_warmstart=jp.zeros((m.nv,), dtype=float),
-      ctrl=jp.zeros((m.nu,), dtype=float),
-      qfrc_applied=jp.zeros((m.nv,), dtype=float),
-      xfrc_applied=jp.zeros((m.nbody, 6), dtype=float),
-      eq_active=jp.zeros((m.neq,), dtype=jp.uint8),
-      mocap_pos=jp.zeros((m.nmocap, 3), dtype=float),
-      mocap_quat=jp.zeros((m.nmocap, 4), dtype=float),
-      qacc=jp.zeros((m.nv,), dtype=float),
-      act_dot=jp.zeros((m.na,), dtype=float),
-      userdata=jp.zeros((m.nuserdata,), dtype=float),
-      sensordata=jp.zeros((m.nsensordata,), dtype=float),
-      xpos=jp.zeros((m.nbody, 3), dtype=float),
-      xquat=jp.zeros((m.nbody, 4), dtype=float),
-      xmat=jp.zeros((m.nbody, 3, 3), dtype=float),
-      xipos=jp.zeros((m.nbody, 3), dtype=float),
-      ximat=jp.zeros((m.nbody, 3, 3), dtype=float),
-      xanchor=jp.zeros((m.njnt, 3), dtype=float),
-      xaxis=jp.zeros((m.njnt, 3), dtype=float),
-      geom_xpos=jp.zeros((m.ngeom, 3), dtype=float),
-      geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=float),
-      site_xpos=jp.zeros((m.nsite, 3), dtype=float),
-      site_xmat=jp.zeros((m.nsite, 3, 3), dtype=float),
-      cam_xpos=jp.zeros((m.ncam, 3), dtype=float),
-      cam_xmat=jp.zeros((m.ncam, 3, 3), dtype=float),
-      light_xpos=jp.zeros((m.nlight, 3), dtype=float),
-      light_xdir=jp.zeros((m.nlight, 3), dtype=float),
-      subtree_com=jp.zeros((m.nbody, 3), dtype=float),
-      cdof=jp.zeros((m.nv, 6), dtype=float),
-      cinert=jp.zeros((m.nbody, 10), dtype=float),
-      flexvert_xpos=jp.zeros((m.nflexvert, 3), dtype=float),
-      flexelem_aabb=jp.zeros((m.nflexelem, 6), dtype=float),
-      flexedge_J_rownnz=jp.zeros((m.nflexedge,), dtype=jp.int32),
-      flexedge_J_rowadr=jp.zeros((m.nflexedge,), dtype=jp.int32),
-      flexedge_J_colind=jp.zeros((m.nflexedge, m.nv), dtype=jp.int32),
-      flexedge_J=jp.zeros((m.nflexedge, m.nv), dtype=float),
-      flexedge_length=jp.zeros((m.nflexedge,), dtype=float),
-      ten_wrapadr=jp.zeros((m.ntendon,), dtype=jp.int32),
-      ten_wrapnum=jp.zeros((m.ntendon,), dtype=jp.int32),
-      ten_J_rownnz=jp.zeros((m.ntendon,), dtype=jp.int32),
-      ten_J_rowadr=jp.zeros((m.ntendon,), dtype=jp.int32),
-      ten_J_colind=jp.zeros((m.ntendon, m.nv), dtype=jp.int32),
-      ten_J=jp.zeros((m.ntendon, m.nv), dtype=float),
-      ten_length=jp.zeros((m.ntendon,), dtype=float),
-      wrap_obj=jp.zeros((m.nwrap, 2), dtype=jp.int32),
-      wrap_xpos=jp.zeros((m.nwrap, 6), dtype=float),
-      actuator_length=jp.zeros((m.nu,), dtype=float),
-      actuator_moment=jp.zeros((m.nu, m.nv), dtype=float),
-      crb=jp.zeros((m.nbody, 10), dtype=float),
-      qM=(
-          jp.zeros((m.nM,), dtype=float)
-          if support.is_sparse(m)
-          else jp.zeros((m.nv, m.nv), dtype=float)
-      ),
-      qLD=(
-          jp.zeros((m.nM,), dtype=float)
-          if support.is_sparse(m)
-          else jp.zeros((m.nv, m.nv), dtype=float)
-      ),
-      qLDiagInv=(
-          jp.zeros((m.nv,), dtype=float) if support.is_sparse(m)
-          else jp.zeros((0,), dtype=float)
-      ),
-      qLDiagSqrtInv=jp.zeros((m.nv,), dtype=float),
-      bvh_aabb_dyn=jp.zeros((m.nbvhdynamic, 6), dtype=float),
-      bvh_active=jp.zeros((m.nbvh,), dtype=jp.uint8),
-      flexedge_velocity=jp.zeros((m.nflexedge,), dtype=float),
-      ten_velocity=jp.zeros((m.ntendon,), dtype=float),
-      actuator_velocity=jp.zeros((m.nu,), dtype=float),
-      cvel=jp.zeros((m.nbody, 6), dtype=float),
-      cdof_dot=jp.zeros((m.nv, 6), dtype=float),
-      qfrc_bias=jp.zeros((m.nv,), dtype=float),
-      qfrc_spring=jp.zeros((m.nv,), dtype=float),
-      qfrc_damper=jp.zeros((m.nv,), dtype=float),
-      qfrc_gravcomp=jp.zeros((m.nv,), dtype=float),
-      qfrc_fluid=jp.zeros((m.nv,), dtype=float),
-      qfrc_passive=jp.zeros((m.nv,), dtype=float),
-      subtree_linvel=jp.zeros((m.nbody, 3), dtype=float),
-      subtree_angmom=jp.zeros((m.nbody, 3), dtype=float),
-      qH=jp.zeros((m.nM,), dtype=float),
-      qHDiagInv=jp.zeros((m.nv,), dtype=float),
-      D_rownnz=jp.zeros((m.nv,), dtype=jp.int32),
-      D_rowadr=jp.zeros((m.nv,), dtype=jp.int32),
-      D_colind=jp.zeros((m.nD,), dtype=jp.int32),
-      B_rownnz=jp.zeros((m.nbody,), dtype=jp.int32),
-      B_rowadr=jp.zeros((m.nbody,), dtype=jp.int32),
-      B_colind=jp.zeros((m.nB,), dtype=jp.int32),
-      qDeriv=jp.zeros((m.nD,), dtype=float),
-      qLU=jp.zeros((m.nD,), dtype=float),
-      actuator_force=jp.zeros((m.nu,), dtype=float),
-      qfrc_actuator=jp.zeros((m.nv,), dtype=float),
-      qfrc_smooth=jp.zeros((m.nv,), dtype=float),
-      qacc_smooth=jp.zeros((m.nv,), dtype=float),
-      qfrc_constraint=jp.zeros((m.nv,), dtype=float),
-      qfrc_inverse=jp.zeros((m.nv,), dtype=float),
-      cacc=jp.zeros((m.nbody, 6), dtype=float),
-      cfrc_int=jp.zeros((m.nbody, 6), dtype=float),
-      cfrc_ext=jp.zeros((m.nbody, 6), dtype=float),
       contact=contact,
       efc_type=efc_type,
-      efc_J=jp.zeros((nefc, m.nv), dtype=float),
-      efc_frictionloss=jp.zeros((nefc,), dtype=float),
-      efc_D=jp.zeros((nefc,), dtype=float),
-      efc_aref=jp.zeros((nefc,), dtype=float),
-      efc_force=jp.zeros((nefc,), dtype=float),
-      _qM_sparse=jp.zeros((m.nM), dtype=float),
-      _qLD_sparse=jp.zeros((m.nM), dtype=float),
-      _qLDiagInv_sparse=jp.zeros((m.nv,), dtype=float),
+      **zero_fields
   )
 
   return d
@@ -347,7 +379,8 @@ def get_data_into(
     result_i.efc_J_colind[:] = np.tile(np.arange(m.nv), nefc)
 
     for field in types.Data.fields():
-      if field.name.startswith('_') and field.name.endswith('_sparse'):
+      restricted_to = field.metadata.get('restricted_to')
+      if restricted_to == 'mjx':
         continue
 
       if field.name == 'contact':
@@ -375,6 +408,8 @@ def get_data_into(
         value = np.ones(m.nv)
 
       if isinstance(value, np.ndarray) and value.shape:
+        if restricted_to in ('mujoco', 'mjx') and value.shape == (0,):
+          continue  # don't copy fields that are mujoco-only or MJX-only
         getattr(result_i, field.name)[:] = value
       else:
         setattr(result_i, field.name, value)
@@ -415,8 +450,22 @@ def _make_contact(
   return types.Contact(**fields), contact_map
 
 
-def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
-  """Puts mujoco.MjData onto a device, resulting in mjx.Data."""
+def put_data(
+    m: mujoco.MjModel, d: mujoco.MjData, device=None, _full_compat: bool = False  # pylint: disable=invalid-name
+) -> types.Data:
+  """Puts mujoco.MjData onto a device, resulting in mjx.Data.
+
+  Args:
+    m: the model to use
+    d: the data to put on device
+    device: which device to use - if unspecified picks the default device
+    _full_compat: put all MjModel fields onto device irrespective of MJX support
+        This is an experimental feature.  Avoid using it for now.
+        If using this flag, also use _full_compat for put_model.
+
+  Returns:
+    an mjx.Data placed on device
+  """
   dim = collision_driver.make_condim(m)
   efc_type = constraint.make_efc_type(m, dim)
   efc_address = constraint.make_efc_address(m, dim, efc_type)
@@ -433,8 +482,11 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
     if d_val > val:
       raise ValueError(f'd.{name} too high, d.{name} = {d_val}, model = {val}')
 
-  fields = {f.name: getattr(d, f.name) for f in types.Data.fields()
-            if not f.name.endswith('_sparse')}
+  fields = {
+      f.name: getattr(d, f.name)
+      for f in types.Data.fields()
+      if f.metadata.get('restricted_to') != 'mjx'
+  }
 
   # MJX prefers square matrices for these fields:
   for fname in ('xmat', 'ximat', 'geom_xmat', 'site_xmat', 'cam_xmat'):
@@ -461,7 +513,14 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
     fields['efc_J'] = fields['efc_J'].reshape((-1 if m.nv else 0, m.nv))
 
   # move efc rows to their correct offsets
-  for fname in ('efc_J', 'efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'):
+  for fname in (
+      'efc_J',
+      'efc_pos',
+      'efc_frictionloss',
+      'efc_D',
+      'efc_aref',
+      'efc_force',
+  ):
     value = np.zeros((nefc, m.nv)) if fname == 'efc_J' else np.zeros(nefc)
     for i in range(3):
       value_beg = sum([ne, nf][:i])
@@ -482,9 +541,6 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
     fields[fname] = value
 
   # convert qM and qLD if jacobian is dense
-  fields['_qM_sparse'] = fields['qM']
-  fields['_qLD_sparse'] = fields['qLD']
-  fields['_qLDiagInv_sparse'] = fields['qLDiagInv']
   if not support.is_sparse(m):
     fields['qM'] = np.zeros((m.nv, m.nv))
     mujoco.mj_fullM(m, fields['qM'], d.qM)
@@ -496,6 +552,20 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
       fields['qLD'] = np.zeros((m.nv, m.nv))
     fields['qLDiagInv'] = np.zeros(0)
 
+  if _full_compat:
+    # full compatibility mode, we store sparse qM regardless of jacobian setting
+    fields['_qM_sparse'] = fields['qM']
+    fields['_qLD_sparse'] = fields['qLD']
+    fields['_qLDiagInv_sparse'] = fields['qLDiagInv']
+  else:
+    fields['_qM_sparse'] = jp.zeros(0, dtype=float)
+    fields['_qLD_sparse'] = jp.zeros(0, dtype=float)
+    fields['_qLDiagInv_sparse'] = jp.zeros(0, dtype=float)
+    # otherwise clear out unused arrays
+    for f in types.Data.fields():
+      if f.metadata.get('restricted_to') == 'mujoco':
+        fields[f.name] = np.zeros(0, dtype=fields[f.name].dtype)
+
   fields['contact'] = contact
   fields.update(ne=ne, nf=nf, nl=nl, nefc=nefc, ncon=ncon, efc_type=efc_type)
 

mjx/mujoco/mjx/_src/sensor.py

@@ -15,11 +15,14 @@
 """Sensor functions."""
 
 import jax
+from jax import numpy as jp
+import mujoco
 # pylint: disable=g-importing-member
+from mujoco.mjx._src import math
 from mujoco.mjx._src.types import Data
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ObjType
 from mujoco.mjx._src.types import SensorType
-from typing import Tuple
 # pylint: enable=g-importing-member
 import numpy as np
 
@@ -27,19 +30,115 @@ import numpy as np
 def sensor_pos(m: Model, d: Data) -> Data:
   """Compute position-dependent sensors values."""
 
-  sensordata = d.sensordata
-  if np.isin(SensorType.JOINTPOS, m.sensor_type):
-    # jointpos
-    i = m.sensor_type == SensorType.JOINTPOS
-    objid = m.sensor_objid[i]
-    adr = m.sensor_adr[i]
-    sensordata = sensordata.at[adr].set(d.qpos[m.jnt_qposadr[objid]])
-  if np.isin(SensorType.ACTUATORPOS, m.sensor_type):
-    # actuatorpos
-    i = m.sensor_type == SensorType.ACTUATORPOS
-    objid = m.sensor_objid[i]
-    adr = m.sensor_adr[i]
-    sensordata = sensordata.at[adr].set(d.actuator_length[objid])
+  # no position-dependent sensors
+  stage_pos = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS
+  if sum(stage_pos) == 0:
+    return d
+
+  # position and orientation by object type
+  objtype_data = {
+      ObjType.UNKNOWN: (
+          np.expand_dims(np.eye(3), axis=0),
+          np.zeros((1, 3)),
+      ),  # world
+      ObjType.BODY: (d.xipos, d.ximat),
+      ObjType.XBODY: (d.xpos, d.xmat),
+      ObjType.GEOM: (d.geom_xpos, d.geom_xmat),
+      ObjType.SITE: (d.site_xpos, d.site_xmat),
+      ObjType.CAMERA: (d.cam_xpos, d.cam_xmat),
+  }
+
+  # frame axis indexing
+  frame_axis = {
+      SensorType.FRAMEXAXIS: 0,
+      SensorType.FRAMEYAXIS: 1,
+      SensorType.FRAMEZAXIS: 2,
+  }
+
+  sensors, adrs = [], []
+
+  for sensor_type in set(m.sensor_type[stage_pos]):
+    idx = m.sensor_type == sensor_type
+    objid = m.sensor_objid[idx]
+    adr = m.sensor_adr[idx]
+
+    if sensor_type == SensorType.MAGNETOMETER:
+      sensor = jax.vmap(lambda xmat: xmat.T @ m.opt.magnetic)(
+          d.site_xmat[objid]
+      ).reshape(-1)
+      adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
+    elif sensor_type == SensorType.JOINTPOS:
+      sensor = d.qpos[m.jnt_qposadr[objid]]
+    elif sensor_type == SensorType.ACTUATORPOS:
+      sensor = d.actuator_length[objid]
+    elif sensor_type == SensorType.BALLQUAT:
+      jnt_qposadr = m.jnt_qposadr[objid, None] + np.arange(4)[None]
+      quat = d.qpos[jnt_qposadr]
+      sensor = jax.vmap(math.normalize)(quat).reshape(-1)
+      adr = (adr[:, None] + np.arange(4)[None]).reshape(-1)
+    elif sensor_type == SensorType.FRAMEPOS:
+
+      def _framepos(xpos, xpos_ref, xmat_ref, refid):
+        return jp.where(refid == -1, xpos, xmat_ref.T @ (xpos - xpos_ref))
+
+      objtype = m.sensor_objtype[idx]
+      reftype = m.sensor_reftype[idx]
+      refid = m.sensor_refid[idx]
+
+      # evaluate for valid object and reference object type pairs
+      for ot, rt in set(zip(objtype, reftype)):
+        id_ = (objtype == ot) & (reftype == rt)
+        refid_ = refid[id_]
+        xpos, _ = objtype_data[ot]
+        xpos_ref, xmat_ref = objtype_data[rt]
+        xpos = xpos[objid[id_]]
+        xpos_ref = xpos_ref[refid_]
+        xmat_ref = xmat_ref[refid_]
+        sensor = jax.vmap(_framepos)(xpos, xpos_ref, xmat_ref, refid_)
+        adr_ = adr[id_, None] + np.arange(3)[None]
+        sensors.append(sensor.reshape(-1))
+        adrs.append(adr_.reshape(-1))
+      continue  # avoid adding to sensors/adrs list a second time
+    elif sensor_type in frame_axis:
+
+      def _frameaxis(xmat, xmat_ref, refid):
+        axis = xmat[:, frame_axis[sensor_type]]
+        return jp.where(refid == -1, axis, xmat_ref.T @ axis)
+
+      objtype = m.sensor_objtype[idx]
+      reftype = m.sensor_reftype[idx]
+      refid = m.sensor_refid[idx]
+
+      # evaluate for valid object and reference object type pairs
+      for ot, rt in set(zip(objtype, reftype)):
+        id_ = (objtype == ot) & (reftype == rt)
+        refid_ = refid[id_]
+        _, xmat = objtype_data[ot]
+        _, xmat_ref = objtype_data[rt]
+        xmat = xmat[objid[id_]]
+        xmat_ref = xmat_ref[refid_]
+        sensor = jax.vmap(_frameaxis)(xmat, xmat_ref, refid_)
+        adr_ = adr[id_, None] + np.arange(3)[None]
+        sensors.append(sensor.reshape(-1))
+        adrs.append(adr_.reshape(-1))
+      continue  # avoid adding to sensors/adrs list a second time
+    elif sensor_type == SensorType.SUBTREECOM:
+      sensor = d.subtree_com[objid].reshape(-1)
+      adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
+    elif sensor_type == SensorType.CLOCK:
+      sensor = jp.repeat(d.time, sum(idx))
+    else:
+      continue  # unsupported sensor type
+
+    sensors.append(sensor)
+    adrs.append(adr)
+
+  if not adrs:
+    return d
+
+  sensordata = d.sensordata.at[np.concatenate(adrs)].set(
+      jp.concatenate(sensors)
+  )
 
   return d.replace(sensordata=sensordata)
 

mjx/mujoco/mjx/_src/sensor_test.py

@@ -62,6 +62,13 @@ class SensorTest(parameterized.TestCase):
     # sensor values
     _assert_eq(d.sensordata, dx.sensordata, 'sensordata')
 
+  def test_unsupported_sensor(self):
+    """Tests MJX sensor functions do not break for unsupported sensors."""
+    m = test_util.load_test_file('unsupported_sensor.xml')
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.forward)(mx, mjx.put_data(m, mujoco.MjData(m)))
+    _assert_eq(np.zeros(m.nsensordata), dx.sensordata, 'sensordata')
+
 
 if __name__ == '__main__':
   absltest.main()

mjx/mujoco/mjx/_src/types.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Base types used in MJX."""
 
+import dataclasses
 import enum
 from typing import Tuple
 import jax
@@ -22,6 +23,13 @@ from mujoco.mjx._src.dataclasses import PyTreeNode  # pylint: disable=g-importin
 import numpy as np
 
 
+def _restricted_to(platform: str):
+  """Specifies whether a field exists in only MuJoCo or MJX."""
+  if platform not in ('mujoco', 'mjx'):
+    raise ValueError(f'unknown platform: {platform}')
+  return dataclasses.field(metadata={'restricted_to': platform})
+
+
 class DisableBit(enum.IntFlag):
   """Disable default feature bitflags.
 
@@ -282,11 +290,46 @@ class SensorType(enum.IntEnum):
   """Type of sensor.
 
   Members:
+    MAGNETOMETER: magnetometer
     JOINTPOS: joint position
     ACTUATORPOS: actuator position
+    BALLQUAT: ball joint orientation
+    FRAMEPOS: frame position
+    FRAMEXAXIS: frame x-axis
+    FRAMEYAXIS: frame y-axis
+    FRAMEZAXIS: frame z-axis
+    SUBTREECOM: subtree centor of mass
+    CLOCK: simulation time
   """
+  MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER
   JOINTPOS = mujoco.mjtSensor.mjSENS_JOINTPOS
   ACTUATORPOS = mujoco.mjtSensor.mjSENS_ACTUATORPOS
+  BALLQUAT = mujoco.mjtSensor.mjSENS_BALLQUAT
+  FRAMEPOS = mujoco.mjtSensor.mjSENS_FRAMEPOS
+  FRAMEXAXIS = mujoco.mjtSensor.mjSENS_FRAMEXAXIS
+  FRAMEYAXIS = mujoco.mjtSensor.mjSENS_FRAMEYAXIS
+  FRAMEZAXIS = mujoco.mjtSensor.mjSENS_FRAMEZAXIS
+  SUBTREECOM = mujoco.mjtSensor.mjSENS_SUBTREECOM
+  CLOCK = mujoco.mjtSensor.mjSENS_CLOCK
+
+
+class ObjType(PyTreeNode):
+  """Type of object.
+
+  Members:
+    UNKNOWN: unknown object type
+    BODY: body
+    XBODY: body, used to access regular frame instead of i-frame
+    GEOM: geom
+    SITE: site
+    CAMERA: camera
+  """
+  UNKNOWN = mujoco.mjtObj.mjOBJ_UNKNOWN
+  BODY = mujoco.mjtObj.mjOBJ_BODY
+  XBODY = mujoco.mjtObj.mjOBJ_XBODY
+  GEOM = mujoco.mjtObj.mjOBJ_GEOM
+  SITE = mujoco.mjtObj.mjOBJ_SITE
+  CAMERA = mujoco.mjtObj.mjOBJ_CAMERA
 
 
 class Option(PyTreeNode):
@@ -329,12 +372,12 @@ class Option(PyTreeNode):
     sdf_iterations:    max number of iterations for gradient descent (not used)
   """
   timestep: jax.Array
-  apirate: jax.Array
+  apirate: jax.Array = _restricted_to('mujoco')
   impratio: jax.Array
   tolerance: jax.Array
   ls_tolerance: jax.Array
-  noslip_tolerance: jax.Array
-  mpr_tolerance: jax.Array
+  noslip_tolerance: jax.Array = _restricted_to('mujoco')
+  mpr_tolerance: jax.Array = _restricted_to('mujoco')
   gravity: jax.Array
   wind: jax.Array
   magnetic: jax.Array
@@ -344,20 +387,20 @@ class Option(PyTreeNode):
   o_solref: jax.Array
   o_solimp: jax.Array
   o_friction: jax.Array
-  has_fluid_params: bool
+  has_fluid_params: bool = _restricted_to('mjx')
   integrator: IntegratorType
   cone: ConeType
   jacobian: JacobianType
   solver: SolverType
   iterations: int
   ls_iterations: int
-  noslip_iterations: int
-  mpr_iterations: int
+  noslip_iterations: int = _restricted_to('mujoco')
+  mpr_iterations: int = _restricted_to('mujoco')
   disableflags: DisableBit
   enableflags: int
   disableactuator: int
-  sdf_initpoints: int
-  sdf_iterations: int
+  sdf_initpoints: int = _restricted_to('mujoco')
+  sdf_iterations: int = _restricted_to('mujoco')
 
 
 class Statistic(PyTreeNode):
@@ -688,22 +731,22 @@ class Model(PyTreeNode):
   nu: int
   na: int
   nbody: int
-  nbvh: int
-  nbvhstatic: int
-  nbvhdynamic: int
+  nbvh: int = _restricted_to('mujoco')
+  nbvhstatic: int = _restricted_to('mujoco')
+  nbvhdynamic: int = _restricted_to('mujoco')
   njnt: int
   ngeom: int
   nsite: int
   ncam: int
   nlight: int
-  nflex: int
-  nflexvert: int
-  nflexedge: int
-  nflexelem: int
-  nflexelemdata: int
-  nflexshelldata: int
-  nflexevpair: int
-  nflextexcoord: int
+  nflex: int = _restricted_to('mujoco')
+  nflexvert: int = _restricted_to('mujoco')
+  nflexedge: int = _restricted_to('mujoco')
+  nflexelem: int = _restricted_to('mujoco')
+  nflexelemdata: int = _restricted_to('mujoco')
+  nflexshelldata: int = _restricted_to('mujoco')
+  nflexevpair: int = _restricted_to('mujoco')
+  nflextexcoord: int = _restricted_to('mujoco')
   nmesh: int
   nmeshvert: int
   nmeshnormal: int
@@ -726,11 +769,11 @@ class Model(PyTreeNode):
   nM: int  # pylint:disable=invalid-name
   nD: int  # pylint:disable=invalid-name
   nB: int  # pylint:disable=invalid-name
-  ntree: int
+  ntree: int = _restricted_to('mujoco')
   ngravcomp: int
   nuserdata: int
   nsensordata: int
-  narena: int
+  narena: int = _restricted_to('mujoco')
   opt: Option
   stat: Statistic
   qpos0: jax.Array
@@ -759,11 +802,11 @@ class Model(PyTreeNode):
   body_margin: np.ndarray
   body_contype: np.ndarray
   body_conaffinity: np.ndarray
-  body_bvhadr: np.ndarray
-  body_bvhnum: np.ndarray
-  bvh_child: np.ndarray
-  bvh_nodeid: np.ndarray
-  bvh_aabb: np.ndarray
+  body_bvhadr: np.ndarray = _restricted_to('mujoco')
+  body_bvhnum: np.ndarray = _restricted_to('mujoco')
+  bvh_child: np.ndarray = _restricted_to('mujoco')
+  bvh_nodeid: np.ndarray = _restricted_to('mujoco')
+  bvh_aabb: np.ndarray = _restricted_to('mujoco')
   body_invweight0: jax.Array
   jnt_type: np.ndarray
   jnt_qposadr: np.ndarray
@@ -809,7 +852,7 @@ class Model(PyTreeNode):
   geom_size: jax.Array
   geom_aabb: np.ndarray
   geom_rbound: jax.Array
-  geom_rbound_hfield: np.ndarray
+  geom_rbound_hfield: np.ndarray = _restricted_to('mjx')
   geom_pos: jax.Array
   geom_quat: jax.Array
   geom_friction: jax.Array
@@ -835,54 +878,54 @@ class Model(PyTreeNode):
   cam_resolution: np.ndarray
   cam_sensorsize: np.ndarray
   cam_intrinsic: np.ndarray
-  light_mode: np.ndarray
-  light_bodyid: np.ndarray
-  light_targetbodyid: np.ndarray
-  light_pos: np.ndarray
-  light_dir: np.ndarray
-  light_poscom0: np.ndarray
-  light_pos0: np.ndarray
-  light_dir0: np.ndarray
-  flex_contype: np.ndarray
-  flex_conaffinity: np.ndarray
-  flex_condim: np.ndarray
-  flex_priority: np.ndarray
-  flex_solmix: np.ndarray
-  flex_solref: np.ndarray
-  flex_solimp: np.ndarray
-  flex_friction: np.ndarray
-  flex_margin: np.ndarray
-  flex_gap: np.ndarray
-  flex_internal: np.ndarray
-  flex_selfcollide: np.ndarray
-  flex_activelayers: np.ndarray
-  flex_dim: np.ndarray
-  flex_vertadr: np.ndarray
-  flex_vertnum: np.ndarray
-  flex_edgeadr: np.ndarray
-  flex_edgenum: np.ndarray
-  flex_elemadr: np.ndarray
-  flex_elemnum: np.ndarray
-  flex_elemdataadr: np.ndarray
-  flex_evpairadr: np.ndarray
-  flex_evpairnum: np.ndarray
-  flex_vertbodyid: np.ndarray
-  flex_edge: np.ndarray
-  flex_elem: np.ndarray
-  flex_elemlayer: np.ndarray
-  flex_evpair: np.ndarray
-  flex_vert: np.ndarray
-  flexedge_length0: np.ndarray
-  flexedge_invweight0: np.ndarray
-  flex_radius: np.ndarray
-  flex_edgestiffness: np.ndarray
-  flex_edgedamping: np.ndarray
-  flex_edgeequality: np.ndarray
-  flex_rigid: np.ndarray
-  flexedge_rigid: np.ndarray
-  flex_centered: np.ndarray
-  flex_bvhadr: np.ndarray
-  flex_bvhnum: np.ndarray
+  light_mode: np.ndarray = _restricted_to('mujoco')
+  light_bodyid: np.ndarray = _restricted_to('mujoco')
+  light_targetbodyid: np.ndarray = _restricted_to('mujoco')
+  light_pos: np.ndarray = _restricted_to('mujoco')
+  light_dir: np.ndarray = _restricted_to('mujoco')
+  light_poscom0: np.ndarray = _restricted_to('mujoco')
+  light_pos0: np.ndarray = _restricted_to('mujoco')
+  light_dir0: np.ndarray = _restricted_to('mujoco')
+  flex_contype: np.ndarray = _restricted_to('mujoco')
+  flex_conaffinity: np.ndarray = _restricted_to('mujoco')
+  flex_condim: np.ndarray = _restricted_to('mujoco')
+  flex_priority: np.ndarray = _restricted_to('mujoco')
+  flex_solmix: np.ndarray = _restricted_to('mujoco')
+  flex_solref: np.ndarray = _restricted_to('mujoco')
+  flex_solimp: np.ndarray = _restricted_to('mujoco')
+  flex_friction: np.ndarray = _restricted_to('mujoco')
+  flex_margin: np.ndarray = _restricted_to('mujoco')
+  flex_gap: np.ndarray = _restricted_to('mujoco')
+  flex_internal: np.ndarray = _restricted_to('mujoco')
+  flex_selfcollide: np.ndarray = _restricted_to('mujoco')
+  flex_activelayers: np.ndarray = _restricted_to('mujoco')
+  flex_dim: np.ndarray = _restricted_to('mujoco')
+  flex_vertadr: np.ndarray = _restricted_to('mujoco')
+  flex_vertnum: np.ndarray = _restricted_to('mujoco')
+  flex_edgeadr: np.ndarray = _restricted_to('mujoco')
+  flex_edgenum: np.ndarray = _restricted_to('mujoco')
+  flex_elemadr: np.ndarray = _restricted_to('mujoco')
+  flex_elemnum: np.ndarray = _restricted_to('mujoco')
+  flex_elemdataadr: np.ndarray = _restricted_to('mujoco')
+  flex_evpairadr: np.ndarray = _restricted_to('mujoco')
+  flex_evpairnum: np.ndarray = _restricted_to('mujoco')
+  flex_vertbodyid: np.ndarray = _restricted_to('mujoco')
+  flex_edge: np.ndarray = _restricted_to('mujoco')
+  flex_elem: np.ndarray = _restricted_to('mujoco')
+  flex_elemlayer: np.ndarray = _restricted_to('mujoco')
+  flex_evpair: np.ndarray = _restricted_to('mujoco')
+  flex_vert: np.ndarray = _restricted_to('mujoco')
+  flexedge_length0: np.ndarray = _restricted_to('mujoco')
+  flexedge_invweight0: np.ndarray = _restricted_to('mujoco')
+  flex_radius: np.ndarray = _restricted_to('mujoco')
+  flex_edgestiffness: np.ndarray = _restricted_to('mujoco')
+  flex_edgedamping: np.ndarray = _restricted_to('mujoco')
+  flex_edgeequality: np.ndarray = _restricted_to('mujoco')
+  flex_rigid: np.ndarray = _restricted_to('mujoco')
+  flexedge_rigid: np.ndarray = _restricted_to('mujoco')
+  flex_centered: np.ndarray = _restricted_to('mujoco')
+  flex_bvhadr: np.ndarray = _restricted_to('mujoco')
+  flex_bvhnum: np.ndarray = _restricted_to('mujoco')
   mesh_vertadr: np.ndarray
   mesh_vertnum: np.ndarray
   mesh_faceadr: np.ndarray
@@ -894,7 +937,7 @@ class Model(PyTreeNode):
   mesh_graph: np.ndarray
   mesh_pos: np.ndarray
   mesh_quat: np.ndarray
-  mesh_convex: Tuple[ConvexMesh, ...]
+  mesh_convex: Tuple[ConvexMesh, ...] = _restricted_to('mjx')
   hfield_size: np.ndarray
   hfield_nrow: np.ndarray
   hfield_ncol: np.ndarray
@@ -934,9 +977,9 @@ class Model(PyTreeNode):
   tendon_lengthspring: jax.Array
   tendon_length0: jax.Array
   tendon_invweight0: jax.Array
-  wrap_type: np.ndarray
-  wrap_objid: np.ndarray
-  wrap_prm: np.ndarray
+  wrap_type: np.ndarray = _restricted_to('mujoco')
+  wrap_objid: np.ndarray = _restricted_to('mujoco')
+  wrap_prm: np.ndarray = _restricted_to('mujoco')
   actuator_trntype: np.ndarray
   actuator_dyntype: np.ndarray
   actuator_gaintype: np.ndarray
@@ -959,7 +1002,7 @@ class Model(PyTreeNode):
   actuator_cranklength: np.ndarray
   actuator_acc0: np.ndarray
   actuator_lengthrange: np.ndarray
-  actuator_plugin: np.ndarray
+  actuator_plugin: np.ndarray = _restricted_to('mujoco')
   sensor_type: np.ndarray
   sensor_datatype: np.ndarray
   sensor_needstage: np.ndarray
@@ -1137,6 +1180,7 @@ class Data(PyTreeNode):
     contact: all detected contacts                              (ncon,)
     efc_type: constraint type                                   (nefc,)
     efc_J: constraint Jacobian                                  (nefc, nv)
+    efc_pos: constraint position (equality, contact)            (nefc,)
     efc_frictionloss: frictionloss (friction)                   (nefc,)
     efc_D: constraint mass                                      (nefc,)
     efc_aref: reference pseudo-acceleration                     (nefc,)
@@ -1166,8 +1210,8 @@ class Data(PyTreeNode):
   xfrc_applied: jax.Array
   eq_active: jax.Array
   # mocap data:
-  mocap_pos: jax.Array
-  mocap_quat: jax.Array
+  mocap_pos: jax.Array = _restricted_to('mujoco')
+  mocap_quat: jax.Array = _restricted_to('mujoco')
   # dynamics:
   qacc: jax.Array
   act_dot: jax.Array
@@ -1188,27 +1232,27 @@ class Data(PyTreeNode):
   site_xmat: jax.Array
   cam_xpos: jax.Array
   cam_xmat: jax.Array
-  light_xpos: jax.Array
-  light_xdir: jax.Array
+  light_xpos: jax.Array = _restricted_to('mujoco')
+  light_xdir: jax.Array = _restricted_to('mujoco')
   subtree_com: jax.Array
   cdof: jax.Array
   cinert: jax.Array
-  flexvert_xpos: jax.Array
+  flexvert_xpos: jax.Array = _restricted_to('mujoco')
   flexelem_aabb: jax.Array
-  flexedge_J_rownnz: jax.Array  # pylint:disable=invalid-name
-  flexedge_J_rowadr: jax.Array  # pylint:disable=invalid-name
-  flexedge_J_colind: jax.Array  # pylint:disable=invalid-name
-  flexedge_J: jax.Array  # pylint:disable=invalid-name
-  flexedge_length: jax.Array
-  ten_wrapadr: jax.Array
-  ten_wrapnum: jax.Array
-  ten_J_rownnz: jax.Array  # pylint:disable=invalid-name
-  ten_J_rowadr: jax.Array  # pylint:disable=invalid-name
-  ten_J_colind: jax.Array  # pylint:disable=invalid-name
+  flexedge_J_rownnz: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  flexedge_J_rowadr: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  flexedge_J_colind: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  flexedge_J: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  flexedge_length: jax.Array = _restricted_to('mujoco')
+  ten_wrapadr: jax.Array = _restricted_to('mujoco')
+  ten_wrapnum: jax.Array = _restricted_to('mujoco')
+  ten_J_rownnz: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  ten_J_rowadr: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  ten_J_colind: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
   ten_J: jax.Array  # pylint:disable=invalid-name
   ten_length: jax.Array
-  wrap_obj: jax.Array
-  wrap_xpos: jax.Array
+  wrap_obj: jax.Array = _restricted_to('mujoco')
+  wrap_xpos: jax.Array = _restricted_to('mujoco')
   actuator_length: jax.Array
   actuator_moment: jax.Array
   crb: jax.Array
@@ -1216,32 +1260,32 @@ class Data(PyTreeNode):
   qLD: jax.Array  # pylint:disable=invalid-name
   qLDiagInv: jax.Array  # pylint:disable=invalid-name
   qLDiagSqrtInv: jax.Array  # pylint:disable=invalid-name
-  bvh_aabb_dyn: jax.Array
-  bvh_active: jax.Array
+  bvh_aabb_dyn: jax.Array = _restricted_to('mujoco')
+  bvh_active: jax.Array = _restricted_to('mujoco')
   # position, velocity dependent:
-  flexedge_velocity: jax.Array
+  flexedge_velocity: jax.Array = _restricted_to('mujoco')
   ten_velocity: jax.Array
   actuator_velocity: jax.Array
   cvel: jax.Array
   cdof_dot: jax.Array
   qfrc_bias: jax.Array
-  qfrc_spring: jax.Array
-  qfrc_damper: jax.Array
+  qfrc_spring: jax.Array = _restricted_to('mujoco')
+  qfrc_damper: jax.Array = _restricted_to('mujoco')
   qfrc_gravcomp: jax.Array
   qfrc_fluid: jax.Array
   qfrc_passive: jax.Array
   subtree_linvel: jax.Array
   subtree_angmom: jax.Array
-  qH: jax.Array  # pylint:disable=invalid-name
-  qHDiagInv: jax.Array  # pylint:disable=invalid-name
-  D_rownnz: jax.Array  # pylint:disable=invalid-name
-  D_rowadr: jax.Array  # pylint:disable=invalid-name
-  D_colind: jax.Array  # pylint:disable=invalid-name
-  B_rownnz: jax.Array  # pylint:disable=invalid-name
-  B_rowadr: jax.Array  # pylint:disable=invalid-name
-  B_colind: jax.Array  # pylint:disable=invalid-name
-  qDeriv: jax.Array  # pylint:disable=invalid-name
-  qLU: jax.Array  # pylint:disable=invalid-name
+  qH: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  qHDiagInv: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  D_rownnz: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  D_rowadr: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  D_colind: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  B_rownnz: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  B_rowadr: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  B_colind: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  qDeriv: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  qLU: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
   # position, velocity, control & acceleration dependent:
   qfrc_actuator: jax.Array
   actuator_force: jax.Array
@@ -1257,6 +1301,7 @@ class Data(PyTreeNode):
   # dynamically sized - position dependent:
   efc_type: jax.Array
   efc_J: jax.Array  # pylint:disable=invalid-name
+  efc_pos: jax.Array
   efc_frictionloss: jax.Array
   efc_D: jax.Array  # pylint:disable=invalid-name
   # dynamically sized - position & velocity dependent:
@@ -1265,6 +1310,6 @@ class Data(PyTreeNode):
   efc_force: jax.Array
   # sparse representation of qM, qLD, qLDiagInv, for compatibility with MuJoCo
   # when in dense mode
-  _qM_sparse: jax.Array  # pylint:disable=invalid-name
-  _qLD_sparse: jax.Array  # pylint:disable=invalid-name
-  _qLDiagInv_sparse: jax.Array  # pylint:disable=invalid-name
+  _qM_sparse: jax.Array = _restricted_to('mjx')  # pylint:disable=invalid-name
+  _qLD_sparse: jax.Array = _restricted_to('mjx')  # pylint:disable=invalid-name
+  _qLDiagInv_sparse: jax.Array = _restricted_to('mjx')  # pylint:disable=invalid-name

mjx/mujoco/mjx/test_data/sensor.xml

@@ -1,28 +1,74 @@
 <!-- For validating sensors:
 
-* position-dependent:
+* position-dependent sensors:
+-magnetometer
 -jointpos
 -actuatorpos
-* velocity-dependent:
-* acceleration/force-dependent:
+-ballquat
+-framepos
+-framexaxis
+-frameyaxis
+-framezaxis
+-subtreecom
+-clock
+* velocity-dependent sensors:
+* acceleration/force-dependent sensors:
 -->
 <mujoco model="sensor">
   <worldbody>
-    <!-- body 0 -->
+    <!-- tree 0 -->
     <body name="body0" pos="1 2 3">
       <joint name="hinge0" type="hinge" axis="1 0 0"/>
       <geom size="1"/>
+      <site name="site0" pos=".1 .2 .3"/>
+      <body name="body1" pos="0.1 0.2 0.3">
+        <joint name="hinge1" type="hinge" axis="0 1 0"/>
+        <geom size="0.25"/>
+        <site name="site1" pos=".2 .4 .6"/>
+      </body>
+    </body>
+
+    <!-- body 2 -->
+    <body name="body2" pos=".1 .1 .1">
+      <joint name="ballquat2" type="ball" pos="0.1 0.1 0.1"/>
+      <geom size="1"/>
+    </body>
+
+    <!-- body 3 -->
+    <body name="body3" pos="-.1 -.1 -.1">
+      <joint name="ballquat3" type="ball" pos="0.1 0.2 0.3"/>
+      <geom size="1"/>
     </body>
   </worldbody>
 
   <actuator>
-    <motor name="motor0" joint="hinge0" ctrlrange="-1 1" ctrllimited="true"/>
+    <motor name="motor0" joint="hinge0" ctrlrange="-1 1" gear="10"
+      ctrllimited="true"/>
+    <motor name="motor1" joint="hinge1" ctrlrange="-1 1" gear="10"
+      ctrllimited="true"/>
   </actuator>
 
   <sensor>
     <!-- position-dependent sensors -->
+    <magnetometer name="magnetometer0" site="site0"/>
     <jointpos name="jointpos0" joint="hinge0"/>
     <actuatorpos name="actuatorpos0" actuator="motor0"/>
+    <ballquat name="ballquat2" joint="ballquat2"/>
+    <framexaxis name="framexaxis0" objtype="site" objname="site0"/>
+    <framezaxis name="framezaxis0" objtype="site" objname="site0"/>
+    <frameyaxis name="frameyaxis0" objtype="site" objname="site0"/>
+    <framepos name="framepos0" objtype="site" objname="site0"/>
+    <subtreecom name="subtreecom0" body="body0"/>
+    <magnetometer name="magnetometer1" site="site1"/>
+    <jointpos name="jointpos1" joint="hinge1"/>
+    <actuatorpos name="actuatorpos1" actuator="motor1"/>
+    <framepos name="framepos1" objtype="site" objname="site1"/>
+    <framezaxis name="framezaxis1" objtype="site" objname="site1"/>
+    <framexaxis name="framexaxis1" objtype="site" objname="site1"/>
+    <ballquat name="ballquat3" joint="ballquat3"/>
+    <frameyaxis name="frameyaxis1" objtype="site" objname="site1"/>
+    <subtreecom name="subtreecom1" body="body1"/>
+    <clock/>
     <!-- velocity-dependent sensors -->
     <!-- acceleration/force-dependent sensors -->
   </sensor>

mjx/mujoco/mjx/test_data/unsupported_sensor.xml

@@ -0,0 +1,19 @@
+<!-- For validating model with unsupported sensor -->
+<mujoco model="unsupported_sensor">
+  <worldbody>
+    <body>
+      <joint type="hinge"/>
+      <geom name="geom0" size="1"/>
+      <site name="site"/>
+    </body>
+    <body>
+      <joint type="hinge"/>
+      <geom name="geom1" size="1"/>
+    </body>
+  </worldbody>
+  <sensor>
+    <distance name="distance" geom1="geom0" geom2="geom1"/>
+    <framelinvel name="framelinvel" objtype="site" objname="site"/>
+    <touch name="touch" site="site"/>
+  </sensor>
+</mujoco>

python/mujoco/codegen/generate_spec_bindings.py

@@ -39,8 +39,10 @@ def _value_binding_code(
     fullvarname = 'ptr->' + varname
   if field.name.startswith('mjs'):  # all other mjs are raw structs
     fulltype = field.name.replace('mjs', 'raw::Mjs')
-    if field.name != 'mjsPlugin' and field.name != 'mjsOrientation':
-      fulltype = fulltype + '*'  # plugin and orientation are pointers
+    if field.name == 'mjsPlugin' or field.name == 'mjsOrientation':
+      fulltype = fulltype + '&'  # plugin and orientation are not pointers
+    else:
+      fulltype = fulltype + '*'
 
   def_property_args = (
       f'"{varname}"',

python/mujoco/specs.cc

@@ -1017,9 +1017,12 @@ PYBIND11_MODULE(_specs, m) {
   mjsTuple.def("delete", [](raw::MjsTuple& self) { mjs_delete(self.element); });
 
   // ============================= MJSPLUGIN ===================================
-  mjsPlugin.def_property_readonly("id", [](raw::MjsPlugin& self) -> int {
-    return mjs_getId(self.instance);
-  });
+  mjsPlugin.def_property(
+      "id",
+      [](raw::MjsPlugin& self) -> int { return mjs_getId(self.instance); },
+      [](raw::MjsPlugin& self, raw::MjsPlugin* other) {
+        self.instance = other->instance;
+      });
   mjsPlugin.def("delete",
                 [](raw::MjsPlugin& self) { mjs_delete(self.instance); });
 

python/mujoco/specs_test.py

@@ -322,5 +322,35 @@ class SpecsTest(absltest.TestCase):
     self.assertEqual(model.nsite, 10)
     self.assertEqual(model.nsensor, 9)
 
+  def test_plugin(self):
+    xml = """
+    <mujoco>
+      <extension>
+        <plugin plugin="mujoco.elasticity.cable"/>
+      </extension>
+    </mujoco>
+    """
+
+    spec = mujoco.MjSpec()
+    spec.from_string(xml)
+    self.assertIsNotNone(spec.worldbody)
+
+    body = spec.worldbody.add_body()
+    body.plugin.name = 'mujoco.elasticity.cable'
+    body.plugin.id = spec.add_plugin()
+    body.plugin.active = True
+    self.assertEqual(body.plugin.id, 0)
+
+    geom = body.add_geom()
+    geom.type = mujoco.mjtGeom.mjGEOM_BOX
+    geom.size[0] = 1
+    geom.size[1] = 1
+    geom.size[2] = 1
+
+    model = spec.compile()
+    self.assertIsNotNone(model)
+    self.assertEqual(model.nplugin, 1)
+    self.assertEqual(model.body_plugin[1], 0)
+
 if __name__ == '__main__':
   absltest.main()

src/engine/engine_collision_convex.c

@@ -360,8 +360,10 @@ static void mju_rotateFrame(const mjtNum origin[3], const mjtNum rot[9],
 int mjc_Convex(const mjModel* m, const mjData* d,
                mjContact* con, int g1, int g2, mjtNum margin) {
   ccd_t ccd;
-  mjCCDObj obj1 = {m, d, g1, -1, -1, -1, -1, margin, {1, 0, 0, 0}, {0, 0, 0}};
-  mjCCDObj obj2 = {m, d, g2, -1, -1, -1, -1, margin, {1, 0, 0, 0}, {0, 0, 0}};
+  mjCCDObj obj1 = {m, d, g1, -1, -1, -1, -1, margin, {1, 0, 0, 0}, {0, 0, 0},
+                   mjc_center, mjc_support};
+  mjCCDObj obj2 = {m, d, g2, -1, -1, -1, -1, margin, {1, 0, 0, 0}, {0, 0, 0},
+                   mjc_center, mjc_support};
 
   // init ccd structure
   mjc_initCCD(&ccd, m);
@@ -586,46 +588,45 @@ int mjc_PlaneConvex(const mjModel* m, const mjData* d,
 
 //----------------------------  heightfield collisions ---------------------------------------------
 
-// ccd prism object type
-struct _mjtPrism {
-  mjtNum v[6][3];
-};
-
-typedef struct _mjtPrism mjtPrism;
-
-
-// ccd prism support function
-static void prism_support(const void *obj, const ccd_vec3_t *dir, ccd_vec3_t *vec) {
+// prism support function
+static void mjc_prism_support(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]) {
   int istart, ibest;
   mjtNum best, tmp;
-  const mjtPrism* p = (const mjtPrism*)obj;
 
   // find best vertex in halfspace determined by dir.z
-  istart = dir->v[2] < 0 ? 0 : 3;
+  istart = dir[2] < 0 ? 0 : 3;
   ibest = istart;
-  best = mju_dot3(p->v[istart], dir->v);
+  best = mju_dot3(obj->prism[istart], dir);
   for (int i=istart+1; i < istart+3; i++) {
-    if ((tmp = mju_dot3(p->v[i], dir->v)) > best) {
+    if ((tmp = mju_dot3(obj->prism[i], dir)) > best) {
       ibest = i;
       best = tmp;
     }
   }
 
   // copy best point
-  mju_copy3(vec->v, p->v[ibest]);
+  mju_copy3(res, obj->prism[ibest]);
+}
+
+// ccd prism support function
+static void mjccd_prism_support(const void *obj, const ccd_vec3_t *dir, ccd_vec3_t *vec) {
+  mjc_prism_support(vec->v, (mjCCDObj*) obj, dir->v);
 }
 
 
-// ccd prism center function
-static void prism_center(const void *obj, ccd_vec3_t *center) {
-  const mjtPrism* p = (const mjtPrism*)obj;
-
+// prism center function
+static void mjc_prism_center(mjtNum res[3], const mjCCDObj* obj) {
   // compute mean
-  mju_zero3(center->v);
+  mju_zero3(res);
   for (int i=0; i < 6; i++) {
-    mju_addTo3(center->v, p->v[i]);
+    mju_addTo3(res, obj->prism[i]);
   }
-  mju_scl3(center->v, center->v, 1.0/6.0);
+  mju_scl3(res, res, 1.0/6.0);
+}
+
+// ccd prism center function
+static void mjccd_prism_center(const void *obj, ccd_vec3_t *center) {
+  mjc_prism_center(center->v, (const mjCCDObj*) obj);
 }
 
 
@@ -636,17 +637,17 @@ static void prism_firstdir(const void* o1, const void* o2, ccd_vec3_t *vec) {
 
 
 // add vertex to prism, count vertices
-static void addVert(int* nvert, mjtPrism* prism, mjtNum x, mjtNum y, mjtNum z) {
+static void addVert(int* nvert, mjCCDObj* obj, mjtNum x, mjtNum y, mjtNum z) {
   // move old data
-  mju_copy3(prism->v[0], prism->v[1]);
-  mju_copy3(prism->v[1], prism->v[2]);
-  mju_copy3(prism->v[3], prism->v[4]);
-  mju_copy3(prism->v[4], prism->v[5]);
+  mju_copy3(obj->prism[0], obj->prism[1]);
+  mju_copy3(obj->prism[1], obj->prism[2]);
+  mju_copy3(obj->prism[3], obj->prism[4]);
+  mju_copy3(obj->prism[4], obj->prism[5]);
 
   // add new vertex at last position
-  prism->v[2][0] = prism->v[5][0] = x;
-  prism->v[2][1] = prism->v[5][1] = y;
-  prism->v[5][2] = z;
+  obj->prism[2][0] = obj->prism[5][0] = x;
+  obj->prism[2][1] = obj->prism[5][1] = y;
+  obj->prism[5][2] = z;
 
   // count
   (*nvert)++;
@@ -665,12 +666,15 @@ int mjc_ConvexHField(const mjModel* m, const mjData* d,
   int ncol = m->hfield_ncol[hid];
   int dr[2], cnt, rmin, rmax, cmin, cmax;
   const float* data = m->hfield_data + m->hfield_adr[hid];
-  mjtPrism prism;
+  mjCCDObj obj1;
+  obj1.center = mjc_prism_center;
+  obj1.support = mjc_prism_support;
 
   // ccd-related
   ccd_vec3_t dirccd, vecccd;
   ccd_real_t depth;
-  mjCCDObj obj = {m, d, g2, -1, -1, -1, -1, 0, {1, 0, 0, 0}, {0, 0, 0}};
+  mjCCDObj obj2 = {m, d, g2, -1, -1, -1, -1, 0, {1, 0, 0, 0}, {0, 0, 0},
+                  mjc_center, mjc_support};
   ccd_t ccd;
 
   // point size1 to hfield size instead of geom1 size
@@ -713,32 +717,32 @@ int mjc_ConvexHField(const mjModel* m, const mjData* d,
 
   // get support point in +X
   ccdVec3Set(&dirccd, 1, 0, 0);
-  mjccd_support(&obj, &dirccd, &vecccd);
+  mjccd_support(&obj2, &dirccd, &vecccd);
   xmax = vecccd.v[0];
 
   // get support point in -X
   ccdVec3Set(&dirccd, -1, 0, 0);
-  mjccd_support(&obj, &dirccd, &vecccd);
+  mjccd_support(&obj2, &dirccd, &vecccd);
   xmin = vecccd.v[0];
 
   // get support point in +Y
   ccdVec3Set(&dirccd, 0, 1, 0);
-  mjccd_support(&obj, &dirccd, &vecccd);
+  mjccd_support(&obj2, &dirccd, &vecccd);
   ymax = vecccd.v[1];
 
   // get support point in -Y
   ccdVec3Set(&dirccd, 0, -1, 0);
-  mjccd_support(&obj, &dirccd, &vecccd);
+  mjccd_support(&obj2, &dirccd, &vecccd);
   ymin = vecccd.v[1];
 
   // get support point in +Z
   ccdVec3Set(&dirccd, 0, 0, 1);
-  mjccd_support(&obj, &dirccd, &vecccd);
+  mjccd_support(&obj2, &dirccd, &vecccd);
   zmax = vecccd.v[2];
 
   // get support point in -Z
   ccdVec3Set(&dirccd, 0, 0, -1);
-  mjccd_support(&obj, &dirccd, &vecccd);
+  mjccd_support(&obj2, &dirccd, &vecccd);
   zmin = vecccd.v[2];
 
   // box-box test
@@ -767,13 +771,13 @@ int mjc_ConvexHField(const mjModel* m, const mjData* d,
   // init ccd structure
   mjc_initCCD(&ccd, m);
   ccd.first_dir = prism_firstdir;
-  ccd.center1 = prism_center;
+  ccd.center1 = mjccd_prism_center;
   ccd.center2 = mjccd_center;
-  ccd.support1 = prism_support;
+  ccd.support1 = mjccd_prism_support;
   ccd.support2 = mjccd_support;
 
   // geom margin needed for actual collision test
-  obj.margin = margin;
+  obj2.margin = margin;
 
   // compute real-valued grid step, and triangulation direction
   dx = (2.0*size1[0]) / (ncol-1);
@@ -782,7 +786,7 @@ int mjc_ConvexHField(const mjModel* m, const mjData* d,
   dr[1] = 0;
 
   // set zbottom value using base size
-  prism.v[0][2] = prism.v[1][2] = prism.v[2][2] = -size1[3];
+  obj1.prism[0][2] = obj1.prism[1][2] = obj1.prism[2][2] = -size1[3];
 
   // process all prisms in sub-grid
   cnt = 0;
@@ -791,19 +795,20 @@ int mjc_ConvexHField(const mjModel* m, const mjData* d,
     for (int c=cmin; c <= cmax; c++) {
       for (int i=0; i < 2; i++) {
         // send vertex to prism constructor
-        addVert(&nvert, &prism, dx*c-size1[0], dy*(r+dr[i])-size1[1],
+        addVert(&nvert, &obj1, dx*c-size1[0], dy*(r+dr[i])-size1[1],
                 data[(r+dr[i])*ncol+c]*size1[2]+margin);
 
         // check for enough vertices
         if (nvert > 2) {
           // prism height test
-          if (prism.v[3][2] < zmin && prism.v[4][2] < zmin && prism.v[5][2] < zmin) {
+          if (obj1.prism[3][2] < zmin && obj1.prism[4][2] < zmin
+              && obj1.prism[5][2] < zmin) {
             continue;
           }
 
           // run MPR, save contact
-          if (_mjCCDPENETRATION(&prism, &obj, &ccd, &depth, &dirccd, &vecccd) == 0 &&
-              !ccdVec3Eq(&dirccd, ccd_vec3_origin)) {
+          if (_mjCCDPENETRATION(&obj1, &obj2, &ccd, &depth, &dirccd, &vecccd) == 0
+              && !ccdVec3Eq(&dirccd, ccd_vec3_origin)) {
             // fill in contact data, transform to global coordinates
             con[cnt].dist = -depth;
             mju_mulMatVec3(con[cnt].frame, mat1, dirccd.v);
@@ -1103,8 +1108,10 @@ void mjc_fixNormal(const mjModel* m, const mjData* d, mjContact* con, int g1, in
 int mjc_ConvexElem(const mjModel* m, const mjData* d, mjContact* con,
                    int g1, int f1, int e1, int v1, int f2, int e2, mjtNum margin) {
   ccd_t ccd;
-  mjCCDObj obj1 = {m, d, g1, -1, f1, e1, v1, margin, {1, 0, 0, 0}, {0, 0, 0}};
-  mjCCDObj obj2 = {m, d, -1, -1, f2, e2, -1, margin, {1, 0, 0, 0}, {0, 0, 0}};
+  mjCCDObj obj1 = {m, d, g1, -1, f1, e1, v1, margin, {1, 0, 0, 0}, {0, 0, 0},
+                   mjc_center, mjc_support};
+  mjCCDObj obj2 = {m, d, -1, -1, f2, e2, -1, margin, {1, 0, 0, 0}, {0, 0, 0},
+                   mjc_center, mjc_support};
 
   // init ccd structure
   mjc_initCCD(&ccd, m);
@@ -1128,7 +1135,9 @@ int mjc_HFieldElem(const mjModel* m, const mjData* d, mjContact* con,
   mjtNum vec[3], dx, dy;
   mjtNum xmin, xmax, ymin, ymax, zmin, zmax;
   int dr[2], cnt, rmin, rmax, cmin, cmax;
-  mjtPrism prism;
+  mjCCDObj obj1;
+  obj1.center = mjc_prism_center;
+  obj1.support = mjc_prism_support;
 
   // get hfield info
   int hid = m->geom_dataid[g];
@@ -1151,7 +1160,8 @@ int mjc_HFieldElem(const mjModel* m, const mjData* d, mjContact* con,
   // ccd-related
   ccd_vec3_t dirccd, vecccd;
   ccd_real_t depth;
-  mjCCDObj obj = {m, d, -1, -1, f, e, -1, margin, {1, 0, 0, 0}, {0, 0, 0}};
+  mjCCDObj obj2 = {m, d, -1, -1, f, e, -1, margin, {1, 0, 0, 0}, {0, 0, 0},
+                  mjc_center, mjc_support};
   ccd_t ccd;
 
   //------------------------------------- AABB computation, box-box test
@@ -1211,9 +1221,9 @@ int mjc_HFieldElem(const mjModel* m, const mjData* d, mjContact* con,
   // init ccd structure
   CCD_INIT(&ccd);
   ccd.first_dir = prism_firstdir;
-  ccd.center1 = prism_center;
+  ccd.center1 = mjccd_prism_center;
   ccd.center2 = mjccd_center;
-  ccd.support1 = prism_support;
+  ccd.support1 = mjccd_prism_support;
   ccd.support2 = mjccd_support;
 
   // set ccd parameters
@@ -1227,7 +1237,7 @@ int mjc_HFieldElem(const mjModel* m, const mjData* d, mjContact* con,
   dr[1] = 0;
 
   // set zbottom value using base size
-  prism.v[0][2] = prism.v[1][2] = prism.v[2][2] = -hsize[3];
+  obj1.prism[0][2] = obj1.prism[1][2] = obj1.prism[2][2] = -hsize[3];
 
   // process all prisms in sub-grid
   cnt = 0;
@@ -1236,18 +1246,18 @@ int mjc_HFieldElem(const mjModel* m, const mjData* d, mjContact* con,
     for (int c=cmin; c <= cmax; c++) {
       for (int k=0; k < 2; k++) {
         // send vertex to prism constructor
-        addVert(&nvert, &prism, dx*c-hsize[0], dy*(r+dr[k])-hsize[1],
+        addVert(&nvert, &obj1, dx*c-hsize[0], dy*(r+dr[k])-hsize[1],
                 hdata[(r+dr[k])*ncol+c]*hsize[2]+margin);
 
         // check for enough vertices
         if (nvert > 2) {
           // prism height test
-          if (prism.v[3][2] < zmin && prism.v[4][2] < zmin && prism.v[5][2] < zmin) {
+          if (obj1.prism[3][2] < zmin && obj1.prism[4][2] < zmin && obj1.prism[5][2] < zmin) {
             continue;
           }
 
           // run MPR, save contact
-          if (_mjCCDPENETRATION(&prism, &obj, &ccd, &depth, &dirccd, &vecccd) == 0) {
+          if (_mjCCDPENETRATION(&obj1, &obj2, &ccd, &depth, &dirccd, &vecccd) == 0) {
             if (!ccdVec3Eq(&dirccd, ccd_vec3_origin)) {
               // fill in contact data, transform to global coordinates
               con[cnt].dist = -depth;

src/engine/engine_collision_convex.h

@@ -52,6 +52,9 @@ struct _mjCCDObj {
   mjtNum margin;
   mjtNum rotate[4];
   mjtNum x0[3];  // initial guess of the witness point
+  void (*center)(mjtNum res[3], const struct _mjCCDObj* obj);
+  void (*support)(mjtNum res[3], struct _mjCCDObj* obj, const mjtNum dir[3]);
+  mjtNum prism[6][3];  // for hfield
 };
 typedef struct _mjCCDObj mjCCDObj;
 

src/engine/engine_collision_gjk.c

@@ -203,8 +203,8 @@ static void gjk_support(mjtNum s1[3], mjtNum s2[3], mjCCDObj* obj1, mjCCDObj* ob
   mju_scl3(dir, dir_neg, -1);
 
   // compute S_{A-B}(dir) = S_A(dir) - S_B(-dir)
-  mjc_support(s1, obj1, dir);
-  mjc_support(s2, obj2, dir_neg);
+  obj1->support(s1, obj1, dir);
+  obj2->support(s2, obj2, dir_neg);
 }
 
 
@@ -218,8 +218,8 @@ static void support(mjtNum s1[3], mjtNum s2[3], mjCCDObj* obj1, mjCCDObj* obj2,
   mju_scl3(dir_neg, dir, -1);
 
   // compute S_{A-B}(dir) = S_A(dir) - S_B(-dir)
-  mjc_support(s1, obj1, dir);
-  mjc_support(s2, obj2, dir_neg);
+  obj1->support(s1, obj1, dir);
+  obj2->support(s2, obj2, dir_neg);
 }
 
 
@@ -246,6 +246,33 @@ static mjtNum det3(const mjtNum v1[3], const mjtNum v2[3], const mjtNum v3[3]) {
 }
 
 
+// res = origin projected onto plane defined by v1, v2, v3
+static inline void projectOriginPlane(mjtNum res[3], mjtNum normal[3], const mjtNum v1[3],
+                                      const mjtNum v2[3], const mjtNum v3[3]) {
+  mjtNum diff1[3], diff2[3], tmp[3];
+  mju_sub3(diff1, v2, v1);
+  mju_sub3(diff2, v3, v1);
+  mju_cross(tmp, diff1, diff2);  // vector normal to the plane
+
+  // res = tmp * <tmp, v1> / ||tmp||^2
+  mjtNum tmp_sqr = mju_dot3(tmp, tmp);
+  mju_scl3(res, tmp, mju_dot3(tmp, v1) / tmp_sqr);
+  if (normal) mju_scl3(normal, tmp, 1/mju_sqrt(tmp_sqr));
+}
+
+
+
+// res = origin projected onto line defined by v1, v2
+static inline void projectOriginLine(mjtNum res[3], const mjtNum v1[3], const mjtNum v2[3]) {
+  // res = v2 - <v2, v2 - v1> / <v2 - v1, v2 - v1> * (v2 - v1)
+  mjtNum diff[3];
+  mju_sub3(diff, v2, v1);
+  mjtNum temp1 = mju_dot3(v2, diff);
+  mjtNum temp2 = mju_dot3(diff, diff);
+  mju_addScl3(res, v2, diff, - temp1 / temp2);
+}
+
+
 
 // returns true only when a and b are both strictly positive or both strictly negative
 static int compareSigns(mjtNum a, mjtNum b) {
@@ -391,32 +418,25 @@ static void S2D(mjtNum lambda[3], const mjtNum simplex[9]) {
   const mjtNum* s2 = simplex + 3;
   const mjtNum* s3 = simplex + 6;
 
-  // compute normal
-  mjtNum diff1[3], diff2[3], n[3];
-  mju_sub3(diff1, s2, s1);
-  mju_sub3(diff2, s3, s1);
-  mju_cross(n, diff1, diff2);
-
-  // project origin
+  // project origin onto affine hull of the simplex
   mjtNum p_o[3];
-  mju_scl3(p_o, n, mju_dot3(n, s1) / mju_dot3(n, n));
+  projectOriginPlane(p_o, NULL, s1, s2, s3);
 
-  mjtNum mu_max = 0;
-
-  //  Below are the minors M_i4 of the matrix M given by
-  //  [[ s1_x, s2_x, s3_x, s4_x ],
-  //   [ s1_y, s2_y, s3_y, s4_y ],
-  //   [ s1_z, s2_z, s3_z, s4_z ],
-  //   [ 1,    1,    1,    1    ]]
+  // Below are the minors M_i4 of the matrix M given by
+  // [[ s1_x, s2_x, s3_x, s4_x ],
+  //  [ s1_y, s2_y, s3_y, s4_y ],
+  //  [ s1_z, s2_z, s3_z, s4_z ],
+  //  [ 1,    1,    1,    1    ]]
   mjtNum M_14 = s2[1]*s3[2] - s2[2]*s3[1] - s1[1]*s3[2] + s1[2]*s3[1] + s1[1]*s2[2] - s1[2]*s2[1];
   mjtNum M_24 = s2[0]*s3[2] - s2[2]*s3[0] - s1[0]*s3[2] + s1[2]*s3[0] + s1[0]*s2[2] - s1[2]*s2[0];
   mjtNum M_34 = s2[0]*s3[1] - s2[1]*s3[0] - s1[0]*s3[1] + s1[1]*s3[0] + s1[0]*s2[1] - s1[1]*s2[0];
 
   // exclude one of the axes with the largest projection of the simplex using the computed minors
+  mjtNum M_max = 0;
   mjtNum s1_2D[2], s2_2D[2], s3_2D[2], p_o_2D[2];
   mjtNum mu1 = mju_abs(M_14), mu2 = mju_abs(M_24), mu3 = mju_abs(M_34);
   if (mu1 >= mu2 && mu1 >= mu3) {
-    mu_max = mu1;
+    M_max = M_14;
     s1_2D[0] = s1[1];
     s1_2D[1] = s1[2];
 
@@ -429,7 +449,7 @@ static void S2D(mjtNum lambda[3], const mjtNum simplex[9]) {
     p_o_2D[0] = p_o[1];
     p_o_2D[1] = p_o[2];
   } else if (mu2 >= mu3) {
-    mu_max = mu2;
+    M_max = M_24;
     s1_2D[0] = s1[0];
     s1_2D[1] = s1[2];
 
@@ -442,7 +462,7 @@ static void S2D(mjtNum lambda[3], const mjtNum simplex[9]) {
     p_o_2D[0] = p_o[0];
     p_o_2D[1] = p_o[2];
   } else {
-    mu_max = mu3;
+    M_max = M_34;
     s1_2D[0] = s1[0];
     s1_2D[1] = s1[1];
 
@@ -456,25 +476,32 @@ static void S2D(mjtNum lambda[3], const mjtNum simplex[9]) {
     p_o_2D[1] = p_o[1];
   }
 
-  // substitute p_o as a vertex in simplex
-  mjtNum C1 = p_o_2D[0]*s2_2D[1] + p_o_2D[1]*s3_2D[0] + s2_2D[0]*s3_2D[1]
-            - p_o_2D[0]*s3_2D[1] - p_o_2D[1]*s2_2D[0] - s3_2D[0]*s2_2D[1];
+  // compute the cofactors C3i of the following matrix:
+  // [[ s1_2D[0] - p_o_2D[0], s2_2D[0] - p_o_2D[0], s3_2D[0] - p_o_2D[0] ],
+  //  [ s1_2D[1] - p_o_2D[1], s2_2D[1] - p_o_2D[1], s3_2D[1] - p_o_2D[1] ],
+  //  [ 1,                    1,                    1                    ]]
 
-  mjtNum C2 = p_o_2D[0]*s3_2D[1] + p_o_2D[1]*s1_2D[0] + s3_2D[0]*s1_2D[1]
-            - p_o_2D[0]*s1_2D[1] - p_o_2D[1]*s3_2D[0] - s1_2D[0]*s3_2D[1];
+  // C31 corresponds to the signed area of 2-simplex: (p_o_2D, s2_2D, s3_2D)
+  mjtNum C31 = p_o_2D[0]*s2_2D[1] + p_o_2D[1]*s3_2D[0] + s2_2D[0]*s3_2D[1]
+             - p_o_2D[0]*s3_2D[1] - p_o_2D[1]*s2_2D[0] - s3_2D[0]*s2_2D[1];
 
-  mjtNum C3 = p_o_2D[0]*s1_2D[1] + p_o_2D[1]*s2_2D[0] + s1_2D[0]*s2_2D[1]
-            - p_o_2D[0]*s2_2D[1] - p_o_2D[1]*s1_2D[0] - s2_2D[0]*s1_2D[1];
+  // C32 corresponds to the signed area of 2-simplex: (_po_2D, s1_2D, s3_2D)
+  mjtNum C32 = p_o_2D[0]*s3_2D[1] + p_o_2D[1]*s1_2D[0] + s3_2D[0]*s1_2D[1]
+             - p_o_2D[0]*s1_2D[1] - p_o_2D[1]*s3_2D[0] - s1_2D[0]*s3_2D[1];
 
-  int comp1 = compareSigns(mu_max, C1),
-      comp2 = compareSigns(mu_max, C2),
-      comp3 = compareSigns(mu_max, C3);
+  // C33 corresponds to the signed area of 2-simplex: (p_o_2D, s1_2D, s2_2D)
+  mjtNum C33 = p_o_2D[0]*s1_2D[1] + p_o_2D[1]*s2_2D[0] + s1_2D[0]*s2_2D[1]
+             - p_o_2D[0]*s2_2D[1] - p_o_2D[1]*s1_2D[0] - s2_2D[0]*s1_2D[1];
 
-  // inside the simplex
+  int comp1 = compareSigns(M_max, C31),
+      comp2 = compareSigns(M_max, C32),
+      comp3 = compareSigns(M_max, C33);
+
+  // all the same sign, p_o is inside the 2-simplex
   if (comp1 && comp2 && comp3) {
-    lambda[0] = C1 / mu_max;
-    lambda[1] = C2 / mu_max;
-    lambda[2] = C3 / mu_max;
+    lambda[0] = C31 / M_max;
+    lambda[1] = C32 / M_max;
+    lambda[2] = C33 / M_max;
     return;
   }
 
@@ -517,9 +544,9 @@ static void S2D(mjtNum lambda[3], const mjtNum simplex[9]) {
     lincomb(x, lambda_1d, verts, 2);
     mjtNum d = mju_norm3(x);
     if (d < dist) {
+      lambda[0] = 0;
       lambda[1] = lambda_1d[0];
       lambda[2] = lambda_1d[1];
-      lambda[0] = 0;
       dist = d;
     }
   }
@@ -533,13 +560,8 @@ static void S1D(mjtNum lambda[2], const mjtNum simplex[6]) {
   const mjtNum* s2 = simplex + 3;
 
   // find projection of origin onto the 1-simplex:
-  //   p_o = s2 - <s2, s2 - s1> / <s2 - s1, s2 - s1> * (s2 - s1)
   mjtNum p_o[3];
-  mjtNum diff[3];
-  mju_sub3(diff, s2, s1);
-  mjtNum temp1 = mju_dot3(s2, diff);
-  mjtNum temp2 = mju_dot3(diff, diff);
-  mju_addScl3(p_o, s2, diff, - temp1 / temp2);
+  projectOriginLine(p_o, s1, s2);
 
   // find the axis with the largest projection "shadow" of the simplex
   mjtNum mu_max = 0;
@@ -785,18 +807,11 @@ static void attachFace(Polytope* pt, int v1, int v2, int v3) {
   face->verts[1] = v2;
   face->verts[2] = v3;
 
-  // compute normal n
+    // compute witness point v
   mjtNum* pv1 = pt->verts[v1].v;
   mjtNum* pv2 = pt->verts[v2].v;
   mjtNum* pv3 = pt->verts[v3].v;
-  mjtNum diff1[3], diff2[3];
-  mju_sub3(diff1, pv2, pv1);
-  mju_sub3(diff2, pv3, pv1);
-  mju_cross(face->n, diff1, diff2);
-  mju_normalize3(face->n);
-
-  // compute witness point v
-  mju_scl3(face->v, face->n, mju_dot3(face->n, pv1));
+  projectOriginPlane(face->v, face->n, pv1, pv2, pv3);
   face->dist = mju_norm3(face->v);
 
   // orientation check
@@ -992,8 +1007,8 @@ int mj_gjkPenetration(const void *obj1, const void *obj2, const ccd_t *ccd,
   mjCCDObj* o2 = (mjCCDObj*) obj2;
   nearest.n[1] = 34;
 
-  mjc_center(o1->x0, o1);
-  mjc_center(o2->x0, o2);
+  o1->center(o1->x0, o1);
+  o2->center(o2->x0, o2);
 
   config.max_iterations = ccd->max_iterations;
   config.tolerance = ccd->mpr_tolerance;

src/engine/engine_derivative.c

@@ -338,7 +338,7 @@ static void mjd_comVel_vel_dense(const mjModel* m, mjData* d, mjtNum* Dcvel, mjt
   mju_zero(Dcvel, nbody*6*nv);
 
   // forward pass over bodies: accumulate Dcvel, set Dcdofdot
-  for (int i=1; i < m->nbody; i++) {
+  for (int i=1; i < nbody; i++) {
     // Dcvel = Dcvel_parent
     mju_copy(Dcvel+i*6*nv, Dcvel+m->body_parentid[i]*6*nv, 6*nv);
 
@@ -450,7 +450,7 @@ void mjd_rne_vel_dense(const mjModel* m, mjData* d) {
   mju_zero(Dcfrcbody, 6*nv);
 
   // backward pass over bodies: accumulate Dcfrcbody
-  for (int i=m->nbody-1; i > 0; i--) {
+  for (int i=nbody-1; i > 0; i--) {
     if (m->body_parentid[i]) {
       mju_addTo(Dcfrcbody+m->body_parentid[i]*6*nv, Dcfrcbody+i*6*nv, 6*nv);
     }
@@ -825,7 +825,7 @@ static mjtNum mjd_muscleGain_vel(mjtNum len, mjtNum vel, const mjtNum lengthrang
 
 // add (d qfrc_actuator / d qvel) to qDeriv
 void mjd_actuator_vel(const mjModel* m, mjData* d) {
-  int nv = m->nv;
+  int nv = m->nv, nu = m->nu;
 
   // disabled: nothing to add
   if (mjDISABLED(mjDSBL_ACTUATION)) {
@@ -833,7 +833,7 @@ void mjd_actuator_vel(const mjModel* m, mjData* d) {
   }
 
   // process actuators
-  for (int i=0; i < m->nu; i++) {
+  for (int i=0; i < nu; i++) {
     // skip if disabled
     if (mj_actuatorDisabled(m, i)) {
       continue;
@@ -867,7 +867,9 @@ void mjd_actuator_vel(const mjModel* m, mjData* d) {
       if (m->actuator_dyntype[i] == mjDYN_NONE) {
         bias_vel += gain_vel * d->ctrl[i];
       } else {
-        bias_vel += gain_vel * d->act[i-(m->nu - m->na)];
+        int act_first = m->actuator_actadr[i];
+        int act_last = act_first + m->actuator_actnum[i] - 1;
+        bias_vel += gain_vel * d->act[act_last];
       }
     }
 

src/engine/engine_forward.c

@@ -264,21 +264,6 @@ static void clampVec(mjtNum* vec, const mjtNum* range, const mjtByte* limited, i
 
 
 
-// return number of dofs given joint type
-static int jnt_dofnum(mjtJoint type) {
-  if (type == mjJNT_FREE) {
-    return 6;
-  }
-
-  if (type == mjJNT_BALL) {
-    return 3;
-  }
-
-  return 1;
-}
-
-
-
 // (qpos, qvel, ctrl, act) => (qfrc_actuator, actuator_force, act_dot)
 void mj_fwdActuation(const mjModel* m, mjData* d) {
   TM_START;
@@ -493,6 +478,8 @@ void mj_fwdActuation(const mjModel* m, mjData* d) {
 
   // actuator-level gravity compensation
   if (m->ngravcomp && !mjDISABLED(mjDSBL_GRAVITY) && mju_norm3(m->opt.gravity)) {
+    // number of dofs for each joint type: {mjJNT_FREE, mjJNT_BALL, mjJNT_SLIDE, mjJNT_HINGE}
+    static const int jnt_dofnum[4] = {6, 3, 1, 1};
     int njnt = m->njnt;
     for (int i=0; i < njnt; i++) {
       // skip if gravcomp added as passive force
@@ -501,7 +488,7 @@ void mj_fwdActuation(const mjModel* m, mjData* d) {
       }
 
       // add gravcomp force
-      int dofnum = jnt_dofnum(m->jnt_type[i]);
+      int dofnum = jnt_dofnum[m->jnt_type[i]];
       int dofadr = m->jnt_dofadr[i];
       mju_addTo(d->qfrc_actuator + dofadr, d->qfrc_gravcomp + dofadr, dofnum);
     }

src/user/user_api.cc

@@ -93,15 +93,17 @@ mjModel* mj_compile(mjSpec* s, const mjVFS* vfs) {
 // recompile spec into existing model and data while preserving the state
 void mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d) {
   mjCModel* modelC = static_cast<mjCModel*>(s->element);
+  std::string state_name = "state";
   mjtNum time = 0;
   if (d) {
     time = d->time;
-    modelC->SaveState(d->qpos, d->qvel, d->act);
+    modelC->SaveState(state_name, d->qpos, d->qvel, d->act, d->ctrl, d->mocap_pos, d->mocap_quat);
   }
   modelC->Compile(vfs, &m);
   if (d) {
     modelC->MakeData(m, &d);
-    modelC->RestoreState(m->qpos0, d->qpos, d->qvel, d->act);
+    modelC->RestoreState(state_name, m->qpos0, m->body_pos, m->body_quat, d->qpos, d->qvel,
+                         d->act, d->ctrl, d->mocap_pos, d->mocap_quat);
     d->time = time;
   }
 }
@@ -604,6 +606,9 @@ const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* s
 
 // get id
 int mjs_getId(mjsElement* element) {
+  if (!element) {
+    return -1;
+  }
   return static_cast<mjCBase*>(element)->id;
 }
 

src/user/user_mesh.cc

@@ -1915,7 +1915,7 @@ void mjCMesh::MakeCenter(void) {
     // compute circumradius
     double norm_a_2 = mjuu_dot3(a, a);
     double norm_b_2 = mjuu_dot3(b, b);
-    double area = mjuu_normvec(nrm, 3);
+    double area = sqrt(mjuu_dot3(nrm, nrm));
 
     // compute circumcenter
     double res[3], vec[3] = {

src/user/user_model.cc

@@ -213,6 +213,7 @@ void mjCModel::SaveDofOffsets() {
   int qposadr = 0;
   int dofadr = 0;
   int actadr = 0;
+  int nmocap = 0;
 
   for (auto joint : joints_) {
     joint->qposadr_ = qposadr;
@@ -230,6 +231,15 @@ void mjCModel::SaveDofOffsets() {
     actuator->actadr_ = actuator->actdim_ ? actadr : -1;
     actadr += actuator->actdim_;
   }
+
+  for (mjCBody* body : bodies_) {
+    if (body->spec.mocap) {
+      body->mocapid = nmocap;
+      nmocap++;
+    } else {
+      body->mocapid = -1;
+    }
+  }
 }
 
 
@@ -315,6 +325,16 @@ void mjCModel::RemoveFromList(std::vector<T*>& list, const mjCModel& other) {
 
 
 
+template <>
+void mjCModel::DeleteAll<mjCKey>(std::vector<mjCKey*>& elements) {
+  for (mjCKey* element : elements) {
+    delete element;
+  }
+  elements.clear();
+}
+
+
+
 mjCModel& mjCModel::operator-=(const mjCBody& subtree) {
   mjCModel oldmodel(*this);
 
@@ -335,6 +355,10 @@ mjCModel& mjCModel::operator-=(const mjCBody& subtree) {
   MakeLists(world);
   ProcessLists(/*checkrepeat=*/false);
 
+  // store keyframes in the old model
+  oldmodel.key_pending_.clear();
+  oldmodel.StoreKeyframes();
+
   // check if we have to remove anything else
   RemoveFromList(pairs_, oldmodel);
   RemoveFromList(excludes_, oldmodel);
@@ -343,6 +367,12 @@ mjCModel& mjCModel::operator-=(const mjCBody& subtree) {
   RemoveFromList(actuators_, oldmodel);
   RemoveFromList(sensors_, oldmodel);
 
+  // move all keyframes to pending so that they will be resized
+  DeleteAll(keys_);
+  for (const auto& key : oldmodel.key_pending_) {
+    key_pending_.push_back(key);
+  }
+
   // restore to the original state
   if (!compiled) {
     ResetTreeLists();
@@ -1153,13 +1183,6 @@ void mjCModel::DeleteAll<mjCTexture>(std::vector<mjCTexture*>& elements) {
   elements.clear();
 }
 
-template <>
-void mjCModel::DeleteAll<mjCKey>(std::vector<mjCKey*>& elements) {
-  for (mjCKey* element : elements) {
-    delete element;
-  }
-  elements.clear();
-}
 
 // set nuser fields
 void mjCModel::SetNuser() {
@@ -2860,16 +2883,36 @@ void mjCModel::CopyObjects(mjModel* m) {
 
 // save the current state
 template <class T>
-void mjCModel::SaveState(const T* qpos, const T* qvel, const T* act) {
+void mjCModel::SaveState(const std::string& state_name, const T* qpos, const T* qvel, const T* act,
+                         const T* ctrl, const T* mpos, const T* mquat) {
   for (auto joint : joints_) {
-    if (qpos) mjuu_copyvec(joint->qpos, qpos + joint->qposadr_, joint->nq());
-    if (qvel) mjuu_copyvec(joint->qvel, qvel + joint->dofadr_, joint->nv());
+    if (joint->qposadr_ == -1 || joint->dofadr_ == -1) {
+      throw mjCError(NULL, "SaveState: joint %s has no address", joint->name.c_str());
+    }
+    if (qpos) mjuu_copyvec(joint->qpos(state_name), qpos + joint->qposadr_, joint->nq());
+    if (qvel) mjuu_copyvec(joint->qvel(state_name), qvel + joint->dofadr_, joint->nv());
   }
 
-  for (auto actuator : actuators_) {
-    if (actuator->actadr_ != -1 && act) {
-      actuator->act.assign(actuator->actdim_, 0);
-      mjuu_copyvec(actuator->act.data(), act + actuator->actadr_, actuator->actdim_);
+  for (unsigned int i=0; i<actuators_.size(); i++) {
+    auto actuator = actuators_[i];
+    if (actuator->actadr_ != -1 && actuator->actdim_ != -1 && act) {
+      actuator->act(state_name).assign(actuator->actdim_, 0);
+      mjuu_copyvec(actuator->act(state_name).data(), act + actuator->actadr_, actuator->actdim_);
+    }
+    if (ctrl) {
+      actuator->ctrl(state_name) = ctrl[i];
+    }
+  }
+
+  for (auto body : bodies_) {
+    if (!body->spec.mocap) {
+      continue;
+    }
+    if (mpos) {
+      mjuu_copyvec(body->mpos(state_name), mpos + 3*body->mocapid, 3);
+    }
+    if (mquat) {
+      mjuu_copyvec(body->mquat(state_name), mquat + 4*body->mocapid, 4);
     }
   }
 }
@@ -2890,59 +2933,93 @@ void mjCModel::MakeData(const mjModel* m, mjData** dest) {
 
 // restore the previous state
 template <class T>
-void mjCModel::RestoreState(const mjtNum* pos0, T* qpos, T* qvel, T* act) {
+void mjCModel::RestoreState(const std::string& state_name, const mjtNum* pos0,
+                            const mjtNum* mpos0, const mjtNum* mquat0, T* qpos,
+                            T* qvel, T* act, T* ctrl, T* mpos, T* mquat) {
   for (auto joint : joints_) {
     if (qpos) {
-      if (mjuu_defined(joint->qpos[0])) {
-        mjuu_copyvec(qpos + joint->qposadr_, joint->qpos, joint->nq());
+      if (mjuu_defined(joint->qpos(state_name)[0])) {
+        mjuu_copyvec(qpos + joint->qposadr_, joint->qpos(state_name), joint->nq());
       } else {
         mjuu_copyvec(qpos + joint->qposadr_, pos0 + joint->qposadr_, joint->nq());
       }
     }
-    if (mjuu_defined(joint->qvel[0]) && qvel) {
-      mjuu_copyvec(qvel + joint->dofadr_, joint->qvel, joint->nv());
+    if (mjuu_defined(joint->qvel(state_name)[0]) && qvel) {
+      mjuu_copyvec(qvel + joint->dofadr_, joint->qvel(state_name), joint->nv());
     }
   }
 
   // restore act
-  for (auto actuator : actuators_) {
-    if (mjuu_defined(actuator->act[0]) && act) {
-      mjuu_copyvec(act + actuator->actadr_, actuator->act.data(), actuator->actdim_);
+  for (unsigned int i=0; i<actuators_.size(); i++) {
+    auto actuator = actuators_[i];
+    if (!actuator->act(state_name).empty() && mjuu_defined(actuator->act(state_name)[0]) && act) {
+      mjuu_copyvec(act + actuator->actadr_, actuator->act(state_name).data(), actuator->actdim_);
+    }
+    if (ctrl) {
+      ctrl[i] = mjuu_defined(actuator->ctrl(state_name)) ? actuator->ctrl(state_name) : 0;
+    }
+  }
+
+  for (unsigned int i=0; i<bodies_.size(); i++) {
+    auto body = bodies_[i];
+    if (!body->mocap) {
+      continue;
+    }
+    if (mpos) {
+      if (mjuu_defined(body->mpos(state_name)[0])) {
+        mjuu_copyvec(mpos + 3*body->mocapid, body->mpos(state_name), 3);
+      } else {
+        mjuu_copyvec(mpos + 3*body->mocapid, mpos0 + 3*i, 3);
+      }
+    }
+    if (mquat) {
+      if (mjuu_defined(body->mquat(state_name)[0])) {
+        mjuu_copyvec(mquat + 4*body->mocapid, body->mquat(state_name), 4);
+      } else {
+        mjuu_copyvec(mquat + 4*body->mocapid, mquat0 + 4*i, 4);
+      }
     }
   }
 }
 
-
-
 // force explicit instantiations
-template void mjCModel::SaveState<mjtNum>(const mjtNum* qpos,
-                                          const mjtNum* qvel,
-                                          const mjtNum* act);
+template void mjCModel::SaveState<mjtNum>(
+    const std::string& name, const mjtNum* qpos, const mjtNum* qvel, const mjtNum* act,
+    const mjtNum* ctrl, const mjtNum* mpos, const mjtNum* mquat);
 
-template void mjCModel::RestoreState<mjtNum>(const mjtNum* qpos0, mjtNum* qpos,
-                                             mjtNum* qvel, mjtNum* act);
+template void mjCModel::RestoreState<mjtNum>(
+    const std::string& name, const mjtNum* qpos0, const mjtNum* mpos0, const mjtNum* mquat0,
+    mjtNum* qpos, mjtNum* qvel, mjtNum* act, mjtNum* ctrl, mjtNum* mpos, mjtNum* mquat);
 
 
 
 // resolve keyframe references
 void mjCModel::StoreKeyframes() {
   bool resetlists = false;
-  if (joints_.empty()) {
+
+  // create tree lists if they are empty, occurs if an uncompiled model is attached
+  if (bodies_.size() == 1 && geoms_.empty() && sites_.empty() && joints_.empty() &&
+      cameras_.empty() && lights_.empty() && frames_.empty()) {
     MakeLists(bodies_[0]);
     resetlists = true;
   }
 
   SaveDofOffsets();
 
-  for (auto key : keys_) {
+  for (auto& key : keys_) {
     mjKeyInfo info;
     info.name = prefix + key->name + suffix;
+    info.time = key->spec.time;
     info.qpos = !key->spec_qpos_.empty();
     info.qvel = !key->spec_qvel_.empty();
     info.act = !key->spec_act_.empty();
+    info.ctrl = !key->spec_ctrl_.empty();
+    info.mpos = !key->spec_mpos_.empty();
+    info.mquat = !key->spec_mquat_.empty();
     key_pending_.push_back(info);
-    SaveState(key->spec_qpos_.data(), key->spec_qvel_.data(), key->spec_act_.data());
-    break;  // (b/350784262) save only the first keyframe for now
+    SaveState(info.name, key->spec_qpos_.data(), key->spec_qvel_.data(),
+              key->spec_act_.data(), key->spec_ctrl_.data(),
+              key->spec_mpos_.data(), key->spec_mquat_.data());
   }
 
   if (resetlists) {
@@ -3484,7 +3561,10 @@ void mjCModel::ResolveKeyframes(const mjModel* m) {
     return;
   }
 
-  // resize non-pending keyframes to the new number of dofs
+  // store dof offsets in joints and actuators
+  SaveDofOffsets();
+
+  // resize existing keyframes to the new state, fill in missing default values
   for (unsigned int i = 0; i < nkey - key_pending_.size(); i++) {
     mjCKey* key = keys_[i];
     if (!key->spec_qpos_.empty()) {
@@ -3500,19 +3580,53 @@ void mjCModel::ResolveKeyframes(const mjModel* m) {
     if (!key->spec_act_.empty()) {
       key->spec_act_.resize(na);
     }
+    if (!key->spec_ctrl_.empty()) {
+      key->spec_ctrl_.resize(nu);
+    }
+    if (!key->spec_mpos_.empty()) {
+      int nmocap0 = key->spec_mpos_.size() / 3;
+      key->spec_mpos_.resize(3*nmocap);
+      for (unsigned int j = 0; j < bodies_.size(); j++) {
+        if (bodies_[j]->mocapid < nmocap0) {
+          continue;
+        }
+        int i = bodies_[j]->mocapid;
+        key->spec_mpos_[3*i+0] = (double)m->body_pos[3*j+0];
+        key->spec_mpos_[3*i+1] = (double)m->body_pos[3*j+1];
+        key->spec_mpos_[3*i+2] = (double)m->body_pos[3*j+2];
+      }
+    }
+    if (!key->spec_mquat_.empty()) {
+      int nmocap0 = key->spec_mquat_.size() / 4;
+      key->spec_mquat_.resize(4*nmocap);
+      for (unsigned int j = 0; j < bodies_.size(); j++) {
+        if (bodies_[j]->mocapid < nmocap0) {
+          continue;
+        }
+        int i = bodies_[j]->mocapid;
+        key->spec_mquat_[4*i+0] = (double)m->body_quat[4*j+0];
+        key->spec_mquat_[4*i+1] = (double)m->body_quat[4*j+1];
+        key->spec_mquat_[4*i+2] = (double)m->body_quat[4*j+2];
+        key->spec_mquat_[4*i+3] = (double)m->body_quat[4*j+3];
+      }
+    }
   }
 
-  // store dof offsets in joints and actuators
-  SaveDofOffsets();
-
-  // copy state stored in joints and actuators to keyframes
+  // create new keyframes, fill in missing default values
   for (const auto& info : key_pending_) {
     mjCKey* key = (mjCKey*)FindObject(mjOBJ_KEY, info.name);
     key->name = info.name;
+    key->spec.time = info.time;
     if (info.qpos) key->spec_qpos_.assign(nq, 0);
     if (info.qvel) key->spec_qvel_.assign(nv, 0);
     if (info.act) key->spec_act_.assign(na, 0);
-    RestoreState(m->qpos0, key->spec_qpos_.data(), key->spec_qvel_.data(), key->spec_act_.data());
+    if (info.ctrl) key->spec_ctrl_.assign(nu, 0);
+    if (info.mpos) key->spec_mpos_.assign(3*nmocap, 0);
+    if (info.mquat) key->spec_mquat_.assign(4*nmocap, 0);
+    RestoreState(info.name, m->qpos0, m->body_pos, m->body_quat,
+                 key->spec_qpos_.data(), key->spec_qvel_.data(),
+                 key->spec_act_.data(), key->spec_ctrl_.data(),
+                 key->spec_mpos_.data(), key->spec_mquat_.data());
   }
 
   // the attached keyframes have been copied into the model

src/user/user_model.h

@@ -36,9 +36,13 @@ typedef std::array<mjKeyMap, mjNOBJECT> mjListKeyMap;
 
 typedef struct mjKeyInfo_ {
   std::string name;
+  double time;
   bool qpos;
   bool qvel;
   bool act;
+  bool ctrl;
+  bool mpos;
+  bool mquat;
 } mjKeyInfo;
 
 class mjCModel_ : public mjsElement {
@@ -282,9 +286,15 @@ class mjCModel : public mjCModel_, private mjSpec {
   template <class T> void DeleteMaterial(std::vector<T*>& list,
                                          std::string_view name = "");
 
-  // save/restore the current state
-  template <class T> void SaveState(const T* qpos, const T* qvel, const T* act);
-  template <class T> void RestoreState(const mjtNum* pos0, T* qpos, T* qvel, T* act);
+  // save the current state
+  template <class T>
+  void SaveState(const std::string& state_name, const T* qpos, const T* qvel, const T* act,
+                 const T* ctrl, const T* mpos, const T* mquat);
+
+  // restore the previously saved state
+  template <class T>
+  void RestoreState(const std::string& state_name, const mjtNum* pos0, const mjtNum* mpos0,
+                    const mjtNum* mquat0, T* qpos, T* qvel, T* act, T* ctrl, T* mpos, T* mquat);
 
   // clear existing data
   void MakeData(const mjModel* m, mjData** dest);

src/user/user_objects.cc

@@ -937,9 +937,6 @@ mjCBody& mjCBody::operator-=(const mjCBody& subtree) {
     *bodies[i] -= subtree;
   }
 
-  // (b/350784262) delete keyframes
-  model->DeleteAll<mjCKey>(model->keys_);
-
   return *this;
 }
 
@@ -1408,9 +1405,11 @@ void mjCBody::ComputeBVH() {
 // reset keyframe references for allowing self-attach
 void mjCBody::ForgetKeyframes() const {
   for (auto joint : joints) {
-    joint->qpos[0] = mjNAN;
-    joint->qvel[0] = mjNAN;
+    joint->qpos_.clear();
+    joint->qvel_.clear();
   }
+  model->FindBody((mjCBody*)this, name)->mpos_.clear();  // this is a hack to avoid const
+  model->FindBody((mjCBody*)this, name)->mquat_.clear();  // this is a hack to avoid const
   for (auto body : bodies) {
     body->ForgetKeyframes();
   }
@@ -1418,6 +1417,24 @@ void mjCBody::ForgetKeyframes() const {
 
 
 
+mjtNum* mjCBody::mpos(const std::string& state_name) {
+  if (mpos_.find(state_name) == mpos_.end()) {
+    mpos_[state_name] = {mjNAN, 0, 0};
+  }
+  return mpos_.at(state_name).data();
+}
+
+
+
+mjtNum* mjCBody::mquat(const std::string& state_name) {
+  if (mquat_.find(state_name) == mquat_.end()) {
+    mquat_[state_name] = {mjNAN, 0, 0, 0};
+  }
+  return mquat_.at(state_name).data();
+}
+
+
+
 // compiler
 void mjCBody::Compile(void) {
   CopyFromSpec();
@@ -1769,8 +1786,8 @@ mjCJoint::mjCJoint(mjCModel* _model, mjCDef* _def) {
   CopyFromSpec();
 
   // no previous state when a joint is created
-  qpos[0] = mjNAN;
-  qvel[0] = mjNAN;
+  qposadr_ = -1;
+  dofadr_ = -1;
 }
 
 
@@ -1828,6 +1845,24 @@ int mjCJoint::nv(mjtJoint joint_type) {
 
 
 
+mjtNum* mjCJoint::qpos(const std::string& state_name) {
+  if (qpos_.find(state_name) == qpos_.end()) {
+    qpos_[state_name] = {mjNAN, 0, 0, 0, 0, 0, 0};
+  }
+  return qpos_.at(state_name).data();
+}
+
+
+
+mjtNum* mjCJoint::qvel(const std::string& state_name) {
+  if (qvel_.find(state_name) == qvel_.end()) {
+    qvel_[state_name] = {mjNAN, 0, 0, 0, 0, 0};
+  }
+  return qvel_.at(state_name).data();
+}
+
+
+
 void mjCJoint::PointToLocal() {
   spec.element = static_cast<mjsElement*>(this);
   spec.name = &name;
@@ -2090,15 +2125,15 @@ double mjCGeom::GetVolume() const {
     return mesh->GetVolumeRef(typeinertia);
   }
 
-  // compute from geom shape (type) and inertia (typeinertia)
+  // compute from geom shape (type) and inertia type (typeinertia)
   switch (type) {
     case mjGEOM_SPHERE: {
       double radius = size[0];
       switch (typeinertia) {
-        case mjINERTIA_SHELL:
-          return 4 * mjPI * radius * radius;
         case mjINERTIA_VOLUME:
           return 4 * mjPI * radius * radius * radius / 3;
+        case mjINERTIA_SHELL:
+          return 4 * mjPI * radius * radius;
       }
       break;
     }
@@ -2106,10 +2141,10 @@ double mjCGeom::GetVolume() const {
       double height = 2 * size[1];
       double radius = size[0];
       switch (typeinertia) {
-        case mjINERTIA_SHELL:
-          return 4 * mjPI * radius * radius + 2 * mjPI * radius * height;
         case mjINERTIA_VOLUME:
           return mjPI * (radius * radius * height + 4 * radius * radius * radius / 3);
+        case mjINERTIA_SHELL:
+          return 4 * mjPI * radius * radius + 2 * mjPI * radius * height;
       }
       break;
     }
@@ -2117,15 +2152,17 @@ double mjCGeom::GetVolume() const {
       double height = 2 * size[1];
       double radius = size[0];
       switch (typeinertia) {
-        case mjINERTIA_SHELL:
-          return 2 * mjPI * radius * radius + 2 * mjPI * radius * height;
         case mjINERTIA_VOLUME:
           return mjPI * radius * radius * height;
+        case mjINERTIA_SHELL:
+          return 2 * mjPI * radius * radius + 2 * mjPI * radius * height;
       }
       break;
     }
     case mjGEOM_ELLIPSOID: {
       switch (typeinertia) {
+        case mjINERTIA_VOLUME:
+          return 4 * mjPI * size[0] * size[1] * size[2] / 3;
         case mjINERTIA_SHELL: {
           // Thomsen approximation
           // https://www.numericana.com/answer/ellipsoid.htm#thomsen
@@ -2135,18 +2172,16 @@ double mjCGeom::GetVolume() const {
                        std::pow(size[2] * size[0], p);
           return 4 * mjPI * std::pow(tmp / 3, 1 / p);
         }
-        case mjINERTIA_VOLUME:
-          return 4 * mjPI * size[0] * size[1] * size[2] / 3;
       }
       break;
     }
     case mjGEOM_HFIELD:
     case mjGEOM_BOX: {
       switch (typeinertia) {
-        case mjINERTIA_SHELL:
-          return 8 * (size[0] * size[1] + size[1] * size[2] + size[2] * size[0]);
         case mjINERTIA_VOLUME:
           return size[0] * size[1] * size[2] * 8;
+        case mjINERTIA_SHELL:
+          return 8 * (size[0] * size[1] + size[1] * size[2] + size[2] * size[0]);
       }
       break;
     }
@@ -2185,16 +2220,16 @@ void mjCGeom::SetInertia(void) {
     return;
   }
 
-  // compute from geom shape (type) and inertia (typeinertia)
+  // compute from geom shape (type) and inertia type (typeinertia)
   switch (type) {
     case mjGEOM_SPHERE: {
       switch (typeinertia) {
-        case mjINERTIA_SHELL:
-          inertia[0] = inertia[1] = inertia[2] = 2 * mass_ * size[0] * size[0] / 3;
-          return;
         case mjINERTIA_VOLUME:
           inertia[0] = inertia[1] = inertia[2] = 2 * mass_ * size[0] * size[0] / 5;
           return;
+        case mjINERTIA_SHELL:
+          inertia[0] = inertia[1] = inertia[2] = 2 * mass_ * size[0] * size[0] / 3;
+          return;
       }
       break;
     }
@@ -2203,6 +2238,22 @@ void mjCGeom::SetInertia(void) {
       double height = 2 * size[1];
       double radius = size[0];
       switch (typeinertia) {
+        case mjINERTIA_VOLUME: {
+          double sphere_mass =
+              mass_ * 4 * radius / (4 * radius + 3 * height);  // mass*(sphere_vol/total_vol)
+          double cylinder_mass = mass_ - sphere_mass;
+
+          // cylinder part
+          inertia[0] = inertia[1] = cylinder_mass * (3 * radius * radius + height * height) / 12;
+          inertia[2] = cylinder_mass * radius * radius / 2;
+
+          // add two hemispheres, displace along third axis
+          double sphere_inertia = 2 * sphere_mass * radius * radius / 5;
+          inertia[0] += sphere_inertia + sphere_mass * height * (3 * radius + 2 * height) / 8;
+          inertia[1] += sphere_inertia + sphere_mass * height * (3 * radius + 2 * height) / 8;
+          inertia[2] += sphere_inertia;
+          return;
+        }
         case mjINERTIA_SHELL: {
           // surface area
           double Asphere = 4 * mjPI * radius * radius;
@@ -2226,22 +2277,6 @@ void mjCGeom::SetInertia(void) {
           inertia[2] += sphere_inertia;
           return;
         }
-        case mjINERTIA_VOLUME: {
-          double sphere_mass =
-              mass_ * 4 * radius / (4 * radius + 3 * height);  // mass*(sphere_vol/total_vol)
-          double cylinder_mass = mass_ - sphere_mass;
-
-          // cylinder part
-          inertia[0] = inertia[1] = cylinder_mass * (3 * radius * radius + height * height) / 12;
-          inertia[2] = cylinder_mass * radius * radius / 2;
-
-          // add two hemispheres, displace along third axis
-          double sphere_inertia = 2 * sphere_mass * radius * radius / 5;
-          inertia[0] += sphere_inertia + sphere_mass * height * (3 * radius + 2 * height) / 8;
-          inertia[1] += sphere_inertia + sphere_mass * height * (3 * radius + 2 * height) / 8;
-          inertia[2] += sphere_inertia;
-          return;
-        }
         break;
       }
       break;
@@ -2251,6 +2286,10 @@ void mjCGeom::SetInertia(void) {
       double height = 2 * halfheight;
       double radius = size[0];
       switch (typeinertia) {
+        case mjINERTIA_VOLUME:
+          inertia[0] = inertia[1] = mass_ * (3 * radius * radius + height * height) / 12;
+          inertia[2] = mass_ * radius * radius / 2;
+          return;
         case mjINERTIA_SHELL: {
           // surface area
           double Adisk = mjPI * radius * radius;
@@ -2276,10 +2315,6 @@ void mjCGeom::SetInertia(void) {
           inertia[2] += 2 * inertia_disk_z;
           return;
         }
-        case mjINERTIA_VOLUME:
-          inertia[0] = inertia[1] = mass_ * (3 * radius * radius + height * height) / 12;
-          inertia[2] = mass_ * radius * radius / 2;
-          return;
       }
       break;
     }
@@ -2288,6 +2323,12 @@ void mjCGeom::SetInertia(void) {
       double s11 = size[1] * size[1];
       double s22 = size[2] * size[2];
       switch (typeinertia) {
+        case mjINERTIA_VOLUME: {
+          inertia[0] = mass_ * (s11 + s22) / 5;
+          inertia[1] = mass_ * (s00 + s22) / 5;
+          inertia[2] = mass_ * (s00 + s11) / 5;
+          return;
+        }
         case mjINERTIA_SHELL: {
           // approximate shell inertia by subtracting ellipsoid from expanded ellipsoid
           double eps = 1e-6;
@@ -2323,12 +2364,6 @@ void mjCGeom::SetInertia(void) {
           inertia[2] = inertia_b[2] - inertia_a[2];
           return;
         }
-        case mjINERTIA_VOLUME: {
-          inertia[0] = mass_ * (s11 + s22) / 5;
-          inertia[1] = mass_ * (s00 + s22) / 5;
-          inertia[2] = mass_ * (s00 + s11) / 5;
-          return;
-        }
       }
       break;
     }
@@ -2338,6 +2373,12 @@ void mjCGeom::SetInertia(void) {
       double s11 = size[1] * size[1];
       double s22 = size[2] * size[2];
       switch (typeinertia) {
+        case mjINERTIA_VOLUME: {
+          inertia[0] = mass_ * (s11 + s22) / 3;
+          inertia[1] = mass_ * (s00 + s22) / 3;
+          inertia[2] = mass_ * (s00 + s11) / 3;
+          return;
+        }
         case mjINERTIA_SHELL: {
           // length
           double lx = 2 * size[0];  // side 0
@@ -2374,12 +2415,6 @@ void mjCGeom::SetInertia(void) {
           inertia[2] = 2 * (mass1 * s00 + mass2 * s11 + Iz0 + Iz1 + Iz2);
           return;
         }
-        case mjINERTIA_VOLUME: {
-          inertia[0] = mass_ * (s11 + s22) / 3;
-          inertia[1] = mass_ * (s00 + s22) / 3;
-          inertia[2] = mass_ * (s00 + s11) / 3;
-          return;
-        }
         break;
       }
       break;
@@ -5349,7 +5384,8 @@ mjCActuator::mjCActuator(mjCModel* _model, mjCDef* _def) {
   PointToLocal();
 
   // no previous state when an actuator is created
-  act.push_back(mjNAN);
+  actadr_ = -1;
+  actdim_ = -1;
 }
 
 
@@ -5374,8 +5410,8 @@ mjCActuator& mjCActuator::operator=(const mjCActuator& other) {
 
 
 void mjCActuator::ForgetKeyframes() {
-  act.clear();
-  act.push_back(mjNAN);
+  act_.clear();
+  ctrl_.clear();
 }
 
 
@@ -5386,6 +5422,24 @@ bool mjCActuator::is_actlimited() const { return islimited(actlimited, actrange)
 
 
 
+std::vector<mjtNum>& mjCActuator::act(const std::string& state_name) {
+  if (act_.find(state_name) == act_.end()) {
+    act_[state_name] = std::vector<mjtNum>(model->nu, mjNAN);
+  }
+  return act_.at(state_name);
+}
+
+
+
+mjtNum& mjCActuator::ctrl(const std::string& state_name) {
+  if (ctrl_.find(state_name) == ctrl_.end()) {
+    ctrl_[state_name] = mjNAN;
+  }
+  return ctrl_.at(state_name);
+}
+
+
+
 void mjCActuator::PointToLocal() {
   spec.element = static_cast<mjsElement*>(this);
   spec.name = &name;

src/user/user_objects.h

@@ -16,6 +16,7 @@
 #define MUJOCO_SRC_USER_USER_OBJECTS_H_
 
 #include <cstddef>
+#include <array>
 #include <cstdlib>
 #include <functional>
 #include <map>
@@ -254,6 +255,10 @@ class mjCBody_ : public mjCBase {
   std::string plugin_instance_name;
   std::vector<double> userdata_;
   std::vector<double> spec_userdata_;
+
+  // variables used for temporarily storing the state of the mocap bodies
+  std::map<std::string, std::array<mjtNum, 3>> mpos_;   // saved mocap_pos
+  std::map<std::string, std::array<mjtNum, 4>> mquat_;  // saved mocap_quat
 };
 
 class mjCBody : public mjCBody_, private mjsBody {
@@ -322,6 +327,10 @@ class mjCBody : public mjCBody_, private mjsBody {
   // reset keyframe references for allowing self-attach
   void ForgetKeyframes() const;
 
+  // get mocap position and quaternion
+  mjtNum* mpos(const std::string& state_name);
+  mjtNum* mquat(const std::string& state_name);
+
  private:
   mjCBody(const mjCBody& other, mjCModel* _model);  // copy constructor
   mjCBody& operator=(const mjCBody& other);         // copy assignment
@@ -401,10 +410,10 @@ class mjCJoint_ : public mjCBase {
   mjCBody* body;                   // joint's body
 
   // variable used for temporarily storing the state of the joint
-  int qposadr_;                    // address of dof in data->qpos
-  int dofadr_;                     // address of dof in data->qvel
-  mjtNum qpos[7];                  // qpos at the previous step
-  mjtNum qvel[6];                  // qvel at the previous step
+  int qposadr_;                                        // address of dof in data->qpos
+  int dofadr_;                                         // address of dof in data->qvel
+  std::map<std::string, std::array<mjtNum, 7>> qpos_;  // qpos at the previous step
+  std::map<std::string, std::array<mjtNum, 6>> qvel_;  // qvel at the previous step
 
   // variable-size data
   std::vector<double> userdata_;
@@ -443,6 +452,9 @@ class mjCJoint : public mjCJoint_, private mjsJoint {
   int nq() const { return nq(spec.type); }
   int nv() const { return nv(spec.type); }
 
+  mjtNum* qpos(const std::string& state_name);
+  mjtNum* qvel(const std::string& state_name);
+
  private:
   int Compile(void);               // compiler; return dofnum
   void PointToLocal(void);
@@ -1390,9 +1402,10 @@ class mjCActuator_ : public mjCBase {
   int trnid[2];                   // id of transmission target
 
   // variable used for temporarily storing the state of the actuator
-  int actadr_;              // address of dof in data->act
-  int actdim_;              // number of dofs in data->act
-  std::vector<mjtNum> act;  // act at the previous step
+  int actadr_;                                      // address of dof in data->act
+  int actdim_;                                      // number of dofs in data->act
+  std::map<std::string, std::vector<mjtNum>> act_;  // act at the previous step
+  std::map<std::string, mjtNum> ctrl_;              // ctrl at the previous step
 
   // variable-size data
   std::string plugin_name;
@@ -1431,6 +1444,9 @@ class mjCActuator : public mjCActuator_, private mjsActuator {
   bool is_forcelimited() const;
   bool is_actlimited() const;
 
+  std::vector<mjtNum>& act(const std::string& state_name);
+  mjtNum& ctrl(const std::string& state_name);
+
  private:
   void Compile(void);                       // compiler
   void CopyFromSpec();

test/engine/engine_collision_gjk_test.cc

@@ -50,8 +50,10 @@ void mjccd_support(const void *obj, const ccd_vec3_t *_dir, ccd_vec3_t *vec) {
 mjtNum run_gjk(mjModel* m, mjData* d, int g1, int g2, mjtNum x1[3],
                mjtNum x2[3]) {
   mjCCDConfig config = {kMaxIterations, kTolerance};
-  mjCCDObj obj1 = {m, d, g1, -1, -1, -1, -1, 0, {1, 0, 0, 0}, {0, 0, 0}};
-  mjCCDObj obj2 = {m, d, g2, -1, -1, -1, -1, 0, {1, 0, 0, 0}, {0, 0, 0}};
+  mjCCDObj obj1 = {m, d, g1, -1, -1, -1, -1, 0, {1, 0, 0, 0}, {0, 0, 0},
+                   mjc_center, mjc_support};
+  mjCCDObj obj2 = {m, d, g2, -1, -1, -1, -1, 0, {1, 0, 0, 0}, {0, 0, 0},
+                   mjc_center, mjc_support};
   mjc_center(obj1.x0, &obj1);
   mjc_center(obj2.x0, &obj2);
   mjtNum dist = mj_gjk(&config, &obj1, &obj2);
@@ -63,8 +65,10 @@ mjtNum run_gjk(mjModel* m, mjData* d, int g1, int g2, mjtNum x1[3],
 
 mjtNum run_gjkPenetration(mjModel* m, mjData* d, int g1, int g2,
                           mjtNum dir[3] = nullptr, mjtNum pos[3] = nullptr) {
-  mjCCDObj obj1 = {m, d, g1, -1, -1, -1, -1, 0, {1, 0, 0, 0}, {0, 0, 0}};
-  mjCCDObj obj2 = {m, d, g2, -1, -1, -1, -1, 0, {1, 0, 0, 0}, {0, 0, 0}};
+  mjCCDObj obj1 = {m, d, g1, -1, -1, -1, -1, 0, {1, 0, 0, 0}, {0, 0, 0},
+                   mjc_center, mjc_support};
+  mjCCDObj obj2 = {m, d, g2, -1, -1, -1, -1, 0, {1, 0, 0, 0}, {0, 0, 0},
+                   mjc_center, mjc_support};
   ccd_t ccd;
   ccd.mpr_tolerance = kTolerance;
   ccd.epa_tolerance = kTolerance;
@@ -198,6 +202,31 @@ TEST_F(MjGjkTest, EllipsoidEllipsoid) {
   mj_deleteModel(model);
 }
 
+TEST_F(MjGjkTest, EllipsoidEllipsoidIntersect) {
+  static constexpr char xml[] = R"(
+  <mujoco>
+  <worldbody>
+    <geom name="geom1" type="ellipsoid" pos="1.5 0 -.5" size="2.25 4.5 3"/>
+    <geom name="geom2" type="ellipsoid" pos="1.5 .5 .5" size="1.5 1.5 2.25"/>
+  </worldbody>
+  </mujoco>)";
+
+  std::array<char, 1000> error;
+  mjModel* model = LoadModelFromString(xml, error.data(), error.size());
+  ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data();
+
+  mjData* data = mj_makeData(model);
+  mj_forward(model, data);
+
+  int geom1 = mj_name2id(model, mjOBJ_GEOM, "geom1");
+  int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2");
+  mjtNum dist = run_gjk(model, data, geom1, geom2, nullptr, nullptr);
+
+  EXPECT_NEAR(dist, 0, kTolerance);
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
 TEST_F(MjGjkTest, CapsuleCapsule) {
   static constexpr char xml[] = R"(
   <mujoco>
@@ -218,7 +247,7 @@ TEST_F(MjGjkTest, CapsuleCapsule) {
   int geom2 = mj_name2id(model, mjOBJ_GEOM, "geom2");
   mjtNum dist = run_gjk(model, data, geom1, geom2, nullptr, nullptr);
 
-  EXPECT_NEAR(dist, 0.4765, .0001);
+  EXPECT_NEAR(dist, 0.4711, .0001);
   mj_deleteData(data);
   mj_deleteModel(model);
 }

test/engine/engine_derivative_test.cc

@@ -223,6 +223,85 @@ TEST_F(DerivativeTest, DisabledActuators) {
   mj_deleteModel(m1);
 }
 
+// actuator order has no effect
+TEST_F(DerivativeTest, ActuatorOrder) {
+  // model with stateful actuator first
+  static constexpr char xml1[] = R"(
+  <mujoco>
+    <option integrator="implicitfast"/>
+
+    <worldbody>
+      <body>
+        <joint name="0" type="slide" range="-1 1"/>
+        <geom size=".1"/>
+      </body>
+      <body pos="1 0 0">
+        <joint name="1" type="slide" range="-1 1"/>
+        <geom size=".1"/>
+      </body>
+    </worldbody>
+
+    <actuator>
+      <muscle joint="0" ctrlrange="0 6"/>
+      <damper joint="1" kv="200" ctrlrange="0 6"/>
+    </actuator>
+  </mujoco>
+  )";
+
+  char error[1024];
+  mjModel* m1 = LoadModelFromString(xml1, error, sizeof(error));
+  ASSERT_THAT(m1, NotNull()) << "Failed to load model: " << error;
+  mjData* d1 = mj_makeData(m1);
+
+  d1->ctrl[0] = 6;
+  d1->ctrl[1] = 6;
+
+  while (d1->time < 1)
+    mj_step(m1, d1);
+
+  // model with stateful actuator second
+  static constexpr char xml2[] = R"(
+  <mujoco>
+    <option integrator="implicitfast"/>
+
+    <worldbody>
+      <body>
+        <joint name="0" type="slide" range="-1 1"/>
+        <geom size=".1"/>
+      </body>
+      <body pos="1 0 0">
+        <joint name="1" type="slide" range="-1 1"/>
+        <geom size=".1"/>
+      </body>
+    </worldbody>
+
+    <actuator>
+      <damper joint="1" kv="200" ctrlrange="0 6"/>
+      <muscle joint="0" ctrlrange="0 6"/>
+    </actuator>
+  </mujoco>
+  )";
+
+  mjModel* m2 = LoadModelFromString(xml2, error, sizeof(error));
+  ASSERT_THAT(m2, NotNull()) << "Failed to load model: " << error;
+  mjData* d2 = mj_makeData(m2);
+
+  d2->ctrl[0] = 6;
+  d2->ctrl[1] = 6;
+
+  while (d2->time < 1)
+    mj_step(m2, d2);
+
+  // expect same qvel in both models
+  EXPECT_EQ(d1->qvel[0], d2->qvel[0]);
+  EXPECT_EQ(d1->qvel[1], d2->qvel[1]);
+
+  mj_deleteData(d2);
+  mj_deleteModel(m2);
+  mj_deleteData(d1);
+  mj_deleteModel(m1);
+}
+
 // compare analytic and fin-diff d_qfrc_passive/d_qvel
 TEST_F(DerivativeTest, PassiveDvel) {
   for (const char* local_path : {kTumblingThinObjectPath,

test/user/user_api_test.cc

@@ -20,6 +20,7 @@
 #include <filesystem>
 #include <memory>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include <gmock/gmock.h>
@@ -104,6 +105,41 @@ TEST_F(MujocoTest, TreeTraversal) {
   mj_deleteSpec(spec);
 }
 
+TEST_F(PluginTest, ActivatePlugin) {
+  std::string plugin_name = "mujoco.elasticity.cable";
+  mjSpec* spec = mj_makeSpec();
+
+  // get slot of requested plugin
+  int plugin_slot = -1;
+  const mjpPlugin* plugin = mjp_getPlugin(plugin_name.c_str(), &plugin_slot);
+  EXPECT_THAT(plugin, NotNull());
+
+  // activated plugin in the slot
+  std::vector<std::pair<const mjpPlugin*, int>> active_plugins;
+  active_plugins.emplace_back(std::make_pair(plugin, plugin_slot));
+  mjs_setActivePlugins(spec, &active_plugins);
+
+  // associate plugin to body
+  mjsBody* body = mjs_addBody(mjs_findBody(spec, "world"), 0);
+  mjs_setString(body->plugin.name, plugin_name.c_str());
+  body->plugin.instance = mjs_addPlugin(spec)->instance;
+  body->plugin.active = true;
+  mjsGeom* geom = mjs_addGeom(body, 0);
+  geom->type = mjGEOM_BOX;
+  geom->size[0] = 1;
+  geom->size[1] = 1;
+  geom->size[2] = 1;
+
+  // compile and check that the plugin is present
+  mjModel* model = mj_compile(spec, NULL);
+  EXPECT_THAT(model, NotNull());
+  EXPECT_THAT(model->nplugin, 1);
+  EXPECT_THAT(model->body_plugin[1], 0);
+
+  mj_deleteSpec(spec);
+  mj_deleteModel(model);
+}
+
 // ------------------- test recompilation multiple files -----------------------
 TEST_F(PluginTest, RecompileCompare) {
   mjtNum tol = 0;
@@ -430,8 +466,8 @@ static constexpr char xml_child[] = R"(
     </contact>
 
     <keyframe>
-      <key name="two" qpos="2" act="2 2"/>
-      <key name="three" qpos="3" act="3 3"/>
+      <key name="two" time="2" qpos="2" act="2 2" ctrl="2 2"/>
+      <key name="three" time="3" qpos="3" act="3 3" ctrl="3 3"/>
     </keyframe>
   </mujoco>)";
 
@@ -501,9 +537,10 @@ TEST_F(MujocoTest, AttachSame) {
     </contact>
 
     <keyframe>
-      <key name="two" qpos="2 0" act="2 2 0 0"/>
-      <key name="three" qpos="3 0" act="3 3 0 0"/>
-      <key name="attached-two-1" qpos="0 2" act="0 0 2 2"/>
+      <key name="two" time="2" qpos="2 0" act="2 2 0 0" ctrl="2 2 0 0"/>
+      <key name="three" time="3" qpos="3 0" act="3 3 0 0" ctrl="3 3 0 0"/>
+      <key name="attached-two-1" time="2" qpos="0 2" act="0 0 2 2" ctrl="0 0 2 2"/>
+      <key name="attached-three-1" time="3" qpos="0 3" act="0 0 3 3" ctrl="0 0 3 3"/>
     </keyframe>
   </mujoco>)";
 
@@ -620,7 +657,8 @@ TEST_F(MujocoTest, AttachDifferent) {
     </contact>
 
     <keyframe>
-      <key name="attached-two-1" qpos="0 0 0 1 0 0 0 2" act="2 2"/>
+      <key name="attached-two-1" time="2" qpos="0 0 0 1 0 0 0 2" act="2 2" ctrl="2 2"/>
+      <key name="attached-three-1" time="3" qpos="0 0 0 1 0 0 0 3" act="3 3" ctrl="3 3"/>
     </keyframe>
   </mujoco>)";
 
@@ -741,7 +779,8 @@ TEST_F(MujocoTest, AttachFrame) {
     </contact>
 
     <keyframe>
-      <key name="attached-two-1" qpos="0 0 0 1 0 0 0 2" act="2 2"/>
+      <key name="attached-two-1" time="2" qpos="0 0 0 1 0 0 0 2" act="2 2" ctrl="2 2"/>
+      <key name="attached-three-1" time="3" qpos="0 0 0 1 0 0 0 3" act="3 3" ctrl="3 3"/>
     </keyframe>
   </mujoco>)";
 
@@ -813,6 +852,11 @@ void TestDetachBody(bool compile) {
     <sensor>
       <framepos name="ignore" objtype="body" objname="ignore"/>
     </sensor>
+
+    <keyframe>
+      <key name="two" time="2"/>
+      <key name="three" time="3"/>
+    </keyframe>
   </mujoco>)";
 
   // model with one cylinder and a hinge
@@ -867,6 +911,8 @@ TEST_F(MujocoTest, PreserveState) {
         <joint type="slide" axis="0 0 1" name="slide"/>
         <geom type="sphere" size=".2"/>
       </body>
+      <body name="mocap_detach" mocap="true"/>
+      <body name="mocap" mocap="true"/>
     </worldbody>
     <actuator>
       <position name="hinge" joint="hinge" timeconst=".01"/>
@@ -885,6 +931,7 @@ TEST_F(MujocoTest, PreserveState) {
         <joint type="slide" axis="0 0 1"/>
         <geom type="sphere" size=".3"/>
       </body>
+      <body name="mocap" mocap="true"/>
     </worldbody>
     <actuator>
       <position name="slide" joint="slide" timeconst=".01"/>
@@ -912,6 +959,14 @@ TEST_F(MujocoTest, PreserveState) {
   data->ctrl[1] = 2;
   d_expected->ctrl[0] = 2;
 
+  // set mocap
+  data->mocap_pos[3] = 1;
+  data->mocap_quat[4] = 0;
+  data->mocap_quat[5] = 1;
+  d_expected->mocap_pos[0] = 1;
+  d_expected->mocap_quat[0] = 0;
+  d_expected->mocap_quat[1] = 1;
+
   // step models
   mj_step(model, data);
   mj_step(m_expected, d_expected);
@@ -922,6 +977,11 @@ TEST_F(MujocoTest, PreserveState) {
   EXPECT_THAT(body, NotNull());
   EXPECT_THAT(mjs_detachBody(spec, body), 0);
 
+  // detach mocap
+  mjsBody* mocap_body = mjs_findBody(spec, "mocap_detach");
+  EXPECT_THAT(mocap_body, NotNull());
+  EXPECT_THAT(mjs_detachBody(spec, mocap_body), 0);
+
   // add body
   mjsBody* newbody = mjs_addBody(mjs_findBody(spec, "world"), 0);
   EXPECT_THAT(newbody, NotNull());
@@ -965,6 +1025,17 @@ TEST_F(MujocoTest, PreserveState) {
     EXPECT_EQ(data->act[i], d_expected->act[i]) << i;
   }
 
+  // compare mocap
+  EXPECT_EQ(model->nmocap, m_expected->nmocap);
+  for (int i = 0; i < model->nmocap; ++i) {
+    for (int j = 0; j < 3; ++j) {
+      EXPECT_EQ(data->mocap_pos[3*i+j], d_expected->mocap_pos[3*i+j]) << i;
+    }
+    for (int j = 0; j < 4; ++j) {
+      EXPECT_EQ(data->mocap_quat[4*i+j], d_expected->mocap_quat[4*i+j]) << i;
+    }
+  }
+
   // check that the function is callable with no data
   mj_deleteData(data);
   mj_recompile(spec, 0, model, nullptr);
@@ -976,5 +1047,68 @@ TEST_F(MujocoTest, PreserveState) {
   mj_deleteModel(m_expected);
 }
 
+TEST_F(MujocoTest, AttachMocap) {
+  std::array<char, 1000> er;
+  mjtNum tol = 0;
+  std::string field = "";
+
+  static constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+      <body pos="1 1 1" quat="0 1 0 0" name="mocap" mocap="true"/>
+    </worldbody>
+    <keyframe>
+      <key name="key" time="1" mpos="2 2 2" mquat="1 0 0 0"/>
+    </keyframe>
+  </mujoco>)";
+
+  static constexpr char xml_expected[] = R"(
+  <mujoco>
+    <worldbody>
+      <body pos="1 1 1" quat="0 1 0 0" name="mocap" mocap="true"/>
+      <body pos="3 3 3" quat="0 0 1 0" name="attached-mocap-1" mocap="true"/>
+    </worldbody>
+    <keyframe>
+      <key name="key" time="1" mpos="2 2 2 3 3 3" mquat="1 0 0 0 0 0 1 0"/>
+      <key name="attached-key-1" time="1" mpos="1 1 1 2 2 2" mquat="0 1 0 0 1 0 0 0"/>
+    </keyframe>
+  </mujoco>)";
+
+  mjSpec* spec = mj_parseXMLString(xml, 0, er.data(), er.size());
+  EXPECT_THAT(spec, NotNull()) << er.data();
+
+  mjsBody* body = mjs_findBody(spec, "mocap");
+  EXPECT_THAT(body, NotNull());
+
+  mjsBody* world = mjs_findBody(spec, "world");
+  EXPECT_THAT(world, NotNull());
+
+  mjsFrame* frame = mjs_addFrame(world, NULL);
+  mjs_attachBody(frame, body, "attached-", "-1");
+
+  mjsBody* attached_body = mjs_findBody(spec, "attached-mocap-1");
+  EXPECT_THAT(attached_body, NotNull());
+  attached_body->pos[0] = 3;
+  attached_body->pos[1] = 3;
+  attached_body->pos[2] = 3;
+  attached_body->quat[0] = 0;
+  attached_body->quat[1] = 0;
+  attached_body->quat[2] = 1;
+  attached_body->quat[3] = 0;
+
+  mjModel* model = mj_compile(spec, 0);
+  EXPECT_THAT(model, NotNull());
+
+  mjModel* m_expected = LoadModelFromString(xml_expected, er.data(), er.size());
+  EXPECT_THAT(m_expected, NotNull()) << er.data();
+  EXPECT_LE(CompareModel(model, m_expected, field), tol)
+            << "Expected and attached models are different!\n"
+            << "Different field: " << field << '\n';
+
+  mj_deleteSpec(spec);
+  mj_deleteModel(model);
+  mj_deleteModel(m_expected);
+}
+
 }  // namespace
 }  // namespace mujoco