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

Remove mjData.qLDiagSqrtInv, add corresponding argument to mj_solveM2.

Browse Source 
- `qLDiagSqrtInv` is only required for the dual solvers. It is now computed as-needed rather than unconditionally.
- `mj_solveM2` now requires a new input array `sqrtInvD` which contains the square root of the inverse diagonal D (formerly saved in `qLDiagSqrtInv`).

PiperOrigin-RevId: 710805133
Change-Id: I0622d6a8da3882916824e9c10bad9223c122c321
This commit is contained in:
Yuval Tassa
2024-12-30 15:23:41 -08:00
committed by Copybara-Service
parent ec322641b7
commit 69c9ac074a
19 changed files with 53 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+8 -2
View File
@@ -2067,6 +2067,12 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
mj_markStack(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]]);
}
// space for backsubM2(J')' and its traspose
mjtNum* JM2 = mjSTACKALLOC(d, nefc*nv, mjtNum);
mjtNum* JM2T = mjSTACKALLOC(d, nv*nefc, mjtNum);
@@ -2140,7 +2146,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
// process if not zero
if (xi) {
// x(i) /= sqrt(L(i,i))
JM2[adr+i] *= d->qLDiagSqrtInv[colind[adr+i]];
JM2[adr+i] *= sqrtInvD[colind[adr+i]];
// x(j) -= L(i,j) * x(i)
int Madr_ij = m->dof_Madr[colind[adr+i]]+1;
@@ -2191,7 +2197,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
// dense
else {
// JM2 = backsubM2(J')'
mj_solveM2(m, d, JM2, d->efc_J, nefc);
mj_solveM2(m, d, JM2, d->efc_J, sqrtInvD, nefc);
// construct JM2T
mju_transpose(JM2T, JM2, nefc, nv);
src/engine/engine_core_smooth.c
+7 -11
View File
@@ -774,7 +774,7 @@ void mj_tendon(const mjModel* m, mjData* d) {
L[i] += (mju_dist3(wpnt, wpnt+3) + wlen + mju_dist3(wpnt+6, wpnt+9))/divisor;
}
// accumulate moments if consequtive points are in different bodies
// accumulate moments if consecutive points are in different bodies
for (int k=0; k < (wlen < 0 ? 1 : 3); k++) {
if (wbody[k] != wbody[k+1]) {
// get 3D position difference, normalize
@@ -1387,8 +1387,7 @@ void mj_crb(const mjModel* m, mjData* d) {
// 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) {
void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv) {
int cnt;
int Madr_kk, Madr_ki;
mjtNum tmp;
@@ -1445,9 +1444,6 @@ void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNu
for (int i=0; i < nv; i++) {
mjtNum qLDi = qLD[dof_Madr[i]];
qLDiagInv[i] = 1.0/qLDi;
if (qLDiagSqrtInv) {
qLDiagSqrtInv[i] = 1.0/mju_sqrt(qLDi);
}
}
}
@@ -1456,7 +1452,7 @@ 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, d->qLDiagSqrtInv);
mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv);
TM_ADD(mjTIMER_POS_INERTIA);
}
@@ -1685,10 +1681,10 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int
// half of sparse backsubstitution: x = sqrt(inv(D))*inv(L')*y
void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {
void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
const mjtNum* sqrtInvD, int n) {
// local copies of key variables
mjtNum* qLD = d->qLD;
mjtNum* qLDiagSqrtInv = d->qLDiagSqrtInv;
int* dof_Madr = m->dof_Madr;
int* dof_parentid = m->dof_parentid;
int nv = m->nv;
@@ -1720,7 +1716,7 @@ void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n)
// x <- sqrt(inv(D)) * x
for (int i=0; i < nv; i++) {
x[i+offset] *= qLDiagSqrtInv[i]; // x(i) /= sqrt(L(i,i))
x[i+offset] *= sqrtInvD[i]; // x(i) /= sqrt(L(i,i))
}
}
}
@@ -1781,7 +1777,7 @@ void mj_comVel(const mjModel* m, mjData* d) {
default:
// in principle we should use the new velocity to compute cdofdot,
// but it makes no difference becase crossMotion(cdof, cdof) = 0,
// but it makes no difference because crossMotion(cdof, cdof) = 0,
// and using the old velocity may be more accurate numerically
mju_crossMotion(cdofdot+6*j, cvel, d->cdof+6*(bda+j));
src/engine/engine_core_smooth.h
+3 -3
View File
@@ -49,8 +49,7 @@ MJAPI void mj_transmission(const mjModel* m, mjData* d);
MJAPI void mj_crb(const mjModel* m, mjData* d);
// sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd
MJAPI void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv,
mjtNum* qLDiagSqrtInv);
MJAPI void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv);
// sparse L'*D*L factorizaton of the inertia matrix M, assumed spd
MJAPI void mj_factorM(const mjModel* m, mjData* d);
@@ -71,7 +70,8 @@ MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, in
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, int n);
MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
const mjtNum* sqrtInvD, int n);
//-------------------------- velocity --------------------------------------------------------------
src/engine/engine_forward.c
+2 -2
View File
@@ -803,7 +803,7 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
}
// factor
mj_factorI(m, d, MhB, d->qH, d->qHDiagInv, 0);
mj_factorI(m, d, MhB, d->qH, d->qHDiagInv);
}
// solve
@@ -986,7 +986,7 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM);
// factorize
mj_factorI(m, d, MhB, d->qH, d->qHDiagInv, NULL);
mj_factorI(m, d, MhB, d->qH, d->qHDiagInv);
}
// solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
src/engine/engine_print.c
-1
View File
@@ -1107,7 +1107,6 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
}
printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format);
printArray("QLDIAGSQRTINV", m->nv, 1, d->qLDiagSqrtInv, fp, float_format);
// B sparse structure
printSparsity("B: body-dof matrix", m->nbody, m->nv, d->B_rowadr, NULL, d->B_rownnz, NULL,
src/engine/engine_support.c
+2 -3
View File
@@ -1087,7 +1087,6 @@ 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 mjtNum* qLDiagSqrtInv = d->qLDiagSqrtInv;
const int* dofMadr = m->dof_Madr;
mju_zero(res, nv);
@@ -1117,9 +1116,9 @@ void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec)
}
}
// res = sqrt(D) * res
// res *= sqrt(D)
for (int i=0; i < nv; i++) {
res[i] /= qLDiagSqrtInv[i];
res[i] *= mju_sqrt(qLD[dofMadr[i]]);
}
}
src/engine/engine_vis_interact.c
+5 -1
View File
@@ -541,6 +541,7 @@ void mjv_initPerturb(const mjModel* m, mjData* d, const mjvScene* scn, mjvPertur
mjtNum* jac = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* jacM2 = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* sqrtInvD = mjSTACKALLOC(d, nv, mjtNum);
// invalid selected body: return
if (sel <= 0 || sel >= m->nbody) {
@@ -554,8 +555,11 @@ void mjv_initPerturb(const mjModel* m, mjData* d, const mjvScene* scn, mjvPertur
mju_addTo3(selpos, d->xpos+3*sel);
// 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]]);
}
mj_jac(m, d, jac, NULL, selpos, sel);
mj_solveM2(m, d, jacM2, jac, 3);
mj_solveM2(m, d, jacM2, jac, sqrtInvD, 3);
mjtNum invmass = mju_dot(jacM2+0*nv, jacM2+0*nv, nv) +
mju_dot(jacM2+1*nv, jacM2+1*nv, nv) +
mju_dot(jacM2+2*nv, jacM2+2*nv, nv);