Fix bugs in engine_island.c

PiperOrigin-RevId: 562519702
Change-Id: I3075e539ccc1cef92485152721a9344f180aa4b1
This commit is contained in:
Yuval Tassa
2023-09-04 04:05:46 -07:00
committed by Copybara-Service
parent 89f47c33e4
commit feb8aa2fbb
3 changed files with 145 additions and 84 deletions
+113 -58
View File
@@ -160,13 +160,12 @@ static int countMaxEdge(const mjModel* m, const mjData* d) {
// return id of next tree in Jacobian row i that is different from tree, -1 if not found
// write the index of the found tree to *index if given
// start search from *index if given, otherwise 0
// if J is (dense/sparse) *index is the (column/nonzro) index, respectively
// start search from *index
// write the index of the found tree to *index
// if J is (dense/sparse) *index is the (column/nonzero) index, respectively
static int treeNext(const mjModel* m, const mjData* d, int tree, int i, int *index) {
int tree_next = -1;
int j0 = index ? *index : 0; // start searching at *index if given, otherwise 0
int j; // loop variable, saved to *index
int j; // local loop variable, saved to *index
// sparse
if (mj_isSparse(m)) {
@@ -174,7 +173,7 @@ static int treeNext(const mjModel* m, const mjData* d, int tree, int i, int *ind
int* colind = d->efc_J_colind + d->efc_J_rowadr[i];
// loop over remaining nonzeros, look for different tree
for (j=j0; j < rownnz; j++) {
for (j=(*index); j < rownnz; j++) {
int tree_j = m->dof_treeid[colind[j]];
if (tree_j != tree) {
// found different tree
@@ -189,7 +188,7 @@ static int treeNext(const mjModel* m, const mjData* d, int tree, int i, int *ind
int nv = m->nv;
// scan row, look for different tree
for (j=j0; j < nv; j++) {
for (j=(*index); j < nv; j++) {
if (d->efc_J[nv*i + j]) {
int tree_j = m->dof_treeid[j];
if (tree_j != tree) {
@@ -202,13 +201,94 @@ static int treeNext(const mjModel* m, const mjData* d, int tree, int i, int *ind
}
// save last index
if (index) *index = j;
*index = j;
return tree_next;
}
// find first and possibly second nonegative tree ids in Jacobian row i
// if row i is special-cased (no more trees), return -1
// otherwise call treeNext, starting scan at index 0, return index
static int treeFirst(const mjModel* m, const mjData* d, int tree[2], int i) {
int efc_type = d->efc_type[i];
int efc_id = d->efc_id[i];
// clear outputs
tree[0] = -1;
tree[1] = -1;
// ==== fast handling of special cases
// joint friction
if (efc_type == mjCNSTR_FRICTION_DOF) {
tree[0] = m->dof_treeid[efc_id];
return -1;
}
// joint limit
if (efc_type == mjCNSTR_LIMIT_JOINT) {
tree[0] = m->dof_treeid[m->jnt_dofadr[efc_id]];
return -1;
}
// contact
if (efc_type == mjCNSTR_CONTACT_FRICTIONLESS ||
efc_type == mjCNSTR_CONTACT_PYRAMIDAL ||
efc_type == mjCNSTR_CONTACT_ELLIPTIC) {
tree[0] = m->body_treeid[m->geom_bodyid[d->contact[efc_id].geom1]];
tree[1] = m->body_treeid[m->geom_bodyid[d->contact[efc_id].geom2]];
// handle static bodies
if (tree[0] < 0) {
if (tree[1] < 0) {
mjERROR("contact %d is between two static bodies", efc_id); // SHOULD NOT OCCUR
} else {
int tmp = tree[0];
tree[0] = tree[1];
tree[1] = tmp;
}
}
return -1;
}
// connect or weld constraints
if (efc_type == mjCNSTR_EQUALITY) {
mjtEq eq_type = m->eq_type[efc_id];
if (eq_type == mjEQ_CONNECT || eq_type == mjEQ_WELD) {
tree[0] = m->body_treeid[m->eq_obj1id[efc_id]];
tree[1] = m->body_treeid[m->eq_obj2id[efc_id]];
// handle static bodies
if (tree[0] < 0) {
if (tree[1] < 0) {
mjERROR("equality %d is between two static bodies", efc_id); // SHOULD NOT OCCUR
} else {
int tmp = tree[0];
tree[0] = tree[1];
tree[1] = tmp;
}
}
return -1;
}
}
// ==== generic case: scan Jacobian
int index = 0;
tree[0] = treeNext(m, d, -1, i, &index);
if (tree[0] < 0) {
mjERROR("no tree found for constraint %d", i); // SHOULD NOT OCCUR
}
return index;
}
// add 0 edges, 1 self-edge or 2 flipped edges to array, increment treenedge
// return current number of edges
static int addEdge(int* treenedge, int* edge, int nedge, int tree1, int tree2, int nedge_max) {
@@ -273,7 +353,6 @@ static int findEdges(const mjModel* m, const mjData* d, int* treenedge, int* edg
int nefc = d->nefc;
int efc_type = -1;
int efc_id = -1;
int tree1, tree2;
// clear treenedge
memset(treenedge, 0, m->ntree*sizeof(int));
@@ -287,60 +366,34 @@ static int findEdges(const mjModel* m, const mjData* d, int* treenedge, int* edg
efc_type = d->efc_type[i];
efc_id = d->efc_id[i];
// ==== fast handling of special cases
int tree[2];
int index = treeFirst(m, d, tree, i);
int tree1 = tree[0];
int tree2 = tree[1];
// joint friction
if (efc_type == mjCNSTR_FRICTION_DOF) {
tree1 = m->dof_treeid[efc_id];
nedge = addEdge(treenedge, edge, nedge, tree1, tree1, nedge_max);
// no more edges to find, add and continue
if (index == -1) {
nedge = addEdge(treenedge, edge, nedge, tree1, tree2 == -1 ? tree1 : tree2, nedge_max);
continue;
}
// joint limit
if (efc_type == mjCNSTR_LIMIT_JOINT) {
tree1 = m->dof_treeid[m->jnt_dofadr[efc_id]];
nedge = addEdge(treenedge, edge, nedge, tree1, tree1, nedge_max);
continue;
}
// possibly more edges, scan Jacobian row
else {
tree2 = treeNext(m, d, tree1, i, &index);
// contact
if (efc_type == mjCNSTR_CONTACT_FRICTIONLESS ||
efc_type == mjCNSTR_CONTACT_PYRAMIDAL ||
efc_type == mjCNSTR_CONTACT_ELLIPTIC) {
tree1 = m->body_treeid[m->geom_bodyid[d->contact[efc_id].geom1]];
tree2 = m->body_treeid[m->geom_bodyid[d->contact[efc_id].geom2]];
nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max);
continue;
}
// connect or weld constraints
if (efc_type == mjCNSTR_EQUALITY) {
mjtEq eq_type = m->eq_type[efc_id];
if (eq_type == mjEQ_CONNECT || eq_type == mjEQ_WELD) {
tree1 = m->body_treeid[m->eq_obj1id[efc_id]];
tree2 = m->body_treeid[m->eq_obj2id[efc_id]];
if (tree2 == -1) {
// 1 tree found: add self-edge
nedge = addEdge(treenedge, edge, nedge, tree1, tree1, nedge_max);
} else {
// 2 trees found: add edge, keep scanning and adding until no more trees
nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max);
continue;
}
}
// ==== generic case: scan Jacobian
int index = 0;
tree1 = treeNext(m, d, -1, i, &index);
tree2 = treeNext(m, d, tree1, i, &index);
if (tree2 == -1) {
// 1 tree found: add self-edge
nedge = addEdge(treenedge, edge, nedge, tree1, tree1, nedge_max);
} else {
// 2 trees found: add edge, keep scanning and adding until no more trees
nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max);
int tree3 = treeNext(m, d, tree2, i, &index);
while (tree3 > -1 && tree3 != tree2) {
tree1 = tree2;
tree2 = tree3;
nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max);
tree3 = treeNext(m, d, tree2, i, &index);
int tree3 = treeNext(m, d, tree2, i, &index);
while (tree3 > -1 && tree3 != tree2) {
tree1 = tree2;
tree2 = tree3;
nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max);
tree3 = treeNext(m, d, tree2, i, &index);
}
}
}
}
@@ -452,7 +505,9 @@ void mj_island(const mjModel* m, mjData* d) {
// compute efc_island, island_efcnum
memset(d->island_efcnum, 0, nisland*sizeof(int));
for (int i=0; i < nefc; i++) {
int island = tree_island[treeNext(m, d, -1, i, NULL)];
int tree[2];
treeFirst(m, d, tree, i);
int island = tree_island[tree[0]];
d->efc_island[i] = island;
d->island_efcnum[island]++;
}
+24 -24
View File
@@ -350,33 +350,33 @@ TEST_F(IslandTest, IslandEfc) {
mj_step(model, data);
}
// sizes
int nv = model->nv;
int nefc = data->nefc;
int nisland = data->nisland;
// expect island structure to correspond to comment at top of xml
EXPECT_EQ(nisland, 3);
EXPECT_EQ(data->ne, 1);
EXPECT_EQ(data->nisland, 4);
EXPECT_EQ(data->ne, 4);
EXPECT_EQ(data->nf, 2);
EXPECT_EQ(data->nl, 1);
EXPECT_EQ(nefc, 24);
EXPECT_THAT(AsVector(data->dof_island, nv),
ElementsAre(0, -1, 0, 0, 0, 1, 2, 0, 1, 1, 1, 1, 1, 1));
EXPECT_THAT(AsVector(data->island_dofnum, nisland), ElementsAre(5, 7, 1));
EXPECT_THAT(AsVector(data->island_dofadr, nisland), ElementsAre(0, 5, 12));
EXPECT_THAT(AsVector(data->island_dofind, nv),
ElementsAre(0, 2, 3, 4, 7, 5, 8, 9, 10, 11, 12, 13, 6, -1));
EXPECT_THAT(AsVector(data->dof_islandind, nv),
ElementsAre(0, -1, 1, 2, 3, 0, 0, 4, 1, 2, 3, 4, 5, 6));
EXPECT_THAT(AsVector(data->efc_island, nefc),
ElementsAre(0, 1, 2, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 1, 1));
EXPECT_THAT(AsVector(data->island_efcnum, nisland), ElementsAre(6, 17, 1));
EXPECT_THAT(AsVector(data->island_efcadr, nisland), ElementsAre(0, 6, 23));
EXPECT_THAT(AsVector(data->island_efcind, nefc),
ElementsAre(0, 3, 4, 5, 6, 7, 1, 8, 9, 10, 11, 12, 13, 14, 15, 16,
17, 18, 19, 20, 21, 22, 23, 2));
EXPECT_EQ(data->nefc, 27);
mj_deleteData(data);
mj_deleteModel(model);
}
TEST_F(IslandTest, IslandEfcElliptic) {
const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath);
mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
mjData* data = mj_makeData(model);
model->opt.cone = mjCONE_ELLIPTIC;
while (data->time < 0.2) {
mj_step(model, data);
}
mj_forward(model, data);
EXPECT_EQ(data->nisland, 4);
EXPECT_EQ(data->ne, 4);
EXPECT_EQ(data->nf, 2);
EXPECT_EQ(data->nl, 1);
EXPECT_EQ(data->nefc, 22);
mj_deleteData(data);
mj_deleteModel(model);
+8 -2
View File
@@ -1,10 +1,10 @@
<!--
test model for constraint islanding:
- 1 equality
- 4 equalities (1:joint, 3:connect)
- 1 limit
- 2 friction constraints
- 1 unconstrained dof
- 3 islands, 2 with non-contiguous dofs
- 4 islands, 2 with non-contiguous dofs
-->
<mujoco>
@@ -64,9 +64,15 @@
<freejoint/>
<geom type="box" size=".03 .03 .03" pos="0.01 0.01 0.01"/>
</body>
<body pos="-.5 0 .3" name="connect">
<freejoint/>
<geom type="box" size=".05 .05 .05"/>
</body>
</worldbody>
<equality>
<joint joint1="hinge1" joint2="hinge2"/>
<connect body1="connect" body2="world" anchor="-.05 -.05 .05"/>
</equality>
</mujoco>