Add mjSAMEFRAME_BODYROT and mjSAMEFRAME_INERTIAROT frame alignment shortcuts.

PiperOrigin-RevId: 666318799
Change-Id: I0af1824d2bd9eaca7b48a9e33e09f873dd33b22d

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

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;
 
 

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',

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;
     }

src/user/user_model.cc

@@ -2013,8 +2013,11 @@ void mjCModel::CopyTree(mjModel* m) {
 
     // set sameframe
     mjtSameFrame sameframe;
+    mjtNum* nullnum = static_cast<mjtNum*>(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<mjtNum*>(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<double*>(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<double*>(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;
       }

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<mjtNum> GetRow(const mjtNum* array, int ncolumn, int row) {
   return std::vector<mjtNum>(array + ncolumn * row,
                              array + ncolumn * (row + 1));
 }
 
+std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
+  return std::vector<mjtNum>(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"(
+   <mujoco>
+     <default>
+      <geom type="box" size="1 2 3"/>
+     </default>
+
+     <worldbody>
+       <body name="body1">
+         <geom name="none"       mass="0" pos="1 1 1" euler="10 10 10"/>
+         <geom name="body"       mass="0"/>
+         <geom name="inertia"    mass="1" pos="3 2 1" euler="20 30 40"/>
+         <geom name="bodyrot"    mass="0" pos="1 1 1"/>
+         <geom name="inertiarot" mass="0" euler="20 30 40"/>
+        </body>
+      </worldbody>
+    </mujoco>)";
+
+  std::array<char, 1024> 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;

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,