From cd8ff4401dfaf05b063606be1c8b210d55bc14ca Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Wed, 4 Sep 2024 21:59:35 -0700 Subject: [PATCH] Add eq_objtype to mjx.Model and support for connect constraint using sites in MJX. PiperOrigin-RevId: 671225109 Change-Id: I8e2e6eb4c539d318658d3d7dfe0a7f03b411aabc --- doc/changelog.rst | 1 + mjx/mujoco/mjx/_src/collision_driver_test.py | 2 +- mjx/mujoco/mjx/_src/constraint.py | 35 ++++++++++++++++---- mjx/mujoco/mjx/_src/constraint_test.py | 14 ++++---- mjx/mujoco/mjx/_src/types.py | 2 ++ mjx/mujoco/mjx/test_data/constraints.xml | 14 ++++++-- 6 files changed, 51 insertions(+), 17 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index d422d1b0..950a4262 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -61,6 +61,7 @@ MJX :doc:`fluid drag `. - Fixed a bug where ``qLDiagInv`` had the wrong size for sparse mass matrices. - Added support for joint and tendon :ref:`frictionloss `. +- Added support for :ref:`connect` equality constraints using two sites. Bug fixes ^^^^^^^^^ diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py index 5492ea7f..f290b5b7 100644 --- a/mjx/mujoco/mjx/_src/collision_driver_test.py +++ b/mjx/mujoco/mjx/_src/collision_driver_test.py @@ -887,7 +887,7 @@ class DimTest(parameterized.TestCase): def test_ncon(self): m = test_util.load_test_file('constraints.xml') dim = collision_driver.make_condim(m) - expected = [1] * 4 + [3] * 20 + [4] * 4 + [6] * 4 + expected = [1] * 4 + [3] * 24 + [4] * 4 + [6] * 4 np.testing.assert_array_equal(dim, np.array(expected)) def test_disable_contact(self): diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py index 4c317539..431f665b 100644 --- a/mjx/mujoco/mjx/_src/constraint.py +++ b/mjx/mujoco/mjx/_src/constraint.py @@ -32,6 +32,7 @@ from mujoco.mjx._src.types import DisableBit from mujoco.mjx._src.types import EqType from mujoco.mjx._src.types import JointType from mujoco.mjx._src.types import Model +from mujoco.mjx._src.types import ObjType # pylint: enable=g-importing-member import numpy as np @@ -106,25 +107,47 @@ def _efc_equality_connect(m: Model, d: Data) -> Optional[_Efc]: return None @jax.vmap - def rows(obj1id, obj2id, data, solref, solimp): + def rows(is_site, obj1id, obj2id, body1id, body2id, data, solref, solimp): anchor1, anchor2 = data[0:3], data[3:6] - # error is difference in global positions pos1 = d.xmat[obj1id] @ anchor1 + d.xpos[obj1id] pos2 = d.xmat[obj2id] @ anchor2 + d.xpos[obj2id] + + if m.nsite: + pos1 = jp.where(is_site, d.site_xpos[obj1id], pos1) + pos2 = jp.where(is_site, d.site_xpos[obj2id], pos2) + + # error is difference in global positions pos = pos1 - pos2 # compute Jacobian difference (opposite of contact: 0 - 1) - jacp1, _ = support.jac(m, d, pos1, obj1id) - jacp2, _ = support.jac(m, d, pos2, obj2id) + jacp1, _ = support.jac(m, d, pos1, body1id) + jacp2, _ = support.jac(m, d, pos2, body2id) j = (jacp1 - jacp2).T pos_imp = math.norm(pos) - invweight = m.body_invweight0[obj1id, 0] + m.body_invweight0[obj2id, 0] + invweight = m.body_invweight0[body1id, 0] + m.body_invweight0[body2id, 0] zero = jp.zeros_like(pos) return _row(j, pos, pos_imp, invweight, solref, solimp, zero, zero) - args = (m.eq_obj1id, m.eq_obj2id, m.eq_data, m.eq_solref, m.eq_solimp) + is_site = m.eq_objtype == ObjType.SITE + + body1id = np.copy(m.eq_obj1id) + body2id = np.copy(m.eq_obj2id) + if m.nsite: + body1id[is_site] = m.site_bodyid[body1id[is_site]] + body2id[is_site] = m.site_bodyid[body2id[is_site]] + + args = ( + is_site, + m.eq_obj1id, + m.eq_obj2id, + body1id, + body2id, + m.eq_data, + m.eq_solref, + m.eq_solimp, + ) args = jax.tree_util.tree_map(lambda x: x[eq_id], args) # concatenate to drop row grouping return jax.tree_util.tree_map(jp.concatenate, rows(*args)) diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py index 8951d76f..76bd9266 100644 --- a/mjx/mujoco/mjx/_src/constraint_test.py +++ b/mjx/mujoco/mjx/_src/constraint_test.py @@ -103,31 +103,31 @@ class ConstraintTest(parameterized.TestCase): self.assertEqual(ne, 0) self.assertEqual(nf, 2) self.assertEqual(nl, 5) - self.assertEqual(nc, 148) + self.assertEqual(nc, 164) dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m)) - self.assertEqual(dx.efc_J.shape[0], 155) # only joint/tendon limit, contact + self.assertEqual(dx.efc_J.shape[0], 171) # only joint/tendon limit, contact def test_disable_contact(self): m = test_util.load_test_file('constraints.xml') m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONTACT ne, nf, nl, nc = constraint.counts(constraint.make_efc_type(m)) - self.assertEqual(ne, 11) + self.assertEqual(ne, 14) self.assertEqual(nf, 2) self.assertEqual(nl, 5) self.assertEqual(nc, 0) dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m)) - self.assertEqual(dx.efc_J.shape[0], 18) # only equality, joint/tendon limit + self.assertEqual(dx.efc_J.shape[0], 21) # only equality, joint/tendon limit def test_disable_frictionloss(self): m = test_util.load_test_file('constraints.xml') m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.FRICTIONLOSS ne, nf, nl, nc = constraint.counts(constraint.make_efc_type(m)) - self.assertEqual(ne, 11) + self.assertEqual(ne, 14) self.assertEqual(nf, 0) self.assertEqual(nl, 5) - self.assertEqual(nc, 148) + self.assertEqual(nc, 164) dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m)) - self.assertEqual(dx.efc_J.shape[0], 164) + self.assertEqual(dx.efc_J.shape[0], 183) def test_margin(self): """Test margin.""" diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index cfa777db..ffe1ad5a 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -681,6 +681,7 @@ class Model(PyTreeNode): eq_type: constraint type (mjtEq) (neq,) eq_obj1id: id of object 1 (neq,) eq_obj2id: id of object 2 (neq,) + eq_objtype: type of both objects (mjtObj) (neq,) eq_active0: initial enable/disable constraint state (neq,) eq_solref: constraint solver reference (neq, mjNREF) eq_solimp: constraint solver impedance (neq, mjNIMP) @@ -991,6 +992,7 @@ class Model(PyTreeNode): eq_type: np.ndarray eq_obj1id: np.ndarray eq_obj2id: np.ndarray + eq_objtype: np.ndarray eq_active0: np.ndarray eq_solref: jax.Array eq_solimp: jax.Array diff --git a/mjx/mujoco/mjx/test_data/constraints.xml b/mjx/mujoco/mjx/test_data/constraints.xml index 08ee2b6a..ff5ff7a8 100644 --- a/mjx/mujoco/mjx/test_data/constraints.xml +++ b/mjx/mujoco/mjx/test_data/constraints.xml @@ -16,6 +16,13 @@ + + + + + + + @@ -76,6 +83,7 @@ + @@ -90,10 +98,10 @@ - + - + - +