From 14c0b0c92b58359675b53d0f30566a5442907c1f Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Thu, 2 Jul 2026 10:48:21 -0700 Subject: [PATCH] Recompute sameframe flags in `mj_setConst`. The `body_sameframe`, `geom_sameframe`, and `site_sameframe` flags are now recomputed within `mj_setConst` based on the current model geometry. This allows these flags to be updated if the model's body/geom/site frames are modified after the initial compilation. Tests are added to verify the correct recomputation and its effect on forward kinematics. The compiler also checks that its initial computation matches the result of `mj_setConst`. Fixes #3029, #1628. PiperOrigin-RevId: 941758649 Change-Id: I7b995f8c67ad305bca5b52732436ba578c905d73 --- doc/changelog.rst | 2 + src/engine/engine_setconst.c | 109 +++++++++++++++++++++ test/engine/engine_setconst_test.cc | 147 ++++++++++++++++++++++++++++ 3 files changed, 258 insertions(+) diff --git a/doc/changelog.rst b/doc/changelog.rst index c8ea1f9d..8cf5803a 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -10,6 +10,8 @@ General - Added Nesterov momentum extrapolation with adaptive gradient restart (O'Donoghue-Candès) to the PGS solver, significantly improving convergence. Overall PGS now requires ~2x fewer iterations. - :ref:`mj_encode` now supports encoding of MJB and TXT files. +- :ref:`mj_setConst` now recomputes the ``mjModel.{body,geom,site}_sameframe`` flags, to account for changes in + body/geom/site frames after compilation. - The :el:`attach` element now supports self-attachment (attaching elements of the current model to itself) by omitting the :at:`model` attribute. It also supports attaching a frame via the new :at:`frame` attribute, which is mutually exclusive with :at:`body`. diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c index 0098feaa..47afec01 100644 --- a/src/engine/engine_setconst.c +++ b/src/engine/engine_setconst.c @@ -82,6 +82,112 @@ static int GetWrapBodyTreeId(const mjModel* m, int wrap_index) { return (bodyid != -1) ? m->body_treeid[bodyid] : -1; } + +// tolerance for frame comparison, must match compiler's kFrameEps +static const mjtNum kSameFrameEps = 1e-6; + + +// return 1 if all 3 elements are near zero +static int isNullVec3(const mjtNum v[3]) { + return mju_abs(v[0]) < kSameFrameEps && + mju_abs(v[1]) < kSameFrameEps && + mju_abs(v[2]) < kSameFrameEps; +} + + +// return 1 if quaternion is near identity (1,0,0,0), accounting for double-cover +static int isNullQuat(const mjtNum q[4]) { + int plus = mju_abs(q[0] - 1) < kSameFrameEps && + mju_abs(q[1]) < kSameFrameEps && + mju_abs(q[2]) < kSameFrameEps && + mju_abs(q[3]) < kSameFrameEps; + int minus = mju_abs(q[0] + 1) < kSameFrameEps && + mju_abs(q[1]) < kSameFrameEps && + mju_abs(q[2]) < kSameFrameEps && + mju_abs(q[3]) < kSameFrameEps; + return plus || minus; +} + + +// return 1 if two quaternions are near equal, accounting for double-cover +static int isSameQuat(const mjtNum q1[4], const mjtNum q2[4]) { + int plus = mju_abs(q1[0] - q2[0]) < kSameFrameEps && + mju_abs(q1[1] - q2[1]) < kSameFrameEps && + mju_abs(q1[2] - q2[2]) < kSameFrameEps && + mju_abs(q1[3] - q2[3]) < kSameFrameEps; + int minus = mju_abs(q1[0] + q2[0]) < kSameFrameEps && + mju_abs(q1[1] + q2[1]) < kSameFrameEps && + mju_abs(q1[2] + q2[2]) < kSameFrameEps && + mju_abs(q1[3] + q2[3]) < kSameFrameEps; + return plus || minus; +} + + +// return 1 if two 3-vectors are near equal +static int isSameVec3(const mjtNum v1[3], const mjtNum v2[3]) { + return mju_abs(v1[0] - v2[0]) < kSameFrameEps && + mju_abs(v1[1] - v2[1]) < kSameFrameEps && + mju_abs(v1[2] - v2[2]) < kSameFrameEps; +} + + +// recompute body_sameframe, geom_sameframe, site_sameframe from model geometry +static void setSameframe(mjModel* m) { + // body_sameframe: compare body inertial frame to body frame + for (int i=1; i < m->nbody; i++) { + mjtNum* ipos = m->body_ipos+3*i; + mjtNum* iquat = m->body_iquat+4*i; + if (isNullVec3(ipos) && isNullQuat(iquat)) { + m->body_sameframe[i] = mjSAMEFRAME_BODY; + } else if (isNullQuat(iquat)) { + m->body_sameframe[i] = mjSAMEFRAME_BODYROT; + } else { + m->body_sameframe[i] = mjSAMEFRAME_NONE; + } + } + + // geom_sameframe: compare geom frame to body and inertial frames + for (int i=0; i < m->ngeom; i++) { + int b = m->geom_bodyid[i]; + mjtNum* gpos = m->geom_pos+3*i; + mjtNum* gquat = m->geom_quat+4*i; + mjtNum* ipos = m->body_ipos+3*b; + mjtNum* iquat = m->body_iquat+4*b; + if (isNullVec3(gpos) && isNullQuat(gquat)) { + m->geom_sameframe[i] = mjSAMEFRAME_BODY; + } else if (isNullQuat(gquat)) { + m->geom_sameframe[i] = mjSAMEFRAME_BODYROT; + } else if (isSameVec3(gpos, ipos) && isSameQuat(gquat, iquat)) { + m->geom_sameframe[i] = mjSAMEFRAME_INERTIA; + } else if (isSameQuat(gquat, iquat)) { + m->geom_sameframe[i] = mjSAMEFRAME_INERTIAROT; + } else { + m->geom_sameframe[i] = mjSAMEFRAME_NONE; + } + } + + // site_sameframe: compare site frame to body and inertial frames + for (int i=0; i < m->nsite; i++) { + int b = m->site_bodyid[i]; + mjtNum* spos = m->site_pos+3*i; + mjtNum* squat = m->site_quat+4*i; + mjtNum* ipos = m->body_ipos+3*b; + mjtNum* iquat = m->body_iquat+4*b; + if (isNullVec3(spos) && isNullQuat(squat)) { + m->site_sameframe[i] = mjSAMEFRAME_BODY; + } else if (isNullQuat(squat)) { + m->site_sameframe[i] = mjSAMEFRAME_BODYROT; + } else if (isSameVec3(spos, ipos) && isSameQuat(squat, iquat)) { + m->site_sameframe[i] = mjSAMEFRAME_INERTIA; + } else if (isSameQuat(squat, iquat)) { + m->site_sameframe[i] = mjSAMEFRAME_INERTIAROT; + } else { + m->site_sameframe[i] = mjSAMEFRAME_NONE; + } + } +} + + // set fixed quantities (do not depend on qpos0) static void setFixed(mjModel* m, mjData* d) { mj_markStack(d); @@ -1215,6 +1321,9 @@ static void setSpring(mjModel* m, mjData* d) { // entry point: set all remaining constant fields of mjModel, except for lengthrange void mj_setConst(mjModel* m, mjData* d) { + // recompute sameframe flags from current model geometry + setSameframe(m); + // set fixed quantities setFixed(m, d); diff --git a/test/engine/engine_setconst_test.cc b/test/engine/engine_setconst_test.cc index b53174b6..43c3eff3 100644 --- a/test/engine/engine_setconst_test.cc +++ b/test/engine/engine_setconst_test.cc @@ -457,5 +457,152 @@ TEST_F(SetConstTest, DofLength) { EXPECT_NEAR(model->dof_length[10], 5, tol); } +TEST_F(SetConstTest, BodySameframeRecomputed) { + constexpr char xml[] = R"( + + + + + + + + + )"; + char error[1024]; + MjModelPtr m = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(m.get(), NotNull()) << error; + MjDataPtr d(mj_makeData(m.get())); + + int b = mj_name2id(m.get(), mjOBJ_BODY, "B1"); + + // initially sameframe should be BODY (ipos=0, iquat=identity) + EXPECT_EQ(m->body_sameframe[b], mjSAMEFRAME_BODY); + + // perturb body_ipos, call mj_setConst + m->body_ipos[3*b+0] = 1.0; + mj_setConst(m.get(), d.get()); + EXPECT_EQ(m->body_sameframe[b], mjSAMEFRAME_BODYROT); + + // also perturb body_iquat + m->body_iquat[4*b+0] = 0.5; + m->body_iquat[4*b+1] = 0.5; + m->body_iquat[4*b+2] = 0.5; + m->body_iquat[4*b+3] = 0.5; + mj_setConst(m.get(), d.get()); + EXPECT_EQ(m->body_sameframe[b], mjSAMEFRAME_NONE); + + // restore to identity, should go back to BODY + m->body_ipos[3*b+0] = 0; + m->body_iquat[4*b+0] = 1; + m->body_iquat[4*b+1] = 0; + m->body_iquat[4*b+2] = 0; + m->body_iquat[4*b+3] = 0; + mj_setConst(m.get(), d.get()); + EXPECT_EQ(m->body_sameframe[b], mjSAMEFRAME_BODY); +} + +TEST_F(SetConstTest, GeomSameframeRecomputed) { + constexpr char xml[] = R"( + + + + + + + + + )"; + char error[1024]; + MjModelPtr m = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(m.get(), NotNull()) << error; + MjDataPtr d(mj_makeData(m.get())); + + int g = mj_name2id(m.get(), mjOBJ_GEOM, "G1"); + + // initially sameframe should be BODY + EXPECT_EQ(m->geom_sameframe[g], mjSAMEFRAME_BODY); + + // perturb geom_pos + m->geom_pos[3*g+1] = 0.5; + mj_setConst(m.get(), d.get()); + EXPECT_EQ(m->geom_sameframe[g], mjSAMEFRAME_BODYROT); + + // restore, should go back to BODY + m->geom_pos[3*g+1] = 0; + mj_setConst(m.get(), d.get()); + EXPECT_EQ(m->geom_sameframe[g], mjSAMEFRAME_BODY); +} + +TEST_F(SetConstTest, SiteSameframeRecomputed) { + constexpr char xml[] = R"( + + + + + + + + + + )"; + char error[1024]; + MjModelPtr m = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(m.get(), NotNull()) << error; + MjDataPtr d(mj_makeData(m.get())); + + int s = mj_name2id(m.get(), mjOBJ_SITE, "S1"); + + // initially sameframe should be BODY + EXPECT_EQ(m->site_sameframe[s], mjSAMEFRAME_BODY); + + // perturb site_pos + m->site_pos[3*s+2] = 0.3; + mj_setConst(m.get(), d.get()); + EXPECT_EQ(m->site_sameframe[s], mjSAMEFRAME_BODYROT); + + // restore + m->site_pos[3*s+2] = 0; + mj_setConst(m.get(), d.get()); + EXPECT_EQ(m->site_sameframe[s], mjSAMEFRAME_BODY); +} + +TEST_F(SetConstTest, SameframeKinematicsCorrect) { + constexpr char xml[] = R"( + + + + + + + + + + )"; + char error[1024]; + MjModelPtr m = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(m.get(), NotNull()) << error; + MjDataPtr d(mj_makeData(m.get())); + + int b = mj_name2id(m.get(), mjOBJ_BODY, "B1"); + int g = mj_name2id(m.get(), mjOBJ_GEOM, "G1"); + + // perturb body inertial offset, breaking sameframe + m->body_ipos[3*b+1] = 0.5; + mj_setConst(m.get(), d.get()); + EXPECT_EQ(m->body_sameframe[b], mjSAMEFRAME_BODYROT); + + // run forward kinematics, check that xipos != xpos + mj_forward(m.get(), d.get()); + EXPECT_NEAR(d->xipos[3*b+1], 0.5, MjTol(1e-10, 1e-6)); + EXPECT_NEAR(d->xpos[3*b+1], 0.0, MjTol(1e-10, 1e-6)); + + // perturb geom_pos, check geom global position + m->geom_pos[3*g+2] = 0.3; + mj_setConst(m.get(), d.get()); + EXPECT_NE(m->geom_sameframe[g], mjSAMEFRAME_BODY); + mj_forward(m.get(), d.get()); + EXPECT_NEAR(d->geom_xpos[3*g+2], 0.3, MjTol(1e-10, 1e-6)); +} + } // namespace } // namespace mujoco