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

PiperOrigin-RevId: 512947632
Change-Id: I9605e81b4e51fbc90854c13153e4fee38ad48fb9

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
+  <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
 ^^^^^^^^^^^^^^^

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_type;         // constraint type (mjtConstraint)                  (nefc x 1)
-  int*      efc_id;               // id of object of specified type           (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_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_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_rowsuper;   // number of subsequent rows in supernode           (nefc x 1)
-  int*      efc_J_colind;         // column indices in Jacobian               (nefc x nv)
+  int*    efc_J_colind;     // column indices in constraint Jacobian            (nnzJ x 1)
-  int*      efc_JT_rownnz;        // number of non-zeros in Jacobian row    T (nv 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_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_rowsuper;  // number of subsequent rows in supernode         T (nv x 1)
-  int*      efc_JT_colind;        // column indices in Jacobian             T (nv x nefc)
+  int*    efc_JT_colind;    // column indices in constraint Jacobian          T (nnzJ x 1)
-  mjtNum*   efc_J;                // constraint Jacobian                      (nefc x nv)
+  mjtNum* efc_J;            // constraint Jacobian                              (nnzJ x 1)
-  mjtNum*   efc_JT;               // constraint Jacobian transposed           (nv x nefc)
+  mjtNum* efc_JT;           // constraint Jacobian transposed                   (nnzJ x 1)
-  mjtNum*   efc_pos;              // constraint position (equality, contact)  (nefc x 1)
+  mjtNum* efc_pos;          // constraint position (equality, contact)          (nefc x 1)
-  mjtNum*   efc_margin;           // inclusion margin (contact)               (nefc x 1)
+  mjtNum* efc_margin;       // inclusion margin (contact)                       (nefc x 1)
-  mjtNum*   efc_frictionloss;     // frictionloss (friction)                  (nefc x 1)
+  mjtNum* efc_frictionloss; // frictionloss (friction)                          (nefc x 1)
-  mjtNum*   efc_diagApprox;       // approximation to diagonal of A           (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_KBIP;         // stiffness, damping, impedance, imp'              (nefc x 4)
-  mjtNum*   efc_D;                // constraint mass                          (nefc x 1)
+  mjtNum* efc_D;            // constraint mass                                  (nefc x 1)
-  mjtNum*   efc_R;                // inverse 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_b;            // linear cost term: J*qacc_smooth - aref           (nefc x 1)
-  mjtNum*   efc_force;            // constraint force in constraint space     (nefc x 1)
+  mjtNum* efc_force;        // constraint force in constraint space             (nefc x 1)
-  int*      efc_state;            // constraint state (mjtConstraintState)    (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_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_rowadr;       // row start address in colind array                (nefc x 1)
-  int*      efc_AR_colind;        // column indices in sparse AR              (nefc x nefc)
+  int* efc_AR_colind;       // column indices in sparse AR                      (nefc x nefc)
-  mjtNum*   efc_AR;               // J*inv(M)*J' + R                          (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

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_type;         // constraint type (mjtConstraint)                  (nefc x 1)
-  int*      efc_id;               // id of object of specified type           (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_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_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_rowsuper;   // number of subsequent rows in supernode           (nefc x 1)
-  int*      efc_J_colind;         // column indices in Jacobian               (nefc x nv)
+  int*    efc_J_colind;     // column indices in constraint Jacobian            (nnzJ x 1)
-  int*      efc_JT_rownnz;        // number of non-zeros in Jacobian row    T (nv 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_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_rowsuper;  // number of subsequent rows in supernode         T (nv x 1)
-  int*      efc_JT_colind;        // column indices in Jacobian             T (nv x nefc)
+  int*    efc_JT_colind;    // column indices in constraint Jacobian          T (nnzJ x 1)
-  mjtNum*   efc_J;                // constraint Jacobian                      (nefc x nv)
+  mjtNum* efc_J;            // constraint Jacobian                              (nnzJ x 1)
-  mjtNum*   efc_JT;               // constraint Jacobian transposed           (nv x nefc)
+  mjtNum* efc_JT;           // constraint Jacobian transposed                   (nnzJ x 1)
-  mjtNum*   efc_pos;              // constraint position (equality, contact)  (nefc x 1)
+  mjtNum* efc_pos;          // constraint position (equality, contact)          (nefc x 1)
-  mjtNum*   efc_margin;           // inclusion margin (contact)               (nefc x 1)
+  mjtNum* efc_margin;       // inclusion margin (contact)                       (nefc x 1)
-  mjtNum*   efc_frictionloss;     // frictionloss (friction)                  (nefc x 1)
+  mjtNum* efc_frictionloss; // frictionloss (friction)                          (nefc x 1)
-  mjtNum*   efc_diagApprox;       // approximation to diagonal of A           (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_KBIP;         // stiffness, damping, impedance, imp'              (nefc x 4)
-  mjtNum*   efc_D;                // constraint mass                          (nefc x 1)
+  mjtNum* efc_D;            // constraint mass                                  (nefc x 1)
-  mjtNum*   efc_R;                // inverse 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_b;            // linear cost term: J*qacc_smooth - aref           (nefc x 1)
-  mjtNum*   efc_force;            // constraint force in constraint space     (nefc x 1)
+  mjtNum* efc_force;        // constraint force in constraint space             (nefc x 1)
-  int*      efc_state;            // constraint state (mjtConstraintState)    (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_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_rowadr;       // row start address in colind array                (nefc x 1)
-  int*      efc_AR_colind;        // column indices in sparse AR              (nefc x nefc)
+  int* efc_AR_colind;       // column indices in sparse AR                      (nefc x nefc)
-  mjtNum*   efc_AR;               // J*inv(M)*J' + R                          (nefc x nefc)
+  mjtNum* efc_AR;           // J*inv(M)*J' + R                                  (nefc x nefc)
 };
 typedef struct mjData_ mjData;
 

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                          \
+#define MJDATA_ARENA_POINTERS_PRIMAL              \
-    X( int,       efc_type,          MJ_D(nefc), 1          ) \
+  X(int,      efc_type,          MJ_D(nefc),  1)  \
-    X( int,       efc_id,            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_rownnz,      MJ_D(nefc),  1)  \
-    X( int,       efc_J_rowadr,      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_rowsuper,    MJ_D(nefc),  1)  \
-    X( int,       efc_J_colind,      MJ_D(nefc), MJ_M(nv)   ) \
+  X(int,      efc_J_colind,      MJ_D(nnzJ),  1)  \
-    X( int,       efc_JT_rownnz,     MJ_M(nv),   1          ) \
+  X(int,      efc_JT_rownnz,     MJ_M(nv),    1)  \
-    X( int,       efc_JT_rowadr,     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_rowsuper,   MJ_M(nv),    1)  \
-    X( int,       efc_JT_colind,     MJ_M(nv),   MJ_D(nefc) ) \
+  X(int,      efc_JT_colind,     MJ_D(nnzJ),  1)  \
-    X( mjtNum,    efc_J,             MJ_D(nefc), MJ_M(nv)   ) \
+  X(mjtNum,   efc_J,             MJ_D(nnzJ),  1)  \
-    X( mjtNum,    efc_JT,            MJ_M(nv),   MJ_D(nefc) ) \
+  X(mjtNum,   efc_JT,            MJ_D(nnzJ),  1)  \
-    X( mjtNum,    efc_pos,           MJ_D(nefc), 1          ) \
+  X(mjtNum,   efc_pos,           MJ_D(nefc),  1)  \
-    X( mjtNum,    efc_margin,        MJ_D(nefc), 1          ) \
+  X(mjtNum,   efc_margin,        MJ_D(nefc),  1)  \
-    X( mjtNum,    efc_frictionloss,  MJ_D(nefc), 1          ) \
+  X(mjtNum,   efc_frictionloss,  MJ_D(nefc),  1)  \
-    X( mjtNum,    efc_diagApprox,    MJ_D(nefc), 1          ) \
+  X(mjtNum,   efc_diagApprox,    MJ_D(nefc),  1)  \
-    X( mjtNum,    efc_KBIP,          MJ_D(nefc), 4          ) \
+  X(mjtNum,   efc_KBIP,          MJ_D(nefc),  4)  \
-    X( mjtNum,    efc_D,             MJ_D(nefc), 1          ) \
+  X(mjtNum,   efc_D,             MJ_D(nefc),  1)  \
-    X( mjtNum,    efc_R,             MJ_D(nefc), 1          ) \
+  X(mjtNum,   efc_R,             MJ_D(nefc),  1)  \
-    X( mjtNum,    efc_vel,           MJ_D(nefc), 1          ) \
+  X(mjtNum,   efc_vel,           MJ_D(nefc),  1)  \
-    X( mjtNum,    efc_aref,          MJ_D(nefc), 1          ) \
+  X(mjtNum,   efc_aref,          MJ_D(nefc),  1)  \
-    X( mjtNum,    efc_b,             MJ_D(nefc), 1          ) \
+  X(mjtNum,   efc_b,             MJ_D(nefc),  1)  \
-    X( mjtNum,    efc_force,         MJ_D(nefc), 1          ) \
+  X(mjtNum,   efc_force,         MJ_D(nefc),  1)  \
-    X( int,       efc_state,         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                            \

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);

src/engine/engine_core_constraint.c

@@ -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
+// count the number of non-zeros in the sum of two sparse vectors
-static inline int mj_ne(const mjModel* m, const mjData* d) {
+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 (nnz) {
-    if (!m->eq_active[i]) {
+    chain = (int*)mj_stackAlloc(d, nv);
-      continue;
+    chain2 = (int*)mj_stackAlloc(d, nv);
-    }
+  }
 
-    // process according to type
+  // find active equality constraints
-    switch (m->eq_type[i]) {
+  for (int i=0; i<neq; i++) {
-    case mjEQ_CONNECT:
+    if (m->eq_active[i]) {
-      ne += 3;
+      id[0] = m->eq_obj1id[i];
-      break;
+      id[1] = m->eq_obj2id[i];
+      size = 0;
+      NV = 0;
+      NV2 = 0;
 
-    case mjEQ_WELD:
+      // process according to type
-      ne += 6;
+      switch (m->eq_type[i]) {
-      break;
+        case mjEQ_CONNECT:
+          size = 3;
+          if (!nnz) {
+            break;
+          }
 
-    case mjEQ_JOINT:
+          NV = mj_jacDifPairCount(m, chain, id[1], id[0]);
-    case mjEQ_TENDON:
+          break;
-      ne++;
-      break;
 
-    default:                        // SHOULD NOT OCCUR
+        case mjEQ_WELD:
-      mju_error_i("Invalid equality constraint type %d", m->eq_type[i]);
+          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
+// count frictional constraints, count Jacobian nonzeros if nnz is not NULL
-static inline int mj_nf(const mjModel* m, const mjData* d) {
+static inline int mj_nf(const mjModel* m, const mjData* d, int *nnz) {
-  // disabled: return
+  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
+// count limit constraints, count Jacobian nonzeros if nnz is not NULL
-static inline int mj_nl(const mjModel* m, const mjData* d) {
+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<m->njnt; i++) {
-  for (int i=0; i<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
+      value = d->qpos[m->jnt_qposadr[i]];
-      mjtNum margin = m->jnt_margin[i];
+      for (side=-1; side<=1; side+=2) {
-
+        dist = side * (m->jnt_range[2*i+(side+1)/2] - value);
-      // 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
         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
+// count contact constraints, count Jacobian nonzeros if nnz is not NULL
-static inline int mj_nc(const mjModel* m, const mjData* d) {
+static inline int mj_nc(const mjModel* m, mjData* d, int* nnz) {
-  // disabled or no contacts: return
+  int nnzc = 0, nc = 0;
-  int ncon = d->ncon;
+  int ispyramid = mj_isPyramidal(m), ncon = d->ncon;
-  if (mjDISABLED(mjDSBL_CONTACT) || ncon==0) {
+
+  if (mjDISABLED(mjDSBL_CONTACT) || !ncon) {
     return 0;
   }
 
-  int nc = 0;
+  mjMARKSTACK;
-  int ispyramid = mj_isPyramidal(m);
+  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 (d->contact[i].exclude) {
-    if (con->exclude) {
       continue;
     }
 
+    mjContact* con = d->contact + i;
     int dim = con->dim;
-
+    int b1 = m->geom_bodyid[con->geom1];
-    // dim 1: single constraint
+    int b2 = m->geom_bodyid[con->geom2];
-    if (dim==1) {
+    int NV = mj_jacDifPairCount(m, chain, b1, b2);
-      nc++;
+    if (!NV) {
+      continue;
     }
 
-    // dim > 1: depends on cone type
+    if (dim==1) {
-    else {
+      nc++;
-      nc += (ispyramid ? 2*(dim-1) : dim);
+      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);

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<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

test/engine/testdata/core_constraint/dofless_contact.xml

@@ -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>

test/engine/testdata/core_constraint/dofless_tendon_frictional.xml

@@ -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>

test/engine/testdata/core_constraint/dofless_tendon_limited.xml

@@ -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>

test/engine/testdata/core_constraint/dofless_tendon_limitedmargin.xml

@@ -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>

test/engine/testdata/core_constraint/dofless_weld.xml

@@ -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>

test/engine/testdata/core_constraint/joint_limited_bilateral_margin.xml

@@ -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>

test/engine/testdata/core_constraint/tendon_limited_bilateral_margin.xml

@@ -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>

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];