Add efc_pos to MJX.

fixes #1388

PiperOrigin-RevId: 663358494
Change-Id: I63885f9208d12edbaa4b7d123ce4c6d16f40dcbe

doc/changelog.rst

@@ -11,6 +11,10 @@ General
 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``.
+
 Version 3.2.2 (Aug 8, 2024)
 ---------------------------
 

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)
+  aref, r, pos = fn(efc)
-  d = d.replace(efc_J=efc.J, efc_D=1 / r, efc_aref=aref)
+  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

@@ -270,6 +270,7 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
       contact=contact,
       efc_type=efc_type,
       efc_J=jp.zeros((nefc, m.nv), dtype=float),
+      efc_pos=jp.zeros((nefc,), dtype=float),
       efc_frictionloss=jp.zeros((nefc,), dtype=float),
       efc_D=jp.zeros((nefc,), dtype=float),
       efc_aref=jp.zeros((nefc,), dtype=float),
@@ -461,7 +462,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])

mjx/mujoco/mjx/_src/types.py

@@ -1137,6 +1137,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,)
@@ -1257,6 +1258,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: