From 57b3d6d0f4a8a785c5fd01901bd4ceaaf5e6d235 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 30 Jun 2025 03:15:27 -0700 Subject: [PATCH] Make island discovery compatible with flex PiperOrigin-RevId: 777490642 Change-Id: Iae2f4e8297cc24bded833e04ebfe32895f3a5460 --- src/engine/engine_forward.c | 4 +- src/engine/engine_island.c | 72 ++++++++++++------------------- test/engine/engine_island_test.cc | 29 +++++++++++-- test/testdata/flex.xml | 4 +- 4 files changed, 58 insertions(+), 51 deletions(-) diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 16d33cae..3acaa713 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -173,9 +173,7 @@ void mj_fwdPosition(const mjModel* m, mjData* d) { TM_RESTART; mj_makeConstraint(m, d); - if (mjENABLED(mjENBL_ISLAND)) { - mj_island(m, d); - } + mj_island(m, d); TM_END(mjTIMER_POS_MAKE); TM_RESTART; diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 05364373..7f898f8c 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -141,29 +141,6 @@ int mj_floodFill(int* island, int nr, const int* rownnz, const int* rowadr, cons -// return upper bound on number of tree-tree edges -static int countMaxEdge(const mjModel* m, const mjData* d) { - int nedge_max = 0; - nedge_max += 2*d->ncon; // contact: 2 edges - nedge_max += 2*d->ne; // equality: 2 edges - nedge_max += d->nl; // limit: 1 edges (always within same tree) - nedge_max += d->nf; // joint friction: 1 edge (always within same tree) - - // tendon limits and friction add up to tendon_num edges - for (int i=0; i < m->ntendon; i++) { - if (m->tendon_frictionloss[i]) { - nedge_max += m->tendon_num[i]; - } - if (m->tendon_limited[i]) { - nedge_max += m->tendon_num[i]; - } - } - - return nedge_max; -} - - - // 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 @@ -242,21 +219,27 @@ static int treeFirst(const mjModel* m, const mjData* d, int tree[2], int i) { 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].geom[0]]]; - tree[1] = m->body_treeid[m->geom_bodyid[d->contact[efc_id].geom[1]]]; + int g1 = d->contact[efc_id].geom[0]; + int g2 = d->contact[efc_id].geom[1]; - // 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; + // 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; + } } - } - return -1; + return -1; + } } // connect or weld constraints @@ -373,9 +356,12 @@ static int findEdges(const mjModel* m, const mjData* d, int* treenedge, int* edg int nedge = 0; for (int i=0; i < nefc; i++) { - // row i is still in the same constraint: skip + // row i is still in the same constraint: skip it, if (efc_type == d->efc_type[i] && efc_id == d->efc_id[i]) { - continue; + // unless it is a flex equality, where the tree pattern changes per dof + if (!(efc_type == mjCNSTR_EQUALITY && m->eq_type[efc_id] == mjEQ_FLEX)) { + continue; + } } efc_type = d->efc_type[i]; efc_id = d->efc_id[i]; @@ -425,21 +411,19 @@ void mj_island(const mjModel* m, mjData* d) { int nv = m->nv, nefc = d->nefc, ntree=m->ntree; // no constraints: quick return - if (!nefc || m->nflex) { // TODO: add flex support to island discovery - d->nisland = 0; - d->nidof = 0; + if (!mjENABLED(mjENBL_ISLAND) || !nefc) { + d->nisland = d->nidof = 0; return; } mj_markStack(d); - // allocate edge array - int nedge_max = countMaxEdge(m, d); - int* edge = mjSTACKALLOC(d, 2*nedge_max, int); + // allocate edge array, nJ is an upper bound + int* edge = mjSTACKALLOC(d, 2*d->nJ, int); // get tree-tree edges and rownnz counts from efc arrays int* rownnz = mjSTACKALLOC(d, ntree, int); // number of edges per tree - int nedge = findEdges(m, d, rownnz, edge, nedge_max); + int nedge = findEdges(m, d, rownnz, edge, d->nJ); // compute starting address of tree's column indices while resetting rownnz int* rowadr = mjSTACKALLOC(d, ntree, int); diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index e59ac3b4..80437e7c 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -28,12 +28,11 @@ namespace mujoco { namespace { +using ::testing::DoubleNear; using ::testing::ElementsAre; +using ::testing::Pointwise; using IslandTest = MujocoTest; -std::vector AsVector(const int* array, int n) { - return std::vector(array, array + n); -} TEST_F(IslandTest, FloodFillSingleton) { // adjacency matrix for the graph 0 1 2 @@ -368,6 +367,30 @@ TEST_F(IslandTest, IslandEfc) { mj_deleteModel(model); } +TEST_F(IslandTest, IslandFlex) { + const std::string xml_path = GetTestDataFilePath("testdata/flex.xml"); + mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + mjData* data1 = mj_makeData(model); + mjData* data2 = mj_makeData(model); + + model->opt.enableflags |= mjENBL_ISLAND; + while (data1->time < 0.2) { + mj_step(model, data1); + } + + model->opt.enableflags &= ~mjENBL_ISLAND; + while (data2->time < 0.2) { + mj_step(model, data2); + } + + EXPECT_THAT(AsVector(data1->qpos, model->nq), + Pointwise(DoubleNear(1e-6), AsVector(data2->qpos, model->nq))); + + mj_deleteData(data2); + mj_deleteData(data1); + mj_deleteModel(model); +} + static const char* const k2H100Path = "engine/testdata/island/2humanoid100.xml"; TEST_F(IslandTest, IslandJacobian) { diff --git a/test/testdata/flex.xml b/test/testdata/flex.xml index 4e243dbf..0bc6403a 100644 --- a/test/testdata/flex.xml +++ b/test/testdata/flex.xml @@ -1,5 +1,7 @@ -