chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Add eq_objtype to mjx.Model and support for connect constraint using sites in MJX.

Browse Source 
PiperOrigin-RevId: 671225109
Change-Id: I8e2e6eb4c539d318658d3d7dfe0a7f03b411aabc
This commit is contained in:
Taylor Howell
2024-09-04 21:59:35 -07:00
committed by Copybara-Service
parent 81618f95f9
commit cd8ff4401d
6 changed files with 51 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+1
View File
@@ -61,6 +61,7 @@ MJX
:doc:`fluid drag <computation/fluid>`.
- Fixed a bug where ``qLDiagInv`` had the wrong size for sparse mass matrices.
- Added support for joint and tendon :ref:`frictionloss <coFriction>`.
- Added support for :ref:`connect<equality-connect>` equality constraints using two sites.
Bug fixes
^^^^^^^^^
mjx/mujoco/mjx/_src/collision_driver_test.py
+1 -1
View File
@@ -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):
mjx/mujoco/mjx/_src/constraint.py
+29 -6
View File
@@ -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))
mjx/mujoco/mjx/_src/constraint_test.py
+7 -7
View File
@@ -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."""
mjx/mujoco/mjx/_src/types.py
+2
View File
@@ -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
mjx/mujoco/mjx/test_data/constraints.xml
+11 -3
View File
@@ -16,6 +16,13 @@
<worldbody>
<geom pos="0 0 -1" type="plane" size="20 20 .01" condim="1"/>
<site name="0" pos="0 0 1"/>
<body pos="0 0 1">
<freejoint/>
<geom class="box"/>
<site name="1"/>
</body>
<body name="anchor1" pos="-3 0 0"/>
<body name="beam1" pos="-3 0 0">
<joint name="joint1" type="ball" range="0 45" solreflimit="0.03 0.9" solimplimit="0.89 0.9 0.01 2.1"/>
@@ -76,6 +83,7 @@
</tendon>
<equality>
<connect name="site" site1="0" site2="1"/>
<connect name="connect" body1="anchor1" body2="beam1" anchor="1 0 -1" />
<weld name="weld" body1="anchor2" body2="beam2" relpose="0 0 0 1 -.3 0 0" torquescale="0.002" anchor="0 -2 0"/>
<joint name="joint" joint1="joint3" joint2="joint4" polycoef="0.5 -1 0.1 0.15 0.2" />
@@ -90,10 +98,10 @@
<keyframe>
<!-- keyframe 0: default position with some motion, zero contacts -->
<key qpos='1 0 0 0 -1 0 0 1 0 0 0 0 0 1 0 0 0 4 0 0 1 0 0 0 5 0 0 1 0 0 0 6 0 0 1 0 0 0 7 0 0 1 0 0 0' qvel='1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1'/>
<key qpos=' -1 -1 1 1 0 0 0 1 0 0 0 -1 0 0 1 0 0 0 0 0 1 0 0 0 4 0 0 1 0 0 0 5 0 0 1 0 0 0 6 0 0 1 0 0 0 7 0 0 1 0 0 0' qvel='1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1'/>
<!-- keyframe 1: some contacts but constraints are quadratic -->
<key qpos='1 0.0011 -8.2e-07 -0.0011 -0.82 1.9 -0.58 0.97 -0.14 0.17 -0.069 0.22 0.29 0.93 0.2 0.15 0.28 4.3 0.23 -0.26 0.97 0.12 0.11 0.18 5.4 0.25 -0.26 0.97 0.11 0.16 0.17 6.4 0.25 -0.26 0.97 0.11 0.16 0.17 7.4 0.25 -0.26 0.97 0.11 0.16 0.17' qvel='-0.14 3.3e-05 0.14 -0.14 -0.93 -3.6 -1.8 1 -0.57 -0.16 0.21 1.1 1 2.2 0.65 1.3 -4.2 -2.6 -5.6 0.54 1.6 2.8 -3.7 -2.3 0.7 -1.8 1.6 2.8 -3.7 -2.3 0.62 -1.8 1.6 2.8 -3.7 -2.3 0.61 -1.8'/>
<key qpos=' -1 -1 1 1 0 0 0 1 0.0011 -8.2e-07 -0.0011 -0.82 1.9 -0.58 0.97 -0.14 0.17 -0.069 0.22 0.29 0.93 0.2 0.15 0.28 4.3 0.23 -0.26 0.97 0.12 0.11 0.18 5.4 0.25 -0.26 0.97 0.11 0.16 0.17 6.4 0.25 -0.26 0.97 0.11 0.16 0.17 7.4 0.25 -0.26 0.97 0.11 0.16 0.17' qvel='1 1 1 1 1 1 -0.14 3.3e-05 0.14 -0.14 -0.93 -3.6 -1.8 1 -0.57 -0.16 0.21 1.1 1 2.2 0.65 1.3 -4.2 -2.6 -5.6 0.54 1.6 2.8 -3.7 -2.3 0.7 -1.8 1.6 2.8 -3.7 -2.3 0.62 -1.8 1.6 2.8 -3.7 -2.3 0.61 -1.8'/>
<!-- keyframe 2: some contacts and some constraints are in cone state (for elliptic) -->
<key qpos='1 0.0087 2.4e-07 -0.0086 -0.89 1.8 -0.77 0.98 -0.2 -0.0022 -0.026 0.19 0.33 0.86 0.32 0.064 0.38 4.4 0.36 -0.81 0.97 -0.0013 -0.0011 0.25 5.6 0.52 -0.75 0.98 -0.018 0.17 0.094 6.6 0.52 -0.75 0.98 -0.017 0.16 0.094 7.6 0.52 -0.76 0.98 -0.017 0.16 0.094' qvel='0.2 -1.8e-05 -0.2 -0.72 0.072 0.025 0.015 -4.9 0.35 -0.22 0.26 0.99 -4.6 1.7 1.1 0.52 -0.73 0.16 7.9 1 0.043 -0.042 1.3 1.6 -1.6 0.98 0.027 -0.024 1.3 1.6 -1.8 0.96 0.025 -0.022 1.3 1.6 -1.8 0.96'/>
<key qpos=' -1 -1 1 1 0 0 0 1 0.0087 2.4e-07 -0.0086 -0.89 1.8 -0.77 0.98 -0.2 -0.0022 -0.026 0.19 0.33 0.86 0.32 0.064 0.38 4.4 0.36 -0.81 0.97 -0.0013 -0.0011 0.25 5.6 0.52 -0.75 0.98 -0.018 0.17 0.094 6.6 0.52 -0.75 0.98 -0.017 0.16 0.094 7.6 0.52 -0.76 0.98 -0.017 0.16 0.094' qvel='1 1 1 1 1 1 0.2 -1.8e-05 -0.2 -0.72 0.072 0.025 0.015 -4.9 0.35 -0.22 0.26 0.99 -4.6 1.7 1.1 0.52 -0.73 0.16 7.9 1 0.043 -0.042 1.3 1.6 -1.6 0.98 0.027 -0.024 1.3 1.6 -1.8 0.96 0.025 -0.022 1.3 1.6 -1.8 0.96'/>
</keyframe>
</mujoco>