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

Fix sparse-path rotational Jacobian misalignment in mj_jacSum

Browse Source 
PiperOrigin-RevId: 951341459
Change-Id: I18bc27765b3147a5eb520e3026317e9a5a2dfc30
This commit is contained in:
Yuval Tassa
2026-07-21 02:01:24 -07:00
committed by Copybara-Service
parent a1f38c8e6e
commit 426cb5481d
6 changed files with 182 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+7 -8
View File
@@ -1533,8 +1533,8 @@ static int mj_instantiateLimit(const mjModel* m, mjData* d, int count_only, int*
// compute Jacobian for contact, return number of DOFs affected
int mj_contactJacobian(const mjModel* m, mjData* d, const mjContact* con, int dim,
mjtNum* jac, mjtNum* jacdif, mjtNum* jacdifp,
mjtNum* jacdifr, mjtNum* jac1p, mjtNum* jac2p,
mjtNum* jacdifp, mjtNum* jacdifr,
mjtNum* jac1p, mjtNum* jac2p,
mjtNum* jac1r, mjtNum* jac2r, int* chain) {
// special case: single body on each side
if ((con->geom[0] >= 0 || (con->vert[0] >= 0 && m->flex_interp[con->flex[0]] == 0)) &&
@@ -1608,7 +1608,7 @@ int mj_contactJacobian(const mjModel* m, mjData* d, const mjContact* con, int di
}
// combine weighted Jacobians
return mj_jacSum(m, d, chain, nb, bid, bweight, con->pos, jacdif, dim > 3);
return mj_jacSum(m, d, chain, nb, bid, bweight, con->pos, jacdifp, jacdifr, dim > 3);
}
}
@@ -1618,7 +1618,7 @@ void mj_instantiateContact(const mjModel* m, mjData* d) {
int ispyramid = mj_isPyramidal(m), issparse = mj_isSparse(m), ncon = d->ncon;
int dim, NV, nv = m->nv, *chain = NULL;
mjContact* con;
mjtNum cpos[6], cmargin[6], *jac, *jacdif, *jacdifp, *jacdifr, *jac1p, *jac2p, *jac1r, *jac2r;
mjtNum cpos[6], cmargin[6], *jac, *jacdifp, *jacdifr, *jac1p, *jac2p, *jac1r, *jac2r;
if (mjDISABLED(mjDSBL_CONTACT) || ncon == 0 || nv == 0) {
return;
@@ -1628,9 +1628,8 @@ void mj_instantiateContact(const mjModel* m, mjData* d) {
// allocate Jacobian
jac = mjSTACKALLOC(d, 6*nv, mjtNum);
jacdif = mjSTACKALLOC(d, 6*nv, mjtNum);
jacdifp = jacdif;
jacdifr = jacdif + 3*nv;
jacdifp = mjSTACKALLOC(d, 3*nv, mjtNum);
jacdifr = mjSTACKALLOC(d, 3*nv, mjtNum);
jac1p = mjSTACKALLOC(d, 3*nv, mjtNum);
jac2p = mjSTACKALLOC(d, 3*nv, mjtNum);
jac1r = mjSTACKALLOC(d, 3*nv, mjtNum);
@@ -1649,7 +1648,7 @@ void mj_instantiateContact(const mjModel* m, mjData* d) {
con = d->contact + i;
dim = con->dim;
con->efc_address = d->nefc;
NV = mj_contactJacobian(m, d, con, dim, jac, jacdif, jacdifp, jacdifr,
NV = mj_contactJacobian(m, d, con, dim, jacdifp, jacdifr,
jac1p, jac2p, jac1r, jac2r, chain);
// skip contact if no DOFs affected
src/engine/engine_core_constraint.h
+18 -2
View File
@@ -66,9 +66,25 @@ void mj_instantiateEquality(const mjModel* m, mjData* d);
void mj_instantiateContact(const mjModel* m, mjData* d);
// compute Jacobian for contact, return number of DOFs affected
//
// Arguments:
// m: model structure (mjModel)
// d: data structure (mjData)
// con: contact information
// dim: contact dimension (1 to 6)
// jacdifp: translational Jacobian difference [3 * (NV or nv)]
// jacdifr: rotational Jacobian difference [3 * (NV or nv)], nullable, unused if dim <= 3
// jac1p: translational Jacobian for body 1 [3 * (NV or nv)], nullable
// jac2p: translational Jacobian for body 2 [3 * (NV or nv)], nullable
// jac1r: rotational Jacobian for body 1 [3 * (NV or nv)], nullable, unused if dim <= 3
// jac2r: rotational Jacobian for body 2 [3 * (NV or nv)], nullable, unused if dim <= 3
// chain: list of DOF indices affecting contact [NV], unused if dense
//
// Returns:
// number of DOFs affected (NV for sparse, nv for dense)
MJAPI int mj_contactJacobian(const mjModel* m, mjData* d, const mjContact* con, int dim,
mjtNum* jac, mjtNum* jacdif, mjtNum* jacdifp,
mjtNum* jacdifr, mjtNum* jac1p, mjtNum* jac2p,
mjtNum* jacdifp, mjtNum* jacdifr,
mjtNum* jac1p, mjtNum* jac2p,
mjtNum* jac1r, mjtNum* jac2r, int* chain);
src/engine/engine_core_util.c
+25 -14
View File
@@ -517,29 +517,29 @@ int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
// dense or sparse weighted sum of multiple body Jacobians at same point
int mj_jacSum(const mjModel* m, mjData* d, int* chain,
int n, const int* body, const mjtNum* weight,
const mjtNum point[3], mjtNum* jac, int flg_rot) {
const mjtNum point[3], mjtNum* jacp, mjtNum* jacr, int flg_rot) {
int nv = m->nv, NV;
mjtNum* jacp = jac;
mj_markStack(d);
mjtNum* jtmp = mjSTACKALLOC(d, flg_rot ? 6*nv : 3*nv, mjtNum);
mjtNum* jp = jtmp;
// sparse
if (mj_isSparse(m)) {
// the sparse merge produces one packed [jacp; jacr] block; split into the outputs at the end
mjtNum* jac = mjSTACKALLOC(d, flg_rot ? 6*nv : 3*nv, mjtNum);
mjtNum* buf = mjSTACKALLOC(d, flg_rot ? 6*nv : 3*nv, mjtNum);
int* buf_ind = mjSTACKALLOC(d, nv, int);
int* bodychain = mjSTACKALLOC(d, nv, int);
// set first
// set first (rotational rows packed right after the translational rows, at offset 3*NV)
NV = mj_bodyChain(m, body[0], chain);
if (NV) {
// get Jacobian
mjtNum* jacr = flg_rot ? jac + 3*NV : NULL;
mjtNum* jr = flg_rot ? jac + 3*NV : NULL;
if (m->body_simple[body[0]]) {
mj_jacSparseSimple(m, d, jacp, jacr, point, body[0], 1, NV, 0);
mj_jacSparseSimple(m, d, jac, jr, point, body[0], 1, NV, 0);
} else {
mj_jacSparse(m, d, jacp, jacr, point, body[0], NV, chain, /*flg_skipcommon=*/0);
mj_jacSparse(m, d, jac, jr, point, body[0], NV, chain, /*flg_skipcommon=*/0);
}
// apply weight
@@ -555,30 +555,41 @@ int mj_jacSum(const mjModel* m, mjData* d, int* chain,
}
mjtNum* jr = flg_rot ? jtmp + 3*bodyNV : NULL;
if (m->body_simple[body[i]]) {
mj_jacSparseSimple(m, d, jp, jr, point, body[i], 1, bodyNV, 0);
mj_jacSparseSimple(m, d, jtmp, jr, point, body[i], 1, bodyNV, 0);
} else {
mj_jacSparse(m, d, jp, jr, point, body[i], bodyNV, bodychain, /*flg_skipcommon=*/0);
mj_jacSparse(m, d, jtmp, jr, point, body[i], bodyNV, bodychain, /*flg_skipcommon=*/0);
}
// combine sparse matrices
NV = mju_addToSparseMat(jac, jtmp, nv, flg_rot ? 6 : 3, weight[i],
NV, bodyNV, chain, bodychain, buf, buf_ind);
}
// split the packed block into the separate output buffers (each NV-packed)
mju_copy(jacp, jac, 3*NV);
if (flg_rot) {
mju_copy(jacr, jac + 3*NV, 3*NV);
}
}
// dense
else {
mjtNum* jacr = flg_rot ? jac + 3*nv : NULL;
mjtNum* jr = flg_rot ? jtmp + 3*nv : NULL;
// set first
mj_jac(m, d, jacp, jacr, point, body[0]);
mju_scl(jac, jac, weight[0], flg_rot ? 6*nv : 3*nv);
mj_jac(m, d, jacp, flg_rot ? jacr : NULL, point, body[0]);
mju_scl(jacp, jacp, weight[0], 3*nv);
if (flg_rot) {
mju_scl(jacr, jacr, weight[0], 3*nv);
}
// accumulate remaining
for (int i=1; i < n; i++) {
mj_jac(m, d, jp, jr, point, body[i]);
mju_addToScl(jac, jtmp, weight[i], flg_rot ? 6*nv : 3*nv);
mj_jac(m, d, jtmp, jr, point, body[i]);
mju_addToScl(jacp, jtmp, weight[i], 3*nv);
if (flg_rot) {
mju_addToScl(jacr, jr, weight[i], 3*nv);
}
}
NV = nv;
src/engine/engine_core_util.h
+1 -1
View File
@@ -101,7 +101,7 @@ MJAPI int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
// dense or sparse weighted sum of multiple body Jacobians at same point
int mj_jacSum(const mjModel* m, mjData* d, int* chain,
int n, const int* body, const mjtNum* weight,
const mjtNum point[3], mjtNum* jac, int flg_rot);
const mjtNum point[3], mjtNum* jacp, mjtNum* jacr, int flg_rot);
// compute 3/6-by-nv Jacobian time derivative of global point attached to given body
MJAPI void mj_jacDot(const mjModel* m, const mjData* d,
src/engine/engine_passive.c
+5 -12
View File
@@ -884,7 +884,7 @@ static int mj_fluid(const mjModel* m, mjData* d) {
int mj_contactPassive(const mjModel* m, mjData* d) {
int ncon = d->ncon, issparse = mj_isSparse(m);
int dim, NV, nv = m->nv, *chain = NULL;
mjtNum *jac, *jacdif, *jacdifp, *jacdifr, *jac1p, *jac2p, *jac1r, *jac2r, *qfrc;
mjtNum *jac, *jacdifp, *jac1p, *jac2p, *qfrc;
mjContact* con;
int has_contact = 0;
@@ -907,13 +907,9 @@ int mj_contactPassive(const mjModel* m, mjData* d) {
// allocate Jacobian
mj_markStack(d);
jac = mjSTACKALLOC(d, 6*nv, mjtNum);
jacdif = mjSTACKALLOC(d, 6*nv, mjtNum);
jacdifp = jacdif;
jacdifr = jacdif + 3*nv;
jacdifp = mjSTACKALLOC(d, 3*nv, mjtNum);
jac1p = mjSTACKALLOC(d, 3*nv, mjtNum);
jac2p = mjSTACKALLOC(d, 3*nv, mjtNum);
jac1r = mjSTACKALLOC(d, 3*nv, mjtNum);
jac2r = mjSTACKALLOC(d, 3*nv, mjtNum);
qfrc = mjSTACKALLOC(d, nv, mjtNum);
if (issparse) {
chain = mjSTACKALLOC(d, nv, int);
@@ -929,8 +925,8 @@ int mj_contactPassive(const mjModel* m, mjData* d) {
con = d->contact + i;
dim = con->dim;
con->efc_address = -1;
NV = mj_contactJacobian(m, d, con, dim, jac, jacdif, jacdifp, jacdifr,
jac1p, jac2p, jac1r, jac2r, chain);
NV = mj_contactJacobian(m, d, con, dim, jacdifp, NULL,
jac1p, jac2p, NULL, NULL, chain);
// skip contact if no DOFs affected
if (NV == 0) {
@@ -941,9 +937,6 @@ int mj_contactPassive(const mjModel* m, mjData* d) {
// rotate Jacobian differences to contact frame
mju_mulMatMat(jac, con->frame, jacdifp, dim > 1 ? 3 : 1, 3, NV);
if (dim > 3) {
mju_mulMatMat(jac + 3*NV, con->frame, jacdifr, dim-3, 3, NV);
}
// compute passive contact force (dim = 1)
mjtNum scl = -kContactStiffness*con->dist;
@@ -1004,7 +997,7 @@ int mj_adhesion(const mjModel* m, mjData* d) {
}
// normal Jacobian
NV = mj_contactJacobian(m, d, con, 1, jac, jacdif, jacdif, NULL,
NV = mj_contactJacobian(m, d, con, 1, jacdif, NULL,
jac1p, jac2p, NULL, NULL, chain);
if (NV == 0) {
continue;
test/engine/engine_core_constraint_test.cc
+126 -4
View File
@@ -17,8 +17,10 @@
#include "src/engine/engine_core_constraint.h"
#include <array>
#include <cstdio>
#include <cstring>
#include <string>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
@@ -1301,10 +1303,8 @@ TEST_F(CoreConstraintTest, ShellModeContactJacobian) {
// buffer for Jacobian
std::vector<mjtNum> jacdif(3 * model->nv, 0.0);
// call mj_contactJacobian
mj_contactJacobian(model.get(), data.get(), &con, 1, nullptr, jacdif.data(),
nullptr, nullptr, nullptr, nullptr, nullptr, nullptr,
nullptr);
mj_contactJacobian(model.get(), data.get(), &con, 1, jacdif.data(),
nullptr, nullptr, nullptr, nullptr, nullptr, nullptr);
// check that boundary nodes have non-zero entries, and central node has zero
@@ -1524,5 +1524,127 @@ TEST_F(CoreConstraintTest, AdhesionPriority) {
EXPECT_EQ(data->contact[0].adhesion, 10);
}
// Verify that the rotational Jacobian from mj_jacSum in the sparse path matches
// the dense path. This is a regression test for the bug where the sparse path
// wrote rotational rows at offset 3*NV inside a packed buffer, but the caller
// expected them at offset 3*nv.
TEST_F(CoreConstraintTest, SparseRotationalJacobianMatchesDense) {
// XML template with a %s placeholder for the jacobian type
constexpr char xml_template[] = R"(
<mujoco>
<option jacobian="%s"/>
<worldbody>
<body name="parent" pos="0 0 1">
<freejoint/>
<geom size=".01" mass="1"/>
<flexcomp name="flex" type="grid" count="3 3 3" spacing=".1 .1 .1" dim="3" dof="trilinear">
<elasticity elastic2d="stretch" thickness="0.01"/>
<contact selfcollide="none"/>
</flexcomp>
</body>
<geom name="plane" type="plane" size="1 1 1" pos="0 0 -1"/>
</worldbody>
</mujoco>
)";
auto run_contact_jacobian = [&](const char* jacobian_type,
std::vector<mjtNum>& jacdifp_dense_out,
std::vector<mjtNum>& jacdifr_dense_out) {
char xml[1024];
snprintf(xml, sizeof(xml), xml_template, jacobian_type);
char error[1024];
MjModelPtr model = LoadModelFromString(xml, error, sizeof(error));
ASSERT_THAT(model.get(), testing::NotNull()) << error;
MjDataPtr data = MakeData(model);
mj_forward(model.get(), data.get());
int nv = model->nv;
// pick a boundary vertex so the Jacobian is non-trivial
int vert_idx = 0;
int plane_geom_id = mj_name2id(model.get(), mjOBJ_GEOM, "plane");
ASSERT_GE(plane_geom_id, 0);
// create a frictional contact with dim=6 (normal + 2 tangent + 3 torsional)
mjContact con;
memset(&con, 0, sizeof(mjContact));
con.flex[0] = -1;
con.flex[1] = -1;
con.vert[0] = -1;
con.vert[1] = -1;
con.geom[0] = plane_geom_id; // plane geom (world body)
con.geom[1] = -1; // trigger flex branch in mj_contactJacobian
con.flex[1] = 0;
con.vert[1] = vert_idx;
con.dim = 6; // full frictional contact: exercises rotational Jacobian
mju_copy3(con.pos, data->flexvert_xpos + 3 * vert_idx);
// set contact frame to identity
mju_zero(con.frame, 9);
con.frame[0] = 1;
con.frame[4] = 1;
con.frame[8] = 1;
// allocate output buffers
std::vector<mjtNum> jacdifp(3 * nv, 0.0);
std::vector<mjtNum> jacdifr(3 * nv, 0.0);
std::vector<int> chain(nv, 0);
int NV = mj_contactJacobian(model.get(), data.get(), &con, con.dim,
jacdifp.data(), jacdifr.data(),
nullptr, nullptr, nullptr, nullptr,
chain.data());
ASSERT_GT(NV, 0);
// for sparse: unpack to dense using chain indices
if (strcmp(jacobian_type, "sparse") == 0) {
std::vector<mjtNum> jacdifp_full(3 * nv, 0.0);
std::vector<mjtNum> jacdifr_full(3 * nv, 0.0);
for (int row = 0; row < 3; row++) {
for (int j = 0; j < NV; j++) {
jacdifp_full[row * nv + chain[j]] = jacdifp[row * NV + j];
jacdifr_full[row * nv + chain[j]] = jacdifr[row * NV + j];
}
}
jacdifp_dense_out = std::move(jacdifp_full);
jacdifr_dense_out = std::move(jacdifr_full);
} else {
jacdifp_dense_out = std::move(jacdifp);
jacdifr_dense_out = std::move(jacdifr);
}
};
// compute Jacobians with dense and sparse paths
std::vector<mjtNum> dense_jacdifp, dense_jacdifr;
std::vector<mjtNum> sparse_jacdifp, sparse_jacdifr;
run_contact_jacobian("dense", dense_jacdifp, dense_jacdifr);
ASSERT_FALSE(HasFatalFailure());
run_contact_jacobian("sparse", sparse_jacdifp, sparse_jacdifr);
ASSERT_FALSE(HasFatalFailure());
ASSERT_EQ(dense_jacdifp.size(), sparse_jacdifp.size());
ASSERT_EQ(dense_jacdifr.size(), sparse_jacdifr.size());
// verify positional Jacobian matches
const mjtNum tol = MjTol(1e-10, 1e-5);
for (int i = 0; i < dense_jacdifp.size(); i++) {
EXPECT_NEAR(dense_jacdifp[i], sparse_jacdifp[i], tol)
<< "jacdifp mismatch at index " << i;
}
// verify rotational Jacobian matches (this is where the bug was)
bool has_nonzero_rot = false;
for (int i = 0; i < dense_jacdifr.size(); i++) {
EXPECT_NEAR(dense_jacdifr[i], sparse_jacdifr[i], tol)
<< "jacdifr mismatch at index " << i;
if (mju_abs(dense_jacdifr[i]) > tol) {
has_nonzero_rot = true;
}
}
EXPECT_TRUE(has_nonzero_rot)
<< "Rotational Jacobian should have non-zero entries";
}
} // namespace
} // namespace mujoco