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

Add support for computing elliptic cone contact forces to MJX. Fixes #2153.

Browse Source 
PiperOrigin-RevId: 691829630
Change-Id: I862ac9c7db07d0a74b718d2a41e19d5ae636a260
This commit is contained in:
Taylor Howell
2024-10-31 09:22:02 -07:00
committed by Copybara-Service
parent 18682aac6a
commit 554641060b
4 changed files with 44 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/sensor_test.py
+8 -2
View File
@@ -22,6 +22,8 @@ from jax import numpy as jp
import mujoco
from mujoco import mjx
from mujoco.mjx._src import test_util
from mujoco.mjx._src.types import ConeType
import numpy as np
# tolerance for difference between MuJoCo and MJX smooth calculations - mostly
@@ -41,10 +43,14 @@ def _assert_attr_eq(a, b, attr):
class SensorTest(parameterized.TestCase):
@parameterized.parameters('sensor/model.xml', 'sensor/sensor.xml')
def test_sensor(self, filename):
@parameterized.product(
filename=['sensor/model.xml', 'sensor/sensor.xml'],
cone_type=list(ConeType),
)
def test_sensor(self, filename, cone_type):
"""Tests MJX sensor functions match MuJoCo sensor functions."""
m = test_util.load_test_file(filename)
m.opt.cone = cone_type
d = mujoco.MjData(m)
# give the system a little kick to ensure we have non-identity rotations
d.qvel = 0.1 * np.random.random(m.nv)
mjx/mujoco/mjx/_src/smooth_test.py
+21 -5
View File
@@ -20,6 +20,7 @@ import jax
import mujoco
from mujoco import mjx
from mujoco.mjx._src import test_util
from mujoco.mjx._src.types import ConeType
import numpy as np
# tolerance for difference between MuJoCo and MJX smooth calculations - mostly
@@ -198,27 +199,42 @@ class SmoothTest(absltest.TestCase):
_assert_attr_eq(d, dx, 'subtree_linvel')
_assert_attr_eq(d, dx, 'subtree_angmom')
def test_rnepostconstraint(self):
class RnePostConstraintTest(parameterized.TestCase):
@parameterized.parameters(ConeType)
def test_rnepostconstraint(self, cone_type):
"""Tests MJX rne_postconstraint function to match MuJoCo mj_rnePostConstraint."""
m = mujoco.MjModel.from_xml_string("""
<mujoco>
<worldbody>
<geom name="floor" size="0 0 .05" type="plane"/>
<geom name="floor" size="10 10 .05" type="plane"/>
<body pos="0 0 1">
<joint type="ball" damping="1"/>
<geom type="capsule" size="0.1 0.5" fromto="0 0 0 0.5 0 0"/>
<geom type="capsule" size="0.1 0.5" fromto="0 0 0 0.5 0 0" condim="1"/>
<body pos="0.5 0 0">
<joint type="ball" damping="1"/>
<geom type="capsule" size="0.1 0.5" fromto="0 0 0 0.5 0 0"/>
<geom type="capsule" size="0.1 0.5" fromto="0 0 0 0.5 0 0" condim="3"/>
</body>
</body>
<body pos="0 1 1">
<joint type="ball" damping="1"/>
<geom type="capsule" size="0.1 0.5" fromto="0 0 0 0.5 0 0" condim="6"/>
<body pos="0.5 0 0">
<joint type="ball" damping="1"/>
<geom type="capsule" size="0.1 0.5" fromto="0 0 0 0.5 0 0" condim="3"/>
</body>
</body>
</worldbody>
<keyframe>
<key qpos='0.424577 0.450592 0.451703 -0.642391 0.729379 0.545151 0.407756 0.0674697'/>
<key qpos='0.424577 0.450592 0.451703 -0.642391 0.729379 0.545151 0.407756 0.0674697 0.424577 1.450592 0.451703 -0.642391 0.729379 0.545151 0.407756 0.0674697'/>
</keyframe>
</mujoco>
""")
# set cone type
m.opt.cone = cone_type
# create data and set to keyframe
d = mujoco.MjData(m)
mujoco.mj_resetDataKeyframe(m, d, 0)
# apply external forces
mjx/mujoco/mjx/_src/support.py
+11 -8
View File
@@ -21,6 +21,7 @@ import mujoco
from mujoco.mjx._src import math
from mujoco.mjx._src import scan
# pylint: disable=g-importing-member
from mujoco.mjx._src.types import ConeType
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import JacobianType
from mujoco.mjx._src.types import Model
@@ -307,12 +308,13 @@ def contact_force(
"""Extract 6D force:torque for one contact, in contact frame by default."""
efc_address = d.contact.efc_address[contact_id]
condim = d.contact.dim[contact_id]
if m.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL:
if m.opt.cone == ConeType.PYRAMIDAL:
force = _decode_pyramid(
d.efc_force[efc_address:], d.contact.friction[contact_id], condim
)
elif m.opt.cone == mujoco.mjtCone.mjCONE_ELLIPTIC:
raise NotImplementedError('Elliptic cone force is not implemented yet.')
elif m.opt.cone == ConeType.ELLIPTIC:
force = d.efc_force[efc_address : efc_address + condim]
force = jp.concatenate([force, jp.zeros((6 - condim))])
else:
raise ValueError(f'Unknown cone type: {m.opt.cone}')
@@ -331,7 +333,7 @@ def contact_force_dim(
idx_dim = (d.contact.efc_address >= 0) & (d.contact.dim == dim)
# contact force from efc
if m.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL:
if m.opt.cone == ConeType.PYRAMIDAL:
efc_address = (
d.contact.efc_address[idx_dim, None]
+ np.arange(np.where(dim == 1, 1, 2 * (dim - 1)))[None]
@@ -340,12 +342,13 @@ def contact_force_dim(
force = jax.vmap(_decode_pyramid, in_axes=(0, 0, None))(
efc_force, d.contact.friction[idx_dim], dim
)
return force, np.where(idx_dim)[0]
elif m.opt.cone == mujoco.mjtCone.mjCONE_ELLIPTIC:
# TODO(taylorhowell): add support for elliptic cone
raise NotImplementedError('Elliptic cone force is not implemented yet.')
elif m.opt.cone == ConeType.ELLIPTIC:
efc_address = d.contact.efc_address[idx_dim, None] + np.arange(dim)[None]
force = d.efc_force[efc_address]
force = jp.hstack([force, jp.zeros((force.shape[0], 6 - dim))])
else:
raise ValueError(f'Unknown cone type: {m.opt.cone}.')
return force, np.where(idx_dim)[0]
def _length_circle(
mjx/mujoco/mjx/test_data/sensor/sensor.xml
+4 -4
View File
@@ -106,22 +106,22 @@
<geom type="plane" size="2 2 .1" pos="-20 -20 -20"/>
<body pos="-20 -20 -20">
<freejoint/>
<geom type="capsule" fromto="-.5 0 0 .5 0 0" size="0.0125"/>
<geom type="capsule" fromto="-.5 0 0 .5 0 0" size="0.0125" condim="1"/>
<site name="touch_sphere" type="sphere" pos="-0.5 0 0" size="0.025"/>
</body>
<body pos="-20 -20.25 -20">
<freejoint/>
<geom type="capsule" fromto="-.5 0 0 .5 0 0" size="0.0125"/>
<geom type="capsule" fromto="-.5 0 0 .5 0 0" size="0.0125" condim="3"/>
<site name="touch_capsule" type="capsule" fromto="0.4 0 0.0 0.5 0 0.0" size="0.025"/>
</body>
<body pos="-20 -20.5 -20">
<freejoint/>
<geom type="capsule" fromto="-.5 0 0 .5 0 0" size="0.0125"/>
<geom type="capsule" fromto="-.5 0 0 .5 0 0" size="0.0125" condim="6"/>
<site name="touch_box" pos="-0.5 0 0" type="box" size="0.025 0.025 0.025"/>
</body>
<body pos="-20 -20.75 -20">
<freejoint/>
<geom type="capsule" fromto="-.5 0 0 .5 0 0" size="0.0125"/>
<geom type="capsule" fromto="-.5 0 0 .5 0 0" size="0.0125" condim="3"/>
<site name="touch_ellipsoid" type="ellipsoid" pos="0.5 0 0.0" size="0.05 0.01 0.02"/>
</body>
</worldbody>