From 2efbd317628bb45ec6cc12104224a9e51feffd10 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Thu, 22 Aug 2024 06:07:30 -0700 Subject: [PATCH] Add `mjSAMEFRAME_BODYROT` and `mjSAMEFRAME_INERTIAROT` frame alignment shortcuts. PiperOrigin-RevId: 666318799 Change-Id: I0af1824d2bd9eaca7b48a9e33e09f873dd33b22d --- doc/includes/references.h | 2 + include/mujoco/mjmodel.h | 2 + introspect/enums.py | 2 + src/engine/engine_support.c | 4 ++ src/user/user_model.cc | 15 +++++++- test/user/user_model_test.cc | 57 ++++++++++++++++++++++++++++ unity/Runtime/Bindings/MjBindings.cs | 2 + 7 files changed, 83 insertions(+), 1 deletion(-) diff --git a/doc/includes/references.h b/doc/includes/references.h index 0ca6cacd..87010fe4 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -688,6 +688,8 @@ typedef enum mjtSameFrame_ { // frame alignment of bodies with their childr mjSAMEFRAME_NONE = 0, // no alignment mjSAMEFRAME_BODY, // frame is same as body frame mjSAMEFRAME_INERTIA, // frame is same as inertial frame + mjSAMEFRAME_BODYROT, // frame orientation is same as body orientation + mjSAMEFRAME_INERTIAROT // frame orientation is same as inertia orientation } mjtSameFrame; typedef enum mjtLRMode_ { // mode for actuator length range computation mjLRMODE_NONE = 0, // do not process any actuators diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index 4fd691c9..16e18ea8 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -377,6 +377,8 @@ typedef enum mjtSameFrame_ { // frame alignment of bodies with their childr mjSAMEFRAME_NONE = 0, // no alignment mjSAMEFRAME_BODY, // frame is same as body frame mjSAMEFRAME_INERTIA, // frame is same as inertial frame + mjSAMEFRAME_BODYROT, // frame orientation is same as body orientation + mjSAMEFRAME_INERTIAROT // frame orientation is same as inertia orientation } mjtSameFrame; diff --git a/introspect/enums.py b/introspect/enums.py index 3c36b5cf..9e24b3b9 100644 --- a/introspect/enums.py +++ b/introspect/enums.py @@ -396,6 +396,8 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjSAMEFRAME_NONE', 0), ('mjSAMEFRAME_BODY', 1), ('mjSAMEFRAME_INERTIA', 2), + ('mjSAMEFRAME_BODYROT', 3), + ('mjSAMEFRAME_INERTIAROT', 4), ]), )), ('mjtLRMode', diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index bd6699d5..35510f8d 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -1620,6 +1620,8 @@ void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9], if (xpos && pos) { switch (sf) { case mjSAMEFRAME_NONE: + case mjSAMEFRAME_BODYROT: + case mjSAMEFRAME_INERTIAROT: mju_mulMatVec3(xpos, d->xmat+9*body, pos); mju_addTo3(xpos, d->xpos+3*body); break; @@ -1641,9 +1643,11 @@ void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9], mju_quat2Mat(xmat, tmp); break; case mjSAMEFRAME_BODY: + case mjSAMEFRAME_BODYROT: mju_copy(xmat, d->xmat+9*body, 9); break; case mjSAMEFRAME_INERTIA: + case mjSAMEFRAME_INERTIAROT: mju_copy(xmat, d->ximat+9*body, 9); break; } diff --git a/src/user/user_model.cc b/src/user/user_model.cc index 140156ba..646ecfe7 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -2013,8 +2013,11 @@ void mjCModel::CopyTree(mjModel* m) { // set sameframe mjtSameFrame sameframe; + mjtNum* nullnum = static_cast(nullptr); if (IsNullPose(m->body_ipos+3*i, m->body_iquat+4*i)) { sameframe = mjSAMEFRAME_BODY; + } else if (IsNullPose(nullnum, m->body_iquat+4*i)) { + sameframe = mjSAMEFRAME_BODYROT; } else { sameframe = mjSAMEFRAME_NONE; } @@ -2064,7 +2067,7 @@ void mjCModel::CopyTree(mjModel* m) { bool axis_aligned = ((std::abs(pj->axis[0]) > mjEPS) + (std::abs(pj->axis[1]) > mjEPS) + (std::abs(pj->axis[2]) > mjEPS)) == 1; - if (rotfound || !IsNullPose(m->jnt_pos+3*jid, static_cast(nullptr)) || + if (rotfound || !IsNullPose(m->jnt_pos+3*jid, nullnum) || ((pj->type == mjJNT_HINGE || pj->type == mjJNT_SLIDE) && !axis_aligned)) { m->body_simple[i] = 0; } @@ -2172,10 +2175,15 @@ void mjCModel::CopyTree(mjModel* m) { mjuu_copyvec(m->geom_rgba+4*gid, pg->rgba, 4); // determine sameframe + double* nulldouble = static_cast(nullptr); if (IsNullPose(m->geom_pos+3*gid, m->geom_quat+4*gid)) { sameframe = mjSAMEFRAME_BODY; + } else if (IsNullPose(nullnum, m->geom_quat+4*gid)) { + sameframe = mjSAMEFRAME_BODYROT; } else if (IsSamePose(pg->pos, pb->ipos, pg->quat, pb->iquat)) { sameframe = mjSAMEFRAME_INERTIA; + } else if (IsSamePose(nulldouble, nulldouble, pg->quat, pb->iquat)) { + sameframe = mjSAMEFRAME_INERTIAROT; } else { sameframe = mjSAMEFRAME_NONE; } @@ -2203,10 +2211,15 @@ void mjCModel::CopyTree(mjModel* m) { mjuu_copyvec(m->site_rgba+4*sid, ps->rgba, 4); // determine sameframe + double* nulldouble = static_cast(nullptr); if (IsNullPose(m->site_pos+3*sid, m->site_quat+4*sid)) { sameframe = mjSAMEFRAME_BODY; + } else if (IsNullPose(nullnum, m->site_quat+4*sid)) { + sameframe = mjSAMEFRAME_BODYROT; } else if (IsSamePose(ps->pos, pb->ipos, ps->quat, pb->iquat)) { sameframe = mjSAMEFRAME_INERTIA; + } else if (IsSamePose(nulldouble, nulldouble, ps->quat, pb->iquat)) { + sameframe = mjSAMEFRAME_INERTIAROT; } else { sameframe = mjSAMEFRAME_NONE; } diff --git a/test/user/user_model_test.cc b/test/user/user_model_test.cc index 8b4f2cdb..6e6f1bbc 100644 --- a/test/user/user_model_test.cc +++ b/test/user/user_model_test.cc @@ -36,12 +36,17 @@ using ::testing::ElementsAre; using ::testing::HasSubstr; using ::testing::IsNull; using ::testing::NotNull; +using ::testing::Pointwise; static std::vector GetRow(const mjtNum* array, int ncolumn, int row) { return std::vector(array + ncolumn * row, array + ncolumn * (row + 1)); } +std::vector AsVector(const mjtNum* array, int n) { + return std::vector(array, array + n); +} + // ----------------------------- test mjCModel -------------------------------- using UserCModelTest = MujocoTest; @@ -64,6 +69,58 @@ TEST_F(UserCModelTest, RepeatedNames) { EXPECT_THAT(error.data(), HasSubstr("repeated name 'geom1' in geom")); } +TEST_F(UserCModelTest, SameFrame) { + static constexpr char xml[] = R"( + + + + + + + + + + + + + + + )"; + + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + ASSERT_THAT(model, NotNull()) << error.data(); + EXPECT_EQ(model->geom_sameframe[0], mjSAMEFRAME_NONE); + EXPECT_EQ(model->geom_sameframe[1], mjSAMEFRAME_BODY); + EXPECT_EQ(model->geom_sameframe[2], mjSAMEFRAME_INERTIA); + EXPECT_EQ(model->geom_sameframe[3], mjSAMEFRAME_BODYROT); + EXPECT_EQ(model->geom_sameframe[4], mjSAMEFRAME_INERTIAROT); + + // make data, get geom_xpos + mjData* data = mj_makeData(model); + mj_kinematics(model, data); + auto geom_xpos = AsVector(data->geom_xpos, model->ngeom*3); + auto geom_xmat = AsVector(data->geom_xmat, model->ngeom*9); + + // set all geom_sameframe to 0, call kinematics again + for (int i = 0; i < model->ngeom; i++) { + model->geom_sameframe[i] = mjSAMEFRAME_NONE; + } + mj_resetData(model, data); + mj_kinematics(model, data); + auto geom_xpos2 = AsVector(data->geom_xpos, model->ngeom*3); + auto geom_xmat2 = AsVector(data->geom_xmat, model->ngeom*9); + + // expect them to be equal + constexpr double eps = 1e-6; + EXPECT_THAT(geom_xpos, Pointwise(DoubleNear(eps), geom_xpos2)); + EXPECT_THAT(geom_xmat, Pointwise(DoubleNear(eps), geom_xmat2)); + + mj_deleteData(data); + mj_deleteModel(model); +} + + // ------------- test automatic inference of nuser_xxx ------------------------- using UserDataTest = MujocoTest; diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 0d62e552..08ebfe4e 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -396,6 +396,8 @@ public enum mjtSameFrame : int{ mjSAMEFRAME_NONE = 0, mjSAMEFRAME_BODY = 1, mjSAMEFRAME_INERTIA = 2, + mjSAMEFRAME_BODYROT = 3, + mjSAMEFRAME_INERTIAROT = 4, } public enum mjtLRMode : int{ mjLRMODE_NONE = 0,