From 554641060b569be8f55237a3ceee2f61007587cf Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Thu, 31 Oct 2024 09:22:02 -0700 Subject: [PATCH] Add support for computing elliptic cone contact forces to MJX. Fixes #2153. PiperOrigin-RevId: 691829630 Change-Id: I862ac9c7db07d0a74b718d2a41e19d5ae636a260 --- mjx/mujoco/mjx/_src/sensor_test.py | 10 +++++++-- mjx/mujoco/mjx/_src/smooth_test.py | 26 +++++++++++++++++----- mjx/mujoco/mjx/_src/support.py | 19 +++++++++------- mjx/mujoco/mjx/test_data/sensor/sensor.xml | 8 +++---- 4 files changed, 44 insertions(+), 19 deletions(-) diff --git a/mjx/mujoco/mjx/_src/sensor_test.py b/mjx/mujoco/mjx/_src/sensor_test.py index 19d4ef46..37b1e13e 100644 --- a/mjx/mujoco/mjx/_src/sensor_test.py +++ b/mjx/mujoco/mjx/_src/sensor_test.py @@ -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) diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index 8b5218e5..fc14ac65 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -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(""" - + - + - + + + + + + + + + - + """) + # 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 diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index 622c6659..f065d9d8 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/mjx/mujoco/mjx/_src/support.py @@ -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( diff --git a/mjx/mujoco/mjx/test_data/sensor/sensor.xml b/mjx/mujoco/mjx/test_data/sensor/sensor.xml index b6d3fa2a..efd855e1 100644 --- a/mjx/mujoco/mjx/test_data/sensor/sensor.xml +++ b/mjx/mujoco/mjx/test_data/sensor/sensor.xml @@ -106,22 +106,22 @@ - + - + - + - +