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

PiperOrigin-RevId: 794813523
Change-Id: Ic5684788eca5a3f4fd86193ec5d035d40b3aad35

Author: Google DeepMind

doc/changelog.rst

@@ -2,19 +2,6 @@
 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)
 -----------------------------------
 

mjx/mujoco/mjx/_src/constraint.py

@@ -322,8 +322,8 @@ def rows(
 
   inv1, inv2 = m.tendon_invweight0[obj1id], m.tendon_invweight0[obj2id]
   jac1, jac2 = d._impl.ten_J[obj1id], d._impl.ten_J[obj2id]
-  pos1 = d.ten_length[obj1id] - m.tendon_length0[obj1id]
-  pos2 = d.ten_length[obj2id] - m.tendon_length0[obj2id]
+  pos1 = d._impl.ten_length[obj1id] - m.tendon_length0[obj1id]
+  pos2 = d._impl.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.ten_length,
+          d._impl.ten_length,
           d._impl.ten_J,
           m.tendon_range,
           m.tendon_margin,

mjx/mujoco/mjx/_src/io.py

@@ -576,7 +576,6 @@ 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()
@@ -638,6 +637,7 @@ 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,6 +742,7 @@ def get(m, name: str):
       '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),

mjx/mujoco/mjx/_src/io_test.py

@@ -478,7 +478,7 @@ def test_put_data(self, impl: str):
     np.testing.assert_allclose(dx.site_xmat.reshape((1, 9)), d.site_xmat)
 
     # tendon data is correct
-    np.testing.assert_allclose(dx.ten_length, d.ten_length)
+    np.testing.assert_allclose(dx._impl.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)))

mjx/mujoco/mjx/_src/passive.py

@@ -75,7 +75,7 @@ def fn(jnt_typs, stiffness, qpos_spring, qpos):
   qfrc -= m.dof_damping * d.qvel
 
   # tendon-level spring-dampers
-  below, above = m.tendon_lengthspring.T - d.ten_length
+  below, above = m.tendon_lengthspring.T - d._impl.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

mjx/mujoco/mjx/_src/sensor.py

@@ -22,6 +22,7 @@
 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
@@ -173,7 +174,7 @@ def _cam_project(
     elif sensor_type == SensorType.JOINTPOS:
       sensor = d.qpos[m.jnt_qposadr[objid]]
     elif sensor_type == SensorType.TENDONPOS:
-      sensor = d.ten_length[objid]
+      sensor = d._impl.ten_length[objid]
     elif sensor_type == SensorType.ACTUATORPOS:
       sensor = d._impl.actuator_length[objid]
     elif sensor_type == SensorType.BALLQUAT:

mjx/mujoco/mjx/_src/smooth.py

@@ -26,24 +26,18 @@
 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
 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
@@ -850,10 +844,6 @@ def _forward(carry, cfrc_ext, cinert, cvel, body_dofadr, body_dofnum):
 
 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.')
 
@@ -1101,7 +1091,7 @@ def _distance(p0, p1):
 
   # assemble length and moment
   ten_length = (
-      jp.zeros_like(d.ten_length).at[tendon_id_jnt].set(length_jnt)
+      jp.zeros_like(d._impl.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)
@@ -1171,7 +1161,7 @@ def _distance(p0, p1):
   ).reshape((m.nwrap, 2))
 
   return d.tree_replace({
-      'ten_length': ten_length,
+      '_impl.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),
@@ -1273,7 +1263,7 @@ def fn(
       wrench = jp.concatenate((frame_xmat @ gear[:3], frame_xmat @ gear[3:]))
       moment = jac @ wrench
     elif trntype == TrnType.TENDON:
-      length = d.ten_length[trnid[0]] * gear[:1]
+      length = d._impl.ten_length[trnid[0]] * gear[:1]
       moment = d._impl.ten_J[trnid[0]] * gear[0]
     else:
       raise RuntimeError(f'unrecognized trntype: {TrnType(trntype)}')

mjx/mujoco/mjx/_src/smooth_test.py

@@ -122,7 +122,7 @@ def test_smooth(self):
     # 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, 'ten_length')
+    _assert_attr_eq(d, dx._impl, 'ten_length')
     # transmission
     dx = jax.jit(mjx.transmission)(mx, dx)
     _assert_attr_eq(d, dx._impl, 'actuator_length')
@@ -394,7 +394,7 @@ def test_tendon(self, filename):
     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_length, dx._impl.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')

mjx/mujoco/mjx/_src/types.py

@@ -998,6 +998,7 @@ 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
@@ -1069,6 +1070,7 @@ 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
@@ -1131,7 +1133,6 @@ 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

mjx/mujoco/mjx/warp/forward.py

@@ -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.ten_length.shape,
+      'ten_length': d._impl.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.ten_length,
+      d._impl.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],
-      'ten_length': out[102],
+      '_impl.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.ten_length.shape,
+      'ten_length': d._impl.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.ten_length,
+      d._impl.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],
-      'ten_length': out[114],
+      '_impl.ten_length': out[114],
       '_impl.ten_velocity': out[115],
       '_impl.ten_wrapadr': out[116],
       '_impl.ten_wrapnum': out[117],