From d3ed16e78ff7614f02aab298f1c0509e5c47f682 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Tue, 29 Aug 2023 11:05:41 -0700 Subject: [PATCH] Add mjData.dof_islandind, indexing from dofs into islands. Also fix bug in allocation for constraint islands, added allocation for limit constraints. PiperOrigin-RevId: 561083128 Change-Id: I9e01d5a5b7454ec3cf60bccadaa0f96f2c0d87d1 --- doc/includes/references.h | 1 + include/mujoco/mjdata.h | 1 + include/mujoco/mjxmacro.h | 1 + introspect/structs.py | 7 +++ src/engine/engine_island.c | 8 +++- src/engine/engine_print.c | 6 +++ test/engine/engine_island_test.cc | 53 +++++++++++++++++++++- test/engine/testdata/island/island_efc.xml | 53 ++++++++++++++++++++++ unity/Runtime/Bindings/MjBindings.cs | 1 + 9 files changed, 127 insertions(+), 4 deletions(-) create mode 100644 test/engine/testdata/island/island_efc.xml diff --git a/doc/includes/references.h b/doc/includes/references.h index 72fc0dc1..8b63b22a 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -350,6 +350,7 @@ struct mjData_ { int* island_dofnum; // number of dofs in island (nisland x 1) int* island_dofadr; // start address in island_dofind (nisland x 1) int* island_dofind; // island dof indices; -1: none (nv x 1) + int* dof_islandind; // dof island indices; -1: none (nv x 1) int* efc_island; // island id of this constraint (nefc x 1) int* island_efcnum; // number of constraints in island (nisland x 1) int* island_efcadr; // start address in island_efcind (nisland x 1) diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 752263ea..756b7675 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -377,6 +377,7 @@ struct mjData_ { int* island_dofnum; // number of dofs in island (nisland x 1) int* island_dofadr; // start address in island_dofind (nisland x 1) int* island_dofind; // island dof indices; -1: none (nv x 1) + int* dof_islandind; // dof island indices; -1: none (nv x 1) int* efc_island; // island id of this constraint (nefc x 1) int* island_efcnum; // number of constraints in island (nisland x 1) int* island_efcadr; // start address in island_efcind (nisland x 1) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index c35c00a6..7d5c5724 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -605,6 +605,7 @@ X( int, island_dofnum, MJ_D(nisland), 1 ) \ X( int, island_dofadr, MJ_D(nisland), 1 ) \ X( int, island_dofind, MJ_M(nv), 1 ) \ + X( int, dof_islandind, MJ_M(nv), 1 ) \ X( int, efc_island, MJ_D(nefc), 1 ) \ X( int, island_efcnum, MJ_D(nisland), 1 ) \ X( int, island_efcadr, MJ_D(nisland), 1 ) \ diff --git a/introspect/structs.py b/introspect/structs.py index efdadbf1..f0cc95d9 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -4372,6 +4372,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), doc='island dof indices; -1: none (nv x 1)', # pylint: disable=line-too-long ), + StructFieldDecl( + name='dof_islandind', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='dof island indices; -1: none (nv x 1)', # pylint: disable=line-too-long + ), StructFieldDecl( name='efc_island', type=PointerType( diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 4f3699dd..3c90f032 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -141,6 +141,7 @@ 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 @@ -425,13 +426,16 @@ void mj_island(const mjModel* m, mjData* d) { // reset island_dofnum memset(d->island_dofnum, 0, nisland*sizeof(int)); - // compute dof_islandind + // compute dof_islandind, island_dofind int num_dof_island = 0; for (int i=0; i < nv; i++) { int island = d->dof_island[i]; if (island >= 0) { - d->island_dofind[d->island_dofadr[island] + (d->island_dofnum[island]++)] = i; + d->island_dofind[d->island_dofadr[island] + d->island_dofnum[island]] = i; + d->dof_islandind[i] = d->island_dofnum[island]++; num_dof_island++; + } else { + d->dof_islandind[i] = -1; } } diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index c1f8a947..8cd95131 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -1083,6 +1083,12 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename, } fprintf(fp, "\n\n"); + fprintf(fp, NAME_FORMAT, "DOF_ISLANDIND"); + for (int i = 0; i < m->nv; i++) { + fprintf(fp, " %d", d->dof_islandind[i]); + } + fprintf(fp, "\n\n"); + fprintf(fp, NAME_FORMAT, "EFC_ISLAND"); for (int i = 0; i < d->nefc; i++) { fprintf(fp, " %d", d->efc_island[i]); diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index 02f09beb..1c631254 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -230,6 +230,11 @@ TEST_F(IslandTest, Abacus) { // last index is unassigned since dof 1 is unconstrained EXPECT_THAT(AsVector(data->island_dofind, nv), ElementsAre(0, 2, 3, -1)); + // dof 0 constitutes first island + // dofs 1 is unassigned + // dofs 2, 3 are second island + EXPECT_THAT(AsVector(data->dof_islandind, nv), ElementsAre(0, -1, 0, 1)); + // island 0 starts at constraint 0 // island 1 starts at constraint 4 EXPECT_THAT(AsVector(data->island_efcadr, nisland), ElementsAre(0, 4)); @@ -242,8 +247,7 @@ TEST_F(IslandTest, Abacus) { EXPECT_THAT(AsVector(data->efc_island, nefc), ElementsAre(0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1)); - // linked list for island 0 - // linked list for island 1 + // index lists for islands 0 and 1 EXPECT_THAT(AsVector(data->island_efcind, nefc), ElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11)); @@ -268,6 +272,7 @@ TEST_F(IslandTest, Abacus) { EXPECT_THAT(AsVector(data->island_dofnum, nisland), ElementsAre(1, 2, 1)); EXPECT_THAT(AsVector(data->dof_island, nv), ElementsAre(0, 1, 1, 2)); EXPECT_THAT(AsVector(data->island_dofind, nv), ElementsAre(0, 1, 2, 3)); + EXPECT_THAT(AsVector(data->dof_islandind, nv), ElementsAre(0, 0, 1, 0)); EXPECT_THAT(AsVector(data->island_efcadr, nisland), ElementsAre(0, 4, 8)); EXPECT_THAT(AsVector(data->island_efcnum, nisland), ElementsAre(4, 4, 4)); EXPECT_THAT(AsVector(data->efc_island, nefc), @@ -317,6 +322,8 @@ TEST_F(IslandTest, DenseSparse) { AsVector(data2->dof_island, nv)); EXPECT_EQ(AsVector(data1->island_dofind, nv), AsVector(data2->island_dofind, nv)); + EXPECT_EQ(AsVector(data1->dof_islandind, nv), + AsVector(data2->dof_islandind, nv)); EXPECT_EQ(AsVector(data1->island_efcadr, nisland), AsVector(data2->island_efcadr, nisland)); EXPECT_EQ(AsVector(data1->island_efcnum, nisland), @@ -331,5 +338,47 @@ TEST_F(IslandTest, DenseSparse) { mj_deleteModel(model); } +static const char* const kIlslandEfcPath = + "engine/testdata/island/island_efc.xml"; + +TEST_F(IslandTest, IslandEfc) { + const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath); + mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + mjData* data = mj_makeData(model); + + while (data->time < 0.3) { + 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, 2); + EXPECT_EQ(data->ne, 1); + EXPECT_EQ(data->nf, 1); + EXPECT_EQ(data->nl, 1); + EXPECT_EQ(nefc, 7); + EXPECT_THAT(AsVector(data->dof_island, nv), + ElementsAre(0, -1, 0, 0, 0, 1, 0)); + EXPECT_THAT(AsVector(data->island_dofnum, nisland), ElementsAre(5, 1)); + EXPECT_THAT(AsVector(data->island_dofadr, nisland), ElementsAre(0, 5)); + EXPECT_THAT(AsVector(data->island_dofind, nv), + ElementsAre(0, 2, 3, 4, 6, 5, -1)); + EXPECT_THAT(AsVector(data->dof_islandind, nv), + ElementsAre(0, -1, 1, 2, 3, 0, 4)); + EXPECT_THAT(AsVector(data->efc_island, nefc), + ElementsAre(0, 1, 0, 0, 0, 0, 0)); + EXPECT_THAT(AsVector(data->island_efcnum, nisland), ElementsAre(6, 1)); + EXPECT_THAT(AsVector(data->island_efcadr, nisland), ElementsAre(0, 6)); + EXPECT_THAT(AsVector(data->island_efcind, nefc), + ElementsAre(0, 2, 3, 4, 5, 6, 1)); + + mj_deleteData(data); + mj_deleteModel(model); +} + } // namespace } // namespace mujoco diff --git a/test/engine/testdata/island/island_efc.xml b/test/engine/testdata/island/island_efc.xml new file mode 100644 index 00000000..0be5e69e --- /dev/null +++ b/test/engine/testdata/island/island_efc.xml @@ -0,0 +1,53 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 25510728..eca6c7dd 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -1721,6 +1721,7 @@ public unsafe struct mjData_ { public int* island_dofnum; public int* island_dofadr; public int* island_dofind; + public int* dof_islandind; public int* efc_island; public int* island_efcnum; public int* island_efcadr;