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

Add site semantics to weld constraint in MJX.

Browse Source 
PiperOrigin-RevId: 677823461
Change-Id: Iab20e042695dfad416d61fa7cc938faba8c7a6b6
This commit is contained in:
Taylor Howell
2024-09-23 09:24:12 -07:00
committed by Copybara-Service
parent 04f8403d8e
commit 76b48b8e61
4 changed files with 69 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
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] * 24 + [4] * 4 + [6] * 4
expected = [1] * 4 + [3] * 28 + [4] * 4 + [6] * 4
np.testing.assert_array_equal(dim, np.array(expected))
def test_disable_contact(self):
mjx/mujoco/mjx/_src/constraint.py
+47 -14
View File
@@ -110,8 +110,8 @@ def _efc_equality_connect(m: Model, d: Data) -> Optional[_Efc]:
def rows(is_site, obj1id, obj2id, body1id, body2id, data, solref, solimp):
anchor1, anchor2 = data[0:3], data[3:6]
pos1 = d.xmat[obj1id] @ anchor1 + d.xpos[obj1id]
pos2 = d.xmat[obj2id] @ anchor2 + d.xpos[obj2id]
pos1 = d.xmat[body1id] @ anchor1 + d.xpos[body1id]
pos2 = d.xmat[body2id] @ anchor2 + d.xpos[body2id]
if m.nsite:
pos1 = jp.where(is_site, d.site_xpos[obj1id], pos1)
@@ -131,12 +131,12 @@ def _efc_equality_connect(m: Model, d: Data) -> Optional[_Efc]:
return _row(j, pos, pos_imp, invweight, solref, solimp, zero, zero)
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]]
body1id[is_site] = m.site_bodyid[m.eq_obj1id[is_site]]
body2id[is_site] = m.site_bodyid[m.eq_obj2id[is_site]]
args = (
is_site,
@@ -161,24 +161,40 @@ def _efc_equality_weld(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]
relpose, torquescale = data[6:10], data[10]
# error is difference in global position and orientation
pos1 = d.xmat[obj1id] @ anchor2 + d.xpos[obj1id]
pos2 = d.xmat[obj2id] @ anchor1 + d.xpos[obj2id]
pos1 = d.xmat[body1id] @ anchor2 + d.xpos[body1id]
pos2 = d.xmat[body2id] @ anchor1 + d.xpos[body2id]
if m.nsite:
pos1 = jp.where(is_site, d.site_xpos[obj1id], pos1)
pos2 = jp.where(is_site, d.site_xpos[obj2id], pos2)
cpos = pos1 - pos2
# compute Jacobian difference (opposite of contact: 0 - 1)
jacp1, jacr1 = support.jac(m, d, pos1, obj1id)
jacp2, jacr2 = support.jac(m, d, pos2, obj2id)
jacp1, jacr1 = support.jac(m, d, pos1, body1id)
jacp2, jacr2 = support.jac(m, d, pos2, body2id)
jacdifp = jacp1 - jacp2
jacdifr = (jacr1 - jacr2) * torquescale
# compute orientation error: neg(q1) * q0 * relpose (axis components only)
quat = math.quat_mul(d.xquat[obj1id], relpose)
quat1 = math.quat_inv(d.xquat[obj2id])
quat = math.quat_mul(d.xquat[body1id], relpose)
quat1 = math.quat_inv(d.xquat[body2id])
if m.nsite:
quat = jp.where(
is_site, math.quat_mul(d.xquat[body1id], m.site_quat[obj1id]), quat
)
quat1 = jp.where(
is_site,
math.quat_inv(math.quat_mul(d.xquat[body2id], m.site_quat[obj2id])),
quat1,
)
crot = math.quat_mul(quat1, quat)[1:] # copy axis components
pos = jp.concatenate((cpos, crot * torquescale))
@@ -188,13 +204,30 @@ def _efc_equality_weld(m: Model, d: Data) -> Optional[_Efc]:
jacdifr = 0.5 * jax.vmap(jac_fn)(jacdifr)
j = jp.concatenate((jacdifp.T, jacdifr.T))
pos_imp = math.norm(pos)
invweight = m.body_invweight0[obj1id] + m.body_invweight0[obj2id]
invweight = m.body_invweight0[body1id] + m.body_invweight0[body2id]
invweight = jp.repeat(invweight, 3, axis=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[m.eq_obj1id[is_site]]
body2id[is_site] = m.site_bodyid[m.eq_obj2id[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, 164)
self.assertEqual(nc, 180)
dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
self.assertEqual(dx.efc_J.shape[0], 171) # only joint/tendon limit, contact
self.assertEqual(dx.efc_J.shape[0], 187) # 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, 14)
self.assertEqual(ne, 20)
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], 21) # only equality, joint/tendon limit
self.assertEqual(dx.efc_J.shape[0], 27) # 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, 14)
self.assertEqual(ne, 20)
self.assertEqual(nf, 0)
self.assertEqual(nl, 5)
self.assertEqual(nc, 164)
self.assertEqual(nc, 180)
dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
self.assertEqual(dx.efc_J.shape[0], 183)
self.assertEqual(dx.efc_J.shape[0], 205)
def test_margin(self):
"""Test margin."""
mjx/mujoco/mjx/test_data/constraints.xml
+14 -6
View File
@@ -16,11 +16,18 @@
<worldbody>
<geom pos="0 0 -1" type="plane" size="20 20 .01" condim="1"/>
<site name="0" pos="0 0 1"/>
<site name="site0" pos="0 0 1"/>
<body pos="0 0 1">
<freejoint/>
<geom class="box"/>
<site name="1"/>
<site name="site1"/>
</body>
<site name="site2" pos="0 -1 1"/>
<body pos="0 -1 1">
<freejoint/>
<geom class="box"/>
<site name="site3"/>
</body>
<body name="anchor1" pos="-3 0 0"/>
@@ -83,8 +90,9 @@
</tendon>
<equality>
<connect name="site" site1="0" site2="1"/>
<connect name="c_site" site1="site0" site2="site1"/>
<connect name="connect" body1="anchor1" body2="beam1" anchor="1 0 -1" />
<weld name="w_site" site1="site2" site2="site3" torquescale="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" />
<tendon name="tendon" tendon1="tendon_1" tendon2="tendon_2" polycoef="0.5 -1 0.1 0.15 0.2"/>
@@ -98,10 +106,10 @@
<keyframe>
<!-- keyframe 0: default position with some motion, zero contacts -->
<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'/>
<key qpos='-1 -1 1 1 0 0 0 -2 -2 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 1 1 1 1 1 1'/>
<!-- keyframe 1: some contacts but constraints are quadratic -->
<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'/>
<key qpos='-1 -1 1 1 0 0 0 -2 -2 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 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 -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'/>
<key qpos='-1 -1 1 1 0 0 0 -2 -2 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 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>