diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index c4aaad09..0ab2360d 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -468,7 +468,8 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { // compute Jacobian difference (opposite of contact: 0 - 1) NV = mj_jacDifPair(m, d, chain, body_id[1], body_id[0], pos[1], pos[0], - jac[1], jac[0], jacdif, NULL, NULL, NULL, issparse); + jac[1], jac[0], jacdif, NULL, NULL, NULL, issparse, + /*flg_skipcommon=*/0); // copy difference into jac[0] mju_copy(jac[0], jacdif, 3*NV); @@ -504,7 +505,8 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { // compute error Jacobian (opposite of contact: 0 - 1) NV = mj_jacDifPair(m, d, chain, body_id[1], body_id[0], pos[1], pos[0], jac[1], jac[0], jacdif, - jac[1]+3*nv, jac[0]+3*nv, jacdif+3*nv, issparse); + jac[1]+3*nv, jac[0]+3*nv, jacdif+3*nv, issparse, + /*flg_skipcommon=*/0); // copy difference into jac[0], compress translation:rotation if sparse mju_copy(jac[0], jacdif, 3*NV); @@ -670,7 +672,8 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { for (int n = 0; n < nodenum; n++) { int chain_nnz = mj_bodyChain(m, bodyid[n], chain_col); mju_zero(blk_jac, 3*nv); - mj_jacSparse(m, d, blk_jac, NULL, xpos + 3*n, bodyid[n], chain_nnz, chain_col); + mj_jacSparse(m, d, blk_jac, NULL, xpos + 3*n, bodyid[n], chain_nnz, chain_col, + /*flg_skipcommon=*/0); // expand sparse Jacobian to dense row format for (int r = 0; r < 3; r++) { @@ -1297,14 +1300,13 @@ int mj_contactJacobian(const mjModel* m, mjData* d, const mjContact* con, int di m->geom_bodyid[con->geom[side]] : m->flex_vertbodyid[m->flex_vertadr[con->flex[side]] + con->vert[side]]; } - - // compute Jacobian differences + // compute Jacobian differences, skipping common dofs if (dim > 3) { return mj_jacDifPair(m, d, chain, bid[0], bid[1], con->pos, con->pos, - jac1p, jac2p, jacdifp, jac1r, jac2r, jacdifr, mj_isSparse(m)); + jac1p, jac2p, jacdifp, jac1r, jac2r, jacdifr, mj_isSparse(m), 1); } else { return mj_jacDifPair(m, d, chain, bid[0], bid[1], con->pos, con->pos, - jac1p, jac2p, jacdifp, NULL, NULL, NULL, mj_isSparse(m)); + jac1p, jac2p, jacdifp, NULL, NULL, NULL, mj_isSparse(m), 1); } } @@ -2067,7 +2069,8 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) { } NV = mj_jacDifPair(m, NULL, chain, id[1], id[0], NULL, NULL, - NULL, NULL, NULL, NULL, NULL, NULL, issparse); + NULL, NULL, NULL, NULL, NULL, NULL, issparse, + /*flg_skipcommon=*/0); break; case mjEQ_WELD: @@ -2083,7 +2086,8 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) { } NV = mj_jacDifPair(m, NULL, chain, id[1], id[0], NULL, NULL, - NULL, NULL, NULL, NULL, NULL, NULL, issparse); + NULL, NULL, NULL, NULL, NULL, NULL, issparse, + /*flg_skipcommon=*/0); break; case mjEQ_JOINT: @@ -2138,7 +2142,8 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) { int b1 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e]]; int b2 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e+1]]; NV += mj_jacDifPair(m, NULL, chain, b1, b2, NULL, NULL, - NULL, NULL, NULL, NULL, NULL, NULL, issparse); + NULL, NULL, NULL, NULL, NULL, NULL, issparse, + /*flg_skipcommon=*/0); } } break; @@ -2252,55 +2257,72 @@ static int mj_nc(const mjModel* m, mjData* d, int* nnz) { // compute NV only if nnz requested int NV = 0; if (nnz) { - // get bodies - int nb = 0, bid[729]; - for (int side=0; side < 2; side++) { - // geom - if (con->geom[side] >= 0) { - bid[nb++] = m->geom_bodyid[con->geom[side]]; - } - - // flex - else { - int nw = 0; - int vid[4]; - mjtNum vweight[4]; - - // flex vert - if (con->vert[side] >= 0) { - vid[nw++] = m->flex_vertadr[con->flex[side]] + con->vert[side]; - vweight[0] = 1; - } - - // flex elem - else { - int f = con->flex[side]; - int fdim = m->flex_dim[f]; - const int* edata = m->flex_elem + m->flex_elemdataadr[f] + con->elem[side]*(fdim+1); - for (int k=0; k <= fdim; k++) { - vid[nw++] = m->flex_vertadr[f] + edata[k]; - } - - if (m->flex_interp[f]) { - nw = mj_elemBodyWeight(m, d, con->flex[side], con->elem[side], - con->vert[1-side], con->pos, vid, vweight); - } - } - - // get body or node ids and weights - if (m->flex_interp[con->flex[side]] == 0) { - for (int k=0; k < nw; k++) { - bid[nb] = m->flex_vertbodyid[vid[k]]; - nb++; - } - } else { - nb += mj_vertBodyWeight(m, d, con->flex[side], vid, bid+nb, NULL, vweight, nw); - } + // single body on each side (geom-geom or flex vert-vert): skip common dofs + if ((con->geom[0] >= 0 || (con->vert[0] >= 0 && m->flex_interp[con->flex[0]] == 0)) && + (con->geom[1] >= 0 || (con->vert[1] >= 0 && m->flex_interp[con->flex[1]] == 0))) { + // get bodies + int bid[2]; + for (int side=0; side < 2; side++) { + bid[side] = (con->geom[side] >= 0) ? + m->geom_bodyid[con->geom[side]] : + m->flex_vertbodyid[m->flex_vertadr[con->flex[side]] + con->vert[side]]; } + NV = mj_jacDifPair(m, NULL, chain, bid[0], bid[1], NULL, NULL, + NULL, NULL, NULL, NULL, NULL, NULL, mj_isSparse(m), 1); } - // count non-zeros in merged chain - NV = mj_jacSumCount(m, d, chain, nb, bid); + // general case: flex elements involved + else { + // get bodies + int nb = 0, bid[729]; + for (int side=0; side < 2; side++) { + // geom + if (con->geom[side] >= 0) { + bid[nb++] = m->geom_bodyid[con->geom[side]]; + } + + // flex + else { + int nw = 0; + int vid[4]; + mjtNum vweight[4]; + + // flex vert + if (con->vert[side] >= 0) { + vid[nw++] = m->flex_vertadr[con->flex[side]] + con->vert[side]; + vweight[0] = 1; + } + + // flex elem + else { + int f = con->flex[side]; + int fdim = m->flex_dim[f]; + const int* edata = m->flex_elem + m->flex_elemdataadr[f] + con->elem[side]*(fdim+1); + for (int k=0; k <= fdim; k++) { + vid[nw++] = m->flex_vertadr[f] + edata[k]; + } + + if (m->flex_interp[f]) { + nw = mj_elemBodyWeight(m, d, con->flex[side], con->elem[side], + con->vert[1-side], con->pos, vid, vweight); + } + } + + // get body or node ids and weights + if (m->flex_interp[con->flex[side]] == 0) { + for (int k=0; k < nw; k++) { + bid[nb] = m->flex_vertbodyid[vid[k]]; + nb++; + } + } else { + nb += mj_vertBodyWeight(m, d, con->flex[side], vid, bid+nb, NULL, vweight, nw); + } + } + } + + // count non-zeros in merged chain + NV = mj_jacSumCount(m, d, chain, nb, bid); + } if (!NV) { continue; } diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 97c26f60..557bbd4f 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -709,7 +709,8 @@ void mj_flex(const mjModel* m, mjData* d) { // get endpoint Jacobians, subtract int NV = mj_jacDifPair(m, d, chain, b1, b2, pos1, pos2, - jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1); + jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1, + /*skipcommon=*/0); // no dofs: skip if (!NV) { @@ -853,9 +854,11 @@ void mj_flex(const mjModel* m, mjData* d) { int b1 = m->flex_vertbodyid[vbase+v1]; int b2 = m->flex_vertbodyid[vbase+v2]; int NV1 = mj_bodyChain(m, b1, chain1); - mj_jacSparse(m, d, jac1, NULL, d->flexvert_xpos + 3*(vbase+v1), b1, NV1, chain1); + mj_jacSparse(m, d, jac1, NULL, d->flexvert_xpos + 3*(vbase+v1), b1, NV1, chain1, + /*flg_skipcommon=*/0); int NV2 = mj_bodyChain(m, b2, chain2); - mj_jacSparse(m, d, jac2, NULL, d->flexvert_xpos + 3*(vbase+v2), b2, NV2, chain2); + mj_jacSparse(m, d, jac2, NULL, d->flexvert_xpos + 3*(vbase+v2), b2, NV2, chain2, + /*flg_skipcommon=*/0); // accumulate dense Jacobians for vertex v mju_mulMatTVec(J_local, jac1, dI1dy1, 3, NV1); @@ -1042,7 +1045,8 @@ void mj_tendon(const mjModel* m, mjData* d) { // get endpoint Jacobians, subtract int NV = mj_jacDifPair(m, d, chain, wbody[k], wbody[k+1], wpnt+3*k, wpnt+3*k+3, - jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1); + jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1, + /*skipcommon=*/0); // no dofs: skip if (!NV) { @@ -1638,7 +1642,8 @@ void mj_transmission(const mjModel* m, mjData* d) { // get Jacobian difference int NV = mj_jacDifPair(m, d, chain, b1, b2, con->pos, con->pos, - jac1p, jac2p, jacdifp, NULL, NULL, NULL, issparse); + jac1p, jac2p, jacdifp, NULL, NULL, NULL, issparse, + /*flg_skipcommon=*/0); // project Jacobian along the normal of the contact frame mju_mulMatMat(jac, con->frame, jacdifp, 1, 3, NV); diff --git a/src/engine/engine_core_util.c b/src/engine/engine_core_util.c index d573f525..6a2b9150 100644 --- a/src/engine/engine_core_util.c +++ b/src/engine/engine_core_util.c @@ -303,7 +303,7 @@ void mj_jacPointAxis(const mjModel* m, mjData* d, mjtNum* jacPoint, mjtNum* jacA // compute 3/6-by-nv sparse Jacobian of global point attached to given body void mj_jacSparse(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body, - int NV, const int* chain) { + int NV, const int* chain, int flg_skipcommon) { // clear jacobians if (jacp) { mju_zero(jacp, 3*NV); @@ -337,8 +337,12 @@ void mj_jacSparse(const mjModel* m, const mjData* d, ci--; } - // make sure we found it; SHOULD NOT OCCUR + // dof not in chain: skip if shared dofs are excluded, otherwise SHOULD NOT OCCUR if (ci < 0 || chain[ci] != da) { + if (flg_skipcommon) { + da = m->dof_parentid[da]; + continue; + } mjERROR("dof index %d not found in chain", da); } @@ -433,7 +437,8 @@ void mj_jacSparseSimple(const mjModel* m, const mjData* d, int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain, int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3], mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp, - mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr, int issparse) { + mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr, + int issparse, int flg_skipcommon) { int issimple = (m->body_simple[b1] && m->body_simple[b2]); int NV = m->nv; @@ -447,7 +452,7 @@ int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain, if (issimple) { NV = mj_mergeChainSimple(m, chain, b1, b2); } else { - NV = mj_mergeChain(m, chain, b1, b2, 0); + NV = mj_mergeChain(m, chain, b1, b2, flg_skipcommon); } } @@ -477,8 +482,8 @@ int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain, // regular processing else { // Jacobians - mj_jacSparse(m, d, jac1p, jac1r, pos1, b1, NV, chain); - mj_jacSparse(m, d, jac2p, jac2r, pos2, b2, NV, chain); + mj_jacSparse(m, d, jac1p, jac1r, pos1, b1, NV, chain, flg_skipcommon); + mj_jacSparse(m, d, jac2p, jac2r, pos2, b2, NV, chain, flg_skipcommon); // differences if (jacdifp) { @@ -535,7 +540,7 @@ int mj_jacSum(const mjModel* m, mjData* d, int* chain, if (m->body_simple[body[0]]) { mj_jacSparseSimple(m, d, jacp, jacr, point, body[0], 1, NV, 0); } else { - mj_jacSparse(m, d, jacp, jacr, point, body[0], NV, chain); + mj_jacSparse(m, d, jacp, jacr, point, body[0], NV, chain, /*flg_skipcommon=*/0); } // apply weight @@ -552,7 +557,7 @@ int mj_jacSum(const mjModel* m, mjData* d, int* chain, if (m->body_simple[body[i]]) { mj_jacSparseSimple(m, d, jp, jr, point, body[i], 1, bodyNV, 0); } else { - mj_jacSparse(m, d, jp, jr, point, body[i], bodyNV, bodychain); + mj_jacSparse(m, d, jp, jr, point, body[i], bodyNV, bodychain, /*flg_skipcommon=*/0); } // combine sparse matrices diff --git a/src/engine/engine_core_util.h b/src/engine/engine_core_util.h index 7f2cb097..2e12e935 100644 --- a/src/engine/engine_core_util.h +++ b/src/engine/engine_core_util.h @@ -78,7 +78,7 @@ MJAPI void mj_jacPointAxis(const mjModel* m, mjData* d, // compute 3/6-by-nv sparse Jacobian of global point attached to given body void mj_jacSparse(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body, - int NV, const int* chain); + int NV, const int* chain, int flg_skipcommon); // sparse Jacobian difference for simple body contacts void mj_jacSparseSimple(const mjModel* m, const mjData* d, @@ -89,7 +89,8 @@ void mj_jacSparseSimple(const mjModel* m, const mjData* d, MJAPI int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain, int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3], mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp, - mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr, int issparse); + mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr, + int issparse, int flg_skipcommon); // dense or sparse weighted sum of multiple body Jacobians at same point int mj_jacSum(const mjModel* m, mjData* d, int* chain, diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index 09dd4903..dea6ba45 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -980,7 +980,8 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op, int chain_nnz = mj_bodyChain(m, bodyid[i], chain_colind); // compute sparse Jacobian for this node (3 rows) - mj_jacSparse(m, d, blk_jac, NULL, xpos+3*i, bodyid[i], chain_nnz, chain_colind); + mj_jacSparse(m, d, blk_jac, NULL, xpos+3*i, bodyid[i], chain_nnz, chain_colind, + /*flg_skipcommon=*/0); // copy to sparse structure for (int r=0; r<3; r++) { @@ -1488,7 +1489,8 @@ void mjd_ellipsoidFluid(const mjModel* m, mjData* d, int bodyid) { // get geom global Jacobian: rotation then translation if (mj_isSparse(m)) { - mj_jacSparse(m, d, J+3*nnz, J, d->geom_xpos+3*geomid, m->geom_bodyid[geomid], nnz, colind); + mj_jacSparse(m, d, J+3*nnz, J, d->geom_xpos+3*geomid, m->geom_bodyid[geomid], nnz, colind, + /*flg_skipcommon=*/0); } else { mj_jacGeom(m, d, J+3*nv, J, geomid); } @@ -1574,7 +1576,7 @@ void mjd_inertiaBoxFluid(const mjModel* m, mjData* d, int i) { nnz = mj_bodyChain(m, i, colind); // get sparse jacBodyCom - mj_jacSparse(m, d, J+3*nnz, J, d->xipos+3*i, i, nnz, colind); + mj_jacSparse(m, d, J+3*nnz, J, d->xipos+3*i, i, nnz, colind, /*flg_skipcommon=*/0); // prepare rownnz, rowadr, colind for all 6 rows rownnz[0] = nnz; diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c index 6473601f..5a256ea9 100644 --- a/src/engine/engine_setconst.c +++ b/src/engine/engine_setconst.c @@ -434,7 +434,7 @@ static void makeFlexSparse(mjModel* m, mjData* d) { // get sparsity int NV = mj_jacDifPair(m, d, chain, b1, b2, dummy_pos, dummy_pos, NULL, - NULL, NULL, NULL, NULL, NULL, /*issparse=*/1); + NULL, NULL, NULL, NULL, NULL, /*issparse=*/1, /*skipcommon=*/0); // copy sparsity info rownnz[ebase + e] = NV; diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 0e2dcd35..a4827ddb 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -464,7 +464,7 @@ void mj_applyFT(const mjModel* m, mjData* d, // construct chain and sparse Jacobians int* chain = mjSTACKALLOC(d, nv, int); int NV = mj_bodyChain(m, body, chain); - mj_jacSparse(m, d, jacp, jacr, point, body, NV, chain); + mj_jacSparse(m, d, jacp, jacr, point, body, NV, chain, /*flg_skipcommon=*/0); // compute J'*f and accumulate if (force) { diff --git a/test/engine/engine_core_constraint_test.cc b/test/engine/engine_core_constraint_test.cc index 2d7b46fb..7bb7f0a1 100644 --- a/test/engine/engine_core_constraint_test.cc +++ b/test/engine/engine_core_constraint_test.cc @@ -134,7 +134,8 @@ TEST_F(CoreConstraintTest, WeldRotJacobian) { // rotational Jacobian difference mj_jacDifPair(model, data, NULL, 2, 1, point, point, - NULL, NULL, NULL, jac0, jac1, jacdif, mj_isSparse(model)); + NULL, NULL, NULL, jac0, jac1, jacdif, mj_isSparse(model), + /*flg_skipcommon=*/0); // formula: 0.5 * neg(quat2) * (jac1-jac2) * quat1 mjtNum axis[3], quat3[4], quat4[4]; @@ -716,8 +717,8 @@ TEST_F(CoreConstraintTest, StrainConstraintNoPinning) { // Rotate by 45 degrees around Z axis via quaternion mjtNum angle = 0.785398; // 45 degrees d->qpos[3] = mju_cos(angle/2); // w - d->qpos[4] = 0; // x - d->qpos[5] = 0; // y + d->qpos[4] = 0; // x + d->qpos[5] = 0; // y d->qpos[6] = mju_sin(angle/2); // z mj_forward(m, d); @@ -743,5 +744,44 @@ TEST_F(CoreConstraintTest, StrainConstraintNoPinning) { mj_deleteModel(m); } +TEST_F(CoreConstraintTest, ContactSharedDofJacobian) { + constexpr char xml[] = R"( + + + )"; + char error[1024]; + mjModel* model = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(model, NotNull()) << error; + ASSERT_EQ(model->nv, 3); + ASSERT_TRUE(mj_isSparse(model)); + mjData* data = mj_makeData(model); + + mj_forward(model, data); + + ASSERT_EQ(data->ncon, 1); + ASSERT_GE(data->nefc, 1); + + EXPECT_EQ(data->efc_J_rownnz[0], 2); + + mj_deleteData(data); + mj_deleteModel(model); +} + + } // namespace } // namespace mujoco