Fix island graph construction for multi-tree constraints, fixes #3073
Replace dynamically-sized edge list with a dense ntree×ntree adjacency matrix for deduplication and a flat CSR representation. PiperOrigin-RevId: 868728432 Change-Id: Iebe97a0870740e5465549b967bbde0c81e658269
This commit is contained in:
committed by
Copybara-Service
parent
cc5d544d15
commit
2b5afce4aa
+54
-92
@@ -280,74 +280,49 @@ static int treeFirst(const mjModel* m, const mjData* d, int tree[2], int i) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
// add 0 edges, 1 self-edge or 2 flipped edges to array, increment treenedge
|
// add 0 edges, 1 self-edge or 2 flipped edges to array
|
||||||
// return current number of edges
|
// increment rownnz using tree_tree to de-dupe; return number of edges added
|
||||||
static int addEdge(int* treenedge, int* edge, int nedge, int tree1, int tree2, int nedge_max) {
|
static int addEdge(int* rownnz, int* edge, mjtByte* tree_tree, int ntree, int tree1, int tree2) {
|
||||||
// handle the static tree
|
|
||||||
if (tree1 == -1 && tree2 == -1) {
|
if (tree1 == -1 && tree2 == -1) {
|
||||||
mjERROR("self-edge of the static tree"); // SHOULD NOT OCCUR
|
mjERROR("self-edge of the static tree"); // SHOULD NOT OCCUR
|
||||||
return 0;
|
return 0;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// handle static trees (treat as self-edge)
|
||||||
if (tree1 == -1) tree1 = tree2;
|
if (tree1 == -1) tree1 = tree2;
|
||||||
if (tree2 == -1) tree2 = tree1;
|
if (tree2 == -1) tree2 = tree1;
|
||||||
|
|
||||||
// previous edge
|
// skip if edge already present
|
||||||
int p1 = nedge ? edge[2*nedge - 2] : -1;
|
if (tree_tree[tree1*ntree + tree2]) {
|
||||||
int p2 = nedge ? edge[2*nedge - 1] : -1;
|
|
||||||
|
|
||||||
// === self edge
|
|
||||||
if (tree1 == tree2) {
|
|
||||||
// same as previous edge, return
|
|
||||||
if (nedge && tree1 == p1 && tree1 == p2) {
|
|
||||||
return nedge;
|
|
||||||
}
|
|
||||||
|
|
||||||
// check size
|
|
||||||
if (nedge >= nedge_max) {
|
|
||||||
mjERROR("edge array too small");
|
|
||||||
return 0;
|
|
||||||
}
|
|
||||||
|
|
||||||
// add tree1-tree1 self-edge
|
|
||||||
edge[2*nedge + 0] = tree1;
|
|
||||||
edge[2*nedge + 1] = tree1;
|
|
||||||
treenedge[tree1]++;
|
|
||||||
return nedge + 1;
|
|
||||||
}
|
|
||||||
|
|
||||||
// === non-self edge
|
|
||||||
if (nedge && ((tree1 == p1 && tree2 == p2) || (tree1 == p2 && tree2 == p1))) {
|
|
||||||
// same as previous edge, return
|
|
||||||
return nedge;
|
|
||||||
}
|
|
||||||
|
|
||||||
// check size
|
|
||||||
if (nedge + 2 > nedge_max) {
|
|
||||||
mjERROR("edge array too small");
|
|
||||||
return 0;
|
return 0;
|
||||||
}
|
}
|
||||||
|
|
||||||
// add tree1-tree2 and tree2-tree1
|
// add edge
|
||||||
edge[2*nedge + 0] = tree1;
|
tree_tree[tree1*ntree + tree2] = 1;
|
||||||
edge[2*nedge + 1] = tree2;
|
edge[tree1*ntree + rownnz[tree1]++] = tree2;
|
||||||
edge[2*nedge + 2] = tree2;
|
|
||||||
edge[2*nedge + 3] = tree1;
|
// add flipped edge (off-diagonal)
|
||||||
treenedge[tree1]++;
|
if (tree1 != tree2) {
|
||||||
treenedge[tree2]++;
|
tree_tree[tree2*ntree + tree1] = 1;
|
||||||
return nedge + 2;
|
edge[tree2*ntree + rownnz[tree2]++] = tree1;
|
||||||
|
return 2;
|
||||||
|
}
|
||||||
|
|
||||||
|
return 1;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
// find tree-tree edges, increment treenedge counters, return total number of edges
|
// find tree-tree edges (column indices), increment row nonzeros, return total number of edges
|
||||||
static int findEdges(const mjModel* m, const mjData* d, int* treenedge, int* edge, int nedge_max) {
|
static int findEdges(const mjModel* m, const mjData* d,
|
||||||
|
int* rownnz, int* colind, mjtByte* tree_tree, int ntree) {
|
||||||
int nefc = d->nefc;
|
int nefc = d->nefc;
|
||||||
|
int nnz = 0;
|
||||||
int efc_type = -1;
|
int efc_type = -1;
|
||||||
int efc_id = -1;
|
int efc_id = -1;
|
||||||
|
|
||||||
// clear treenedge
|
// clear row nonzeros
|
||||||
mju_zeroInt(treenedge, m->ntree);
|
mju_zeroInt(rownnz, ntree);
|
||||||
|
|
||||||
int nedge = 0;
|
|
||||||
for (int i=0; i < nefc; i++) {
|
for (int i=0; i < nefc; i++) {
|
||||||
// row i is still in the same constraint: skip it,
|
// row i is still in the same constraint: skip it,
|
||||||
if (efc_type == d->efc_type[i] && efc_id == d->efc_id[i]) {
|
if (efc_type == d->efc_type[i] && efc_id == d->efc_id[i]) {
|
||||||
@@ -366,34 +341,29 @@ static int findEdges(const mjModel* m, const mjData* d, int* treenedge, int* edg
|
|||||||
int tree1 = tree[0];
|
int tree1 = tree[0];
|
||||||
int tree2 = tree[1];
|
int tree2 = tree[1];
|
||||||
|
|
||||||
// no more edges to find, add and continue
|
// simple constraint (no more edges), add edge and continue
|
||||||
if (index == -1) {
|
if (index == -1) {
|
||||||
nedge = addEdge(treenedge, edge, nedge, tree1, tree2 == -1 ? tree1 : tree2, nedge_max);
|
nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, tree2);
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
// possibly more edges, scan Jacobian row
|
// add first edge
|
||||||
else {
|
tree2 = treeNext(m, d, tree1, i, &index);
|
||||||
tree2 = treeNext(m, d, tree1, i, &index);
|
nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, tree2);
|
||||||
|
|
||||||
if (tree2 == -1) {
|
// possibly more edges, scan Jacobian row
|
||||||
// 1 tree found: add self-edge
|
if (tree2 != -1) {
|
||||||
nedge = addEdge(treenedge, edge, nedge, tree1, tree1, nedge_max);
|
int tree3 = treeNext(m, d, tree2, i, &index);
|
||||||
} else {
|
while (tree3 > -1 && tree3 != tree2) {
|
||||||
// 2 trees found: add edge, keep scanning and adding until no more trees
|
tree1 = tree2;
|
||||||
nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max);
|
tree2 = tree3;
|
||||||
int tree3 = treeNext(m, d, tree2, i, &index);
|
nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, tree2);
|
||||||
while (tree3 > -1 && tree3 != tree2) {
|
tree3 = treeNext(m, d, tree2, i, &index);
|
||||||
tree1 = tree2;
|
|
||||||
tree2 = tree3;
|
|
||||||
nedge = addEdge(treenedge, edge, nedge, tree1, tree2, nedge_max);
|
|
||||||
tree3 = treeNext(m, d, tree2, i, &index);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
return nedge;
|
return nnz;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
@@ -402,9 +372,9 @@ static int findEdges(const mjModel* m, const mjData* d, int* treenedge, int* edg
|
|||||||
// discover islands:
|
// discover islands:
|
||||||
// nisland, island_idofadr, dof_island, dof_islandnext, island_efcadr, efc_island, efc_islandnext
|
// nisland, island_idofadr, dof_island, dof_islandnext, island_efcadr, efc_island, efc_islandnext
|
||||||
void mj_island(const mjModel* m, mjData* d) {
|
void mj_island(const mjModel* m, mjData* d) {
|
||||||
int nv = m->nv, nefc = d->nefc, ntree = m->ntree, nJ = d->nJ;
|
int nv = m->nv, nefc = d->nefc, ntree = m->ntree;
|
||||||
|
|
||||||
// no constraints: quick return
|
// no constraints or islands disabled: quick return
|
||||||
if (mjDISABLED(mjDSBL_ISLAND) || !nefc) {
|
if (mjDISABLED(mjDSBL_ISLAND) || !nefc) {
|
||||||
d->nisland = d->nidof = 0;
|
d->nisland = d->nidof = 0;
|
||||||
return;
|
return;
|
||||||
@@ -412,33 +382,25 @@ void mj_island(const mjModel* m, mjData* d) {
|
|||||||
|
|
||||||
mj_markStack(d);
|
mj_markStack(d);
|
||||||
|
|
||||||
// allocate edge array, nJ is an upper bound
|
// dense tree-tree adjacency matrix
|
||||||
int* edge = mjSTACKALLOC(d, 2*nJ, int);
|
int ntree2 = ntree * ntree;
|
||||||
|
mjtByte* tree_tree = mjSTACKALLOC(d, ntree2, mjtByte);
|
||||||
|
memset(tree_tree, 0, ntree2);
|
||||||
|
|
||||||
// get tree-tree edges and rownnz counts from efc arrays
|
// CSR representation of tree-tree adjacency matrix (uncompressed)
|
||||||
int* rownnz = mjSTACKALLOC(d, ntree, int); // number of edges per tree
|
int* colind = mjSTACKALLOC(d, ntree2, int);
|
||||||
int nedge = findEdges(m, d, rownnz, edge, nJ);
|
int* rownnz = mjSTACKALLOC(d, ntree, int);
|
||||||
|
|
||||||
// compute starting address of tree's column indices while resetting rownnz
|
|
||||||
int* rowadr = mjSTACKALLOC(d, ntree, int);
|
int* rowadr = mjSTACKALLOC(d, ntree, int);
|
||||||
rowadr[0] = 0;
|
for (int r=0; r < ntree; r++) {
|
||||||
for (int r=1; r < ntree; r++) {
|
rowadr[r] = r * ntree;
|
||||||
rowadr[r] = rowadr[r-1] + rownnz[r-1];
|
|
||||||
rownnz[r-1] = 0;
|
|
||||||
}
|
}
|
||||||
rownnz[ntree-1] = 0;
|
|
||||||
|
|
||||||
// copy column indices: list each tree's neighbors
|
// compute tree-tree adjacency matrix: fill rownnz and colind
|
||||||
int* colind = mjSTACKALLOC(d, nedge, int);
|
int nnz = findEdges(m, d, rownnz, colind, tree_tree, ntree);
|
||||||
for (int e=0; e < nedge; e++) {
|
|
||||||
int row = edge[2*e];
|
|
||||||
int col = edge[2*e + 1];
|
|
||||||
colind[rowadr[row] + rownnz[row]++] = col;
|
|
||||||
}
|
|
||||||
|
|
||||||
// discover islands
|
// discover islands
|
||||||
int* tree_island = mjSTACKALLOC(d, ntree, int); // id of island assigned to tree
|
int* tree_island = mjSTACKALLOC(d, ntree, int);
|
||||||
int* stack = mjSTACKALLOC(d, nedge, int);
|
int* stack = mjSTACKALLOC(d, nnz, int);
|
||||||
d->nisland = mj_floodFill(tree_island, ntree, rownnz, rowadr, colind, stack);
|
d->nisland = mj_floodFill(tree_island, ntree, rownnz, rowadr, colind, stack);
|
||||||
|
|
||||||
// no islands found: quick return
|
// no islands found: quick return
|
||||||
|
|||||||
@@ -555,5 +555,52 @@ TEST_F(IslandTest, IslandEfcElliptic) {
|
|||||||
mj_deleteModel(model);
|
mj_deleteModel(model);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
TEST_F(IslandTest, EqualityConstraintOfTendons) {
|
||||||
|
static const char xml[] = R"(
|
||||||
|
<mujoco>
|
||||||
|
<worldbody>
|
||||||
|
<body name="b1">
|
||||||
|
<inertial pos="0 0 0" mass="1" diaginertia="1 1 1"/>
|
||||||
|
<joint name="j1" type="slide" axis="1 0 0"/>
|
||||||
|
</body>
|
||||||
|
<body name="b2">
|
||||||
|
<inertial pos="0 0 0" mass="1" diaginertia="1 1 1"/>
|
||||||
|
<joint name="j2" type="slide" axis="1 0 0"/>
|
||||||
|
</body>
|
||||||
|
<body name="b3">
|
||||||
|
<inertial pos="0 0 0" mass="1" diaginertia="1 1 1"/>
|
||||||
|
<joint name="j3" type="slide" axis="1 0 0"/>
|
||||||
|
</body>
|
||||||
|
<body name="b4">
|
||||||
|
<inertial pos="0 0 0" mass="1" diaginertia="1 1 1"/>
|
||||||
|
<joint name="j4" type="slide" axis="1 0 0"/>
|
||||||
|
</body>
|
||||||
|
</worldbody>
|
||||||
|
|
||||||
|
<tendon>
|
||||||
|
<fixed name="t12">
|
||||||
|
<joint joint="j1" coef="1"/>
|
||||||
|
<joint joint="j2" coef="1"/>
|
||||||
|
</fixed>
|
||||||
|
<fixed name="t34">
|
||||||
|
<joint joint="j3" coef="1"/>
|
||||||
|
<joint joint="j4" coef="1"/>
|
||||||
|
</fixed>
|
||||||
|
</tendon>
|
||||||
|
|
||||||
|
<equality>
|
||||||
|
<tendon name="eq" tendon1="t12" tendon2="t34"/>
|
||||||
|
</equality>
|
||||||
|
</mujoco>
|
||||||
|
)";
|
||||||
|
char error[1024];
|
||||||
|
mjModel* model = LoadModelFromString(xml, error, sizeof(error));
|
||||||
|
ASSERT_THAT(model, NotNull()) << error;
|
||||||
|
mjData* data = mj_makeData(model);
|
||||||
|
mj_forward(model, data);
|
||||||
|
mj_deleteData(data);
|
||||||
|
mj_deleteModel(model);
|
||||||
|
}
|
||||||
|
|
||||||
} // namespace
|
} // namespace
|
||||||
} // namespace mujoco
|
} // namespace mujoco
|
||||||
|
|||||||
Reference in New Issue
Block a user