From a74c184f9d6ead05b0c2193a79746ee0fa19c019 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Thu, 15 Aug 2024 10:32:52 -0700 Subject: [PATCH] Add efc_pos to MJX. fixes #1388 PiperOrigin-RevId: 663358494 Change-Id: I63885f9208d12edbaa4b7d123ce4c6d16f40dcbe --- doc/changelog.rst | 4 ++++ mjx/mujoco/mjx/_src/constraint.py | 8 ++++---- mjx/mujoco/mjx/_src/constraint_test.py | 2 ++ mjx/mujoco/mjx/_src/io.py | 10 +++++++++- mjx/mujoco/mjx/_src/types.py | 2 ++ 5 files changed, 21 insertions(+), 5 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 80294d25..8a29045c 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -11,6 +11,10 @@ General 2. :ref:`shellinertia ` is now supported by all geom types. 3. Added support for :ref:`attaching` keyframes. +MJX +^^^ +4. Added ``efc_pos`` to ``mjx.Data``. + Version 3.2.2 (Aug 8, 2024) --------------------------- diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py index eb16863d..8f0bbf92 100644 --- a/mjx/mujoco/mjx/_src/constraint.py +++ b/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 diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py index db13e699..07a57475 100644 --- a/mjx/mujoco/mjx/_src/constraint_test.py +++ b/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') diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 51001559..d10a4760 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/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]) diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index d5c3bdcc..568cec28 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/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: