Preserve frames when saving and recompiling.

Frames were half transient on save: the writer emitted <frame> with only its
name and childclass, while the frame's transformation had already been
accumulated into the children. A saved model reloaded with the same frame names
but identity transforms and moved children.

The writer now emits every frame with its pos and quat, relative to its parent
frame, and writes the frame's contents in frame-relative coordinates: bodies,
geoms, sites, cameras, lights, joint anchors and axes, and nested frames.

Fixes that this required or exposed:

- mjCFrame cached its accumulated pose behind a `compiled` flag which was never
  reset, so editing the pos or quat of a frame after a first compilation was
  silently ignored. The flag is now cleared for all frames at the start of every
  compilation.
- Free-joint alignment moved a body and counter-transformed its children, but
  not its frames. It now carries the frames along, so a saved frame stays in
  place and its contents keep their authored poses. This mutates the compiled
  pose of the frames, hence the reset above.
- A body's <plugin> was written again inside each of the body's frames.
- childclass was written on every frame whose class is not "main". It is now
  written exactly when it differs from the class active around the frame.

Tests: frames holding every kind of element round-trip through save and reload
with their authored poses and the same compiled model, also inside a
free-aligned body; a body plugin is written once; frame pose edits take effect
on recompilation.
PiperOrigin-RevId: 985340056
Change-Id: Ic73f2e63e56e22915de2dfd32bd2a73438815167

Author: yuvaltassa

doc/XMLreference.rst

@@ -169,7 +169,8 @@ Meta elements
 
 These elements are not strictly part of the low-level MJCF format definition, but rather instruct the compiler to
 perform some operation on the model. A general property of meta-elements is that they disappear from the model upon
-saving the XML. There are currently six meta-elements in MJCF:
+saving the XML; the exception is :ref:`frame<frame>`, which is preserved. There are currently six meta-elements in
+MJCF:
 
 - :ref:`include<include>`, :ref:`frame<frame>`, and :ref:`replicate<replicate>` which are outside of the schema.
 - :ref:`composite<body-composite>`, :ref:`flexcomp<body-flexcomp>` and :ref:`attach<body-attach>` which are part of the
@@ -181,46 +182,10 @@ saving the XML. There are currently six meta-elements in MJCF:
 ^^^^^^^^^^^^^
 
 The frame meta-element is a pure coordinate transformation that can wrap any group of elements in the kinematic tree
-(under :ref:`worldbody<body>`). After compilation, frame elements disappear and their transformation is accumulated
-in their direct children. The attributes of the frame meta-element are documented :ref:`below<body-frame>`.
-
-.. collapse:: Usage example of frame
-
-   Loading this model and saving it:
-
-   .. code-block:: xml
-
-      <mujoco>
-        <worldbody>
-          <frame quat="0 0 1 0">
-             <geom name="Alice" quat="0 1 0 0" size="1"/>
-          </frame>
-
-          <frame pos="0 1 0">
-            <geom name="Bob" pos="0 1 0" size="1"/>
-            <body name="Carl" pos="1 0 0">
-              ...
-            </body>
-          </frame>
-        </worldbody>
-      </mujoco>
-
-   Results in this model:
-
-   .. code-block:: xml
-
-      <mujoco>
-        <worldbody>
-          <geom name="Alice" quat="0 0 0 1" size="1"/>
-          <geom name="Bob" pos="0 2 0" size="1"/>
-          <body name="Carl" pos="1 1 0">
-            ...
-          </body>
-        </worldbody>
-      </mujoco>
-
-   Note that in the saved model, the frame elements have disappeared but their transformation was accumulated with those
-   of their child elements.
+(under :ref:`worldbody<body>`). At compile time the transformation is accumulated into the frame's direct children;
+frames have no counterpart in :ref:`mjModel`. Unlike the other meta-elements, frames are preserved when the model is
+saved: the frame is written with its pose and its contents in frame-relative coordinates, so a saved model reloads
+with the same frames. The attributes of the frame meta-element are documented :ref:`below<body-frame>`.
 
 .. _replicate:
 
@@ -4229,8 +4194,8 @@ the saved XML file. Note that this element is a subset of the functionality of t
 :el-prefix:`body/` |-| **frame** |*|
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-Frames specify a coordinate transformation which is applied to all child elements. They disappear during compilation
-and the transformation they encode is accumulated in their direct children. See :ref:`frame<frame>` for examples.
+Frames specify a coordinate transformation which is applied to all child elements. During compilation the
+transformation they encode is accumulated in their direct children; frames are preserved when the model is saved.
 
 .. _frame-name:
 

doc/changelog.rst

@@ -21,6 +21,10 @@ General
   :ref:`mjpResourceProvider` interface to mount and read archive containers (such as
   ``.mjz``/``.zip``), decoupling container handling from format decoders and enabling
   on-demand asset extraction without requiring a pre-allocated :ref:`VFS<Virtualfilesystem>`.
+- :ref:`Frames<frame>` are now preserved when saving MJCF: ``<frame>`` elements are written with their ``pos`` and
+  ``quat`` and their contents in frame-relative coordinates, so a saved model reloads with the same frames. Previously
+  the frame transformation was accumulated into the children on save and only the frame's name was kept. Changes to the
+  ``pos`` or ``quat`` of a frame in an :ref:`mjSpec` now also take effect when the spec is recompiled.
 
 
 Engine

src/user/user_model.cc

@@ -5324,6 +5324,9 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
     timer[mjCTIMER_ASSETS] = Seconds(Clock::now() - t0).count();
   }
 
+  // frames cache their accumulated pose, recompute it in every compile
+  for (mjCFrame* frame : frames_) { frame->compiled = false; }
+
   // compile objects in kinematic tree
   for (int i = 0; i < bodies_.size(); i++) {
     bodies_[i]->Compile();  // also compiles joints, geoms, sites, cameras, lights, frames

src/user/user_objects.cc

@@ -2784,6 +2784,11 @@ void mjCBody::Compile(void) {
       mjuu_frameaccumChild(ipos_inverse, iquat_inverse, lights[i]->pos, qunit);
       mjuu_rotVecQuat(lights[i]->dir, lights[i]->dir, iquat_inverse);
     }
+
+    // frames: keep them in place relative to their contents, for the writer
+    for (int i = 0; i < frames.size(); i++) {
+      mjuu_frameaccumChild(ipos_inverse, iquat_inverse, frames[i]->pos, frames[i]->quat);
+    }
   }
 }
 

src/user/user_objects.h

@@ -655,6 +655,7 @@ class mjCFrame : public mjCFrame_, private mjsFrame {
   friend class mjCCamera;
   friend class mjCLight;
   friend class mjCModel;
+  friend class mjXWriter;
 
  public:
   mjCFrame(mjCModel* = 0, mjCFrame* = 0);

src/xml/xml_native_writer.cc

@@ -382,9 +382,15 @@ void mjXWriter::OneJoint(XMLElement*     elem,
                  &def->Joint().spec,
                  kJointAttrs,
                  kJointAttrsN);
-  if (joint->type != mjJNT_FREE) { WriteAttr(elem, "pos", 3, joint->pos, def->Joint().spec.pos); }
+  // pos and axis relative to the joint's frame
+  double pos[3], axis[3], iquat[4] = {1, 0, 0, 0};
+  mjuu_copyvec(pos, joint->pos, 3);
+  mjuu_copyvec(axis, joint->axis, 3);
+  FrameLocal(joint->frame, pos, iquat);
+  mjuu_rotVecQuat(axis, axis, iquat);
+  if (joint->type != mjJNT_FREE) { WriteAttr(elem, "pos", 3, pos, def->Joint().spec.pos); }
   if (joint->type != mjJNT_FREE && joint->type != mjJNT_BALL) {
-    WriteAttr(elem, "axis", 3, joint->axis, def->Joint().spec.axis);
+    WriteAttr(elem, "axis", 3, axis, def->Joint().spec.axis);
   }
   if (joint->type != mjJNT_FREE) {
     WriteAttrKey(elem, "limited", FalseTrueAuto_map, 3, joint->limited, def->Joint().limited);
@@ -422,30 +428,16 @@ void mjXWriter::OneGeom(XMLElement* elem, const mjCGeom* geom, mjCDef* def, stri
     }
     if (mjuu_defined(geom->mass)) { mass = geom->GetVolume() * def->Geom().density; }
 
-    // mesh geom
-    if (geom->type == mjGEOM_MESH || geom->type == mjGEOM_SDF) {
-      mjCMesh* mesh = geom->mesh;
-
-      // write pos/quat if there is a difference
-      if (!SameVector(geom->pos, mesh->GetPosPtr(), 3) ||
-          !SameVector(geom->quat, mesh->GetQuatPtr(), 4)) {
-        // recover geom pos/quat before mesh frame transformation
-        double p[3], q[4];
-        mjuu_copyvec(p, geom->pos, 3);
-        mjuu_copyvec(q, geom->quat, 4);
-        mjuu_frameaccuminv(p, q, mesh->GetPosPtr(), mesh->GetQuatPtr());
-
-        // write
-        WriteAttr(elem, "pos", 3, p, unitq + 1);
-        WriteAttr(elem, "quat", 4, q, unitq);
-      }
-    }
-
-    // non-mesh geom
-    else {
-      WriteAttr(elem, "pos", 3, geom->pos, unitq + 1);
-      WriteAttr(elem, "quat", 4, geom->quat, unitq);
+    // pose: undo the mesh transformation, then the frame
+    double pos[3], quat[4];
+    mjuu_copyvec(pos, geom->pos, 3);
+    mjuu_copyvec(quat, geom->quat, 4);
+    if ((geom->type == mjGEOM_MESH || geom->type == mjGEOM_SDF) && geom->mesh) {
+      mjuu_frameaccuminv(pos, quat, geom->mesh->GetPosPtr(), geom->mesh->GetQuatPtr());
     }
+    FrameLocal(geom->frame, pos, quat);
+    WriteAttr(elem, "pos", 3, pos, unitq + 1);
+    WriteAttr(elem, "quat", 4, quat, unitq);
   } else {
     WriteAttr(elem, "size", 3, geom->size, def->Geom().size);
   }
@@ -505,30 +497,16 @@ void mjXWriter::OneSite(XMLElement* elem, const mjCSite* site, mjCDef* def, stri
       WriteAttr(elem, "size", mjGEOMINFO[site->type], site->size, def->Site().size);
     }
 
-    // mesh site
-    if (site->type == mjGEOM_MESH) {
-      mjCMesh* mesh = site->mesh;
-
-      // write pos/quat if there is a difference
-      if (!SameVector(site->pos, mesh->GetPosPtr(), 3) ||
-          !SameVector(site->quat, mesh->GetQuatPtr(), 4)) {
-        // recover site pos/quat before mesh frame transformation
-        double p[3], q[4];
-        mjuu_copyvec(p, site->pos, 3);
-        mjuu_copyvec(q, site->quat, 4);
-        mjuu_frameaccuminv(p, q, mesh->GetPosPtr(), mesh->GetQuatPtr());
-
-        // write
-        WriteAttr(elem, "pos", 3, p, unitq + 1);
-        WriteAttr(elem, "quat", 4, q, unitq);
-      }
-    }
-
-    // non-mesh site
-    else {
-      WriteAttr(elem, "pos", 3, site->pos, unitq + 1);
-      WriteAttr(elem, "quat", 4, site->quat, unitq);
+    // pose: undo the mesh transformation, then the frame
+    double pos[3], quat[4];
+    mjuu_copyvec(pos, site->pos, 3);
+    mjuu_copyvec(quat, site->quat, 4);
+    if (site->type == mjGEOM_MESH && site->mesh) {
+      mjuu_frameaccuminv(pos, quat, site->mesh->GetPosPtr(), site->mesh->GetQuatPtr());
     }
+    FrameLocal(site->frame, pos, quat);
+    WriteAttr(elem, "pos", 3, pos, unitq + 1);
+    WriteAttr(elem, "quat", 4, quat, unitq);
   } else {
     WriteAttr(elem, "size", 3, site->size, def->Site().size);
   }
@@ -562,22 +540,22 @@ void mjXWriter::OneCamera(XMLElement*      elem,
                           string_view      classname) {
   double unitq[4] = {1, 0, 0, 0};
 
+  // pose relative to the camera's frame
+  mjsCamera local = *static_cast<const mjsCamera*>(camera);
+  FrameLocal(camera->frame, local.pos, local.quat);
+
   // regular
   if (!writingdefaults) {
     WriteAttrTxt(elem, "name", camera->name);
     if (classname != camera->classname && camera->classname != "main") {
       WriteAttrTxt(elem, "class", camera->classname);
     }
     WriteAttrTxt(elem, "target", camera->get_targetbody());
-    WriteAttr(elem, "quat", 4, camera->quat, unitq);
+    WriteAttr(elem, "quat", 4, local.quat, unitq);
   }
 
   // defaults and regular
-  WriteAttrTable(elem,
-                 static_cast<const mjsCamera*>(camera),
-                 &def->Camera().spec,
-                 kCameraAttrs,
-                 kCameraAttrsN);
+  WriteAttrTable(elem, &local, &def->Camera().spec, kCameraAttrs, kCameraAttrsN);
 
   // camera intrinsics if specified
   if (camera->sensor_size[0] > 0 && camera->sensor_size[1] > 0) {
@@ -612,12 +590,14 @@ void mjXWriter::OneLight(XMLElement*     elem,
     WriteAttrTxt(elem, "target", light->get_targetbody());
   }
 
+  // pos and dir relative to the light's frame
+  mjsLight local    = *static_cast<const mjsLight*>(light);
+  double   iquat[4] = {1, 0, 0, 0};
+  FrameLocal(light->frame, local.pos, iquat);
+  mjuu_rotVecQuat(local.dir, local.dir, iquat);
+
   // defaults and regular
-  WriteAttrTable(elem,
-                 static_cast<const mjsLight*>(light),
-                 &def->Light().spec,
-                 kLightAttrs,
-                 kLightAttrsN);
+  WriteAttrTable(elem, &local, &def->Light().spec, kLightAttrs, kLightAttrsN);
   WriteAttrKey(elem, "type", lighttype_map, lighttype_sz, light->type, def->Light().type);
   WriteAttrTxt(elem, "texture", light->get_texture());
 }
@@ -1544,17 +1524,42 @@ void mjXWriter::Asset(XMLElement* root) {
 }
 
 
-XMLElement* mjXWriter::OneFrame(XMLElement* elem, mjCFrame* frame) {
+// strip a frame from a compiled pose: pos/quat become relative to the frame
+void mjXWriter::FrameLocal(const mjCFrame* frame, double pos[3], double quat[4]) {
+  if (!frame) { return; }
+  double ipos[3], iquat[4];
+  mjuu_frameinvert(ipos, iquat, frame->pos, frame->quat);
+  mjuu_frameaccumChild(ipos, iquat, pos, quat);
+}
+
+
+XMLElement* mjXWriter::OneFrame(XMLElement* elem, mjCFrame* frame, string_view childclass) {
   if (!frame) { return elem; }
+  double unitq[4] = {1, 0, 0, 0};
 
-  // TODO: empty classname should not occur (but does)
-  if (frame->name.empty() && (frame->classname.empty() || frame->classname == "main")) {
+  // pose relative to the parent frame
+  double pos[3], quat[4];
+  mjuu_copyvec(pos, frame->pos, 3);
+  mjuu_copyvec(quat, frame->quat, 4);
+  FrameLocal(frame->frame, pos, quat);
+
+  // omit unnamed identity frame with no childclass
+  bool has_class = !frame->classname.empty() && frame->classname != childclass;
+  if (frame->name.empty() && !has_class &&
+      SameVector(pos, unitq + 1, 3) && SameVector(quat, unitq, 4)) {
     return elem;
   }
 
   XMLElement* frame_elem = InsertEnd(elem, "frame");
   WriteAttrTxt(frame_elem, "name", frame->name);
-  if (frame->classname != "main") { WriteAttrTxt(frame_elem, "childclass", frame->classname); }
+
+  // childclass, unless inherited from the enclosing frame or body
+  if (has_class) {
+    WriteAttrTxt(frame_elem, "childclass", frame->classname);
+  }
+
+  WriteAttr(frame_elem, "pos", 3, pos, unitq + 1);
+  WriteAttr(frame_elem, "quat", 4, quat, unitq);
   return frame_elem;
 }
 
@@ -1574,9 +1579,13 @@ void mjXWriter::Body(XMLElement* elem, mjCBody* body, mjCFrame* frame, string_vi
       WriteAttrTxt(elem, "childclass", body->classname);
     }
 
-    // write pos if it's not {0, 0, 0}
-    if (body->pos[0] || body->pos[1] || body->pos[2]) { WriteAttr(elem, "pos", 3, body->pos); }
-    WriteAttr(elem, "quat", 4, body->quat, unitq);
+    // pose relative to the body's frame
+    double pos[3], quat[4];
+    mjuu_copyvec(pos, body->pos, 3);
+    mjuu_copyvec(quat, body->quat, 4);
+    FrameLocal(body->frame, pos, quat);
+    WriteAttr(elem, "pos", 3, pos, unitq + 1);
+    WriteAttr(elem, "quat", 4, quat, unitq);
     if (body->mocap) { WriteAttrKey(elem, "mocap", bool_map, 2, 1); }
 
     // gravity compensation
@@ -1666,8 +1675,8 @@ void mjXWriter::Body(XMLElement* elem, mjCBody* body, mjCFrame* frame, string_vi
              classname.empty() ? childclass : classname);
   }
 
-  // write plugin
-  if (body->plugin.active) { OnePlugin(InsertEnd(elem, "plugin"), &body->plugin); }
+  // write plugin, once: not again inside the body's frames
+  if (!frame && body->plugin.active) { OnePlugin(InsertEnd(elem, "plugin"), &body->plugin); }
 
   // write children recursively
   int i = 0, j = 0;
@@ -1694,9 +1703,12 @@ void mjXWriter::Body(XMLElement* elem, mjCBody* body, mjCFrame* frame, string_vi
 
       // write frame if its frame matches the current frame
       if (fframe->frame == frame) {
-        string classname =
-            fframe && !fframe->classname.empty() ? fframe->classname : body->classname;
-        Body(OneFrame(elem, fframe), body, fframe, childclass);
+        // class active around the new frame: the current frame's, else the body's
+        string classname = frame && !frame->classname.empty() ? frame->classname : body->classname;
+        Body(OneFrame(elem, fframe, classname.empty() ? childclass : classname),
+             body,
+             fframe,
+             childclass);
       }
     }
 

src/xml/xml_native_writer.h

@@ -105,7 +105,12 @@ class mjXWriter : public mjXBase {
   void OneTendon(tinyxml2::XMLElement* elem, const mjCTendon* ptendon, mjCDef* def);
   void OneActuator(tinyxml2::XMLElement* elem, const mjCActuator* pactuator, mjCDef* def);
   void OnePlugin(tinyxml2::XMLElement* elem, const mjsPlugin* plugin);
-  tinyxml2::XMLElement* OneFrame(tinyxml2::XMLElement* elem, mjCFrame* frame);
+  tinyxml2::XMLElement* OneFrame(tinyxml2::XMLElement* elem,
+                                 mjCFrame*             frame,
+                                 std::string_view      childclass);
+
+  // strip a frame from a compiled pose: pos/quat become relative to the frame
+  static void FrameLocal(const mjCFrame* frame, double pos[3], double quat[4]);
 
   bool writingdefaults;  // true during defaults write
 };