Do not overwrite existing parent frames for worldbody elements when attaching an mjSpec to a frame or a site.

PiperOrigin-RevId: 737662046
Change-Id: I7028c21c803c01845ae1d53cb476f7b3df172313

doc/APIreference/functions.rst

@@ -4202,6 +4202,15 @@ Find child body by name.
 
 Get parent body.
 
+.. _mjs_getFrame:
+
+`mjs_getFrame <#mjs_getFrame>`__
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mjs_getFrame
+
+Get parent frame.
+
 .. _mjs_findFrame:
 
 `mjs_findFrame <#mjs_findFrame>`__

doc/changelog.rst

@@ -15,6 +15,8 @@ Bug fixes
 ^^^^^^^^^
 - :ref:`mj_jacDot` was missing a term that accounts for the motion of the point with respect to
   which the Jacobian is computed, now fixed.
+- Fixed a bug that caused the parent frame of elements in the child worldbody to be incorrectly set when attaching an
+  mjSpec to a frame or a site.
 
 Version 3.3.0 (Feb 26, 2025)
 ----------------------------

doc/includes/references.h

@@ -3649,6 +3649,7 @@ mjsBody* mjs_findBody(mjSpec* s, const char* name);
 mjsElement* mjs_findElement(mjSpec* s, mjtObj type, const char* name);
 mjsBody* mjs_findChild(mjsBody* body, const char* name);
 mjsBody* mjs_getParent(mjsElement* element);
+mjsFrame* mjs_getFrame(mjsElement* element);
 mjsFrame* mjs_findFrame(mjSpec* s, const char* name);
 mjsDefault* mjs_getDefault(mjsElement* element);
 const mjsDefault* mjs_findDefault(mjSpec* s, const char* classname);

include/mujoco/mujoco.h

@@ -1556,6 +1556,9 @@ MJAPI mjsBody* mjs_findChild(mjsBody* body, const char* name);
 // Get parent body.
 MJAPI mjsBody* mjs_getParent(mjsElement* element);
 
+// Get parent frame.
+MJAPI mjsFrame* mjs_getFrame(mjsElement* element);
+
 // Find frame by name.
 MJAPI mjsFrame* mjs_findFrame(mjSpec* s, const char* name);
 

python/mujoco/introspect/functions.py

@@ -9896,6 +9896,22 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Get parent body.',
      )),
+    ('mjs_getFrame',
+     FunctionDecl(
+         name='mjs_getFrame',
+         return_type=PointerType(
+             inner_type=ValueType(name='mjsFrame'),
+         ),
+         parameters=(
+             FunctionParameterDecl(
+                 name='element',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjsElement'),
+                 ),
+             ),
+         ),
+         doc='Get parent frame.',
+     )),
     ('mjs_findFrame',
      FunctionDecl(
          name='mjs_findFrame',

python/mujoco/specs.cc

@@ -272,7 +272,7 @@ py::list FindAllImpl(raw::MjsBody& body, mjtObj objtype, bool recursive) {
 void SetFrame(raw::MjsBody* body, mjtObj objtype, raw::MjsFrame* frame) {
   mjsElement* el = mjs_firstChild(body, objtype, 0);
   while (el) {
-    if (frame->element != el) {
+    if (frame->element != el && mjs_getFrame(el) == nullptr) {
       mjs_setFrame(el, frame);
     }
     el = mjs_nextChild(body, el, 0);

python/mujoco/specs_test.py

@@ -1072,6 +1072,7 @@ class SpecsTest(absltest.TestCase):
     child2 = mujoco.MjSpec()
     child2.assets = {'cube2.obj': 'cube2_content'}
     body2 = child2.worldbody.add_body(name='body')
+    body2.set_frame(child2.worldbody.add_frame(pos=[-1, -1, 1]))
     self.assertIsNotNone(parent.attach(child2, frame=frame, prefix='child-'))
     self.assertIsNotNone(child2.worldbody)
     self.assertEqual(child2.parent, parent)
@@ -1080,7 +1081,7 @@ class SpecsTest(absltest.TestCase):
     self.assertIsNotNone(model2)
     self.assertEqual(model2.nbody, 3)
     np.testing.assert_array_equal(model2.body_pos[1], [0, 1, 4])
-    np.testing.assert_array_equal(model2.body_pos[2], [2, 3, 2])
+    np.testing.assert_array_equal(model2.body_pos[2], [3, 4, 3])
     np.testing.assert_array_equal(model2.body_quat[1], [0, 0, 0, 1])
     np.testing.assert_array_equal(model2.body_quat[2], [0, 0, 0, 1])
     self.assertEqual(parent.assets['cube.obj'], 'cube_content')
@@ -1090,6 +1091,7 @@ class SpecsTest(absltest.TestCase):
     child3 = mujoco.MjSpec()
     child3.assets = {'cube2.obj': 'new_cube2_content'}
     body3 = child3.worldbody.add_body(name='body')
+    body3.set_frame(child3.worldbody.add_frame(pos=[-1, -1, 1]))
     self.assertIsNotNone(parent.attach(child3, frame='frame', prefix='child3-'))
     self.assertIsNotNone(child3.worldbody)
     self.assertEqual(child3.parent, parent)
@@ -1098,8 +1100,8 @@ class SpecsTest(absltest.TestCase):
     self.assertIsNotNone(model3)
     self.assertEqual(model3.nbody, 4)
     np.testing.assert_array_equal(model3.body_pos[1], [0, 1, 4])
-    np.testing.assert_array_equal(model3.body_pos[2], [2, 3, 2])
-    np.testing.assert_array_equal(model3.body_pos[3], [3, 4, 1])
+    np.testing.assert_array_equal(model3.body_pos[2], [3, 4, 3])
+    np.testing.assert_array_equal(model3.body_pos[3], [4, 5, 2])
     np.testing.assert_array_equal(model3.body_quat[1], [0, 0, 0, 1])
     np.testing.assert_array_equal(model3.body_quat[2], [0, 0, 0, 1])
     np.testing.assert_array_equal(model3.body_quat[3], [0, 0, 0, 1])

src/user/user_api.cc

@@ -711,6 +711,25 @@ mjsBody* mjs_getParent(mjsElement* element) {
 
 
 
+// get parent frame
+mjsFrame* mjs_getFrame(mjsElement* element) {
+  mjCBase* base = static_cast<mjCBase*>(element);
+  switch (element->elemtype) {
+    case mjOBJ_BODY:
+    case mjOBJ_FRAME:
+    case mjOBJ_JOINT:
+    case mjOBJ_GEOM:
+    case mjOBJ_SITE:
+    case mjOBJ_CAMERA:
+    case mjOBJ_LIGHT:
+      return base->frame ? &(base->frame->spec) : nullptr;
+    default:
+      return nullptr;
+  }
+}
+
+
+
 // find frame by name
 mjsFrame* mjs_findFrame(mjSpec* s, const char* name) {
   mjsElement* frame = mjs_findElement(s, mjOBJ_FRAME, name);

src/user/user_api.h

@@ -211,6 +211,9 @@ MJAPI mjsBody* mjs_findChild(mjsBody* body, const char* name);
 // Get parent body.
 MJAPI mjsBody* mjs_getParent(mjsElement* element);
 
+// Get parent frame.
+MJAPI mjsFrame* mjs_getFrame(mjsElement* element);
+
 // Find frame by name.
 MJAPI mjsFrame* mjs_findFrame(mjSpec* s, const char* name);