Implement sparse Jacobian time derivative

PiperOrigin-RevId: 897609541
Change-Id: Ic86b6026cfc369c584d3edfe12a9a2fd14cc4357

src/engine/engine_core_smooth.c

@@ -1113,6 +1113,8 @@ void mj_tendonDot(const mjModel* m, mjData* d, int id, mjtNum* Jdot) {
 
   // allocate stack arrays
   mj_markStack(d);
+  int issparse = mj_isSparse(m);
+  int* chain = issparse ? mjSTACKALLOC(d, nv, int) : NULL;
   mjtNum* jac1 = mjSTACKALLOC(d, 3*nv, mjtNum);
   mjtNum* jac2 = mjSTACKALLOC(d, 3*nv, mjtNum);
   mjtNum* jacdif = mjSTACKALLOC(d, 3*nv, mjtNum);
@@ -1178,30 +1180,64 @@ void mj_tendonDot(const mjModel* m, mjData* d, int id, mjtNum* Jdot) {
       mju_addToScl3(dvel, dpnt, -dot);
       mju_scl3(dvel, dvel, norm > mjMINVAL ? 1/norm : 0);
 
-      // TODO(tassa ) write sparse branch, requires mj_jacDotSparse
-      // if (mj_isSparse(m)) { ... }
+      // sparse
+      if (issparse) {
+        // construct merged chain
+        int NV = mj_mergeChain(m, chain, wbody[0], wbody[1], /*flg_skipcommon=*/0);
 
-      // get endpoint JacobianDots, subtract
-      mj_jacDot(m, d, jac1, 0, wpnt, wbody[0]);
-      mj_jacDot(m, d, jac2, 0, wpnt+3, wbody[1]);
-      mju_sub(jacdif, jac2, jac1, 3*nv);
+        if (NV) {
+          // get endpoint JacobianDots, subtract
+          mj_jacDotSparse(m, d, jac1, 0, wpnt, wbody[0], NV, chain);
+          mj_jacDotSparse(m, d, jac2, 0, wpnt+3, wbody[1], NV, chain);
+          mju_sub(jacdif, jac2, jac1, 3*NV);
 
-      // chain rule, first term: Jdot += d/dt(jac2 - jac1) * dpnt
-      mju_mulMatTVec(tmp, jacdif, dpnt, 3, nv);
+          // 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);
+          // scatter into dense output
+          for (int k=0; k < NV; k++) {
+            Jdot[chain[k]] += tmp[k] / divisor;
+          }
 
-      // get endpoint Jacobians, subtract
-      mj_jac(m, d, jac1, 0, wpnt, wbody[0]);
-      mj_jac(m, d, jac2, 0, wpnt+3, wbody[1]);
-      mju_sub(jacdif, jac2, jac1, 3*nv);
+          // get endpoint Jacobians, subtract
+          mj_jacSparse(m, d, jac1, 0, wpnt, wbody[0], NV, chain, /*flg_skipcommon=*/0);
+          mj_jacSparse(m, d, jac2, 0, wpnt+3, wbody[1], NV, chain, /*flg_skipcommon=*/0);
+          mju_sub(jacdif, jac2, jac1, 3*NV);
 
-      // chain rule, second term: Jdot += (jac2 - jac1) * d/dt(dpnt)
-      mju_mulMatTVec(tmp, jacdif, dvel, 3, nv);
+          // 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);
+          // scatter into dense output
+          for (int k=0; k < NV; k++) {
+            Jdot[chain[k]] += tmp[k] / divisor;
+          }
+        }
+      }
+
+      // dense
+      else {
+        // get endpoint JacobianDots, subtract
+        mj_jacDot(m, d, jac1, 0, wpnt, wbody[0]);
+        mj_jacDot(m, d, jac2, 0, wpnt+3, wbody[1]);
+        mju_sub(jacdif, jac2, jac1, 3*nv);
+
+        // 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);
+
+        // get endpoint Jacobians, subtract
+        mj_jac(m, d, jac1, 0, wpnt, wbody[0]);
+        mj_jac(m, d, jac2, 0, wpnt+3, wbody[1]);
+        mju_sub(jacdif, jac2, jac1, 3*nv);
+
+        // 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);
+      }
     }
 
     // advance

src/engine/engine_core_util.c

@@ -660,6 +660,92 @@ void mj_jacDot(const mjModel* m, const mjData* d,
 }
 
 
+// compute 3/6-by-NV sparse Jacobian time derivative of global point attached to given body
+void mj_jacDotSparse(const mjModel* m, const mjData* d,
+                     mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body,
+                     int NV, const int* chain) {
+  mjtNum offset[3];
+  mjtNum pvel[6];
+
+  // clear jacobians, compute offset and pvel if required
+  if (jacp) {
+    mju_zero(jacp, 3*NV);
+    const mjtNum* com = d->subtree_com+3*m->body_rootid[body];
+    mju_sub3(offset, point, com);
+    mju_transformSpatial(pvel, d->cvel+6*body, 0, point, com, 0);
+  }
+  if (jacr) {
+    mju_zero(jacr, 3*NV);
+  }
+
+  // skip fixed bodies
+  body = m->body_weldid[body];
+
+  // no movable body found: nothing to do
+  if (!body) {
+    return;
+  }
+
+  // get last dof that affects this body
+  int da = m->body_dofadr[body] + m->body_dofnum[body] - 1;
+
+  // start at end of chain (chain is in increasing order)
+  int ci = NV-1;
+
+  // backward pass over dof ancestor chain
+  while (da >= 0) {
+    // find chain index for this dof
+    while (ci >= 0 && chain[ci] > da) {
+      ci--;
+    }
+
+    // dof not in chain: SHOULD NOT OCCUR
+    if (ci < 0 || chain[ci] != da) {
+      mjERROR("dof index %d not found in chain", da);
+    }
+
+    mjtNum cdof_dot[6];
+    mji_copy6(cdof_dot, d->cdof_dot+6*da);
+    mjtNum* cdof = d->cdof+6*da;
+
+    // check for quaternion
+    mjtJoint type = m->jnt_type[m->dof_jntid[da]];
+    int dofadr = m->jnt_dofadr[m->dof_jntid[da]];
+    int is_quat = type == mjJNT_BALL || (type == mjJNT_FREE && da >= dofadr + 3);
+
+    // compute cdof_dot for quaternion (use current body cvel)
+    if (is_quat) {
+      mji_crossMotion(cdof_dot, d->cvel+6*m->dof_bodyid[da], cdof);
+    }
+
+    // construct rotation jacobian
+    if (jacr) {
+      jacr[ci+0*NV] += cdof_dot[0];
+      jacr[ci+1*NV] += cdof_dot[1];
+      jacr[ci+2*NV] += cdof_dot[2];
+    }
+
+    // construct translation jacobian (correct for rotation)
+    if (jacp) {
+      // first correction term, account for varying cdof
+      mjtNum tmp1[3];
+      mji_cross(tmp1, cdof_dot, offset);
+
+      // second correction term, account for point translational velocity
+      mjtNum tmp2[3];
+      mji_cross(tmp2, cdof, pvel + 3);
+
+      jacp[ci+0*NV] += cdof_dot[3] + tmp1[0] + tmp2[0];
+      jacp[ci+1*NV] += cdof_dot[4] + tmp1[1] + tmp2[1];
+      jacp[ci+2*NV] += cdof_dot[5] + tmp1[2] + tmp2[2];
+    }
+
+    // advance to parent dof
+    da = m->dof_parentid[da];
+  }
+}
+
+
 // compute subtree angular momentum matrix
 void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body) {
   int nv = m->nv;

src/engine/engine_core_util.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" {
@@ -85,6 +86,11 @@ void mj_jacSparseSimple(const mjModel* m, const mjData* d,
                         mjtNum* jacdifp, mjtNum* jacdifr, const mjtNum* point,
                         int body, int flg_second, int NV, int start);
 
+// compute 3/6-by-NV sparse Jacobian time derivative of global point attached to given body
+MJAPI void mj_jacDotSparse(const mjModel* m, const mjData* d,
+                           mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body,
+                           int NV, const int* chain);
+
 // dense or sparse Jacobian difference for two body points: pos2 - pos1, global
 MJAPI int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
                         int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3],