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

Convert qLD to CSR format.

Browse Source 
PiperOrigin-RevId: 723955038
Change-Id: I30c3dc7f59739e89ae5fff8841432bc74717ec1b
This commit is contained in:
Yuval Tassa
2025-02-06 08:58:45 -08:00
committed by Copybara-Service
parent cb1696eb34
commit c27d3758c2
22 changed files with 151 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+3 -9
View File
@@ -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);
}
src/engine/engine_core_smooth.c
+15 -26
View File
@@ -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,11 +1771,9 @@ 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);
}
@@ -1785,23 +1781,18 @@ void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int islan
// 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);
}
src/engine/engine_core_smooth.h
+1 -2
View File
@@ -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,
src/engine/engine_forward.c
+21 -13
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, 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");
}
src/engine/engine_print.c
+2 -1
View File
@@ -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
src/engine/engine_support.c
+8 -21
View File
@@ -1087,38 +1087,25 @@ 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 int* dofMadr = m->dof_Madr;
mju_zero(res, nv);
// res = L * vec
for (int i=0; i < nv; i++) {
// simple: diagonal
if (m->dof_simplenum[i]) {
res[i] = vec[i];
}
// diagonal
res[i] = vec[i];
// regular: full multiplication
else {
// diagonal
res[i] += vec[i];
// off-diagonal
int j = m->dof_parentid[i];
int adr = dofMadr[i] + 1;
while (j >= 0) {
res[i] += qLD[adr]*vec[j];
// advance to next element
j = m->dof_parentid[j];
adr++;
}
// 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);
}
}
// res *= sqrt(D)
for (int i=0; i < nv; i++) {
res[i] *= mju_sqrt(qLD[dofMadr[i]]);
int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
res[i] *= mju_sqrt(qLD[diag]);
}
}
src/engine/engine_vis_interact.c
+2 -1
View File
@@ -556,7 +556,8 @@ 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++) {
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]);
}
mj_jac(m, d, jac, NULL, selpos, sel);
mj_solveM2(m, d, jacM2, jac, sqrtInvD, 3);