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

Switch mjData.{qH,qLD} from full ("M") to reduced ("C") inertia matrix structure.

Browse Source 
PiperOrigin-RevId: 758273074
Change-Id: If1a2e663ea70044694af985e0119afd6d58115ac
This commit is contained in:
Yuval Tassa
2025-05-13 10:19:22 -07:00
committed by Copybara-Service
parent 436b5a8e1f
commit 627fffdef9
18 changed files with 99 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/includes/references.h
+5 -5
View File
@@ -272,7 +272,7 @@ struct mjData_ {
mjtNum* qM; // total inertia (sparse) (nM x 1) mjtNum* qM; // total inertia (sparse) (nM x 1)
// computed by mj_fwdPosition/mj_factorM // computed by mj_fwdPosition/mj_factorM
mjtNum* qLD; // L'*D*L factorization of M (sparse) (nM x 1) mjtNum* qLD; // L'*D*L factorization of M (sparse) (nC x 1)
mjtNum* qLDiagInv; // 1/diag(D) (nv x 1) mjtNum* qLDiagInv; // 1/diag(D) (nv x 1)
// computed by mj_collisionTree // computed by mj_collisionTree
@@ -305,7 +305,7 @@ struct mjData_ {
mjtNum* subtree_angmom; // angular momentum about subtree com (nbody x 3) mjtNum* subtree_angmom; // angular momentum about subtree com (nbody x 3)
// computed by mj_Euler or mj_implicit // computed by mj_Euler or mj_implicit
mjtNum* qH; // L'*D*L factorization of modified M (nM x 1) mjtNum* qH; // L'*D*L factorization of modified M (nC x 1)
mjtNum* qHDiagInv; // 1/diag(D) of modified M (nv x 1) mjtNum* qHDiagInv; // 1/diag(D) of modified M (nv x 1)
// computed by mj_resetData // computed by mj_resetData
@@ -396,9 +396,9 @@ struct mjData_ {
int* iM_rownnz; // inertia: non-zeros in each row (nidof x 1) int* iM_rownnz; // inertia: non-zeros in each row (nidof x 1)
int* iM_rowadr; // inertia: address of each row in iM_colind (nidof x 1) int* iM_rowadr; // inertia: address of each row in iM_colind (nidof x 1)
int* iM_diagnum; // inertia: num of consecutive diagonal elements (nidof x 1) int* iM_diagnum; // inertia: num of consecutive diagonal elements (nidof x 1)
int* iM_colind; // inertia: column indices of non-zeros (nM x 1) int* iM_colind; // inertia: column indices of non-zeros (nC x 1)
mjtNum* iM; // total inertia (sparse) (nM x 1) mjtNum* iM; // total inertia (sparse) (nC x 1)
mjtNum* iLD; // L'*D*L factorization of M (sparse) (nM x 1) mjtNum* iLD; // L'*D*L factorization of M (sparse) (nC x 1)
mjtNum* iLDiagInv; // 1/diag(D) (nidof x 1) mjtNum* iLDiagInv; // 1/diag(D) (nidof x 1)
mjtNum* iacc; // acceleration (nidof x 1) mjtNum* iacc; // acceleration (nidof x 1)
include/mujoco/mjdata.h
+5 -5
View File
@@ -300,7 +300,7 @@ struct mjData_ {
mjtNum* qM; // total inertia (sparse) (nM x 1) mjtNum* qM; // total inertia (sparse) (nM x 1)
// computed by mj_fwdPosition/mj_factorM // computed by mj_fwdPosition/mj_factorM
mjtNum* qLD; // L'*D*L factorization of M (sparse) (nM x 1) mjtNum* qLD; // L'*D*L factorization of M (sparse) (nC x 1)
mjtNum* qLDiagInv; // 1/diag(D) (nv x 1) mjtNum* qLDiagInv; // 1/diag(D) (nv x 1)
// computed by mj_collisionTree // computed by mj_collisionTree
@@ -333,7 +333,7 @@ struct mjData_ {
mjtNum* subtree_angmom; // angular momentum about subtree com (nbody x 3) mjtNum* subtree_angmom; // angular momentum about subtree com (nbody x 3)
// computed by mj_Euler or mj_implicit // computed by mj_Euler or mj_implicit
mjtNum* qH; // L'*D*L factorization of modified M (nM x 1) mjtNum* qH; // L'*D*L factorization of modified M (nC x 1)
mjtNum* qHDiagInv; // 1/diag(D) of modified M (nv x 1) mjtNum* qHDiagInv; // 1/diag(D) of modified M (nv x 1)
// computed by mj_resetData // computed by mj_resetData
@@ -424,9 +424,9 @@ struct mjData_ {
int* iM_rownnz; // inertia: non-zeros in each row (nidof x 1) int* iM_rownnz; // inertia: non-zeros in each row (nidof x 1)
int* iM_rowadr; // inertia: address of each row in iM_colind (nidof x 1) int* iM_rowadr; // inertia: address of each row in iM_colind (nidof x 1)
int* iM_diagnum; // inertia: num of consecutive diagonal elements (nidof x 1) int* iM_diagnum; // inertia: num of consecutive diagonal elements (nidof x 1)
int* iM_colind; // inertia: column indices of non-zeros (nM x 1) int* iM_colind; // inertia: column indices of non-zeros (nC x 1)
mjtNum* iM; // total inertia (sparse) (nM x 1) mjtNum* iM; // total inertia (sparse) (nC x 1)
mjtNum* iLD; // L'*D*L factorization of M (sparse) (nM x 1) mjtNum* iLD; // L'*D*L factorization of M (sparse) (nC x 1)
mjtNum* iLDiagInv; // 1/diag(D) (nidof x 1) mjtNum* iLDiagInv; // 1/diag(D) (nidof x 1)
mjtNum* iacc; // acceleration (nidof x 1) mjtNum* iacc; // acceleration (nidof x 1)
include/mujoco/mjxmacro.h
+5 -5
View File
@@ -653,7 +653,7 @@
X ( mjtNum, actuator_moment, nJmom, 1 ) \ X ( mjtNum, actuator_moment, nJmom, 1 ) \
X ( mjtNum, crb, nbody, 10 ) \ X ( mjtNum, crb, nbody, 10 ) \
X ( mjtNum, qM, nM, 1 ) \ X ( mjtNum, qM, nM, 1 ) \
X ( mjtNum, qLD, nM, 1 ) \ X ( mjtNum, qLD, nC, 1 ) \
X ( mjtNum, qLDiagInv, nv, 1 ) \ X ( mjtNum, qLDiagInv, nv, 1 ) \
XMJV( mjtNum, bvh_aabb_dyn, nbvhdynamic, 6 ) \ XMJV( mjtNum, bvh_aabb_dyn, nbvhdynamic, 6 ) \
XMJV( mjtByte, bvh_active, nbvh, 1 ) \ XMJV( mjtByte, bvh_active, nbvh, 1 ) \
@@ -670,7 +670,7 @@
X ( mjtNum, qfrc_passive, nv, 1 ) \ X ( mjtNum, qfrc_passive, nv, 1 ) \
X ( mjtNum, subtree_linvel, nbody, 3 ) \ X ( mjtNum, subtree_linvel, nbody, 3 ) \
X ( mjtNum, subtree_angmom, nbody, 3 ) \ X ( mjtNum, subtree_angmom, nbody, 3 ) \
X ( mjtNum, qH, nM, 1 ) \ X ( mjtNum, qH, nC, 1 ) \
X ( mjtNum, qHDiagInv, nv, 1 ) \ X ( mjtNum, qHDiagInv, nv, 1 ) \
X ( int, B_rownnz, nbody, 1 ) \ X ( int, B_rownnz, nbody, 1 ) \
X ( int, B_rowadr, nbody, 1 ) \ X ( int, B_rowadr, nbody, 1 ) \
@@ -758,9 +758,9 @@
X( int, iM_rownnz, MJ_D(nidof), 1 ) \ X( int, iM_rownnz, MJ_D(nidof), 1 ) \
X( int, iM_rowadr, MJ_D(nidof), 1 ) \ X( int, iM_rowadr, MJ_D(nidof), 1 ) \
X( int, iM_diagnum, MJ_D(nidof), 1 ) \ X( int, iM_diagnum, MJ_D(nidof), 1 ) \
X( int, iM_colind, MJ_M(nM), 1 ) \ X( int, iM_colind, MJ_M(nC), 1 ) \
X( mjtNum, iM, MJ_M(nM), 1 ) \ X( mjtNum, iM, MJ_M(nC), 1 ) \
X( mjtNum, iLD, MJ_M(nM), 1 ) \ X( mjtNum, iLD, MJ_M(nC), 1 ) \
X( mjtNum, iLDiagInv, MJ_D(nidof), 1 ) \ X( mjtNum, iLDiagInv, MJ_D(nidof), 1 ) \
X( mjtNum, iacc, MJ_D(nidof), 1 ) \ X( mjtNum, iacc, MJ_D(nidof), 1 ) \
X( int, efc_island, MJ_D(nefc), 1 ) \ X( int, efc_island, MJ_D(nefc), 1 ) \
mjx/mujoco/mjx/_src/io.py
+3 -3
View File
@@ -509,7 +509,7 @@ def _make_data_jax(
'actuator_moment': (m.nu, m.nv, float_), 'actuator_moment': (m.nu, m.nv, float_),
'crb': (m.nbody, 10, float_), 'crb': (m.nbody, 10, float_),
'qM': (m.nM, float_) if support.is_sparse(m) else (m.nv, m.nv, float_), 'qM': (m.nM, float_) if support.is_sparse(m) else (m.nv, m.nv, float_),
'qLD': (m.nM, float_) if support.is_sparse(m) else (m.nv, m.nv, float_), 'qLD': (m.nC, float_) if support.is_sparse(m) else (m.nv, m.nv, float_),
'qLDiagInv': (m.nv, float_) if support.is_sparse(m) else (0, float_), 'qLDiagInv': (m.nv, float_) if support.is_sparse(m) else (0, float_),
'ten_velocity': (m.ntendon, float_), 'ten_velocity': (m.ntendon, float_),
'actuator_velocity': (m.nu, float_), 'actuator_velocity': (m.nu, float_),
@@ -618,8 +618,8 @@ def _make_data_c(
'flexedge_velocity': (nflexedge, float_), 'flexedge_velocity': (nflexedge, float_),
'crb': (m.nbody, 10, float_), 'crb': (m.nbody, 10, float_),
'qM': (m.nM, float_), 'qM': (m.nM, float_),
'qLD': (m.nM, float_), 'qLD': (m.nC, float_),
'qH': (m.nM, float_), 'qH': (m.nC, float_),
'qHDiagInv': (m.nv, float_), 'qHDiagInv': (m.nv, float_),
'qLDiagInv': (m.nv, float_), 'qLDiagInv': (m.nv, float_),
'ten_velocity': (m.ntendon, float_), 'ten_velocity': (m.ntendon, float_),
mjx/mujoco/mjx/_src/smooth_test.py
+2 -2
View File
@@ -91,8 +91,8 @@ class SmoothTest(absltest.TestCase):
# factor_m # factor_m
dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d)) dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
qLDLegacy = np.zeros(mx.nM) # pylint:disable=invalid-name qLDLegacy = np.zeros(mx.nM) # pylint:disable=invalid-name
for i in range(m.nM): for i in range(m.nC):
qLDLegacy[d.mapM2M[i]] = d.qLD[i] qLDLegacy[d.mapM2C[i]] = d.qLD[i]
_assert_eq(qLDLegacy, dx._impl.qLD, 'qLD') _assert_eq(qLDLegacy, dx._impl.qLD, 'qLD')
_assert_attr_eq(d, dx._impl, 'qLDiagInv') _assert_attr_eq(d, dx._impl, 'qLDiagInv')
# com_vel # com_vel
python/mujoco/introspect/structs.py
+5 -5
View File
@@ -5398,7 +5398,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
inner_type=ValueType(name='mjtNum'), inner_type=ValueType(name='mjtNum'),
), ),
doc="L'*D*L factorization of M (sparse)", doc="L'*D*L factorization of M (sparse)",
array_extent=('nM',), array_extent=('nC',),
), ),
StructFieldDecl( StructFieldDecl(
name='qLDiagInv', name='qLDiagInv',
@@ -5534,7 +5534,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
inner_type=ValueType(name='mjtNum'), inner_type=ValueType(name='mjtNum'),
), ),
doc="L'*D*L factorization of modified M", doc="L'*D*L factorization of modified M",
array_extent=('nM',), array_extent=('nC',),
), ),
StructFieldDecl( StructFieldDecl(
name='qHDiagInv', name='qHDiagInv',
@@ -6030,7 +6030,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
inner_type=ValueType(name='int'), inner_type=ValueType(name='int'),
), ),
doc='inertia: column indices of non-zeros', doc='inertia: column indices of non-zeros',
array_extent=('nM',), array_extent=('nC',),
), ),
StructFieldDecl( StructFieldDecl(
name='iM', name='iM',
@@ -6038,7 +6038,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
inner_type=ValueType(name='mjtNum'), inner_type=ValueType(name='mjtNum'),
), ),
doc='total inertia (sparse)', doc='total inertia (sparse)',
array_extent=('nM',), array_extent=('nC',),
), ),
StructFieldDecl( StructFieldDecl(
name='iLD', name='iLD',
@@ -6046,7 +6046,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
inner_type=ValueType(name='mjtNum'), inner_type=ValueType(name='mjtNum'),
), ),
doc="L'*D*L factorization of M (sparse)", doc="L'*D*L factorization of M (sparse)",
array_extent=('nM',), array_extent=('nC',),
), ),
StructFieldDecl( StructFieldDecl(
name='iLDiagInv', name='iLDiagInv',
src/engine/engine_core_constraint.c
+9 -9
View File
@@ -2039,7 +2039,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
// inverse square root of D from inertia LDL decomposition // inverse square root of D from inertia LDL decomposition
mjtNum* sqrtInvD = mjSTACKALLOC(d, nv, mjtNum); mjtNum* sqrtInvD = mjSTACKALLOC(d, nv, mjtNum);
for (int i=0; i < nv; i++) { for (int i=0; i < nv; i++) {
int diag = d->M_rowadr[i] + d->M_rownnz[i] - 1; int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]); sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]);
} }
@@ -2075,11 +2075,11 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
continue; continue;
} }
// traverse row j of M, marking new unique nonzeros // traverse row j of C, marking new unique nonzeros
int nnzM = d->M_rownnz[j]; int nnzC = d->C_rownnz[j];
int adrM = d->M_rowadr[j]; int adrC = d->C_rowadr[j];
for (int k=0; k < nnzM; k++) { for (int k=0; k < nnzC; k++) {
int c = d->M_colind[adrM + k]; int c = d->C_colind[adrC + k];
if (marker[c] != r) { if (marker[c] != r) {
marker[c] = r; marker[c] = r;
nnz++; nnz++;
@@ -2159,10 +2159,10 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
continue; continue;
} }
int j = B_colind[i]; int j = B_colind[i];
int adrM = d->M_rowadr[j]; int adrC = d->C_rowadr[j];
mju_addToSclSparseInc(B + adrB, d->qLD + adrM, mju_addToSclSparseInc(B + adrB, d->qLD + adrC,
nnzB, B_colind + adrB, nnzB, B_colind + adrB,
d->M_rownnz[j]-1, d->M_colind + adrM, -b); d->C_rownnz[j]-1, d->C_colind + adrC, -b);
} }
// B(r,:) <- sqrt(inv(D)) * B(r,:) // B(r,:) <- sqrt(inv(D)) * B(r,:)
src/engine/engine_core_smooth.c
+6 -7
View File
@@ -1653,9 +1653,8 @@ void mj_factorM(const mjModel* m, mjData* d) {
TM_START; TM_START;
// gather LD <- M (legacy to CSR) and factorize in-place // gather LD <- M (legacy to CSR) and factorize in-place
mju_gather(d->qLD, d->qM, d->mapM2M, m->nM); mju_gather(d->qLD, d->qM, d->mapM2C, m->nC);
mj_factorI(d->qLD, d->qLDiagInv, m->nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); 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); TM_ADD(mjTIMER_POS_INERTIA);
} }
@@ -1897,7 +1896,7 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {
mju_copy(x, y, n*m->nv); mju_copy(x, y, n*m->nv);
} }
mj_solveLD(x, d->qLD, d->qLDiagInv, m->nv, n, mj_solveLD(x, d->qLD, d->qLDiagInv, m->nv, n,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
} }
@@ -1908,9 +1907,9 @@ void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
int nv = m->nv; int nv = m->nv;
// local copies of key variables // local copies of key variables
const int* rownnz = d->M_rownnz; const int* rownnz = d->C_rownnz;
const int* rowadr = d->M_rowadr; const int* rowadr = d->C_rowadr;
const int* colind = d->M_colind; const int* colind = d->C_colind;
const int* diagnum = m->dof_simplenum; const int* diagnum = m->dof_simplenum;
const mjtNum* qLD = d->qLD; const mjtNum* qLD = d->qLD;
src/engine/engine_forward.c
+9 -9
View File
@@ -836,7 +836,7 @@ static void mj_advance(const mjModel* m, mjData* d,
// Euler integrator, semi-implicit in velocity, possibly skipping factorisation // Euler integrator, semi-implicit in velocity, possibly skipping factorisation
void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) { void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
TM_START; TM_START;
int nv = m->nv, nM = m->nM; int nv = m->nv, nC = m->nC;
mj_markStack(d); mj_markStack(d);
mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum); mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum);
mjtNum* qacc = mjSTACKALLOC(d, nv, mjtNum); mjtNum* qacc = mjSTACKALLOC(d, nv, mjtNum);
@@ -861,20 +861,20 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
else { else {
if (!skipfactor) { if (!skipfactor) {
// qH = M + h*diag(B) // qH = M + h*diag(B)
mju_gather(d->qH, d->qM, d->mapM2M, nM); mju_gather(d->qH, d->qM, d->mapM2C, nC);
for (int i=0; i < nv; i++) { for (int i=0; i < nv; i++) {
d->qH[d->M_rowadr[i] + d->M_rownnz[i] - 1] += m->opt.timestep * m->dof_damping[i]; d->qH[d->C_rowadr[i] + d->C_rownnz[i] - 1] += m->opt.timestep * m->dof_damping[i];
} }
// factorize in-place // factorize in-place
mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); mj_factorI(d->qH, d->qHDiagInv, nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
} }
// solve // solve
mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv); mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv);
mju_copy(qacc, qfrc, m->nv); mju_copy(qacc, qfrc, m->nv);
mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1, mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
} }
// advance state and time // advance state and time
@@ -1003,7 +1003,7 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) {
// fully implicit in velocity, possibly skipping factorization // fully implicit in velocity, possibly skipping factorization
void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
TM_START; TM_START;
int nv = m->nv, nM = m->nM, nD = m->nD; int nv = m->nv, nM = m->nM, nD = m->nD, nC = m->nC;
mj_markStack(d); mj_markStack(d);
mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum); mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum);
@@ -1047,16 +1047,16 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM); mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM);
// gather qH <- MhB (legacy to CSR) // gather qH <- MhB (legacy to CSR)
mju_gather(d->qH, MhB, d->mapM2M, nM); mju_gather(d->qH, MhB, d->mapM2C, nC);
// factorize in-place // factorize in-place
mj_factorI(d->qH, d->qHDiagInv, nv, d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); mj_factorI(d->qH, d->qHDiagInv, nv, d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
} }
// solve for qacc: (qM - dt*qDeriv) * qacc = qfrc // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
mju_copy(qacc, qfrc, nv); mju_copy(qacc, qfrc, nv);
mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1, mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
} else { } else {
mjERROR("integrator must be implicit or implicitfast"); mjERROR("integrator must be implicit or implicitfast");
src/engine/engine_island.c
+3 -3
View File
@@ -536,12 +536,12 @@ void mj_island(const mjModel* m, mjData* d) {
} }
// local CSR copy of qM // local CSR copy of qM
mjtNum* qM = mjSTACKALLOC(d, m->nM, mjtNum); mjtNum* qM = mjSTACKALLOC(d, m->nC, mjtNum);
mju_gather(qM, d->qM, d->mapM2M, m->nM); mju_gather(qM, d->qM, d->mapM2C, m->nC);
// inertia: block-diagonalize both iLD <- qLD and iM <- qM // inertia: block-diagonalize both iLD <- qLD and iM <- qM
mju_blockDiagSparse(d->iLD, d->iM_rownnz, d->iM_rowadr, d->iM_colind, mju_blockDiagSparse(d->iLD, d->iM_rownnz, d->iM_rowadr, d->iM_colind,
d->qLD, d->M_rownnz, d->M_rowadr, d->M_colind, d->qLD, d->C_rownnz, d->C_rowadr, d->C_colind,
nidof, nisland, nidof, nisland,
d->map_idof2dof, d->map_dof2idof, d->map_idof2dof, d->map_dof2idof,
d->island_idofadr, d->island_idofadr, d->island_idofadr, d->island_idofadr,
src/engine/engine_print.c
+3 -3
View File
@@ -1126,12 +1126,12 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
printInertia("QM", d->qM, m, fp, float_format); printInertia("QM", d->qM, m, fp, float_format);
printSparse("QLD", d->qLD, m->nv, d->M_rownnz, printSparse("QLD", d->qLD, m->nv, d->C_rownnz,
d->M_rowadr, d->M_colind, fp, float_format); d->C_rowadr, d->C_colind, fp, float_format);
printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format); printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format);
if (!mju_isZero(d->qHDiagInv, m->nv)) { if (!mju_isZero(d->qHDiagInv, m->nv)) {
printSparse("QH", d->qH, m->nv, d->M_rownnz, d->M_rowadr, d->M_colind, fp, float_format); printSparse("QH", d->qH, m->nv, d->C_rownnz, d->C_rowadr, d->C_colind, fp, float_format);
printArray("QHDIAGINV", m->nv, 1, d->qHDiagInv, fp, float_format); printArray("QHDIAGINV", m->nv, 1, d->qHDiagInv, fp, float_format);
} }
src/engine/engine_solver.c
+3 -3
View File
@@ -881,10 +881,10 @@ static void CGpointers(const mjModel* m, const mjData* d, mjCGContext* ctx, int
ctx->qacc = d->qacc; ctx->qacc = d->qacc;
// inertia // inertia
ctx->M_rownnz = d->M_rownnz; ctx->M_rownnz = d->C_rownnz;
ctx->M_rowadr = d->M_rowadr; ctx->M_rowadr = d->C_rowadr;
ctx->M_diagnum = m->dof_simplenum; ctx->M_diagnum = m->dof_simplenum;
ctx->M_colind = d->M_colind; ctx->M_colind = d->C_colind;
ctx->dof_Madr = m->dof_Madr; ctx->dof_Madr = m->dof_Madr;
ctx->dof_parentid = m->dof_parentid; ctx->dof_parentid = m->dof_parentid;
ctx->qM = d->qM; ctx->qM = d->qM;
src/engine/engine_support.c
+3 -3
View File
@@ -1047,14 +1047,14 @@ void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec)
// non-simple: add off-diagonals // non-simple: add off-diagonals
if (!m->dof_simplenum[i]) { if (!m->dof_simplenum[i]) {
int adr = d->M_rowadr[i]; int adr = d->C_rowadr[i];
res[i] += mju_dotSparse(qLD+adr, vec, d->M_rownnz[i] - 1, d->M_colind+adr); res[i] += mju_dotSparse(qLD+adr, vec, d->C_rownnz[i] - 1, d->C_colind+adr);
} }
} }
// res *= sqrt(D) // res *= sqrt(D)
for (int i=0; i < nv; i++) { for (int i=0; i < nv; i++) {
int diag = d->M_rowadr[i] + d->M_rownnz[i] - 1; int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
res[i] *= mju_sqrt(qLD[diag]); res[i] *= mju_sqrt(qLD[diag]);
} }
} }
test/benchmark/factorI_benchmark_test.cc
+3 -3
View File
@@ -45,8 +45,8 @@ static void BM_factorI(benchmark::State& state, bool legacy, bool coil) {
mj_markStack(d); mj_markStack(d);
// M: mass matrix in CSR format // M: mass matrix in CSR format
mjtNum* M = mj_stackAllocNum(d, m->nM); mjtNum* M = mj_stackAllocNum(d, m->nC);
mju_gather(M, d->qM, d->mapM2M, m->nM); mju_gather(M, d->qM, d->mapM2C, m->nC);
// LDlegacy: legacy LD matrix (size nM) // LDlegacy: legacy LD matrix (size nM)
mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM); mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM);
@@ -59,7 +59,7 @@ static void BM_factorI(benchmark::State& state, bool legacy, bool coil) {
} else { } else {
mju_copy(d->qLD, M, m->nC); mju_copy(d->qLD, M, m->nC);
mj_factorI(d->qLD, d->qLDiagInv, m->nv, mj_factorI(d->qLD, d->qLDiagInv, m->nv,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
} }
} }
} }
test/benchmark/inertia_benchmark_test.cc
+4 -4
View File
@@ -47,8 +47,8 @@ static void BM_solve(benchmark::State& state, SolveType type) {
mj_markStack(d); mj_markStack(d);
// M: mass matrix in CSR format // M: mass matrix in CSR format
mjtNum* M = mj_stackAllocNum(d, m->nM); mjtNum* M = mj_stackAllocNum(d, m->nC);
mju_gather(M, d->qM, d->mapM2M, m->nM); mju_gather(M, d->qM, d->mapM2C, m->nC);
// LDlegacy: legacy LD matrix (size nM) // LDlegacy: legacy LD matrix (size nM)
mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM); mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM);
@@ -73,9 +73,9 @@ static void BM_solve(benchmark::State& state, SolveType type) {
case SolveType::kCsr: case SolveType::kCsr:
mju_copy(d->qLD, M, m->nC); mju_copy(d->qLD, M, m->nC);
mj_factorI(d->qLD, d->qLDiagInv, m->nv, mj_factorI(d->qLD, d->qLDiagInv, m->nv,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
mj_solveLD(res, d->qLD, d->qLDiagInv, m->nv, 1, mj_solveLD(res, d->qLD, d->qLDiagInv, m->nv, 1,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
} }
} }
} }
test/benchmark/solveLD_benchmark_test.cc
+3 -2
View File
@@ -53,7 +53,8 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) {
// scatter into legacy matrix // scatter into legacy matrix
mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM); mjtNum* LDlegacy = mj_stackAllocNum(d, m->nM);
mju_scatter(LDlegacy, d->qLD, d->mapM2M, m->nM); mju_zero(LDlegacy, m->nM);
mju_scatter(LDlegacy, d->qLD, d->mapM2C, m->nC);
// benchmark // benchmark
while (state.KeepRunningBatch(kNumBenchmarkSteps)) { while (state.KeepRunningBatch(kNumBenchmarkSteps)) {
@@ -63,7 +64,7 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) {
mj_solveLD_legacy(m, res, 1, LDlegacy, d->qLDiagInv); mj_solveLD_legacy(m, res, 1, LDlegacy, d->qLDiagInv);
} else { } else {
mj_solveLD(res, d->qLD, d->qLDiagInv, m->nv, 1, mj_solveLD(res, d->qLD, d->qLDiagInv, m->nv, 1,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
} }
} }
} }
test/engine/engine_core_smooth_test.cc
+26 -28
View File
@@ -649,7 +649,7 @@ TEST_F(CoreSmoothTest, FactorI) {
int nv = model->nv; int nv = model->nv;
vector<mjtNum> Ldense(nv*nv, 0); vector<mjtNum> Ldense(nv*nv, 0);
mju_sparse2dense(Ldense.data(), data->qLD, nv, nv, mju_sparse2dense(Ldense.data(), data->qLD, nv, nv,
data->M_rownnz, data->M_rowadr, data->M_colind); data->C_rownnz, data->C_rowadr, data->C_colind);
for (int i=0; i < nv; i++) { for (int i=0; i < nv; i++) {
// set diagonal to 1 // set diagonal to 1
Ldense[i*nv+i] = 1; Ldense[i*nv+i] = 1;
@@ -658,7 +658,7 @@ TEST_F(CoreSmoothTest, FactorI) {
// dense D matrix // dense D matrix
vector<mjtNum> Ddense(nv*nv); vector<mjtNum> Ddense(nv*nv);
mju_sparse2dense(Ddense.data(), data->qLD, nv, nv, mju_sparse2dense(Ddense.data(), data->qLD, nv, nv,
data->M_rownnz, data->M_rowadr, data->M_colind); data->C_rownnz, data->C_rowadr, data->C_colind);
for (int i=0; i < nv; i++) { for (int i=0; i < nv; i++) {
for (int j=0; j < nv; j++) { for (int j=0; j < nv; j++) {
// zero everything except the diagonal // zero everything except the diagonal
@@ -694,15 +694,16 @@ TEST_F(CoreSmoothTest, SolveLDs) {
int nv = m->nv; int nv = m->nv;
int nM = m->nM; int nM = m->nM;
int nC = m->nC;
// scatter M into LD: Legacy format // copy M into LD: Legacy format
vector<mjtNum> LDlegacy(nM); vector<mjtNum> LDlegacy(nM, 0);
mju_scatter(LDlegacy.data(), d->qLD, d->mapM2M, nM); mju_scatter(LDlegacy.data(), d->qLD, d->mapM2C, nC);
// compare LD and LDs densified matrices // compare LD and LDs densified matrices
vector<mjtNum> LDdense(nv*nv); vector<mjtNum> LDdense(nv*nv);
mju_sparse2dense(LDdense.data(), d->qLD, nv, nv, mju_sparse2dense(LDdense.data(), d->qLD, nv, nv,
d->M_rownnz, d->M_rowadr, d->M_colind); d->C_rownnz, d->C_rowadr, d->C_colind);
vector<mjtNum> LDdense2(nv*nv); vector<mjtNum> LDdense2(nv*nv);
mj_fullM(m, LDdense2.data(), LDlegacy.data()); mj_fullM(m, LDdense2.data(), LDlegacy.data());
@@ -721,7 +722,7 @@ TEST_F(CoreSmoothTest, SolveLDs) {
mj_solveLD_legacy(m, vec.data(), 1, LDlegacy.data(), d->qLDiagInv); mj_solveLD_legacy(m, vec.data(), 1, LDlegacy.data(), d->qLDiagInv);
mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, 1, mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, 1,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
// expect vectors to match up to floating point precision // expect vectors to match up to floating point precision
for (int i=0; i < nv; i++) { for (int i=0; i < nv; i++) {
@@ -742,11 +743,10 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) {
mj_forward(m, d); mj_forward(m, d);
int nv = m->nv; int nv = m->nv;
int nM = m->nM;
// scatter LD into LDlegacy: Legacy format // copy LD into LDlegacy: Legacy format
vector<mjtNum> LDlegacy(nM); vector<mjtNum> LDlegacy(m->nM, 0);
mju_scatter(LDlegacy.data(), d->qLD, d->mapM2M, nM); mju_scatter(LDlegacy.data(), d->qLD, d->mapM2C, m->nC);
// compare n LD and LDs vector solve // compare n LD and LDs vector solve
int n = 3; int n = 3;
@@ -757,7 +757,7 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) {
mj_solveLD_legacy(m, vec.data(), n, LDlegacy.data(), d->qLDiagInv); mj_solveLD_legacy(m, vec.data(), n, LDlegacy.data(), d->qLDiagInv);
mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, n, mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, n,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
// expect vectors to match up to floating point precision // expect vectors to match up to floating point precision
for (int i=0; i < nv*n; i++) { for (int i=0; i < nv*n; i++) {
@@ -781,7 +781,7 @@ TEST_F(CoreSmoothTest, SolveM2) {
int nv = m->nv; int nv = m->nv;
vector<mjtNum> sqrtInvD(nv); vector<mjtNum> sqrtInvD(nv);
for (int i=0; i < nv; i++) { for (int i=0; i < nv; i++) {
int diag = d->M_rowadr[i] + d->M_rownnz[i] - 1; int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]); sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]);
} }
@@ -795,7 +795,7 @@ TEST_F(CoreSmoothTest, SolveM2) {
mj_solveM2(m, d, res.data(), vec.data(), sqrtInvD.data(), n); mj_solveM2(m, d, res.data(), vec.data(), sqrtInvD.data(), n);
mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, n, mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, n,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
// expect equality of dot(v, M^-1 * v) and dot(M^-1/2 * v, M^-1/2 * v) // expect equality of dot(v, M^-1 * v) and dot(M^-1/2 * v, M^-1/2 * v)
for (int i=0; i < n; i++) { for (int i=0; i < n; i++) {
@@ -816,27 +816,25 @@ TEST_F(CoreSmoothTest, FactorIs) {
mjData* d = mj_makeData(m); mjData* d = mj_makeData(m);
mj_forward(m, d); mj_forward(m, d);
int nM = m->nM, nv = m->nv; int nC = m->nC, nM = m->nM, nv = m->nv;
// copy qM into into qLDlegacy and factorize // copy qM into into qLDlegacy and factorize
vector<mjtNum> qLDlegacy(nM); vector<mjtNum> qLDlegacy(nM);
mj_factorI_legacy(m, d, d->qM, qLDlegacy.data(), d->qLDiagInv); mj_factorI_legacy(m, d, d->qM, qLDlegacy.data(), d->qLDiagInv);
// copy qLDlegacy into qLDexpected: CSR format // copy qLDlegacy into qLDexpected: CSR format
vector<mjtNum> qLDexpected(nM); vector<mjtNum> qLDexpected(nC);
for (int i=0; i < nM; i++) { mju_gather(qLDexpected.data(), qLDlegacy.data(), d->mapM2C, nC);
qLDexpected[i] = qLDlegacy[d->mapM2M[i]];
}
// gather qM into qLD: CSR format // copy qM into qLD: CSR format
vector<mjtNum> qLD(nM); vector<mjtNum> qLD(nC);
mju_gather(qLD.data(), d->qM, d->mapM2M, nM); mju_gather(qLD.data(), d->qM, d->mapM2C, nC);
vector<mjtNum> qLDiagInvExpected(d->qLDiagInv, d->qLDiagInv + nv); vector<mjtNum> qLDiagInvExpected(d->qLDiagInv, d->qLDiagInv + nv);
vector<mjtNum> qLDiagInv(nv, 0); vector<mjtNum> qLDiagInv(nv, 0);
mj_factorI(qLD.data(), qLDiagInv.data(), nv, mj_factorI(qLD.data(), qLDiagInv.data(), nv,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
// expect outputs to match to floating point precision // expect outputs to match to floating point precision
EXPECT_THAT(qLD, Pointwise(DoubleNear(1e-12), qLDexpected)); EXPECT_THAT(qLD, Pointwise(DoubleNear(1e-12), qLDexpected));
@@ -846,12 +844,12 @@ TEST_F(CoreSmoothTest, FactorIs) {
vector<mjtNum> LDdense(nv*nv); vector<mjtNum> LDdense(nv*nv);
mju_sparse2dense(LDdense.data(), qLDexpected.data(), nv, nv, mju_sparse2dense(LDdense.data(), qLDexpected.data(), nv, nv,
d->M_rownnz, d->M_rowadr, d->M_colind); d->C_rownnz, d->C_rowadr, d->C_colind);
PrintMatrix(LDdense.data(), nv, nv, 2, "qLDexpected"); PrintMatrix(LDdense.data(), nv, nv, 2);
mju_sparse2dense(LDdense.data(), qLD.data(), nv, nv, mju_sparse2dense(LDdense.data(), qLDs.data(), nv, nv,
d->M_rownnz, d->M_rowadr, d->M_colind); d->C_rownnz, d->C_rowadr, d->C_colind);
PrintMatrix(LDdense.data(), nv, nv, 2, "qLD"); PrintMatrix(LDdense.data(), nv, nv, 2);
*/ */
mj_deleteData(d); mj_deleteData(d);
test/engine/engine_derivative_test.cc
+2 -2
View File
@@ -436,7 +436,7 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) {
Ac[nv*nv + i*nv + i] = -m->dof_damping[i]; Ac[nv*nv + i*nv + i] = -m->dof_damping[i];
} }
mj_solveLD(Ac, d->qH, d->qHDiagInv, nv, 2*nv, mj_solveLD(Ac, d->qH, d->qHDiagInv, nv, 2*nv,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
// A = [dt*Ac; Ac] // A = [dt*Ac; Ac]
mju_transpose(A, Ac, 2*nv, nv); mju_transpose(A, Ac, 2*nv, nv);
@@ -464,7 +464,7 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) {
mju_sparse2dense(Bc, d->actuator_moment, nu, nv, d->moment_rownnz, mju_sparse2dense(Bc, d->actuator_moment, nu, nv, d->moment_rownnz,
d->moment_rowadr, d->moment_colind); d->moment_rowadr, d->moment_colind);
mj_solveLD(Bc, d->qH, d->qHDiagInv, nv, nu, mj_solveLD(Bc, d->qH, d->qHDiagInv, nv, nu,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind); d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
mju_transpose(BcT, Bc, nu, nv); mju_transpose(BcT, Bc, nu, nv);
mju_scl(B, BcT, dt*dt, nu*nv); mju_scl(B, BcT, dt*dt, nu*nv);
mju_scl(B+nu*nv, BcT, dt, nu*nv); mju_scl(B+nu*nv, BcT, dt, nu*nv);