diff --git a/doc/changelog.rst b/doc/changelog.rst index c175ccd3..6e813d81 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -17,7 +17,7 @@ General MJX ^^^ - Added ``mocap_pos`` and ``mocap_quat`` in kinematics. -- Added support for :ref:`spatial tendons ` with external sphere and cylinder wrapping. +- Added support for :ref:`spatial tendons ` with pulleys and external sphere and cylinder wrapping. Bug fixes ^^^^^^^^^ diff --git a/doc/mjx.rst b/doc/mjx.rst index 8322b62d..270e8ca9 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -196,7 +196,7 @@ The following features are **fully supported** in MJX: * - :ref:`Actuator Bias ` - ``NONE``, ``AFFINE`` * - :ref:`Tendon Wrapping ` - - ``JOINT``, ``SITE`` + - ``JOINT``, ``SITE``, ``PULLEY`` * - :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. @@ -247,7 +247,7 @@ The following features are **in development** and coming soon: * - :ref:`Actuator Bias ` - ``MUSCLE`` * - :ref:`Tendon Wrapping ` - - ``PULLEY``, ``SPHERE``, ``CYLINDER`` + - ``SPHERE``, ``CYLINDER`` * - Fluid Model - :ref:`flEllipsoid` * - :ref:`Tendons ` diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py index 27b32510..6eb8491b 100644 --- a/mjx/mujoco/mjx/_src/smooth.py +++ b/mjx/mujoco/mjx/_src/smooth.py @@ -713,6 +713,21 @@ def tendon(m: Model, d: Data) -> Data: adr_moment_jnt = np.repeat(tendon_id_jnt, tendon_num_jnt) dofadr_moment_jnt = m.jnt_dofadr[wrap_objid_jnt] + # process pulleys + (wrap_id_pulley,) = np.nonzero(m.wrap_type == WrapType.PULLEY) + nwrap_pulley = wrap_id_pulley.size + + tendon_wrapnum_pulley = np.array([ + sum((wrap_id_pulley >= adr) & (wrap_id_pulley < adr + num)) + for adr, num in zip(m.tendon_adr, m.tendon_num) + ]) + + 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] = m.wrap_prm[id_pulley] + # process spatial tendon sites (wrap_id_site,) = np.nonzero(m.wrap_type == WrapType.SITE) nwrap_site = wrap_id_site.size @@ -720,7 +735,6 @@ def tendon(m: Model, d: Data) -> Data: # 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] @@ -746,6 +760,11 @@ def tendon(m: Model, d: Data) -> Data: m.site_bodyid[wrap_objid_site1], ) + if wrap_id_site_pair.size: + divisor_site_pair = divisor[wrap_id_site_pair] + lengths_site /= divisor_site_pair + moments_site /= 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) @@ -848,6 +867,11 @@ def tendon(m: Model, d: Data) -> Data: moments_sitegeom + moments_geomgeom + moments_geomsite, ) + if wrap_id_geom.size: + divisor_geom = divisor[wrap_id_geom] + lengths_geom /= divisor_geom + moments_geom /= divisor_geom[:, None] + # construct number of site-geom-site instances per tendon tendon_ngeom = np.array([ sum((wrap_id_geom >= adr) & (wrap_id_geom < adr + num)) @@ -886,12 +910,16 @@ def tendon(m: Model, d: Data) -> Data: ten_moment = ten_moment.at[tendon_id_geom].add(moment_geom) # construct wrap addresses + wrap_adr_pulley = [] wrap_adr_site = [] wrap_adr_geom = [] count = 0 for wrap_type in m.wrap_type: - if wrap_type == WrapType.SITE: + if wrap_type == WrapType.PULLEY: + wrap_adr_pulley.append(count) + count += 1 + elif wrap_type == WrapType.SITE: wrap_adr_site.append(count) count += 1 elif wrap_type in (WrapType.SPHERE, WrapType.CYLINDER): @@ -899,11 +927,12 @@ def tendon(m: Model, d: Data) -> Data: wrap_adr_geom.append(count + 1) count += 2 + wrap_adr_pulley = np.array(wrap_adr_pulley).astype(int) wrap_adr_site = np.array(wrap_adr_site).astype(int) wrap_adr_geom = np.array(wrap_adr_geom).astype(int) - wrap_adr_sitegeom = np.concatenate([wrap_adr_site, wrap_adr_geom]) + wrap_adr = np.concatenate([wrap_adr_pulley, wrap_adr_site, wrap_adr_geom]) - ten_wrapnum = jp.array(tendon_wrapnum_site) + ten_wrapnum = jp.array(tendon_wrapnum_pulley + tendon_wrapnum_site) ten_wrapnum = ten_wrapnum.at[tendon_id_geom].add(tendon_wrapnum_geom) ten_wrapadr = jp.concatenate([jp.array([0]), jp.cumsum(ten_wrapnum)[:-1]]) @@ -912,9 +941,7 @@ def tendon(m: Model, d: Data) -> Data: xpos_geom = jp.hstack([geom_pnt0, geom_pnt1]).reshape((-1, 3)) # sort objects, moving no wrap geoms to bottom rows - nwrap_sitegeom = wrap_adr_sitegeom.size - wrap_adr_sitegeom_sort = np.argsort(wrap_adr_sitegeom) - + wrap_adr_sort = np.argsort(wrap_adr) skipped = ( jp.zeros(count, dtype=bool) .at[wrap_adr_geom] @@ -922,17 +949,20 @@ def tendon(m: Model, d: Data) -> Data: ) sort = jp.argsort(skipped) - wrap_xpos = jp.concatenate([xpos_site, xpos_geom])[wrap_adr_sitegeom_sort] wrap_xpos = jp.concatenate( - [wrap_xpos[sort], jp.zeros((2 * m.nwrap - nwrap_sitegeom, 3))] + [jp.zeros((nwrap_pulley, 3)), xpos_site, xpos_geom] + )[wrap_adr_sort] + wrap_xpos = jp.concatenate( + [wrap_xpos[sort], jp.zeros((2 * m.nwrap - count, 3))] ).reshape((m.nwrap, 6)) wrap_obj = jp.concatenate([ + -2 * jp.ones(nwrap_pulley, dtype=int), -1 * jp.ones(nwrap_site, dtype=int), jp.repeat(wrap_objid_geom_skip, 2).reshape(-1), - ])[wrap_adr_sitegeom_sort] + ])[wrap_adr_sort] wrap_obj = jp.concatenate( - [wrap_obj[sort], jp.zeros(2 * m.nwrap - nwrap_sitegeom, dtype=int)] + [wrap_obj[sort], jp.zeros(2 * m.nwrap - count, dtype=int)] ).reshape((m.nwrap, 2)) return d.replace( diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index 9892e740..46f1fa2e 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -240,9 +240,14 @@ class TendonTest(parameterized.TestCase): 'tendon/fixed.xml', 'tendon/fixed_site.xml', 'tendon/fixed_site_wrap.xml', + 'tendon/pulley_fixed_site_wrap.xml', + 'tendon/pulley_site.xml', + 'tendon/pulley_site_wrap.xml', + 'tendon/pulley_wrap.xml', 'tendon/no_tendon.xml', 'tendon/site.xml', 'tendon/site_wrap.xml', + 'tendon/tendon.xml', 'tendon/wrap_sidesite.xml', ) def test_tendon(self, filename): diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index a9ba040b..b0be4743 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -197,15 +197,16 @@ class WrapType(enum.IntEnum): Members: JOINT: constant moment arm + PULLEY: pulley used to split tendon SITE: pass through site SPHERE: wrap around sphere CYLINDER: wrap around (infinite) cylinder """ JOINT = mujoco.mjtWrap.mjWRAP_JOINT + PULLEY = mujoco.mjtWrap.mjWRAP_PULLEY SITE = mujoco.mjtWrap.mjWRAP_SITE SPHERE = mujoco.mjtWrap.mjWRAP_SPHERE CYLINDER = mujoco.mjtWrap.mjWRAP_CYLINDER - # unsupported: NONE, PULLEY class TrnType(enum.IntEnum): diff --git a/mjx/mujoco/mjx/test_data/tendon/pulley_fixed_site_wrap.xml b/mjx/mujoco/mjx/test_data/tendon/pulley_fixed_site_wrap.xml new file mode 100644 index 00000000..a3e553cd --- /dev/null +++ b/mjx/mujoco/mjx/test_data/tendon/pulley_fixed_site_wrap.xml @@ -0,0 +1,103 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/mjx/mujoco/mjx/test_data/tendon/pulley_site.xml b/mjx/mujoco/mjx/test_data/tendon/pulley_site.xml new file mode 100644 index 00000000..6ae3e137 --- /dev/null +++ b/mjx/mujoco/mjx/test_data/tendon/pulley_site.xml @@ -0,0 +1,84 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/mjx/mujoco/mjx/test_data/tendon/pulley_site_wrap.xml b/mjx/mujoco/mjx/test_data/tendon/pulley_site_wrap.xml new file mode 100644 index 00000000..6621b6fa --- /dev/null +++ b/mjx/mujoco/mjx/test_data/tendon/pulley_site_wrap.xml @@ -0,0 +1,98 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/mjx/mujoco/mjx/test_data/tendon/pulley_wrap.xml b/mjx/mujoco/mjx/test_data/tendon/pulley_wrap.xml new file mode 100644 index 00000000..c9929d6e --- /dev/null +++ b/mjx/mujoco/mjx/test_data/tendon/pulley_wrap.xml @@ -0,0 +1,78 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/mjx/mujoco/mjx/test_data/tendon/tendon.xml b/mjx/mujoco/mjx/test_data/tendon/tendon.xml new file mode 100644 index 00000000..ab1c4227 --- /dev/null +++ b/mjx/mujoco/mjx/test_data/tendon/tendon.xml @@ -0,0 +1,66 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +