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

Precompute mappings between M <-> D sparse representations.

Browse Source 
PiperOrigin-RevId: 670898353
Change-Id: I4b15477b1b6eb4a64cf7d4a0bb5af8fcb95d4e45
This commit is contained in:
Yuval Tassa
2024-09-04 02:56:05 -07:00
committed by Copybara-Service
parent 303e5e7718
commit 42eb669e4e
11 changed files with 156 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/includes/references.h
+2
View File
@@ -309,6 +309,8 @@ struct mjData_ {
int* D_rownnz; // non-zeros in each row (nv x 1)
int* D_rowadr; // address of each row in D_colind (nv x 1)
int* D_colind; // column indices of non-zeros (nD x 1)
int* mapM2D; // index mapping from M to D (nD x 1)
int* mapD2M; // index mapping from D to M (nM x 1)
int* B_rownnz; // non-zeros in each row (nbody x 1)
int* B_rowadr; // address of each row in B_colind (nbody x 1)
int* B_colind; // column indices of non-zeros (nB x 1)
include/mujoco/mjdata.h
+2
View File
@@ -337,6 +337,8 @@ struct mjData_ {
int* D_rownnz; // non-zeros in each row (nv x 1)
int* D_rowadr; // address of each row in D_colind (nv x 1)
int* D_colind; // column indices of non-zeros (nD x 1)
int* mapM2D; // index mapping from M to D (nD x 1)
int* mapD2M; // index mapping from D to M (nM x 1)
int* B_rownnz; // non-zeros in each row (nbody x 1)
int* B_rowadr; // address of each row in B_colind (nbody x 1)
int* B_colind; // column indices of non-zeros (nB x 1)
include/mujoco/mjxmacro.h
+2
View File
@@ -643,6 +643,8 @@
X ( int, D_rownnz, nv, 1 ) \
X ( int, D_rowadr, nv, 1 ) \
X ( int, D_colind, nD, 1 ) \
X ( int, mapM2D, nD, 1 ) \
X ( int, mapD2M, nM, 1 ) \
X ( int, B_rownnz, nbody, 1 ) \
X ( int, B_rowadr, nbody, 1 ) \
X ( int, B_colind, nB, 1 ) \
introspect/structs.py
+14
View File
@@ -4866,6 +4866,20 @@ STRUCTS: Mapping[str, StructDecl] = dict([
),
doc='column indices of non-zeros (nD x 1)', # pylint: disable=line-too-long
),
StructFieldDecl(
name='mapM2D',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='index mapping from M to D (nD x 1)', # pylint: disable=line-too-long
),
StructFieldDecl(
name='mapD2M',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='index mapping from D to M (nM x 1)', # pylint: disable=line-too-long
),
StructFieldDecl(
name='B_rownnz',
type=PointerType(
src/engine/engine_forward.c
+11 -7
View File
@@ -792,7 +792,7 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
mjtNum* MhB = mj_stackAllocNum(d, nM);
// MhB = M + h*diag(B)
mju_copy(MhB, d->qM, m->nM);
mju_copy(MhB, d->qM, nM);
for (int i=0; i < nv; i++) {
MhB[m->dof_Madr[i]] += m->opt.timestep * m->dof_damping[i];
}
@@ -933,7 +933,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;
int nv = m->nv, nM = m->nM, nD = m->nD;
mj_markStack(d);
mjtNum* qfrc = mj_stackAllocNum(d, nv);
@@ -949,7 +949,9 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
mjd_smooth_vel(m, d, /* flg_bias = */ 1);
// set qLU = qM
mj_copyM2DSparse(m, d, d->qLU, d->qM);
for (int i=0; i < nD; i++) {
d->qLU[i] = d->qM[d->mapM2D[i]];
}
// set qLU = qM - dt*qDeriv
mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD);
@@ -970,18 +972,20 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
mjd_smooth_vel(m, d, /* flg_bias = */ 0);
// modified mass matrix MhB = qDeriv[Lower]
mjtNum* MhB = mj_stackAllocNum(d, m->nM);
mj_copyD2MSparse(m, d, MhB, d->qDeriv);
mjtNum* MhB = mj_stackAllocNum(d, nM);
for (int i=0; i < nM; i++) {
MhB[i] = d->qDeriv[d->mapD2M[i]];
}
// set MhB = M - dt*qDeriv
mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, m->nM);
mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM);
// factorize
mj_factorI(m, d, MhB, d->qH, d->qHDiagInv, NULL);
}
// solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
mju_copy(qacc, qfrc, m->nv);
mju_copy(qacc, qfrc, nv);
mj_solveLD(m, qacc, 1, d->qH, d->qHDiagInv);
} else {
mjERROR("integrator must be implicit or implicitfast");
src/engine/engine_inverse.c
+7 -3
View File
@@ -71,7 +71,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, dof_damping;
int nv = m->nv, nM = m->nM, nD = m->nD, dof_damping;
mjtNum *qacc = d->qacc;
mj_markStack(d);
@@ -114,7 +114,9 @@ static void mj_discreteAcc(const mjModel* m, mjData* d) {
mjd_smooth_vel(m, d, /* flg_bias = */ 1);
// set qLU = qM
mj_copyM2DSparse(m, d, d->qLU, d->qM);
for (int i=0; i < nD; i++) {
d->qLU[i] = d->qM[d->mapM2D[i]];
}
// set qLU = qM - dt*qDeriv
mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD);
@@ -134,7 +136,9 @@ static void mj_discreteAcc(const mjModel* m, mjData* d) {
// set M = M - dt*qDeriv (reduced to M nonzeros)
mjtNum* qDerivReduced = mj_stackAllocNum(d, m->nM);
mj_copyD2MSparse(m, d, qDerivReduced, d->qDeriv);
for (int i=0; i < nM; i++) {
qDerivReduced[i] = d->qDeriv[d->mapD2M[i]];
}
mju_addToScl(d->qM, qDerivReduced, -m->opt.timestep, m->nM);
// set qfrc = (M - dt*qDeriv) * qacc
src/engine/engine_io.c
+102
View File
@@ -1059,6 +1059,107 @@ static void checkDBSparse(const mjModel* m, mjData* d) {
// integer valued dst[D] = src[M], handle different sparsity representations
static void copyM2DSparse(const mjModel* m, mjData* d, int* dst, const int* src) {
int nv = m->nv;
mj_markStack(d);
// init remaining
int* remaining = mj_stackAllocInt(d, nv);
mju_copyInt(remaining, d->D_rownnz, nv);
// copy data
for (int i = nv - 1; i >= 0; i--) {
// init at diagonal
int adr = m->dof_Madr[i];
remaining[i]--;
dst[d->D_rowadr[i] + remaining[i]] = src[adr];
adr++;
// process below diagonal
int j = i;
while ((j = m->dof_parentid[j]) >= 0) {
remaining[i]--;
dst[d->D_rowadr[i] + remaining[i]] = src[adr];
remaining[j]--;
dst[d->D_rowadr[j] + remaining[j]] = src[adr];
adr++;
}
}
// check that none remaining
for (int i=0; i < nv; i++) {
if (remaining[i]) {
mjERROR("unassigned index");
}
}
mj_freeStack(d);
}
// integer valued dst[M] = src[D lower], handle different sparsity representations
static void copyD2MSparse(const mjModel* m, const mjData* d, int* dst, const int* src) {
int nv = m->nv;
// copy data
for (int i = nv - 1; i >= 0; i--) {
// find diagonal in qDeriv
int j = 0;
while (d->D_colind[d->D_rowadr[i] + j] < i) {
j++;
}
// copy
int adr = m->dof_Madr[i];
while (j >= 0) {
dst[adr] = src[d->D_rowadr[i] + j];
adr++;
j--;
}
}
}
// construct index mappings between D <-> M
static void makeDmap(const mjModel* m, mjData* d) {
int nM = m->nM, nD = m->nD;
mj_markStack(d);
// make mapM2D
int* M = mj_stackAllocInt(d, nM);
for (int i=0; i < nM; i++) M[i] = i;
for (int i=0; i < nD; i++) d->mapM2D[i] = -1;
copyM2DSparse(m, d, d->mapM2D, M);
// check that all indices are filled in
for (int i=0; i < nD; i++) {
if (d->mapM2D[i] < 0) {
mjERROR("unassigned index in mapM2D");
}
}
// make mapD2M
int* D = mj_stackAllocInt(d, nD);
for (int i=0; i < nD; i++) D[i] = i;
for (int i=0; i < nM; i++) d->mapD2M[i] = -1;
copyD2MSparse(m, d, d->mapD2M, D);
// check that all indices are filled in
for (int i=0; i < nM; i++) {
if (d->mapD2M[i] < 0) {
mjERROR("unassigned index in mapD2M");
}
}
mj_freeStack(d);
}
//----------------------------------- mjData construction ------------------------------------------
// set pointers into mjData buffer
@@ -1707,6 +1808,7 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
makeDSparse(m, d);
makeBSparse(m, d);
checkDBSparse(m, d);
makeDmap(m, d);
}
// restore pluginstate and plugindata
src/engine/engine_print.c
+14
View File
@@ -1000,6 +1000,20 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
}
fprintf(fp, "\n\n");
// mapM2D
fprintf(fp, NAME_FORMAT, "mapM2D");
for (int i = 0; i < m->nD; i++) {
fprintf(fp, " %d", d->mapM2D[i]);
}
fprintf(fp, "\n\n");
// mapD2M
fprintf(fp, NAME_FORMAT, "mapD2M");
for (int i = 0; i < m->nM; i++) {
fprintf(fp, " %d", d->mapD2M[i]);
}
fprintf(fp, "\n\n");
// B_rownnz
fprintf(fp, NAME_FORMAT, "B_rownnz");
for (int i = 0; i < m->nbody; i++) {
src/engine/engine_support.c
-60
View File
@@ -1285,66 +1285,6 @@ void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst) {
}
//-------------------------- sparse system matrix conversion ---------------------------------------
// dst[D] = src[M], handle different sparsity representations
void mj_copyM2DSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* src) {
int nv = m->nv;
mj_markStack(d);
// init remaining
int* remaining = mj_stackAllocInt(d, nv);
mju_copyInt(remaining, d->D_rownnz, nv);
// copy data
for (int i = nv - 1; i >= 0; i--) {
// init at diagonal
int adr = m->dof_Madr[i];
remaining[i]--;
dst[d->D_rowadr[i] + remaining[i]] = src[adr];
adr++;
// process below diagonal
int j = i;
while ((j = m->dof_parentid[j]) >= 0) {
remaining[i]--;
dst[d->D_rowadr[i] + remaining[i]] = src[adr];
remaining[j]--;
dst[d->D_rowadr[j] + remaining[j]] = src[adr];
adr++;
}
}
mj_freeStack(d);
}
// dst[M] = src[D lower], handle different sparsity representations
void mj_copyD2MSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* src) {
int nv = m->nv;
// copy data
for (int i = nv - 1; i >= 0; i--) {
// find diagonal in qDeriv
int j = 0;
while (d->D_colind[d->D_rowadr[i] + j] < i) {
j++;
}
// copy
int adr = m->dof_Madr[i];
while (j >= 0) {
dst[adr] = src[d->D_rowadr[i] + j];
adr++;
j--;
}
}
}
//-------------------------- perturbations ---------------------------------------------------------
src/engine/engine_support.h
-9
View File
@@ -148,15 +148,6 @@ MJAPI void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
MJAPI void mj_addMDense(const mjModel* m, mjData* d, mjtNum* dst);
//-------------------------- sparse system matrix conversion ---------------------------------------
// dst[D] = src[M], handle different sparsity representations
MJAPI void mj_copyM2DSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* src);
// dst[M] = src[D lower], handle different sparsity representations
MJAPI void mj_copyD2MSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* src);
//-------------------------- perturbations ---------------------------------------------------------
// apply Cartesian force and torque
unity/Runtime/Bindings/MjBindings.cs
+2
View File
@@ -4924,6 +4924,8 @@ public unsafe struct mjData_ {
public int* D_rownnz;
public int* D_rowadr;
public int* D_colind;
public int* mapM2D;
public int* mapD2M;
public int* B_rownnz;
public int* B_rowadr;
public int* B_colind;