Add sites to MJX.

PiperOrigin-RevId: 585691582
Change-Id: I6cc594daf436231f91534f137ff259e7b39811d9

doc/changelog.rst

@@ -23,6 +23,11 @@ Plugins
 - Added the `pid <https://github.com/deepmind/mujoco/blob/main/plugin/actuator/README.md>`__ actuator plugin, a
   configurable PID controller that implements the Integral term, which is not available with native MuJoCo actuators.
 
+MJX
+^^^
+
+- Added ``site_xpos`` and ``site_xmat`` to MJX.
+
 Bug fixes
 ^^^^^^^^^
 - Fix bug in Cartesian actuation with movable refsite, as when using body-centric Cartesian actuators on a quadruped.

mjx/mujoco/mjx/_src/device.py

@@ -67,6 +67,7 @@ _TRANSFORMS = {
     (types.Data, 'ximat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
     (types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
     (types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
+    (types.Data, 'site_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
     (types.Model, 'actuator_trnid'): lambda x: x[:, 0],
     (types.Contact, 'frame'): (
         lambda x: x.reshape(x.shape[:-1] + (3, 3))  # pylint: disable=g-long-lambda
@@ -78,6 +79,7 @@ _INVERSE_TRANSFORMS = {
     (types.Data, 'ximat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
     (types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
     (types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
+    (types.Data, 'site_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
     (types.Contact, 'frame'): (
         lambda x: x.reshape(x.shape[:-2] + (9,))  # pylint: disable=g-long-lambda
         if x is not None and x.shape[0] else jp.zeros((0, 9))

mjx/mujoco/mjx/_src/io.py

@@ -56,6 +56,8 @@ def make_data(m: Model) -> Data:
       xaxis=jp.zeros((m.njnt, 3), dtype=jp.float32),
       geom_xpos=jp.zeros((m.ngeom, 3), dtype=jp.float32),
       geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=jp.float32),
+      site_xpos=jp.zeros((m.nsite, 3), dtype=jp.float32),
+      site_xmat=jp.zeros((m.nsite, 3, 3), dtype=jp.float32),
       subtree_com=jp.zeros((m.nbody, 3), dtype=jp.float32),
       cdof=jp.zeros((m.nv, 6), dtype=jp.float32),
       cinert=jp.zeros((m.nbody, 10), dtype=jp.float32),

mjx/mujoco/mjx/_src/smooth.py

@@ -101,13 +101,20 @@ def kinematics(m: Model, d: Data) -> Data:
 
   # TODO(erikfrey): confirm that quats are more performant for mjx than mats
   xipos, ximat = local_to_global(xpos, xquat, m.body_ipos, m.body_iquat)
-  geom_xpos, geom_xmat = local_to_global(
-      xpos[m.geom_bodyid], xquat[m.geom_bodyid], m.geom_pos, m.geom_quat
-  )
-
   d = d.replace(qpos=qpos, xanchor=xanchor, xaxis=xaxis, xpos=xpos)
   d = d.replace(xquat=xquat, xmat=xmat, xipos=xipos, ximat=ximat)
-  d = d.replace(geom_xpos=geom_xpos, geom_xmat=geom_xmat)
+
+  if m.ngeom:
+    geom_xpos, geom_xmat = local_to_global(
+        xpos[m.geom_bodyid], xquat[m.geom_bodyid], m.geom_pos, m.geom_quat
+    )
+    d = d.replace(geom_xpos=geom_xpos, geom_xmat=geom_xmat)
+
+  if m.nsite:
+    site_xpos, site_xmat = local_to_global(
+        xpos[m.site_bodyid], xquat[m.site_bodyid], m.site_pos, m.site_quat
+    )
+    d = d.replace(site_xpos=site_xpos, site_xmat=site_xmat)
 
   return d
 

mjx/mujoco/mjx/_src/smooth_test.py

@@ -162,5 +162,40 @@ class DisableGravityTest(absltest.TestCase):
     )
 
 
+class SiteTest(absltest.TestCase):
+
+  def test_site(self):
+    """Tests that site positions and orientations match MuJoCo."""
+    m = mujoco.MjModel.from_xml_string("""
+      <mujoco>
+        <worldbody>
+          <site name="origin"/>
+          <body>
+            <joint type="free"/>
+            <geom pos="1 0 1" type="box" size="0.1 0.01 0.01"/>
+            <site name="s1" pos="1.3 0 0"/>     <!-- pos only -->
+            <site name="s2"/>                   <!-- no pos, no quat -->
+            <site name="s3" quat="1 0 1 0"/>    <!- quat only -->
+            <site name="s4" pos="0 1.5 0" quat="1 0 1 0"/>
+            <site name="s5" pos="1 0 1"/>       <!-- same as ipos -->
+            <site/>
+          </body>
+        </worldbody>
+      </mujoco>
+      """)
+    d = mujoco.MjData(m)
+
+    mx = mjx.device_put(m)
+    dx = mjx.device_put(d)
+
+    mujoco.mj_forward(m, d)
+    dx = mjx.forward(mx, dx)
+
+    np.testing.assert_array_almost_equal(dx.site_xpos, d.site_xpos)
+    np.testing.assert_array_almost_equal(
+        dx.site_xmat, d.site_xmat.reshape((-1, 3, 3))
+    )
+
+
 if __name__ == '__main__':
   absltest.main()

mjx/mujoco/mjx/_src/types.py

@@ -260,6 +260,7 @@ class Model(PyTreeNode):
     nbody: number of bodies
     njnt: number of joints
     ngeom: number of geoms
+    nsite: number of sites
     nmesh: number of meshes
     npair: number of predefined geom pairs
     nexclude: number of excluded geom pairs
@@ -327,6 +328,9 @@ class Model(PyTreeNode):
     geom_friction: friction for (slide, spin, roll)           (ngeom, 3)
     geom_margin: include in solver if dist<margin-gap         (ngeom,)
     geom_gap: include in solver if dist<margin-gap            (ngeom,)
+    site_bodyid: id of site's body                            (nsite,)
+    site_pos: local position offset rel. to body              (nsite,3)
+    site_quat: local orientation offset rel. to body          (nsite,4)
     geom_convex_face: vertex face data, MJX only              (ngeom,)
     geom_convex_vert: vertex data, MJX only                   (ngeom,)
     geom_convex_edge: unique edge data, MJX only              (ngeom,)
@@ -377,6 +381,7 @@ class Model(PyTreeNode):
   nbody: int
   njnt: int
   ngeom: int
+  nsite: int
   nmesh: int
   npair: int
   nexclude: int
@@ -444,6 +449,9 @@ class Model(PyTreeNode):
   geom_friction: jax.Array
   geom_margin: jax.Array
   geom_gap: jax.Array
+  site_bodyid: np.ndarray
+  site_pos: jax.Array
+  site_quat: jax.Array
   pair_dim: np.ndarray
   pair_geom1: np.ndarray
   pair_geom2: np.ndarray
@@ -569,6 +577,8 @@ class Data(PyTreeNode):
     xaxis: Cartesian joint axis                                   (njnt, 3)
     geom_xpos: Cartesian geom position                            (ngeom, 3)
     geom_xmat: Cartesian geom orientation                         (ngeom, 3, 3)
+    site_xpos: Cartesian site position                            (nsite, 3)
+    site_xmat: Cartesian site orientation                         (nsite, 9)
     subtree_com: center of mass of each subtree                   (nbody, 3)
     cdof: com-based motion axis of each dof                       (nv, 6)
     cinert: com-based body inertia and mass                       (nbody, 10)
@@ -631,6 +641,8 @@ class Data(PyTreeNode):
   xaxis: jax.Array
   geom_xpos: jax.Array
   geom_xmat: jax.Array
+  site_xpos: jax.Array
+  site_xmat: jax.Array
   subtree_com: jax.Array
   cdof: jax.Array
   cinert: jax.Array