From 426cb5481df9dbeaf0a6c01bb15369e77869f84f Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Tue, 21 Jul 2026 02:01:24 -0700 Subject: [PATCH] Fix sparse-path rotational Jacobian misalignment in mj_jacSum PiperOrigin-RevId: 951341459 Change-Id: I18bc27765b3147a5eb520e3026317e9a5a2dfc30 --- src/engine/engine_core_constraint.c | 15 ++- src/engine/engine_core_constraint.h | 20 +++- src/engine/engine_core_util.c | 39 ++++--- src/engine/engine_core_util.h | 2 +- src/engine/engine_passive.c | 17 +-- test/engine/engine_core_constraint_test.cc | 130 ++++++++++++++++++++- 6 files changed, 182 insertions(+), 41 deletions(-) diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index f0e002e9..24ec73a5 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -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 diff --git a/src/engine/engine_core_constraint.h b/src/engine/engine_core_constraint.h index da585555..cfe506b5 100644 --- a/src/engine/engine_core_constraint.h +++ b/src/engine/engine_core_constraint.h @@ -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); diff --git a/src/engine/engine_core_util.c b/src/engine/engine_core_util.c index dc1031b8..009a0a33 100644 --- a/src/engine/engine_core_util.c +++ b/src/engine/engine_core_util.c @@ -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; diff --git a/src/engine/engine_core_util.h b/src/engine/engine_core_util.h index 0f02d4a0..b31c15a6 100644 --- a/src/engine/engine_core_util.h +++ b/src/engine/engine_core_util.h @@ -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, diff --git a/src/engine/engine_passive.c b/src/engine/engine_passive.c index bd6680bf..ad4916fb 100644 --- a/src/engine/engine_passive.c +++ b/src/engine/engine_passive.c @@ -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; diff --git a/test/engine/engine_core_constraint_test.cc b/test/engine/engine_core_constraint_test.cc index f463ac54..1f169625 100644 --- a/test/engine/engine_core_constraint_test.cc +++ b/test/engine/engine_core_constraint_test.cc @@ -17,8 +17,10 @@ #include "src/engine/engine_core_constraint.h" #include +#include #include #include +#include #include #include @@ -1301,10 +1303,8 @@ TEST_F(CoreConstraintTest, ShellModeContactJacobian) { // buffer for Jacobian std::vector 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"( + + + )"; + + auto run_contact_jacobian = [&](const char* jacobian_type, + std::vector& jacdifp_dense_out, + std::vector& 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 jacdifp(3 * nv, 0.0); + std::vector jacdifr(3 * nv, 0.0); + std::vector 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 jacdifp_full(3 * nv, 0.0); + std::vector 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 dense_jacdifp, dense_jacdifr; + std::vector 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