Remove mjData.qLDiagSqrtInv, add corresponding argument to mj_solveM2.

yuvaltassa · copybara-github · commit 69c9ac074a7f · 2024-12-30T15:24:36.000-08:00
- `qLDiagSqrtInv` is only required for the dual solvers. It is now computed as-needed rather than unconditionally.
- `mj_solveM2` now requires a new input array `sqrtInvD` which contains the square root of the inverse diagonal D (formerly saved in `qLDiagSqrtInv`).

PiperOrigin-RevId: 710805133
Change-Id: I0622d6a8da3882916824e9c10bad9223c122c321
diff --git a/doc/changelog.rst b/doc/changelog.rst
@@ -13,6 +13,15 @@ Python bindings
   from multiple threads. To run multiple threaded rollouts simultaneously, use the new class ``Rollout`` which
   encapsulates the thread pool. Contribution by :github:user:`aftersomemath`.
 
+General
+^^^^^^^
+
+.. admonition:: Breaking API changes (minor)
+   :class: attention
+
+   - The field ``mjData.qLDiagSqrtInv`` has been removed. This field is only required for the dual solvers. It is now
+     computed as-needed rather than unconditionally. Relatedly, added the corresponding argument to :ref:`mj_solveM2`.
+
 Bug fixes
 ^^^^^^^^^
 - Fixed a bug in the box-sphere collider, depth was incorrect for deep penetrations (:github:issue:`2206`).
diff --git a/doc/includes/references.h b/doc/includes/references.h
@@ -273,7 +273,6 @@ struct mjData_ {
   // computed by mj_fwdPosition/mj_factorM
   mjtNum* qLD;               // L'*D*L factorization of M (sparse)               (nM x 1)
   mjtNum* qLDiagInv;         // 1/diag(D)                                        (nv x 1)
-  mjtNum* qLDiagSqrtInv;     // 1/sqrt(diag(D))                                  (nv x 1)
 
   // computed by mj_collisionTree
   mjtNum*  bvh_aabb_dyn;     // global bounding box (center, size)               (nbvhdynamic x 6)
@@ -3221,7 +3220,8 @@ void mj_transmission(const mjModel* m, mjData* d);
 void mj_crb(const mjModel* m, mjData* d);
 void mj_factorM(const mjModel* m, mjData* d);
 void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
-void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
+void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
+                const mjtNum* sqrtInvD, int n);
 void mj_comVel(const mjModel* m, mjData* d);
 void mj_passive(const mjModel* m, mjData* d);
 void mj_subtreeVel(const mjModel* m, mjData* d);
diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h
@@ -301,7 +301,6 @@ struct mjData_ {
   // computed by mj_fwdPosition/mj_factorM
   mjtNum* qLD;               // L'*D*L factorization of M (sparse)               (nM x 1)
   mjtNum* qLDiagInv;         // 1/diag(D)                                        (nv x 1)
-  mjtNum* qLDiagSqrtInv;     // 1/sqrt(diag(D))                                  (nv x 1)
 
   // computed by mj_collisionTree
   mjtNum*  bvh_aabb_dyn;     // global bounding box (center, size)               (nbvhdynamic x 6)
diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h
@@ -632,7 +632,6 @@
     X   ( mjtNum,    qM,                nM,          1           ) \
     X   ( mjtNum,    qLD,               nM,          1           ) \
     X   ( mjtNum,    qLDiagInv,         nv,          1           ) \
-    X   ( mjtNum,    qLDiagSqrtInv,     nv,          1           ) \
     XMJV( mjtNum,    bvh_aabb_dyn,      nbvhdynamic, 6           ) \
     XMJV( mjtByte,   bvh_active,        nbvh,        1           ) \
     X   ( mjtNum,    flexedge_velocity, nflexedge,   1           ) \
diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h
@@ -365,7 +365,8 @@ MJAPI void mj_factorM(const mjModel* m, mjData* d);
 MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
 
 // Half of linear solve:  x = sqrt(inv(D))*inv(L')*y
-MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
+MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
+                      const mjtNum* sqrtInvD, int n);
 
 // Compute cvel, cdof_dot.
 MJAPI void mj_comVel(const mjModel* m, mjData* d);
diff --git a/introspect/functions.py b/introspect/functions.py
@@ -1868,6 +1868,12 @@
                      inner_type=ValueType(name='mjtNum', is_const=True),
                  ),
              ),
+             FunctionParameterDecl(
+                 name='sqrtInvD',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtNum', is_const=True),
+                 ),
+             ),
              FunctionParameterDecl(
                  name='n',
                  type=ValueType(name='int'),
diff --git a/introspect/structs.py b/introspect/structs.py
@@ -5213,14 +5213,6 @@
                  doc='1/diag(D)',
                  array_extent=('nv',),
              ),
-             StructFieldDecl(
-                 name='qLDiagSqrtInv',
-                 type=PointerType(
-                     inner_type=ValueType(name='mjtNum'),
-                 ),
-                 doc='1/sqrt(diag(D))',
-                 array_extent=('nv',),
-             ),
              StructFieldDecl(
                  name='bvh_aabb_dyn',
                  type=PointerType(
diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
@@ -310,7 +310,6 @@ def make_data(
         'qM': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float),
         'qLD': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float),
         'qLDiagInv': (m.nv, float) if support.is_sparse(m) else (0, float),
-        'qLDiagSqrtInv': (m.nv, float),
         'bvh_aabb_dyn': (m.nbvhdynamic, 6, float),
         'bvh_active': (m.nbvh, jp.uint8),
         'flexedge_velocity': (m.nflexedge, float),
diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
@@ -1274,7 +1274,6 @@ class Data(PyTreeNode):
                                                      if dense:  (nv, nv)
     qLDiagInv: 1/diag(D)                             if sparse: (nv,)
                                                      if dense:  (0,)
-    qLDiagSqrtInv: 1/sqrt(diag(D))                              (nv,)
     bvh_aabb_dyn: global bounding box (center, size)            (nbvhdynamic, 6)
     bvh_active: volume has been added to collisions             (nbvh,)
     flexedge_velocity: flex edge velocities                     (nflexedge,)
@@ -1404,7 +1403,6 @@ class Data(PyTreeNode):
   qM: jax.Array  # pylint:disable=invalid-name
   qLD: jax.Array  # pylint:disable=invalid-name
   qLDiagInv: jax.Array  # pylint:disable=invalid-name
-  qLDiagSqrtInv: jax.Array  # pylint:disable=invalid-name
   bvh_aabb_dyn: jax.Array = _restricted_to('mujoco')
   bvh_active: jax.Array = _restricted_to('mujoco')
   # position, velocity dependent:
diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc
@@ -234,7 +234,7 @@ PYBIND11_MODULE(_functions, pymodule) {
   DEF_WITH_OMITTED_PY_ARGS(traits::mj_solveM2, "n")(
       pymodule,
       [](const raw::MjModel* m, raw::MjData* d, Eigen::Ref<EigenArrayXX> x,
-         Eigen::Ref<const EigenArrayXX> y) {
+         Eigen::Ref<const EigenArrayXX> y, Eigen::Ref<const EigenArrayXX> sqrtInvD) {
         if (x.rows() != y.rows()) {
           throw py::type_error(
               "the first dimension of x and y should be of the same size");
@@ -247,8 +247,12 @@ PYBIND11_MODULE(_functions, pymodule) {
           throw py::type_error(
               "the last dimension of y should be of size nv");
         }
+        if (sqrtInvD.size() != m->nv) {
+          throw py::type_error(
+              "the size of sqrtInvD should be nv");
+        }
         return InterceptMjErrors(::mj_solveM2)(
-            m, d, x.data(), y.data(), y.rows());
+            m, d, x.data(), y.data(), sqrtInvD.data(), y.rows());
       });
   Def<traits::mj_comVel>(pymodule);
   Def<traits::mj_passive>(pymodule);
diff --git a/python/mujoco/indexer_xmacro.h b/python/mujoco/indexer_xmacro.h
@@ -379,7 +379,6 @@
   X( mjtNum, , xaxis,           njnt, 3           ) \
   X( mjtNum, , cdof,            nv,   6           ) \
   X( mjtNum, , qLDiagInv,       nv,   1           ) \
-  X( mjtNum, , qLDiagSqrtInv,   nv,   1           ) \
   X( mjtNum, , cdof_dot,        nv,   6           ) \
   X( mjtNum, , qfrc_bias,       nv,   1           ) \
   X( mjtNum, , qfrc_passive,    nv,   1           ) \
diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c
@@ -2067,6 +2067,12 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
 
   mj_markStack(d);
 
+  // inverse square root of D from inertia LDL decomposition
+  mjtNum* sqrtInvD = mjSTACKALLOC(d, nv, mjtNum);
+  for (int i=0; i < nv; i++) {
+    sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]);
+  }
+
   // space for backsubM2(J')' and its traspose
   mjtNum* JM2 = mjSTACKALLOC(d, nefc*nv, mjtNum);
   mjtNum* JM2T = mjSTACKALLOC(d, nv*nefc, mjtNum);
@@ -2140,7 +2146,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
         // process if not zero
         if (xi) {
           // x(i) /= sqrt(L(i,i))
-          JM2[adr+i] *= d->qLDiagSqrtInv[colind[adr+i]];
+          JM2[adr+i] *= sqrtInvD[colind[adr+i]];
 
           // x(j) -= L(i,j) * x(i)
           int Madr_ij = m->dof_Madr[colind[adr+i]]+1;
@@ -2191,7 +2197,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
   // dense
   else {
     // JM2 = backsubM2(J')'
-    mj_solveM2(m, d, JM2, d->efc_J, nefc);
+    mj_solveM2(m, d, JM2, d->efc_J, sqrtInvD, nefc);
 
     // construct JM2T
     mju_transpose(JM2T, JM2, nefc, nv);
diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c
@@ -774,7 +774,7 @@ void mj_tendon(const mjModel* m, mjData* d) {
         L[i] += (mju_dist3(wpnt, wpnt+3) + wlen + mju_dist3(wpnt+6, wpnt+9))/divisor;
       }
 
-      // accumulate moments if consequtive points are in different bodies
+      // accumulate moments if consecutive points are in different bodies
       for (int k=0; k < (wlen < 0 ? 1 : 3); k++) {
         if (wbody[k] != wbody[k+1]) {
           // get 3D position difference, normalize
@@ -1387,8 +1387,7 @@ void mj_crb(const mjModel* m, mjData* d) {
 
 
 // sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd
-void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv,
-                mjtNum* qLDiagSqrtInv) {
+void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv) {
   int cnt;
   int Madr_kk, Madr_ki;
   mjtNum tmp;
@@ -1445,9 +1444,6 @@ void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNu
   for (int i=0; i < nv; i++) {
     mjtNum qLDi = qLD[dof_Madr[i]];
     qLDiagInv[i] = 1.0/qLDi;
-    if (qLDiagSqrtInv) {
-      qLDiagSqrtInv[i] = 1.0/mju_sqrt(qLDi);
-    }
   }
 }
 
@@ -1456,7 +1452,7 @@ void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNu
 // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd
 void mj_factorM(const mjModel* m, mjData* d) {
   TM_START;
-  mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv, d->qLDiagSqrtInv);
+  mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv);
   TM_ADD(mjTIMER_POS_INERTIA);
 }
 
@@ -1685,10 +1681,10 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int
 
 
 // half of sparse backsubstitution:  x = sqrt(inv(D))*inv(L')*y
-void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {
+void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
+                const mjtNum* sqrtInvD, int n) {
   // local copies of key variables
   mjtNum* qLD = d->qLD;
-  mjtNum* qLDiagSqrtInv = d->qLDiagSqrtInv;
   int* dof_Madr = m->dof_Madr;
   int* dof_parentid = m->dof_parentid;
   int nv = m->nv;
@@ -1720,7 +1716,7 @@ void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n)
 
     // x <- sqrt(inv(D)) * x
     for (int i=0; i < nv; i++) {
-      x[i+offset] *= qLDiagSqrtInv[i];  // x(i) /= sqrt(L(i,i))
+      x[i+offset] *= sqrtInvD[i];  // x(i) /= sqrt(L(i,i))
     }
   }
 }
@@ -1781,7 +1777,7 @@ void mj_comVel(const mjModel* m, mjData* d) {
 
       default:
         // in principle we should use the new velocity to compute cdofdot,
-        // but it makes no difference becase crossMotion(cdof, cdof) = 0,
+        // but it makes no difference because crossMotion(cdof, cdof) = 0,
         // and using the old velocity may be more accurate numerically
         mju_crossMotion(cdofdot+6*j, cvel, d->cdof+6*(bda+j));
 
diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h
@@ -49,8 +49,7 @@ MJAPI void mj_transmission(const mjModel* m, mjData* d);
 MJAPI void mj_crb(const mjModel* m, mjData* d);
 
 // sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd
-MJAPI void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv,
-                      mjtNum* qLDiagSqrtInv);
+MJAPI void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv);
 
 // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd
 MJAPI void mj_factorM(const mjModel* m, mjData* d);
@@ -71,7 +70,8 @@ MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, in
 MJAPI void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* x, int island);
 
 // half of sparse backsubstitution:  x = sqrt(inv(D))*inv(L')*y
-MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
+MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
+                      const mjtNum* sqrtInvD, int n);
 
 
 //-------------------------- velocity --------------------------------------------------------------
diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c
@@ -803,7 +803,7 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
       }
 
       // factor
-      mj_factorI(m, d, MhB, d->qH, d->qHDiagInv, 0);
+      mj_factorI(m, d, MhB, d->qH, d->qHDiagInv);
     }
 
     // solve
@@ -986,7 +986,7 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
       mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM);
 
       // factorize
-      mj_factorI(m, d, MhB, d->qH, d->qHDiagInv, NULL);
+      mj_factorI(m, d, MhB, d->qH, d->qHDiagInv);
     }
 
     // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c
@@ -1107,7 +1107,6 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
   }
 
   printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format);
-  printArray("QLDIAGSQRTINV", m->nv, 1, d->qLDiagSqrtInv, fp, float_format);
 
   // B sparse structure
   printSparsity("B: body-dof matrix", m->nbody, m->nv, d->B_rowadr, NULL, d->B_rownnz, NULL,
diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c
@@ -1087,7 +1087,6 @@ void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum
 void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) {
   int  nv = m->nv;
   const mjtNum* qLD = d->qLD;
-  const mjtNum* qLDiagSqrtInv = d->qLDiagSqrtInv;
   const int* dofMadr = m->dof_Madr;
 
   mju_zero(res, nv);
@@ -1117,9 +1116,9 @@ void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec)
     }
   }
 
-  // res = sqrt(D) * res
+  // res *= sqrt(D)
   for (int i=0; i < nv; i++) {
-    res[i] /= qLDiagSqrtInv[i];
+    res[i] *= mju_sqrt(qLD[dofMadr[i]]);
   }
 }
 
diff --git a/src/engine/engine_vis_interact.c b/src/engine/engine_vis_interact.c
@@ -541,6 +541,7 @@ void mjv_initPerturb(const mjModel* m, mjData* d, const mjvScene* scn, mjvPertur
 
   mjtNum* jac = mjSTACKALLOC(d, 3*nv, mjtNum);
   mjtNum* jacM2 = mjSTACKALLOC(d, 3*nv, mjtNum);
+  mjtNum* sqrtInvD = mjSTACKALLOC(d, nv, mjtNum);
 
   // invalid selected body: return
   if (sel <= 0 || sel >= m->nbody) {
@@ -554,8 +555,11 @@ void mjv_initPerturb(const mjModel* m, mjData* d, const mjvScene* scn, mjvPertur
   mju_addTo3(selpos, d->xpos+3*sel);
 
   // compute average spatial inertia at selection point
+  for (int i=0; i < nv; i++) {
+    sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]);
+  }
   mj_jac(m, d, jac, NULL, selpos, sel);
-  mj_solveM2(m, d, jacM2, jac, 3);
+  mj_solveM2(m, d, jacM2, jac, sqrtInvD, 3);
   mjtNum invmass = mju_dot(jacM2+0*nv, jacM2+0*nv, nv) +
                    mju_dot(jacM2+1*nv, jacM2+1*nv, nv) +
                    mju_dot(jacM2+2*nv, jacM2+2*nv, nv);
diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs
@@ -4917,7 +4917,6 @@ public unsafe struct mjData_ {
   public double* qM;
   public double* qLD;
   public double* qLDiagInv;
-  public double* qLDiagSqrtInv;
   public double* bvh_aabb_dyn;
   public byte* bvh_active;
   public double* flexedge_velocity;
@@ -6679,7 +6678,7 @@ public unsafe struct mjvSceneState_ {
 public static unsafe extern void mj_solveM(mjModel_* m, mjData_* d, double* x, double* y, int n);
 
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
-public static unsafe extern void mj_solveM2(mjModel_* m, mjData_* d, double* x, double* y, int n);
+public static unsafe extern void mj_solveM2(mjModel_* m, mjData_* d, double* x, double* y, double* sqrtInvD, int n);
 
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mj_comVel(mjModel_* m, mjData_* d);
