Convert qLD to CSR format.

PiperOrigin-RevId: 723955038 Change-Id: I30c3dc7f59739e89ae5fff8841432bc74717ec1b
google-deepmind · Feb 6, 2025 · c27d375 · c27d375
1 parent cb1696e
commit c27d375
Show file tree

Hide file tree

Showing 22 changed files with 151 additions and 153 deletions.
diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst
@@ -423,7 +423,8 @@ Get name of object with the specified mjtObj type and id, returns NULL if name n
 
 .. mujoco-include:: mj_fullM
 
-Convert sparse inertia matrix M into full (i.e. dense) matrix.
+Convert sparse inertia matrix ``M`` into full (i.e. dense) matrix.
+|br| ``dst`` must be of size ``nv x nv``, ``M`` must be of the same size as ``mjData.qM``.
 
 .. _mj_mulM:
 

diff --git a/doc/APIreference/functions_override.rst b/doc/APIreference/functions_override.rst
@@ -233,6 +233,11 @@ found, the function will return ``distmax`` and ``fromto``, if given, will be se
    In order to determine whether a geom pair uses ``mjc_Convex``, inspect the table at the top of
    `engine_collision_driver.c <https://github.com/google-deepmind/mujoco/blob/main/src/engine/engine_collision_driver.c>`__.
 
+.. _mj_fullM:
+
+Convert sparse inertia matrix ``M`` into full (i.e. dense) matrix.
+|br| ``dst`` must be of size ``nv x nv``, ``M`` must be of the same size as ``mjData.qM``.
+
 .. _mj_mulM:
 
 This function multiplies the joint-space inertia matrix stored in mjData.qM by a vector. qM has a custom sparse format

diff --git a/doc/includes/references.h b/doc/includes/references.h
@@ -272,7 +272,7 @@ struct mjData_ {
   mjtNum* qM;                // total inertia (sparse)                           (nM x 1)
 
   // computed by mj_fwdPosition/mj_factorM
-  mjtNum* qLD;               // L'*D*L factorization of M (sparse)               (nM x 1)
+  mjtNum* qLD;               // L'*D*L factorization of M (sparse)               (nC x 1)
   mjtNum* qLDiagInv;         // 1/diag(D)                                        (nv x 1)
 
   // computed by mj_collisionTree
@@ -305,7 +305,7 @@ struct mjData_ {
   mjtNum* subtree_angmom;    // angular momentum about subtree com               (nbody x 3)
 
   // computed by mj_Euler or mj_implicit
-  mjtNum* qH;                // L'*D*L factorization of modified M               (nM x 1)
+  mjtNum* qH;                // L'*D*L factorization of modified M               (nC x 1)
   mjtNum* qHDiagInv;         // 1/diag(D) of modified M                          (nv x 1)
 
   // computed by mj_resetData

diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h
@@ -300,7 +300,7 @@ struct mjData_ {
   mjtNum* qM;                // total inertia (sparse)                           (nM x 1)
 
   // computed by mj_fwdPosition/mj_factorM
-  mjtNum* qLD;               // L'*D*L factorization of M (sparse)               (nM x 1)
+  mjtNum* qLD;               // L'*D*L factorization of M (sparse)               (nC x 1)
   mjtNum* qLDiagInv;         // 1/diag(D)                                        (nv x 1)
 
   // computed by mj_collisionTree
@@ -333,7 +333,7 @@ struct mjData_ {
   mjtNum* subtree_angmom;    // angular momentum about subtree com               (nbody x 3)
 
   // computed by mj_Euler or mj_implicit
-  mjtNum* qH;                // L'*D*L factorization of modified M               (nM x 1)
+  mjtNum* qH;                // L'*D*L factorization of modified M               (nC x 1)
   mjtNum* qHDiagInv;         // 1/diag(D) of modified M                          (nv x 1)
 
   // computed by mj_resetData

diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h
@@ -637,7 +637,7 @@
     X   ( mjtNum,    actuator_moment,   nJmom,       1           ) \
     X   ( mjtNum,    crb,               nbody,       10          ) \
     X   ( mjtNum,    qM,                nM,          1           ) \
-    X   ( mjtNum,    qLD,               nM,          1           ) \
+    X   ( mjtNum,    qLD,               nC,          1           ) \
     X   ( mjtNum,    qLDiagInv,         nv,          1           ) \
     XMJV( mjtNum,    bvh_aabb_dyn,      nbvhdynamic, 6           ) \
     XMJV( mjtByte,   bvh_active,        nbvh,        1           ) \
@@ -654,7 +654,7 @@
     X   ( mjtNum,    qfrc_passive,      nv,          1           ) \
     X   ( mjtNum,    subtree_linvel,    nbody,       3           ) \
     X   ( mjtNum,    subtree_angmom,    nbody,       3           ) \
-    X   ( mjtNum,    qH,                nM,          1           ) \
+    X   ( mjtNum,    qH,                nC,          1           ) \
     X   ( mjtNum,    qHDiagInv,         nv,          1           ) \
     X   ( int,       B_rownnz,          nbody,       1           ) \
     X   ( int,       B_rowadr,          nbody,       1           ) \

diff --git a/introspect/structs.py b/introspect/structs.py
@@ -5261,7 +5261,7 @@
                      inner_type=ValueType(name='mjtNum'),
                  ),
                  doc="L'*D*L factorization of M (sparse)",
-                 array_extent=('nM',),
+                 array_extent=('nC',),
              ),
              StructFieldDecl(
                  name='qLDiagInv',
@@ -5397,7 +5397,7 @@
                      inner_type=ValueType(name='mjtNum'),
                  ),
                  doc="L'*D*L factorization of modified M",
-                 array_extent=('nM',),
+                 array_extent=('nC',),
              ),
              StructFieldDecl(
                  name='qHDiagInv',

diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
@@ -368,7 +368,7 @@ def make_data(
         'efc_aref': (nefc, float),
         'efc_force': (nefc, float),
         '_qM_sparse': (m.nM, float),
-        '_qLD_sparse': (m.nM, float),
+        '_qLD_sparse': (m.nC, float),
         '_qLDiagInv_sparse': (m.nv, float),
     }
 

diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py
@@ -91,7 +91,10 @@ def test_smooth(self):
     _assert_eq(dx._qM_sparse, np.zeros(0), '_qM_sparse')
     # factor_m
     dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
-    _assert_attr_eq(d, dx, 'qLD')
+    qLDLegacy = np.zeros(mx.nM)  # pylint:disable=invalid-name
+    for i in range(m.nC):
+      qLDLegacy[d.mapM2C[i]] = d.qLD[i]
+    _assert_eq(qLDLegacy, dx.qLD, 'qLD')
     _assert_attr_eq(d, dx, 'qLDiagInv')
     _assert_eq(dx._qLD_sparse, np.zeros(0), '_qLD_sparse')
     _assert_eq(dx._qLDiagInv_sparse, np.zeros(0), '_qLDiagInv_sparse')

diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
@@ -1348,7 +1348,7 @@ class Data(PyTreeNode):
     efc_aref: reference pseudo-acceleration                     (nefc,)
     efc_force: constraint force in constraint space             (nefc,)
     _qM_sparse: qM in sparse representation                     (nM,)
-    _qLD_sparse: qLD in sparse representation                   (nM,)
+    _qLD_sparse: qLD in sparse representation                   (nC,)
     _qLDiagInv_sparse: qLDiagInv in sparse representation       (nv,)
   """  # fmt: skip
   # constant sizes:

diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c
@@ -2131,7 +2131,8 @@ void mj_projectConstraint(const mjModel* m, mjData* 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]]);
+    int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
+    sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]);
   }
 
   // sparse
@@ -2238,13 +2239,6 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
 
     // === in-place sparse back-substitution:  B <- B * M^-1/2
 
-    // make qLD
-    int nC = m->nC;
-    mjtNum* qLD = mjSTACKALLOC(d, nC, mjtNum);
-    for (int i=0; i < nC; i++) {
-      qLD[i] = d->qLD[d->mapM2C[i]];
-    }
-
     // sparse backsubM2 (half of LD back-substitution)
     for (int r=0; r < nefc; r++) {
       int nnzB = B_rownnz[r];
@@ -2258,7 +2252,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
         }
         int j = B_colind[i];
         int adrC = d->C_rowadr[j];
-        mju_addToSclSparseInc(B + adrB, qLD + adrC,
+        mju_addToSclSparseInc(B + adrB, d->qLD + adrC,
                               nnzB, B_colind + adrB,
                               d->C_rownnz[j]-1, d->C_colind + adrC, -b);
       }

diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c
@@ -1465,7 +1465,11 @@ 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);
+  int nC = m->nC;
+  for (int i=0; i < nC; i++) {
+    d->qLD[i] = d->qM[d->mapM2C[i]];
+  }
+  mj_factorIs(d->qLD, d->qLDiagInv, m->nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
   TM_ADD(mjTIMER_POS_INERTIA);
 }
 
@@ -1709,18 +1713,20 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {
   if (x != y) {
     mju_copy(x, y, n*m->nv);
   }
-  mj_solveLD(m, x, n, d->qLD, d->qLDiagInv);
+  mj_solveLDs(x, d->qLD, d->qLDiagInv, m->nv, n,
+              d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
 }
 
 
 // in-place sparse backsubstitution for one island:  x = inv(L'*D*L)*x
 //  L is in lower triangle of qLD; D is on diagonal of qLD
-void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int island) {
+void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int island) {
   // if no islands, call mj_solveLD
   const mjtNum* qLD = d->qLD;
   const mjtNum* qLDiagInv = d->qLDiagInv;
   if (island < 0) {
-    mj_solveLD(m, x, 1, qLD, qLDiagInv);
+    mj_solveLDs(x, qLD, qLDiagInv, m->nv, 1,
+                d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
     return;
   }
 
@@ -1730,14 +1736,6 @@ void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int islan
   const int* colind = d->C_colind;
   const int* diagnum = m->dof_simplenum;
 
-  // temporary: make local CSR version of qLD
-  int nC = m->nC;
-  mj_markStack(d);
-  mjtNum* qLDs = mjSTACKALLOC(d, nC, mjtNum);
-  for (int i=0; i < nC; i++) {
-    qLDs[i] = d->qLD[d->mapM2C[i]];
-  }
-
   // local constants: island specific
   int ndof = d->island_dofnum[island];
   const int* dofind = d->island_dofind + d->island_dofadr[island];
@@ -1751,7 +1749,7 @@ void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int islan
       int start = rowadr[i];
       int end = start + rownnz[i] - 1;
       for (int adr=end-1; adr >= start; adr--) {
-        x[islandind[colind[adr]]] -= qLDs[adr] * x_k;
+        x[islandind[colind[adr]]] -= qLD[adr] * x_k;
       }
     }
   }
@@ -1773,35 +1771,28 @@ void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int islan
     int start = rowadr[i];
     int end = start + rownnz[i] - 1;
     for (int adr=end-1; adr >= start; adr--) {
-      x[k] -= x[islandind[colind[adr]]] * qLDs[adr];
+      x[k] -= x[islandind[colind[adr]]] * qLD[adr];
     }
   }
-
-  mj_freeStack(d);
 }
 
 
 
 // half of sparse backsubstitution:  x = sqrt(inv(D))*inv(L')*y
 void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
                 const mjtNum* sqrtInvD, int n) {
+  int nv = m->nv;
+
   // local copies of key variables
-  int nv = m->nv, nC = m->nC;
   const int* rownnz = d->C_rownnz;
   const int* rowadr = d->C_rowadr;
   const int* colind = d->C_colind;
   const int* diagnum = m->dof_simplenum;
+  const mjtNum* qLD = d->qLD;
 
   // x = y
   mju_copy(x, y, n * nv);
 
-  // temporary: make local CSR version of qLD
-  mj_markStack(d);
-  mjtNum* qLD = mjSTACKALLOC(d, nC, mjtNum);
-  for (int i=0; i < nC; i++) {
-    qLD[i] = d->qLD[d->mapM2C[i]];
-  }
-
   // x <- L^-T x
   for (int i=nv-1; i > 0; i--) {
     // skip diagonal rows
@@ -1831,8 +1822,6 @@ void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
       x[i+offset] *= invD_i;
     }
   }
-
-  mj_freeStack(d);
 }
 
 

diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h
@@ -71,9 +71,8 @@ MJAPI void mj_solveLDs(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, i
 // sparse backsubstitution:  x = inv(L'*D*L)*y, use factorization in d
 MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
 
-// TODO(tassa): Restore mjData const-ness.
 // sparse backsubstitution for one island:  x = inv(L'*D*L)*x, use factorization in d
-MJAPI void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* x, int island);
+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,

diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c
@@ -770,7 +770,7 @@ static void mj_advance(const mjModel* m, mjData* d,
 // Euler integrator, semi-implicit in velocity, possibly skipping factorisation
 void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
   TM_START;
-  int nv = m->nv, nM = m->nM;
+  int nv = m->nv, nC = m->nC;
   mj_markStack(d);
   mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum);
   mjtNum* qacc = mjSTACKALLOC(d, nv, mjtNum);
@@ -794,22 +794,23 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
   // damping: integrate implicitly
   else {
     if (!skipfactor) {
-      mjtNum* MhB = mjSTACKALLOC(d, nM, mjtNum);
-
-      // MhB = M + h*diag(B)
-      mju_copy(MhB, d->qM, nM);
+      // qH = M + h*diag(B)
+      for (int i=0; i < nC; i++) {
+        d->qH[i] = d->qM[d->mapM2C[i]];
+      }
       for (int i=0; i < nv; i++) {
-        MhB[m->dof_Madr[i]] += m->opt.timestep * m->dof_damping[i];
+        d->qH[d->C_rowadr[i] + d->C_rownnz[i] - 1] += m->opt.timestep * m->dof_damping[i];
       }
 
-      // factor
-      mj_factorI(m, d, MhB, d->qH, d->qHDiagInv);
+      // factorize in-place
+      mj_factorIs(d->qH, d->qHDiagInv, nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
     }
 
     // solve
     mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv);
     mju_copy(qacc, qfrc, m->nv);
-    mj_solveLD(m, qacc, 1, d->qH, d->qHDiagInv);
+    mj_solveLDs(qacc, d->qH, d->qHDiagInv, nv, 1,
+                d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
   }
 
   // advance state and time
@@ -938,7 +939,7 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) {
 // fully implicit in velocity, possibly skipping factorization
 void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
   TM_START;
-  int nv = m->nv, nM = m->nM, nD = m->nD;
+  int nv = m->nv, nM = m->nM, nD = m->nD, nC = m->nC;
 
   mj_markStack(d);
   mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum);
@@ -985,13 +986,20 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
       // set MhB = M - dt*qDeriv
       mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM);
 
-      // factorize
-      mj_factorI(m, d, MhB, d->qH, d->qHDiagInv);
+      // copy into qH
+      for (int i=0; i < nC; i++) {
+        d->qH[i] = MhB[d->mapM2C[i]];
+      }
+
+      // factorize in-place
+      mj_factorIs(d->qH, d->qHDiagInv, nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
     }
 
     // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
     mju_copy(qacc, qfrc, nv);
-    mj_solveLD(m, qacc, 1, d->qH, d->qHDiagInv);
+    mj_solveLDs(qacc, d->qH, d->qHDiagInv, nv, 1,
+                d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
+
   } else {
     mjERROR("integrator must be implicit or implicitfast");
   }

diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c
@@ -1127,7 +1127,8 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
 
   printInertia("QM", d->qM, m, fp, float_format);
 
-  printInertia("QLD", d->qLD, m, fp, float_format);
+  printSparse("QLD", d->qLD, m->nv, d->C_rownnz,
+              d->C_rowadr, d->C_colind, fp, float_format);
   printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format);
 
   // B sparse structure