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

Skip shared dofs for contact constraints

Browse Source 
PiperOrigin-RevId: 884455239
Change-Id: I7ee36a13c28be88e12380b4790f150ad62f268c1
This commit is contained in:
Taylor Howell
2026-03-16 08:36:14 -07:00
committed by Copybara-Service
parent eca5758bf9
commit bb38a34869
8 changed files with 154 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+78 -56
View File
@@ -468,7 +468,8 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
// compute Jacobian difference (opposite of contact: 0 - 1)
NV = mj_jacDifPair(m, d, chain, body_id[1], body_id[0], pos[1], pos[0],
jac[1], jac[0], jacdif, NULL, NULL, NULL, issparse);
jac[1], jac[0], jacdif, NULL, NULL, NULL, issparse,
/*flg_skipcommon=*/0);
// copy difference into jac[0]
mju_copy(jac[0], jacdif, 3*NV);
@@ -504,7 +505,8 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
// compute error Jacobian (opposite of contact: 0 - 1)
NV = mj_jacDifPair(m, d, chain, body_id[1], body_id[0], pos[1], pos[0],
jac[1], jac[0], jacdif,
jac[1]+3*nv, jac[0]+3*nv, jacdif+3*nv, issparse);
jac[1]+3*nv, jac[0]+3*nv, jacdif+3*nv, issparse,
/*flg_skipcommon=*/0);
// copy difference into jac[0], compress translation:rotation if sparse
mju_copy(jac[0], jacdif, 3*NV);
@@ -670,7 +672,8 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
for (int n = 0; n < nodenum; n++) {
int chain_nnz = mj_bodyChain(m, bodyid[n], chain_col);
mju_zero(blk_jac, 3*nv);
mj_jacSparse(m, d, blk_jac, NULL, xpos + 3*n, bodyid[n], chain_nnz, chain_col);
mj_jacSparse(m, d, blk_jac, NULL, xpos + 3*n, bodyid[n], chain_nnz, chain_col,
/*flg_skipcommon=*/0);
// expand sparse Jacobian to dense row format
for (int r = 0; r < 3; r++) {
@@ -1297,14 +1300,13 @@ int mj_contactJacobian(const mjModel* m, mjData* d, const mjContact* con, int di
m->geom_bodyid[con->geom[side]] :
m->flex_vertbodyid[m->flex_vertadr[con->flex[side]] + con->vert[side]];
}
// compute Jacobian differences
// compute Jacobian differences, skipping common dofs
if (dim > 3) {
return mj_jacDifPair(m, d, chain, bid[0], bid[1], con->pos, con->pos,
jac1p, jac2p, jacdifp, jac1r, jac2r, jacdifr, mj_isSparse(m));
jac1p, jac2p, jacdifp, jac1r, jac2r, jacdifr, mj_isSparse(m), 1);
} else {
return mj_jacDifPair(m, d, chain, bid[0], bid[1], con->pos, con->pos,
jac1p, jac2p, jacdifp, NULL, NULL, NULL, mj_isSparse(m));
jac1p, jac2p, jacdifp, NULL, NULL, NULL, mj_isSparse(m), 1);
}
}
@@ -2067,7 +2069,8 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) {
}
NV = mj_jacDifPair(m, NULL, chain, id[1], id[0], NULL, NULL,
NULL, NULL, NULL, NULL, NULL, NULL, issparse);
NULL, NULL, NULL, NULL, NULL, NULL, issparse,
/*flg_skipcommon=*/0);
break;
case mjEQ_WELD:
@@ -2083,7 +2086,8 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) {
}
NV = mj_jacDifPair(m, NULL, chain, id[1], id[0], NULL, NULL,
NULL, NULL, NULL, NULL, NULL, NULL, issparse);
NULL, NULL, NULL, NULL, NULL, NULL, issparse,
/*flg_skipcommon=*/0);
break;
case mjEQ_JOINT:
@@ -2138,7 +2142,8 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) {
int b1 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e]];
int b2 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e+1]];
NV += mj_jacDifPair(m, NULL, chain, b1, b2, NULL, NULL,
NULL, NULL, NULL, NULL, NULL, NULL, issparse);
NULL, NULL, NULL, NULL, NULL, NULL, issparse,
/*flg_skipcommon=*/0);
}
}
break;
@@ -2252,55 +2257,72 @@ static int mj_nc(const mjModel* m, mjData* d, int* nnz) {
// compute NV only if nnz requested
int NV = 0;
if (nnz) {
// get bodies
int nb = 0, bid[729];
for (int side=0; side < 2; side++) {
// geom
if (con->geom[side] >= 0) {
bid[nb++] = m->geom_bodyid[con->geom[side]];
}
// flex
else {
int nw = 0;
int vid[4];
mjtNum vweight[4];
// flex vert
if (con->vert[side] >= 0) {
vid[nw++] = m->flex_vertadr[con->flex[side]] + con->vert[side];
vweight[0] = 1;
}
// flex elem
else {
int f = con->flex[side];
int fdim = m->flex_dim[f];
const int* edata = m->flex_elem + m->flex_elemdataadr[f] + con->elem[side]*(fdim+1);
for (int k=0; k <= fdim; k++) {
vid[nw++] = m->flex_vertadr[f] + edata[k];
}
if (m->flex_interp[f]) {
nw = mj_elemBodyWeight(m, d, con->flex[side], con->elem[side],
con->vert[1-side], con->pos, vid, vweight);
}
}
// get body or node ids and weights
if (m->flex_interp[con->flex[side]] == 0) {
for (int k=0; k < nw; k++) {
bid[nb] = m->flex_vertbodyid[vid[k]];
nb++;
}
} else {
nb += mj_vertBodyWeight(m, d, con->flex[side], vid, bid+nb, NULL, vweight, nw);
}
// single body on each side (geom-geom or flex vert-vert): skip common dofs
if ((con->geom[0] >= 0 || (con->vert[0] >= 0 && m->flex_interp[con->flex[0]] == 0)) &&
(con->geom[1] >= 0 || (con->vert[1] >= 0 && m->flex_interp[con->flex[1]] == 0))) {
// get bodies
int bid[2];
for (int side=0; side < 2; side++) {
bid[side] = (con->geom[side] >= 0) ?
m->geom_bodyid[con->geom[side]] :
m->flex_vertbodyid[m->flex_vertadr[con->flex[side]] + con->vert[side]];
}
NV = mj_jacDifPair(m, NULL, chain, bid[0], bid[1], NULL, NULL,
NULL, NULL, NULL, NULL, NULL, NULL, mj_isSparse(m), 1);
}
// count non-zeros in merged chain
NV = mj_jacSumCount(m, d, chain, nb, bid);
// general case: flex elements involved
else {
// get bodies
int nb = 0, bid[729];
for (int side=0; side < 2; side++) {
// geom
if (con->geom[side] >= 0) {
bid[nb++] = m->geom_bodyid[con->geom[side]];
}
// flex
else {
int nw = 0;
int vid[4];
mjtNum vweight[4];
// flex vert
if (con->vert[side] >= 0) {
vid[nw++] = m->flex_vertadr[con->flex[side]] + con->vert[side];
vweight[0] = 1;
}
// flex elem
else {
int f = con->flex[side];
int fdim = m->flex_dim[f];
const int* edata = m->flex_elem + m->flex_elemdataadr[f] + con->elem[side]*(fdim+1);
for (int k=0; k <= fdim; k++) {
vid[nw++] = m->flex_vertadr[f] + edata[k];
}
if (m->flex_interp[f]) {
nw = mj_elemBodyWeight(m, d, con->flex[side], con->elem[side],
con->vert[1-side], con->pos, vid, vweight);
}
}
// get body or node ids and weights
if (m->flex_interp[con->flex[side]] == 0) {
for (int k=0; k < nw; k++) {
bid[nb] = m->flex_vertbodyid[vid[k]];
nb++;
}
} else {
nb += mj_vertBodyWeight(m, d, con->flex[side], vid, bid+nb, NULL, vweight, nw);
}
}
}
// count non-zeros in merged chain
NV = mj_jacSumCount(m, d, chain, nb, bid);
}
if (!NV) {
continue;
}
src/engine/engine_core_smooth.c
+10 -5
View File
@@ -709,7 +709,8 @@ void mj_flex(const mjModel* m, mjData* d) {
// get endpoint Jacobians, subtract
int NV = mj_jacDifPair(m, d, chain, b1, b2, pos1, pos2,
jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1);
jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1,
/*skipcommon=*/0);
// no dofs: skip
if (!NV) {
@@ -853,9 +854,11 @@ void mj_flex(const mjModel* m, mjData* d) {
int b1 = m->flex_vertbodyid[vbase+v1];
int b2 = m->flex_vertbodyid[vbase+v2];
int NV1 = mj_bodyChain(m, b1, chain1);
mj_jacSparse(m, d, jac1, NULL, d->flexvert_xpos + 3*(vbase+v1), b1, NV1, chain1);
mj_jacSparse(m, d, jac1, NULL, d->flexvert_xpos + 3*(vbase+v1), b1, NV1, chain1,
/*flg_skipcommon=*/0);
int NV2 = mj_bodyChain(m, b2, chain2);
mj_jacSparse(m, d, jac2, NULL, d->flexvert_xpos + 3*(vbase+v2), b2, NV2, chain2);
mj_jacSparse(m, d, jac2, NULL, d->flexvert_xpos + 3*(vbase+v2), b2, NV2, chain2,
/*flg_skipcommon=*/0);
// accumulate dense Jacobians for vertex v
mju_mulMatTVec(J_local, jac1, dI1dy1, 3, NV1);
@@ -1042,7 +1045,8 @@ void mj_tendon(const mjModel* m, mjData* d) {
// get endpoint Jacobians, subtract
int NV = mj_jacDifPair(m, d, chain,
wbody[k], wbody[k+1], wpnt+3*k, wpnt+3*k+3,
jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1);
jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1,
/*skipcommon=*/0);
// no dofs: skip
if (!NV) {
@@ -1638,7 +1642,8 @@ void mj_transmission(const mjModel* m, mjData* d) {
// get Jacobian difference
int NV = mj_jacDifPair(m, d, chain, b1, b2, con->pos, con->pos,
jac1p, jac2p, jacdifp, NULL, NULL, NULL, issparse);
jac1p, jac2p, jacdifp, NULL, NULL, NULL, issparse,
/*flg_skipcommon=*/0);
// project Jacobian along the normal of the contact frame
mju_mulMatMat(jac, con->frame, jacdifp, 1, 3, NV);
src/engine/engine_core_util.c
+13 -8
View File
@@ -303,7 +303,7 @@ void mj_jacPointAxis(const mjModel* m, mjData* d, mjtNum* jacPoint, mjtNum* jacA
// compute 3/6-by-nv sparse Jacobian of global point attached to given body
void mj_jacSparse(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body,
int NV, const int* chain) {
int NV, const int* chain, int flg_skipcommon) {
// clear jacobians
if (jacp) {
mju_zero(jacp, 3*NV);
@@ -337,8 +337,12 @@ void mj_jacSparse(const mjModel* m, const mjData* d,
ci--;
}
// make sure we found it; SHOULD NOT OCCUR
// dof not in chain: skip if shared dofs are excluded, otherwise SHOULD NOT OCCUR
if (ci < 0 || chain[ci] != da) {
if (flg_skipcommon) {
da = m->dof_parentid[da];
continue;
}
mjERROR("dof index %d not found in chain", da);
}
@@ -433,7 +437,8 @@ void mj_jacSparseSimple(const mjModel* m, const mjData* d,
int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3],
mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp,
mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr, int issparse) {
mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr,
int issparse, int flg_skipcommon) {
int issimple = (m->body_simple[b1] && m->body_simple[b2]);
int NV = m->nv;
@@ -447,7 +452,7 @@ int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
if (issimple) {
NV = mj_mergeChainSimple(m, chain, b1, b2);
} else {
NV = mj_mergeChain(m, chain, b1, b2, 0);
NV = mj_mergeChain(m, chain, b1, b2, flg_skipcommon);
}
}
@@ -477,8 +482,8 @@ int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
// regular processing
else {
// Jacobians
mj_jacSparse(m, d, jac1p, jac1r, pos1, b1, NV, chain);
mj_jacSparse(m, d, jac2p, jac2r, pos2, b2, NV, chain);
mj_jacSparse(m, d, jac1p, jac1r, pos1, b1, NV, chain, flg_skipcommon);
mj_jacSparse(m, d, jac2p, jac2r, pos2, b2, NV, chain, flg_skipcommon);
// differences
if (jacdifp) {
@@ -535,7 +540,7 @@ int mj_jacSum(const mjModel* m, mjData* d, int* chain,
if (m->body_simple[body[0]]) {
mj_jacSparseSimple(m, d, jacp, jacr, point, body[0], 1, NV, 0);
} else {
mj_jacSparse(m, d, jacp, jacr, point, body[0], NV, chain);
mj_jacSparse(m, d, jacp, jacr, point, body[0], NV, chain, /*flg_skipcommon=*/0);
}
// apply weight
@@ -552,7 +557,7 @@ int mj_jacSum(const mjModel* m, mjData* d, int* chain,
if (m->body_simple[body[i]]) {
mj_jacSparseSimple(m, d, jp, jr, point, body[i], 1, bodyNV, 0);
} else {
mj_jacSparse(m, d, jp, jr, point, body[i], bodyNV, bodychain);
mj_jacSparse(m, d, jp, jr, point, body[i], bodyNV, bodychain, /*flg_skipcommon=*/0);
}
// combine sparse matrices
src/engine/engine_core_util.h
+3 -2
View File
@@ -78,7 +78,7 @@ MJAPI void mj_jacPointAxis(const mjModel* m, mjData* d,
// compute 3/6-by-nv sparse Jacobian of global point attached to given body
void mj_jacSparse(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body,
int NV, const int* chain);
int NV, const int* chain, int flg_skipcommon);
// sparse Jacobian difference for simple body contacts
void mj_jacSparseSimple(const mjModel* m, const mjData* d,
@@ -89,7 +89,8 @@ void mj_jacSparseSimple(const mjModel* m, const mjData* d,
MJAPI int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3],
mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp,
mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr, int issparse);
mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr,
int issparse, int flg_skipcommon);
// dense or sparse weighted sum of multiple body Jacobians at same point
int mj_jacSum(const mjModel* m, mjData* d, int* chain,
src/engine/engine_derivative.c
+5 -3
View File
@@ -980,7 +980,8 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
int chain_nnz = mj_bodyChain(m, bodyid[i], chain_colind);
// compute sparse Jacobian for this node (3 rows)
mj_jacSparse(m, d, blk_jac, NULL, xpos+3*i, bodyid[i], chain_nnz, chain_colind);
mj_jacSparse(m, d, blk_jac, NULL, xpos+3*i, bodyid[i], chain_nnz, chain_colind,
/*flg_skipcommon=*/0);
// copy to sparse structure
for (int r=0; r<3; r++) {
@@ -1488,7 +1489,8 @@ void mjd_ellipsoidFluid(const mjModel* m, mjData* d, int bodyid) {
// get geom global Jacobian: rotation then translation
if (mj_isSparse(m)) {
mj_jacSparse(m, d, J+3*nnz, J, d->geom_xpos+3*geomid, m->geom_bodyid[geomid], nnz, colind);
mj_jacSparse(m, d, J+3*nnz, J, d->geom_xpos+3*geomid, m->geom_bodyid[geomid], nnz, colind,
/*flg_skipcommon=*/0);
} else {
mj_jacGeom(m, d, J+3*nv, J, geomid);
}
@@ -1574,7 +1576,7 @@ void mjd_inertiaBoxFluid(const mjModel* m, mjData* d, int i) {
nnz = mj_bodyChain(m, i, colind);
// get sparse jacBodyCom
mj_jacSparse(m, d, J+3*nnz, J, d->xipos+3*i, i, nnz, colind);
mj_jacSparse(m, d, J+3*nnz, J, d->xipos+3*i, i, nnz, colind, /*flg_skipcommon=*/0);
// prepare rownnz, rowadr, colind for all 6 rows
rownnz[0] = nnz;
src/engine/engine_setconst.c
+1 -1
View File
@@ -434,7 +434,7 @@ static void makeFlexSparse(mjModel* m, mjData* d) {
// get sparsity
int NV = mj_jacDifPair(m, d, chain, b1, b2, dummy_pos, dummy_pos, NULL,
NULL, NULL, NULL, NULL, NULL, /*issparse=*/1);
NULL, NULL, NULL, NULL, NULL, /*issparse=*/1, /*skipcommon=*/0);
// copy sparsity info
rownnz[ebase + e] = NV;
src/engine/engine_support.c
+1 -1
View File
@@ -464,7 +464,7 @@ void mj_applyFT(const mjModel* m, mjData* d,
// construct chain and sparse Jacobians
int* chain = mjSTACKALLOC(d, nv, int);
int NV = mj_bodyChain(m, body, chain);
mj_jacSparse(m, d, jacp, jacr, point, body, NV, chain);
mj_jacSparse(m, d, jacp, jacr, point, body, NV, chain, /*flg_skipcommon=*/0);
// compute J'*f and accumulate
if (force) {
test/engine/engine_core_constraint_test.cc
+43 -3
View File
@@ -134,7 +134,8 @@ TEST_F(CoreConstraintTest, WeldRotJacobian) {
// rotational Jacobian difference
mj_jacDifPair(model, data, NULL, 2, 1, point, point,
NULL, NULL, NULL, jac0, jac1, jacdif, mj_isSparse(model));
NULL, NULL, NULL, jac0, jac1, jacdif, mj_isSparse(model),
/*flg_skipcommon=*/0);
// formula: 0.5 * neg(quat2) * (jac1-jac2) * quat1
mjtNum axis[3], quat3[4], quat4[4];
@@ -716,8 +717,8 @@ TEST_F(CoreConstraintTest, StrainConstraintNoPinning) {
// Rotate by 45 degrees around Z axis via quaternion
mjtNum angle = 0.785398; // 45 degrees
d->qpos[3] = mju_cos(angle/2); // w
d->qpos[4] = 0; // x
d->qpos[5] = 0; // y
d->qpos[4] = 0; // x
d->qpos[5] = 0; // y
d->qpos[6] = mju_sin(angle/2); // z
mj_forward(m, d);
@@ -743,5 +744,44 @@ TEST_F(CoreConstraintTest, StrainConstraintNoPinning) {
mj_deleteModel(m);
}
TEST_F(CoreConstraintTest, ContactSharedDofJacobian) {
constexpr char xml[] = R"(
<mujoco>
<option jacobian="sparse"/>
<worldbody>
<body pos="0 0 0.5">
<joint type="slide" axis="0 0 1"/>
<geom size="0.01"/>
<body pos="0 0.04 0">
<joint type="slide" axis="0 1 0"/>
<geom type="sphere" size="0.05" condim="1"/>
</body>
<body pos="0 -0.04 0">
<joint type="slide" axis="0 1 0"/>
<geom type="sphere" size="0.05" condim="1"/>
</body>
</body>
</worldbody>
</mujoco>
)";
char error[1024];
mjModel* model = LoadModelFromString(xml, error, sizeof(error));
ASSERT_THAT(model, NotNull()) << error;
ASSERT_EQ(model->nv, 3);
ASSERT_TRUE(mj_isSparse(model));
mjData* data = mj_makeData(model);
mj_forward(model, data);
ASSERT_EQ(data->ncon, 1);
ASSERT_GE(data->nefc, 1);
EXPECT_EQ(data->efc_J_rownnz[0], 2);
mj_deleteData(data);
mj_deleteModel(model);
}
} // namespace
} // namespace mujoco