Reduce flexedge_J size from nflexedge x nv to the effective sparse size.

PiperOrigin-RevId: 856290141
Change-Id: Ide297d1b6f0e3aa07e3e20bf9ab44b6501097695

doc/includes/references.h

@@ -272,7 +272,7 @@ struct mjData_ {
   // computed by mj_fwdPosition/mj_flex
   mjtNum* flexvert_xpos;     // Cartesian flex vertex positions                  (nflexvert x 3)
   mjtNum* flexelem_aabb;     // flex element bounding boxes (center, size)       (nflexelem x 6)
-  mjtNum* flexedge_J;        // flex edge Jacobian                               (nflexedge x nv)
+  mjtNum* flexedge_J;        // flex edge Jacobian                               (nJfe x 1)
   mjtNum* flexedge_length;   // flex edge lengths                                (nflexedge x 1)
   mjtNum* bvh_aabb_dyn;      // global bounding box (center, size)               (nbvhdynamic x 6)
 
@@ -1033,6 +1033,7 @@ struct mjModel_ {
   int nflexshelldata;             // number of shell fragment vertex ids in all flexes
   int nflexevpair;                // number of element-vertex pairs in all flexes
   int nflextexcoord;              // number of vertices with texture coordinates
+  int nJfe;                       // number of non-zeros in sparse flexedge Jacobian matrix
   int nmesh;                      // number of meshes
   int nmeshvert;                  // number of vertices in all meshes
   int nmeshnormal;                // number of normals in all meshes
@@ -1341,7 +1342,7 @@ struct mjModel_ {
   int*      flex_bvhnum;          // number of bounding volumes               (nflex x 1)
   int*      flexedge_J_rownnz;    // number of non-zeros in Jacobian row      (nflexedge x 1)
   int*      flexedge_J_rowadr;    // row start address in colind array        (nflexedge x 1)
-  int*      flexedge_J_colind;    // column indices in sparse Jacobian        (nflexedge x nv)
+  int*      flexedge_J_colind;    // column indices in sparse Jacobian        (nJfe x 1)
   float*    flex_rgba;            // rgba when material is omitted            (nflex x 4)
   float*    flex_texcoord;        // vertex texture coordinates               (nflextexcoord x 2)
 

include/mujoco/mjdata.h

@@ -306,7 +306,7 @@ struct mjData_ {
   // computed by mj_fwdPosition/mj_flex
   mjtNum* flexvert_xpos;     // Cartesian flex vertex positions                  (nflexvert x 3)
   mjtNum* flexelem_aabb;     // flex element bounding boxes (center, size)       (nflexelem x 6)
-  mjtNum* flexedge_J;        // flex edge Jacobian                               (nflexedge x nv)
+  mjtNum* flexedge_J;        // flex edge Jacobian                               (nJfe x 1)
   mjtNum* flexedge_length;   // flex edge lengths                                (nflexedge x 1)
   mjtNum* bvh_aabb_dyn;      // global bounding box (center, size)               (nbvhdynamic x 6)
 

include/mujoco/mjmodel.h

@@ -701,6 +701,7 @@ struct mjModel_ {
   int nflexshelldata;             // number of shell fragment vertex ids in all flexes
   int nflexevpair;                // number of element-vertex pairs in all flexes
   int nflextexcoord;              // number of vertices with texture coordinates
+  int nJfe;                       // number of non-zeros in sparse flexedge Jacobian matrix
   int nmesh;                      // number of meshes
   int nmeshvert;                  // number of vertices in all meshes
   int nmeshnormal;                // number of normals in all meshes
@@ -1009,7 +1010,7 @@ struct mjModel_ {
   int*      flex_bvhnum;          // number of bounding volumes               (nflex x 1)
   int*      flexedge_J_rownnz;    // number of non-zeros in Jacobian row      (nflexedge x 1)
   int*      flexedge_J_rowadr;    // row start address in colind array        (nflexedge x 1)
-  int*      flexedge_J_colind;    // column indices in sparse Jacobian        (nflexedge x nv)
+  int*      flexedge_J_colind;    // column indices in sparse Jacobian        (nJfe x 1)
   float*    flex_rgba;            // rgba when material is omitted            (nflex x 4)
   float*    flex_texcoord;        // vertex texture coordinates               (nflextexcoord x 2)
 

include/mujoco/mjxmacro.h

@@ -96,6 +96,7 @@
     X( nflexshelldata )     \
     X( nflexevpair )        \
     X( nflextexcoord )      \
+    X( nJfe )               \
     X( nmesh )              \
     X( nmeshvert )          \
     X( nmeshnormal )        \
@@ -399,7 +400,7 @@
     X   ( int,     flex_bvhnum,           nflex,         1                    ) \
     X   ( int,     flexedge_J_rownnz,     nflexedge,     1                    ) \
     X   ( int,     flexedge_J_rowadr,     nflexedge,     1                    ) \
-    X   ( int,     flexedge_J_colind,     nflexedge,     MJ_M(nv)             ) \
+    X   ( int,     flexedge_J_colind,     nJfe,          1                    ) \
     X   ( float,   flex_rgba,             nflex,         4                    ) \
     X   ( float,   flex_texcoord,         nflextexcoord, 2                    )
 
@@ -733,7 +734,7 @@
     X   ( mjtNum,    cinert,            nbody,       10          ) \
     X   ( mjtNum,    flexvert_xpos,     nflexvert,   3           ) \
     X   ( mjtNum,    flexelem_aabb,     nflexelem,   6           ) \
-    X   ( mjtNum,    flexedge_J,        nflexedge,   MJ_M(nv)    ) \
+    X   ( mjtNum,    flexedge_J,        nJfe,        1           ) \
     X   ( mjtNum,    flexedge_length,   nflexedge,   1           ) \
     X   ( mjtNum,    bvh_aabb_dyn,      nbvhdynamic, 6           ) \
     X   ( int,       ten_wrapadr,       ntendon,     1           ) \

mjx/mujoco/mjx/_src/io.py

@@ -754,7 +754,7 @@ def _make_data_c(
       'light_xdir': (m.nlight, 3, float_),
       'flexvert_xpos': (nflexvert, 3, float_),
       'flexelem_aabb': (nflexelem, 6, float_),
-      'flexedge_J': (nflexedge, m.nv, float_),
+      'flexedge_J': (m.nJfe, float_),
       'flexedge_length': (nflexedge, float_),
       'ten_J_rownnz': (m.ntendon, np.int32),
       'ten_J_rowadr': (m.ntendon, np.int32),

python/mujoco/introspect/structs.py

@@ -992,6 +992,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  type=ValueType(name='int'),
                  doc='number of vertices with texture coordinates',
              ),
+             StructFieldDecl(
+                 name='nJfe',
+                 type=ValueType(name='int'),
+                 doc='number of non-zeros in sparse flexedge Jacobian matrix',
+             ),
              StructFieldDecl(
                  name='nmesh',
                  type=ValueType(name='int'),
@@ -2993,7 +2998,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                      inner_type=ValueType(name='int'),
                  ),
                  doc='column indices in sparse Jacobian',
-                 array_extent=('nflexedge', 'nv'),
+                 array_extent=('nJfe',),
              ),
              StructFieldDecl(
                  name='flex_rgba',
@@ -5594,7 +5599,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                      inner_type=ValueType(name='mjtNum'),
                  ),
                  doc='flex edge Jacobian',
-                 array_extent=('nflexedge', 'nv'),
+                 array_extent=('nJfe',),
              ),
              StructFieldDecl(
                  name='flexedge_length',

src/engine/engine_core_constraint.c

@@ -447,7 +447,7 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
 
       // compute Jacobian difference (opposite of contact: 0 - 1)
       NV = mj_jacDifPair(m, d, chain, body_id[1], body_id[0], pos[1], pos[0],
-                          jac[1], jac[0], jacdif, NULL, NULL, NULL);
+                          jac[1], jac[0], jacdif, NULL, NULL, NULL, issparse);
 
       // copy difference into jac[0]
       mju_copy(jac[0], jacdif, 3*NV);
@@ -483,7 +483,7 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
       // compute error Jacobian (opposite of contact: 0 - 1)
       NV = mj_jacDifPair(m, d, chain, body_id[1], body_id[0], pos[1], pos[0],
                           jac[1], jac[0], jacdif,
-                          jac[1]+3*nv, jac[0]+3*nv, jacdif+3*nv);
+                          jac[1]+3*nv, jac[0]+3*nv, jacdif+3*nv, issparse);
 
       // copy difference into jac[0], compress translation:rotation if sparse
       mju_copy(jac[0], jacdif, 3*NV);
@@ -628,13 +628,17 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
         // add constraint: sparse or dense
         if (issparse) {
           mj_addConstraint(m, d, d->flexedge_J+m->flexedge_J_rowadr[e], cpos, 0, 0,
-                            1, mjCNSTR_EQUALITY, i,
+                           1, mjCNSTR_EQUALITY, i,
-                            m->flexedge_J_rownnz[e],
+                           m->flexedge_J_rownnz[e],
-                            m->flexedge_J_colind+m->flexedge_J_rowadr[e]);
+                           m->flexedge_J_colind+m->flexedge_J_rowadr[e]);
         } else {
-          mj_addConstraint(m, d, d->flexedge_J+e*nv, cpos, 0, 0,
+          mju_zero(jac[0], nv);  // reuse first row of jac[0]
-                            1, mjCNSTR_EQUALITY, i,
+          int rowadr = m->flexedge_J_rowadr[e];
-                            0, NULL);
+          int rownnz = m->flexedge_J_rownnz[e];
+          for (int k=0; k<rownnz; k++) {
+            jac[0][m->flexedge_J_colind[rowadr+k]] = d->flexedge_J[rowadr+k];
+          }
+          mj_addConstraint(m, d, jac[0], cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, 0, NULL);
         }
       }
       break;
@@ -891,10 +895,10 @@ int mj_contactJacobian(const mjModel* m, mjData* d, const mjContact* con, int di
     // compute Jacobian differences
     if (dim > 3) {
       return mj_jacDifPair(m, d, chain, bid[0], bid[1], con->pos, con->pos,
-                           jac1p, jac2p, jacdifp, jac1r, jac2r, jacdifr);
+                           jac1p, jac2p, jacdifp, jac1r, jac2r, jacdifr, mj_isSparse(m));
     } else {
       return mj_jacDifPair(m, d, chain, bid[0], bid[1], con->pos, con->pos,
-                           jac1p, jac2p, jacdifp, NULL, NULL, NULL);
+                           jac1p, jac2p, jacdifp, NULL, NULL, NULL, mj_isSparse(m));
     }
   }
 

src/engine/engine_core_smooth.c

@@ -535,9 +535,8 @@ void mj_updateDynamicBVH(const mjModel* m, mjData* d, int bvhadr, int bvhnum) {
 
 // compute flex-related quantities
 void mj_flex(const mjModel* m, mjData* d) {
-  int nv = m->nv, issparse = mj_isSparse(m);
+  int nv = m->nv;
   int* rowadr = m->flexedge_J_rowadr, *rownnz = m->flexedge_J_rownnz;
-  mjtNum* J = d->flexedge_J;
 
   // skip if no flexes
   if (!m->nflex) {
@@ -655,15 +654,11 @@ void mj_flex(const mjModel* m, mjData* d) {
   mjtNum* jac1 = mjSTACKALLOC(d, 3*nv, mjtNum);
   mjtNum* jac2 = mjSTACKALLOC(d, 3*nv, mjtNum);
   mjtNum* jacdif = mjSTACKALLOC(d, 3*nv, mjtNum);
-  int* chain = issparse ? mjSTACKALLOC(d, nv, int) : NULL;
+  int* chain = mjSTACKALLOC(d, nv, int);
 
-  // clear Jacobian: sparse or dense
+  // clear Jacobian
-  if (issparse) {
+  mju_zeroInt(rowadr, m->nflexedge);
-    mju_zeroInt(rowadr, m->nflexedge);
+  mju_zeroInt(rownnz, m->nflexedge);
-    mju_zeroInt(rownnz, m->nflexedge);
-  } else {
-    mju_zero(J, m->nflexedge*nv);
-  }
 
   // compute lengths and Jacobians of edges
   for (int f=0; f < m->nflex; f++) {
@@ -699,40 +694,26 @@ void mj_flex(const mjModel* m, mjData* d) {
         continue;
       }
 
-      // sparse edge Jacobian
+      // set rowadr
-      if (issparse) {
+      if (ebase+e > 0) {
-        // set rowadr
+        rowadr[ebase+e] = rowadr[ebase+e-1] + rownnz[ebase+e-1];
-        if (ebase+e > 0) {
-          rowadr[ebase+e] = rowadr[ebase+e-1] + rownnz[ebase+e-1];
-        }
-
-        // get endpoint Jacobians, subtract
-        int NV = mj_jacDifPair(m, d, chain, b1, b2, pos1, pos2,
-                               jac1, jac2, jacdif, NULL, NULL, NULL);
-
-        // no dofs: skip
-        if (!NV) {
-          continue;
-        }
-
-        // apply chain rule to compute edge Jacobian
-        mju_mulMatTVec(J + rowadr[ebase+e], jacdif, vec, 3, NV);
-
-        // copy sparsity info
-        rownnz[ebase+e] = NV;
-        mju_copyInt(m->flexedge_J_colind + rowadr[ebase+e], chain, NV);
       }
 
-      // dense edge Jacobian
+      // get endpoint Jacobians, subtract
-      else {
+      int NV = mj_jacDifPair(m, d, chain, b1, b2, pos1, pos2,
-        // get endpoint Jacobians, subtract
+                              jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1);
-        mj_jac(m, d, jac1, NULL, pos1, b1);
-        mj_jac(m, d, jac2, NULL, pos2, b2);
-        mju_sub(jacdif, jac2, jac1, 3*nv);
 
-        // apply chain rule to compute edge Jacobian
+      // no dofs: skip
-        mju_mulMatTVec(J + (ebase+e)*nv, jacdif, vec, 3, nv);
+      if (!NV) {
+        continue;
       }
+
+      // apply chain rule to compute edge Jacobian
+      mju_mulMatTVec(d->flexedge_J + rowadr[ebase+e], jacdif, vec, 3, NV);
+
+      // copy sparsity info
+      rownnz[ebase+e] = NV;
+      mju_copyInt(m->flexedge_J_colind + rowadr[ebase+e], chain, NV);
     }
   }
 
@@ -903,7 +884,7 @@ void mj_tendon(const mjModel* m, mjData* d) {
             // get endpoint Jacobians, subtract
             int NV = mj_jacDifPair(m, d, chain,
                                    wbody[k], wbody[k+1], wpnt+3*k, wpnt+3*k+3,
-                                   jac1, jac2, jacdif, NULL, NULL, NULL);
+                                   jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1);
 
             // no dofs: skip
             if (!NV) {
@@ -1526,7 +1507,7 @@ void mj_transmission(const mjModel* m, mjData* d) {
 
             // get Jacobian difference
             int NV = mj_jacDifPair(m, d, chain, b1, b2, con->pos, con->pos,
-                                   jac1p, jac2p, jacdifp, NULL, NULL, NULL);
+                                   jac1p, jac2p, jacdifp, NULL, NULL, NULL, issparse);
 
             // project Jacobian along the normal of the contact frame
             mju_mulMatMat(jac, con->frame, jacdifp, 1, 3, NV);

src/engine/engine_core_util.c

@@ -438,9 +438,8 @@ void mj_jacSparseSimple(const mjModel* m, const mjData* d,
 int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
                   int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3],
                   mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp,
-                  mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr) {
+                  mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr, int issparse) {
   int issimple = (m->body_simple[b1] && m->body_simple[b2]);
-  int issparse = mj_isSparse(m);
   int NV = m->nv;
 
   // skip if no DOFs

src/engine/engine_core_util.h

@@ -89,7 +89,7 @@ void mj_jacSparseSimple(const mjModel* m, const mjData* d,
 MJAPI int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
                         int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3],
                         mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp,
-                        mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr);
+                        mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr, int issparse);
 
 // dense or sparse weighted sum of multiple body Jacobians at same point
 int mj_jacSum(const mjModel* m, mjData* d, int* chain,

src/engine/engine_derivative.c

@@ -1495,13 +1495,9 @@ void mjd_passive_vel(const mjModel* m, mjData* d) {
         continue;
       }
 
-      // add sparse or dense
+      // always sparse
-      if (mj_isSparse(m)) {
+      addJTBJSparse(m, d, d->flexedge_J, &B, 1, e,
-        addJTBJSparse(m, d, d->flexedge_J, &B, 1, e,
+                    m->flexedge_J_rownnz, m->flexedge_J_rowadr, m->flexedge_J_colind);
-                      m->flexedge_J_rownnz, m->flexedge_J_rowadr, m->flexedge_J_colind);
-      } else {
-        addJTBJ(m, d, d->flexedge_J+e*nv, &B, 1);
-      }
     }
   }
 

src/engine/engine_forward.c

@@ -218,13 +218,9 @@ void mj_fwdPosition(const mjModel* m, mjData* d) {
 void mj_fwdVelocity(const mjModel* m, mjData* d) {
   TM_START;
 
-  // flexedge velocity: dense or sparse
+  // flexedge velocity: always sparse
-  if (mj_isSparse(m)) {
+  mju_mulMatVecSparse(d->flexedge_velocity, d->flexedge_J, d->qvel, m->nflexedge,
-    mju_mulMatVecSparse(d->flexedge_velocity, d->flexedge_J, d->qvel, m->nflexedge,
+                      m->flexedge_J_rownnz, m->flexedge_J_rowadr, m->flexedge_J_colind, NULL);
-                        m->flexedge_J_rownnz, m->flexedge_J_rowadr, m->flexedge_J_colind, NULL);
-  } else {
-    mju_mulMatVec(d->flexedge_velocity, d->flexedge_J, d->qvel, m->nflexedge, m->nv);
-  }
 
   // tendon velocity: dense or sparse
   if (mj_isSparse(m)) {

src/engine/engine_io.c

@@ -218,9 +218,9 @@ static void freeModelBuffers(mjModel* m) {
 void mj_makeModel(mjModel** dest,
     int nq, int nv, int nu, int na, int nbody, int nbvh,
     int nbvhstatic, int nbvhdynamic, int noct, int njnt, int ntree,
-    int nM, int nB, int nC, int nD, int ngeom, int nsite, int ncam,
+    int nM, int nB, int nC, int nD, int ngeom, int nsite, int ncam, int nlight,
-    int nlight, int nflex, int nflexnode, int nflexvert, int nflexedge, int nflexelem,
+    int nflex, int nflexnode, int nflexvert, int nflexedge, int nflexelem, int nflexelemdata,
-    int nflexelemdata, int nflexelemedge, int nflexshelldata, int nflexevpair, int nflextexcoord,
+    int nflexelemedge, int nflexshelldata, int nflexevpair, int nflextexcoord, int nJfe,
     int nmesh, int nmeshvert, int nmeshnormal, int nmeshtexcoord, int nmeshface,
     int nmeshgraph, int nmeshpoly, int nmeshpolyvert, int nmeshpolymap, int nskin, int nskinvert,
     int nskintexvert, int nskinface,
@@ -278,6 +278,7 @@ void mj_makeModel(mjModel** dest,
   m->nflexshelldata = nflexshelldata;
   m->nflexevpair = nflexevpair;
   m->nflextexcoord = nflextexcoord;
+  m->nJfe = nJfe;
   m->nmesh = nmesh;
   m->nmeshvert = nmeshvert;
   m->nmeshnormal = nmeshnormal;
@@ -406,7 +407,7 @@ mjModel* mj_copyModel(mjModel* dest, const mjModel* src) {
         src->nM, src->nB, src->nC, src->nD, src->ngeom, src->nsite, src->ncam,
         src->nlight, src->nflex, src->nflexnode, src->nflexvert, src->nflexedge,
         src->nflexelem, src->nflexelemdata, src->nflexelemedge,
-        src->nflexshelldata, src->nflexevpair, src->nflextexcoord, src->nmesh,
+        src->nflexshelldata, src->nflexevpair, src->nflextexcoord, src->nJfe, src->nmesh,
         src->nmeshvert, src->nmeshnormal, src->nmeshtexcoord, src->nmeshface,
         src->nmeshgraph, src->nmeshpoly, src->nmeshpolyvert, src->nmeshpolymap,
         src->nskin, src->nskinvert, src->nskintexvert, src->nskinface,
@@ -597,7 +598,7 @@ mjModel* mj_loadModelBuffer(const void* buffer, int buffer_sz) {
                ints[49], ints[50], ints[51], ints[52], ints[53], ints[54], ints[55],
                ints[56], ints[57], ints[58], ints[59], ints[60], ints[61], ints[62],
                ints[63], ints[64], ints[65], ints[66], ints[67], ints[68], ints[69],
-               ints[70], ints[71], ints[72], ints[73], ints[74]);
+               ints[70], ints[71], ints[72], ints[73], ints[74], ints[75]);
 
   // read mjModel mjtSize fields
   mjtSize sizes[8];

src/engine/engine_io.h

@@ -51,7 +51,7 @@ void mj_makeModel(mjModel** dest,
     int nq, int nv, int nu, int na, int nbody, int nbvh, int nbvhstatic, int nbvhdynamic, int noct,
     int njnt, int ntree, int nM, int nB, int nC, int nD, int ngeom, int nsite, int ncam, int nlight,
     int nflex, int nflexnode, int nflexvert, int nflexedge, int nflexelem, int nflexelemdata,
-    int nflexelemedge, int nflexshelldata, int nflexevpair, int nflextexcoord, int nmesh,
+    int nJfe, int nflexelemedge, int nflexshelldata, int nflexevpair, int nflextexcoord, int nmesh,
     int nmeshvert, int nmeshnormal, int nmeshtexcoord, int nmeshface, int nmeshgraph, int nmeshpoly,
     int nmeshpolyvert, int nmeshpolymap, int nskin, int nskinvert, int nskintexvert, int nskinface,
     int nskinbone, int nskinbonevert, int nhfield, int nhfielddata, int ntex, int ntexdata,

src/engine/engine_passive.c

@@ -412,18 +412,13 @@ static void mj_springdamper(const mjModel* m, mjData* d) {
       mjtNum frc_spring = stiffness * (m->flexedge_length0[e] - d->flexedge_length[e]);
       mjtNum frc_damper = -damping * d->flexedge_velocity[e];
 
-      // transform to joint torque, add to qfrc_{spring, damper}: dense or sparse
+      // transform to joint torque, add to qfrc_{spring, damper}: always sparse
-      if (issparse) {
+      int end = m->flexedge_J_rowadr[e] + m->flexedge_J_rownnz[e];
-        int end = m->flexedge_J_rowadr[e] + m->flexedge_J_rownnz[e];
+      for (int j=m->flexedge_J_rowadr[e]; j < end; j++) {
-        for (int j=m->flexedge_J_rowadr[e]; j < end; j++) {
+        int colind = m->flexedge_J_colind[j];
-          int colind = m->flexedge_J_colind[j];
+        mjtNum J = d->flexedge_J[j];
-          mjtNum J = d->flexedge_J[j];
+        d->qfrc_spring[colind] += J * frc_spring;
-          d->qfrc_spring[colind] += J * frc_spring;
+        d->qfrc_damper[colind] += J * frc_damper;
-          d->qfrc_damper[colind] += J * frc_damper;
-        }
-      } else {
-        if (frc_spring) mju_addToScl(d->qfrc_spring, d->flexedge_J+e*nv, frc_spring, nv);
-        if (frc_damper) mju_addToScl(d->qfrc_damper, d->flexedge_J+e*nv, frc_damper, nv);
       }
     }
   }

src/engine/engine_print.c

@@ -1280,15 +1280,11 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
 
   printArray2d("FLEXVERT_XPOS", m->nflexvert, 3, d->flexvert_xpos, fp, float_format);
   printArray2d("FLEXELEM_AABB", m->nflexelem, 6, d->flexelem_aabb, fp, float_format);
-  if (!mj_isSparse(m)) {
+  mj_printSparsity("FLEXEDGE_J: flex edge connectivity", m->nflexedge, m->nv,
-    printArray2d("FLEXEDGE_J", m->nflexedge, m->nv, d->flexedge_J, fp, float_format);
+                    m->flexedge_J_rowadr, NULL, m->flexedge_J_rownnz, NULL, m->flexedge_J_colind,
-  } else {
+                    fp);
-    mj_printSparsity("FLEXEDGE_J: flex edge connectivity", m->nflexedge, m->nv,
+  printSparse("FLEXEDGE_J", d->flexedge_J, m->nflexedge, m->flexedge_J_rownnz,
-                     m->flexedge_J_rowadr, NULL, m->flexedge_J_rownnz, NULL, m->flexedge_J_colind,
+                            m->flexedge_J_rowadr, m->flexedge_J_colind, fp, float_format);
-                     fp);
-    printSparse("FLEXEDGE_J", d->flexedge_J, m->nflexedge, m->flexedge_J_rownnz,
-                              m->flexedge_J_rowadr, m->flexedge_J_colind, fp, float_format);
-  }
   printArray2d("FLEXEDGE_LENGTH", m->nflexedge, 1, d->flexedge_length, fp, float_format);
 
   printArray2d("TEN_LENGTH", m->ntendon, 1, d->ten_length, fp, float_format);

src/engine/engine_setconst.c

@@ -468,14 +468,10 @@ static void set0(mjModel* m, mjData* d) {
         // handle general edge
         else {
           // make dense vector into tmp
-          if (mj_isSparse(m)) {
+          mju_zero(tmp, nv);
-            mju_zero(tmp, nv);
+          int end = m->flexedge_J_rowadr[i] + m->flexedge_J_rownnz[i];
-            int end = m->flexedge_J_rowadr[i] + m->flexedge_J_rownnz[i];
+          for (int j=m->flexedge_J_rowadr[i]; j < end; j++) {
-            for (int j=m->flexedge_J_rowadr[i]; j < end; j++) {
+            tmp[m->flexedge_J_colind[j]] = d->flexedge_J[j];
-              tmp[m->flexedge_J_colind[j]] = d->flexedge_J[j];
-            }
-          } else {
-            mju_copy(tmp, d->flexedge_J+i*nv, nv);
           }
 
           // solve into tmp+nv

src/user/user_model.cc

@@ -34,6 +34,7 @@
 #include <unordered_map>
 #include <utility>
 #include <vector>
+#include <unordered_set>
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmacro.h>
@@ -1153,6 +1154,7 @@ void mjCModel::Clear() {
   nflexshelldata = 0;
   nflexevpair = 0;
   nflextexcoord = 0;
+  nJfe = 0;
   nmeshvert = 0;
   nmeshnormal = 0;
   nmeshtexcoord = 0;
@@ -2141,6 +2143,29 @@ void mjCModel::SetSizes() {
     nflexshelldata += (int)flexes_[i]->shell.size();
     nflexevpair += (int)flexes_[i]->evpair.size()/2;
     nflextexcoord += (flexes_[i]->HasTexcoord() ? flexes_[i]->get_texcoord().size()/2 : 0);
+    if (flexes_[i]->interpolated || flexes_[i]->rigid) {
+      continue;
+    }
+
+    // count number of non-zero elements in the edge Jacobian matrix
+    for (const auto& edge : flexes_[i]->edge) {
+      mjCBody* b1 = bodies_[flexes_[i]->vertbodyid[edge.first]];
+      mjCBody* b2 = bodies_[flexes_[i]->vertbodyid[edge.second]];
+      std::unordered_set<mjCBody*> bodies_in_jac;
+      while (b1 || b2) {
+        if (b1) {
+          bodies_in_jac.insert(b1);
+          b1 = b1->parent;
+        }
+        if (b2) {
+          bodies_in_jac.insert(b2);
+          b2 = b2->parent;
+        }
+      }
+      for (mjCBody* b : bodies_in_jac) {
+        nJfe += b->dofnum;
+      }
+    }
   }
 
   // mesh counts
@@ -4831,7 +4856,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
   mj_makeModel(&m,
                nq, nv, nu, na, nbody, nbvh, nbvhstatic, nbvhdynamic, noct, njnt, ntree, nM, nB, nC,
                nD, ngeom, nsite, ncam, nlight, nflex, nflexnode, nflexvert, nflexedge, nflexelem,
-               nflexelemdata, nflexelemedge, nflexshelldata, nflexevpair, nflextexcoord,
+               nflexelemdata, nflexelemedge, nflexshelldata, nflexevpair, nflextexcoord, nJfe,
                nmesh, nmeshvert, nmeshnormal, nmeshtexcoord, nmeshface, nmeshgraph, nmeshpoly,
                nmeshpolyvert, nmeshpolymap, nskin, nskinvert, nskintexvert, nskinface, nskinbone,
                nskinbonevert, nhfield, nhfielddata, ntex, ntexdata, nmat, npair, nexclude,

src/user/user_model.h

@@ -99,6 +99,7 @@ class mjCModel_ : public mjsElement {
   int nflexshelldata;  // number of shell fragment vertex ids in all flexes
   int nflexevpair;     // number of element-vertex pairs in all flexes
   int nflextexcoord;   // number of vertex texture coordinates in all flexes
+  int nJfe;            // number of non-zeros in sparse flex constraint Jacobian
   int nmeshvert;       // number of vertices in all meshes
   int nmeshnormal;     // number of normals in all meshes
   int nmeshtexcoord;   // number of texture coordinates in all meshes

test/engine/engine_core_constraint_test.cc

@@ -133,7 +133,7 @@ TEST_F(CoreConstraintTest, WeldRotJacobian) {
 
   // rotational Jacobian difference
   mj_jacDifPair(model, data, NULL, 2, 1, point, point,
-                NULL, NULL, NULL, jac0, jac1, jacdif);
+                NULL, NULL, NULL, jac0, jac1, jacdif, mj_isSparse(model));
 
   // formula: 0.5 * neg(quat2) * (jac1-jac2) * quat1
   mjtNum axis[3], quat3[4], quat4[4];

unity/Runtime/Bindings/MjBindings.cs

@@ -5313,6 +5313,7 @@ public unsafe struct mjModel_ {
   public int nflexshelldata;
   public int nflexevpair;
   public int nflextexcoord;
+  public int nJfe;
   public int nmesh;
   public int nmeshvert;
   public int nmeshnormal;

wasm/codegen/generated/bindings.cc

@@ -3579,6 +3579,12 @@ struct MjModel {
   void set_nflextexcoord(int value) {
     ptr_->nflextexcoord = value;
   }
+  int nJfe() const {
+    return ptr_->nJfe;
+  }
+  void set_nJfe(int value) {
+    ptr_->nJfe = value;
+  }
   int nmesh() const {
     return ptr_->nmesh;
   }
@@ -4558,7 +4564,7 @@ struct MjModel {
     return emscripten::val(emscripten::typed_memory_view(ptr_->nflexedge, ptr_->flexedge_J_rowadr));
   }
   emscripten::val flexedge_J_colind() const {
-    return emscripten::val(emscripten::typed_memory_view(ptr_->nflexedge * ptr_->nv, ptr_->flexedge_J_colind));
+    return emscripten::val(emscripten::typed_memory_view(ptr_->nJfe, ptr_->flexedge_J_colind));
   }
   emscripten::val flex_rgba() const {
     return emscripten::val(emscripten::typed_memory_view(ptr_->nflex * 4, ptr_->flex_rgba));
@@ -6295,7 +6301,7 @@ struct MjData {
     return emscripten::val(emscripten::typed_memory_view(model->nflexelem * 6, ptr_->flexelem_aabb));
   }
   emscripten::val flexedge_J() const {
-    return emscripten::val(emscripten::typed_memory_view(model->nflexedge * model->nv, ptr_->flexedge_J));
+    return emscripten::val(emscripten::typed_memory_view(model->nJfe, ptr_->flexedge_J));
   }
   emscripten::val flexedge_length() const {
     return emscripten::val(emscripten::typed_memory_view(model->nflexedge, ptr_->flexedge_length));
@@ -11470,6 +11476,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
     .property("nB", &MjModel::nB, &MjModel::set_nB, reference())
     .property("nC", &MjModel::nC, &MjModel::set_nC, reference())
     .property("nD", &MjModel::nD, &MjModel::set_nD, reference())
+    .property("nJfe", &MjModel::nJfe, &MjModel::set_nJfe, reference())
     .property("nJmom", &MjModel::nJmom, &MjModel::set_nJmom, reference())
     .property("nM", &MjModel::nM, &MjModel::set_nM, reference())
     .property("na", &MjModel::na, &MjModel::set_na, reference())