Add mjENBL_DIAGEXACT for exact constraint diagonal. Fixes #2472

PiperOrigin-RevId: 916932908
Change-Id: Id23ac39b5cd996afc52990719a4e07c0cc7de600

src/engine/engine_core_constraint.c

@@ -2812,15 +2812,10 @@ void mj_makeConstraint(const mjModel* m, mjData* d) {
 }
 
 
-// compute efc_AR
-void mj_projectConstraint(const mjModel* m, mjData* d) {
+// compute Y = J*M^{-1/2}; if flg_diagexact, overwrite efc_diagApprox with ||Y_i||^2
+static void mj_makeY(const mjModel* m, mjData* d, int flg_diagexact) {
   int nefc = d->nefc, nv = m->nv;
 
-  // nothing to do
-  if (nefc == 0 || !mj_isDual(m)) {
-    return;
-  }
-
   mj_markStack(d);
 
   // inverse square root of D from inertia LDL decomposition
@@ -2843,10 +2838,8 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
       return;
     }
 
-    // markers for merged dofs, initialized to -1
-    int* marker = mjSTACKALLOC(d, nv, int);
-
     // pre-count Y_rownnz, Y_rowadr, nY (total nonzeros)
+    int* marker = mjSTACKALLOC(d, nv, int);
     d->nY = computeY_precount(d->efc_Y_rownnz, d->efc_Y_rowadr, nefc, nv,
                               d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
                               m->M_rownnz, m->M_rowadr, m->M_colind, marker);
@@ -2867,12 +2860,57 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
                   d->efc_J, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
                   m->dof_parentid);
 
-
     // 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);
 
+    // overwrite diagApprox with exact diagonal: diagApprox[i] = ||Y_i||^2
+    if (flg_diagexact) {
+      for (int i=0; i < nefc; i++) {
+        int adr = d->efc_Y_rowadr[i];
+        int nnz = d->efc_Y_rownnz[i];
+        d->efc_diagApprox[i] = mju_dot(d->efc_Y+adr, d->efc_Y+adr, nnz);
+      }
+    }
+  }
+
+  // dense Y = backsubM2(J')' and its transpose
+  else {
+    // 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);
+
+    // overwrite diagApprox with exact diagonal: diagApprox[i] = ||Y_i||^2
+    if (flg_diagexact) {
+      for (int i=0; i < nefc; i++) {
+        d->efc_diagApprox[i] = mju_dot(d->efc_Y+i*nv, d->efc_Y+i*nv, nv);
+      }
+    }
+  }
+
+  mj_freeStack(d);
+}
+
+
+// assemble AR = Y*Y' + diag(R) for dual solver
+static void mj_makeAR(const mjModel* m, mjData* d) {
+  int nefc = d->nefc, nv = m->nv;
+
+  mj_markStack(d);
+
+  // sparse
+  if (mj_isSparse(m)) {
     // Y supernodes are identical to J supernodes
     const int* Y_rowsuper = d->efc_J_rowsuper;
 
@@ -2934,20 +2972,6 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
 
   // dense Y = backsubM2(J')' and its transpose
   else {
-    // 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));
@@ -2976,6 +3000,41 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
 }
 
 
+// compute efc_Y, optionally efc_diagApprox, optionally efc_AR
+void mj_projectConstraint(const mjModel* m, mjData* d) {
+  int nefc = d->nefc;
+
+  // nothing to do
+  if (!nefc) {
+    return;
+  }
+
+  int isDual = mj_isDual(m);
+  int diagexact = mjENABLED(mjENBL_DIAGEXACT);
+
+  // compute Y = J*M^{-1/2}; overwrite diagApprox if diagexact
+  if (isDual || diagexact) {
+    mj_makeY(m, d, diagexact);
+  }
+
+  // recompute impedance from exact diagonal
+  if (diagexact && d->nefc) {
+    mj_makeImpedance(m, d);
+
+    // re-gather island D/R
+    if (d->nisland) {
+      mju_gather(d->iefc_D, d->efc_D, d->map_iefc2efc, d->nefc);
+      mju_gather(d->iefc_R, d->efc_R, d->map_iefc2efc, d->nefc);
+    }
+  }
+
+  // assemble AR for dual solver
+  if (isDual && d->nefc) {
+    mj_makeAR(m, d);
+  }
+}
+
+
 // compute efc_vel, efc_aref
 void mj_referenceConstraint(const mjModel* m, mjData* d) {
   int nefc = d->nefc;

src/engine/engine_forward.c

@@ -207,14 +207,14 @@ void mj_fwdPosition(const mjModel* m, mjData* d) {
   mj_island(m, d);
   TM_END(mjTIMER_POS_MAKE);
 
-  TM_RESTART;
-  mj_transmission(m, d);
-  TM_ADD(mjTIMER_POS_KINEMATICS);
-
   TM_RESTART;
   mj_projectConstraint(m, d);
   TM_END(mjTIMER_POS_PROJECT);
 
+  TM_RESTART;
+  mj_transmission(m, d);
+  TM_ADD(mjTIMER_POS_KINEMATICS);
+
   TM_END1(mjTIMER_POSITION);
 }
 

src/engine/engine_inverse.c

@@ -59,6 +59,13 @@ void mj_invPosition(const mjModel* m, mjData* d) {
   mj_makeConstraint(m, d);
   TM_END(mjTIMER_POS_MAKE);
 
+  // compute exact diagonal if enabled
+  if (mjENABLED(mjENBL_DIAGEXACT)) {
+    TM_RESTART;
+    mj_projectConstraint(m, d);
+    TM_END(mjTIMER_POS_PROJECT);
+  }
+
   TM_RESTART;
   mj_transmission(m, d);
   TM_ADD(mjTIMER_POS_KINEMATICS);

src/engine/engine_memory.h

@@ -76,6 +76,7 @@ static inline void mj_clearEfc(mjData* d) {
 #undef X
   d->nefc = 0;
   d->nisland = 0;
+  d->nJ = d->nY = d->nA = 0;
   d->contact = (mjContact*) d->arena;
 
   // if any contacts are allocated, clear their efc_address

src/engine/engine_support.c

@@ -77,7 +77,8 @@ const char* mjENABLESTRING[mjNENABLE] = {
   "Energy",
   "Fwdinv",
   "InvDiscrete",
-  "Sleep"
+  "Sleep",
+  "DiagExact"
 };