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

Use mjData.M where appropriate

Browse Source 
PiperOrigin-RevId: 758374391
Change-Id: I9de7af7be8e41b5c300d0a04ea99082b4cefdec6
This commit is contained in:
Yuval Tassa
2025-05-13 14:21:41 -07:00
committed by Copybara-Service
parent 79c74d7eae
commit 4186589738
5 changed files with 23 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_smooth.c
+2 -3
View File
@@ -1579,6 +1579,7 @@ void mj_makeM(const mjModel* m, mjData* d) {
TM_START;
mj_crb(m, d);
mj_tendonArmature(m, d);
mju_gather(d->M, d->qM, d->mapM2C, m->nC);
TM_END(mjTIMER_POS_INERTIA);
}
@@ -1651,9 +1652,7 @@ void mj_factorI_legacy(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD
// sparse L'*D*L factorizaton of the inertia matrix M, assumed spd
void mj_factorM(const mjModel* m, mjData* d) {
TM_START;
// gather LD <- M (legacy to CSR) and factorize in-place
mju_gather(d->qLD, d->qM, d->mapM2C, m->nC);
mju_copy(d->qLD, d->M, m->nC);
mj_factorI(d->qLD, d->qLDiagInv, m->nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
TM_ADD(mjTIMER_POS_INERTIA);
}
src/engine/engine_forward.c
+1 -1
View File
@@ -861,7 +861,7 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
else {
if (!skipfactor) {
// qH = M + h*diag(B)
mju_gather(d->qH, d->qM, d->mapM2C, nC);
mju_copy(d->qH, d->M, nC);
for (int i=0; i < nv; i++) {
d->qH[d->C_rowadr[i] + d->C_rownnz[i] - 1] += m->opt.timestep * m->dof_damping[i];
}
src/engine/engine_island.c
+1 -5
View File
@@ -535,17 +535,13 @@ void mj_island(const mjModel* m, mjData* d) {
d->island_dofadr[i] = d->map_idof2dof[d->island_idofadr[i]];
}
// local CSR copy of qM
mjtNum* qM = mjSTACKALLOC(d, m->nC, mjtNum);
mju_gather(qM, d->qM, d->mapM2C, m->nC);
// inertia: block-diagonalize both iLD <- qLD and iM <- qM
mju_blockDiagSparse(d->iLD, d->iM_rownnz, d->iM_rowadr, d->iM_colind,
d->qLD, d->C_rownnz, d->C_rowadr, d->C_colind,
nidof, nisland,
d->map_idof2dof, d->map_dof2idof,
d->island_idofadr, d->island_idofadr,
d->iM, qM);
d->iM, d->M);
mju_gather(d->iLDiagInv, d->qLDiagInv, d->map_idof2dof, nidof);
// compute iM_diagnum (dof_simplenum per island)
src/engine/engine_print.c
+2 -2
View File
@@ -1123,9 +1123,9 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
printSparse("ACTUATOR_MOMENT", d->actuator_moment, m->nu, d->moment_rownnz,
d->moment_rowadr, d->moment_colind, fp, float_format);
printArray("CRB", m->nbody, 10, d->crb, fp, float_format);
printInertia("QM", d->qM, m, fp, float_format);
printSparse("M", d->M, m->nv, d->C_rownnz,
d->C_rowadr, d->C_colind, fp, float_format);
printSparse("QLD", d->qLD, m->nv, d->C_rownnz,
d->C_rowadr, d->C_colind, fp, float_format);
printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format);
src/engine/engine_solver.c
+17 -31
View File
@@ -787,9 +787,7 @@ struct _mjCGContext {
const int* M_rowadr;
const int* M_diagnum;
const int* M_colind;
const int* dof_Madr;
const int* dof_parentid;
const mjtNum* qM;
const mjtNum* M;
const mjtNum* qLD;
const mjtNum* qLDiagInv;
@@ -827,7 +825,6 @@ struct _mjCGContext {
// Newton arrays, known-size (CGallocate)
mjtNum* D; // constraint inertia (nefc x 1)
mjtNum* C; // reduced sparse inertia matrix (nC x 1)
int* H_rowadr; // Hessian row addresses (nv x 1)
int* H_rownnz; // Hessian row nonzeros (nv x 1)
int* H_lowernnz; // Hessian lower triangle row nonzeros (nv x 1)
@@ -885,9 +882,7 @@ static void CGpointers(const mjModel* m, const mjData* d, mjCGContext* ctx, int
ctx->M_rowadr = d->C_rowadr;
ctx->M_diagnum = m->dof_simplenum;
ctx->M_colind = d->C_colind;
ctx->dof_Madr = m->dof_Madr;
ctx->dof_parentid = m->dof_parentid;
ctx->qM = d->qM;
ctx->M = d->M;
ctx->qLD = d->qLD;
ctx->qLDiagInv = d->qLDiagInv;
@@ -936,7 +931,7 @@ static void CGpointers(const mjModel* m, const mjData* d, mjCGContext* ctx, int
ctx->M_rowadr = d->iM_rowadr + idofadr;
ctx->M_diagnum = d->iM_diagnum + idofadr;
ctx->M_colind = d->iM_colind;
ctx->qM = d->iM;
ctx->M = d->iM;
ctx->qLD = d->iLD;
ctx->qLDiagInv = d->iLDiagInv + idofadr;
@@ -1001,7 +996,6 @@ static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx, int island
// sparse Newton only
if (mj_isSparse(m)) {
ctx->C = mjSTACKALLOC(d, m->nC, mjtNum);
ctx->H_rowadr = mjSTACKALLOC(d, nv, int);
ctx->H_rownnz = mjSTACKALLOC(d, nv, int);
ctx->H_lowernnz = mjSTACKALLOC(d, nv, int);
@@ -1359,14 +1353,9 @@ static mjtNum CGsearch(mjCGContext* ctx, mjtNum tolerance, mjtNum ls_iterations)
mjtNum gtol = tolerance * snorm / ctx->scale;
mjtNum slopescl = ctx->scale / snorm;
// compute Mv = M * v (island or monolithic)
if (ctx->island >= 0) {
mju_mulSymVecSparse(ctx->Mv, ctx->qM, ctx->search, nv,
ctx->M_rownnz, ctx->M_rowadr, ctx->M_diagnum, ctx->M_colind);
} else {
mj_mulM_impl(ctx->Mv, ctx->search, nv, ctx->qM,
ctx->dof_Madr, ctx->dof_parentid, ctx->M_diagnum);
}
// compute Mv = M * v
mju_mulSymVecSparse(ctx->Mv, ctx->M, ctx->search, nv,
ctx->M_rownnz, ctx->M_rowadr, ctx->M_diagnum, ctx->M_colind);
// compute Jv = J * search (dense or sparse)
if (!ctx->J_rowadr) {
@@ -1545,9 +1534,6 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) {
// sparse
if (mj_isSparse(m)) {
// gather C <- qM (legacy to CSR)
mju_gather(ctx->C, d->qM, d->mapM2C, m->nC);
// initialize Hessian rowadr, rownnz
mju_sqrMatTDSparseCount(ctx->H_rownnz, ctx->H_rowadr, nv,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
@@ -1577,7 +1563,7 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) {
// 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);
ctx->M, d->C_rownnz, d->C_rowadr, d->C_colind);
// transiently compute H'; mju_cholFactorNNZ is memory-contiguous in upper triangle layout
mj_markStack(d);
@@ -1637,7 +1623,9 @@ static void MakeHessian(const mjModel* m, mjData* d, mjCGContext* ctx) {
// compute H = M + J'*D*J
mju_sqrMatTD(ctx->L, d->efc_J, ctx->D, nefc, nv);
mj_addMDense(m, d, ctx->L);
mju_addToSymSparse(ctx->L, ctx->M, ctx->nv,
ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind,
/*flg_upper=*/ 1);
}
}
@@ -1671,7 +1659,7 @@ static void FactorizeHessian(const mjModel* m, mjData* d, mjCGContext* ctx,
// 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);
ctx->M, d->C_rownnz, d->C_rowadr, d->C_colind);
}
// copy H lower-triangle into L, fill-in already accounted for
@@ -1702,7 +1690,9 @@ static void FactorizeHessian(const mjModel* m, mjData* d, mjCGContext* ctx,
// maybe compute H = M + J'*D*J
if (flg_recompute) {
mju_sqrMatTD(ctx->L, d->efc_J, ctx->D, nefc, nv);
mj_addMDense(m, d, ctx->L);
mju_addToSymSparse(ctx->L, ctx->M, ctx->nv,
ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind,
/*flg_upper=*/ 1);
}
// factorize H
@@ -1891,13 +1881,9 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
int* oldstate = mjSTACKALLOC(d, nefc, int);
// compute Ma = M * qacc (island or monolithic)
if (island >= 0) {
mju_mulSymVecSparse(ctx.Ma, ctx.qM, ctx.qacc, nv,
ctx.M_rownnz, ctx.M_rowadr, ctx.M_diagnum, ctx.M_colind);
} else {
mj_mulM_impl(ctx.Ma, ctx.qacc, nv, ctx.qM,
ctx.dof_Madr, ctx.dof_parentid, ctx.M_diagnum);
}
mju_mulSymVecSparse(ctx.Ma, ctx.M, ctx.qacc, nv,
ctx.M_rownnz, ctx.M_rowadr, ctx.M_diagnum, ctx.M_colind);
// compute Jaref = J * qacc - aref (dense or sparse)
if (!ctx.J_rownnz) {