Move square root of Delassus matrix from stack to arena

PiperOrigin-RevId: 916852244
Change-Id: I4379553f808d0b238a1421606f806ea2d0a33d7c

doc/includes/references.h

@@ -191,6 +191,7 @@ struct mjData_ {
   int     nl;                // number of limit constraints
   int     nefc;              // number of constraints
   int     nJ;                // number of non-zeros in constraint Jacobian
+  int     nY;                // number of non-zeros in constraint inverse inertia square root
   int     nA;                // number of non-zeros in constraint inverse inertia matrix
   int     nisland;           // number of detected constraint islands
   int     nidof;             // number of dofs in all islands
@@ -445,8 +446,12 @@ struct mjData_ {
   mjtNum* iefc_R;            // inverse constraint mass                          (nefc x 1)
 
   // computed by mj_projectConstraint (PGS solver)
+  int*    efc_Y_rownnz;      // number of non-zeros in Y row                     (nefc x 1)
+  int*    efc_Y_rowadr;      // row start address in Y colind array              (nefc x 1)
+  int*    efc_Y_colind;      // column indices in sparse Y                       (nY x 1)
+  mjtNum* efc_Y;             // whitened Jacobian Y = J*M^(-1/2)                 (nY 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 AR colind array             (nefc x 1)
   int*    efc_AR_colind;     // column indices in sparse AR                      (nA x 1)
   mjtNum* efc_AR;            // J*inv(M)*J' + R                                  (nA x 1)
 

include/mujoco/mjdata.h

@@ -225,6 +225,7 @@ struct mjData_ {
   int     nl;                // number of limit constraints
   int     nefc;              // number of constraints
   int     nJ;                // number of non-zeros in constraint Jacobian
+  int     nY;                // number of non-zeros in constraint inverse inertia square root
   int     nA;                // number of non-zeros in constraint inverse inertia matrix
   int     nisland;           // number of detected constraint islands
   int     nidof;             // number of dofs in all islands
@@ -479,8 +480,12 @@ struct mjData_ {
   mjtNum* iefc_R;            // inverse constraint mass                          (nefc x 1)
 
   // computed by mj_projectConstraint (PGS solver)
+  int*    efc_Y_rownnz;      // number of non-zeros in Y row                     (nefc x 1)
+  int*    efc_Y_rowadr;      // row start address in Y colind array              (nefc x 1)
+  int*    efc_Y_colind;      // column indices in sparse Y                       (nY x 1)
+  mjtNum* efc_Y;             // whitened Jacobian Y = J*M^(-1/2)                 (nY 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 AR colind array             (nefc x 1)
   int*    efc_AR_colind;     // column indices in sparse AR                      (nA x 1)
   mjtNum* efc_AR;            // J*inv(M)*J' + R                                  (nA x 1)
 

include/mujoco/mjxmacro.h

@@ -946,6 +946,10 @@
 
 // array fields of mjData that are used in the dual problem
 #define MJDATA_ARENA_POINTERS_DUAL                       \
+    XNV( int,      efc_Y_rownnz,      MJ_D(nefc),    1 ) \
+    XNV( int,      efc_Y_rowadr,      MJ_D(nefc),    1 ) \
+    XNV( int,      efc_Y_colind,      MJ_D(nY),      1 ) \
+    XNV( mjtNum,   efc_Y,             MJ_D(nY),      1 ) \
     XNV( int,      efc_AR_rownnz,     MJ_D(nefc),    1 ) \
     XNV( int,      efc_AR_rowadr,     MJ_D(nefc),    1 ) \
     XNV( int,      efc_AR_colind,     MJ_D(nA),      1 ) \
@@ -1020,6 +1024,7 @@
     X( int,       nl                 ) \
     X( int,       nefc               ) \
     X( int,       nJ                 ) \
+    X( int,       nY                 ) \
     X( int,       nA                 ) \
     X( int,       nisland            ) \
     X( int,       nidof              ) \

python/mujoco/introspect/structs.py

@@ -5492,6 +5492,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  type=ValueType(name='int'),
                  doc='number of non-zeros in constraint Jacobian',
              ),
+             StructFieldDecl(
+                 name='nY',
+                 type=ValueType(name='int'),
+                 doc='number of non-zeros in constraint inverse inertia square root',  # pylint: disable=line-too-long
+             ),
              StructFieldDecl(
                  name='nA',
                  type=ValueType(name='int'),
@@ -6726,6 +6731,38 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  doc='inverse constraint mass',
                  array_extent=('nefc',),
              ),
+             StructFieldDecl(
+                 name='efc_Y_rownnz',
+                 type=PointerType(
+                     inner_type=ValueType(name='int'),
+                 ),
+                 doc='number of non-zeros in Y row',
+                 array_extent=('nefc',),
+             ),
+             StructFieldDecl(
+                 name='efc_Y_rowadr',
+                 type=PointerType(
+                     inner_type=ValueType(name='int'),
+                 ),
+                 doc='row start address in Y colind array',
+                 array_extent=('nefc',),
+             ),
+             StructFieldDecl(
+                 name='efc_Y_colind',
+                 type=PointerType(
+                     inner_type=ValueType(name='int'),
+                 ),
+                 doc='column indices in sparse Y',
+                 array_extent=('nY',),
+             ),
+             StructFieldDecl(
+                 name='efc_Y',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtNum'),
+                 ),
+                 doc='whitened Jacobian Y = J*M^(-1/2)',
+                 array_extent=('nY',),
+             ),
              StructFieldDecl(
                  name='efc_AR_rownnz',
                  type=PointerType(
@@ -6739,7 +6776,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  type=PointerType(
                      inner_type=ValueType(name='int'),
                  ),
-                 doc='row start address in colind array',
+                 doc='row start address in AR colind array',
                  array_extent=('nefc',),
              ),
              StructFieldDecl(

python/mujoco/structs_wrappers.cc

@@ -812,6 +812,7 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
   X(ne);
   X(nf);
   X(nJ);
+  X(nY);
   X(nA);
   X(nefc);
   X(nisland);
@@ -891,6 +892,7 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
   X(ne);
   X(nf);
   X(nJ);
+  X(nY);
   X(nA);
   X(nefc);
   X(nisland);

src/engine/engine_core_constraint.c

@@ -2591,12 +2591,137 @@ static int mj_nc(const mjModel* m, mjData* d, int* nnz) {
 }
 
 
+// pre-count Y_rownnz, Y_rowadr, return total nonzeros nY
+//   Y has the sparsity of J * inv(L'), where L is the Cholesky factor of M
+static int computeY_precount(int* Y_rownnz, int* Y_rowadr, int nefc, int nv,
+                             const int* J_rownnz, const int* J_rowadr, const int* J_colind,
+                             const int* M_rownnz, const int* M_rowadr, const int* M_colind,
+                             int* marker) {
+  mju_fillInt(marker, -1, nv);
+
+  Y_rowadr[0] = 0;
+  for (int r=0; r < nefc; r++) {
+    int nnz = 0;  // nonzeros in row r of Y
+
+    // traverse row r of J in reverse, count unique nonzeros
+    int start = J_rowadr[r];
+    int end = start + J_rownnz[r];
+    for (int i=end-1; i >= start; i--) {
+      int j = J_colind[i];
+
+      // if dof j is marked, it was already counted by a child dof: skip it
+      if (marker[j] == r) {
+        continue;
+      }
+
+      // traverse row j of M, marking new unique nonzeros
+      int nnzM = M_rownnz[j];
+      int adrM = M_rowadr[j];
+      for (int k=0; k < nnzM; k++) {
+        int c = M_colind[adrM + k];
+        if (marker[c] != r) {
+          marker[c] = r;
+          nnz++;
+        }
+      }
+    }
+
+    // update rownnz and rowadr
+    Y_rownnz[r] = nnz;
+    if (r < nefc - 1) {
+      Y_rowadr[r+1] = Y_rowadr[r] + nnz;
+    }
+  }
+
+  // total non-zeros in Y
+  return Y_rowadr[nefc-1] + Y_rownnz[nefc-1];
+}
+
+
+// fill Y column indices and values from J, chaining up the kinematic tree
+static void computeY_fill(mjtNum* Y, int* Y_colind,
+                          const int* Y_rownnz, const int* Y_rowadr, int nefc,
+                          const mjtNum* J, const int* J_rownnz, const int* J_rowadr,
+                          const int* J_colind, const int* dof_parentid) {
+  for (int r=0; r < nefc; r++) {
+    // init row
+    int end = Y_rowadr[r] + Y_rownnz[r];
+    int adrJ = J_rowadr[r];
+    int remainJ = J_rownnz[r];
+    int nnzY = 0;
+
+    // complete chain in reverse
+    while (1) {
+      // get previous dof in src and dst
+      int prev_src = (remainJ > 0 ? J_colind[adrJ + remainJ - 1] : -1);
+      int prev_dst = (nnzY > 0 ? dof_parentid[Y_colind[end - nnzY]] : -1);
+
+      // both finished: break
+      if (prev_src < 0 && prev_dst < 0) {
+        break;
+      }
+
+      // add src
+      else if (prev_src >= prev_dst) {
+        nnzY++;
+        remainJ--;
+        Y_colind[end - nnzY] = prev_src;
+        Y[end - nnzY] = J[adrJ + remainJ];
+      }
+
+      // add dst
+      else {
+        nnzY++;
+        Y_colind[end - nnzY] = prev_dst;
+        Y[end - nnzY] = 0;
+      }
+    }
+
+    // compare with Y_rownnz: SHOULD NOT OCCUR
+    if (nnzY != Y_rownnz[r]) {
+      mjERROR("pre and post-count of Y_rownnz are not equal on row %d", r);
+    }
+  }
+}
+
+
+// in-place sparse back-substitution:  Y <- Y * M^{-1/2}
+static void computeY_backsub(mjtNum* Y, const int* Y_rownnz, const int* Y_rowadr,
+                             const int* Y_colind, int nefc,
+                             const mjtNum* qLD, const int* M_rownnz, const int* M_rowadr,
+                             const int* M_colind, const mjtNum* sqrtInvD) {
+  for (int r=0; r < nefc; r++) {
+    int nnzY = Y_rownnz[r];
+    int adrY = Y_rowadr[r];
+
+    // Y(r,:) <- inv(L') * Y(r,:), exploit sparsity of input vector
+    for (int i=adrY + nnzY-1; i >= adrY; i--) {
+      mjtNum val = Y[i];
+      if (val == 0) {
+        continue;
+      }
+      int j = Y_colind[i];
+      int adrM = M_rowadr[j];
+      mju_addToSclSparseInc(Y + adrY, qLD + adrM,
+                            nnzY, Y_colind + adrY,
+                            M_rownnz[j]-1, M_colind + adrM, -val);
+    }
+
+    // Y(r,:) <- sqrt(inv(D)) * Y(r,:)
+    for (int i=adrY; i < adrY + nnzY; i++) {
+      int j = Y_colind[i];
+      Y[i] *= sqrtInvD[j];
+    }
+  }
+}
+
+
 //---------------------------- 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->nl = d->nefc = d->nJ = d->nA = 0;
+  d->ne = d->nf = d->nl = d->nefc = d->nJ = d->nA = d->nY = 0;
 
   // disabled or Jacobian not allocated: return
   if (mjDISABLED(mjDSBL_CONSTRAINT)) {
@@ -2705,177 +2830,60 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
     sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]);
   }
 
-  // sparse
+  // sparse Y = backsubM2(J')' and its transpose
   if (mj_isSparse(m)) {
-    // compute B = backsubM2(J')' and its transpose
+    // arena-allocate Y rownnz and rowadr
-
+    d->efc_Y_rownnz = mj_arenaAllocByte(d, sizeof(int) * nefc, _Alignof(int));
-
+    d->efc_Y_rowadr = mj_arenaAllocByte(d, sizeof(int) * nefc, _Alignof(int));
-    // === pre-count B_rownnz, B_rowadr, nB (total nonzeros)
+    if (!d->efc_Y_rownnz || !d->efc_Y_rowadr) {
-
+      mj_warning(d, mjWARN_CNSTRFULL, d->narena);
-    // allocate B rownnz and rowadr
+      mj_clearEfc(d);
-    int* B_rownnz = mjSTACKALLOC(d, nefc, int);
+      d->parena = d->ncon * sizeof(mjContact);
-    int* B_rowadr = mjSTACKALLOC(d, nefc, int);
+      mj_freeStack(d);
+      return;
+    }
 
     // markers for merged dofs, initialized to -1
     int* marker = mjSTACKALLOC(d, nv, int);
-    mju_fillInt(marker, -1, nv);
 
-    B_rowadr[0] = 0;
+    // pre-count Y_rownnz, Y_rowadr, nY (total nonzeros)
-    for (int r=0; r < nefc; r++) {
+    d->nY = computeY_precount(d->efc_Y_rownnz, d->efc_Y_rowadr, nefc, nv,
-      // supernode: same sparsity as previous row
+                              d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
-      if (r > 0 && d->efc_J_rowsuper[r-1] > 0) {
+                              m->M_rownnz, m->M_rowadr, m->M_colind, marker);
-        B_rownnz[r] = B_rownnz[r-1];
-      }
 
-      // first row in supernode block: full chain traversal
+    // arena-allocate values and column indices
-      else {
+    d->efc_Y = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nY, _Alignof(mjtNum));
-        int nnz = 0;
+    d->efc_Y_colind = mj_arenaAllocByte(d, sizeof(int) * d->nY, _Alignof(int));
-
+    if (!d->efc_Y || !d->efc_Y_colind) {
-        // traverse row r of J in reverse, count unique nonzeros
+      mj_warning(d, mjWARN_CNSTRFULL, d->narena);
-        int start = d->efc_J_rowadr[r];
+      mj_clearEfc(d);
-        int end = start + d->efc_J_rownnz[r];
+      d->parena = d->ncon * sizeof(mjContact);
-        for (int i=end-1; i >= start; i--) {
+      mj_freeStack(d);
-          int j = d->efc_J_colind[i];
+      return;
-
-          // if dof j is marked, it was already counted by a child dof: skip it
-          if (marker[j] == r) {
-            continue;
-          }
-
-          // traverse row j of M, marking new unique nonzeros
-          int nnzM = m->M_rownnz[j];
-          int adrM = m->M_rowadr[j];
-          for (int k=0; k < nnzM; k++) {
-            int c = m->M_colind[adrM + k];
-            if (marker[c] != r) {
-              marker[c] = r;
-              nnz++;
-            }
-          }
-        }
-        B_rownnz[r] = nnz;
-      }
-
-      // update rowadr
-      if (r < nefc - 1) {
-        B_rowadr[r+1] = B_rowadr[r] + B_rownnz[r];
-      }
     }
 
-    // total non-zeros in B
+    // fill in Y column indices, copy values from J
-    int nB = B_rowadr[nefc-1] + B_rownnz[nefc-1];
+    computeY_fill(d->efc_Y, d->efc_Y_colind, d->efc_Y_rownnz, d->efc_Y_rowadr, nefc,
+                  d->efc_J, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
+                  m->dof_parentid);
 
 
-    // === fill in B column indices, copy values from J
+    // in-place sparse back-substitution:  Y <- Y * M^-1/2
+    computeY_backsub(d->efc_Y, d->efc_Y_rownnz, d->efc_Y_rowadr,
+                     d->efc_Y_colind, nefc,
+                     d->qLD, m->M_rownnz, m->M_rowadr, m->M_colind, sqrtInvD);
 
-    // allocate values and column indices
+    // Y supernodes are identical to J supernodes
-    mjtNum* B = mjSTACKALLOC(d, nB, mjtNum);
+    const int* Y_rowsuper = d->efc_J_rowsuper;
-    int* B_colind = mjSTACKALLOC(d, nB, int);
 
-    for (int r=0; r < nefc; r++) {
+    // construct Y transposed
-      // supernode: copy column indices, only update values from J
+    int* YT_rownnz = mjSTACKALLOC(d, nv, int);
-      if (r > 0 && d->efc_J_rowsuper[r-1] > 0) {
+    int* YT_rowadr = mjSTACKALLOC(d, nv, int);
-        int prevAdr = B_rowadr[r-1];
+    int* YT_colind = mjSTACKALLOC(d, d->nY, int);
-        int adrB = B_rowadr[r];
+    mjtNum* YT = mjSTACKALLOC(d, d->nY, mjtNum);
-        int nnzB = B_rownnz[r];
+    mju_transposeSparse(YT, d->efc_Y, nefc, nv,
-        mju_copyInt(B_colind + adrB, B_colind + prevAdr, nnzB);
+                        YT_rownnz, YT_rowadr, YT_colind, NULL,
-        mju_zero(B + adrB, nnzB);
+                        d->efc_Y_rownnz, d->efc_Y_rowadr, d->efc_Y_colind);
-
-        // copy J values into correct positions
-        int adrJ = d->efc_J_rowadr[r];
-        int jnnz = d->efc_J_rownnz[r];
-        int bi = 0, ji = 0;
-        while (ji < jnnz && bi < nnzB) {
-          if (B_colind[adrB+bi] == d->efc_J_colind[adrJ+ji]) {
-            B[adrB+bi] = d->efc_J[adrJ+ji];
-            bi++;
-            ji++;
-          } else {
-            bi++;
-          }
-        }
-      }
-
-      // first row in supernode block: full chain completion
-      else {
-        int end = B_rowadr[r] + B_rownnz[r];
-        int adrJ = d->efc_J_rowadr[r];
-        int remainJ = d->efc_J_rownnz[r];
-        int nnzB = 0;
-
-        // complete chain in reverse
-        while (1) {
-          // get previous dof in src and dst
-          int prev_src = (remainJ > 0 ? d->efc_J_colind[adrJ + remainJ - 1] : -1);
-          int prev_dst = (nnzB > 0 ? m->dof_parentid[B_colind[end - nnzB]] : -1);
-
-          // both finished: break
-          if (prev_src < 0 && prev_dst < 0) {
-            break;
-          }
-
-          // add src
-          else if (prev_src >= prev_dst) {
-            nnzB++;
-            remainJ--;
-            B_colind[end - nnzB] = prev_src;
-            B[end - nnzB] = d->efc_J[adrJ + remainJ];
-          }
-
-          // add dst
-          else {
-            nnzB++;
-            B_colind[end - nnzB] = prev_dst;
-            B[end - nnzB] = 0;
-          }
-        }
-
-        // compare with B_rownnz: SHOULD NOT OCCUR
-        if (nnzB != B_rownnz[r]) {
-          mjERROR("pre and post-count of B_rownnz are not equal on row %d", r);
-        }
-      }
-    }
-
-
-    // === in-place sparse back-substitution:  B <- B * M^-1/2
-
-    // sparse backsubM2 (half of LD back-substitution)
-    for (int r=0; r < nefc; r++) {
-      int nnzB = B_rownnz[r];
-      int adrB = B_rowadr[r];
-
-      // B(r,:) <- inv(L') * B(r,:), exploit sparsity of input vector
-      for (int i=adrB + nnzB-1; i >= adrB; i--) {
-        mjtNum b = B[i];
-        if (b == 0) {
-          continue;
-        }
-        int j = B_colind[i];
-        int adrC = m->M_rowadr[j];
-        mju_addToSclSparseInc(B + adrB, d->qLD + adrC,
-                              nnzB, B_colind + adrB,
-                              m->M_rownnz[j]-1, m->M_colind + adrC, -b);
-      }
-
-      // B(r,:) <- sqrt(inv(D)) * B(r,:)
-      for (int i=adrB; i < adrB + nnzB; i++) {
-        int j = B_colind[i];
-        B[i] *= sqrtInvD[j];
-      }
-    }
-
-    // B supernodes are identical to J supernodes
-    const int* B_rowsuper = d->efc_J_rowsuper;
-
-    // construct B transposed
-    int* BT_rownnz = mjSTACKALLOC(d, nv, int);
-    int* BT_rowadr = mjSTACKALLOC(d, nv, int);
-    int* BT_colind = mjSTACKALLOC(d, nB, int);
-    mjtNum* BT = mjSTACKALLOC(d, nB, mjtNum);
-    mju_transposeSparse(BT, B, nefc, nv,
-                        BT_rownnz, BT_rowadr, BT_colind, NULL,
-                        B_rownnz, B_rowadr, B_colind);
 
     // allocate AR row nonzeros and addresses on arena
     d->efc_AR_rownnz = mj_arenaAllocByte(d, sizeof(int) * nefc, _Alignof(int));
@@ -2891,8 +2899,8 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
     int* diagind = mjSTACKALLOC(d, nefc, int);
     d->nA = mju_sqrMatTDSparseSymbolic(
         d->efc_AR_rownnz, d->efc_AR_rowadr, NULL, diagind,
-        nv, nefc, BT_rownnz, BT_rowadr, BT_colind,
+        nv, nefc, YT_rownnz, YT_rowadr, YT_colind,
-        B_rownnz, B_rowadr, B_colind, B_rowsuper, d);
+        d->efc_Y_rownnz, d->efc_Y_rowadr, d->efc_Y_colind, Y_rowsuper, d);
 
     // allocate A values and column indices on arena
     d->efc_AR = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nA, _Alignof(mjtNum));
@@ -2905,17 +2913,18 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
       return;
     }
 
-    // A = B * B': symbolic phase
+    // A = Y * Y': symbolic phase
     mju_sqrMatTDSparseSymbolic(
         d->efc_AR_rownnz, d->efc_AR_rowadr, d->efc_AR_colind, diagind,
-        nv, nefc, BT_rownnz, BT_rowadr, BT_colind,
+        nv, nefc, YT_rownnz, YT_rowadr, YT_colind,
-        B_rownnz, B_rowadr, B_colind, B_rowsuper, d);
+        d->efc_Y_rownnz, d->efc_Y_rowadr, d->efc_Y_colind, Y_rowsuper, d);
 
-    // A = B * B': numeric phase
+    // A = Y * Y': numeric phase
     mju_sqrMatTDSparseNumeric(
         d->efc_AR, nefc, d->efc_AR_rownnz, d->efc_AR_rowadr,
-        d->efc_AR_colind, diagind, BT, BT_rownnz, BT_rowadr,
+        d->efc_AR_colind, diagind, YT, YT_rownnz, YT_rowadr,
-        BT_colind, B, B_rownnz, B_rowadr, B_colind, B_rowsuper, NULL, d);
+        YT_colind, d->efc_Y, d->efc_Y_rownnz, d->efc_Y_rowadr,
+        d->efc_Y_colind, Y_rowsuper, NULL, d);
 
     // AR = A + diag(R)
     for (int i=0; i < nefc; i++) {
@@ -2923,11 +2932,24 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
     }
   }
 
-  // dense
+  // dense Y = backsubM2(J')' and its transpose
   else {
-    d->nA = nefc * nefc;
+    // arena-allocate efc_Y
+    d->nY = nefc * nv;
+    d->efc_Y = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nY, _Alignof(mjtNum));
+    if (!d->efc_Y) {
+      mj_warning(d, mjWARN_CNSTRFULL, d->narena);
+      mj_clearEfc(d);
+      d->parena = d->ncon * sizeof(mjContact);
+      mj_freeStack(d);
+      return;
+    }
+
+    // Y = backsubM2(J')'
+    mj_solveM2(m, d, d->efc_Y, d->efc_J, sqrtInvD, nefc);
 
     // arena-allocate efc_AR
+    d->nA = nefc * nefc;
     d->efc_AR = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nA, _Alignof(mjtNum));
     if (!d->efc_AR) {
       mj_warning(d, mjWARN_CNSTRFULL, d->narena);
@@ -2937,18 +2959,12 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
       return;
     }
 
-    // space for B = backsubM2(J')' and its transpose
+    // construct YT on stack
-    mjtNum* B = mjSTACKALLOC(d, nefc*nv, mjtNum);
+    mjtNum* YT = mjSTACKALLOC(d, nv*nefc, mjtNum);
-    mjtNum* BT = mjSTACKALLOC(d, nv*nefc, mjtNum);
+    mju_transpose(YT, d->efc_Y, nefc, nv);
 
-    // B = backsubM2(J')'
+    // AR = Y * Y'
-    mj_solveM2(m, d, B, d->efc_J, sqrtInvD, nefc);
+    mju_sqrMatTD(d->efc_AR, YT, NULL, nv, nefc);
-
-    // construct BT
-    mju_transpose(BT, B, nefc, nv);
-
-    // AR = B * B'
-    mju_sqrMatTD(d->efc_AR, BT, NULL, nv, nefc);
 
     // add R to diagonal of AR
     for (int r=0; r < nefc; r++) {

src/engine/engine_io.c

@@ -1320,6 +1320,7 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
   d->nl = 0;
   d->nefc = 0;
   d->nJ = 0;
+  d->nY = 0;
   d->nA = 0;
   d->nisland = 0;
   d->nidof = 0;

src/engine/engine_print.c

@@ -1192,7 +1192,7 @@ void mj_printFormattedModel(const mjModel* m, const char* filename, const char*
 
   // BVHs
   fprintf(fp, "BVH:\n");
-  fprintf(fp, "  %-8s%-8s%-8s%-10s%-s\n","id", "depth", "nodeid", "child[0]" ,"child[1]");
+  fprintf(fp, "  %-8s%-8s%-8s%-10s%-s\n", "id", "depth", "nodeid", "child[0]", "child[1]");
   for (int i=0; i < m->nbvh; i++) {
     fprintf(fp, "  %-8d%-8d% -8d% -10d% -d\n",
             i, m->bvh_depth[i], m->bvh_nodeid[i], m->bvh_child[2*i], m->bvh_child[2*i+1]);
@@ -1204,7 +1204,6 @@ void mj_printFormattedModel(const mjModel* m, const char* filename, const char*
   }
 }
 
-
 // print mjModel to text file
 void mj_printModel(const mjModel* m, const char* filename) {
   mj_printFormattedModel(m, filename, FLOAT_FORMAT);
@@ -1545,7 +1544,6 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
     mjtNum force[6] = {0};
     mj_contactForce(m, d, i, force);
     printVector("     force        ", force, 6, fp, float_format);
-
   }
   if (d->ncon) fprintf(fp, "\n");
 
@@ -1568,6 +1566,11 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
                 d->efc_J_rowadr, d->efc_J_colind, fp, float_format);
     mj_printSparsity("J: constraint Jacobian", d->nefc, m->nv, d->efc_J_rowadr, NULL,
                      d->efc_J_rownnz, d->efc_J_rowsuper, d->efc_J_colind, fp);
+    if (d->nY) {
+      mj_printSparsity("EFC_Y: inverse constraint inertia square root", d->nefc, m->nv,
+                       d->efc_Y_rowadr, NULL, d->efc_Y_rownnz, d->efc_J_rowsuper,
+                       d->efc_Y_colind, fp);
+    }
     if (d->nisland) {
       mj_printBlockSparsity("IEFC_J: block-diagonalized constraint Jacobian (nnzs are island ids)",
                             d->nefc, d->nidof, d->nisland,
@@ -1582,7 +1585,7 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
       printArray2dInt("EFC_AR_ROWADR", d->nefc, 1, d->efc_AR_rowadr, fp);
       printSparse("EFC_AR", d->efc_AR, d->nefc, d->efc_AR_rownnz,
                   d->efc_AR_rowadr, d->efc_AR_colind, fp, float_format);
-      mj_printSparsity("efc_AR: inverse constraint inertia", d->nefc, d->nefc, d->efc_AR_rowadr,
+      mj_printSparsity("EFC_AR: inverse constraint inertia", d->nefc, d->nefc, d->efc_AR_rowadr,
                        NULL, d->efc_AR_rownnz, NULL, d->efc_AR_colind, fp);
     }
   }

unity/Runtime/Bindings/MjBindings.cs

@@ -5724,6 +5724,7 @@ public unsafe struct mjData_ {
   public int nl;
   public int nefc;
   public int nJ;
+  public int nY;
   public int nA;
   public int nisland;
   public int nidof;
@@ -5883,6 +5884,10 @@ public unsafe struct mjData_ {
   public double* iefc_frictionloss;
   public double* iefc_D;
   public double* iefc_R;
+  public int* efc_Y_rownnz;
+  public int* efc_Y_rowadr;
+  public int* efc_Y_colind;
+  public double* efc_Y;
   public int* efc_AR_rownnz;
   public int* efc_AR_rowadr;
   public int* efc_AR_colind;

wasm/codegen/generated/bindings.cc

@@ -6494,6 +6494,12 @@ struct MjData {
   void set_nJ(int value) {
     ptr_->nJ = value;
   }
+  int nY() const {
+    return ptr_->nY;
+  }
+  void set_nY(int value) {
+    ptr_->nY = value;
+  }
   int nA() const {
     return ptr_->nA;
   }
@@ -7005,6 +7011,18 @@ struct MjData {
   emscripten::val iefc_R() const {
     return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->iefc_R));
   }
+  emscripten::val efc_Y_rownnz() const {
+    return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->efc_Y_rownnz));
+  }
+  emscripten::val efc_Y_rowadr() const {
+    return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->efc_Y_rowadr));
+  }
+  emscripten::val efc_Y_colind() const {
+    return emscripten::val(emscripten::typed_memory_view(ptr_->nY, ptr_->efc_Y_colind));
+  }
+  emscripten::val efc_Y() const {
+    return emscripten::val(emscripten::typed_memory_view(ptr_->nY, ptr_->efc_Y));
+  }
   emscripten::val efc_AR_rownnz() const {
     return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->efc_AR_rownnz));
   }
@@ -11547,6 +11565,10 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
     .property("efc_J_rowsuper", &MjData::efc_J_rowsuper)
     .property("efc_KBIP", &MjData::efc_KBIP)
     .property("efc_R", &MjData::efc_R)
+    .property("efc_Y", &MjData::efc_Y)
+    .property("efc_Y_colind", &MjData::efc_Y_colind)
+    .property("efc_Y_rowadr", &MjData::efc_Y_rowadr)
+    .property("efc_Y_rownnz", &MjData::efc_Y_rownnz)
     .property("efc_aref", &MjData::efc_aref)
     .property("efc_b", &MjData::efc_b)
     .property("efc_diagApprox", &MjData::efc_diagApprox)
@@ -11626,6 +11648,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
     .property("moment_rownnz", &MjData::moment_rownnz)
     .property("nA", &MjData::nA, &MjData::set_nA, reference())
     .property("nJ", &MjData::nJ, &MjData::set_nJ, reference())
+    .property("nY", &MjData::nY, &MjData::set_nY, reference())
     .property("narena", &MjData::narena, &MjData::set_narena, reference())
     .property("nbody_awake", &MjData::nbody_awake, &MjData::set_nbody_awake, reference())
     .property("nbuffer", &MjData::nbuffer, &MjData::set_nbuffer, reference())

wasm/codegen/generators/constants.py

@@ -315,6 +315,10 @@ MJDATA_SIZES: tuple[str, ...] = (
     "efc_AR_colind",
     "efc_AR_rowadr",
     "efc_AR_rownnz",
+    "efc_Y",
+    "efc_Y_colind",
+    "efc_Y_rowadr",
+    "efc_Y_rownnz",
     "efc_D",
     "efc_J",
     "efc_JT",