From caaf7b3a69d674c98572c0244dce1081abe49ca1 Mon Sep 17 00:00:00 2001 From: Erik Frey Date: Mon, 9 Jun 2025 14:26:18 -0700 Subject: [PATCH] Add tendon armature to MJX. PiperOrigin-RevId: 769294284 Change-Id: Idfbd4f355eb26fba5035973ab989a9d3af05c85f --- doc/changelog.rst | 4 + mjx/mujoco/mjx/__init__.py | 2 + mjx/mujoco/mjx/_src/forward.py | 2 + mjx/mujoco/mjx/_src/smooth.py | 152 +++++++++++++++++++++++++++++ mjx/mujoco/mjx/_src/smooth_test.py | 53 ++++++++++ mjx/mujoco/mjx/_src/support.py | 37 +++++++ 6 files changed, 250 insertions(+) diff --git a/doc/changelog.rst b/doc/changelog.rst index 8efd9e54..f82052b1 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -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) ------------------------------ diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index c60c785c..d71d6a24 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/mjx/mujoco/mjx/__init__.py @@ -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 diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py index 62535122..d8faf293 100644 --- a/mjx/mujoco/mjx/_src/forward.py +++ b/mjx/mujoco/mjx/_src/forward.py @@ -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 diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py index f87c2ed3..fba3d3ce 100644 --- a/mjx/mujoco/mjx/_src/smooth.py +++ b/mjx/mujoco/mjx/_src/smooth.py @@ -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.') diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index 96470a21..6c038522 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -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(""" + + + + + + + + + + + + + + + + + + + + + + + + + """) + + 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() diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index 15715021..ba62d9e3 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/mjx/mujoco/mjx/_src/support.py @@ -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,