Refactor constraint tree iteration and island assignment.
- Simplified constraint tree discovery with a unified iterator. - Filling in a new `efc_tree` array to avoid calling tree discovery again in the constraint phase. - Implemented dof skipping in the dense case for efficiency. PiperOrigin-RevId: 875440387 Change-Id: Id44a5a6fe2a80edaaf30935506a023a307e21303
This commit is contained in:
Yuval Tassa
committed by
Copybara-Service
Copybara-Service
parent
b7e192e514
commit
7232ec3b47
+125
-117
@@ -137,24 +137,42 @@ int mj_floodFill(int* island, int nr, const int* rownnz, const int* rowadr, cons
|
||||
}
|
||||
|
||||
|
||||
// return id of next tree in Jacobian row i that is different from tree, -1 if not found
|
||||
// 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 j; // local loop variable, saved to *index
|
||||
// state of iterator for finding trees involved in a constraint
|
||||
typedef struct {
|
||||
int trees[2]; // pre-calculated trees (special-cased constraints); -2: empty/sentinel
|
||||
int jac_idx; // generic scan: current lookup index in Jacobian row; -1: scan disabled
|
||||
int tree_prev; // generic scan: previous tree in ongoing scan
|
||||
} mjTreeIter;
|
||||
|
||||
|
||||
// return next tree of constraint i from iterator; -2: no more trees
|
||||
static int treeNext(const mjModel* m, const mjData* d, int i, mjTreeIter* iter) {
|
||||
// handle special cases
|
||||
if (iter->trees[0] != -2) {
|
||||
// get first tree, queue up second tree, return first tree
|
||||
int tree = iter->trees[0];
|
||||
iter->trees[0] = iter->trees[1];
|
||||
iter->trees[1] = -2;
|
||||
return tree;
|
||||
}
|
||||
|
||||
// special case mode complete
|
||||
if (iter->jac_idx == -1) {
|
||||
return -2;
|
||||
}
|
||||
|
||||
// generic scan mode
|
||||
int j;
|
||||
int tree_next = -2;
|
||||
|
||||
// sparse
|
||||
if (mj_isSparse(m)) {
|
||||
int rownnz = d->efc_J_rownnz[i];
|
||||
int* colind = d->efc_J_colind + d->efc_J_rowadr[i];
|
||||
|
||||
// loop over remaining nonzeros, look for different tree
|
||||
for (j=(*index); j < rownnz; j++) {
|
||||
const int* colind = d->efc_J_colind + d->efc_J_rowadr[i];
|
||||
for (j = iter->jac_idx; j < rownnz; j++) {
|
||||
int tree_j = m->dof_treeid[colind[j]];
|
||||
if (tree_j != tree) {
|
||||
// found different tree
|
||||
if (tree_j != iter->tree_prev) {
|
||||
// found new tree
|
||||
tree_next = tree_j;
|
||||
break;
|
||||
}
|
||||
@@ -164,125 +182,107 @@ static int treeNext(const mjModel* m, const mjData* d, int tree, int i, int *ind
|
||||
// dense
|
||||
else {
|
||||
int nv = m->nv;
|
||||
|
||||
// scan row, look for different tree
|
||||
for (j=(*index); j < nv; j++) {
|
||||
if (d->efc_J[nv*i + j]) {
|
||||
const mjtNum* J = d->efc_J + nv * i;
|
||||
for (j = iter->jac_idx; j < nv; j++) {
|
||||
if (J[j]) {
|
||||
int tree_j = m->dof_treeid[j];
|
||||
if (tree_j != tree) {
|
||||
// found different tree
|
||||
if (tree_j != iter->tree_prev) {
|
||||
// found new tree
|
||||
tree_next = tree_j;
|
||||
break;
|
||||
}
|
||||
|
||||
// skip to end of tree's dof block
|
||||
j = m->tree_dofadr[tree_j] + m->tree_dofnum[tree_j] - 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// save last index
|
||||
*index = j;
|
||||
// update iterator state
|
||||
iter->jac_idx = j;
|
||||
if (tree_next != -2) {
|
||||
iter->tree_prev = tree_next;
|
||||
}
|
||||
|
||||
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) {
|
||||
// initialize tree iterator, handle special cases
|
||||
static void treeIterInit(const mjModel* m, const mjData* d, int i, mjTreeIter* iter) {
|
||||
iter->trees[0] = -2;
|
||||
iter->trees[1] = -2;
|
||||
iter->jac_idx = -1;
|
||||
iter->tree_prev = -1;
|
||||
|
||||
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
|
||||
// ==== special cases: fill iter->trees where possible
|
||||
|
||||
// joint friction
|
||||
if (efc_type == mjCNSTR_FRICTION_DOF) {
|
||||
tree[0] = m->dof_treeid[efc_id];
|
||||
return -1;
|
||||
iter->trees[0] = m->dof_treeid[efc_id];
|
||||
}
|
||||
|
||||
// joint limit
|
||||
if (efc_type == mjCNSTR_LIMIT_JOINT) {
|
||||
tree[0] = m->dof_treeid[m->jnt_dofadr[efc_id]];
|
||||
return -1;
|
||||
else if (efc_type == mjCNSTR_LIMIT_JOINT) {
|
||||
iter->trees[0] = m->dof_treeid[m->jnt_dofadr[efc_id]];
|
||||
}
|
||||
|
||||
// contact
|
||||
if (efc_type == mjCNSTR_CONTACT_FRICTIONLESS ||
|
||||
efc_type == mjCNSTR_CONTACT_PYRAMIDAL ||
|
||||
efc_type == mjCNSTR_CONTACT_ELLIPTIC) {
|
||||
else if (efc_type == mjCNSTR_CONTACT_FRICTIONLESS ||
|
||||
efc_type == mjCNSTR_CONTACT_PYRAMIDAL ||
|
||||
efc_type == mjCNSTR_CONTACT_ELLIPTIC) {
|
||||
int g1 = d->contact[efc_id].geom[0];
|
||||
int g2 = d->contact[efc_id].geom[1];
|
||||
|
||||
// no shortcut for flex contacts (handled in the generic case)
|
||||
if (g1 >=0 && g2 >= 0) {
|
||||
tree[0] = m->body_treeid[m->geom_bodyid[g1]];
|
||||
tree[1] = m->body_treeid[m->geom_bodyid[g2]];
|
||||
|
||||
// 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;
|
||||
}
|
||||
// geom-geom contact
|
||||
if (g1 >= 0 && g2 >= 0) {
|
||||
iter->trees[0] = m->body_treeid[m->geom_bodyid[g1]];
|
||||
iter->trees[1] = m->body_treeid[m->geom_bodyid[g2]];
|
||||
if (iter->trees[0] < 0 && iter->trees[1] < 0) {
|
||||
mjERROR("contact %d is between two static bodies", efc_id); // SHOULD NOT OCCUR
|
||||
}
|
||||
}
|
||||
|
||||
return -1;
|
||||
// no shortcut for flex contacts: enable generic scan
|
||||
else {
|
||||
iter->jac_idx = 0;
|
||||
}
|
||||
}
|
||||
|
||||
// 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) {
|
||||
int b1 = m->eq_obj1id[efc_id];
|
||||
int b2 = m->eq_obj2id[efc_id];
|
||||
else if (efc_type == mjCNSTR_EQUALITY &&
|
||||
(m->eq_type[efc_id] == mjEQ_CONNECT ||
|
||||
m->eq_type[efc_id] == mjEQ_WELD)) {
|
||||
int b1 = m->eq_obj1id[efc_id];
|
||||
int b2 = m->eq_obj2id[efc_id];
|
||||
|
||||
// get body ids if using site semantics
|
||||
if (m->eq_objtype[efc_id] == mjOBJ_SITE) {
|
||||
b1 = m->site_bodyid[b1];
|
||||
b2 = m->site_bodyid[b2];
|
||||
}
|
||||
// get body ids if using site semantics
|
||||
if (m->eq_objtype[efc_id] == mjOBJ_SITE) {
|
||||
b1 = m->site_bodyid[b1];
|
||||
b2 = m->site_bodyid[b2];
|
||||
}
|
||||
|
||||
tree[0] = m->body_treeid[b1];
|
||||
tree[1] = m->body_treeid[b2];
|
||||
|
||||
// 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;
|
||||
// get trees
|
||||
iter->trees[0] = m->body_treeid[b1];
|
||||
iter->trees[1] = m->body_treeid[b2];
|
||||
if (iter->trees[0] < 0 && iter->trees[1] < 0) {
|
||||
mjERROR("equality %d is between two static bodies", efc_id); // SHOULD NOT OCCUR
|
||||
}
|
||||
}
|
||||
|
||||
// ==== 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
|
||||
// otherwise enable generic scan
|
||||
else {
|
||||
iter->jac_idx = 0;
|
||||
}
|
||||
|
||||
return index;
|
||||
}
|
||||
|
||||
|
||||
// add 0 edges, 1 self-edge or 2 flipped edges to array
|
||||
// add 0, 1 or 2 edges to uncompressed CSR adjacency matrix
|
||||
// increment rownnz using tree_tree to de-dupe; return number of edges added
|
||||
static int addEdge(int* rownnz, int* edge, mjtByte* tree_tree, int ntree, int tree1, int tree2) {
|
||||
static int addEdge(int* rownnz, int* colind, mjtByte* tree_tree, int ntree, int tree1, int tree2) {
|
||||
if (tree1 == -1 && tree2 == -1) {
|
||||
mjERROR("self-edge of the static tree"); // SHOULD NOT OCCUR
|
||||
return 0;
|
||||
@@ -299,12 +299,12 @@ static int addEdge(int* rownnz, int* edge, mjtByte* tree_tree, int ntree, int tr
|
||||
|
||||
// add edge
|
||||
tree_tree[tree1*ntree + tree2] = 1;
|
||||
edge[tree1*ntree + rownnz[tree1]++] = tree2;
|
||||
colind[tree1*ntree + rownnz[tree1]++] = tree2; // uncompressed format, rowadr is known
|
||||
|
||||
// add flipped edge (off-diagonal)
|
||||
if (tree1 != tree2) {
|
||||
tree_tree[tree2*ntree + tree1] = 1;
|
||||
edge[tree2*ntree + rownnz[tree2]++] = tree1;
|
||||
colind[tree2*ntree + rownnz[tree2]++] = tree1; // uncompressed format, rowadr is known
|
||||
return 2;
|
||||
}
|
||||
|
||||
@@ -312,9 +312,10 @@ static int addEdge(int* rownnz, int* edge, mjtByte* tree_tree, int ntree, int tr
|
||||
}
|
||||
|
||||
|
||||
// find tree-tree edges (column indices), increment row nonzeros, return total number of edges
|
||||
// find tree-tree edges (column indices), return total number of edges
|
||||
// efc_tree: first nonegative tree index of each constraint
|
||||
static int findEdges(const mjModel* m, const mjData* d,
|
||||
int* rownnz, int* colind, mjtByte* tree_tree, int ntree) {
|
||||
int* rownnz, int* colind, mjtByte* tree_tree, int* efc_tree, int ntree) {
|
||||
int nefc = d->nefc;
|
||||
int nnz = 0;
|
||||
int efc_type = -1;
|
||||
@@ -323,6 +324,7 @@ static int findEdges(const mjModel* m, const mjData* d,
|
||||
// clear row nonzeros
|
||||
mju_zeroInt(rownnz, ntree);
|
||||
|
||||
// iterate over constraints, compute tree-tree edges, assign efc_tree
|
||||
for (int i=0; i < nefc; i++) {
|
||||
// row i is still in the same constraint: skip it,
|
||||
if (efc_type == d->efc_type[i] && efc_id == d->efc_id[i]) {
|
||||
@@ -330,36 +332,41 @@ static int findEdges(const mjModel* m, const mjData* d,
|
||||
if (!(efc_type == mjCNSTR_EQUALITY &&
|
||||
(m->eq_type[efc_id] == mjEQ_FLEX ||
|
||||
m->eq_type[efc_id] == mjEQ_FLEXVERT))) {
|
||||
// copy tree assignment from previous constraint and continue
|
||||
efc_tree[i] = efc_tree[i-1];
|
||||
continue;
|
||||
}
|
||||
}
|
||||
efc_type = d->efc_type[i];
|
||||
efc_id = d->efc_id[i];
|
||||
|
||||
int tree[2];
|
||||
int index = treeFirst(m, d, tree, i);
|
||||
int tree1 = tree[0];
|
||||
int tree2 = tree[1];
|
||||
// initialize tree iterator
|
||||
mjTreeIter iter;
|
||||
treeIterInit(m, d, i, &iter);
|
||||
|
||||
// simple constraint (no more edges), add edge and continue
|
||||
if (index == -1) {
|
||||
nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, tree2);
|
||||
continue;
|
||||
}
|
||||
// iterate over trees involved in constraint i
|
||||
int tree1 = treeNext(m, d, i, &iter);
|
||||
if (tree1 != -2) {
|
||||
int tree2 = treeNext(m, d, i, &iter);
|
||||
|
||||
// add first edge
|
||||
tree2 = treeNext(m, d, tree1, i, &index);
|
||||
nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, tree2);
|
||||
|
||||
// possibly more edges, scan Jacobian row
|
||||
if (tree2 != -1) {
|
||||
int tree3 = treeNext(m, d, tree2, i, &index);
|
||||
while (tree3 > -1 && tree3 != tree2) {
|
||||
tree1 = tree2;
|
||||
tree2 = tree3;
|
||||
nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, tree2);
|
||||
tree3 = treeNext(m, d, tree2, i, &index);
|
||||
// assign tree to constraint, one of (tree1, tree2) must be non-negative
|
||||
efc_tree[i] = tree1 >= 0 ? tree1 : tree2;
|
||||
if (efc_tree[i] < 0) {
|
||||
mjERROR("constraint %d is between two static bodies", i); // SHOULD NOT OCCUR
|
||||
}
|
||||
|
||||
// add one edge or continue to search for more edges
|
||||
if (tree2 == -2) {
|
||||
nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, -1);
|
||||
} else {
|
||||
while (tree2 != -2) {
|
||||
nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, tree2);
|
||||
tree1 = tree2;
|
||||
tree2 = treeNext(m, d, i, &iter);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
mjERROR("no tree found for constraint %d", i); // SHOULD NOT OCCUR
|
||||
}
|
||||
}
|
||||
|
||||
@@ -395,8 +402,11 @@ void mj_island(const mjModel* m, mjData* d) {
|
||||
rowadr[r] = r * ntree;
|
||||
}
|
||||
|
||||
// first non-negative tree index of each constraint, used later for computing efc_island
|
||||
int* efc_tree = mjSTACKALLOC(d, nefc, int);
|
||||
|
||||
// compute tree-tree adjacency matrix: fill rownnz and colind
|
||||
int nnz = findEdges(m, d, rownnz, colind, tree_tree, ntree);
|
||||
int nnz = findEdges(m, d, rownnz, colind, tree_tree, efc_tree, ntree);
|
||||
|
||||
// discover islands
|
||||
int* tree_island = mjSTACKALLOC(d, ntree, int);
|
||||
@@ -532,9 +542,7 @@ void mj_island(const mjModel* m, mjData* d) {
|
||||
mju_zeroInt(d->island_nf, nisland);
|
||||
mju_zeroInt(d->island_nefc, nisland);
|
||||
for (int i=0; i < nefc; i++) {
|
||||
int tree[2];
|
||||
treeFirst(m, d, tree, i);
|
||||
int island = tree_island[tree[0]];
|
||||
int island = tree_island[efc_tree[i]];
|
||||
d->efc_island[i] = island;
|
||||
d->island_nefc[island]++;
|
||||
switch (d->efc_type[i]) {
|
||||
|
||||
Reference in New Issue
Block a user