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
doc/includes/references.h
+2 -2
View File
@@ -1544,8 +1544,8 @@ struct mjModel_ {
int* D_rowadr; // full inertia: row addresses (nv x 1)
int* D_diag; // full inertia: index of diagonal element (nv x 1)
int* D_colind; // full inertia: column indices (nD x 1)
int* mapM2D; // index mapping from qM to D (nD x 1)
int* mapD2M; // index mapping from D to qM (nM x 1)
int* mapM2D; // index mapping from M to D (nD x 1)
int* mapD2M; // index mapping from D to M (nC x 1)
// compilation signature
uint64_t signature; // also held by the mjSpec that compiled this model
include/mujoco/mjmodel.h
+2 -2
View File
@@ -1234,8 +1234,8 @@ struct mjModel_ {
int* D_rowadr; // full inertia: row addresses (nv x 1)
int* D_diag; // full inertia: index of diagonal element (nv x 1)
int* D_colind; // full inertia: column indices (nD x 1)
int* mapM2D; // index mapping from qM to D (nD x 1)
int* mapD2M; // index mapping from D to qM (nM x 1)
int* mapM2D; // index mapping from M to D (nD x 1)
int* mapD2M; // index mapping from D to M (nC x 1)
// compilation signature
uint64_t signature; // also held by the mjSpec that compiled this model
include/mujoco/mjxmacro.h
+1 -1
View File
@@ -609,7 +609,7 @@
X ( int, D_diag, nv, 1 ) \
X ( int, D_colind, nD, 1 ) \
X ( int, mapM2D, nD, 1 ) \
X ( int, mapD2M, nM, 1 )
X ( int, mapD2M, nC, 1 )
//-------------------------------- mjData ----------------------------------------------------------
python/mujoco/introspect/structs.py
+3 -3
View File
@@ -4698,7 +4698,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='index mapping from qM to D',
doc='index mapping from M to D',
array_extent=('nD',),
),
StructFieldDecl(
@@ -4706,8 +4706,8 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='index mapping from D to qM',
array_extent=('nM',),
doc='index mapping from D to M',
array_extent=('nC',),
),
StructFieldDecl(
name='signature',
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");
src/user/user_model.cc
+9 -8
View File
@@ -4876,29 +4876,30 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
// sparsity structures
{
std::vector<int> remaining(m->nv);
std::vector<int> scratch(m->nv);
std::vector<int> count(m->nbody);
std::vector<int> M(m->nM);
std::vector<int> D(m->nD);
// make D
mj_makeDofDofSparse(m->nv, m->nC, m->nD, m->nM, m->dof_parentid, m->dof_simplenum,
m->D_rownnz, m->D_rowadr, m->D_diag, m->D_colind,
/*reduced=*/0, /*upper=*/1, remaining.data());
/*reduced=*/0, /*upper=*/1, scratch.data());
// make B
mj_makeBSparse(m->nv, m->nbody, m->nB, m->body_dofnum, m->body_parentid,
m->body_dofadr, m->B_rownnz, m->B_rowadr, m->B_colind, count.data());
// make C
// make M
mj_makeDofDofSparse(m->nv, m->nC, m->nD, m->nM, m->dof_parentid, m->dof_simplenum,
m->M_rownnz, m->M_rowadr, NULL, m->M_colind,
/*reduced=*/1, /*upper=*/0, remaining.data());
/*reduced=*/1, /*upper=*/0, scratch.data());
// make index mappings: mapM2D, mapD2M, mapM2M
mj_makeDofDofMaps(m->nv, m->nM, m->nC, m->nD, m->dof_Madr, m->dof_simplenum, m->dof_parentid,
m->D_rownnz, m->D_rowadr, m->D_colind, m->M_rownnz, m->M_rowadr,
m->mapM2D, m->mapD2M, m->mapM2M, remaining.data(), M.data(), D.data());
mj_makeDofDofMaps(m->nv, m->nM, m->nC, m->nD,
m->dof_Madr, m->dof_simplenum, m->dof_parentid,
m->D_rownnz, m->D_rowadr, m->D_colind,
m->M_rownnz, m->M_rowadr, m->M_colind,
m->mapM2D, m->mapD2M, m->mapM2M, M.data(), scratch.data());
}
// create data