diff --git a/doc/changelog.rst b/doc/changelog.rst index a5f8578c..0151a01f 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -12,6 +12,12 @@ General - Corrected the spelling of the ``inteval`` attribute to ``interval`` in the ``mjLROpt`` struct. - Mesh texture and normal mappings are now 3-per-triangle rather than 1-per-vertex. Mesh vertices are no longer duplicated in order to circumvent this limitation as they previously were. +- The non-zeros for the sparse constraint Jacobian matrix are now precounted and used for matrix memory allocation. + For instance, the constraint Jacobian matrix from the `humanoid100.xml + `_ model, which previously required + ~500,000 ``mjtNum``'s, now only requires ~6000. Very large models can now load and run with the CG solver. + + Python bindings ^^^^^^^^^^^^^^^ diff --git a/doc/includes/references.h b/doc/includes/references.h index 312dfcda..9a318e10 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -128,6 +128,7 @@ struct mjData_ { int ne; // number of equality constraints int nf; // number of friction constraints int nefc; // number of constraints + int nnzJ; // number of non-zeros in constraint Jacobian int ncon; // number of detected contacts // global properties @@ -284,39 +285,39 @@ struct mjData_ { //-------------------------------- ARENA-ALLOCATED ARRAYS // computed by mj_collision - mjContact* contact; // list of all detected contacts (ncon x 1) + mjContact* contact; // list of all detected contacts (ncon x 1) // computed by mj_makeConstraint - int* efc_type; // constraint type (mjtConstraint) (nefc x 1) - int* efc_id; // id of object of specified type (nefc x 1) - int* efc_J_rownnz; // number of non-zeros in Jacobian row (nefc x 1) - int* efc_J_rowadr; // row start address in colind array (nefc x 1) - int* efc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1) - int* efc_J_colind; // column indices in Jacobian (nefc x nv) - int* efc_JT_rownnz; // number of non-zeros in Jacobian row T (nv x 1) - int* efc_JT_rowadr; // row start address in colind array T (nv x 1) - int* efc_JT_rowsuper; // number of subsequent rows in supernode T (nv x 1) - int* efc_JT_colind; // column indices in Jacobian T (nv x nefc) - mjtNum* efc_J; // constraint Jacobian (nefc x nv) - mjtNum* efc_JT; // constraint Jacobian transposed (nv x nefc) - mjtNum* efc_pos; // constraint position (equality, contact) (nefc x 1) - mjtNum* efc_margin; // inclusion margin (contact) (nefc x 1) - mjtNum* efc_frictionloss; // frictionloss (friction) (nefc x 1) - mjtNum* efc_diagApprox; // approximation to diagonal of A (nefc x 1) - mjtNum* efc_KBIP; // stiffness, damping, impedance, imp' (nefc x 4) - mjtNum* efc_D; // constraint mass (nefc x 1) - mjtNum* efc_R; // inverse constraint mass (nefc x 1) + int* efc_type; // constraint type (mjtConstraint) (nefc x 1) + int* efc_id; // id of object of specified type (nefc x 1) + int* efc_J_rownnz; // number of non-zeros in constraint Jacobian row (nefc x 1) + int* efc_J_rowadr; // row start address in colind array (nefc x 1) + int* efc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1) + int* efc_J_colind; // column indices in constraint Jacobian (nnzJ x 1) + int* efc_JT_rownnz; // number of non-zeros in constraint Jacobian row T (nv x 1) + int* efc_JT_rowadr; // row start address in colind array T (nv x 1) + int* efc_JT_rowsuper; // number of subsequent rows in supernode T (nv x 1) + int* efc_JT_colind; // column indices in constraint Jacobian T (nnzJ x 1) + mjtNum* efc_J; // constraint Jacobian (nnzJ x 1) + mjtNum* efc_JT; // constraint Jacobian transposed (nnzJ x 1) + mjtNum* efc_pos; // constraint position (equality, contact) (nefc x 1) + mjtNum* efc_margin; // inclusion margin (contact) (nefc x 1) + mjtNum* efc_frictionloss; // frictionloss (friction) (nefc x 1) + mjtNum* efc_diagApprox; // approximation to diagonal of A (nefc x 1) + mjtNum* efc_KBIP; // stiffness, damping, impedance, imp' (nefc x 4) + mjtNum* efc_D; // constraint mass (nefc x 1) + mjtNum* efc_R; // inverse constraint mass (nefc x 1) // computed by mj_fwdConstraint/mj_inverse - mjtNum* efc_b; // linear cost term: J*qacc_smooth - aref (nefc x 1) - mjtNum* efc_force; // constraint force in constraint space (nefc x 1) - int* efc_state; // constraint state (mjtConstraintState) (nefc x 1) + mjtNum* efc_b; // linear cost term: J*qacc_smooth - aref (nefc x 1) + mjtNum* efc_force; // constraint force in constraint space (nefc x 1) + int* efc_state; // constraint state (mjtConstraintState) (nefc x 1) // computed by mj_projectConstraint - 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_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) }; typedef struct mjData_ mjData; typedef enum mjtDisableBit_ { // disable default feature bitflags diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 858d0b1e..66e96662 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -153,6 +153,7 @@ struct mjData_ { int ne; // number of equality constraints int nf; // number of friction constraints int nefc; // number of constraints + int nnzJ; // number of non-zeros in constraint Jacobian int ncon; // number of detected contacts // global properties @@ -309,39 +310,39 @@ struct mjData_ { //-------------------------------- ARENA-ALLOCATED ARRAYS // computed by mj_collision - mjContact* contact; // list of all detected contacts (ncon x 1) + mjContact* contact; // list of all detected contacts (ncon x 1) // computed by mj_makeConstraint - int* efc_type; // constraint type (mjtConstraint) (nefc x 1) - int* efc_id; // id of object of specified type (nefc x 1) - int* efc_J_rownnz; // number of non-zeros in Jacobian row (nefc x 1) - int* efc_J_rowadr; // row start address in colind array (nefc x 1) - int* efc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1) - int* efc_J_colind; // column indices in Jacobian (nefc x nv) - int* efc_JT_rownnz; // number of non-zeros in Jacobian row T (nv x 1) - int* efc_JT_rowadr; // row start address in colind array T (nv x 1) - int* efc_JT_rowsuper; // number of subsequent rows in supernode T (nv x 1) - int* efc_JT_colind; // column indices in Jacobian T (nv x nefc) - mjtNum* efc_J; // constraint Jacobian (nefc x nv) - mjtNum* efc_JT; // constraint Jacobian transposed (nv x nefc) - mjtNum* efc_pos; // constraint position (equality, contact) (nefc x 1) - mjtNum* efc_margin; // inclusion margin (contact) (nefc x 1) - mjtNum* efc_frictionloss; // frictionloss (friction) (nefc x 1) - mjtNum* efc_diagApprox; // approximation to diagonal of A (nefc x 1) - mjtNum* efc_KBIP; // stiffness, damping, impedance, imp' (nefc x 4) - mjtNum* efc_D; // constraint mass (nefc x 1) - mjtNum* efc_R; // inverse constraint mass (nefc x 1) + int* efc_type; // constraint type (mjtConstraint) (nefc x 1) + int* efc_id; // id of object of specified type (nefc x 1) + int* efc_J_rownnz; // number of non-zeros in constraint Jacobian row (nefc x 1) + int* efc_J_rowadr; // row start address in colind array (nefc x 1) + int* efc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1) + int* efc_J_colind; // column indices in constraint Jacobian (nnzJ x 1) + int* efc_JT_rownnz; // number of non-zeros in constraint Jacobian row T (nv x 1) + int* efc_JT_rowadr; // row start address in colind array T (nv x 1) + int* efc_JT_rowsuper; // number of subsequent rows in supernode T (nv x 1) + int* efc_JT_colind; // column indices in constraint Jacobian T (nnzJ x 1) + mjtNum* efc_J; // constraint Jacobian (nnzJ x 1) + mjtNum* efc_JT; // constraint Jacobian transposed (nnzJ x 1) + mjtNum* efc_pos; // constraint position (equality, contact) (nefc x 1) + mjtNum* efc_margin; // inclusion margin (contact) (nefc x 1) + mjtNum* efc_frictionloss; // frictionloss (friction) (nefc x 1) + mjtNum* efc_diagApprox; // approximation to diagonal of A (nefc x 1) + mjtNum* efc_KBIP; // stiffness, damping, impedance, imp' (nefc x 4) + mjtNum* efc_D; // constraint mass (nefc x 1) + mjtNum* efc_R; // inverse constraint mass (nefc x 1) // computed by mj_fwdConstraint/mj_inverse - mjtNum* efc_b; // linear cost term: J*qacc_smooth - aref (nefc x 1) - mjtNum* efc_force; // constraint force in constraint space (nefc x 1) - int* efc_state; // constraint state (mjtConstraintState) (nefc x 1) + mjtNum* efc_b; // linear cost term: J*qacc_smooth - aref (nefc x 1) + mjtNum* efc_force; // constraint force in constraint space (nefc x 1) + int* efc_state; // constraint state (mjtConstraintState) (nefc x 1) // computed by mj_projectConstraint - 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_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) }; typedef struct mjData_ mjData; diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index c0ce2d2b..04ec723f 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -536,31 +536,31 @@ X( mjContact, contact, MJ_D(ncon), 1 ) // array fields of mjData that are used in the primal problem -#define MJDATA_ARENA_POINTERS_PRIMAL \ - 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 ) \ - X( int, efc_J_rowadr, MJ_D(nefc), 1 ) \ - X( int, efc_J_rowsuper, MJ_D(nefc), 1 ) \ - X( int, efc_J_colind, MJ_D(nefc), MJ_M(nv) ) \ - X( int, efc_JT_rownnz, MJ_M(nv), 1 ) \ - X( int, efc_JT_rowadr, MJ_M(nv), 1 ) \ - X( int, efc_JT_rowsuper, MJ_M(nv), 1 ) \ - X( int, efc_JT_colind, MJ_M(nv), MJ_D(nefc) ) \ - X( mjtNum, efc_J, MJ_D(nefc), MJ_M(nv) ) \ - X( mjtNum, efc_JT, MJ_M(nv), MJ_D(nefc) ) \ - X( mjtNum, efc_pos, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_margin, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_frictionloss, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_diagApprox, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_KBIP, MJ_D(nefc), 4 ) \ - X( mjtNum, efc_D, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_R, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_vel, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_aref, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_b, MJ_D(nefc), 1 ) \ - X( mjtNum, efc_force, MJ_D(nefc), 1 ) \ - X( int, efc_state, MJ_D(nefc), 1 ) \ +#define MJDATA_ARENA_POINTERS_PRIMAL \ + 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) \ + X(int, efc_J_rowadr, MJ_D(nefc), 1) \ + X(int, efc_J_rowsuper, MJ_D(nefc), 1) \ + X(int, efc_J_colind, MJ_D(nnzJ), 1) \ + X(int, efc_JT_rownnz, MJ_M(nv), 1) \ + X(int, efc_JT_rowadr, MJ_M(nv), 1) \ + X(int, efc_JT_rowsuper, MJ_M(nv), 1) \ + X(int, efc_JT_colind, MJ_D(nnzJ), 1) \ + X(mjtNum, efc_J, MJ_D(nnzJ), 1) \ + X(mjtNum, efc_JT, MJ_D(nnzJ), 1) \ + X(mjtNum, efc_pos, MJ_D(nefc), 1) \ + X(mjtNum, efc_margin, MJ_D(nefc), 1) \ + X(mjtNum, efc_frictionloss, MJ_D(nefc), 1) \ + X(mjtNum, efc_diagApprox, MJ_D(nefc), 1) \ + X(mjtNum, efc_KBIP, MJ_D(nefc), 4) \ + X(mjtNum, efc_D, MJ_D(nefc), 1) \ + X(mjtNum, efc_R, MJ_D(nefc), 1) \ + X(mjtNum, efc_vel, MJ_D(nefc), 1) \ + X(mjtNum, efc_aref, MJ_D(nefc), 1) \ + X(mjtNum, efc_b, MJ_D(nefc), 1) \ + 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_DUAL \ diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc index 7e16bade..c10a29b6 100644 --- a/python/mujoco/structs.cc +++ b/python/mujoco/structs.cc @@ -707,6 +707,7 @@ void MjDataWrapper::Serialize(std::ostream& output) const { X(warning); X(ne); X(nf); + X(nnzJ); X(nefc); X(ncon); X(time); @@ -797,6 +798,7 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) { X(warning); X(ne); X(nf); + X(nnzJ); X(nefc); X(ncon); X(time); diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 39517228..7a47b00a 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -13,8 +13,8 @@ // limitations under the License. #include "engine/engine_core_constraint.h" -#include +#include #include #include @@ -1251,64 +1251,186 @@ void mj_makeImpedance(const mjModel* m, mjData* d) { //------------------------------------- constraint counting ---------------------------------------- -// count equality constraints -static inline int mj_ne(const mjModel* m, const mjData* d) { +// count the number of non-zeros in the sum of two sparse vectors +static int mju_combineSparseCount(int a_nnz, int b_nnz, const int* a_ind, const int* b_ind) { + int c_nnz, d_nnz; + const int* c_ind; + const int* d_ind; + + // choose c to have the least number of non-zeros + if (b_nnzc_ind[i]) { + nnz++; + } + } + + return nnz; +} + + + +// count the non-zero columns in the Jacobian difference of two bodies +static int mj_jacDifPairCount(const mjModel* m, int* chain, int b1, int b2) { + if (!m->nv) { + return 0; + } + + if (m->body_simple[b1] && m->body_simple[b2]) { + return mj_mergeChainSimple(m, chain, b1, b2); + } + return mj_mergeChain(m, chain, b1, b2); +} + + + +// return number of constraint non-zeros, handle dense and dof-less cases +static inline int mj_addConstraintCount(const mjModel* m, int size, int NV) { + // over count for dense allocation + if (!mj_isSparse(m)) { + return m->nv ? size : 0; + } + return mjMAX(0, NV) ? size : 0; +} + + + +// count equality constraints, count Jacobian nonzeros if nnz is not NULL +static inline int mj_ne(const mjModel* m, mjData* d, int* nnz) { + int ne = 0, nnze = 0; + int nv = m->nv, neq = m->neq; + int id[2], size, NV, NV2, *chain = NULL, *chain2 = NULL; + // disabled or no equality constraints: return if (mjDISABLED(mjDSBL_EQUALITY) || m->nemax==0) { return 0; } - int ne = 0; + mjMARKSTACK; - for (int i=0; ineq; i++) { - if (!m->eq_active[i]) { - continue; - } + if (nnz) { + chain = (int*)mj_stackAlloc(d, nv); + chain2 = (int*)mj_stackAlloc(d, nv); + } - // process according to type - switch (m->eq_type[i]) { - case mjEQ_CONNECT: - ne += 3; - break; + // find active equality constraints + for (int i=0; ieq_active[i]) { + id[0] = m->eq_obj1id[i]; + id[1] = m->eq_obj2id[i]; + size = 0; + NV = 0; + NV2 = 0; - case mjEQ_WELD: - ne += 6; - break; + // process according to type + switch (m->eq_type[i]) { + case mjEQ_CONNECT: + size = 3; + if (!nnz) { + break; + } - case mjEQ_JOINT: - case mjEQ_TENDON: - ne++; - break; + NV = mj_jacDifPairCount(m, chain, id[1], id[0]); + break; - default: // SHOULD NOT OCCUR - mju_error_i("Invalid equality constraint type %d", m->eq_type[i]); + case mjEQ_WELD: + size = 6; + if (!nnz) { + break; + } + + NV = mj_jacDifPairCount(m, chain, id[1], id[0]); + break; + + case mjEQ_JOINT: + case mjEQ_TENDON: + size = 1; + if (!nnz) { + break; + } + + for (int j=0; j<1+(id[1]>=0); j++) { + if (m->eq_type[i]==mjEQ_JOINT) { + if (!j) { + NV = 1; + chain[0] = m->jnt_dofadr[id[j]]; + } else { + NV2 = 1; + chain2[0] = m->jnt_dofadr[id[j]]; + } + } else { + if (!j) { + NV = d->ten_J_rownnz[id[j]]; + memcpy(chain, d->ten_J_colind+d->ten_J_rowadr[id[j]], NV*sizeof(int)); + } else { + NV2 = d->ten_J_rownnz[id[j]]; + memcpy(chain2, d->ten_J_colind+d->ten_J_rowadr[id[j]], NV2*sizeof(int)); + } + } + } + + if (id[1]>=0) { + NV = mju_combineSparseCount(NV, NV2, chain, chain2); + NV = 2; + } + break; + } + ne += mj_addConstraintCount(m, size, NV); + nnze += size*NV; } } + if (nnz) { + *nnz += nnze; + } + + mjFREESTACK; return ne; } -// count frictional constraints -static inline int mj_nf(const mjModel* m, const mjData* d) { - // disabled: return +// count frictional constraints, count Jacobian nonzeros if nnz is not NULL +static inline int mj_nf(const mjModel* m, const mjData* d, int *nnz) { + int nf = 0, nnzf = 0; + int nv = m->nv, ntendon = m->ntendon; + if (mjDISABLED(mjDSBL_FRICTIONLOSS)) { return 0; } - int nf = 0; - const int nv = m->nv; - const int ntendon = m->ntendon; - - // count frictional dofs for (int i=0; idof_frictionloss[i] > 0); + if (m->dof_frictionloss[i]>0) { + nf += mj_addConstraintCount(m, 1, 1); + nnzf++; + } } - // count frictional tendons for (int i=0; itendon_frictionloss[i] > 0); + if (m->tendon_frictionloss[i]>0) { + nf += mj_addConstraintCount(m, 1, d->ten_J_rownnz[i]); + nnzf += d->ten_J_rownnz[i]; + } + } + + if (nnz) { + *nnz += nnzf; } return nf; @@ -1316,110 +1438,145 @@ static inline int mj_nf(const mjModel* m, const mjData* d) { -// count limit constraints -static inline int mj_nl(const mjModel* m, const mjData* d) { +// count limit constraints, count Jacobian nonzeros if nnz is not NULL +static inline int mj_nl(const mjModel* m, const mjData* d, int *nnz) { + int nnzl = 0, nl = 0; + int ntendon = m->ntendon; + int side; + mjtNum margin, value, dist; + // disabled: return if (mjDISABLED(mjDSBL_LIMIT)) { return 0; } - int nl = 0; - const int njnt = m->njnt; - const int ntendon = m->ntendon; - // count limited joints - for (int i=0; injnt; i++) { if (!m->jnt_limited[i]) { continue; } - // slides and hinges can have active limits on two sides, check both + margin = m->jnt_margin[i]; + + // slider and hinge joint limits can be bilateral, check both side if (m->jnt_type[i]==mjJNT_SLIDE || m->jnt_type[i]==mjJNT_HINGE) { - // get margin - mjtNum margin = m->jnt_margin[i]; - - // get joint value - mjtNum value = d->qpos[m->jnt_qposadr[i]]; - - // check lower and upper limits - for (int side=-1; side<=1; side+=2) { - // compute distance (negative: penetration) - mjtNum dist = side * (m->jnt_range[2*i+(side+1)/2] - value); - - // detect joint limit + value = d->qpos[m->jnt_qposadr[i]]; + for (side=-1; side<=1; side+=2) { + dist = side * (m->jnt_range[2*i+(side+1)/2] - value); if (distjnt_type[i]==mjJNT_BALL) { + mjtNum angleAxis[3]; + mju_quat2Vel(angleAxis, d->qpos+m->jnt_qposadr[i], 1); + value = mju_normalize3(angleAxis); + dist = mju_max(m->jnt_range[2*i], m->jnt_range[2*i+1]) - value; + if (disttendon_limited[i]; + if (m->tendon_limited[i]) { + value = d->ten_length[i]; + margin = m->tendon_margin[i]; + + // tendon limits can be bilateral, check both sides + for (side=-1; side<=1; side+=2) { + dist = side * (m->tendon_range[2*i+(side+1)/2] - value); + if (distten_J_rownnz[i]); + nnzl += d->ten_J_rownnz[i]; + } + } + } } + if (nnz) { + *nnz += nnzl; + } return nl; } -// count contact constraints -static inline int mj_nc(const mjModel* m, const mjData* d) { - // disabled or no contacts: return - int ncon = d->ncon; - if (mjDISABLED(mjDSBL_CONTACT) || ncon==0) { +// count contact constraints, count Jacobian nonzeros if nnz is not NULL +static inline int mj_nc(const mjModel* m, mjData* d, int* nnz) { + int nnzc = 0, nc = 0; + int ispyramid = mj_isPyramidal(m), ncon = d->ncon; + + if (mjDISABLED(mjDSBL_CONTACT) || !ncon) { return 0; } - int nc = 0; - int ispyramid = mj_isPyramidal(m); + mjMARKSTACK; + int *chain = (int*)mj_stackAlloc(d, m->nv); - // find contacts to be counted for (int i=0; icontact + i; - if (con->exclude) { + if (d->contact[i].exclude) { continue; } + mjContact* con = d->contact + i; int dim = con->dim; - - // dim 1: single constraint - if (dim==1) { - nc++; + int b1 = m->geom_bodyid[con->geom1]; + int b2 = m->geom_bodyid[con->geom2]; + int NV = mj_jacDifPairCount(m, chain, b1, b2); + if (!NV) { + continue; } - // dim > 1: depends on cone type - else { - nc += (ispyramid ? 2*(dim-1) : dim); + if (dim==1) { + nc++; + nnzc += NV; + } else if (ispyramid) { + nc += 2*(dim-1); + nnzc += 2*(dim-1)*NV; + } else { + nc += dim; + nnzc += dim*NV; } } + if (nnz) { + *nnz += nnzc; + } + + mjFREESTACK; return nc; } -// count all constraints -static inline int mj_nefc(const mjModel* m, const mjData* d) { - return mj_ne(m, d) + mj_nf(m, d) + mj_nl(m, d) + mj_nc(m, d); -} - //---------------------------- top-level API for constraint construction --------------------------- + + // driver: call all functions above void mj_makeConstraint(const mjModel* m, mjData* d) { // clear sizes - d->ne = d->nf = d->nefc = 0; + d->ne = d->nf = d->nefc = d->nnzJ = 0; // disabled or Jacobian not allocated: return if (mjDISABLED(mjDSBL_CONSTRAINT)) { return; } - int nefc_allocated = mj_nefc(m, d); + // precount sizes for constraint Jacobian matrices + int *nnz = mj_isSparse(m) ? &(d->nnzJ) : NULL; + + int ne_allocated = mj_ne(m, d, nnz); + int nf_allocated = mj_nf(m, d, nnz); + + int nefc_allocated = ne_allocated + nf_allocated + mj_nl(m, d, nnz) + mj_nc(m, d, nnz); + if (!mj_isSparse(m)) { + d->nnzJ = nefc_allocated * m->nv; + } d->nefc = nefc_allocated; #undef MJ_M @@ -1451,16 +1608,61 @@ void mj_makeConstraint(const mjModel* m, mjData* d) { #undef MJ_D #define MJ_D(n) n + // reset nefc for the instantiation functions, + // and instantiate all elements of Jacobian d->nefc = 0; - - // instantiate all elements of Jacobian mj_instantiateEquality(m, d); mj_instantiateFriction(m, d); mj_instantiateLimit(m, d); mj_instantiateContact(m, d); - if (d->nefc > nefc_allocated) { + + // check sparse allocation + if (mj_isSparse(m)) { + if (d->ne != ne_allocated) { + char msg[1024]; + + // TODO(b/270530821): add var argument support to mju_error + mjSNPRINTF( + msg, "ne mis-allocation: found ne=%d but allocated %d", d->ne, ne_allocated); + mju_error(msg); + } + + if (d->nf != nf_allocated) { + char msg[1024]; + + // TODO(b/270530821): add var argument support to mju_error + mjSNPRINTF( + msg, "nf mis-allocation: found nf=%d but allocated %d", d->nf, nf_allocated); + mju_error(msg); + } + + // check that nefc was computed correctly + if (d->nefc != nefc_allocated) { + char msg[1024]; + + // TODO(b/270530821): add var argument support to mju_error + mjSNPRINTF( + msg, "nefc mis-allocation: found nefc=%d but allocated %d", d->nefc, nefc_allocated); + mju_error(msg); + } + + // check that nnzJ was computed correctly + if (d->nefc > 0) { + int nnz = d->efc_J_rownnz[d->nefc - 1] + d->efc_J_rowadr[d->nefc - 1]; + if (d->nnzJ != nnz) { + char msg[1024]; + + // TODO(b/270530821): add var argument support to mju_error + mjSNPRINTF( + msg, "constraint Jacobian mis-allocation: found nnzJ=%d but allocated %d", nnz, d->nnzJ); + mju_error(msg); + } + } + } else if (d->nefc > nefc_allocated) { char msg[1024]; + + // TODO(b/270530821): add var argument support to mju_error mjSNPRINTF( msg, "nefc under-allocation: found nefc=%d but allocated only %d", d->nefc, nefc_allocated); mju_error(msg); diff --git a/test/engine/engine_core_constraint_test.cc b/test/engine/engine_core_constraint_test.cc index f5623f29..2b61667f 100644 --- a/test/engine/engine_core_constraint_test.cc +++ b/test/engine/engine_core_constraint_test.cc @@ -27,8 +27,8 @@ namespace mujoco { namespace { -using ::testing::Pointwise; using ::testing::DoubleNear; +using ::testing::Pointwise; using CoreConstraintTest = MujocoTest; std::vector AsVector(const mjtNum* array, int n) { @@ -158,5 +158,42 @@ TEST_F(CoreConstraintTest, WeldRotJacobian) { mj_deleteModel(model); } +static const char* const kDoflessContactPath = + "engine/testdata/core_constraint/dofless_contact.xml"; +static const char* const kDoflessTendonFrictionalPath = + "engine/testdata/core_constraint/dofless_tendon_frictional.xml"; +static const char* const kDoflessTendonLimitedPath = + "engine/testdata/core_constraint/dofless_tendon_limited.xml"; +static const char* const kDoflessTendonLimitedMarginPath = + "engine/testdata/core_constraint/dofless_tendon_limitedmargin.xml"; +static const char* const kDoflessWeldPath = + "engine/testdata/core_constraint/dofless_weld.xml"; +static const char* const kJointLimitedBilateralMarginPath = + "engine/testdata/core_constraint/joint_limited_bilateral_margin.xml"; +static const char* const kTendonLimitedBilateralMarginPath = + "engine/testdata/core_constraint/tendon_limited_bilateral_margin.xml"; + +TEST_F(CoreConstraintTest, JacobianPreAllocate) { + for (const char* local_path : + {kDoflessContactPath, kDoflessTendonFrictionalPath, + kDoflessTendonLimitedPath, kDoflessTendonLimitedMarginPath, + kDoflessWeldPath, kJointLimitedBilateralMarginPath, + kTendonLimitedBilateralMarginPath}) { + const std::string xml_path = GetTestDataFilePath(local_path); + + // iterate through dense and sparse + for (mjtJacobian sparsity : {mjJAC_DENSE, mjJAC_SPARSE}) { + mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + model->opt.jacobian = sparsity; + mjData* data = mj_makeData(model); + + mj_step(model, data); + + mj_deleteData(data); + mj_deleteModel(model); + } + } +} + } // namespace } // namespace mujoco diff --git a/test/engine/testdata/core_constraint/dofless_contact.xml b/test/engine/testdata/core_constraint/dofless_contact.xml new file mode 100644 index 00000000..6a033946 --- /dev/null +++ b/test/engine/testdata/core_constraint/dofless_contact.xml @@ -0,0 +1,18 @@ + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_constraint/dofless_tendon_frictional.xml b/test/engine/testdata/core_constraint/dofless_tendon_frictional.xml new file mode 100644 index 00000000..15e856fb --- /dev/null +++ b/test/engine/testdata/core_constraint/dofless_tendon_frictional.xml @@ -0,0 +1,13 @@ + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_constraint/dofless_tendon_limited.xml b/test/engine/testdata/core_constraint/dofless_tendon_limited.xml new file mode 100644 index 00000000..45972208 --- /dev/null +++ b/test/engine/testdata/core_constraint/dofless_tendon_limited.xml @@ -0,0 +1,13 @@ + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_constraint/dofless_tendon_limitedmargin.xml b/test/engine/testdata/core_constraint/dofless_tendon_limitedmargin.xml new file mode 100644 index 00000000..86d1caee --- /dev/null +++ b/test/engine/testdata/core_constraint/dofless_tendon_limitedmargin.xml @@ -0,0 +1,13 @@ + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_constraint/dofless_weld.xml b/test/engine/testdata/core_constraint/dofless_weld.xml new file mode 100644 index 00000000..c3b14169 --- /dev/null +++ b/test/engine/testdata/core_constraint/dofless_weld.xml @@ -0,0 +1,24 @@ + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/core_constraint/joint_limited_bilateral_margin.xml b/test/engine/testdata/core_constraint/joint_limited_bilateral_margin.xml new file mode 100644 index 00000000..df6bed0a --- /dev/null +++ b/test/engine/testdata/core_constraint/joint_limited_bilateral_margin.xml @@ -0,0 +1,8 @@ + + + + + + + + diff --git a/test/engine/testdata/core_constraint/tendon_limited_bilateral_margin.xml b/test/engine/testdata/core_constraint/tendon_limited_bilateral_margin.xml new file mode 100644 index 00000000..fa9b09cd --- /dev/null +++ b/test/engine/testdata/core_constraint/tendon_limited_bilateral_margin.xml @@ -0,0 +1,17 @@ + + + + + + + + + + + + + + + + + diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index a866d56a..9fdfe589 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -1587,6 +1587,7 @@ public unsafe struct mjData_ { public int ne; public int nf; public int nefc; + public int nnzJ; public int ncon; public double time; public fixed double energy[2];