Add tendon armature to MJX.
PiperOrigin-RevId: 769294284 Change-Id: Idfbd4f355eb26fba5035973ab989a9d3af05c85f
This commit is contained in:
Erik Frey
committed by
Copybara-Service
Copybara-Service
parent
b8768aa1cd
commit
caaf7b3a69
@@ -36,6 +36,10 @@ Python bindings
|
||||
|
||||
- Added examples of procedural terrain generation to the Model Editing tutorial: |mjspec_colab|
|
||||
|
||||
MJX
|
||||
^^^
|
||||
- Added tendon armature.
|
||||
|
||||
Version 3.3.2 (April 28, 2025)
|
||||
------------------------------
|
||||
|
||||
|
||||
@@ -48,6 +48,8 @@ from mujoco.mjx._src.smooth import rne
|
||||
from mujoco.mjx._src.smooth import rne_postconstraint
|
||||
from mujoco.mjx._src.smooth import subtree_vel
|
||||
from mujoco.mjx._src.smooth import tendon
|
||||
from mujoco.mjx._src.smooth import tendon_armature
|
||||
from mujoco.mjx._src.smooth import tendon_bias
|
||||
from mujoco.mjx._src.smooth import transmission
|
||||
from mujoco.mjx._src.solver import solve
|
||||
from mujoco.mjx._src.support import apply_ft
|
||||
|
||||
@@ -73,6 +73,7 @@ def fwd_position(m: Model, d: Data) -> Data:
|
||||
d = smooth.camlight(m, d)
|
||||
d = smooth.tendon(m, d)
|
||||
d = smooth.crb(m, d)
|
||||
d = smooth.tendon_armature(m, d)
|
||||
d = smooth.factor_m(m, d)
|
||||
d = collision_driver.collision(m, d)
|
||||
d = constraint.make_constraint(m, d)
|
||||
@@ -93,6 +94,7 @@ def fwd_velocity(m: Model, d: Data) -> Data:
|
||||
d = smooth.com_vel(m, d)
|
||||
d = passive.passive(m, d)
|
||||
d = smooth.rne(m, d)
|
||||
d = smooth.tendon_bias(m, d)
|
||||
return d
|
||||
|
||||
|
||||
|
||||
@@ -1184,3 +1184,155 @@ def transmission(m: Model, d: Data) -> Data:
|
||||
{'_impl.actuator_length': length, '_impl.actuator_moment': moment}
|
||||
)
|
||||
return d
|
||||
|
||||
|
||||
def tendon_armature(m: Model, d: Data) -> Data:
|
||||
"""Add tendon armature to qM."""
|
||||
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
|
||||
raise ValueError('tendon_armature requires JAX backend implementation.')
|
||||
|
||||
if not support.is_sparse(m):
|
||||
return d.tree_replace({
|
||||
'_impl.qM': (
|
||||
d._impl.qM
|
||||
+ d._impl.ten_J.T
|
||||
@ jax.vmap(jp.multiply)(d._impl.ten_J, m.tendon_armature)
|
||||
)
|
||||
})
|
||||
else:
|
||||
# TODO(taylorhowell): implement tendon armature with sparse qM
|
||||
raise NotImplementedError(
|
||||
'Tendon armature with sparse qM is not implemented.'
|
||||
)
|
||||
|
||||
|
||||
def tendon_dot(m: Model, d: Data) -> jax.Array:
|
||||
"""Compute time derivative of dense tendon Jacobian for one tendon."""
|
||||
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
|
||||
raise ValueError('tendon_dot requires JAX backend implementation.')
|
||||
|
||||
ten_Jdot = jp.zeros((m.ntendon, m.nv)) # pylint: disable=invalid-name
|
||||
|
||||
if not m.ntendon:
|
||||
return ten_Jdot
|
||||
|
||||
# process pulleys
|
||||
(wrap_id_pulley,) = np.nonzero(m.wrap_type == WrapType.PULLEY)
|
||||
|
||||
divisor = np.ones(m.nwrap)
|
||||
for adr, num in zip(m.tendon_adr, m.tendon_num):
|
||||
for id_pulley in wrap_id_pulley:
|
||||
if adr <= id_pulley < adr + num:
|
||||
divisor[id_pulley : adr + num] = np.maximum(
|
||||
mujoco.mjMINVAL, m.wrap_prm[id_pulley]
|
||||
)
|
||||
|
||||
# process spatial tendon sites
|
||||
(wrap_id_site,) = np.nonzero(m.wrap_type == WrapType.SITE)
|
||||
|
||||
# find consecutive sites, skipping tendon transitions
|
||||
(pair_id,) = np.nonzero(np.diff(wrap_id_site) == 1)
|
||||
wrap_id_site_pair = np.setdiff1d(wrap_id_site[pair_id], m.tendon_adr[1:] - 1)
|
||||
wrap_objid_site0 = m.wrap_objid[wrap_id_site_pair]
|
||||
wrap_objid_site1 = m.wrap_objid[wrap_id_site_pair + 1]
|
||||
site_bodyid0 = m.site_bodyid[wrap_objid_site0]
|
||||
site_bodyid1 = m.site_bodyid[wrap_objid_site1]
|
||||
site_xpos0 = d.site_xpos[wrap_objid_site0]
|
||||
site_xpos1 = d.site_xpos[wrap_objid_site1]
|
||||
subtree_com0 = d.subtree_com[m.body_rootid[site_bodyid0]]
|
||||
subtree_com1 = d.subtree_com[m.body_rootid[site_bodyid1]]
|
||||
site_vel0 = jax.vmap(lambda a, b: a[3:] - jp.cross(b, a[:3]))(
|
||||
d.cvel[site_bodyid0], site_xpos0 - subtree_com0
|
||||
)
|
||||
site_vel1 = jax.vmap(lambda a, b: a[3:] - jp.cross(b, a[:3]))(
|
||||
d.cvel[site_bodyid1], site_xpos1 - subtree_com1
|
||||
)
|
||||
|
||||
@jax.vmap
|
||||
def _momentdot(wpnt0, wpnt1, wvel0, wvel1, body0, body1):
|
||||
# dpnt = 3D position difference, normalize
|
||||
dpnt = wpnt1 - wpnt0
|
||||
norm = math.norm(dpnt)
|
||||
dpnt = jp.where(
|
||||
norm < mujoco.mjMINVAL, jp.array([1.0, 0.0, 0.0]), dpnt / norm
|
||||
)
|
||||
|
||||
# dvel = d / dt(dpnt)
|
||||
dvel = wvel1 - wvel0
|
||||
dot = jp.dot(dpnt, dvel)
|
||||
dvel += dpnt * -dot
|
||||
dvel = jp.where(norm > mujoco.mjMINVAL, dvel / norm, 0.0)
|
||||
|
||||
# get endpoint JacobianDots, subtract
|
||||
jacp1, _ = support.jac_dot(m, d, wpnt0, body0)
|
||||
jacp2, _ = support.jac_dot(m, d, wpnt1, body1)
|
||||
jacdif = jacp2 - jacp1
|
||||
|
||||
# chain rule, first term: Jdot += d / dt(jac2 - jac1) * dpnt
|
||||
tmp0 = jacdif @ dpnt
|
||||
|
||||
# get endpoint Jacobians, subtract
|
||||
jacp1, _ = support.jac(m, d, wpnt0, body0)
|
||||
jacp2, _ = support.jac(m, d, wpnt1, body1)
|
||||
jacdif = jacp2 - jacp1
|
||||
|
||||
# chain rule, second term: Jdot += (jac2 - jac1) * d/dt (dpnt)
|
||||
tmp1 = jacdif @ dvel
|
||||
|
||||
return jp.where(body0 != body1, tmp0 + tmp1, jp.zeros(m.nv))
|
||||
|
||||
momentdots = _momentdot(
|
||||
site_xpos0,
|
||||
site_xpos1,
|
||||
site_vel0,
|
||||
site_vel1,
|
||||
site_bodyid0,
|
||||
site_bodyid1,
|
||||
)
|
||||
|
||||
if wrap_id_site_pair.size:
|
||||
divisor_site_pair = divisor[wrap_id_site_pair]
|
||||
momentdots /= divisor_site_pair[:, None]
|
||||
|
||||
tendon_nsite = np.array([
|
||||
sum((wrap_id_site_pair >= adr) & (wrap_id_site_pair < adr + num))
|
||||
for adr, num in zip(m.tendon_adr, m.tendon_num)
|
||||
])
|
||||
tendon_has_site = tendon_nsite > 0
|
||||
(tendon_id_site,) = np.nonzero(tendon_has_site)
|
||||
tendon_nsite = tendon_nsite[tendon_has_site]
|
||||
tendon_with_site = tendon_nsite.size
|
||||
ten_site_id = np.repeat(np.arange(tendon_with_site), tendon_nsite)
|
||||
|
||||
momentdot = jax.ops.segment_sum(momentdots, ten_site_id, tendon_with_site)
|
||||
ten_Jdot = ten_Jdot.at[tendon_id_site].set(momentdot) # pylint: disable=invalid-name
|
||||
|
||||
# TODO(taylorhowell): time derivatives for geoms
|
||||
|
||||
return ten_Jdot
|
||||
|
||||
|
||||
def tendon_bias(m: Model, d: Data) -> Data:
|
||||
"""Add bias force due to tendon armature."""
|
||||
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
|
||||
raise ValueError('tendon_bias requires JAX backend implementation.')
|
||||
|
||||
if not m.ntendon:
|
||||
return d
|
||||
|
||||
# get dense d/dt(tendon Jacobian) for each tendon
|
||||
ten_Jdot = tendon_dot(m, d) # pylint: disable=invalid-name
|
||||
|
||||
# add bias term: qfrc += ten_J * armature * ten_Jdot @ qvel
|
||||
coef = m.tendon_armature * jp.dot(ten_Jdot, d.qvel)
|
||||
|
||||
if not support.is_sparse(m):
|
||||
return d.tree_replace({
|
||||
'qfrc_bias': (
|
||||
d.qfrc_bias
|
||||
+ jp.sum(jax.vmap(jp.multiply)(d._impl.ten_J, coef), axis=0)
|
||||
)
|
||||
})
|
||||
else:
|
||||
# TODO(taylorhowell): implement tendon bias with sparse qM
|
||||
raise NotImplementedError('Tendon bias with sparse qM is not implemented.')
|
||||
|
||||
@@ -17,6 +17,7 @@
|
||||
from absl.testing import absltest
|
||||
from absl.testing import parameterized
|
||||
import jax
|
||||
from jax import numpy as jp
|
||||
import mujoco
|
||||
from mujoco import mjx
|
||||
from mujoco.mjx._src import test_util
|
||||
@@ -316,6 +317,58 @@ class TendonTest(parameterized.TestCase):
|
||||
_assert_eq(d.wrap_obj, dx._impl.wrap_obj, 'wrap_obj')
|
||||
_assert_eq(d.wrap_xpos, dx._impl.wrap_xpos, 'wrap_xpos')
|
||||
|
||||
def test_tendon_armature(self):
|
||||
"""Tests MJX tendon armature matches MuJoCo."""
|
||||
m = mujoco.MjModel.from_xml_string("""
|
||||
<mujoco>
|
||||
<worldbody>
|
||||
<site name="site0" pos="1 0 1"/>
|
||||
<body>
|
||||
<joint type="slide" axis="0 0 1"/>
|
||||
<joint type="hinge" axis="0 1 0"/>
|
||||
<geom type="box" size="0.1 0.1 0.1" mass="1" pos="1 0 0"/>
|
||||
<site name="site1"/>
|
||||
</body>
|
||||
</worldbody>
|
||||
<tendon>
|
||||
<spatial armature="123">
|
||||
<site site="site0"/>
|
||||
<site site="site1"/>
|
||||
</spatial>
|
||||
<spatial armature="456">
|
||||
<site site="site0"/>
|
||||
<site site="site1"/>
|
||||
</spatial>
|
||||
</tendon>
|
||||
<keyframe>
|
||||
<key qpos="1.2345 1.2345" qvel="1.2345 1.2345"/>
|
||||
</keyframe>
|
||||
</mujoco>
|
||||
""")
|
||||
|
||||
d = mujoco.MjData(m)
|
||||
mujoco.mj_resetDataKeyframe(m, d, 0)
|
||||
mujoco.mj_forward(m, d)
|
||||
|
||||
qM = np.zeros((m.nv, m.nv)) # pylint: disable=invalid-name
|
||||
mujoco.mj_fullM(m, qM, d.qM)
|
||||
|
||||
mx = mjx.put_model(m)
|
||||
dx = mjx.put_data(m, d)
|
||||
|
||||
dx = dx.tree_replace(
|
||||
{'_impl.qM': jp.zeros((m.nv, m.nv)), 'qfrc_bias': jp.zeros(m.nv)}
|
||||
)
|
||||
|
||||
dx = mjx.crb(mx, dx)
|
||||
dx = mjx.tendon_armature(mx, dx)
|
||||
|
||||
_assert_eq(dx._impl.qM, qM, 'qM')
|
||||
|
||||
dx = mjx.rne(mx, dx)
|
||||
dx = mjx.tendon_bias(mx, dx)
|
||||
_assert_eq(dx.qfrc_bias, d.qfrc_bias, 'qfrc_bias')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
|
||||
@@ -142,6 +142,7 @@ def jac(
|
||||
m: Model, d: Data, point: jax.Array, body_id: jax.Array
|
||||
) -> Tuple[jax.Array, jax.Array]:
|
||||
"""Compute pair of (NV, 3) Jacobians of global point attached to body."""
|
||||
# TODO(taylorhowell): statically construct mask
|
||||
fn = lambda carry, b: b if carry is None else b + carry
|
||||
mask = (jp.arange(m.nbody) == body_id) * 1
|
||||
mask = scan.body_tree(m, fn, 'b', 'b', mask, reverse=True)
|
||||
@@ -155,6 +156,42 @@ def jac(
|
||||
return jacp, jacr
|
||||
|
||||
|
||||
def jac_dot(
|
||||
m: Model, d: Data, point: jax.Array, body_id: jax.Array
|
||||
) -> Tuple[jax.Array, jax.Array]:
|
||||
"""Compute pair of (NV, 3) Jacobian time derivatives of global point attached to body."""
|
||||
# TODO(taylorhowell): statically construct mask
|
||||
fn = lambda carry, b: b if carry is None else b + carry
|
||||
mask = (jp.arange(m.nbody) == body_id) * 1
|
||||
mask = scan.body_tree(m, fn, 'b', 'b', mask, reverse=True)
|
||||
mask = mask[jp.array(m.dof_bodyid)] > 0
|
||||
|
||||
offset = point - d.subtree_com[jp.array(m.body_rootid)[body_id]]
|
||||
pvel_lin = d.cvel[body_id][3:] - jp.cross(offset, d.cvel[body_id][:3])
|
||||
|
||||
cdof = d._impl.cdof
|
||||
cdof_dot = d._impl.cdof_dot
|
||||
|
||||
# check for quaternion
|
||||
jnt_type = m.jnt_type[m.dof_jntid]
|
||||
dof_adr = m.jnt_dofadr[m.dof_jntid]
|
||||
is_quat = (jnt_type == JointType.BALL) | (
|
||||
jnt_type == JointType.FREE & (np.arange(m.nv) >= dof_adr + 3)
|
||||
)
|
||||
|
||||
# compute cdof_dot for quaternion (use current body cvel)
|
||||
cdof_dot_quat = jax.vmap(math.motion_cross)(d.cvel[m.dof_bodyid], cdof)
|
||||
cdof_dot = jp.where(is_quat[:, None], cdof_dot_quat, cdof_dot)
|
||||
|
||||
jacp = jax.vmap(
|
||||
lambda a, b: a[3:] + jp.cross(a[:3], offset) + jp.cross(b[:3], pvel_lin)
|
||||
)(cdof_dot, cdof)
|
||||
jacp = jax.vmap(jp.multiply)(jacp, mask)
|
||||
jacr = jax.vmap(jp.multiply)(cdof_dot[:, :3], mask) # pytype: disable=attribute-error
|
||||
|
||||
return jacp, jacr
|
||||
|
||||
|
||||
def apply_ft(
|
||||
m: Model,
|
||||
d: Data,
|
||||
|
||||
Reference in New Issue
Block a user