diff --git a/doc/changelog.rst b/doc/changelog.rst index 17164571..fd3269de 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -19,6 +19,7 @@ General for any object type. - The experimental use of 2D/3D elasticity plugins with :ref:`composite` has been removed. Users should instead use :ref:`flexcomp`, which provides the correct collision behavior. + - Added the :ref:`nativeccd` flag. When this flag is enabled, general convex collision detection is handled natively, as opposed to using `libccd `__. This feature is in early stages of testing. @@ -63,6 +64,7 @@ MJX - 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. +- Added support for :ref:`spatial tendons ` with site wrapping. Bug fixes ^^^^^^^^^ @@ -73,7 +75,7 @@ Bug fixes - Fixed a bug in tendon wrapping around spheres. Before this fix, tendons that wrapped around spheres with an externally-placed :ref:`sidesite` could jump inside the sphere instead of wrapping around it. - Fixed a bug that caused :at:`meshdir` and :at:`texturedir` to be overwritten during model - :ref:`attachment. + :ref:`attachment`. Python bindings ^^^^^^^^^^^^^^^ diff --git a/doc/mjx.rst b/doc/mjx.rst index 99fc1b8f..1a82c892 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -195,6 +195,8 @@ The following features are **fully supported** in MJX: - ``FIXED``, ``AFFINE`` * - :ref:`Actuator Bias ` - ``NONE``, ``AFFINE`` + * - :ref:`Tendon Wrapping ` + - ``JOINT``, ``SITE`` * - :ref:`Geom ` - ``PLANE``, ``HFIELD``, ``SPHERE``, ``CAPSULE``, ``BOX``, ``MESH`` are fully implemented. ``ELLIPSOID`` and ``CYLINDER`` are implemented but only collide with other primitives, note that ``BOX`` is implemented as a mesh. @@ -244,7 +246,7 @@ The following features are **in development** and coming soon: * - :ref:`Actuator Bias ` - ``MUSCLE`` * - :ref:`Tendon Wrapping ` - - ``NONE``, ``JOINT``, ``PULLEY``, ``SITE``, ``SPHERE``, ``CYLINDER`` + - ``PULLEY``, ``SPHERE``, ``CYLINDER`` * - Fluid Model - :ref:`flEllipsoid` * - :ref:`Tendons ` diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index fe4de217..3fa5c7b6 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -168,6 +168,7 @@ class ModelIOTest(parameterized.TestCase): + diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py index 6d503406..54de7b62 100644 --- a/mjx/mujoco/mjx/_src/smooth.py +++ b/mjx/mujoco/mjx/_src/smooth.py @@ -27,6 +27,7 @@ from mujoco.mjx._src.types import DisableBit from mujoco.mjx._src.types import JointType from mujoco.mjx._src.types import Model from mujoco.mjx._src.types import TrnType +from mujoco.mjx._src.types import WrapType # pylint: enable=g-importing-member import numpy as np @@ -686,13 +687,106 @@ def tendon(m: Model, d: Data) -> Data: if not m.ntendon: return d - ten_id = np.repeat(np.arange(m.ntendon), m.tendon_num) - length = m.wrap_prm * d.qpos[m.jnt_qposadr[m.wrap_objid]] - ten_length = jax.ops.segment_sum(length, ten_id, m.ntendon) - ten_j = jp.zeros((m.ntendon, m.nv)) - ten_j = ten_j.at[ten_id, m.jnt_dofadr[m.wrap_objid]].set(m.wrap_prm) + # process joint tendons + (wrap_id_jnt,) = np.nonzero(m.wrap_type == WrapType.JOINT) + (tendon_id_jnt,) = np.nonzero(np.isin(m.tendon_adr, wrap_id_jnt)) - return d.replace(ten_length=ten_length, ten_J=ten_j) + ntendon_jnt = tendon_id_jnt.size + wrap_objid_jnt = m.wrap_objid[wrap_id_jnt] + tendon_num_jnt = m.tendon_num[tendon_id_jnt] + + moment_jnt = m.wrap_prm[wrap_id_jnt] + length_jnt = jax.ops.segment_sum( + moment_jnt * d.qpos[m.jnt_qposadr[wrap_objid_jnt]], + np.repeat(np.arange(ntendon_jnt), tendon_num_jnt), + ntendon_jnt, + ) + + adr_moment_jnt = np.repeat(tendon_id_jnt, tendon_num_jnt) + dofadr_moment_jnt = m.jnt_dofadr[wrap_objid_jnt] + + # process spatial tendon sites + (wrap_id_site,) = np.nonzero(m.wrap_type == WrapType.SITE) + nwrap_site = wrap_id_site.size + + # 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) + (tendon_id_site,) = np.nonzero(np.isin(m.tendon_adr, wrap_id_site_pair)) + + id0 = m.wrap_objid[wrap_id_site_pair] + id1 = m.wrap_objid[wrap_id_site_pair + 1] + + @jax.vmap + def _length_moment(pnt0, pnt1, body0, body1): + dif = pnt1 - pnt0 + length = jp.linalg.norm(dif) + vec = jp.where( + length < mujoco.mjMINVAL, jp.array([1.0, 0.0, 0.0]), dif / length + ) + + jacp1, _ = support.jac(m, d, pnt0, body0) + jacp2, _ = support.jac(m, d, pnt1, body1) + jacdif = jacp2 - jacp1 + moment = jp.where(body0 != body1, jacdif @ vec, jp.zeros(m.nv)) + + return length, moment + + lengths_site, moments_site = _length_moment( + d.site_xpos[id0], d.site_xpos[id1], m.site_bodyid[id0], m.site_bodyid[id1] + ) + + 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_nsite = tendon_nsite[tendon_nsite > 0] + tendon_wrapnum_site = tendon_nsite + 1 + tendon_with_site = sum([s > 0 for s in tendon_nsite]) + ten_site_id = np.repeat(np.arange(tendon_with_site), tendon_nsite) + + length_site = jax.ops.segment_sum(lengths_site, ten_site_id, tendon_with_site) + moment_site = jax.ops.segment_sum(moments_site, ten_site_id, tendon_with_site) + + # assemble length and moment + ten_length = ( + jp.zeros_like(d.ten_length) + .at[np.concatenate([tendon_id_jnt, tendon_id_site])] + .set(jp.concatenate([length_jnt, length_site])) + ) + ten_moment = ( + jp.zeros_like(d.ten_J) + .at[adr_moment_jnt, dofadr_moment_jnt] + .set(moment_jnt) + ) + ten_moment = ten_moment.at[tendon_id_site].set(moment_site) + + # wrap + wrap_xpos = jp.concatenate([ + d.site_xpos[m.wrap_objid[wrap_id_site]], + jp.zeros((2 * m.nwrap - nwrap_site, 3)), + ]).reshape((m.nwrap, 6)) + + ten_wrapnum = np.zeros(m.ntendon) + ten_wrapnum[tendon_id_site] = tendon_wrapnum_site + + ten_wrapadr = [0] + for wn in ten_wrapnum[:-1]: + ten_wrapadr.append(ten_wrapadr[-1] + wn) + ten_wrapadr = np.array(ten_wrapadr).astype(int) + + wrap_obj = np.zeros(m.nwrap * 2, dtype=int) + wrap_obj[:nwrap_site] = -1 + wrap_obj = wrap_obj.reshape((-1, 2)) + + return d.replace( + ten_length=ten_length, + ten_J=ten_moment, + ten_wrapadr=jp.array(ten_wrapadr), + ten_wrapnum=jp.array(ten_wrapnum), + wrap_xpos=wrap_xpos, + wrap_obj=jp.array(wrap_obj), + ) def _site_dof_mask(m: Model) -> np.ndarray: diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index c94347ce..782c5d16 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -15,6 +15,7 @@ """Tests for smooth dynamics functions.""" from absl.testing import absltest +from absl.testing import parameterized import jax import mujoco from mujoco import mjx @@ -230,5 +231,36 @@ class SmoothTest(absltest.TestCase): _assert_eq(d.cfrc_int, dx.cfrc_int, 'cfrc_int') +class TendonTest(parameterized.TestCase): + + @parameterized.parameters( + 'tendon/fixed.xml', + 'tendon/site.xml', + 'tendon/fixed_site.xml', + 'tendon/no_tendon.xml', + ) + def test_tendon(self, filename): + """Tests MJX tendon function matches MuJoCo mj_tendon.""" + m = test_util.load_test_file(filename) + d = mujoco.MjData(m) + # give the system a little kick to ensure we have non-identity rotations + d.qvel = np.random.random(m.nv) + mujoco.mj_step(m, d, 10) # let dynamics get state significantly non-zero + mx = mjx.put_model(m) + dx = mjx.put_data(m, d) + + mujoco.mj_forward(m, d) + dx = jax.jit(mjx.forward)(mx, dx) + + _assert_eq(d.ten_length, dx.ten_length, 'ten_length') + _assert_eq(d.ten_J, dx.ten_J, 'ten_J') + _assert_eq(d.ten_wrapnum, dx.ten_wrapnum, 'ten_wrapnum') + _assert_eq(d.ten_wrapadr, dx.ten_wrapadr, 'ten_wrapadr') + if d.wrap_obj.shape == dx.wrap_obj.shape: + _assert_eq(d.wrap_obj, dx.wrap_obj, 'wrap_obj') + if d.wrap_xpos.shape == dx.wrap_xpos.shape: + _assert_eq(d.wrap_xpos, dx.wrap_xpos, 'wrap_xpos') + + if __name__ == '__main__': absltest.main() diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 6f1da229..6a200e1c 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -197,9 +197,11 @@ class WrapType(enum.IntEnum): Members: JOINT: constant moment arm + SITE: pass through site """ JOINT = mujoco.mjtWrap.mjWRAP_JOINT - # unsupported: NONE, PULLEY, SITE, SPHERE, CYLINDER + SITE = mujoco.mjtWrap.mjWRAP_SITE + # unsupported: NONE, PULLEY, SPHERE, CYLINDER class TrnType(enum.IntEnum): diff --git a/mjx/mujoco/mjx/test_data/tendon/fixed.xml b/mjx/mujoco/mjx/test_data/tendon/fixed.xml new file mode 100644 index 00000000..8e3b3e3f --- /dev/null +++ b/mjx/mujoco/mjx/test_data/tendon/fixed.xml @@ -0,0 +1,23 @@ + + + + + + + + + + + + + + + + + + + + + + + diff --git a/mjx/mujoco/mjx/test_data/tendon/fixed_site.xml b/mjx/mujoco/mjx/test_data/tendon/fixed_site.xml new file mode 100644 index 00000000..d8842b49 --- /dev/null +++ b/mjx/mujoco/mjx/test_data/tendon/fixed_site.xml @@ -0,0 +1,39 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/mjx/mujoco/mjx/test_data/tendon/no_tendon.xml b/mjx/mujoco/mjx/test_data/tendon/no_tendon.xml new file mode 100644 index 00000000..7eaf3ddf --- /dev/null +++ b/mjx/mujoco/mjx/test_data/tendon/no_tendon.xml @@ -0,0 +1,16 @@ + + + + + + + + + + + + + + + + diff --git a/mjx/mujoco/mjx/test_data/tendon/site.xml b/mjx/mujoco/mjx/test_data/tendon/site.xml new file mode 100644 index 00000000..c2fcb50d --- /dev/null +++ b/mjx/mujoco/mjx/test_data/tendon/site.xml @@ -0,0 +1,26 @@ + + + + + + + + + + + + + + + + + + + + + + + + + +