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

Remove transposed Jacobian fields from mjData

Browse Source 
No longer computed unconditionally, only ad-hoc where required.

PiperOrigin-RevId: 797753297
Change-Id: I2030fc342c98ff33575b0526dd72e2110c4fcb74
This commit is contained in:
Yuval Tassa
2025-08-21 07:20:29 -07:00
committed by Copybara-Service
parent 5fe4224ae7
commit 2f28473bc1
10 changed files with 33 additions and 178 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/includes/references.h
-10
View File
@@ -349,12 +349,7 @@ struct mjData_ {
int* efc_J_rowadr; // row start address in colind array (nefc x 1)
int* efc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1)
int* efc_J_colind; // column indices in constraint Jacobian (nJ x 1)
int* efc_JT_rownnz; // number of non-zeros in constraint Jacobian row T (nv x 1)
int* efc_JT_rowadr; // row start address in colind array T (nv x 1)
int* efc_JT_rowsuper; // number of subsequent rows in supernode T (nv x 1)
int* efc_JT_colind; // column indices in constraint Jacobian T (nJ x 1)
mjtNum* efc_J; // constraint Jacobian (nJ x 1)
mjtNum* efc_JT; // constraint Jacobian transposed (nJ x 1)
mjtNum* efc_pos; // constraint position (equality, contact) (nefc x 1)
mjtNum* efc_margin; // inclusion margin (contact) (nefc x 1)
mjtNum* efc_frictionloss; // frictionloss (friction) (nefc x 1)
@@ -399,12 +394,7 @@ struct mjData_ {
int* iefc_J_rowadr; // row start address in colind array (nefc x 1)
int* iefc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1)
int* iefc_J_colind; // column indices in constraint Jacobian (nJ x 1)
int* iefc_JT_rownnz; // number of non-zeros in constraint Jacobian row T (nidof x 1)
int* iefc_JT_rowadr; // row start address in colind array T (nidof x 1)
int* iefc_JT_rowsuper; // number of subsequent rows in supernode T (nidof x 1)
int* iefc_JT_colind; // column indices in constraint Jacobian T (nJ x 1)
mjtNum* iefc_J; // constraint Jacobian (nJ x 1)
mjtNum* iefc_JT; // constraint Jacobian transposed (nJ x 1)
mjtNum* iefc_frictionloss; // frictionloss (friction) (nefc x 1)
mjtNum* iefc_D; // constraint mass (nefc x 1)
mjtNum* iefc_R; // inverse constraint mass (nefc x 1)
include/mujoco/mjdata.h
-10
View File
@@ -377,12 +377,7 @@ struct mjData_ {
int* efc_J_rowadr; // row start address in colind array (nefc x 1)
int* efc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1)
int* efc_J_colind; // column indices in constraint Jacobian (nJ x 1)
int* efc_JT_rownnz; // number of non-zeros in constraint Jacobian row T (nv x 1)
int* efc_JT_rowadr; // row start address in colind array T (nv x 1)
int* efc_JT_rowsuper; // number of subsequent rows in supernode T (nv x 1)
int* efc_JT_colind; // column indices in constraint Jacobian T (nJ x 1)
mjtNum* efc_J; // constraint Jacobian (nJ x 1)
mjtNum* efc_JT; // constraint Jacobian transposed (nJ x 1)
mjtNum* efc_pos; // constraint position (equality, contact) (nefc x 1)
mjtNum* efc_margin; // inclusion margin (contact) (nefc x 1)
mjtNum* efc_frictionloss; // frictionloss (friction) (nefc x 1)
@@ -427,12 +422,7 @@ struct mjData_ {
int* iefc_J_rowadr; // row start address in colind array (nefc x 1)
int* iefc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1)
int* iefc_J_colind; // column indices in constraint Jacobian (nJ x 1)
int* iefc_JT_rownnz; // number of non-zeros in constraint Jacobian row T (nidof x 1)
int* iefc_JT_rowadr; // row start address in colind array T (nidof x 1)
int* iefc_JT_rowsuper; // number of subsequent rows in supernode T (nidof x 1)
int* iefc_JT_colind; // column indices in constraint Jacobian T (nJ x 1)
mjtNum* iefc_J; // constraint Jacobian (nJ x 1)
mjtNum* iefc_JT; // constraint Jacobian transposed (nJ x 1)
mjtNum* iefc_frictionloss; // frictionloss (friction) (nefc x 1)
mjtNum* iefc_D; // constraint mass (nefc x 1)
mjtNum* iefc_R; // inverse constraint mass (nefc x 1)
include/mujoco/mjxmacro.h
-10
View File
@@ -729,12 +729,7 @@
XNV( int, efc_J_rowadr, MJ_D(nefc), 1 ) \
XNV( int, efc_J_rowsuper, MJ_D(nefc), 1 ) \
XNV( int, efc_J_colind, MJ_D(nJ), 1 ) \
XNV( int, efc_JT_rownnz, MJ_M(nv), 1 ) \
XNV( int, efc_JT_rowadr, MJ_M(nv), 1 ) \
XNV( int, efc_JT_rowsuper, MJ_M(nv), 1 ) \
XNV( int, efc_JT_colind, MJ_D(nJ), 1 ) \
XNV( mjtNum, efc_J, MJ_D(nJ), 1 ) \
XNV( mjtNum, efc_JT, MJ_D(nJ), 1 ) \
X ( mjtNum, efc_pos, MJ_D(nefc), 1 ) \
X ( mjtNum, efc_margin, MJ_D(nefc), 1 ) \
X ( mjtNum, efc_frictionloss, MJ_D(nefc), 1 ) \
@@ -786,12 +781,7 @@
XNV( int, iefc_J_rowadr, MJ_D(nefc), 1 ) \
XNV( int, iefc_J_rowsuper, MJ_D(nefc), 1 ) \
XNV( int, iefc_J_colind, MJ_D(nJ), 1 ) \
XNV( int, iefc_JT_rownnz, MJ_D(nidof), 1 ) \
XNV( int, iefc_JT_rowadr, MJ_D(nidof), 1 ) \
XNV( int, iefc_JT_rowsuper, MJ_D(nidof), 1 ) \
XNV( int, iefc_JT_colind, MJ_D(nJ), 1 ) \
XNV( mjtNum, iefc_J, MJ_D(nJ), 1 ) \
XNV( mjtNum, iefc_JT, MJ_D(nJ), 1 ) \
X ( mjtNum, iefc_frictionloss, MJ_D(nefc), 1 ) \
X ( mjtNum, iefc_D, MJ_D(nefc), 1 ) \
X ( mjtNum, iefc_R, MJ_D(nefc), 1 ) \
python/mujoco/introspect/structs.py
-80
View File
@@ -5914,38 +5914,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([
doc='column indices in constraint Jacobian',
array_extent=('nJ',),
),
StructFieldDecl(
name='efc_JT_rownnz',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='number of non-zeros in constraint Jacobian row T',
array_extent=('nv',),
),
StructFieldDecl(
name='efc_JT_rowadr',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='row start address in colind array T',
array_extent=('nv',),
),
StructFieldDecl(
name='efc_JT_rowsuper',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='number of subsequent rows in supernode T',
array_extent=('nv',),
),
StructFieldDecl(
name='efc_JT_colind',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='column indices in constraint Jacobian T',
array_extent=('nJ',),
),
StructFieldDecl(
name='efc_J',
type=PointerType(
@@ -5954,14 +5922,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([
doc='constraint Jacobian',
array_extent=('nJ',),
),
StructFieldDecl(
name='efc_JT',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='constraint Jacobian transposed',
array_extent=('nJ',),
),
StructFieldDecl(
name='efc_pos',
type=PointerType(
@@ -6250,38 +6210,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([
doc='column indices in constraint Jacobian',
array_extent=('nJ',),
),
StructFieldDecl(
name='iefc_JT_rownnz',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='number of non-zeros in constraint Jacobian row T',
array_extent=('nidof',),
),
StructFieldDecl(
name='iefc_JT_rowadr',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='row start address in colind array T',
array_extent=('nidof',),
),
StructFieldDecl(
name='iefc_JT_rowsuper',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='number of subsequent rows in supernode T',
array_extent=('nidof',),
),
StructFieldDecl(
name='iefc_JT_colind',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='column indices in constraint Jacobian T',
array_extent=('nJ',),
),
StructFieldDecl(
name='iefc_J',
type=PointerType(
@@ -6290,14 +6218,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([
doc='constraint Jacobian',
array_extent=('nJ',),
),
StructFieldDecl(
name='iefc_JT',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='constraint Jacobian transposed',
array_extent=('nJ',),
),
StructFieldDecl(
name='iefc_frictionloss',
type=PointerType(
src/engine/engine_core_constraint.c
+4 -10
View File
@@ -386,10 +386,10 @@ void mj_mulJacTVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum*
}
// sparse Jacobian
if (mj_isSparse(m))
mju_mulMatVecSparse(res, d->efc_JT, vec, m->nv,
d->efc_JT_rownnz, d->efc_JT_rowadr,
d->efc_JT_colind, d->efc_JT_rowsuper);
if (mj_isSparse(m)) {
mju_mulMatTVecSparse(res, d->efc_J, vec, d->nefc, m->nv,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
}
// dense Jacobian
else {
@@ -1996,12 +1996,6 @@ void mj_makeConstraint(const mjModel* m, mjData* d) {
// transpose sparse Jacobian, make row supernodes
if (mj_isSparse(m)) {
// transpose
mju_transposeSparse(d->efc_JT, d->efc_J, d->nefc, m->nv,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
#ifdef mjUSEAVX
// compute supernodes of J; used by mju_mulMatVecSparse_avx
mju_superSparse(d->nefc, d->efc_J_rowsuper,
src/engine/engine_island.c
-14
View File
@@ -597,20 +597,6 @@ void mj_island(const mjModel* m, mjData* d) {
mju_superSparse(d->island_nefc[island], d->iefc_J_rowsuper + adr,
d->iefc_J_rownnz + adr, d->iefc_J_rowadr + adr, d->iefc_J_colind);
}
// block-diagonalize Jacobian-transpose
mju_blockDiagSparse(d->iefc_JT, d->iefc_JT_rownnz, d->iefc_JT_rowadr, d->iefc_JT_colind,
d->efc_JT, d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind,
nidof, nisland,
d->map_idof2dof, d->map_efc2iefc,
d->island_idofadr, d->island_iefcadr, NULL, NULL);
// recompute rowsuper per island
for (int island=0; island < nisland; island++) {
int adr = d->island_idofadr[island];
mju_superSparse(d->island_nv[island], d->iefc_JT_rowsuper + adr,
d->iefc_JT_rownnz + adr, d->iefc_JT_rowadr + adr, d->iefc_JT_colind);
}
}
// copy position-dependent efc vectors required by solver
src/engine/engine_print.c
-2
View File
@@ -1236,8 +1236,6 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
printArrayInt("EFC_J_ROWADR", d->nefc, 1, d->efc_J_rowadr, fp);
printSparse("EFC_J", d->efc_J, d->nefc, d->efc_J_rownnz,
d->efc_J_rowadr, d->efc_J_colind, fp, float_format);
mj_printSparsity("JT: constraint Jacobian transposed", m->nv, d->nefc, d->efc_JT_rowadr, NULL,
d->efc_JT_rownnz, d->efc_JT_rowsuper, d->efc_JT_colind, fp);
if (mj_isDual(m)) {
printArrayInt("EFC_AR_ROWNNZ", d->nefc, 1, d->efc_AR_rownnz, fp);
printArrayInt("EFC_AR_ROWADR", d->nefc, 1, d->efc_AR_rowadr, fp);
test/benchmark/engine_util_sparse_benchmark_test.cc
+29 -11
View File
@@ -448,14 +448,23 @@ static void BM_combineSparse(benchmark::State& state, CombineFuncPtr func) {
}
}
int* JT_rownnz = mj_stackAllocInt(d, m->nv);
int* JT_rowadr = mj_stackAllocInt(d, m->nv);
int* JT_rowsuper = mj_stackAllocInt(d, m->nv);
int* JT_colind = mj_stackAllocInt(d, d->nJ);
mjtNum* JT = mj_stackAllocNum(d, d->nJ);
mju_transposeSparse(JT, d->efc_J, d->nefc, m->nv,
JT_rownnz, JT_rowadr, JT_colind, JT_rowsuper,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
// compute H = J'*D*J, uncompressed layout
mju_sqrMatTDUncompressedInit(rowadr, m->nv);
mju_sqrMatTDSparse(H, d->efc_J, d->efc_JT, D, d->nefc, m->nv,
mju_sqrMatTDSparse(H, d->efc_J, JT, D, d->nefc, m->nv,
rownnz, rowadr, colind,
d->efc_J_rownnz, d->efc_J_rowadr,
d->efc_J_colind, d->efc_J_rowsuper,
d->efc_JT_rownnz, d->efc_JT_rowadr,
d->efc_JT_colind, d->efc_JT_rowsuper, d,
JT_rownnz, JT_rowadr,
JT_colind, JT_rowsuper, d,
diagind);
// compute H = M + J'*D*J
@@ -605,19 +614,28 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
}
}
int* JT_rownnz = mj_stackAllocInt(d, m->nv);
int* JT_rowadr = mj_stackAllocInt(d, m->nv);
int* JT_rowsuper = mj_stackAllocInt(d, m->nv);
int* JT_colind = mj_stackAllocInt(d, d->nJ);
mjtNum* JT = mj_stackAllocNum(d, d->nJ);
mju_transposeSparse(JT, d->efc_J, d->nefc, m->nv,
JT_rownnz, JT_rowadr, JT_colind, JT_rowsuper,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
// time benchmark
if (func) {
mju_sqrMatTDSparseCount(rownnz, rowadr, m->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, nullptr, d, 1);
JT_rownnz, JT_rowadr,
JT_colind, nullptr, d, 1);
for (auto s : state) {
// compute H = J'*D*J, compressed layout
func(H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, rownnz, rowadr, colind,
func(H, d->efc_J, JT, D, d->nefc, m->nv, rownnz, rowadr, 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, diagind);
JT_rownnz, JT_rowadr, JT_colind,
JT_rowsuper, d, diagind);
}
} else {
for (auto s : state) {
@@ -627,10 +645,10 @@ static void BM_sqrMatTDSparse(benchmark::State& state, SqrMatTDFuncPtr func) {
// compute H = J'*D*J, uncompressed layout
mju_sqrMatTDSparse_baseline(
H, d->efc_J, d->efc_JT, D, d->nefc, m->nv, rownnz, rowadr, colind,
H, d->efc_J, JT, D, d->nefc, m->nv, rownnz, rowadr, colind,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, d->efc_J_rowsuper,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind,
d->efc_JT_rowsuper, d, /*unused=*/nullptr);
JT_rownnz, JT_rowadr, JT_colind,
JT_rowsuper, d, /*unused=*/nullptr);
}
}
test/engine/engine_island_test.cc
-21
View File
@@ -467,27 +467,6 @@ TEST_F(IslandTest, IslandJacobian) {
EXPECT_EQ(J_island[i * nv_island + j], J[efc * nv + dof]);
}
}
// === test JT (if sparse)
// get pointer to J_island, dense (nefc_island x nv_island) submatrix
if (jac == mjJAC_SPARSE) {
// dense copy of island in iJ (here used as scratch)
mju_sparse2dense(iJ, d->iefc_JT, nv_island, nefc_island,
d->iefc_JT_rownnz + idof,
d->iefc_JT_rowadr + idof,
d->iefc_JT_colind);
J_island = iJ;
// sequential memory in J_island equals random access memory in J
for (int i=0; i < nv_island; i++) {
for (int j=0; j < nefc_island; j++) {
int dof = d->map_idof2dof[idof + i];
int efc = d->map_iefc2efc[iefc + j];
EXPECT_EQ(J_island[i * nefc_island + j], J[efc * nv + dof]);
}
}
}
}
mju_free(iJ);
unity/Runtime/Bindings/MjBindings.cs
-10
View File
@@ -5001,12 +5001,7 @@ public unsafe struct mjData_ {
public int* efc_J_rowadr;
public int* efc_J_rowsuper;
public int* efc_J_colind;
public int* efc_JT_rownnz;
public int* efc_JT_rowadr;
public int* efc_JT_rowsuper;
public int* efc_JT_colind;
public double* efc_J;
public double* efc_JT;
public double* efc_pos;
public double* efc_margin;
public double* efc_frictionloss;
@@ -5043,12 +5038,7 @@ public unsafe struct mjData_ {
public int* iefc_J_rowadr;
public int* iefc_J_rowsuper;
public int* iefc_J_colind;
public int* iefc_JT_rownnz;
public int* iefc_JT_rowadr;
public int* iefc_JT_rowsuper;
public int* iefc_JT_colind;
public double* iefc_J;
public double* iefc_JT;
public double* iefc_frictionloss;
public double* iefc_D;
public double* iefc_R;