diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index ef4b3f7f..a0b4e3aa 100644 --- a/src/engine/engine_core_smooth.c +++ b/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 diff --git a/src/engine/engine_core_util.c b/src/engine/engine_core_util.c index 642dc896..c49f3ee5 100644 --- a/src/engine/engine_core_util.c +++ b/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; diff --git a/src/engine/engine_core_util.h b/src/engine/engine_core_util.h index 98093d6f..39633ff0 100644 --- a/src/engine/engine_core_util.h +++ b/src/engine/engine_core_util.h @@ -18,6 +18,7 @@ #include #include #include +#include #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], diff --git a/test/engine/engine_support_test.cc b/test/engine/engine_support_test.cc index 8fb8502e..629517ae 100644 --- a/test/engine/engine_support_test.cc +++ b/test/engine/engine_support_test.cc @@ -17,6 +17,7 @@ #include "src/engine/engine_core_util.h" #include "src/engine/engine_support.h" +#include #include #include #include @@ -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 jacp_dense(3*nv); + vector jacr_dense(3*nv); + mj_jacDot(model, data, jacp_dense.data(), jacr_dense.data(), point, bodyid); + + // compute body chain using public mjModel fields + vector 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 jacp_sparse(3*NV); + vector 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 jacp_expanded(3*nv, 0); + vector 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"(