chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Switch mjData.{qH,qLD} from reduced ("C") to full ("M") inertia matrix structure. No performance impact of extra zeros because of existing "simple dof" skipping mechanism.

Browse Source 
PiperOrigin-RevId: 733523931
Change-Id: Ic8d8a152dda5532331c239cb6b4ce7d8d09b7fff
This commit is contained in:
Yuval Tassa
2025-03-04 17:33:48 -08:00
committed by Copybara-Service
parent 51f6aa8b43
commit 0f563ecf31
20 changed files with 95 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+9 -9
View File
@@ -2131,7 +2131,7 @@ 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++) {
int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
int diag = d->M_rowadr[i] + d->M_rownnz[i] - 1;
sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]);
}
@@ -2167,11 +2167,11 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
continue;
}
// traverse row j of C, marking new unique nonzeros
int nnzC = d->C_rownnz[j];
int adrC = d->C_rowadr[j];
for (int k=0; k < nnzC; k++) {
int c = d->C_colind[adrC + k];
// traverse row j of M, marking new unique nonzeros
int nnzM = d->M_rownnz[j];
int adrM = d->M_rowadr[j];
for (int k=0; k < nnzM; k++) {
int c = d->M_colind[adrM + k];
if (marker[c] != r) {
marker[c] = r;
nnz++;
@@ -2251,10 +2251,10 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
continue;
}
int j = B_colind[i];
int adrC = d->C_rowadr[j];
mju_addToSclSparseInc(B + adrB, d->qLD + adrC,
int adrM = d->M_rowadr[j];
mju_addToSclSparseInc(B + adrB, d->qLD + adrM,
nnzB, B_colind + adrB,
d->C_rownnz[j]-1, d->C_colind + adrC, -b);
d->M_rownnz[j]-1, d->M_colind + adrM, -b);
}
// B(r,:) <- sqrt(inv(D)) * B(r,:)
src/engine/engine_core_smooth.c
+12 -12
View File
@@ -1469,11 +1469,11 @@ void mj_factorI_legacy(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD
// sparse L'*D*L factorizaton of the inertia matrix M, assumed spd
void mj_factorM(const mjModel* m, mjData* d) {
TM_START;
int nC = m->nC;
for (int i=0; i < nC; i++) {
d->qLD[i] = d->qM[d->mapM2C[i]];
int nM = m->nM;
for (int i=0; i < nM; i++) {
d->qLD[i] = d->qM[d->mapM2M[i]];
}
mj_factorI(d->qLD, d->qLDiagInv, m->nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
mj_factorI(d->qLD, d->qLDiagInv, m->nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
TM_ADD(mjTIMER_POS_INERTIA);
}
@@ -1715,7 +1715,7 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {
mju_copy(x, y, n*m->nv);
}
mj_solveLD(x, d->qLD, d->qLDiagInv, m->nv, n,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
}
@@ -1727,14 +1727,14 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int
const mjtNum* qLDiagInv = d->qLDiagInv;
if (island < 0) {
mj_solveLD(x, qLD, qLDiagInv, m->nv, 1,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
return;
}
// local copies of key variables
const int* rownnz = d->C_rownnz;
const int* rowadr = d->C_rowadr;
const int* colind = d->C_colind;
const int* rownnz = d->M_rownnz;
const int* rowadr = d->M_rowadr;
const int* colind = d->M_colind;
const int* diagnum = m->dof_simplenum;
// local constants: island specific
@@ -1785,9 +1785,9 @@ void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
int nv = m->nv;
// local copies of key variables
const int* rownnz = d->C_rownnz;
const int* rowadr = d->C_rowadr;
const int* colind = d->C_colind;
const int* rownnz = d->M_rownnz;
const int* rowadr = d->M_rowadr;
const int* colind = d->M_colind;
const int* diagnum = m->dof_simplenum;
const mjtNum* qLD = d->qLD;
src/engine/engine_forward.c
+11 -11
View File
@@ -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, nC = m->nC;
int nv = m->nv, nM = m->nM;
mj_markStack(d);
mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum);
mjtNum* qacc = mjSTACKALLOC(d, nv, mjtNum);
@@ -795,22 +795,22 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
else {
if (!skipfactor) {
// 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 < nM; i++) {
d->qH[i] = d->qM[d->mapM2M[i]];
}
for (int i=0; i < nv; i++) {
d->qH[d->C_rowadr[i] + d->C_rownnz[i] - 1] += m->opt.timestep * m->dof_damping[i];
d->qH[d->M_rowadr[i] + d->M_rownnz[i] - 1] += m->opt.timestep * m->dof_damping[i];
}
// factorize in-place
mj_factorI(d->qH, d->qHDiagInv, nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
}
// solve
mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv);
mju_copy(qacc, qfrc, m->nv);
mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
}
// advance state and time
@@ -939,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, nC = m->nC;
int nv = m->nv, nM = m->nM, nD = m->nD;
mj_markStack(d);
mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum);
@@ -987,18 +987,18 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM);
// copy into qH
for (int i=0; i < nC; i++) {
d->qH[i] = MhB[d->mapM2C[i]];
for (int i=0; i < nM; i++) {
d->qH[i] = MhB[d->mapM2M[i]];
}
// factorize in-place
mj_factorI(d->qH, d->qHDiagInv, nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
}
// solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
mju_copy(qacc, qfrc, nv);
mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
} else {
mjERROR("integrator must be implicit or implicitfast");
src/engine/engine_print.c
+3 -3
View File
@@ -1127,12 +1127,12 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
printInertia("QM", d->qM, m, fp, float_format);
printSparse("QLD", d->qLD, m->nv, d->C_rownnz,
d->C_rowadr, d->C_colind, fp, float_format);
printSparse("QLD", d->qLD, m->nv, d->M_rownnz,
d->M_rowadr, d->M_colind, fp, float_format);
printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format);
if (!mju_isZero(d->qHDiagInv, m->nv)) {
printSparse("QH", d->qH, m->nv, d->C_rownnz, d->C_rowadr, d->C_colind, fp, float_format);
printSparse("QH", d->qH, m->nv, d->M_rownnz, d->M_rowadr, d->M_colind, fp, float_format);
printArray("QHDIAGINV", m->nv, 1, d->qHDiagInv, fp, float_format);
}
src/engine/engine_support.c
+3 -3
View File
@@ -1097,14 +1097,14 @@ void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec)
// non-simple: add off-diagonals
if (!m->dof_simplenum[i]) {
int adr = d->C_rowadr[i];
res[i] += mju_dotSparse(qLD+adr, vec, d->C_rownnz[i] - 1, d->C_colind+adr, /*flg_unc1=*/0);
int adr = d->M_rowadr[i];
res[i] += mju_dotSparse(qLD+adr, vec, d->M_rownnz[i] - 1, d->M_colind+adr, /*flg_unc1=*/0);
}
}
// res *= sqrt(D)
for (int i=0; i < nv; i++) {
int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
int diag = d->M_rowadr[i] + d->M_rownnz[i] - 1;
res[i] *= mju_sqrt(qLD[diag]);
}
}
src/engine/engine_vis_interact.c
+1 -1
View File
@@ -556,7 +556,7 @@ void mjv_initPerturb(const mjModel* m, mjData* d, const mjvScene* scn, mjvPertur
// compute average spatial inertia at selection point
for (int i=0; i < nv; i++) {
int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
int diag = d->M_rowadr[i] + d->M_rownnz[i] - 1;
sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]);
}
mj_jac(m, d, jac, NULL, selpos, sel);