Add pulley support to spatial tendons in MJX.

PiperOrigin-RevId: 683103659
Change-Id: I22544071a3fc6fbd09f9a59690106701b101ee6a
This commit is contained in:
Taylor Howell
2024-10-07 03:00:02 -07:00
committed by Copybara-Service
parent bce43b63ed
commit 160ed9bfe2
10 changed files with 480 additions and 15 deletions
+1 -1
View File
@@ -17,7 +17,7 @@ General
MJX
^^^
- Added ``mocap_pos`` and ``mocap_quat`` in kinematics.
- Added support for :ref:`spatial tendons <tendon-spatial>` with external sphere and cylinder wrapping.
- Added support for :ref:`spatial tendons <tendon-spatial>` with pulleys and external sphere and cylinder wrapping.
Bug fixes
^^^^^^^^^
+2 -2
View File
@@ -196,7 +196,7 @@ The following features are **fully supported** in MJX:
* - :ref:`Actuator Bias <mjtBias>`
- ``NONE``, ``AFFINE``
* - :ref:`Tendon Wrapping <mjtWrap>`
- ``JOINT``, ``SITE``
- ``JOINT``, ``SITE``, ``PULLEY``
* - :ref:`Geom <mjtGeom>`
- ``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 <mjtBias>`
- ``MUSCLE``
* - :ref:`Tendon Wrapping <mjtWrap>`
- ``PULLEY``, ``SPHERE``, ``CYLINDER``
- ``SPHERE``, ``CYLINDER``
* - Fluid Model
- :ref:`flEllipsoid`
* - :ref:`Tendons <tendon>`
+41 -11
View File
@@ -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(
+5
View File
@@ -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):
+2 -1
View File
@@ -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):
@@ -0,0 +1,103 @@
<mujoco model="pulley_fixed_site_wrap">
<option>
<flag contact="disable" gravity="disable"/>
</option>
<worldbody>
<body>
<joint name="joint0" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site0" pos="0.25 0 0.1" size="0.025"/>
<geom name="sphere0" type="sphere" pos="0.5 0 0.125" size="0.085"/>
<body pos="0.5 0 0">
<joint name="joint1" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site1" pos="0.25 0 0.1" size="0.025"/>
<geom name="sphere1" type="sphere" pos="0.5 0 0.125" size="0.085"/>
<body pos="0.5 0 0">
<joint name="joint2" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site2" pos="0.25 0 0.1" size="0.025"/>
<geom name="cylinder2" type="cylinder" contype="0" conaffinity="0" pos="0.5 0 0.25" euler="90 0 0" size="0.085 0.1"/>
<geom name="sphere2" type="sphere" contype="0" conaffinity="0" pos="0.5 0 0.25" size="0.085"/>
<body pos="0.5 0 0">
<joint name="joint3" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site3" pos="0.25 0 0.1" size="0.025"/>
</body>
</body>
</body>
</body>
</worldbody>
<tendon>
<spatial width="0.0125">
<site site="site2"/>
<site site="site3"/>
</spatial>
<spatial width="0.0125">
<pulley divisor="0.2"/>
<site site="site2"/>
<site site="site3"/>
</spatial>
<fixed>
<joint joint="joint0" coef=".1"/>
<joint joint="joint1" coef=".2"/>
<joint joint="joint2" coef=".3"/>
</fixed>
<spatial width="0.0125">
<pulley divisor="3"/>
<site site="site0"/>
<geom geom="sphere0"/>
<site site="site1"/>
<geom geom="sphere1"/>
<site site="site2"/>
<geom geom="cylinder2"/>
<site site="site3"/>
<pulley divisor="2"/>
<site site="site0"/>
<geom geom="sphere0"/>
<site site="site1"/>
<geom geom="sphere1"/>
<site site="site2"/>
<geom geom="cylinder2"/>
<site site="site3"/>
</spatial>
<spatial width="0.0125">
<pulley divisor="3"/>
<site site="site0"/>
<geom geom="sphere0"/>
<site site="site1"/>
<geom geom="sphere1"/>
<site site="site2"/>
<pulley divisor="7"/>
<site site="site2"/>
<geom geom="cylinder2"/>
<site site="site3"/>
<pulley divisor="2"/>
<site site="site0"/>
<geom geom="sphere0"/>
<site site="site1"/>
<geom geom="sphere1"/>
<site site="site2"/>
<geom geom="cylinder2"/>
<site site="site3"/>
</spatial>
<spatial width="0.0125">
<site site="site0"/>
<geom geom="sphere0"/>
<site site="site2"/>
<geom geom="cylinder2"/>
<site site="site3"/>
</spatial>
<spatial width="0.0125">
<pulley divisor="4"/>
<site site="site0"/>
<site site="site2"/>
<site site="site3"/>
<pulley divisor="3"/>
<site site="site0"/>
<site site="site1"/>
<site site="site2"/>
<site site="site3"/>
</spatial>
</tendon>
</mujoco>
@@ -0,0 +1,84 @@
<mujoco model="pulley_site">
<option>
<flag contact="disable" gravity="disable"/>
</option>
<worldbody>
<body>
<joint name="joint0" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site0" pos="0.25 0 0.1" size="0.025"/>
<geom name="sphere0" type="sphere" pos="0.5 0 0.125" size="0.085"/>
<body pos="0.5 0 0">
<joint name="joint1" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site1" pos="0.25 0 0.1" size="0.025"/>
<geom name="sphere1" type="sphere" pos="0.5 0 0.125" size="0.085"/>
<body pos="0.5 0 0">
<joint name="joint2" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site2" pos="0.25 0 0.1" size="0.025"/>
<geom name="cylinder2" type="cylinder" contype="0" conaffinity="0" pos="0.5 0 0.25" euler="90 0 0" size="0.085 0.1"/>
<geom name="sphere2" type="sphere" contype="0" conaffinity="0" pos="0.5 0 0.25" size="0.085"/>
<body pos="0.5 0 0">
<joint name="joint3" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site3" pos="0.25 0 0.1" size="0.025"/>
</body>
</body>
</body>
</body>
</worldbody>
<tendon>
<spatial width="0.0125">
<site site="site2"/>
<site site="site3"/>
</spatial>
<spatial width="0.0125">
<pulley divisor="0.2"/>
<site site="site2"/>
<site site="site3"/>
</spatial>
<spatial width="0.0125">
<pulley divisor="3"/>
<site site="site0"/>
<site site="site1"/>
<site site="site2"/>
<site site="site3"/>
<pulley divisor="2"/>
<site site="site0"/>
<site site="site1"/>
<site site="site2"/>
<site site="site3"/>
</spatial>
<spatial width="0.0125">
<pulley divisor="3"/>
<site site="site0"/>
<site site="site1"/>
<site site="site2"/>
<pulley divisor="7"/>
<site site="site2"/>
<site site="site3"/>
<pulley divisor="2"/>
<site site="site0"/>
<site site="site1"/>
<site site="site2"/>
<site site="site3"/>
</spatial>
<spatial width="0.0125">
<site site="site0"/>
<site site="site2"/>
<site site="site3"/>
</spatial>
<spatial width="0.0125">
<pulley divisor="4"/>
<site site="site0"/>
<site site="site2"/>
<site site="site3"/>
<pulley divisor="3"/>
<site site="site0"/>
<site site="site1"/>
<site site="site2"/>
<site site="site3"/>
</spatial>
</tendon>
</mujoco>
@@ -0,0 +1,98 @@
<mujoco model="pulley_site_wrap">
<option>
<flag contact="disable" gravity="disable"/>
</option>
<worldbody>
<body>
<joint name="joint0" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site0" pos="0.25 0 0.1" size="0.025"/>
<geom name="sphere0" type="sphere" pos="0.5 0 0.125" size="0.085"/>
<body pos="0.5 0 0">
<joint name="joint1" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site1" pos="0.25 0 0.1" size="0.025"/>
<geom name="sphere1" type="sphere" pos="0.5 0 0.125" size="0.085"/>
<body pos="0.5 0 0">
<joint name="joint2" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site2" pos="0.25 0 0.1" size="0.025"/>
<geom name="cylinder2" type="cylinder" contype="0" conaffinity="0" pos="0.5 0 0.25" euler="90 0 0" size="0.085 0.1"/>
<geom name="sphere2" type="sphere" contype="0" conaffinity="0" pos="0.5 0 0.25" size="0.085"/>
<body pos="0.5 0 0">
<joint name="joint3" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site3" pos="0.25 0 0.1" size="0.025"/>
</body>
</body>
</body>
</body>
</worldbody>
<tendon>
<spatial width="0.0125">
<site site="site2"/>
<site site="site3"/>
</spatial>
<spatial width="0.0125">
<pulley divisor="0.2"/>
<site site="site2"/>
<site site="site3"/>
</spatial>
<spatial width="0.0125">
<pulley divisor="3"/>
<site site="site0"/>
<geom geom="sphere0"/>
<site site="site1"/>
<geom geom="sphere1"/>
<site site="site2"/>
<geom geom="cylinder2"/>
<site site="site3"/>
<pulley divisor="2"/>
<site site="site0"/>
<geom geom="sphere0"/>
<site site="site1"/>
<geom geom="sphere1"/>
<site site="site2"/>
<geom geom="cylinder2"/>
<site site="site3"/>
</spatial>
<spatial width="0.0125">
<pulley divisor="3"/>
<site site="site0"/>
<geom geom="sphere0"/>
<site site="site1"/>
<geom geom="sphere1"/>
<site site="site2"/>
<pulley divisor="7"/>
<site site="site2"/>
<geom geom="cylinder2"/>
<site site="site3"/>
<pulley divisor="2"/>
<site site="site0"/>
<geom geom="sphere0"/>
<site site="site1"/>
<geom geom="sphere1"/>
<site site="site2"/>
<geom geom="cylinder2"/>
<site site="site3"/>
</spatial>
<spatial width="0.0125">
<site site="site0"/>
<geom geom="sphere0"/>
<site site="site2"/>
<geom geom="cylinder2"/>
<site site="site3"/>
</spatial>
<spatial width="0.0125">
<pulley divisor="4"/>
<site site="site0"/>
<site site="site2"/>
<site site="site3"/>
<pulley divisor="3"/>
<site site="site0"/>
<site site="site1"/>
<site site="site2"/>
<site site="site3"/>
</spatial>
</tendon>
</mujoco>
@@ -0,0 +1,78 @@
<mujoco model="pulley_wrap">
<option>
<flag contact="disable" gravity="disable"/>
</option>
<worldbody>
<body>
<joint name="joint0" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site0" pos="0.25 0 0.1" size="0.025"/>
<geom name="sphere0" type="sphere" pos="0.5 0 0.125" size="0.085"/>
<body pos="0.5 0 0">
<joint name="joint1" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site1" pos="0.25 0 0.1" size="0.025"/>
<geom name="sphere1" type="sphere" pos="0.5 0 0.125" size="0.085"/>
<body pos="0.5 0 0">
<joint name="joint2" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site2" pos="0.25 0 0.1" size="0.025"/>
<geom name="cylinder2" type="cylinder" contype="0" conaffinity="0" pos="0.5 0 0.25" euler="90 0 0" size="0.085 0.1"/>
<geom name="sphere2" type="sphere" contype="0" conaffinity="0" pos="0.5 0 0.25" size="0.085"/>
<body pos="0.5 0 0">
<joint name="joint3" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site3" pos="0.25 0 0.1" size="0.025"/>
</body>
</body>
</body>
</body>
</worldbody>
<tendon>
<spatial width="0.0125">
<pulley divisor="3"/>
<site site="site0"/>
<geom geom="sphere0"/>
<site site="site1"/>
<geom geom="sphere1"/>
<site site="site2"/>
<geom geom="cylinder2"/>
<site site="site3"/>
<pulley divisor="2"/>
<site site="site0"/>
<geom geom="sphere0"/>
<site site="site1"/>
<geom geom="sphere1"/>
<site site="site2"/>
<geom geom="cylinder2"/>
<site site="site3"/>
</spatial>
<spatial width="0.0125">
<pulley divisor="3"/>
<site site="site0"/>
<geom geom="sphere0"/>
<site site="site1"/>
<geom geom="sphere1"/>
<site site="site2"/>
<pulley divisor="7"/>
<site site="site2"/>
<geom geom="cylinder2"/>
<site site="site3"/>
<pulley divisor="2"/>
<site site="site0"/>
<geom geom="sphere0"/>
<site site="site1"/>
<geom geom="sphere1"/>
<site site="site2"/>
<geom geom="cylinder2"/>
<site site="site3"/>
</spatial>
<spatial width="0.0125">
<site site="site0"/>
<geom geom="sphere0"/>
<site site="site2"/>
<geom geom="cylinder2"/>
<site site="site3"/>
</spatial>
</tendon>
</mujoco>
@@ -0,0 +1,66 @@
<mujoco model="tendon">
<default>
<joint axis="0 1 0" range="0 60" limited="true"/>
<geom type="capsule" size=".01" rgba=".9 .7 .1 1"/>
<site rgba=".9 .9 .9 1"/>
<default class="visual_cylinder">
<geom type="cylinder" fromto="0 .015 0 0 -.015 0" size=".02" rgba=".3 .9 .3 .4"/>
</default>
</default>
<worldbody>
<light pos="0 0 1"/>
<body name="weight" pos="-.03 0 .175">
<joint axis="0 0 1" pos="0 0 .025" type="slide"/>
<geom type="cylinder" size=".03 .025" density="5000" rgba=".2 .2 .5 1"/>
<site name="s1" pos="0 0 .025"/>
</body>
<site name="s2" pos="-.03 0 .33"/>
<body pos="0 0 .3">
<joint/>
<geom name="g1" class="visual_cylinder"/>
<geom fromto="0 0 0 .1 0 0"/>
<site name="s3" pos=".02 0 .03"/>
<body pos=".1 0 0">
<joint/>
<geom name="g2" class="visual_cylinder"/>
<geom fromto="0 0 0 .1 0 0"/>
<site name="s4" pos=".03 0 .01"/>
<site name="s5" pos=".05 0 .02"/>
<site name="side2" pos="0 0 .03"/>
<body pos=".1 0 0">
<joint/>
<geom name="g3" class="visual_cylinder"/>
<geom fromto="0 0 0 .1 0 0"/>
<site name="s6" pos=".03 0 .01"/>
<site name="side3" pos="0 0 .03"/>
</body>
</body>
</body>
</worldbody>
<tendon>
<spatial range="0 .33" limited="true" width=".002" rgba=".95 .3 .3 1">
<site site="s1"/>
<site site="s2"/>
<geom geom="g1"/>
<site site="s3"/>
<pulley divisor="2"/>
<site site="s3"/>
<geom geom="g2" sidesite="side2"/>
<site site="s4"/>
<pulley divisor="2"/>
<site site="s3"/>
<geom geom="g2" sidesite="side2"/>
<site site="s5"/>
<geom geom="g3" sidesite="side3"/>
<site site="s6"/>
</spatial>
</tendon>
</mujoco>