From 2707eadc00fd12c351305d3e110fcc2277c83f6e Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Sat, 21 Jun 2025 04:48:07 -0700 Subject: [PATCH] Refactor `mj_crb` and `mj_tendonArmature` to use CSR `M`. PiperOrigin-RevId: 774077388 Change-Id: Id57118a5d66c074766f5c69fd3b889ebc625cbf7 --- src/engine/engine_core_smooth.c | 76 ++++++++++++++------------ src/engine/engine_io.c | 3 + test/engine/engine_core_smooth_test.cc | 1 + 3 files changed, 45 insertions(+), 35 deletions(-) diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index d4deb5b3..14b2dfe6 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -1469,9 +1469,12 @@ void mj_transmission(const mjModel* m, mjData* d) { //-------------------------- inertia --------------------------------------------------------------- -// add tendon armature to qM +// add tendon armature to M void mj_tendonArmature(const mjModel* m, mjData* d) { int nv = m->nv, ntendon = m->ntendon, issparse = mj_isSparse(m); + const int* M_rownnz = d->M_rownnz; + const int* M_rowadr = d->M_rowadr; + const int* M_colind = d->M_colind; for (int k=0; k < ntendon; k++) { mjtNum armature = m->tendon_armature[k]; @@ -1482,11 +1485,19 @@ void mj_tendonArmature(const mjModel* m, mjData* d) { // dense if (!issparse) { + // M += armature * ten_J' * ten_J mjtNum* ten_J = d->ten_J + nv*k; - for (int i=0; i < m->nv; i++) { - int Madr = m->dof_Madr[i]; - for (int j = i; j >= 0; j = m->dof_parentid[j]) { - d->qM[Madr++] += armature * ten_J[j] * ten_J[i]; + for (int i=0; i < nv; i++) { + mjtNum ten_J_i = ten_J[i]; + if (!ten_J_i) { + continue; + } + + // M[i,:] += armature * ten_J[i] * ten_J + int start = M_rowadr[i]; + int end = start + M_rownnz[i]; + for (int adr = start; adr < end; adr++) { + d->M[adr] += armature * ten_J_i * ten_J[M_colind[adr]]; } } } @@ -1494,29 +1505,24 @@ void mj_tendonArmature(const mjModel* m, mjData* d) { // sparse else { // get sparse info for tendon k - int rowadr = d->ten_J_rowadr[k]; - int rownnz = d->ten_J_rownnz[k]; - const int* colind = d->ten_J_colind + rowadr; - mjtNum* ten_J = d->ten_J + rowadr; + int J_rowadr = d->ten_J_rowadr[k]; + int J_rownnz = d->ten_J_rownnz[k]; + const int* J_colind = d->ten_J_colind + J_rowadr; + mjtNum* ten_J = d->ten_J + J_rowadr; - // iterate forward on nonzero rows i - for (int adr_i=0; adr_i < rownnz; adr_i++) { - int i = colind[adr_i]; - int Madr = m->dof_Madr[i]; - int adr_j = rownnz - 1; - - // iterate backward on ancestors of i, find matching column j - for (int j = i; j >= 0; j = m->dof_parentid[j]) { - // reduce adr_j until column index is no bigger than j - while (colind[adr_j] > j && adr_j >= 0) { - adr_j--; - } - - // found match, update qM - if (colind[adr_j] == j) { - d->qM[Madr++] += armature * ten_J[adr_j] * ten_J[adr_i]; - } + // M += armature * ten_J' * ten_J + for (int j=0; j < J_rownnz; j++) { + mjtNum ten_J_i = ten_J[j]; + if (!ten_J_i) { + continue; } + + // M[i,:] += armature * ten_J[i] * ten_J + int i = J_colind[j]; + int M_adr = M_rowadr[i]; + mju_addToSclSparseInc(d->M + M_adr, ten_J, + M_rownnz[i], M_colind + M_adr, + J_rownnz, J_colind, armature * ten_J_i); } } } @@ -1526,22 +1532,22 @@ void mj_tendonArmature(const mjModel* m, mjData* d) { // composite rigid body inertia algorithm void mj_crb(const mjModel* m, mjData* d) { + int nv = m->nv; mjtNum buf[6]; mjtNum* crb = d->crb; - int last_body = m->nbody - 1, nv = m->nv; // crb = cinert mju_copy(crb, d->cinert, 10*m->nbody); // backward pass over bodies, accumulate composite inertias - for (int i=last_body; i > 0; i--) { + for (int i=m->nbody - 1; i > 0; i--) { if (m->body_parentid[i] > 0) { mju_addTo(crb+10*m->body_parentid[i], crb+10*i, 10); } } - // clear qM - mju_zero(d->qM, m->nM); + // clear M + mju_zero(d->M, m->nC); // dense forward pass over dofs for (int i=0; i < nv; i++) { @@ -1549,7 +1555,7 @@ void mj_crb(const mjModel* m, mjData* d) { if (m->dof_simplenum[i]) { int n = i + m->dof_simplenum[i]; for (; i < n; i++) { - d->qM[m->dof_Madr[i]] = m->dof_M0[i]; + d->M[d->M_rowadr[i]] = m->dof_M0[i]; } // finish or else fall through with next row @@ -1559,8 +1565,8 @@ void mj_crb(const mjModel* m, mjData* d) { } // init M(i,i) with armature inertia - int Madr_ij = m->dof_Madr[i]; - d->qM[Madr_ij] = m->dof_armature[i]; + int Madr_ij = d->M_rowadr[i] + d->M_rownnz[i] - 1; + d->M[Madr_ij] = m->dof_armature[i]; // precompute buf = crb_body_i * cdof_i mju_mulInertVec(buf, crb+10*m->dof_bodyid[i], d->cdof+6*i); @@ -1568,7 +1574,7 @@ void mj_crb(const mjModel* m, mjData* d) { // sparse backward pass over ancestors for (int j=i; j >= 0; j = m->dof_parentid[j]) { // M(i,j) += cdof_j * (crb_body_i * cdof_i) - d->qM[Madr_ij++] += mju_dot(d->cdof+6*j, buf, 6); + d->M[Madr_ij--] += mju_dot(d->cdof+6*j, buf, 6); } } } @@ -1579,7 +1585,7 @@ void mj_makeM(const mjModel* m, mjData* d) { TM_START; mj_crb(m, d); mj_tendonArmature(m, d); - mju_gather(d->M, d->qM, d->mapM2M, m->nC); + mju_scatter(d->qM, d->M, d->mapM2M, m->nC); TM_END(mjTIMER_POS_INERTIA); } diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 54262a18..fdad5ec5 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -1990,6 +1990,9 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { mju_zero(d->mocap_pos, 3*m->nmocap); mju_zero(d->mocap_quat, 4*m->nmocap); + // zero out qM, special case because scattring from M skips simple body off-diagonals + mju_zero(d->qM, m->nM); + // copy qpos0 from model if (m->qpos0) { memcpy(d->qpos, m->qpos0, m->nq*sizeof(mjtNum)); diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index 880f3e7e..59e4ca3e 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/test/engine/engine_core_smooth_test.cc @@ -237,6 +237,7 @@ TEST_F(CoreSmoothTest, TendonArmature) { // put only CRB inertia in M2 mj_crb(m, d); + mju_scatter(d->qM, d->M, d->mapM2M, m->nC); vector M2(nv*nv); mj_fullM(m, M2.data(), d->qM);