chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add spatial tendons with site wrapping to MJX.

Browse Source 
PiperOrigin-RevId: 673783473
Change-Id: I09b13787e3e0914c8de21d81db0347c4638b2f1d
This commit is contained in:
Taylor Howell
2024-09-12 04:15:54 -07:00
committed by Copybara-Service
parent 14f420069e
commit e3d3a24b36
10 changed files with 246 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+3 -1
View File
@@ -19,6 +19,7 @@ General
for any object type.
- The experimental use of 2D/3D elasticity plugins with :ref:`composite<body-composite>` has been removed. Users
should instead use :ref:`flexcomp<body-flexcomp>`, which provides the correct collision behavior.
- Added the :ref:`nativeccd<option-flag-nativeccd>` flag. When this flag is enabled, general convex collision
detection is handled natively, as opposed to using `libccd <https://github.com/danfis/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 <coFriction>`.
- Added support for :ref:`connect<equality-connect>` equality constraints using two sites.
- Added support for :ref:`spatial tendons <tendon-spatial>` 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<spatial-geom-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<meAttachment>.
:ref:`attachment<meAttachment>`.
Python bindings
^^^^^^^^^^^^^^^
doc/mjx.rst
+3 -1
View File
@@ -195,6 +195,8 @@ The following features are **fully supported** in MJX:
- ``FIXED``, ``AFFINE``
* - :ref:`Actuator Bias <mjtBias>`
- ``NONE``, ``AFFINE``
* - :ref:`Tendon Wrapping <mjtWrap>`
- ``JOINT``, ``SITE``
* - :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.
@@ -244,7 +246,7 @@ The following features are **in development** and coming soon:
* - :ref:`Actuator Bias <mjtBias>`
- ``MUSCLE``
* - :ref:`Tendon Wrapping <mjtWrap>`
- ``NONE``, ``JOINT``, ``PULLEY``, ``SITE``, ``SPHERE``, ``CYLINDER``
- ``PULLEY``, ``SPHERE``, ``CYLINDER``
* - Fluid Model
- :ref:`flEllipsoid`
* - :ref:`Tendons <tendon>`
mjx/mujoco/mjx/_src/io_test.py
+1
View File
@@ -168,6 +168,7 @@ class ModelIOTest(parameterized.TestCase):
<tendon>
<spatial name="rope" range="0 .35">
<site site="slider"/>
<geom geom="shoulder"/>
<site site="arm"/>
</spatial>
</tendon>
mjx/mujoco/mjx/_src/smooth.py
+100 -6
View File
@@ -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:
mjx/mujoco/mjx/_src/smooth_test.py
+32
View File
@@ -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()
mjx/mujoco/mjx/_src/types.py
+3 -1
View File
@@ -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):
mjx/mujoco/mjx/test_data/tendon/fixed.xml
+23
View File
@@ -0,0 +1,23 @@
<mujoco model="fixed">
<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"/>
<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"/>
<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"/>
</body>
</body>
</body>
</worldbody>
<tendon>
<fixed>
<joint joint="joint0" coef=".1"/>
<joint joint="joint1" coef=".2"/>
<joint joint="joint2" coef=".3"/>
</fixed>
</tendon>
</mujoco>
mjx/mujoco/mjx/test_data/tendon/fixed_site.xml
+39
View File
@@ -0,0 +1,39 @@
<mujoco model="fixed_site">
<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"/>
<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"/>
<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"/>
</body>
</body>
</body>
</worldbody>
<tendon>
<spatial width="0.0125">
<site site="site0"/>
<site site="site1"/>
<site site="site2"/>
</spatial>
<fixed>
<joint joint="joint0" coef=".1"/>
<joint joint="joint1" coef=".2"/>
<joint joint="joint2" coef=".3"/>
</fixed>
<fixed>
<joint joint="joint0" coef=".1"/>
<joint joint="joint2" coef=".3"/>
</fixed>
<spatial width="0.0125">
<site site="site1"/>
<site site="site2"/>
</spatial>
</tendon>
</mujoco>
mjx/mujoco/mjx/test_data/tendon/no_tendon.xml
+16
View File
@@ -0,0 +1,16 @@
<mujoco model="no_tendon">
<worldbody>
<body>
<joint type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<body pos="0.5 0 0">
<joint type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<body pos="0.5 0 0">
<joint type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
</body>
</body>
</body>
</worldbody>
</mujoco>
mjx/mujoco/mjx/test_data/tendon/site.xml
+26
View File
@@ -0,0 +1,26 @@
<mujoco model="site">
<worldbody>
<body>
<joint 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"/>
<body pos="0.5 0 0">
<joint 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"/>
<body pos="0.5 0 0">
<joint 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"/>
</body>
</body>
</body>
</worldbody>
<tendon>
<spatial width="0.0125">
<site site="site0"/>
<site site="site1"/>
<site site="site2"/>
</spatial>
</tendon>
</mujoco>