Add JOINTINPARENT actuator transmission type to MJX.

PiperOrigin-RevId: 686083761
Change-Id: I5159d050812d32e15ced06df132d3e434fced814

doc/changelog.rst

@@ -26,6 +26,7 @@ MJX
 - Fixed a bug with frictionloss constraints.
 - Added ``TENDONPOS`` and ``TENDONVEL`` sensors.
 - Fixed a bug with the computation of tangential contact forces in ``_decode_pyramid``.
+- Added ``JOINTINPARENT`` actuator transmission type.
 
 Python bindings
 ^^^^^^^^^^^^^^^

doc/mjx.rst

@@ -188,7 +188,7 @@ The following features are **fully supported** in MJX:
    * - :ref:`Joint <mjtJoint>`
      - ``FREE``, ``BALL``, ``SLIDE``, ``HINGE``
    * - :ref:`Transmission <mjtTrn>`
-     - ``TRN_JOINT``, ``TRN_SITE``, ``TRN_TENDON``
+     - ``JOINT``, ``JOINTINPARENT``, ``SITE``, ``TENDON``
    * - :ref:`Actuator Dynamics <mjtDyn>`
      - ``NONE``, ``INTEGRATOR``, ``FILTER``, ``FILTEREXACT``
    * - :ref:`Actuator Gain <mjtGain>`
@@ -270,7 +270,7 @@ The following features are **unsupported**:
    * - :ref:`margin<body-geom-margin>` and :ref:`gap<body-geom-gap>`
      - Unimplemented for collisions with ``Mesh`` :ref:`Geom <mjtGeom>`.
    * - :ref:`Transmission <mjtTrn>`
-     - ``TRN_JOINTINPARENT``, ``TRN_SLIDERCRANK``, ``TRN_BODY``
+     - ``SLIDERCRANK``, ``BODY``
    * - :ref:`Actuator Dynamics <mjtDyn>`
      - ``USER``
    * - :ref:`Actuator Gain <mjtGain>`

mjx/mujoco/mjx/_src/passive_test.py

@@ -42,7 +42,9 @@ class PassiveTest(absltest.TestCase):
     m = test_util.load_test_file('pendula.xml')
     d = mujoco.MjData(m)
     # give the system a little kick to ensure we have non-identity rotations
-    d.ctrl = np.array([0.1, -0.1, 0.2, 0.3, -0.4, 0.5, -0.6, 0.1, -0.2])
+    d.ctrl = np.array(
+        [0.1, -0.1, 0.2, 0.3, -0.4, 0.5, -0.6, 0.1, -0.2, 0.1, 0.2]
+    )
     mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
     mujoco.mj_forward(m, d)
     mx = mjx.put_model(m)

mjx/mujoco/mjx/_src/scan.py

@@ -240,7 +240,7 @@ def flat(
         'a': m.actuator_actadr[i],
         'j': (
             m.actuator_trnid[i, 0]
-            if m.actuator_trntype[i] == TrnType.JOINT
+            if m.actuator_trntype[i] in (TrnType.JOINT, TrnType.JOINTINPARENT)
             else -1
         ),
         's': (
@@ -250,7 +250,7 @@ def flat(
         ),
     }
     v, q = np.array([-1]), np.array([-1])
-    if m.actuator_trntype[i] == TrnType.JOINT:
+    if m.actuator_trntype[i] in (TrnType.JOINT, TrnType.JOINTINPARENT):
       # v/q are associated with the joint transmissions only
       v = np.nonzero(m.dof_jntid == typ_ids['j'])[0]
       q = np.nonzero(_q_jointid(m) == typ_ids['j'])[0]

mjx/mujoco/mjx/_src/smooth.py

@@ -1022,15 +1022,23 @@ def transmission(m: Model, d: Data) -> Data:
       site_xmat,
       site_quat,
   ):
-    if trntype == TrnType.JOINT:
+    if trntype in (TrnType.JOINT, TrnType.JOINTINPARENT):
       if jnt_typ == JointType.FREE:
         length = jp.zeros(1)
         moment = gear
+        if trntype == TrnType.JOINTINPARENT:
+          quat_neg = math.quat_inv(qpos[3:])
+          gearaxis = math.rotate(gear[3:], quat_neg)
+          moment = moment.at[3:].set(gearaxis)
         m_j = m_j + jp.arange(6)
       elif jnt_typ == JointType.BALL:
         axis, angle = math.quat_to_axis_angle(qpos)
-        length = jp.dot(axis * angle, gear[:3])[None]
-        moment = gear[:3]
+        gearaxis = gear[:3]
+        if trntype == TrnType.JOINTINPARENT:
+          quat_neg = math.quat_inv(qpos)
+          gearaxis = math.rotate(gear[:3], quat_neg)
+        length = jp.dot(axis * angle, gearaxis)[None]
+        moment = gearaxis
         m_j = m_j + jp.arange(3)
       elif jnt_typ in (JointType.SLIDE, JointType.HINGE):
         length = qpos * gear[0]

mjx/mujoco/mjx/_src/types.py

@@ -215,12 +215,15 @@ class TrnType(enum.IntEnum):
 
   Members:
     JOINT: force on joint
+    JOINTINPARENT: force on joint, expressed in parent frame
+    TENDON: force on tendon
     SITE: force on site
   """
   JOINT = mujoco.mjtTrn.mjTRN_JOINT
+  JOINTINPARENT = mujoco.mjtTrn.mjTRN_JOINTINPARENT
   SITE = mujoco.mjtTrn.mjTRN_SITE
   TENDON = mujoco.mjtTrn.mjTRN_TENDON
-  # unsupported: JOINTINPARENT, SLIDERCRANK, BODY
+  # unsupported: SLIDERCRANK, BODY
 
 
 class DynType(enum.IntEnum):

mjx/mujoco/mjx/test_data/pendula.xml

@@ -99,7 +99,7 @@
 
     <!-- falling pendulum -->
     <body pos="4.0 0 0">
-      <freejoint/>
+      <freejoint name="freejoint"/>
       <geom/>
       <body pos="0 0 -0.8">
         <joint name="joint13" axis="0.4 0.5 0.6" type="slide" armature="0.02" range="-0.4 0.6"/>
@@ -154,5 +154,7 @@
     <motor gear="150" joint="joint16" name="act7"/>
     <motor gear="150" joint="joint17" name="act8"/>
     <position tendon="tendon_2" kp="100"/>
+    <motor gear="2 3 4" jointinparent="joint1"/>
+    <motor gear="5 6 7 8 9 10" jointinparent="freejoint"/>
   </actuator>
 </mujoco>