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 @@
-
+
-
+
-
+
-
+