Add tendon actuator force limits to MJX.

PiperOrigin-RevId: 745529685
Change-Id: I43c372116daefc3dd2ed6ad1bb9c6e3978bc8527

doc/changelog.rst

@@ -23,6 +23,10 @@ General
 - Added :ref:`tendon actuator force limits<tendon-spatial-actuatorfrclimited>` and
   :ref:`tendon actuator force sensor<sensor-tendonactuatorfrc>`.
 
+MJX
+^^^
+- Added tendon actuator force limits.
+
 Bug fixes
 ^^^^^^^^^
 - :ref:`mj_jacDot` was missing a term that accounts for the motion of the point with respect to

mjx/mujoco/mjx/_src/forward.py

@@ -38,6 +38,7 @@ from mujoco.mjx._src.types import GainType
 from mujoco.mjx._src.types import IntegratorType
 from mujoco.mjx._src.types import JointType
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import TrnType
 # pylint: enable=g-importing-member
 import numpy as np
 
@@ -178,6 +179,34 @@ def fwd_actuation(m: Model, d: Data) -> Data:
       jp.array(m.actuator_acc0),
       group_by='u',
   )
+
+  # tendon total force clamping
+  if np.any(m.tendon_actfrclimited):
+    (tendon_actfrclimited_id,) = np.nonzero(m.tendon_actfrclimited)
+    actuator_tendon = m.actuator_trntype == TrnType.TENDON
+
+    force_mask = [
+        actuator_tendon & (m.actuator_trnid[:, 0] == tendon_id)
+        for tendon_id in tendon_actfrclimited_id
+    ]
+    force_ids = np.concatenate([np.nonzero(mask)[0] for mask in force_mask])
+    force_mat = np.array(force_mask)[:, force_ids]
+    tendon_total_force = force_mat @ force[force_ids]
+
+    force_scaling = jp.where(
+        tendon_total_force < m.tendon_actfrcrange[tendon_actfrclimited_id, 0],
+        m.tendon_actfrcrange[tendon_actfrclimited_id, 0] / tendon_total_force,
+        1,
+    )
+    force_scaling = jp.where(
+        tendon_total_force > m.tendon_actfrcrange[tendon_actfrclimited_id, 1],
+        m.tendon_actfrcrange[tendon_actfrclimited_id, 1] / tendon_total_force,
+        force_scaling,
+    )
+
+    tendon_forces = force[force_ids] * (force_mat.T @ force_scaling)
+    force = force.at[force_ids].set(tendon_forces)
+
   forcerange = jp.where(
       m.actuator_forcelimited[:, None],
       m.actuator_forcerange,

mjx/mujoco/mjx/_src/forward_test.py

@@ -17,6 +17,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
+from jax import numpy as jp
 import mujoco
 from mujoco import mjx
 from mujoco.mjx._src import test_util
@@ -196,6 +197,21 @@ class ActuatorTest(parameterized.TestCase):
     dx = jax.jit(mjx.euler)(mx, dx)
     _assert_attr_eq(d, dx, 'act')
 
+  def test_tendon_force_clamp(self):
+    m = test_util.load_test_file('actuator/tendon_force_clamp.xml')
+    d = mujoco.MjData(m)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+
+    dx = dx.replace(ctrl=jp.array([1.0, 1.0, 1.0, -1.0, 1.0, -20.0, 5.0, -5.0]))
+    dx = mjx.forward(mx, dx)
+
+    _assert_eq(
+        dx.actuator_force,
+        jp.array([1.0, 1.0, 1.0, -1.0, 1.0, -10.0, 5.0, -5.0]),
+        'actuator_force',
+    )
+
 
 if __name__ == '__main__':
   absltest.main()

mjx/mujoco/mjx/test_data/actuator/tendon_force_clamp.xml

@@ -0,0 +1,49 @@
+<mujoco model="fixed_site">
+  <worldbody>
+    <body>
+      <joint name="joint0" type="hinge" axis="0 1 0"/>
+      <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+      <site name="site0" pos="0.25 0 0.1" size="0.025"/>
+      <body pos="0.5 0 0">
+        <joint name="joint1" type="hinge" axis="0 1 0"/>
+        <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+        <site name="site1" pos="0.25 0 0.1" size="0.025"/>
+        <body pos="0.5 0 0">
+          <joint name="joint2" type="hinge" axis="0 1 0"/>
+          <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+          <site name="site2" pos="0.25 0 0.1" size="0.025"/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+  <tendon>
+    <spatial name="spatial0" width="0.0125">
+      <site site="site0"/>
+      <site site="site1"/>
+      <site site="site2"/>
+    </spatial>
+    <spatial name="spatial1" width="0.0125" actuatorfrclimited="true" actuatorfrcrange="-1 1">
+      <site site="site1"/>
+      <site site="site2"/>
+    </spatial>
+    <fixed name="fixed0" actuatorfrclimited="true" actuatorfrcrange="-10 10">
+      <joint joint="joint0" coef=".1"/>
+      <joint joint="joint1" coef=".2"/>
+      <joint joint="joint2" coef=".3"/>
+    </fixed>
+    <fixed name="fixed1" actuatorfrclimited="true" actuatorfrcrange="0 1">
+      <joint joint="joint0" coef=".1"/>
+      <joint joint="joint2" coef=".3"/>
+    </fixed>
+  </tendon>
+  <actuator>
+    <motor tendon="spatial0"/>
+    <motor tendon="spatial0"/>
+    <motor tendon="spatial0"/>
+    <motor tendon="spatial1"/>
+    <motor tendon="spatial1"/>
+    <motor tendon="fixed0"/>
+    <motor tendon="fixed1"/>
+    <motor tendon="fixed1"/>
+  </actuator>
+</mujoco>