Precount number of non-zeros for constraint Jacobian sparse matrix.

PiperOrigin-RevId: 512947632
Change-Id: I9605e81b4e51fbc90854c13153e4fee38ad48fb9
This commit is contained in:
Kyle Bayes
2023-02-28 08:57:01 -08:00
committed by Copybara-Service
parent 90b5dd14ca
commit 50bebcb4ee
15 changed files with 521 additions and 165 deletions
+6
View File
@@ -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
<https://github.com/deepmind/mujoco/blob/main/model/humanoid100/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
^^^^^^^^^^^^^^^
+28 -27
View File
@@ -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
+28 -27
View File
@@ -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;
+25 -25
View File
@@ -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 \
+2
View File
@@ -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);
+287 -85
View File
@@ -13,8 +13,8 @@
// limitations under the License.
#include "engine/engine_core_constraint.h"
#include <stdio.h>
#include <stdio.h>
#include <stddef.h>
#include <string.h>
@@ -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_nnz<a_nnz) {
c_nnz = b_nnz;
c_ind = b_ind;
d_nnz = a_nnz;
d_ind = a_ind;
} else {
c_nnz = a_nnz;
c_ind = a_ind;
d_nnz = b_nnz;
d_ind = b_ind;
}
int nnz=d_nnz, j=0;
for (int i=0; i<c_nnz; i++) {
while (d_ind[j]<c_ind[i]) {
j++;
}
if (d_ind[j]>c_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; i<m->neq; 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; i<neq; i++) {
if (m->eq_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; i<nv; i++) {
nf += (m->dof_frictionloss[i] > 0);
if (m->dof_frictionloss[i]>0) {
nf += mj_addConstraintCount(m, 1, 1);
nnzf++;
}
}
// count frictional tendons
for (int i=0; i<ntendon; i++) {
nf += (m->tendon_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; i<njnt; i++) {
for (int i=0; i<m->njnt; 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 (dist<margin) {
nl++;
nl += mj_addConstraintCount(m, 1, 1);
nnzl++;
}
}
} else {
nl++;
}
else if (m->jnt_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 (dist<margin) {
nl += mj_addConstraintCount(m, 1, 3);
nnzl += 3;
}
}
}
// count limited tendons
for (int i=0; i<ntendon; i++) {
nl += m->tendon_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 (dist<margin) {
nl += mj_addConstraintCount(m, 1, d->ten_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; i<ncon; i++) {
mjContact* con = d->contact + 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);
+38 -1
View File
@@ -27,8 +27,8 @@
namespace mujoco {
namespace {
using ::testing::Pointwise;
using ::testing::DoubleNear;
using ::testing::Pointwise;
using CoreConstraintTest = MujocoTest;
std::vector<mjtNum> 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
@@ -0,0 +1,18 @@
<mujoco>
<default>
<geom rgba=".5 .5 .5 .5"/>
</default>
<worldbody>
<body>
<geom name="1" size="1"/>
</body>
<body>
<geom name="2" size="1" pos="1.5 0 0"/>
</body>
</worldbody>
<contact>
<pair geom1="1" geom2="2"/>
</contact>
</mujoco>
@@ -0,0 +1,13 @@
<mujoco>
<worldbody>
<site name="1"/>
<site name="2" pos="1 0 0"/>
</worldbody>
<tendon>
<spatial frictionloss="1">
<site site="1"/>
<site site="2"/>
</spatial>
</tendon>
</mujoco>
@@ -0,0 +1,13 @@
<mujoco>
<worldbody>
<site name="1"/>
<site name="2" pos="1 0 0"/>
</worldbody>
<tendon>
<spatial limited="true" range="0 0.5">
<site site="1"/>
<site site="2"/>
</spatial>
</tendon>
</mujoco>
@@ -0,0 +1,13 @@
<mujoco>
<worldbody>
<site name="1"/>
<site name="2" pos="1 0 0"/>
</worldbody>
<tendon>
<spatial limited="true" range="0.99 1.01" margin="0.1">
<site site="1"/>
<site site="2"/>
</spatial>
</tendon>
</mujoco>
+24
View File
@@ -0,0 +1,24 @@
<mujoco>
<worldbody>
<body name="1" >
<freejoint />
<geom size="1" pos="7 0 0"/>
</body>
<body name="2" >
<freejoint />
<geom size="1" pos="8.5 0 0"/>
</body>
<body name="3" >
<geom size="1"/>
</body>
<body name="4" >
<geom size="1" pos="1.5 0 0"/>
</body>
</worldbody>
<equality>
<weld body1="1" body2="2"/>
<weld body1="3" body2="4"/>
</equality>
</mujoco>
@@ -0,0 +1,8 @@
<mujoco>
<worldbody>
<body>
<joint axis="1 0 0" range="0 0.1" limited="true" margin="0.2"/>
<geom size="1"/>
</body>
</worldbody>
</mujoco>
@@ -0,0 +1,17 @@
<mujoco>
<worldbody>
<site name="1"/>
<body pos="1 0 0">
<site name="2"/>
<freejoint/>
<geom size=".01"/>
</body>
</worldbody>
<tendon>
<spatial limited="true" range="0.99 1.01" margin="0.1">
<site site="1"/>
<site site="2"/>
</spatial>
</tendon>
</mujoco>
+1
View File
@@ -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];