Optimize mj_tendonBias by computing Jdot * qvel directly

PiperOrigin-RevId: 897628283
Change-Id: Iee26a95d6aaa379730b89be14746308742111e2d

src/engine/engine_core_smooth.c

@@ -1093,22 +1093,20 @@ void mj_tendon(const mjModel* m, mjData* d) {
 }
 
 
-// compute time derivative of dense tendon Jacobian for one tendon
-void mj_tendonDot(const mjModel* m, mjData* d, int id, mjtNum* Jdot) {
+// return dot product of tendon Jacobian time derivative with vector
+mjtNum mj_tendonDot(const mjModel* m, mjData* d, int id, const mjtNum* vec) {
   int nv = m->nv;
+  mjtNum res = 0;
 
   // tendon id is invalid: return
   if (id < 0 || id >= m->ntendon) {
-    return;
+    return 0;
   }
 
-  // clear output
-  mju_zero(Jdot, nv);
-
   // fixed tendon has zero Jdot: return
   int adr = m->tendon_adr[id];
   if (m->wrap_type[adr] == mjWRAP_JOINT) {
-    return;
+    return 0;
   }
 
   // allocate stack arrays
@@ -1194,9 +1192,8 @@ void mj_tendonDot(const mjModel* m, mjData* d, int id, mjtNum* Jdot) {
           // chain rule, first term: Jdot += d/dt(jac2 - jac1) * dpnt
           mju_mulMatTVec(tmp, jacdif, dpnt, 3, NV);
 
-          // scatter into dense output
           for (int k=0; k < NV; k++) {
-            Jdot[chain[k]] += tmp[k] / divisor;
+            res += (tmp[k] / divisor) * vec[chain[k]];
           }
 
           // get endpoint Jacobians, subtract
@@ -1207,9 +1204,8 @@ void mj_tendonDot(const mjModel* m, mjData* d, int id, mjtNum* Jdot) {
           // chain rule, second term: Jdot += (jac2 - jac1) * d/dt(dpnt)
           mju_mulMatTVec(tmp, jacdif, dvel, 3, NV);
 
-          // scatter into dense output
           for (int k=0; k < NV; k++) {
-            Jdot[chain[k]] += tmp[k] / divisor;
+            res += (tmp[k] / divisor) * vec[chain[k]];
           }
         }
       }
@@ -1224,8 +1220,7 @@ void mj_tendonDot(const mjModel* m, mjData* d, int id, mjtNum* Jdot) {
         // chain rule, first term: Jdot += d/dt(jac2 - jac1) * dpnt
         mju_mulMatTVec(tmp, jacdif, dpnt, 3, nv);
 
-        // add to existing
-        mju_addToScl(Jdot, tmp, 1/divisor, nv);
+        res += mju_dot(tmp, vec, nv) / divisor;
 
         // get endpoint Jacobians, subtract
         mj_jac(m, d, jac1, 0, wpnt, wbody[0]);
@@ -1235,8 +1230,7 @@ void mj_tendonDot(const mjModel* m, mjData* d, int id, mjtNum* Jdot) {
         // chain rule, second term: Jdot += (jac2 - jac1) * d/dt(dpnt)
         mju_mulMatTVec(tmp, jacdif, dvel, 3, nv);
 
-        // add to existing
-        mju_addToScl(Jdot, tmp, 1/divisor, nv);
+        res += mju_dot(tmp, vec, nv) / divisor;
       }
     }
 
@@ -1245,6 +1239,7 @@ void mj_tendonDot(const mjModel* m, mjData* d, int id, mjtNum* Jdot) {
   }
 
   mj_freeStack(d);
+  return res;
 }
 
 
@@ -2668,9 +2663,7 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) {
 // add bias force due to tendon armature
 void mj_tendonBias(const mjModel* m, mjData* d, mjtNum* qfrc) {
   int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree;
-  int ntendon = m->ntendon, nv = m->nv;
-  mjtNum* ten_Jdot = NULL;
-  mj_markStack(d);
+  int ntendon = m->ntendon;
 
   // add bias term due to tendon armature
   for (int i=0; i < ntendon; i++) {
@@ -2686,16 +2679,11 @@ void mj_tendonBias(const mjModel* m, mjData* d, mjtNum* qfrc) {
       continue;
     }
 
-    // allocate if required
-    if (!ten_Jdot) {
-      ten_Jdot = mjSTACKALLOC(d, nv, mjtNum);
-    }
-
-    // get dense d/dt(tendon Jacobian) for tendon i
-    mj_tendonDot(m, d, i, ten_Jdot);
+    // get d/dt(tendon Jacobian) dotted with qvel for tendon i
+    mjtNum dot = mj_tendonDot(m, d, i, d->qvel);
 
     // add bias term:  qfrc += ten_J * armature * dot(ten_Jdot, qvel)
-    mjtNum coef = armature * mju_dot(ten_Jdot, d->qvel, nv);
+    mjtNum coef = armature * dot;
 
     if (coef) {
       // sparse
@@ -2708,6 +2696,4 @@ void mj_tendonBias(const mjModel* m, mjData* d, mjtNum* qfrc) {
       }
     }
   }
-
-  mj_freeStack(d);
 }

src/engine/engine_core_smooth.h

@@ -18,6 +18,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
+#include <mujoco/mjtnum.h>
 
 #ifdef __cplusplus
 extern "C" {
@@ -45,8 +46,8 @@ MJAPI void mj_flex(const mjModel* m, mjData* d);
 // compute tendon lengths, velocities and moment arms
 MJAPI void mj_tendon(const mjModel* m, mjData* d);
 
-// compute time derivative of dense tendon Jacobian for one tendon
-MJAPI void mj_tendonDot(const mjModel* m, mjData* d, int id, mjtNum* Jdot);
+// return dot product of tendon Jacobian time derivative with vector
+MJAPI mjtNum mj_tendonDot(const mjModel* m, mjData* d, int id, const mjtNum* vec);
 
 // compute actuator transmission lengths and moments
 MJAPI void mj_transmission(const mjModel* m, mjData* d);

test/engine/engine_core_smooth_test.cc

@@ -191,10 +191,8 @@ TEST_F(CoreSmoothTest, TendonJdot) {
 
       mj_forward(m, d);
 
-      // get current J and Jdot for the tendon
+      // get current J for the tendon
       vector<mjtNum> ten_J(d->ten_J, d->ten_J + nv);
-      vector<mjtNum> ten_Jdot(nv, 0);
-      mj_tendonDot(m, d, 0, ten_Jdot.data());
 
       // compute finite-differenced Jdot
       mjtNum h = MjTol(1e-7, 5e-4);
@@ -206,7 +204,10 @@ TEST_F(CoreSmoothTest, TendonJdot) {
       mju_subFrom(ten_Jh.data(), ten_J.data(), nv);
       mju_scl(ten_Jh.data(), ten_Jh.data(), 1.0 / h, nv);
 
-      EXPECT_THAT(ten_Jdot, Pointwise(MjNear(1e-6, 2e-3), ten_Jh));
+      // test dot product against finite differences
+      mjtNum dot = mj_tendonDot(m, d, 0, d->qvel);
+      mjtNum expected_dot = mju_dot(ten_Jh.data(), d->qvel, nv);
+      EXPECT_NEAR(dot, expected_dot, MjTol(1e-5, 2e-3));
     }
 
     mj_deleteData(d);