Add mj_jacSubtreeCom for computing subtree center-of-mass Jacobians.

PiperOrigin-RevId: 464075727
Change-Id: If679c0fdd05b7dc84154cec00963aeea7bbdacc3

src/engine/engine_support.c

@@ -154,6 +154,35 @@ void mj_jacBodyCom(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr
 
 
 
+// compute subtree-com Jacobian
+void mj_jacSubtreeCom(const mjModel* m, mjData* d, mjtNum* jacp, int body) {
+  int nv = m->nv;
+  mjMARKSTACK;
+  mjtNum* jacp_b = mj_stackAlloc(d, 3*nv);
+
+  // clear output
+  mju_zero(jacp, 3*nv);
+
+  // forward pass starting from body
+  for (int b=body; b<m->nbody; b++) {
+    // end of body subtree, break from the loop
+    if (b > body && m->body_parentid[b] < body) {
+      break;
+    }
+
+    // b is in the body subtree, add mass-weighted Jacobian into jacp
+    mj_jac(m, d, jacp_b, NULL, d->xipos+3*b, b);
+    mju_addToScl(jacp, jacp_b, m->body_mass[b], 3*nv);
+  }
+
+  // normalize by subtree mass
+  mju_scl(jacp, jacp, 1/m->body_subtreemass[body], 3*nv);
+
+  mjFREESTACK;
+}
+
+
+
 // compute geom Jacobian
 void mj_jacGeom(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int geom) {
   mj_jac(m, d, jacp, jacr, d->geom_xpos + 3*geom, m->geom_bodyid[geom]);