Add a way to replace collision meshes upon MJCF conversion

mujoco_ros2_control/docs/tools.rst

@@ -148,6 +148,10 @@ Rough outline of the automated conversion process
     correct per link even when MuJoCo fuses fixed-jointed bodies together. A mesh shared by a visual
     and a collision is converted once and reused for both (and, if that link requests decomposition,
     the whole mesh renders while its decomposed pieces collide).
+  - A link named in a ``replace_collision`` input skips all of the above: its collision geometry
+    (authored or synthesized) is dropped entirely and never converted, and a user-authored MJCF
+    fragment (one or more primitives, optionally grouped under nested sub-bodies) is inserted in
+    its place. See the :ref:`replace_collision attribute reference <threshold-attribute>`.
   - Both kinds of geom get explicit attributes written directly onto them, so visual/collision
     separation does not depend on a user-supplied ``<default class="...">`` block:
 
@@ -231,6 +235,12 @@ Main sub-elements
      For ``geom`` type, elements are identified by ``mesh`` and ``class`` instead of ``name``.
      Useful to tweak physics properties like ``frictionloss``, ``stiffness``, ``damping``,
      ``gravcomp``, etc.
+   - ``replace_collision`` (attribute: ``link``; children: one or more ``geom``/``body`` elements)
+     — drops **all** of the named link's collision geometry (authored or synthesized from its
+     visual) and replaces it with the given MJCF fragment, copied verbatim into that link's
+     ``<body>``. Useful for swapping an expensive collision mesh for cheap primitives (for
+     example capsules), including several primitives, or grouping them under nested named
+     sub-bodies. See the :ref:`replace_collision attribute reference <threshold-attribute>` below.
 
 
 
@@ -381,3 +391,36 @@ attributes the demo converter recognizes; converters may extend this list.
 - Additional attributes: any MJCF attributes you want to set or overwrite (for example ``frictionloss``, ``damping``, ``gravcomp``, ``solimp``, ``solref``, ...).
 - Example: ``<modify_element type="joint" name="joint1" frictionloss="1.0" damping="2.0"/>``.
 - Example: ``<modify_element type="geom" mesh="link1_mesh" class="collision" friction = "0.1 0.005 0.0001"/>``
+
+``replace_collision``
+
+- Required: ``link`` (string) — the URDF link name whose collision geometry should be replaced.
+- Required: one or more child ``geom`` and/or ``body`` elements, using plain MJCF attribute
+  syntax (``type``, ``size``, ``pos``/``quat``, or ``fromto`` for capsules/cylinders). A nested
+  ``body`` is a fixed sub-body used to group its own geoms under a name; it is copied verbatim,
+  including any ``joint`` it contains.
+- Every child ``geom`` (including those inside a nested ``body``) is copied into the link's
+  MJCF ``<body>`` exactly as written; the only attribute the converter injects is
+  ``class="collision"``, and only when the ``geom`` does not already specify a ``class``.
+- All of the link's original collision geometry (authored ``<collision>`` tags, or one
+  synthesized from its ``<visual>``) is dropped and never converted, regardless of
+  ``--use_collision_tags``.
+- Example — replace a link's collision with two capsules:
+
+  .. code-block:: xml
+
+     <replace_collision link="forearm">
+       <geom type="capsule" fromto="0 0 0 0 0 0.3" size="0.04"/>
+       <geom type="capsule" fromto="0 0 0.3 0.1 0 0.3" size="0.03"/>
+     </replace_collision>
+
+- Example — group geoms under a named nested sub-body:
+
+  .. code-block:: xml
+
+     <replace_collision link="leg_left_1_link">
+       <geom name="left_foot6_collision" class="foot_capsule" fromto="-0.0985 0.007 0 0.122 0.0065 0"/>
+       <body name="leg_left_ankle_link">
+         <geom name="leg_left_ankle_collision" class="collision" size="0.03" fromto="-0.04 0 0.03 0.02 0 0.03"/>
+       </body>
+     </replace_collision>

mujoco_ros2_control/mujoco_ros2_control/__init__.py

@@ -16,6 +16,7 @@
     add_mujoco_info,
     remove_tag,
     add_missing_collisions,
+    add_replaced_collisions,
     extract_mesh_info,
     replace_package_names,
     get_images_from_dae,
@@ -57,6 +58,7 @@
     "add_mujoco_info",
     "remove_tag",
     "add_missing_collisions",
+    "add_replaced_collisions",
     "extract_mesh_info",
     "replace_package_names",
     "get_images_from_dae",

mujoco_ros2_control/mujoco_ros2_control/urdf_to_mujoco_utils.py

@@ -113,7 +113,7 @@ def remove_tag(xml_string, tag_to_remove):
     return xmldoc.toprettyxml()
 
 
-def add_missing_collisions(xml_string):
+def add_missing_collisions(xml_string, exclude_links=None):
     """
     Ensures every link that can be rendered can also collide, while respecting any
     collision geometry the URDF author already provided.
@@ -129,14 +129,25 @@ def add_missing_collisions(xml_string):
     when present, copied from the visual otherwise) is used for physics.
 
     :param xml_string: the URDF as a string
+    :param exclude_links: optional iterable of link names to leave collision-free - e.g.
+        links whose collision will be replaced by a user-authored fragment (the
+        replace_collision processed_inputs tag). Any collision already present on such a
+        link (authored or otherwise) is stripped, and none is synthesized from its
+        visuals, so its original collision mesh never enters mesh_info_dict.
     :returns: the URDF string with synthesized collisions added where they were missing
     """
+    exclude_links = set(exclude_links) if exclude_links else set()
     dom = minidom.parseString(xml_string)
 
     for link in dom.getElementsByTagName("link"):
         visuals = [c for c in link.childNodes if c.nodeType == c.ELEMENT_NODE and c.tagName == "visual"]
         collisions = [c for c in link.childNodes if c.nodeType == c.ELEMENT_NODE and c.tagName == "collision"]
 
+        if link.getAttribute("name") in exclude_links:
+            for collision in collisions:
+                link.removeChild(collision)
+            continue
+
         # Respect authored collisions and skip links with nothing to render.
         if collisions or not visuals:
             continue
@@ -949,9 +960,10 @@ def get_processed_mujoco_inputs(processed_inputs_element):
     cameras_dict = dict()
     modify_element_dict = dict()
     lidar_dict = dict()
+    replace_collision_dict = dict()
 
     if not processed_inputs_element:
-        return decompose_dict, cameras_dict, modify_element_dict, lidar_dict
+        return decompose_dict, cameras_dict, modify_element_dict, lidar_dict, replace_collision_dict
 
     for child in processed_inputs_element.childNodes:
         if child.nodeType != child.ELEMENT_NODE:
@@ -1049,7 +1061,25 @@ def get_processed_mujoco_inputs(processed_inputs_element):
             for key_attr, value in attr_dict.items():
                 print(f"  {key_attr}: {value}")
 
-    return decompose_dict, cameras_dict, modify_element_dict, lidar_dict
+        # Grab collision-replacement fragments
+        if child.tagName == "replace_collision":
+            link_name = child.getAttribute("link")
+            if not link_name:
+                raise ValueError("'link' must be in the attributes of a 'replace_collision' tag!")
+            if link_name in replace_collision_dict:
+                raise ValueError(f"Multiple 'replace_collision' tags found for link '{link_name}'")
+
+            fragment = [c for c in child.childNodes if c.nodeType == c.ELEMENT_NODE]
+            if not fragment:
+                raise ValueError(
+                    f"'replace_collision' tag for link '{link_name}' must contain at least one "
+                    "child element (geom or body)!"
+                )
+
+            replace_collision_dict[link_name] = fragment
+            print(f"Will replace collision(s) on link '{link_name}' with {len(fragment)} element(s)")
+
+    return decompose_dict, cameras_dict, modify_element_dict, lidar_dict, replace_collision_dict
 
 
 def parse_inputs_xml(filename=None):
@@ -1504,6 +1534,47 @@ def add_lidar_from_sites(dom, lidar_dict):
     return dom
 
 
+def add_replaced_collisions(dom, replace_collision_dict):
+    """
+    Inserts each link's replace_collision fragment (one or more <geom>/<body> elements,
+    parsed by get_processed_mujoco_inputs) into that link's <body> in the MJCF, verbatim.
+
+    The fragment's elements are deep-imported as-is (name/type/size/fromto/pos/nested
+    <body>/... all preserved exactly as authored). The one default applied: any <geom> in
+    the fragment that does not already carry a class attribute is given class="collision",
+    so it still picks up sane group/contype/conaffinity defaults.
+
+    Raises ValueError if a link named in replace_collision_dict has no matching <body>.
+    """
+    if not replace_collision_dict:
+        return dom
+
+    matched_links = set()
+
+    for body in dom.getElementsByTagName("body"):
+        link_name = body.getAttribute("name")
+        fragment = replace_collision_dict.get(link_name)
+        if fragment is None:
+            continue
+        matched_links.add(link_name)
+
+        for element in fragment:
+            imported = dom.importNode(element, True)
+            body.appendChild(imported)
+
+            # a <geom> can't have child geoms, so these two cases are mutually exclusive
+            geoms = [imported] if imported.tagName == "geom" else imported.getElementsByTagName("geom")
+            for geom in geoms:
+                if not geom.hasAttribute("class"):
+                    geom.setAttribute("class", "collision")
+
+    unmatched = set(replace_collision_dict.keys()) - matched_links
+    if unmatched:
+        raise ValueError(f"replace_collision link(s) not found in the MJCF: {', '.join(sorted(unmatched))}")
+
+    return dom
+
+
 def add_modifiers(dom, modify_element_dict):
     """
     Modify elements that are a part of the worldbody tag by adding attributes.

mujoco_ros2_control/scripts/make_mjcf_from_robot_description.py

@@ -308,6 +308,7 @@ def fix_mujoco_description(
     cameras_dict,
     modify_element_dict,
     lidar_dict,
+    replace_collision_dict,
     request_add_free_joint,
 ):
     """
@@ -334,6 +335,9 @@ def fix_mujoco_description(
     dom = mrc.update_obj_assets(dom, output_filepath, mesh_info_dict)
     dom = mrc.update_non_obj_assets(dom, output_filepath, mesh_info_dict)
 
+    # Insert any user-authored collision replacement fragments (replace_collision tags)
+    dom = mrc.add_replaced_collisions(dom, replace_collision_dict)
+
     # Add the MuJoCo input elements
     dom = mrc.add_mujoco_inputs(dom, raw_inputs, scene_inputs)
 
@@ -488,7 +492,9 @@ def main(args=None):
         mrc.write_mujoco_scene(scene_inputs, output_filepath)
         scene_inputs = None
 
-    decompose_dict, cameras_dict, modify_element_dict, lidar_dict = mrc.get_processed_mujoco_inputs(processed_inputs)
+    decompose_dict, cameras_dict, modify_element_dict, lidar_dict, replace_collision_dict = (
+        mrc.get_processed_mujoco_inputs(processed_inputs)
+    )
 
     if parsed_args.asset_dir:
         assets_filepath = parsed_args.asset_dir
@@ -514,7 +520,12 @@ def main(args=None):
     if not parsed_args.use_collision_tags:
         xml_data = mrc.remove_tag(xml_data, "collision")
 
-    xml_data = mrc.add_missing_collisions(xml_data)
+    # Links with a replace_collision input get their collision(s) fully replaced by a
+    # user-authored fragment (see add_replaced_collisions), so they are excluded here:
+    # any collision they ended up with (authored or otherwise) is dropped and none is
+    # synthesized from their visuals, keeping their original collision mesh out of
+    # mesh_info_dict entirely.
+    xml_data = mrc.add_missing_collisions(xml_data, exclude_links=replace_collision_dict.keys())
 
     xml_data = mrc.replace_package_names(xml_data)
     mesh_info_dict, xml_data = mrc.extract_mesh_info(xml_data, parsed_args.asset_dir, decompose_dict)
@@ -541,6 +552,7 @@ def main(args=None):
         cameras_dict,
         modify_element_dict,
         lidar_dict,
+        replace_collision_dict,
         request_add_free_joint,
     )