Fix bug in frame saving. Fixes #1802

PiperOrigin-RevId: 651696068
Change-Id: Ie1e2e1de4e83c42f2b5ebcc3afd0540b9afe7bbd

src/xml/xml_native_writer.cc

@@ -867,7 +867,7 @@ string mjXWriter::Write(char *error, size_t error_sz) {
   Extension(root);
   Custom(root);
   Asset(root);
-  Body(InsertEnd(root, "worldbody"), model->GetWorld());
+  Body(InsertEnd(root, "worldbody"), model->GetWorld(), nullptr);
   Contact(root);
   Deformable(root);
   Equality(root);
@@ -1526,14 +1526,17 @@ XMLElement* mjXWriter::OneFrame(XMLElement* elem, mjCFrame* frame) {
 
   XMLElement* frame_elem = InsertEnd(elem, "frame");
   WriteAttrTxt(frame_elem, "name", frame->name);
-  WriteAttrTxt(frame_elem, "childclass", frame->classname);
+  if (frame->classname != "main") {
+    WriteAttrTxt(frame_elem, "childclass", frame->classname);
+  }
   return frame_elem;
 }
 
 
 
 // recursive body and frame writer
-void mjXWriter::Body(XMLElement* elem, mjCBody* body, std::string_view childclass) {
+void mjXWriter::Body(XMLElement* elem, mjCBody* body, mjCFrame* frame,
+                     std::string_view childclass) {
   double unitq[4] = {1, 0, 0, 0};
 
   if (!body) {
@@ -1541,7 +1544,7 @@ void mjXWriter::Body(XMLElement* elem, mjCBody* body, std::string_view childclas
   }
 
   // write body attributes and inertial
-  else if (body!=model->GetWorld()) {
+  else if (!frame && body!=model->GetWorld()) {
     WriteAttrTxt(elem, "name", body->name);
     if (childclass != body->classname && body->classname != "main") {
       WriteAttrTxt(elem, "childclass", body->classname);
@@ -1574,57 +1577,67 @@ void mjXWriter::Body(XMLElement* elem, mjCBody* body, std::string_view childclas
     }
   }
 
-  // write joints
+  // joints in this frame
-  for (int i=0; i<body->joints.size(); i++) {
+  for (int i = 0; i < body->joints.size(); i++) {
-    XMLElement* celem = OneFrame(elem, body->joints[i]->frame);
+    if (body->joints[i]->frame != frame) {
+      continue;
+    }
     std::string classname = body->joints[i]->frame && !body->joints[i]->frame->classname.empty()
                                 ? body->joints[i]->frame->classname
                                 : body->classname;
-    OneJoint(InsertEnd(celem, "joint"), body->joints[i],
+    OneJoint(InsertEnd(elem, "joint"), body->joints[i],
              model->def_map[body->joints[i]->classname],
              classname.empty() ? childclass : classname);
   }
 
-  // write geoms
+  // geoms in this frame
-  for (int i=0; i<body->geoms.size(); i++) {
+  for (int i = 0; i < body->geoms.size(); i++) {
-    XMLElement* celem = OneFrame(elem, body->geoms[i]->frame);
+    if (body->geoms[i]->frame != frame) {
+      continue;
+    }
     std::string classname = body->geoms[i]->frame && !body->geoms[i]->frame->classname.empty()
                                 ? body->geoms[i]->frame->classname
                                 : body->classname;
-    OneGeom(InsertEnd(celem, "geom"), body->geoms[i],
+    OneGeom(InsertEnd(elem, "geom"), body->geoms[i],
             model->def_map[body->geoms[i]->classname],
             classname.empty() ? childclass : classname);
   }
 
-  // write sites
+  // sites in this frame
-  for (int i=0; i<body->sites.size(); i++) {
+  for (int i = 0; i < body->sites.size(); i++) {
-    XMLElement* celem = OneFrame(elem, body->sites[i]->frame);
+    if (body->sites[i]->frame != frame) {
+      continue;
+    }
     std::string classname = body->sites[i]->frame && !body->sites[i]->frame->classname.empty()
                                 ? body->sites[i]->frame->classname
                                 : body->classname;
-    OneSite(InsertEnd(celem, "site"), body->sites[i],
+    OneSite(InsertEnd(elem, "site"), body->sites[i],
             model->def_map[body->sites[i]->classname],
             classname.empty() ? childclass : classname);
   }
 
-  // write cameras
+  // cameras in this frame
-  for (int i=0; i<body->cameras.size(); i++) {
+  for (int i = 0; i < body->cameras.size(); i++) {
-    XMLElement* celem = OneFrame(elem, body->cameras[i]->frame);
+    if (body->cameras[i]->frame != frame) {
+      continue;
+    }
     std::string classname = body->cameras[i]->frame && !body->cameras[i]->frame->classname.empty()
                                 ? body->cameras[i]->frame->classname
                                 : body->classname;
-    OneCamera(InsertEnd(celem, "camera"), body->cameras[i],
+    OneCamera(InsertEnd(elem, "camera"), body->cameras[i],
               model->def_map[body->cameras[i]->classname],
               classname.empty() ? childclass : classname);
   }
 
-  // write lights
+  // lights in this frame
-  for (int i=0; i<body->lights.size(); i++) {
+  for (int i = 0; i < body->lights.size(); i++) {
-    XMLElement* celem = OneFrame(elem, body->lights[i]->frame);
+    if (body->lights[i]->frame != frame) {
+      continue;
+    }
     std::string classname = body->lights[i]->frame && !body->lights[i]->frame->classname.empty()
                                 ? body->lights[i]->frame->classname
                                 : body->classname;
-    OneLight(InsertEnd(celem, "light"), body->lights[i],
+    OneLight(InsertEnd(elem, "light"), body->lights[i],
              model->def_map[body->lights[i]->classname],
              classname.empty() ? childclass : classname);
   }
@@ -1634,13 +1647,49 @@ void mjXWriter::Body(XMLElement* elem, mjCBody* body, std::string_view childclas
     OnePlugin(InsertEnd(elem, "plugin"), &body->plugin);
   }
 
-  // write child bodies recursively
+  // write children recursively
-  for (int i=0; i<body->bodies.size(); i++) {
+  int i = 0, j = 0;
-    XMLElement* celem = OneFrame(elem, body->bodies[i]->frame);
+  while (i < body->bodies.size() || body->bodies.empty()) {
-    std::string classname = body->bodies[i]->frame && !body->bodies[i]->frame->classname.empty()
+    mjCFrame* bframe = body->bodies.empty() ? nullptr : body->bodies[i]->frame;
-                                ? body->bodies[i]->frame->classname
+
-                                : body->classname;
+    // write body if its frame matches the current frame, avoid access if there are no bodies
-    Body(InsertEnd(celem, "body"), body->bodies[i], classname.empty() ? childclass : classname);
+    if (bframe == frame && !body->bodies.empty()) {
+      std::string classname = bframe && !bframe->classname.empty()
+                                  ? bframe->classname
+                                  : body->classname;
+      Body(InsertEnd(elem, "body"), body->bodies[i], nullptr,
+            classname.empty() ? childclass : classname);
+    }
+
+    i++;
+
+    // do not go to frames until we reach a body with a frame or we are done with bodies
+    if (!bframe && i < body->bodies.size()) {
+      continue;
+    }
+
+    // loop over the remaining frames in the current body
+    while (j < body->frames.size()) {
+      mjCFrame* fframe = body->frames[j++];
+
+      // write frame if its frame matches the current frame
+      if (fframe->frame == frame) {
+        std::string classname = fframe && !fframe->classname.empty()
+                                    ? fframe->classname
+                                    : body->classname;
+        Body(OneFrame(elem, fframe), body, fframe, childclass);
+      }
+
+      // stop if we reached the frame of the current child body, ignore if there are no bodies
+      if (bframe && bframe == fframe) {
+        break;
+      }
+    }
+
+    // if there are no bodies, we only want to run the loop once
+    if (body->bodies.empty()) {
+      break;
+    }
   }
 }
 

src/xml/xml_native_writer.h

@@ -59,7 +59,7 @@ class mjXWriter : public mjXBase {
   void Keyframe(tinyxml2::XMLElement* root);                              // keyframe section
 
   // body/world section
-  void Body(tinyxml2::XMLElement* elem, mjCBody* body, std::string_view childclass = "");
+  void Body(tinyxml2::XMLElement* elem, mjCBody* body, mjCFrame* frame, std::string_view childclass = "");
 
   // single element writers, used in defaults and main body
   void OneFlex(tinyxml2::XMLElement* elem, const mjCFlex* pflex);

test/user/testdata/frame1.xml

@@ -0,0 +1,24 @@
+<mujoco model="load save load test model with 2 frames">
+
+  <worldbody>
+    <frame name="frame1" pos=".1 .2 .3">
+
+      <geom name="geom1" type="sphere" size=".1"/>
+      <body name="root1">
+        <joint name="joint1"/>
+        <geom type="sphere" size=".1"/>
+      </body>
+
+      <body name="root2">
+        <frame name="frame2" pos=".01 .02 .03">
+          <geom type="sphere" size=".1"/>
+          <body name="root22">
+            <geom type="sphere" size=".1"/>
+          </body>
+        </frame><!-- end frame2 -->
+      </body>
+
+    </frame> <!-- end frame1 -->
+  </worldbody>
+
+</mujoco>

test/user/testdata/frame2.xml

@@ -0,0 +1,31 @@
+<mujoco model="load save load test model with 2 frames">
+
+  <worldbody>
+    <frame name="frame1" pos=".1 .2 .3">
+
+      <geom name="geom1" type="sphere" size=".1"/>
+      <body name="root1">
+        <geom type="sphere" size=".1"/>
+      </body>
+
+      <frame name="frame2" pos=".01 .02 .03">
+
+        <geom name="geom2" type="sphere" size=".1"/>
+        <body name="root2">
+          <geom type="sphere" size=".1"/>
+          <body name="root22">
+            <geom type="sphere" size=".1"/>
+          </body>
+        </body>
+
+        <geom name="geom3" type="sphere" size=".1"/>
+        <body name="root3">
+          <geom type="sphere" size=".1"/>
+          <body name="root5">
+            <geom type="sphere" size=".1"/>
+          </body>
+        </body>
+      </frame> <!-- end frame1 -->
+    </frame><!-- end frame2 -->
+  </worldbody>
+</mujoco>

test/xml/xml_native_writer_test.cc

@@ -726,7 +726,7 @@ TEST_F(XMLWriterTest, WritesActuatorDefaults) {
 
 TEST_F(XMLWriterTest, WritesFrameDefaults) {
   static constexpr char xml[] = R"(
-  <mujoco>
+  <mujoco model="test">
     <default>
       <default class="dframe">
         <geom size=".1"/>
@@ -741,24 +741,53 @@ TEST_F(XMLWriterTest, WritesFrameDefaults) {
       <body name="body">
         <frame pos="0 1 0" name="f2" childclass="dframe">
           <geom pos="0 1 0"/>
-          <body pos="1 0 0">
+          <frame pos="0 1 0" name="f3">
-            <geom pos="0 0 1"/>
+            <frame pos="0 1 0" name="f4">
-          </body>
+              <body pos="1 0 0">
+                <geom pos="0 0 1"/>
+              </body>
+            </frame>
+          </frame>
         </frame>
       </body>
     </worldbody>
   </mujoco>
   )";
+
+  static constexpr char xml_expected[] = R"(<mujoco model="test">
+  <compiler angle="radian"/>
+
+  <default>
+    <default class="dframe">
+      <geom size="0.1 0 0"/>
+    </default>
+  </default>
+
+  <worldbody>
+    <body name="body">
+      <frame name="f2" childclass="dframe">
+        <geom pos="0 2 0"/>
+        <frame name="f3" childclass="dframe">
+          <frame name="f4" childclass="dframe">
+            <body pos="1 3 0">
+              <geom pos="0 0 1"/>
+            </body>
+          </frame>
+        </frame>
+      </frame>
+    </body>
+    <frame name="f1">
+      <geom size="0.5" quat="0.906308 0 0 0.422618"/>
+    </frame>
+  </worldbody>
+</mujoco>
+)";
+
   std::array<char, 1024> error;
   mjModel* model = LoadModelFromString(xml, error.data(), error.size());
   EXPECT_THAT(model, NotNull()) << error.data();
   std::string saved_xml = SaveAndReadXml(model);
-  EXPECT_THAT(saved_xml, HasSubstr("frame name=\"f1\""));
+  EXPECT_STREQ(saved_xml.c_str(), xml_expected);
-  EXPECT_THAT(saved_xml, HasSubstr("frame name=\"f2\" childclass=\"dframe\">"));
-  EXPECT_THAT(saved_xml, HasSubstr("geom pos=\"0 2 0\""));
-  EXPECT_THAT(saved_xml, HasSubstr("body pos=\"1 1 0\""));
-  EXPECT_THAT(saved_xml, HasSubstr("geom pos=\"0 0 1\""));
-  EXPECT_THAT(saved_xml, Not(HasSubstr("<frame>")));
   mj_deleteModel(model);
 }