Enable frames to create the same children as bodies.

PiperOrigin-RevId: 822097998
Change-Id: I8477d8abe2d6f711421126825fd94d39af1a4e07

python/mujoco/codegen/generate_spec_bindings.py

@@ -52,6 +52,15 @@ SPECS = [
     ('mjsExclude',  'Spec', False, 'excludes',   'mjOBJ_EXCLUDE'),
     ('mjsPlugin',   'Spec', False, 'plugins',    'mjOBJ_PLUGIN'),
 ]
+SPECS_ADD = SPECS + [
+    ('mjsBody',     'Frame', True, 'bodies',     'mjOBJ_BODY'),
+    ('mjsSite',     'Frame', True, 'sites',      'mjOBJ_SITE'),
+    ('mjsGeom',     'Frame', True, 'geoms',      'mjOBJ_GEOM'),
+    ('mjsJoint',    'Frame', True, 'joints',     'mjOBJ_JOINT'),
+    ('mjsCamera',   'Frame', True, 'cameras',    'mjOBJ_CAMERA'),
+    ('mjsFrame',    'Frame', True, 'frames',     'mjOBJ_FRAME'),
+    ('mjsLight',    'Frame', True, 'lights',     'mjOBJ_LIGHT'),
+]
 # pylint: enable=bad-whitespace
 
 
@@ -301,7 +310,7 @@ def generate() -> None:
 
 def generate_add() -> None:
   """Generate add constructors with optional keyword arguments."""
-  for key, parent, default, listname, objtype in SPECS:
+  for key, parent, default, listname, objtype in SPECS_ADD:
 
     def _field(f: ast_nodes.StructFieldDecl):
       if f.type == ast_nodes.PointerType(
@@ -390,6 +399,22 @@ def generate_add() -> None:
             raw::MjsDefault* default_, py::kwargs kwargs) -> raw::{titlecase}* {{
             auto out = mjs_add{elem}(&self, default_);
         """
+    elif parent == 'Frame':
+      if key == 'mjsFrame':
+        code = f"""
+          {'mjs' + parent}.def("add_{elemlower}", []({'raw::Mjs' + parent}& self,
+            raw::MjsFrame* parentframe_, py::kwargs kwargs) -> raw::{titlecase}* {{
+            raw::MjsBody* body = mjs_getParent(self.element);
+            auto out = mjs_add{elem}(body, &self);
+        """
+      else:
+        code = f"""
+          {'mjs' + parent}.def("add_{elemlower}", []({'raw::Mjs' + parent}& self,
+            raw::MjsDefault* default_, py::kwargs kwargs) -> raw::{titlecase}* {{
+            raw::MjsBody* body = mjs_getParent(self.element);
+            auto out = mjs_add{elem}(body, default_);
+            mjs_setFrame(out->element, &self);
+        """
     else:
       raise NotImplementedError(f'{parent} parent is not implement.')
 

python/mujoco/specs_test.py

@@ -318,6 +318,14 @@ class SpecsTest(absltest.TestCase):
     np.testing.assert_array_equal(frameb1.pos, frameb0.pos)
     np.testing.assert_array_equal(frameb1.quat, frameb0.quat)
 
+    # Add frame in frame.
+    framec0 = frameb0.add_frame(name='framec', pos=[7, 8, 9], quat=[0, 0, 0, 1])
+    self.assertEqual(framec0.name, 'framec')
+    self.assertEqual(framec0.parent, frameb0.parent)
+    self.assertEqual(framec0.frame, frameb0)
+    np.testing.assert_array_equal(framec0.pos, [7, 8, 9])
+    np.testing.assert_array_equal(framec0.quat, [0, 0, 0, 1])
+
     # Add joint.
     joint = body.add_joint(type=mujoco.mjtJoint.mjJNT_HINGE, axis=[0, 1, 0])
     self.assertEqual(joint.type, mujoco.mjtJoint.mjJNT_HINGE)
@@ -340,6 +348,12 @@ class SpecsTest(absltest.TestCase):
     light = body.add_light(attenuation=[1, 2, 3])
     np.testing.assert_array_equal(light.attenuation, [1, 2, 3])
 
+    # Add light in a frame.
+    light_in_frame = framea0.add_light(cutoff=10)
+    self.assertEqual(light_in_frame.cutoff, 10)
+    self.assertEqual(light_in_frame.parent, framea0.parent)
+    self.assertEqual(light_in_frame.frame, framea0)
+
     # Invalid input for valid keyword argument.
     with self.assertRaises(ValueError) as cm:
       body.add_geom(pos='pos')