chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Add mj_constraintUpdate_island.

Browse Source 
PiperOrigin-RevId: 562581620
Change-Id: If156a02873168127e2c5f2f377532fd45f7eec1c
This commit is contained in:
Yuval Tassa
2023-09-04 10:32:45 -07:00
committed by Copybara-Service
parent b7686440d1
commit 600c12533d
3 changed files with 155 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+43 -22
View File
@@ -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);
}
src/engine/engine_core_constraint.h
+4
View File
@@ -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
test/engine/engine_core_constraint_test.cc
+108
View File
@@ -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