Add ten_length as MJX public field and warp tendon function call to MJX.

PiperOrigin-RevId: 794801198 Change-Id: I970ce665339c75910a767a1b24d8ac9c2bd68b1d
2025-08-13 17:26:17 -07:00
parent 1a7ec97b07
commit d7e925bd8a
15 changed files with 258 additions and 28 deletions
@@ -2,6 +2,19 @@
 Changelog
 =========

+Upcoming version (not yet released)
+-----------------------------------
+
+MJX
+^^^
+- Promote ``ten_length`` to the public MJX API. Add Warp support for ``mjx.tendon``.
+
+.. admonition:: Breaking API changes
+   :class: attention
+
+   - ``ten_length`` was moved from ``mjx.Data._impl.ten_length`` to a public field ``mjx.Data.ten_length``.
+
+
 Version 3.3.5 (August 8, 2025)
 -----------------------------------

@@ -322,8 +322,8 @@ def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]:

  inv1, inv2 = m.tendon_invweight0[obj1id], m.tendon_invweight0[obj2id]
  jac1, jac2 = d._impl.ten_J[obj1id], d._impl.ten_J[obj2id]
-  pos1 = d._impl.ten_length[obj1id] - m.tendon_length0[obj1id]
-  pos2 = d._impl.ten_length[obj2id] - m.tendon_length0[obj2id]
+  pos1 = d.ten_length[obj1id] - m.tendon_length0[obj1id]
+  pos2 = d.ten_length[obj2id] - m.tendon_length0[obj2id]
  invweight = inv1 + inv2 * (obj2id > -1)

  return rows(
@@ -436,7 +436,7 @@ def _efc_limit_tendon(m: Model, d: Data) -> Optional[_Efc]:
  length, j, range_, margin, invweight, solref, solimp = jax.tree_util.tree_map(
      lambda x: x[tendon_id],
      (
-          d._impl.ten_length,
+          d.ten_length,
          d._impl.ten_J,
          m.tendon_range,
          m.tendon_margin,
@@ -576,6 +576,7 @@ def _make_data_public_fields(m: types.Model) -> Dict[str, Any]:
      'qfrc_constraint': (m.nv, float_),
      'qfrc_inverse': (m.nv, float_),
      'cvel': (m.nbody, 6, float_),
+      'ten_length': (m.ntendon, float_),
  }
  zero_fields = {
      k: np.zeros(v[:-1], dtype=v[-1]) for k, v in zero_fields.items()
@@ -637,7 +638,6 @@ def _make_data_jax(
      'ten_wrapadr': (m.ntendon, np.int32),
      'ten_wrapnum': (m.ntendon, np.int32),
      'ten_J': (m.ntendon, m.nv, float_),
-      'ten_length': (m.ntendon, float_),
      'wrap_obj': (m.nwrap, 2, np.int32),
      'wrap_xpos': (m.nwrap, 6, float_),
      'actuator_length': (m.nu, float_),
@@ -742,7 +742,6 @@ def _make_data_c(
      'ten_J_rowadr': (m.ntendon, np.int32),
      'ten_J_colind': (m.ntendon, m.nv, np.int32),
      'ten_J': (m.ntendon, m.nv, float_),
-      'ten_length': (m.ntendon, float_),
      'ten_wrapadr': (m.ntendon, np.int32),
      'ten_wrapnum': (m.ntendon, np.int32),
      'wrap_obj': (m.nwrap, 2, np.int32),
@@ -478,7 +478,7 @@ class DataIOTest(parameterized.TestCase):
    np.testing.assert_allclose(dx.site_xmat.reshape((1, 9)), d.site_xmat)

    # tendon data is correct
-    np.testing.assert_allclose(dx._impl.ten_length, d.ten_length)
+    np.testing.assert_allclose(dx.ten_length, d.ten_length)
    np.testing.assert_equal(dx._impl.ten_wrapadr, np.zeros((1,)))
    np.testing.assert_equal(dx._impl.ten_wrapnum, np.zeros((1,)))
    np.testing.assert_equal(dx._impl.wrap_obj, np.zeros((2, 2)))
@@ -75,7 +75,7 @@ def _spring_damper(m: Model, d: Data) -> jax.Array:
  qfrc -= m.dof_damping * d.qvel

  # tendon-level spring-dampers
-  below, above = m.tendon_lengthspring.T - d._impl.ten_length
+  below, above = m.tendon_lengthspring.T - d.ten_length
  frc_spring = jp.where(below > 0, m.tendon_stiffness * below, 0)
  frc_spring = jp.where(above < 0, m.tendon_stiffness * above, frc_spring)
  frc_damper = -m.tendon_damping * d._impl.ten_velocity
@@ -22,7 +22,6 @@ from mujoco.mjx._src import math
 from mujoco.mjx._src import ray
 from mujoco.mjx._src import smooth
 from mujoco.mjx._src import support
-from mujoco.mjx._src.types import Impl
 from mujoco.mjx._src.types import Data
 from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
@@ -174,7 +173,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
    elif sensor_type == SensorType.JOINTPOS:
      sensor = d.qpos[m.jnt_qposadr[objid]]
    elif sensor_type == SensorType.TENDONPOS:
-      sensor = d._impl.ten_length[objid]
+      sensor = d.ten_length[objid]
    elif sensor_type == SensorType.ACTUATORPOS:
      sensor = d._impl.actuator_length[objid]
    elif sensor_type == SensorType.BALLQUAT:
@@ -26,6 +26,7 @@ from mujoco.mjx._src.types import Data
 from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import EqType
+from mujoco.mjx._src.types import Impl
 from mujoco.mjx._src.types import JointType
 from mujoco.mjx._src.types import Model
 from mujoco.mjx._src.types import ModelJAX
@@ -33,11 +34,16 @@ from mujoco.mjx._src.types import ObjType
 from mujoco.mjx._src.types import TrnType
 from mujoco.mjx._src.types import WrapType
 # pylint: enable=g-importing-member
+import mujoco.mjx.warp as mjxw
 import numpy as np


 def kinematics(m: Model, d: Data) -> Data:
  """Converts position/velocity from generalized coordinates to maximal."""
+  if m.impl == Impl.WARP and d.impl == Impl.WARP and mjxw.WARP_INSTALLED:
+    from mujoco.mjx.warp import smooth as mjxw_smooth  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
+    return mjxw_smooth.kinematics(m, d)
+
  def fn(carry, jnt_typs, jnt_pos, jnt_axis, qpos, qpos0, pos, quat):
    # calculate joint anchors, axes, body pos and quat in global frame
    # also normalize qpos while we're at it
@@ -844,6 +850,10 @@ def rne_postconstraint(m: Model, d: Data) -> Data:

 def tendon(m: Model, d: Data) -> Data:
  """Computes tendon lengths and moments."""
+  if m.impl == Impl.WARP and d.impl == Impl.WARP and mjxw.WARP_INSTALLED:
+    from mujoco.mjx.warp import smooth as mjxw_smooth  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
+    return mjxw_smooth.tendon(m, d)
+
  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
    raise ValueError('tendon requires JAX backend implementation.')

@@ -1091,7 +1101,7 @@ def tendon(m: Model, d: Data) -> Data:

  # assemble length and moment
  ten_length = (
-      jp.zeros_like(d._impl.ten_length).at[tendon_id_jnt].set(length_jnt)
+      jp.zeros_like(d.ten_length).at[tendon_id_jnt].set(length_jnt)
  )
  ten_length = ten_length.at[tendon_id_site].add(length_site)
  ten_length = ten_length.at[tendon_id_geom].add(length_geom)
@@ -1161,7 +1171,7 @@ def tendon(m: Model, d: Data) -> Data:
  ).reshape((m.nwrap, 2))

  return d.tree_replace({
-      '_impl.ten_length': ten_length,
+      'ten_length': ten_length,
      '_impl.ten_J': ten_moment,
      '_impl.ten_wrapadr': jp.array(ten_wrapadr, dtype=int),
      '_impl.ten_wrapnum': jp.array(ten_wrapnum, dtype=int),
@@ -1263,7 +1273,7 @@ def transmission(m: Model, d: Data) -> Data:
      wrench = jp.concatenate((frame_xmat @ gear[:3], frame_xmat @ gear[3:]))
      moment = jac @ wrench
    elif trntype == TrnType.TENDON:
-      length = d._impl.ten_length[trnid[0]] * gear[:1]
+      length = d.ten_length[trnid[0]] * gear[:1]
      moment = d._impl.ten_J[trnid[0]] * gear[0]
    else:
      raise RuntimeError(f'unrecognized trntype: {TrnType(trntype)}')
@@ -122,7 +122,7 @@ class SmoothTest(absltest.TestCase):
    # tendon
    dx = jax.jit(mjx.tendon)(mx, mjx.put_data(m, d))
    _assert_attr_eq(d, dx._impl, 'ten_J')
-    _assert_attr_eq(d, dx._impl, 'ten_length')
+    _assert_attr_eq(d, dx, 'ten_length')
    # transmission
    dx = jax.jit(mjx.transmission)(mx, dx)
    _assert_attr_eq(d, dx._impl, 'actuator_length')
@@ -394,7 +394,7 @@ class TendonTest(parameterized.TestCase):
    mujoco.mj_forward(m, d)
    dx = jax.jit(mjx.forward)(mx, dx)

-    _assert_eq(d.ten_length, dx._impl.ten_length, 'ten_length')
+    _assert_eq(d.ten_length, dx.ten_length, 'ten_length')
    _assert_eq(d.ten_J, dx._impl.ten_J, 'ten_J')
    _assert_eq(d.ten_wrapnum, dx._impl.ten_wrapnum, 'ten_wrapnum')
    _assert_eq(d.ten_wrapadr, dx._impl.ten_wrapadr, 'ten_wrapadr')
@@ -998,7 +998,6 @@ class DataC(PyTreeNode):
  ten_J_rowadr: jax.Array  # pylint:disable=invalid-name
  ten_J_colind: jax.Array  # pylint:disable=invalid-name
  ten_J: jax.Array  # pylint:disable=invalid-name
-  ten_length: jax.Array
  wrap_obj: jax.Array
  wrap_xpos: jax.Array
  actuator_length: jax.Array
@@ -1070,7 +1069,6 @@ class DataJAX(PyTreeNode):
  ten_wrapadr: jax.Array
  ten_wrapnum: jax.Array
  ten_J: jax.Array  # pylint:disable=invalid-name
-  ten_length: jax.Array
  wrap_obj: jax.Array
  wrap_xpos: jax.Array
  actuator_length: jax.Array
@@ -1133,6 +1131,7 @@ class Data(PyTreeNode):
  ximat: jax.Array
  xanchor: jax.Array
  xaxis: jax.Array
+  ten_length: jax.Array
  geom_xpos: jax.Array
  geom_xmat: jax.Array
  site_xpos: jax.Array
@@ -1113,7 +1113,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
      'ten_Jdot': d._impl.ten_Jdot.shape,
      'ten_actfrc': d._impl.ten_actfrc.shape,
      'ten_bias_coef': d._impl.ten_bias_coef.shape,
-      'ten_length': d._impl.ten_length.shape,
+      'ten_length': d.ten_length.shape,
      'ten_velocity': d._impl.ten_velocity.shape,
      'ten_wrapadr': d._impl.ten_wrapadr.shape,
      'ten_wrapnum': d._impl.ten_wrapnum.shape,
@@ -1777,7 +1777,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
      d._impl.ten_Jdot,
      d._impl.ten_actfrc,
      d._impl.ten_bias_coef,
-      d._impl.ten_length,
+      d.ten_length,
      d._impl.ten_velocity,
      d._impl.ten_wrapadr,
      d._impl.ten_wrapnum,
@@ -1956,7 +1956,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
      '_impl.ten_Jdot': out[99],
      '_impl.ten_actfrc': out[100],
      '_impl.ten_bias_coef': out[101],
-      '_impl.ten_length': out[102],
+      'ten_length': out[102],
      '_impl.ten_velocity': out[103],
      '_impl.ten_wrapadr': out[104],
      '_impl.ten_wrapnum': out[105],
@@ -3176,7 +3176,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
      'ten_Jdot': d._impl.ten_Jdot.shape,
      'ten_actfrc': d._impl.ten_actfrc.shape,
      'ten_bias_coef': d._impl.ten_bias_coef.shape,
-      'ten_length': d._impl.ten_length.shape,
+      'ten_length': d.ten_length.shape,
      'ten_velocity': d._impl.ten_velocity.shape,
      'ten_wrapadr': d._impl.ten_wrapadr.shape,
      'ten_wrapnum': d._impl.ten_wrapnum.shape,
@@ -3866,7 +3866,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
      d._impl.ten_Jdot,
      d._impl.ten_actfrc,
      d._impl.ten_bias_coef,
-      d._impl.ten_length,
+      d.ten_length,
      d._impl.ten_velocity,
      d._impl.ten_wrapadr,
      d._impl.ten_wrapnum,
@@ -4057,7 +4057,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
      '_impl.ten_Jdot': out[111],
      '_impl.ten_actfrc': out[112],
      '_impl.ten_bias_coef': out[113],
-      '_impl.ten_length': out[114],
+      'ten_length': out[114],
      '_impl.ten_velocity': out[115],
      '_impl.ten_wrapadr': out[116],
      '_impl.ten_wrapnum': out[117],
@@ -142,7 +142,7 @@ class ForwardTest(parameterized.TestCase):
      if m.ncam:
        tu.assert_attr_eq(dx, d, 'cam_xpos')
        tu.assert_eq(dx.cam_xmat, d.cam_xmat.reshape((-1, 3, 3)), 'cam_xmat')
-      tu.assert_attr_eq(dx._impl, d, 'ten_length')
+      tu.assert_attr_eq(dx, d, 'ten_length')
      tu.assert_attr_eq(dx._impl, d, 'ten_J')
      tu.assert_attr_eq(dx._impl, d, 'ten_wrapadr')
      tu.assert_attr_eq(dx._impl, d, 'ten_wrapnum')
@@ -285,3 +285,213 @@ def kinematics(m: types.Model, d: types.Data):
 def kinematics_vmap(unused_axis_size, is_batched, m, d):
  d = kinematics(m, d)
  return d, is_batched[1]
+
+
+_m = mjwarp.Model(
+    **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
+)
+_d = mjwarp.Data(
+    **{f.name: None for f in dataclasses.fields(mjwarp.Data) if f.init}
+)
+_o = mjwarp.Option(
+    **{f.name: None for f in dataclasses.fields(mjwarp.Option) if f.init}
+)
+_s = mjwarp.Statistic(
+    **{f.name: None for f in dataclasses.fields(mjwarp.Statistic) if f.init}
+)
+_c = mjwarp.Contact(
+    **{f.name: None for f in dataclasses.fields(mjwarp.Contact) if f.init}
+)
+_e = mjwarp.Constraint(
+    **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
+)
+
+
+@ffi.format_args_for_warp
+def _tendon_shim(
+    # Model
+    nworld: int,
+    body_parentid: wp.array(dtype=int),
+    body_rootid: wp.array(dtype=int),
+    dof_bodyid: wp.array(dtype=int),
+    geom_bodyid: wp.array(dtype=int),
+    geom_size: wp.array2d(dtype=wp.vec3),
+    jnt_dofadr: wp.array(dtype=int),
+    jnt_qposadr: wp.array(dtype=int),
+    ntendon: int,
+    nv: int,
+    site_bodyid: wp.array(dtype=int),
+    tendon_adr: wp.array(dtype=int),
+    tendon_geom_adr: wp.array(dtype=int),
+    tendon_jnt_adr: wp.array(dtype=int),
+    tendon_num: wp.array(dtype=int),
+    tendon_site_pair_adr: wp.array(dtype=int),
+    wrap_geom_adr: wp.array(dtype=int),
+    wrap_jnt_adr: wp.array(dtype=int),
+    wrap_objid: wp.array(dtype=int),
+    wrap_prm: wp.array(dtype=float),
+    wrap_pulley_scale: wp.array(dtype=float),
+    wrap_site_pair_adr: wp.array(dtype=int),
+    wrap_type: wp.array(dtype=int),
+    # Data
+    cdof: wp.array2d(dtype=wp.spatial_vector),
+    geom_xmat: wp.array2d(dtype=wp.mat33),
+    geom_xpos: wp.array2d(dtype=wp.vec3),
+    qpos: wp.array2d(dtype=float),
+    site_xpos: wp.array2d(dtype=wp.vec3),
+    subtree_com: wp.array2d(dtype=wp.vec3),
+    ten_J: wp.array3d(dtype=float),
+    ten_length: wp.array2d(dtype=float),
+    ten_wrapadr: wp.array2d(dtype=int),
+    ten_wrapnum: wp.array2d(dtype=int),
+    wrap_geom_xpos: wp.array2d(dtype=wp.spatial_vector),
+    wrap_obj: wp.array2d(dtype=wp.vec2i),
+    wrap_xpos: wp.array2d(dtype=wp.spatial_vector),
+):
+  _m.stat = _s
+  _m.opt = _o
+  _d.efc = _e
+  _d.contact = _c
+  _m.body_parentid = body_parentid
+  _m.body_rootid = body_rootid
+  _m.dof_bodyid = dof_bodyid
+  _m.geom_bodyid = geom_bodyid
+  _m.geom_size = geom_size
+  _m.jnt_dofadr = jnt_dofadr
+  _m.jnt_qposadr = jnt_qposadr
+  _m.ntendon = ntendon
+  _m.nv = nv
+  _m.site_bodyid = site_bodyid
+  _m.tendon_adr = tendon_adr
+  _m.tendon_geom_adr = tendon_geom_adr
+  _m.tendon_jnt_adr = tendon_jnt_adr
+  _m.tendon_num = tendon_num
+  _m.tendon_site_pair_adr = tendon_site_pair_adr
+  _m.wrap_geom_adr = wrap_geom_adr
+  _m.wrap_jnt_adr = wrap_jnt_adr
+  _m.wrap_objid = wrap_objid
+  _m.wrap_prm = wrap_prm
+  _m.wrap_pulley_scale = wrap_pulley_scale
+  _m.wrap_site_pair_adr = wrap_site_pair_adr
+  _m.wrap_type = wrap_type
+  _d.cdof = cdof
+  _d.geom_xmat = geom_xmat
+  _d.geom_xpos = geom_xpos
+  _d.qpos = qpos
+  _d.site_xpos = site_xpos
+  _d.subtree_com = subtree_com
+  _d.ten_J = ten_J
+  _d.ten_length = ten_length
+  _d.ten_wrapadr = ten_wrapadr
+  _d.ten_wrapnum = ten_wrapnum
+  _d.wrap_geom_xpos = wrap_geom_xpos
+  _d.wrap_obj = wrap_obj
+  _d.wrap_xpos = wrap_xpos
+  _d.nworld = nworld
+  mjwarp.tendon(_m, _d)
+
+
+def _tendon_jax_impl(m: types.Model, d: types.Data):
+  output_dims = {
+      'cdof': d._impl.cdof.shape,
+      'geom_xmat': d.geom_xmat.shape,
+      'geom_xpos': d.geom_xpos.shape,
+      'qpos': d.qpos.shape,
+      'site_xpos': d.site_xpos.shape,
+      'subtree_com': d.subtree_com.shape,
+      'ten_J': d._impl.ten_J.shape,
+      'ten_length': d.ten_length.shape,
+      'ten_wrapadr': d._impl.ten_wrapadr.shape,
+      'ten_wrapnum': d._impl.ten_wrapnum.shape,
+      'wrap_geom_xpos': d._impl.wrap_geom_xpos.shape,
+      'wrap_obj': d._impl.wrap_obj.shape,
+      'wrap_xpos': d._impl.wrap_xpos.shape,
+  }
+  jf = ffi.jax_callable_variadic_tuple(
+      _tendon_shim,
+      num_outputs=13,
+      output_dims=output_dims,
+      vmap_method=None,
+      in_out_argnames={
+          'cdof',
+          'geom_xmat',
+          'geom_xpos',
+          'qpos',
+          'site_xpos',
+          'subtree_com',
+          'ten_J',
+          'ten_length',
+          'ten_wrapadr',
+          'ten_wrapnum',
+          'wrap_geom_xpos',
+          'wrap_obj',
+          'wrap_xpos',
+      },
+  )
+  out = jf(
+      d.qpos.shape[0],
+      m.body_parentid,
+      m.body_rootid,
+      m.dof_bodyid,
+      m.geom_bodyid,
+      m.geom_size,
+      m.jnt_dofadr,
+      m.jnt_qposadr,
+      m.ntendon,
+      m.nv,
+      m.site_bodyid,
+      m.tendon_adr,
+      m._impl.tendon_geom_adr,
+      m._impl.tendon_jnt_adr,
+      m.tendon_num,
+      m._impl.tendon_site_pair_adr,
+      m._impl.wrap_geom_adr,
+      m._impl.wrap_jnt_adr,
+      m.wrap_objid,
+      m.wrap_prm,
+      m._impl.wrap_pulley_scale,
+      m._impl.wrap_site_pair_adr,
+      m.wrap_type,
+      d._impl.cdof,
+      d.geom_xmat,
+      d.geom_xpos,
+      d.qpos,
+      d.site_xpos,
+      d.subtree_com,
+      d._impl.ten_J,
+      d.ten_length,
+      d._impl.ten_wrapadr,
+      d._impl.ten_wrapnum,
+      d._impl.wrap_geom_xpos,
+      d._impl.wrap_obj,
+      d._impl.wrap_xpos,
+  )
+  d = d.tree_replace({
+      '_impl.cdof': out[0],
+      'geom_xmat': out[1],
+      'geom_xpos': out[2],
+      'qpos': out[3],
+      'site_xpos': out[4],
+      'subtree_com': out[5],
+      '_impl.ten_J': out[6],
+      'ten_length': out[7],
+      '_impl.ten_wrapadr': out[8],
+      '_impl.ten_wrapnum': out[9],
+      '_impl.wrap_geom_xpos': out[10],
+      '_impl.wrap_obj': out[11],
+      '_impl.wrap_xpos': out[12],
+  })
+  return d
+
+
+@jax.custom_batching.custom_vmap
+@ffi.marshal_jax_warp_callable
+def tendon(m: types.Model, d: types.Data):
+  return _tendon_jax_impl(m, d)
+
+
+@tendon.def_vmap
+@ffi.marshal_custom_vmap
+def tendon_vmap(unused_axis_size, is_batched, m, d):
+  d = tendon(m, d)
+  return d, is_batched[1]
@@ -19,6 +19,7 @@ import os
 import tempfile

 from absl.testing import absltest
+from absl.testing import parameterized
 import jax
 from jax import numpy as jp
 import mujoco
@@ -38,7 +39,7 @@ except ImportError:
 _FORCE_TEST = os.environ.get('MJX_WARP_FORCE_TEST', '0') == '1'


-class SmoothTest(absltest.TestCase):
+class SmoothTest(parameterized.TestCase):

  def setUp(self):
    super().setUp()
@@ -237,6 +238,5 @@ class SmoothTest(absltest.TestCase):
      tu.assert_attr_eq(d, dx, 'site_xpos')
      tu.assert_eq(d.site_xmat.reshape((-1, 3, 3)), dx.site_xmat, 'site_xmat')

-
 if __name__ == '__main__':
  absltest.main()
@@ -152,7 +152,8 @@ def _mjx_efc(dx, worldid: int):
    return 0, empty, empty, np.zeros((0, dx.qvel.shape[0])), empty, empty
  efc_pos = select(dx._impl.efc__pos[:nefc])
  efc_type = select(dx._impl.efc__type[:nefc])
-  keys_sorted = np.lexsort((-efc_pos, efc_type))
+  efc_d = select(dx._impl.efc__D[:nefc])
+  keys_sorted = np.lexsort((-efc_pos, efc_type, efc_d))
  keys = keys[keys_sorted]

  nefc = len(keys)
@@ -178,7 +179,7 @@ def _mj_efc(d):
  else:
    efc_j = d.efc_J.reshape((-1, d.qvel.shape[0]))

-  keys = np.lexsort((-d.efc_pos, d.efc_type))
+  keys = np.lexsort((-d.efc_pos, d.efc_type, d.efc_D))
  type_ = d.efc_type[keys]
  pos = d.efc_pos[keys]
  efc_j = efc_j[keys]
@@ -361,7 +361,6 @@ class DataWarp(PyTreeNode):
  ten_Jdot: jax.Array
  ten_actfrc: jax.Array
  ten_bias_coef: jax.Array
-  ten_length: jax.Array
  ten_velocity: jax.Array
  ten_wrapadr: jax.Array
  ten_wrapnum: jax.Array