diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index c95a3d65..98857258 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -2003,18 +2003,25 @@ void mj_referenceConstraint(const mjModel* m, mjData* d) { //---------------------------- update constraint state --------------------------------------------- -// compute efc_state, efc_force, qfrc_constraint -// optional: cost(qacc) = shat(jar) where jar = Jac*qacc-aref; cone Hessians -void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, - mjtNum cost[1], int flg_coneHessian) { - int ne = d->ne, nf = d->nf, nefc = d->nefc, nv = m->nv; +// compute efc_state, efc_force, qfrc_constraint, optionally restricted to one island +// island < 0: update all d->nefc constraints +// island >= 0: update only d->island_efcnum[island] constraints +// jar = Jac*qacc-aref is restricted to the island, in the above sense +// optional: cost(qacc) = shat(jar); cone Hessians +void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, + mjtNum cost[1], int flg_coneHessian, int island) { + int ne = d->ne, nf = d->nf; const mjtNum *D = d->efc_D, *R = d->efc_R, *floss = d->efc_frictionloss; mjtNum* force = d->efc_force; mjtNum s = 0; + int nefc = island < 0 ? d->nefc : d->island_efcnum[island]; + int* efcind = island < 0 ? NULL : d->island_efcind + d->island_efcadr[island]; + // no constraints: clear qfrc_constraint and cost, return if (!nefc) { - mju_zero(d->qfrc_constraint, nv); + // can only occur for island == -1 + mju_zero(d->qfrc_constraint, m->nv); if (cost) { *cost = 0; } @@ -2022,16 +2029,19 @@ void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, } // compute unconstrained efc_force - for (int i=0; i < nefc; i++) { - force[i] = -D[i]*jar[i]; + for (int c=0; c < nefc; c++) { + int i = efcind ? efcind[c] : c; + force[i] = -D[i]*jar[c]; } // update constraints - for (int i=0; i < nefc; i++) { + for (int c=0; c < nefc; c++) { + int i = efcind ? efcind[c] : c; + // ==== equality if (i < ne) { if (cost) { - s += 0.5*D[i]*jar[i]*jar[i]; + s += 0.5*D[i]*jar[c]*jar[c]; } d->efc_state[i] = mjCNSTRSTATE_QUADRATIC; continue; @@ -2040,9 +2050,9 @@ void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, // ==== friction if (i < ne + nf) { // linear negative - if (jar[i] <= -R[i]*floss[i]) { + if (jar[c] <= -R[i]*floss[i]) { if (cost) { - s += -0.5*R[i]*floss[i]*floss[i] - floss[i]*jar[i]; + s += -0.5*R[i]*floss[i]*floss[i] - floss[i]*jar[c]; } force[i] = floss[i]; @@ -2051,9 +2061,9 @@ void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, } // linear positive - else if (jar[i] >= R[i]*floss[i]) { + else if (jar[c] >= R[i]*floss[i]) { if (cost) { - s += -0.5*R[i]*floss[i]*floss[i] + floss[i]*jar[i]; + s += -0.5*R[i]*floss[i]*floss[i] + floss[i]*jar[c]; } force[i] = -floss[i]; @@ -2064,7 +2074,7 @@ void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, // quadratic else { if (cost) { - s += 0.5*D[i]*jar[i]*jar[i]; + s += 0.5*D[i]*jar[c]*jar[c]; } d->efc_state[i] = mjCNSTRSTATE_QUADRATIC; @@ -2077,7 +2087,7 @@ void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, // non-negative constraint if (d->efc_type[i] != mjCNSTR_CONTACT_ELLIPTIC) { // constraint is satisfied: no cost - if (jar[i] >= 0) { + if (jar[c] >= 0) { force[i] = 0; d->efc_state[i] = mjCNSTRSTATE_SATISFIED; @@ -2086,7 +2096,7 @@ void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, // quadratic else { if (cost) { - s += 0.5*D[i]*jar[i]*jar[i]; + s += 0.5*D[i]*jar[c]*jar[c]; } d->efc_state[i] = mjCNSTRSTATE_QUADRATIC; @@ -2102,9 +2112,9 @@ void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, // map to regular dual cone space mjtNum U[6]; - U[0] = jar[i]*mu; + U[0] = jar[c]*mu; for (int j=1; j < dim; j++) { - U[j] = jar[i+j]*friction[j-1]; + U[j] = jar[c+j]*friction[j-1]; } // decompose into normal and tangent @@ -2122,7 +2132,7 @@ void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, else if (mu*N+T <= 0 || (T <= 0 && N < 0)) { if (cost) { for (int j=0; j < dim; j++) { - s += 0.5*D[i+j]*jar[i+j]*jar[i+j]; + s += 0.5*D[i+j]*jar[c+j]*jar[c+j]; } } @@ -2196,15 +2206,26 @@ void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, } // advance to end of contact - i += (dim-1); + c += (dim-1); } } // compute qfrc_constraint - mj_mulJacTVec(m, d, d->qfrc_constraint, d->efc_force); + int flg_vecunc = 1; + int flg_resunc = 1; + mj_mulJacTVec_island(m, d, d->qfrc_constraint, d->efc_force, island, flg_vecunc, flg_resunc); // assign cost if (cost) { *cost = s; } } + + + +// compute efc_state, efc_force, qfrc_constraint +// optional: cost(qacc) = shat(jar) where jar = Jac*qacc-aref; cone Hessians +void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, + mjtNum cost[1], int flg_coneHessian) { + mj_constraintUpdate_island(m, d, jar, cost, flg_coneHessian, -1); +} diff --git a/src/engine/engine_core_constraint.h b/src/engine/engine_core_constraint.h index 5c6b353a..ce04b120 100644 --- a/src/engine/engine_core_constraint.h +++ b/src/engine/engine_core_constraint.h @@ -116,6 +116,10 @@ MJAPI void mj_referenceConstraint(const mjModel* m, mjData* d); MJAPI void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar, mjtNum cost[1], int flg_coneHessian); +// compute efc_state, efc_force, qfrc_constraint for one island +MJAPI void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar, + mjtNum cost[1], int flg_coneHessian, int island); + #ifdef __cplusplus } #endif diff --git a/test/engine/engine_core_constraint_test.cc b/test/engine/engine_core_constraint_test.cc index 32c5fb0b..6f62a739 100644 --- a/test/engine/engine_core_constraint_test.cc +++ b/test/engine/engine_core_constraint_test.cc @@ -390,5 +390,113 @@ TEST_F(CoreConstraintTest, MulJacTVecIsland) { mj_deleteModel(model); } +TEST_F(CoreConstraintTest, ConstraintUpdateIsland) { + const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath); + mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + mjData* data1 = mj_makeData(model); + mjData* data2 = mj_makeData(model); + + // iterate over sparsity and cone + for (mjtJacobian sparsity : {mjJAC_SPARSE, mjJAC_DENSE}) { + for (mjtCone cone : {mjCONE_PYRAMIDAL, mjCONE_ELLIPTIC}) { + model->opt.jacobian = sparsity; + model->opt.cone = cone; + + // simulate for 0.2 seconds + mj_resetData(model, data1); + mj_resetData(model, data2); + while (data1->time < 0.2) { + mj_step(model, data1); + mj_step(model, data2); + } + mj_forward(model, data1); + mj_forward(model, data2); + + // get sizes + int nefc = data1->nefc; + int nv = model->nv; + int nisland = data1->nisland; + EXPECT_GT(nisland, 0); + + // get jar = J*a - aref + mjtNum* jar = (mjtNum*)mju_malloc(sizeof(mjtNum) * nefc); + mj_mulJacVec(model, data1, jar, data1->qacc); + mju_subFrom(jar, data1->efc_aref, nefc); + + // constraint update for data1 given jar + mjtNum cost1; + mj_constraintUpdate(model, data1, jar, &cost1, /*flg_coneHessian=*/1); + + // iterate over islands, check match + mjtNum cost2 = 0; + for (int island=0; island < nisland; island++) { + // clear outputs from data2 + for (int i=0; i < nefc; i++) data2->efc_state[i] = -1; + mju_zero(data2->efc_force, nefc); + mju_zero(data2->qfrc_constraint, nv); + for (int i=0; i < data2->ncon; i++) mju_zero(data2->contact[i].H, 36); + + // sizes and indices, in this island + int dofnum = data2->island_dofnum[island]; + int efcnum = data2->island_efcnum[island]; + int* dofind = data2->island_dofind + data2->island_dofadr[island]; + int* efcind = data2->island_efcind + data2->island_efcadr[island]; + + // get jar restricted to island + mjtNum* jari = (mjtNum*)mju_malloc(sizeof(mjtNum) * efcnum); + for (int c=0; c < efcnum; c++) { + jari[c] = jar[efcind[c]]; + } + + // update constraints for this island + mjtNum cost2i; + mj_constraintUpdate_island(model, data2, jari, &cost2i, + /*flg_coneHessian=*/1, island); + + // compare nefc vectors + for (int c=0; c < efcnum; c++) { + int i = efcind[c]; + EXPECT_EQ(data2->efc_island[i], island); + EXPECT_EQ(data2->efc_state[i], data1->efc_state[i]); + EXPECT_THAT(data2->efc_force[i], + DoubleNear(data1->efc_force[i], 1e-12)); + } + + // compare qfrc_constraint + for (int c=0; c < dofnum; c++) { + int i = dofind[c]; + EXPECT_THAT(data2->qfrc_constraint[i], + DoubleNear(data1->qfrc_constraint[i], 1e-12)); + } + + // compare cone Hessians + for (int c=0; c < data2->ncon; c++) { + int efcadr = data2->contact[c].efc_address; + if (data2->efc_island[efcadr] == island) { + for (int j=0; j < 36; j++) { + EXPECT_THAT(data2->contact[c].H[j], + DoubleNear(data2->contact[c].H[j], 1e-12)); + } + } + } + + // add island cost to total cost + cost2 += cost2i; + + mju_free(jari); + } + + // expect monolithic total cost + EXPECT_THAT(cost1, DoubleNear(cost2, 1e-12)); + + mju_free(jar); + } + } + + mj_deleteData(data2); + mj_deleteData(data1); + mj_deleteModel(model); +} + } // namespace } // namespace mujoco