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

Represent only the lower triangle in Newton solver's reduced dof-dof matrix.

Browse Source 
PiperOrigin-RevId: 712488529
Change-Id: Iad91c72654376539791d7856765a0d0ac9088251
This commit is contained in:
Yuval Tassa
2025-01-06 04:58:06 -08:00
committed by Copybara-Service
parent ee6f4837f3
commit 7eb8231fda
11 changed files with 112 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+2 -2
View File
@@ -2176,12 +2176,12 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
// AR = JM2 * JM2'
mju_sqrMatTDSparseInit(d->efc_AR_rownnz, d->efc_AR_rowadr, nefc, rownnzT,
rowadrT, colindT, rownnz, rowadr, colind, rowsuper, d);
rowadrT, colindT, rownnz, rowadr, colind, rowsuper, d, /*flg_upper=*/1);
mju_sqrMatTDSparse(d->efc_AR, JM2T, JM2, NULL, nv, nefc,
d->efc_AR_rownnz, d->efc_AR_rowadr, d->efc_AR_colind,
rownnzT, rowadrT, colindT, NULL,
rownnz, rowadr, colind, rowsuper, d);
rownnz, rowadr, colind, rowsuper, d, /*flg_upper=*/1);
// add R to diagonal of AR
for (int i=0; i < nefc; i++) {
src/engine/engine_io.c
+17 -8
View File
@@ -943,8 +943,11 @@ static void makeDofDofSparse(const mjModel* m, mjData* d,
// process below diagonal unless reduced and dof is simple
if (!reduced || !m->dof_simplenum[i]) {
while ((j = m->dof_parentid[j]) >= 0) {
// both reduced and non-reduced have lower triangle
rownnz[i]++;
rownnz[j]++;
// only non-reduced has upper triangle
if (!reduced) rownnz[j]++;
}
}
}
@@ -969,8 +972,11 @@ static void makeDofDofSparse(const mjModel* m, mjData* d,
remaining[i]--;
colind[rowadr[i] + remaining[i]] = j;
remaining[j]--;
colind[rowadr[j] + remaining[j]] = i;
// only non-reduced has upper triangle
if (!reduced) {
remaining[j]--;
colind[rowadr[j] + remaining[j]] = i;
}
}
}
}
@@ -1152,8 +1158,11 @@ static void copyM2Sparse(const mjModel* m, mjData* d, int* dst, const int* src,
remaining[i]--;
dst[rowadr[i] + remaining[i]] = src[adr];
remaining[j]--;
dst[rowadr[j] + remaining[j]] = src[adr];
// only non-reduced has upper triangle
if (!reduced) {
remaining[j]--;
dst[rowadr[j] + remaining[j]] = src[adr];
}
adr++;
}
@@ -1172,7 +1181,7 @@ static void copyM2Sparse(const mjModel* m, mjData* d, int* dst, const int* src,
// integer valued dst[M] = src[D lower], handle different sparsity representations
// integer valued dst[M] = src[D lower]
static void copyD2MSparse(const mjModel* m, const mjData* d, int* dst, const int* src) {
int nv = m->nv;
@@ -1197,7 +1206,7 @@ static void copyD2MSparse(const mjModel* m, const mjData* d, int* dst, const int
// construct index mappings between M <-> D and M -> C
static void makeDmap(const mjModel* m, mjData* d) {
static void makeDofDofmap(const mjModel* m, mjData* d) {
int nM = m->nM, nC = m->nC, nD = m->nD;
mj_markStack(d);
@@ -1955,7 +1964,7 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
// make C
makeDofDofSparse(m, d, d->C_rownnz, d->C_rowadr, d->C_diag, d->C_colind, /*reduced=*/1);
makeDmap(m, d);
makeDofDofmap(m, d);
}
// restore pluginstate and plugindata
src/engine/engine_solver.c
+14 -4
View File
@@ -1403,7 +1403,7 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) {
mju_sqrMatTDSparseInit(ctx->H_rownnz, ctx->H_rowadr, nv,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
d);
d, /*flg_upper=*/0);
// add nC to Hessian total nonzeros (unavoidable overcounting since H_colind is still unknown)
ctx->nH = m->nC + ctx->H_rowadr[nv - 1] + ctx->H_rownnz[nv - 1];
@@ -1424,14 +1424,24 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) {
ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
d);
d, /*flg_upper=*/0);
// add mass matrix: H = J'*D*J + C
mj_addMSparse(m, d, ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
ctx->C, d->C_rownnz, d->C_rowadr, d->C_colind);
// transiently compute H'; mju_cholFactorNNZ is memory-contiguous in upper triangle layout
mj_markStack(d);
int* HT_rownnz = mjSTACKALLOC(d, nv, int);
int* HT_rowadr = mjSTACKALLOC(d, nv, int);
int* HT_colind = mjSTACKALLOC(d, ctx->nH, int);
mju_transposeSparse(NULL, NULL, nv, nv,
HT_rownnz, HT_rowadr, HT_colind,
ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind);
// count total and row non-zeros of reverse-Cholesky factor L
ctx->nL = mju_cholFactorNNZ(ctx->L_rownnz, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind, nv, d);
ctx->nL = mju_cholFactorNNZ(ctx->L_rownnz, HT_rownnz, HT_rowadr, HT_colind, nv, d);
mj_freeStack(d);
// compute L row adresses: rowadr = cumsum(rownnz)
ctx->L_rowadr[0] = 0;
@@ -1508,7 +1518,7 @@ static void FactorizeHessian(const mjModel* m, mjData* d, mjCGContext* ctx,
ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
d);
d, /*flg_upper=*/0);
// add mass matrix: H = J'*D*J + C
mj_addMSparse(m, d, ctx->H, ctx->H_rownnz, ctx->H_rowadr, ctx->H_colind,
src/engine/engine_util_sparse.c
+24 -18
View File
@@ -628,7 +628,7 @@ void mju_superSparse(int nr, int* rowsuper,
void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr, int nr,
const int* rownnz, const int* rowadr, const int* colind,
const int* rownnzT, const int* rowadrT, const int* colindT,
const int* rowsuperT, mjData* d) {
const int* rowsuperT, mjData* d, int flg_upper) {
mj_markStack(d);
int* chain = mjSTACKALLOC(d, 2*nr, int);
int nchain = 0;
@@ -640,8 +640,10 @@ void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr, int nr,
res_rownnz[r] = res_rownnz[r - 1];
// fill in upper triangle
for (int j=0; j < nchain; j++) {
res_rownnz[res_colind[j]]++;
if (flg_upper) {
for (int j=0; j < nchain; j++) {
res_rownnz[res_colind[j]]++;
}
}
// update chain with diagonal
@@ -691,15 +693,17 @@ void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr, int nr,
res_colind = chain + inew;
// update upper triangle
int nchain_end = nchain;
if (flg_upper) {
int nchain_end = nchain;
// avoid double counting.
if (nchain > 0 && res_colind[nchain-1] == r) {
nchain_end = nchain - 1;
}
// avoid double counting
if (nchain > 0 && res_colind[nchain-1] == r) {
nchain_end = nchain - 1;
}
for (int j=0; j < nchain_end; j++) {
res_rownnz[res_colind[j]]++;
for (int j=0; j < nchain_end; j++) {
res_rownnz[res_colind[j]]++;
}
}
}
}
@@ -731,7 +735,7 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
const int* colind, const int* rowsuper,
const int* rownnzT, const int* rowadrT,
const int* colindT, const int* rowsuperT,
mjData* d) {
mjData* d, int flg_upper) {
// allocate space for accumulation buffer and matT
mj_markStack(d);
@@ -846,13 +850,15 @@ void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
// fill upper triangle
for (int i=0; i < nc; i++) {
int start = res_rowadr[i];
int end = start + res_rownnz[i] - 1;
for (int j=start; j < end; j++) {
int adr = res_rowadr[res_colind[j]] + res_rownnz[res_colind[j]]++;
res[adr] = res[j];
res_colind[adr] = i;
if (flg_upper) {
for (int i=0; i < nc; i++) {
int start = res_rowadr[i];
int end = start + res_rownnz[i] - 1;
for (int j=start; j < end; j++) {
int adr = res_rowadr[res_colind[j]] + res_rownnz[res_colind[j]]++;
res[adr] = res[j];
res_colind[adr] = i;
}
}
}
src/engine/engine_util_sparse.h
+2 -2
View File
@@ -96,13 +96,13 @@ MJAPI void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT
const int* colind, const int* rowsuper,
const int* rownnzT, const int* rowadrT,
const int* colindT, const int* rowsuperT,
mjData* d);
mjData* d, int flg_upper);
// precount res_rownnz and precompute res_rowadr for mju_sqrMatTDSparse
MJAPI void mju_sqrMatTDSparseInit(int* res_rownnz, int* res_rowadr, int nr,
const int* rownnz, const int* rowadr, const int* colind,
const int* rownnzT, const int* rowadrT, const int* colindT,
const int* rowsuperT, mjData* d);
const int* rowsuperT, mjData* d, int flg_upper);
// precompute res_rowadr for mju_sqrMatTDSparse using uncompressed memory
MJAPI void mju_sqrMatTDUncompressedInit(int* res_rowadr, int nc);