From b66da3d3691cdf6759782f943a1618233d03af08 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Tue, 8 Oct 2024 10:12:10 -0700 Subject: [PATCH] Add `mjData.L`, sparse Cholesky factor of Newton solver Hessian. PiperOrigin-RevId: 683669555 Change-Id: I18d7530b5d68d1d48cb95d4faa0150c434109638 --- doc/includes/references.h | 12 ++++++-- include/mujoco/mjdata.h | 12 ++++++-- include/mujoco/mjxmacro.h | 16 ++++++++-- introspect/structs.py | 44 ++++++++++++++++++++++++++-- python/mujoco/functions.cc | 4 +-- python/mujoco/structs.cc | 16 ++++++---- src/engine/engine_core_constraint.c | 5 ++-- src/engine/engine_io.c | 3 +- unity/Runtime/Bindings/MjBindings.cs | 8 ++++- 9 files changed, 99 insertions(+), 21 deletions(-) diff --git a/doc/includes/references.h b/doc/includes/references.h index de2d533b..8b4ad04e 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -164,12 +164,13 @@ struct mjData_ { mjTimerStat timer[mjNTIMER]; // timer statistics // variable sizes + int ncon; // number of detected contacts int ne; // number of equality constraints int nf; // number of friction constraints int nl; // number of limit constraints int nefc; // number of constraints int nnzJ; // number of non-zeros in constraint Jacobian - int ncon; // number of detected contacts + int nnzL; // number of non-zeros in Newton Cholesky factor int nisland; // number of detected constraint islands // global properties @@ -385,12 +386,19 @@ struct mjData_ { int* island_efcadr; // start address in island_efcind (nisland x 1) int* island_efcind; // island constraint indices (nefc x 1) - // computed by mj_projectConstraint (dual solver) + // 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) + // computed by mj_fwdConstraint (Newton solver) + int* L_rownnz; // number of non-zeros in Hessian factor L rows (nv x 1) + int* L_rowadr; // row start address in colind array (nv x 1) + int* L_colind; // column indices in sparse AR (nnzL x 1) + mjtNum* L; // chol(M + J'*diag(efc_D)*J) (nnzL x 1) + mjtNum* Lcone; // L with cone contributions (nnzL x 1) + //-------------------- arena-allocated: POSITION, VELOCITY dependent // computed by mj_fwdVelocity/mj_referenceConstraint diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index ed753914..54b3ae5f 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -192,12 +192,13 @@ struct mjData_ { mjTimerStat timer[mjNTIMER]; // timer statistics // variable sizes + int ncon; // number of detected contacts int ne; // number of equality constraints int nf; // number of friction constraints int nl; // number of limit constraints int nefc; // number of constraints int nnzJ; // number of non-zeros in constraint Jacobian - int ncon; // number of detected contacts + int nnzL; // number of non-zeros in Newton Cholesky factor int nisland; // number of detected constraint islands // global properties @@ -413,12 +414,19 @@ struct mjData_ { int* island_efcadr; // start address in island_efcind (nisland x 1) int* island_efcind; // island constraint indices (nefc x 1) - // computed by mj_projectConstraint (dual solver) + // 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) + // computed by mj_fwdConstraint (Newton solver) + int* L_rownnz; // number of non-zeros in Hessian factor L rows (nv x 1) + int* L_rowadr; // row start address in colind array (nv x 1) + int* L_colind; // column indices in sparse AR (nnzL x 1) + mjtNum* L; // chol(M + J'*diag(efc_D)*J) (nnzL x 1) + mjtNum* Lcone; // L with cone contributions (nnzL x 1) + //-------------------- arena-allocated: POSITION, VELOCITY dependent // computed by mj_fwdVelocity/mj_referenceConstraint diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 8bab84ad..c999b855 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -680,7 +680,7 @@ X( mjContact, contact, MJ_D(ncon), 1 ) // array fields of mjData that are used in the primal problem -#define MJDATA_ARENA_POINTERS_PRIMAL \ +#define MJDATA_ARENA_POINTERS_SOLVER \ X( int, efc_type, MJ_D(nefc), 1 ) \ X( int, efc_id, MJ_D(nefc), 1 ) \ X( int, efc_J_rownnz, MJ_D(nefc), 1 ) \ @@ -707,6 +707,14 @@ X( mjtNum, efc_force, MJ_D(nefc), 1 ) \ X( int, efc_state, MJ_D(nefc), 1 ) +// array fields of mjData that are used in the dual problem +#define MJDATA_ARENA_POINTERS_NEWTON \ + X( int, L_rownnz, MJ_M(nv), 1 ) \ + X( int, L_rowadr, MJ_M(nv), 1 ) \ + X( int, L_colind, MJ_D(nnzL), 1 ) \ + X( mjtNum, L, MJ_D(nnzL), 1 ) \ + X( mjtNum, Lcone, MJ_D(nnzL), 1 ) + // array fields of mjData that are used in the dual problem #define MJDATA_ARENA_POINTERS_DUAL \ X( int, efc_AR_rownnz, MJ_D(nefc), 1 ) \ @@ -729,7 +737,8 @@ // array fields of mjData that live in d->arena #define MJDATA_ARENA_POINTERS \ MJDATA_ARENA_POINTERS_CONTACT \ - MJDATA_ARENA_POINTERS_PRIMAL \ + MJDATA_ARENA_POINTERS_SOLVER \ + MJDATA_ARENA_POINTERS_NEWTON \ MJDATA_ARENA_POINTERS_DUAL \ MJDATA_ARENA_POINTERS_ISLAND @@ -747,12 +756,13 @@ X( int, maxuse_con ) \ X( int, maxuse_efc ) \ X( int, solver_nisland ) \ + X( int, ncon ) \ X( int, ne ) \ X( int, nf ) \ X( int, nl ) \ X( int, nefc ) \ X( int, nnzJ ) \ - X( int, ncon ) \ + X( int, nnzL ) \ X( int, nisland ) \ X( mjtNum, time ) \ X( uintptr_t, threadpool ) diff --git a/introspect/structs.py b/introspect/structs.py index 9455d2a7..c4076a03 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -4296,6 +4296,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), doc='timer statistics', ), + StructFieldDecl( + name='ncon', + type=ValueType(name='int'), + doc='number of detected contacts', + ), StructFieldDecl( name='ne', type=ValueType(name='int'), @@ -4322,9 +4327,9 @@ STRUCTS: Mapping[str, StructDecl] = dict([ doc='number of non-zeros in constraint Jacobian', ), StructFieldDecl( - name='ncon', + name='nnzL', type=ValueType(name='int'), - doc='number of detected contacts', + doc='number of non-zeros in Newton Cholesky factor', ), StructFieldDecl( name='nisland', @@ -5282,6 +5287,41 @@ STRUCTS: Mapping[str, StructDecl] = dict([ ), doc="J*inv(M)*J' + R (nefc x nefc)", # pylint: disable=line-too-long ), + StructFieldDecl( + name='L_rownnz', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='number of non-zeros in Hessian factor L rows (nv x 1)', # pylint: disable=line-too-long + ), + StructFieldDecl( + name='L_rowadr', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='row start address in colind array (nv x 1)', # pylint: disable=line-too-long + ), + StructFieldDecl( + name='L_colind', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='column indices in sparse AR (nnzL x 1)', # pylint: disable=line-too-long + ), + StructFieldDecl( + name='L', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc="chol(M + J'*diag(efc_D)*J) (nnzL x 1)", # pylint: disable=line-too-long + ), + StructFieldDecl( + name='Lcone', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc='L with cone contributions (nnzL x 1)', # pylint: disable=line-too-long + ), StructFieldDecl( name='efc_vel', type=PointerType( diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index a82005b2..1dd47eed 100644 --- a/python/mujoco/functions.cc +++ b/python/mujoco/functions.cc @@ -1441,7 +1441,7 @@ PYBIND11_MODULE(_functions, pymodule) { data->nefc = 0; data->contact = static_cast(data->arena); #define X(type, name, nr, nc) data->name = nullptr; - MJDATA_ARENA_POINTERS_PRIMAL + MJDATA_ARENA_POINTERS_SOLVER MJDATA_ARENA_POINTERS_DUAL #undef X }; @@ -1469,7 +1469,7 @@ PYBIND11_MODULE(_functions, pymodule) { throw FatalError("insufficient arena memory available"); \ } - MJDATA_ARENA_POINTERS_PRIMAL + MJDATA_ARENA_POINTERS_SOLVER if (mj_isDual(d.model().get())) { MJDATA_ARENA_POINTERS_DUAL } diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc index 93b9f014..3fe757a8 100644 --- a/python/mujoco/structs.cc +++ b/python/mujoco/structs.cc @@ -41,6 +41,7 @@ #include #include "errors.h" #include "function_traits.h" +#include "indexer_xmacro.h" #include "indexers.h" #include "private.h" #include "raw.h" @@ -743,11 +744,12 @@ void MjDataWrapper::Serialize(std::ostream& output) const { X(solver); X(timer); X(warning); + X(ncon); X(ne); X(nf); X(nnzJ); + X(nnzL); X(nefc); - X(ncon); X(nisland); X(time); X(energy); @@ -772,7 +774,8 @@ void MjDataWrapper::Serialize(std::ostream& output) const { } MJDATA_ARENA_POINTERS_CONTACT - MJDATA_ARENA_POINTERS_PRIMAL + MJDATA_ARENA_POINTERS_SOLVER + // MJDATA_ARENA_POINTERS_NEWTON // TODO: tassa - Add after allocation exists if (mj_isDual(this->model_->get())) { MJDATA_ARENA_POINTERS_DUAL } @@ -819,11 +822,12 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) { X(solver); X(timer); X(warning); + X(ncon); X(ne); X(nf); X(nnzJ); + X(nnzL); X(nefc); - X(ncon); X(nisland); X(time); X(energy); @@ -850,7 +854,8 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) { } MJDATA_ARENA_POINTERS_CONTACT - MJDATA_ARENA_POINTERS_PRIMAL + MJDATA_ARENA_POINTERS_SOLVER + // MJDATA_ARENA_POINTERS_NEWTON // TODO: tassa - Add after allocation exists if (is_dual) { MJDATA_ARENA_POINTERS_DUAL } @@ -2009,7 +2014,8 @@ This is useful for example when the MJB is not available as a file on disk.)")); return InitPyArray(X_ARRAY_SHAPE(dim0, dim1), d.get()->var, d.owner()); \ }); - MJDATA_ARENA_POINTERS_PRIMAL + MJDATA_ARENA_POINTERS_SOLVER + MJDATA_ARENA_POINTERS_NEWTON MJDATA_ARENA_POINTERS_DUAL MJDATA_ARENA_POINTERS_ISLAND diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 66c30541..8c775d1d 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -71,11 +71,10 @@ static int arenaAllocEfc(const mjModel* m, mjData* d) { return 0; \ } - MJDATA_ARENA_POINTERS_PRIMAL + MJDATA_ARENA_POINTERS_SOLVER if (mj_isDual(m)) { MJDATA_ARENA_POINTERS_DUAL } - #undef X #undef MJ_M @@ -1918,7 +1917,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->nnzJ = 0; + d->ne = d->nf = d->nl = d->nefc = d->nnzJ = d->nnzL = 0; // disabled or Jacobian not allocated: return if (mjDISABLED(mjDSBL_CONSTRAINT)) { diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index ef460544..9e721066 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -1768,12 +1768,13 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { mju_zero(d->solver_fwdinv, 2); // clear variable sizes + d->ncon = 0; d->ne = 0; d->nf = 0; d->nl = 0; d->nefc = 0; d->nnzJ = 0; - d->ncon = 0; + d->nnzL = 0; d->nisland = 0; // clear global properties diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index b2fdae2d..b3eefdbf 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -4841,12 +4841,13 @@ public unsafe struct mjData_ { public mjTimerStat_ timer12; public mjTimerStat_ timer13; public mjTimerStat_ timer14; + public int ncon; public int ne; public int nf; public int nl; public int nefc; public int nnzJ; - public int ncon; + public int nnzL; public int nisland; public double time; public fixed double energy[2]; @@ -4984,6 +4985,11 @@ public unsafe struct mjData_ { public int* efc_AR_rowadr; public int* efc_AR_colind; public double* efc_AR; + public int* L_rownnz; + public int* L_rowadr; + public int* L_colind; + public double* L; + public double* Lcone; public double* efc_vel; public double* efc_aref; public double* efc_b;