From 2409a6e5378cf06b3cf01de0cfaf264a365930fb Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Fri, 10 Jan 2025 09:22:58 -0800 Subject: [PATCH] Add `mjData.nA` the number of non-zeros in the constraint inverse inertia matrix. PiperOrigin-RevId: 714076988 Change-Id: I4c323a6e81dfb23f7972b662719bbcb688bacf4a --- doc/includes/references.h | 5 +++-- include/mujoco/mjdata.h | 5 +++-- include/mujoco/mjxmacro.h | 5 +++-- introspect/structs.py | 9 +++++++-- python/mujoco/structs.cc | 2 ++ src/engine/engine_core_constraint.c | 3 ++- src/engine/engine_io.c | 1 + test/engine/engine_io_test.cc | 1 + unity/Runtime/Bindings/MjBindings.cs | 1 + 9 files changed, 23 insertions(+), 9 deletions(-) diff --git a/doc/includes/references.h b/doc/includes/references.h index e1432bd2..0303015e 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -170,6 +170,7 @@ struct mjData_ { int nl; // number of limit constraints int nefc; // number of constraints int nJ; // number of non-zeros in constraint Jacobian + int nA; // number of non-zeros in constraint inverse inertia matrix int nisland; // number of detected constraint islands // global properties @@ -391,8 +392,8 @@ struct mjData_ { // computed by mj_projectConstraint (PGS solver) int* efc_AR_rownnz; // number of non-zeros in AR (nefc x 1) int* efc_AR_rowadr; // row start address in colind array (nefc x 1) - int* efc_AR_colind; // column indices in sparse AR (nefc x nefc) - mjtNum* efc_AR; // J*inv(M)*J' + R (nefc x nefc) + int* efc_AR_colind; // column indices in sparse AR (nA x 1) + mjtNum* efc_AR; // J*inv(M)*J' + R (nA x 1) //-------------------- arena-allocated: POSITION, VELOCITY dependent diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index b51b19ae..3c4d889e 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -198,6 +198,7 @@ struct mjData_ { int nl; // number of limit constraints int nefc; // number of constraints int nJ; // number of non-zeros in constraint Jacobian + int nA; // number of non-zeros in constraint inverse inertia matrix int nisland; // number of detected constraint islands // global properties @@ -419,8 +420,8 @@ struct mjData_ { // computed by mj_projectConstraint (PGS solver) int* efc_AR_rownnz; // number of non-zeros in AR (nefc x 1) int* efc_AR_rowadr; // row start address in colind array (nefc x 1) - int* efc_AR_colind; // column indices in sparse AR (nefc x nefc) - mjtNum* efc_AR; // J*inv(M)*J' + R (nefc x nefc) + int* efc_AR_colind; // column indices in sparse AR (nA x 1) + mjtNum* efc_AR; // J*inv(M)*J' + R (nA x 1) //-------------------- arena-allocated: POSITION, VELOCITY dependent diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 2d78fc02..6de6198d 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -715,8 +715,8 @@ #define MJDATA_ARENA_POINTERS_DUAL \ X( int, efc_AR_rownnz, MJ_D(nefc), 1 ) \ X( int, efc_AR_rowadr, MJ_D(nefc), 1 ) \ - X( int, efc_AR_colind, MJ_D(nefc), MJ_D(nefc) ) \ - X( mjtNum, efc_AR, MJ_D(nefc), MJ_D(nefc) ) + X( int, efc_AR_colind, MJ_D(nA), 1 ) \ + X( mjtNum, efc_AR, MJ_D(nA), 1 ) // array fields of mjData that are used for constraint islands #define MJDATA_ARENA_POINTERS_ISLAND \ @@ -757,6 +757,7 @@ X( int, nl ) \ X( int, nefc ) \ X( int, nJ ) \ + X( int, nA ) \ X( int, nisland ) \ X( mjtNum, time ) \ X( uintptr_t, threadpool ) diff --git a/introspect/structs.py b/introspect/structs.py index be60293c..f988d597 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -4701,6 +4701,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=ValueType(name='int'), doc='number of non-zeros in constraint Jacobian', ), + StructFieldDecl( + name='nA', + type=ValueType(name='int'), + doc='number of non-zeros in constraint inverse inertia matrix', + ), StructFieldDecl( name='nisland', type=ValueType(name='int'), @@ -5797,7 +5802,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ inner_type=ValueType(name='int'), ), doc='column indices in sparse AR', - array_extent=('nefc', 'nefc'), + array_extent=('nA',), ), StructFieldDecl( name='efc_AR', @@ -5805,7 +5810,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ inner_type=ValueType(name='mjtNum'), ), doc="J*inv(M)*J' + R", - array_extent=('nefc', 'nefc'), + array_extent=('nA',), ), StructFieldDecl( name='efc_vel', diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc index 53d1a72b..f3fa7a0b 100644 --- a/python/mujoco/structs.cc +++ b/python/mujoco/structs.cc @@ -748,6 +748,7 @@ void MjDataWrapper::Serialize(std::ostream& output) const { X(ne); X(nf); X(nJ); + X(nA); X(nefc); X(nisland); X(time); @@ -825,6 +826,7 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) { X(ne); X(nf); X(nJ); + X(nA); X(nefc); X(nisland); X(time); diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 02efb04f..a65a467e 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -1943,7 +1943,7 @@ static int mj_nc(const mjModel* m, mjData* d, int* nnz) { // driver: call all functions above void mj_makeConstraint(const mjModel* m, mjData* d) { // clear sizes - d->ne = d->nf = d->nl = d->nefc = d->nJ = 0; + d->ne = d->nf = d->nl = d->nefc = d->nJ = d->nA = 0; // disabled or Jacobian not allocated: return if (mjDISABLED(mjDSBL_CONSTRAINT)) { @@ -1960,6 +1960,7 @@ void mj_makeConstraint(const mjModel* m, mjData* d) { d->nJ = nefc_allocated * m->nv; } d->nefc = nefc_allocated; + d->nA = d->nefc * d->nefc; // allocate efc arrays on arena if (!arenaAllocEfc(m, d)) { diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index e32a9add..ddbb3557 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -1893,6 +1893,7 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { d->nl = 0; d->nefc = 0; d->nJ = 0; + d->nA = 0; d->nisland = 0; // clear global properties diff --git a/test/engine/engine_io_test.cc b/test/engine/engine_io_test.cc index 3e6eb9af..470f6f02 100644 --- a/test/engine/engine_io_test.cc +++ b/test/engine/engine_io_test.cc @@ -157,6 +157,7 @@ TEST_F(EngineIoTest, ResetVariableSizes) { EXPECT_EQ(data->nf, 0); EXPECT_EQ(data->nefc, 0); EXPECT_EQ(data->nJ, 0); + EXPECT_EQ(data->nA, 0); EXPECT_EQ(data->ncon, 0); mj_deleteData(data); diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 7ad3d1b6..19c46aa0 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -4852,6 +4852,7 @@ public unsafe struct mjData_ { public int nl; public int nefc; public int nJ; + public int nA; public int nisland; public double time; public fixed double energy[2];