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
+      // sparse
-      // if (mj_isSparse(m)) { ... }
+      if (issparse) {
+        // construct merged chain
+        int NV = mj_mergeChain(m, chain, wbody[0], wbody[1], /*flg_skipcommon=*/0);
 
-      // get endpoint JacobianDots, subtract
+        if (NV) {
-      mj_jacDot(m, d, jac1, 0, wpnt, wbody[0]);
+          // get endpoint JacobianDots, subtract
-      mj_jacDot(m, d, jac2, 0, wpnt+3, wbody[1]);
+          mj_jacDotSparse(m, d, jac1, 0, wpnt, wbody[0], NV, chain);
-      mju_sub(jacdif, jac2, jac1, 3*nv);
+          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
+          // chain rule, first term: Jdot += d/dt(jac2 - jac1) * dpnt
-      mju_mulMatTVec(tmp, jacdif, dpnt, 3, nv);
+          mju_mulMatTVec(tmp, jacdif, dpnt, 3, NV);
 
-      // add to existing
+          // scatter into dense output
-      mju_addToScl(Jdot, tmp, 1/divisor, nv);
+          for (int k=0; k < NV; k++) {
+            Jdot[chain[k]] += tmp[k] / divisor;
+          }
 
-      // get endpoint Jacobians, subtract
+          // get endpoint Jacobians, subtract
-      mj_jac(m, d, jac1, 0, wpnt, wbody[0]);
+          mj_jacSparse(m, d, jac1, 0, wpnt, wbody[0], NV, chain, /*flg_skipcommon=*/0);
-      mj_jac(m, d, jac2, 0, wpnt+3, wbody[1]);
+          mj_jacSparse(m, d, jac2, 0, wpnt+3, wbody[1], NV, chain, /*flg_skipcommon=*/0);
-      mju_sub(jacdif, jac2, jac1, 3*nv);
+          mju_sub(jacdif, jac2, jac1, 3*NV);
 
-      // chain rule, second term: Jdot += (jac2 - jac1) * d/dt(dpnt)
+          // chain rule, second term: Jdot += (jac2 - jac1) * d/dt(dpnt)
-      mju_mulMatTVec(tmp, jacdif, dvel, 3, nv);
+          mju_mulMatTVec(tmp, jacdif, dvel, 3, NV);
 
-      // add to existing
+          // scatter into dense output
-      mju_addToScl(Jdot, tmp, 1/divisor, nv);
+          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],

test/engine/engine_support_test.cc

@@ -17,6 +17,7 @@
 #include "src/engine/engine_core_util.h"
 #include "src/engine/engine_support.h"
 
+#include <algorithm>
 #include <cstring>
 #include <limits>
 #include <random>
@@ -483,6 +484,79 @@ TEST_F(JacobianTest, JacDot) {
   }
 }
 
+// compare mj_jacDotSparse with dense mj_jacDot
+TEST_F(JacobianTest, JacDotSparse) {
+  for (auto xml : {kHinge, kQuat, kTelescope, kFreeBall, kQuatlessPendulum}) {
+    char error[1024];
+    mjModel* model = LoadModelFromString(xml, error, sizeof(error));
+    ASSERT_THAT(model, NotNull()) << error;
+    int nv = model->nv;
+    mjData* data = mj_makeData(model);
+
+    // load keyframe if present, step for a bit
+    if (model->nkey) mj_resetDataKeyframe(model, data, 0);
+    while (data->time < 0.1) {
+      mj_step(model, data);
+    }
+
+    // minimal call required for mj_jacDot outputs to be valid
+    mj_kinematics(model, data);
+    mj_comPos(model, data);
+    mj_comVel(model, data);
+
+    // get bodyid and site position
+    int bodyid = mj_name2id(model, mjOBJ_BODY, "query");
+    EXPECT_GT(bodyid, 0);
+    int siteid = mj_name2id(model, mjOBJ_SITE, "query");
+    EXPECT_GT(siteid, -1);
+    mjtNum point[3];
+    mju_copy3(point, data->site_xpos+3*siteid);
+
+    // dense jacDot
+    vector<mjtNum> jacp_dense(3*nv);
+    vector<mjtNum> jacr_dense(3*nv);
+    mj_jacDot(model, data, jacp_dense.data(), jacr_dense.data(), point, bodyid);
+
+    // compute body chain using public mjModel fields
+    vector<int> chain(nv);
+    int NV = 0;
+    int weldbody = model->body_weldid[bodyid];
+    if (weldbody) {
+      int da = model->body_dofadr[weldbody] + model->body_dofnum[weldbody] - 1;
+      while (da >= 0) {
+        chain[NV++] = da;
+        da = model->dof_parentid[da];
+      }
+      std::reverse(chain.begin(), chain.begin() + NV);
+    }
+    EXPECT_GT(NV, 0);
+
+    // sparse jacDot
+    vector<mjtNum> jacp_sparse(3*NV);
+    vector<mjtNum> jacr_sparse(3*NV);
+    mj_jacDotSparse(model, data, jacp_sparse.data(), jacr_sparse.data(),
+                    point, bodyid, NV, chain.data());
+
+    // expand sparse to dense and compare
+    vector<mjtNum> jacp_expanded(3*nv, 0);
+    vector<mjtNum> jacr_expanded(3*nv, 0);
+    for (int ci = 0; ci < NV; ci++) {
+      int di = chain[ci];
+      for (int r = 0; r < 3; r++) {
+        jacp_expanded[di+r*nv] = jacp_sparse[ci+r*NV];
+        jacr_expanded[di+r*nv] = jacr_sparse[ci+r*NV];
+      }
+    }
+
+    // expect bitwise equality
+    EXPECT_EQ(jacp_expanded, jacp_dense);
+    EXPECT_EQ(jacr_expanded, jacr_dense);
+
+    mj_deleteData(data);
+    mj_deleteModel(model);
+  }
+}
+
 using Name2idTest = MujocoTest;
 
 static constexpr char name2idTestingModel[] = R"(