Support nested frames in mjs_attachFrame.

PiperOrigin-RevId: 626001651
Change-Id: I3dc52baa6aad875504dd4c7d969a4091f52c4c16

src/user/user_objects.cc

@@ -739,13 +739,15 @@ mjCBody& mjCBody::operator+=(const mjCFrame& other) {
   frames.back()->body = this;
   frames.back()->model = model;
   frames.back()->frame = other.frame;
+  int i = frames.size();
 
   // map input frames to index in this->frames
   std::map<mjCFrame*, int> fmap;
   for (auto frame : subtree->frames) {
     if (frame == static_cast<const mjCFrame*>(&other)) {
       fmap[frame] = frames.size() - 1;
-      break;
+    } else if (other.IsAncestor(frame)) {
+      fmap[frame] = i++;
     }
   }
 
@@ -758,7 +760,7 @@ mjCBody& mjCBody::operator+=(const mjCFrame& other) {
   CopyList(lights, subtree->lights, fmap, &other);
 
   for (int i=0; i<subtree->bodies.size(); i++) {
-    if (subtree->bodies[i]->frame != &other) {
+    if (!other.IsAncestor(subtree->bodies[i]->frame)) {
       continue;
     }
     bodies.push_back(new mjCBody(*subtree->bodies[i], model));  // triggers recursive call
@@ -786,7 +788,7 @@ void mjCBody::CopyList(std::vector<T*>& dst, const std::vector<T*>& src,
                        std::map<mjCFrame*, int>& fmap, const mjCFrame* pframe) {
   int nsrc = (int)src.size();
   for (int i=0; i<nsrc; i++) {
-    if (pframe && src[i]->frame != pframe) {
+    if (pframe && !pframe->IsAncestor(src[i]->frame)) {
       continue;  // skip if the element is not inside pframe
     }
     dst.push_back(new T(*src[i]));
@@ -1473,6 +1475,21 @@ mjCFrame& mjCFrame::operator+=(const mjCBody& other) {
 
 
 
+// return true if child is descendent of this frame
+bool mjCFrame::IsAncestor(const mjCFrame* child) const {
+  if (!child) {
+    return false;
+  }
+
+  if (child == this) {
+    return true;
+  }
+
+  return IsAncestor(child->frame);
+}
+
+
+
 void mjCFrame::SetParent(mjCBody* _body) {
   body = _body;
 }

src/user/user_objects.h

@@ -362,6 +362,8 @@ class mjCFrame : public mjCFrame_, private mjsFrame {
 
   mjCFrame& operator+=(const mjCBody& other);
 
+  bool IsAncestor(const mjCFrame* child) const;  // true if child is contained in this frame
+
  private:
   mjCFrame(mjCModel* = 0, mjCFrame* = 0);      // constructor
   mjCFrame(const mjCFrame& other);             // copy constructor

test/user/user_api_test.cc

@@ -200,13 +200,15 @@ static constexpr char xml_child[] = R"(
   <mujoco>
     <worldbody>
       <frame name="pframe">
-        <body name="body">
+        <frame name="cframe">
-          <joint type="hinge" name="hinge"/>
+          <body name="body">
-          <geom type="cylinder" size=".1 1 0"/>
+            <joint type="hinge" name="hinge"/>
-          <light mode="targetbody" target="targetbody"/>
+            <geom type="cylinder" size=".1 1 0"/>
-          <body name="targetbody"/>
+            <light mode="targetbody" target="targetbody"/>
-          <body/>
+            <body name="targetbody"/>
-        </body>
+            <body/>
+          </body>
+        </frame>
       </frame>
       <body name="ignore"/>
       <frame name="frame" pos=".1 0 0" euler="0 90 0"/>
@@ -437,13 +439,15 @@ TEST_F(MujocoTest, AttachFrame) {
         <geom size=".1"/>
         <frame name="frame" pos=".1 0 0" euler="0 90 0"/>
         <frame name="pframe">
-          <body name="attached-body-1">
+          <frame name="cframe">
-            <joint type="hinge" name="attached-hinge-1"/>
+            <body name="attached-body-1">
-            <geom type="cylinder" size=".1 1 0"/>
+              <joint type="hinge" name="attached-hinge-1"/>
-            <light mode="targetbody" target="attached-targetbody-1"/>
+              <geom type="cylinder" size=".1 1 0"/>
-              <body name="attached-targetbody-1"/>
+              <light mode="targetbody" target="attached-targetbody-1"/>
-              <body/>
+                <body name="attached-targetbody-1"/>
-          </body>
+                <body/>
+            </body>
+          </frame>
         </frame>
       </body>
     </worldbody>