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

Use CSR matrices for dof-dof index mappings

Browse Source 
PiperOrigin-RevId: 795895311
Change-Id: I1000e0ee98906ef35990e73236e85d6354487b82
This commit is contained in:
Yuval Tassa
2025-08-16 13:05:45 -07:00
committed by Copybara-Service
parent c0096b51a4
commit 6b41fbe72a
11 changed files with 54 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_smooth.h
+1 -1
View File
@@ -51,7 +51,7 @@ MJAPI void mj_transmission(const mjModel* m, mjData* d);
// composite rigid body inertia algorithm
MJAPI void mj_crb(const mjModel* m, mjData* d);
// add tendon armature to qM
// add tendon armature to M
MJAPI void mj_tendonArmature(const mjModel* m, mjData* d);
// make inertia matrix
src/engine/engine_forward.c
+11 -15
View File
@@ -1009,7 +1009,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, nD = m->nD, nC = m->nC;
mj_markStack(d);
mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum);
@@ -1024,18 +1024,18 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
// compute analytical derivative qDeriv
mjd_smooth_vel(m, d, /* flg_bias = */ 1);
// gather qLU <- qM (lower to full)
mju_gather(d->qLU, d->qM, m->mapM2D, nD);
// gather qLU <- M (lower to full)
mju_gatherMasked(d->qLU, d->M, m->mapM2D, nD);
// set qLU = qM - dt*qDeriv
mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD);
// set qLU = M - dt*qDeriv
mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, nD);
// factorize qLU
int* scratch = mjSTACKALLOC(d, nv, int);
mju_factorLUSparse(d->qLU, nv, scratch, m->D_rownnz, m->D_rowadr, m->D_colind);
}
// solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
// solve for qacc: (M - dt*qDeriv) * qacc = qfrc
mju_solveLUSparse(qacc, d->qLU, qfrc, nv, m->D_rownnz, m->D_rowadr, m->D_diag, m->D_colind);
}
@@ -1045,21 +1045,17 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
// compute analytical derivative qDeriv; skip rne derivative
mjd_smooth_vel(m, d, /* flg_bias = */ 0);
// modified mass matrix: gather MhB <- qDeriv (full to lower)
mjtNum* MhB = mjSTACKALLOC(d, nM, mjtNum);
mju_gather(MhB, d->qDeriv, m->mapD2M, nM);
// modified mass matrix: gather qH <- qDeriv (full to lower)
mju_gather(d->qH, d->qDeriv, m->mapD2M, nC);
// set MhB = M - dt*qDeriv
mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM);
// gather qH <- MhB (legacy to CSR)
mju_gather(d->qH, MhB, m->mapM2M, nC);
// set qH = M - dt*qDeriv
mju_addScl(d->qH, d->M, d->qH, -m->opt.timestep, nC);
// factorize in-place
mj_factorI(d->qH, d->qHDiagInv, nv, m->M_rownnz, m->M_rowadr, m->M_colind);
}
// solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
// solve for qacc: (M - dt*qDeriv) * qacc = qfrc
mju_copy(qacc, qfrc, nv);
mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1,
m->M_rownnz, m->M_rowadr, m->M_colind);
src/engine/engine_inverse.c
+10 -14
View File
@@ -73,7 +73,7 @@ void mj_invVelocity(const mjModel* m, mjData* d) {
// convert discrete-time qacc to continuous-time qacc
static void mj_discreteAcc(const mjModel* m, mjData* d) {
int nv = m->nv, nM = m->nM, nD = m->nD, dof_damping;
int nv = m->nv, nC = m->nC, nD = m->nD, dof_damping;
mjtNum *qacc = d->qacc;
mj_markStack(d);
@@ -116,7 +116,7 @@ static void mj_discreteAcc(const mjModel* m, mjData* d) {
mjd_smooth_vel(m, d, /* flg_bias = */ 1);
// gather qLU <- qM (lower to full)
mju_gather(d->qLU, d->qM, m->mapM2D, nD);
mju_gatherMasked(d->qLU, d->M, m->mapM2D, nD);
// set qLU = qM - dt*qDeriv
mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD);
@@ -131,21 +131,17 @@ static void mj_discreteAcc(const mjModel* m, mjData* d) {
mjd_smooth_vel(m, d, /* flg_bias = */ 0);
// save mass matrix
mjtNum* qMsave = mjSTACKALLOC(d, m->nM, mjtNum);
mju_copy(qMsave, d->qM, m->nM);
mjtNum* Msave = mjSTACKALLOC(d, m->nC, mjtNum);
mju_copy(Msave, d->M, m->nC);
// set M = M - dt*qDeriv (reduced to M nonzeros)
mjtNum* qDerivReduced = mjSTACKALLOC(d, m->nM, mjtNum);
for (int i=0; i < nM; i++) {
qDerivReduced[i] = d->qDeriv[m->mapD2M[i]];
}
mju_addToScl(d->qM, qDerivReduced, -m->opt.timestep, m->nM);
// modified mass matrix: gather qH <- qDeriv (full to lower)
mju_gather(d->qH, d->qDeriv, m->mapD2M, nC);
// set qH = M - dt*qDeriv
mju_addScl(d->qH, d->M, d->qH, -m->opt.timestep, nC);
// set qfrc = (M - dt*qDeriv) * qacc
mj_mulM(m, d, qfrc, qacc);
// restore mass matrix
mju_copy(d->qM, qMsave, m->nM);
mju_mulSymVecSparse(qfrc, d->qH, qacc, m->nv, m->M_rownnz, m->M_rowadr, m->M_colind);
break;
}
src/engine/engine_io.c
+12 -53
View File
@@ -1183,70 +1183,29 @@ static void copyM2Sparse(int nv,
}
// integer valued dst[M] = src[D lower]
static void copyD2MSparse(int nv, const int* dof_Madr, const int* D_colind,
const int* D_rowadr, const int* src, int* dst) {
// copy data
for (int i = nv - 1; i >= 0; i--) {
// find diagonal in qDeriv
int j = 0;
while (D_colind[D_rowadr[i] + j] < i) {
j++;
}
// copy
int adr = dof_Madr[i];
while (j >= 0) {
dst[adr] = src[D_rowadr[i] + j];
adr++;
j--;
}
}
}
// construct index mappings between M <-> D, M -> C, M (legacy) -> M (CSR)
void mj_makeDofDofMaps(int nv, int nM, int nC, int nD,
const int* dof_Madr, const int* dof_simplenum, const int* dof_parentid,
const int* D_rownnz, const int* D_rowadr, const int* D_colind,
const int* M_rownnz, const int* M_rowadr,
const int* M_rownnz, const int* M_rowadr, const int* M_colind,
int* mapM2D, int* mapD2M, int* mapM2M,
int* remaining, int* M, int* D) {
// make mapM2D
int* M, int* scratch) {
// make mapM2D: M -> D (lower to symmetric)
mju_lower2SymMap(mapM2D, nv, D_rowadr, D_rownnz, D_colind, M_rowadr, M_rownnz, M_colind, scratch);
// make mapD2M: D -> M (symmetric to lower)
mju_sparseMap(mapD2M, nv, M_rowadr, M_rownnz, M_colind, D_rowadr, D_rownnz, D_colind);
// make mapM2M
for (int i=0; i < nM; i++) M[i] = i;
for (int i=0; i < nD; i++) mapM2D[i] = -1;
copyM2Sparse(nv, dof_Madr, dof_simplenum, dof_parentid, D_rownnz,
D_rowadr, M, mapM2D, /*reduced=*/0, /*upper=*/1, remaining);
// check that all indices are filled in
for (int i=0; i < nD; i++) {
if (mapM2D[i] < 0) {
mjERROR("unassigned index in mapM2D");
}
}
// make mapD2M
for (int i=0; i < nD; i++) D[i] = i;
for (int i=0; i < nM; i++) mapD2M[i] = -1;
copyD2MSparse(nv, dof_Madr, D_colind, D_rowadr, D, mapD2M);
// check that all indices are filled in
for (int i=0; i < nM; i++) {
if (mapD2M[i] < 0) {
mjERROR("unassigned index in mapD2M");
}
}
// make mapM2C
for (int i=0; i < nC; i++) mapM2M[i] = -1;
copyM2Sparse(nv, dof_Madr, dof_simplenum, dof_parentid, M_rownnz,
M_rowadr, M, mapM2M, /*reduced=*/1, /*upper=*/0, remaining);
M_rowadr, M, mapM2M, /*reduced=*/1, /*upper=*/0, scratch);
// check that all indices are filled in
for (int i=0; i < nC; i++) {
if (mapM2M[i] < 0) {
mjERROR("unassigned index in mapM2C");
mjERROR("unassigned index in mapM2M");
}
}
}
src/engine/engine_io.h
+2 -2
View File
@@ -101,9 +101,9 @@ MJAPI void mj_makeBSparse(int nv, int nbody, int nB,
MJAPI void mj_makeDofDofMaps(int nv, int nM, int nC, int nD,
const int* dof_Madr, const int* dof_simplenum, const int* dof_parentid,
const int* D_rownnz, const int* D_rowadr, const int* D_colind,
const int* M_rownnz, const int* M_rowadr,
const int* M_rownnz, const int* M_rowadr, const int* M_colind,
int* mapM2D, int* mapD2M, int* mapM2M,
int* remaining, int* M, int* D);
int* M, int* scratch);
//------------------------------- mjData -----------------------------------------------------------
src/engine/engine_print.c
+1 -1
View File
@@ -897,7 +897,7 @@ void mj_printFormattedModel(const mjModel* m, const char* filename, const char*
printArrayInt("D_ROWADR", 1, m->nv, m->D_rowadr, fp);
printArrayInt("D_COLIND", 1, m->nD, m->D_colind, fp);
printArrayInt("MAPM2D", 1, m->nD, m->mapM2D, fp);
printArrayInt("MAPD2M", 1, m->nM, m->mapD2M, fp);
printArrayInt("MAPD2M", 1, m->nC, m->mapD2M, fp);
// signature
fprintf(fp, "\nSIGNATURE\n");