From 7e1c65b10ebc68063792e88b23427b0aca09c0d5 Mon Sep 17 00:00:00 2001 From: Kyle Bayes Date: Mon, 14 Aug 2023 08:31:35 -0700 Subject: [PATCH] Improve CRB code: - remove mj_crbSkip since this was only utilized during compile time to compute m->dof_m0, replaced with mj_crbDiag to compute only the diagonals for dof_m0, - change row backward pass to forward pass to optimize processing diagonal blocks, - and various code refactoring. PiperOrigin-RevId: 556799924 Change-Id: I0a18d7a45862860fdddd7150540ddc8d64b01b29 --- src/engine/engine_core_smooth.c | 64 +++++++++++++++------------------ src/engine/engine_core_smooth.h | 3 -- src/engine/engine_setconst.c | 38 ++++++++++++++++---- src/user/user_model.cc | 17 ++++----- 4 files changed, 67 insertions(+), 55 deletions(-) diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index ea50fe0d..5c71f69e 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -957,16 +957,17 @@ void mj_transmission(const mjModel* m, mjData* d) { //-------------------------- inertia --------------------------------------------------------------- -// composite rigid body inertia algorithm, with skipsimple -void mj_crbSkip(const mjModel* m, mjData* d, int skipsimple) { - mjtNum tmp[6]; +// composite rigid body inertia algorithm +void mj_crb(const mjModel* m, mjData* d) { + 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=m->nbody-1; i > 0; i--) { + for (int i=last_body; i > 0; i--) { if (m->body_parentid[i] > 0) { mju_addTo(crb+10*m->body_parentid[i], crb+10*i, 10); } @@ -975,45 +976,38 @@ void mj_crbSkip(const mjModel* m, mjData* d, int skipsimple) { // clear qM mju_zero(d->qM, m->nM); - // dense backward pass over dofs - for (int i=m->nv-1; i >= 0; i--) { - // copy - if (skipsimple && m->dof_simplenum[i]) { - d->qM[m->dof_Madr[i]] = m->dof_M0[i]; - } + // dense forward pass over dofs + for (int i=0; i < nv; i++) { + // process block of diagonals (simple bodies) + 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]; + } - // compute - else { - // init M(i,i) with armature inertia - int Madr_ij = m->dof_Madr[i]; - d->qM[Madr_ij] = m->dof_armature[i]; - - // precompute tmp = crb * cdof - mju_mulInertVec(tmp, crb+10*m->dof_bodyid[i], d->cdof+6*i); - - // sparse backward pass over ancestors - int j = i; - while (j >= 0) { - // M(i,j) += cdof_j * crb_body(i) * cdof_i = cdof_j * tmp - d->qM[Madr_ij] += mju_dot(d->cdof+6*j, tmp, 6); - - // advance to parent - j = m->dof_parentid[j]; - Madr_ij++; + // finish or else fall through with next row + if (i == nv) { + break; } } + + // init M(i,i) with armature inertia + int Madr_ij = m->dof_Madr[i]; + d->qM[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); + + // 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); + } } } -// composite rigid body inertia algorithm -void mj_crb(const mjModel* m, mjData* d) { - mj_crbSkip(m, d, 1); -} - - - // 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) { diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h index b1fadc25..edaa4133 100644 --- a/src/engine/engine_core_smooth.h +++ b/src/engine/engine_core_smooth.h @@ -42,9 +42,6 @@ MJAPI void mj_transmission(const mjModel* m, mjData* d); //-------------------------- inertia --------------------------------------------------------------- -// composite rigid body inertia algorithm, with skip -void mj_crbSkip(const mjModel* m, mjData* d, int skipsimple); - // composite rigid body inertia algorithm MJAPI void mj_crb(const mjModel* m, mjData* d); diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c index 8470c23f..2b2b6e76 100644 --- a/src/engine/engine_setconst.c +++ b/src/engine/engine_setconst.c @@ -29,6 +29,34 @@ #include "engine/engine_util_misc.h" #include "engine/engine_util_spatial.h" + +// compute dof_M0 via composite rigid body algorithm +static void mj_setM0(mjModel* m, mjData* d) { + mjtNum buf[6]; + mjtNum* crb = d->crb; + int last_body = m->nbody - 1, nv = m->nv; + + // copy cinert into crb + mju_copy(crb, d->cinert, 10*m->nbody); + + // backward pass over bodies, accumulate composite inertias + for (int i=last_body; i > 0; i--) { + if (m->body_parentid[i] > 0) { + mju_addTo(crb+10*m->body_parentid[i], crb+10*i, 10); + } + } + + for (int i=0; i < nv; i++) { + // precomute buf = crb_body_i * cdof_i + mju_mulInertVec(buf, crb+10*m->dof_bodyid[i], d->cdof+6*i); + + // dof_M0(i) = armature inertia + cdof_i * (crb_body_i * cdof_i) + m->dof_M0[i] = m->dof_armature[i] + mju_dot(d->cdof+6*i, buf, 6); + } +} + + + // set quantities that depend on qpos0 static void set0(mjModel* m, mjData* d) { int id, id1, id2, dnum, nv = m->nv; @@ -60,14 +88,12 @@ static void set0(mjModel* m, mjData* d) { mj_kinematics(m, d); mj_comPos(m, d); mj_camlight(m, d); - mj_crbSkip(m, d, 0); - // save dof_M0 - for (int i=0; i < nv; i++) { - m->dof_M0[i] = d->qM[m->dof_Madr[i]]; - } + // compute dof_M0 for CRB algorithm + mj_setM0(m, d); - // run remaining computations (factorM needs dof_M0) + // run remaining computations + mj_crb(m, d); mj_factorM(m, d); mj_tendon(m, d); mj_transmission(m, d); diff --git a/src/user/user_model.cc b/src/user/user_model.cc index eb985cbe..18d0c33e 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -1667,19 +1667,14 @@ void mjCModel::CopyTree(mjModel* m) { m->nB = nB; // set dof_simplenum - int scnt = 0; - for (int i=nv-1; i>=0; i--) { - // dof in simple body + int count = 0; + for (int i=nv-1; i >= 0; i--) { if (m->body_simple[m->dof_bodyid[i]]) { - scnt++; - m->dof_simplenum[i] = scnt; - } - - // dof in regular body - else { - scnt = 0; - m->dof_simplenum[i] = 0; + count++; // increment counter + } else { + count = 0; // reset } + m->dof_simplenum[i] = count; } }