Repository: gazebosim/gz-mujoco (language: python)

Issue:
MJCF to SDF: Added geometry defaults
Signed-off-by: ahcorde <ahcorde@gmail.com>
# 🎉 New feature

## Summary

You can define some defaul values in mujoco:

```xml
   <default>
     <geom type="capsule" mass="1"/>
   </default>
```

This PR allows to read these values and set them

## Checklist
- [x] Signed all commits for DCO
- [ ] Added tests
- [ ] Added example and/or tutorial
- [ ] Updated documentation (as needed)
- [ ] Updated migration guide (as needed)
- [ ] Consider updating Python bindings (if the library has them)
- [x] `codecheck` passed (See [contributing](https://ignitionrobotics.org/docs/all/contributing#contributing-code))
- [x] All tests passed (See [test coverage](https://ignitionrobotics.org/docs/all/contributing#test-coverage))
- [x] While waiting for a review on your PR, please help review [another open pull request](https://github.com/pulls?q=is%3Aopen+is%3Apr+user%3Aignitionrobotics+repo%3Aosrf%2Fsdformat+archived%3Afalse+) to support the maintainers

**Note to maintainers**: Remember to use **Squash-Merge** and edit the commit message to match the pull request summary while retaining `Signed-off-by` messages.



Discussion / hints:

Fix patch:
diff --git a/mjcf_to_sdformat/mjcf_to_sdformat/converters/link.py b/mjcf_to_sdformat/mjcf_to_sdformat/converters/link.py
index 447bc0c..48fcb7d 100644
--- a/mjcf_to_sdformat/mjcf_to_sdformat/converters/link.py
+++ b/mjcf_to_sdformat/mjcf_to_sdformat/converters/link.py
@@ -27,62 +27,15 @@
 VISUAL_GEOM_GROUP = 0
 
 
-def _set_attribute(geom, key, value):
-    """
-    Set attribute to a MJCF Element
-    :param mjcf.Element geom: The MJCF geom to set the attributes
-    :param str key: Key name
-    :param str value: Value
-    :return: The modfied MJCF geometry.
-    :rtype: mjcf.Element
-    """
-    try:
-        geom.get_attributes()[key]
-    except KeyError:
-        geom.set_attributes(**{key: value})
-    return geom
-
-
-def _set_defaults(geom, default_classes=None):
-    """
-    Set default values to the MCJF geom
-    :param mjcf.Element geom: The MJCF geom
-    :param list[mjcf.Default] default_classes: List of default classes setted
-    to the body and parent bodies
-    :return: The modfied MJCF geometry.
-    :rtype: mjcf.Element
-    """
-    if geom.dclass is not None:
-        if geom.dclass.geom is not None:
-            print(geom.dclass.geom)
-
-    if geom.dclass is not None:
-        if geom.dclass.geom is not None:
-            geom.set_attributes(**geom.dclass.geom.get_attributes())
-    elif default_classes is not None:
-        for default_class in default_classes:
-            if default_class.geom is not None:
-                for k, v in default_class.geom.get_attributes().items():
-                    geom = _set_attribute(geom, k, v)
-    elif geom.root.default.geom is not None:
-        for k, v in geom.root.default.geom.get_attributes().items():
-            geom = _set_attribute(geom, k, v)
-
-    _set_attribute(geom, "type", "sphere")
-
-    return geom
-
-
-def mjcf_body_to_sdf(body, physics, body_parent_name=None,
-                     default_classes=None):
+def mjcf_body_to_sdf(body, physics, body_parent_name=None, modifiers=None):
     """
     Converts an MJCF body to a SDFormat.
 
     :param mjcf.Element body: The MJCF body
     :param mujoco.Physics physics: Mujoco Physics
     :param mjcf.Element inertial: Inertial of the body
-    :param list[mjcf.Default] default_classes: List of default classes set
-    to the body and parent bodies
+    :param sdformat_mjcf_utils.MjcfModifiers modifiers: Modifiers that apply
+    default classes to elements.
     :return: The newly created SDFormat link.
     :rtype: sdf.Link
     """
@@ -218,7 +171,8 @@ def set_collision(geom):
             link.add_collision(col)
 
     for geom in body.geom:
-        geom = _set_defaults(geom, default_classes)
+        if modifiers is not None:
+            modifiers.apply_modifiers_to_element(geom)
         # If the group is not defined then visual and collision is added
         if geom.group is None:
             set_visual(geom)
@@ -229,6 +183,8 @@ def set_collision(geom):
             set_collision(geom)
 
     for light in body.light:
+        if modifiers is not None:
+            modifiers.apply_modifiers_to_element(light)
         light_sdf = mjcf_light_to_sdf(light)
         link.add_light(light_sdf)
     return link
diff --git a/mjcf_to_sdformat/mjcf_to_sdformat/converters/world.py b/mjcf_to_sdformat/mjcf_to_sdformat/converters/world.py
index c50f161..7f40f8a 100644
--- a/mjcf_to_sdformat/mjcf_to_sdformat/converters/world.py
+++ b/mjcf_to_sdformat/mjcf_to_sdformat/converters/world.py
@@ -18,6 +18,7 @@
 from mjcf_to_sdformat.converters.link import mjcf_body_to_sdf
 
 import sdformat_mjcf_utils.sdf_utils as su
+from sdformat_mjcf_utils.defaults import MjcfModifiers
 
 import sdformat as sdf
 
@@ -36,13 +37,16 @@ def mjcf_worldbody_to_sdf(mjcf_root, physics, world):
     else:
         model.set_name("model")
 
+    modifiers = MjcfModifiers(mjcf_root)
+
     model_static = sdf.Model()
 
     for light in mjcf_root.worldbody.light:
+        modifiers.apply_modifiers_to_element(light)
         light_sdf = mjcf_light_to_sdf(light)
         world.add_light(light_sdf)
 
-    link = mjcf_body_to_sdf(mjcf_root.worldbody, physics)
+    link = mjcf_body_to_sdf(mjcf_root.worldbody, physics, modifiers=modifiers)
     model_static.set_name(su.find_unique_name(
         mjcf_root.worldbody, "geom", "static"))
     model_static.add_link(link)
@@ -65,22 +69,16 @@ def mjcf_worldbody_to_sdf(mjcf_root, physics, world):
 
     def iterate_bodies(input_body,
                        model,
-                       body_parent_name=None,
-                       default_classes=None):
+                       body_parent_name=None):
         for body in input_body:
-            if body.childclass is not None:
-                if default_classes is None:
-                    default_classes = []
-                default_classes.append(body.childclass)
             link = mjcf_body_to_sdf(body,
                                     physics,
                                     body_parent_name=body_parent_name,
-                                    default_classes=default_classes)
+                                    modifiers=modifiers)
             model.add_link(link)
             iterate_bodies(body.body,
                            model,
-                           body.name,
-                           default_classes=default_classes)
+                           body.name)
     iterate_bodies(body, model)
 
     world.add_model(model)
diff --git a/sdformat_mjcf_utils/sdformat_mjcf_utils/defaults.py b/sdformat_mjcf_utils/sdformat_mjcf_utils/defaults.py
new file mode 100644
index 0000000..0bfc76e
--- /dev/null
+++ b/sdformat_mjcf_utils/sdformat_mjcf_utils/defaults.py
@@ -0,0 +1,83 @@
+# Copyright (C) 2022 Open Source Robotics Foundation
+#
+# Licensed under the Apache License, version 2.0 (the "License")
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dm_control import mjcf
+# Based on https://github.com/deepmind/mujoco/blob/d7ae7f54338a8ea02c34f4a969722182481eb72b/unity/Runtime/Importer/MjXmlModifiers.cs  # noqa
+
+
+class MjcfModifiers:
+    """Process default classes and their inheritance and apply the resulting
+    default class to a given element."""
+
+    def __init__(self, root):
+        """
+        Initialize the class with the root element. Note that this class
+        mutates the default element (root.default) and its descendants in order
+        to processes the inheritance hierarchy.
+        :param mjcf.RootElement root: The MJCF root element
+        """
+        self.root = root
+        self._apply_inherited_default(self.root.default, None)
+
+    def apply_modifiers_to_element(self, elem):
+        """
+        Given an element, e.g., geom, body, etc, this copies into the element
+        the attributes of the default class associated with the element.
+        :param mjcf.Element elem: The MJCF element to process.
+        """
+        dclass = self._get_default_class(elem)
+        def_elem = dclass.get_children(elem.tag)
+        self._copy_attributes(def_elem, elem)
+
+    def _get_default_class(self, elem):
+        if elem.dclass is not None:
+            return elem.dclass
+
+        cur_elem = elem.parent
+        while cur_elem is not None:
+            if hasattr(cur_elem,
+                       "childclass") and cur_elem.childclass is not None:
+                return cur_elem.childclass
+            else:
+                cur_elem = cur_elem.parent
+
+        # No childclass applies, so use the main default class
+        return elem.root.default
+
+    def _apply_inherited_default(self, child_def, parent_def):
+        if parent_def is not None:
+            for elem in child_def.all_children():
+                if elem.tag != "default":
+                    parent_elem = parent_def.get_children(elem.tag)
+                    self._copy_attributes(parent_elem, elem)
+        for grand_child_def in child_def.default:
+            self._apply_inherited_default(grand_child_def, child_def)
+
+    def _copy_attributes(self, from_elem, to_elem):
+        to_attrs = to_elem.get_attributes().keys()
+        for attr, val in from_elem.get_attributes().items():
+            if attr not in to_attrs:
+                to_elem.set_attributes(**{attr: val})
+
+
+# Applies the defaults on all geoms and prints the resulting xml. Only used for
+# degugging purposes.
+if __name__ == "__main__":
+    import sys
+    model = mjcf.from_path(sys.argv[1])
+    modifiers = MjcfModifiers(model)
+    for geom in model.find_all("geom"):
+        modifiers.apply_modifiers_to_element(geom)
+    model.default.remove()
+    print(model.to_xml_string())


Test patch:
diff --git a/mjcf_to_sdformat/tests/resources/test_defaults_inheritance.xml b/mjcf_to_sdformat/tests/resources/test_defaults_inheritance.xml
new file mode 100644
index 0000000..98b2ff4
--- /dev/null
+++ b/mjcf_to_sdformat/tests/resources/test_defaults_inheritance.xml
@@ -0,0 +1,25 @@
+<mujoco>
+    <default class="main">
+        <geom rgba="0 1 0 1" type="sphere" size="2"/>
+        <default class="sub">
+            <geom type="cylinder" size="2 3"/>
+        </default>
+        <default class="small_shape">
+            <geom size="0.5 0.5 0.5" />
+        </default>
+    </default>
+
+    <worldbody>
+        <body name="body1" childclass="sub">
+            <geom name="shape1" />
+            <geom name="shape2" pos="5 0 0" class="main"/>
+            <geom name="shape3" pos="10 0 0" class="small_shape"/>
+        </body>
+        <body name="body2" pos="0 -5 0">
+            <geom name="shape4" />
+            <geom name="shape5" pos="5 0 0" class="main"/>
+            <geom name="shape6" pos="10 0 0" class="small_shape"/>
+        </body>
+    </worldbody>
+</mujoco>
+
diff --git a/mjcf_to_sdformat/tests/test_defaults.py b/mjcf_to_sdformat/tests/test_defaults.py
index acc14a3..0a6e97e 100644
--- a/mjcf_to_sdformat/tests/test_defaults.py
+++ b/mjcf_to_sdformat/tests/test_defaults.py
@@ -58,10 +58,81 @@ def test_defaults(self):
         link = model.link_by_index(1)
         self.assertNotEqual(None, link)
         visual = link.visual_by_index(0)
-        self.assertEqual(sdf.GeometryType.SPHERE, visual.geometry().type())
-        sphere_shape = visual.geometry().sphere_shape()
-        self.assertNotEqual(None, sphere_shape)
-        self.assertEqual(0.01, sphere_shape.radius())
+        self.assertNotEqual(None, link)
+        self.assertEqual(sdf.GeometryType.CAPSULE, visual.geometry().type())
+        capsule_shape2 = visual.geometry().capsule_shape()
+        self.assertNotEqual(None, capsule_shape2)
+        self.assertEqual(0.6, capsule_shape2.length())
+        self.assertEqual(0.07, capsule_shape2.radius())
+
+    def test_inheritance(self):
+        filename = str(TEST_RESOURCES_DIR / "test_defaults_inheritance.xml")
+        mjcf_model = mjcf.from_path(filename)
+        self.assertIsNotNone(mjcf_model)
+        physics = mjcf.Physics.from_xml_path(filename)
+        self.assertIsNotNone(physics)
+
+        world = sdf.World()
+        world.set_name("default")
+
+        mjcf_worldbody_to_sdf(mjcf_model, physics, world)
+        root = sdf.Root()
+        root.add_world(world)
+        model = world.model_by_index(1)
+        self.assertNotEqual(None, model)
+        self.assertEqual(2, model.link_count())
+
+        link1 = model.link_by_name("body1")
+        self.assertIsNotNone(link1)
+
+        # shape1's class is "sub" based on the parent body's childclass
+        visual_shape1 = link1.visual_by_name("visual_shape1")
+        self.assertEqual(sdf.GeometryType.CYLINDER,
+                         visual_shape1.geometry().type())
+        shape1_cylinder = visual_shape1.geometry().cylinder_shape()
+        self.assertEqual(6, shape1_cylinder.length())
+        self.assertEqual(2, shape1_cylinder.radius())
+
+        # shape2's class is "main" based shape2's own class regardles of its
+        # parent body's childclass.
+        visual_shape2 = link1.visual_by_name("visual_shape2")
+        self.assertEqual(sdf.GeometryType.SPHERE,
+                         visual_shape2.geometry().type())
+        shape2_sphere = visual_shape2.geometry().sphere_shape()
+        self.assertEqual(2, shape2_sphere.radius())
+
+        # shape3's class is "small_shape" based shape3's own class regardless
+        # of its parent body's childclass.
+        visual_shape3 = link1.visual_by_name("visual_shape3")
+        self.assertEqual(sdf.GeometryType.SPHERE,
+                         visual_shape3.geometry().type())
+        shape3_sphere = visual_shape3.geometry().sphere_shape()
+        self.assertEqual(0.5, shape3_sphere.radius())
+
+        link2 = model.link_by_name("body2")
+        self.assertIsNotNone(link2)
+
+        # shape4's class is "main" since it doesn't have a parent body with a
+        # childclass nor its own class.
+        visual_shape4 = link2.visual_by_name("visual_shape4")
+        self.assertEqual(sdf.GeometryType.SPHERE,
+                         visual_shape4.geometry().type())
+        shape4_cylinder = visual_shape4.geometry().sphere_shape()
+        self.assertEqual(2, shape4_cylinder.radius())
+
+        # shape5's class is "main" based shape5's own class.
+        visual_shape5 = link2.visual_by_name("visual_shape5")
+        self.assertEqual(sdf.GeometryType.SPHERE,
+                         visual_shape5.geometry().type())
+        shape5_sphere = visual_shape5.geometry().sphere_shape()
+        self.assertEqual(2, shape5_sphere.radius())
+
+        # shape3's class is "small_shape" based shape3's own class.
+        visual_shape6 = link2.visual_by_name("visual_shape6")
+        self.assertEqual(sdf.GeometryType.SPHERE,
+                         visual_shape6.geometry().type())
+        shape6_sphere = visual_shape6.geometry().sphere_shape()
+        self.assertEqual(0.5, shape6_sphere.radius())
 
 
 if __name__ == "__main__":