Merge pull request #9 from ami-iit/fix-model-generation

Fix mass reducing after conversion

mujoco_urdf_loader/loader.py

@@ -135,20 +135,21 @@ def connect_root_to_world(root):
         parent_links = {joint.find("parent").attrib["link"] for joint in joints}
         # The root link is a parent link that is not a child link
         root_link = next(link for link in links if link not in child_links)
-
-        # Add a floating joint that connects the root link to the world
-        floating_joint = ET.Element("joint", attrib={
-            "name": f"{root_link}_floating_joint",
-            "type": "floating"
-        })
-        # Populate the floating joint
-        parent_element = ET.SubElement(floating_joint, "parent", attrib={"link": "world"})
-        child_element = ET.SubElement(floating_joint, "child", attrib={"link": root_link})
-        # Add the floating joint to the urdf root
-        root.insert(0, floating_joint)
-        # Add a link called "world"
-        world_link = ET.Element("link", attrib={"name": "world"})
-        root.insert(0, world_link)
+        # Find the joint that has the root link as parent
+        fixed_joint_found = False
+        for joint in joints:
+            parent_link = joint.find("parent").attrib["link"]
+            if parent_link == root_link:
+                # Check if the joint is fixed
+                if joint.attrib["type"] == "fixed":
+                    # Change the joint type to floating
+                    joint.attrib["type"] = "floating"
+                    print(f"Modified joint {joint.attrib['name']} to type floating.")
+                    fixed_joint_found = True
+                break
+
+        if not fixed_joint_found:
+            raise ValueError("No fixed joint found that can be modified to floating.")
 
     def set_control_mode(self, joint: Union[str, List[str]], mode: ControlMode):
         """

tests/test_urdf_to_mujoco_loader_class.py

@@ -1,55 +1,120 @@
 from xml.etree import ElementTree as ET
 
+import idyntree.bindings as idyntree
+import mujoco
+import pytest
 import resolve_robotics_uri_py as rru
 
+# Import the ANYmal robot description
+from robot_descriptions import anymal_c_description
+
 from mujoco_urdf_loader import ControlMode, URDFtoMuJoCoLoader, URDFtoMuJoCoLoaderCfg
 from mujoco_urdf_loader.urdf_fcn import get_mesh_path
 
+# Define robot configurations in a dictionary
+ROBOTS = {
+    "ergoCub": {
+        "controlled_joints": [
+            "l_hip_pitch",
+            "r_hip_pitch",
+            "torso_roll",
+            "l_hip_roll",
+            "r_hip_roll",
+            "torso_pitch",
+            "torso_yaw",
+            "l_hip_yaw",
+            "r_hip_yaw",
+            "l_shoulder_pitch",
+            "neck_pitch",
+            "r_shoulder_pitch",
+            "l_knee",
+            "r_knee",
+            "l_shoulder_roll",
+            "neck_roll",
+            "r_shoulder_roll",
+            "l_ankle_pitch",
+            "r_ankle_pitch",
+            "neck_yaw",
+            "l_ankle_roll",
+            "r_ankle_roll",
+            "l_shoulder_yaw",
+            "r_shoulder_yaw",
+            "l_elbow",
+            "r_elbow",
+        ],
+        "urdf_path": str(
+            rru.resolve_robotics_uri("package://ergoCub/robots/ergoCubSN002/model.urdf")
+        ),
+        "mesh_path": None,  # Will be set later
+    },
+    "ANYmal": {
+        "controlled_joints": [
+            "LF_HAA",
+            "LF_HFE",
+            "LF_KFE",
+            "RF_HAA",
+            "RF_HFE",
+            "RF_KFE",
+            "LH_HAA",
+            "LH_HFE",
+            "LH_KFE",
+            "RH_HAA",
+            "RH_HFE",
+            "RH_KFE",
+        ],
+        "urdf_path": anymal_c_description.URDF_PATH,
+        "mesh_path": f"{anymal_c_description.PACKAGE_PATH}/meshes",
+    },
+}
+
+# Update mesh paths for robots that need it
+for robot in ROBOTS.values():
+    if robot["mesh_path"] is None:
+        robot["mesh_path"] = get_mesh_path(ET.parse(robot["urdf_path"]).getroot())
+
 
-def test_load_urdf_into_mjcf():
-
-    controlled_joints = [
-        "l_hip_pitch",
-        "r_hip_pitch",
-        "torso_roll",
-        "l_hip_roll",
-        "r_hip_roll",
-        "torso_pitch",
-        "torso_yaw",
-        "l_hip_yaw",
-        "r_hip_yaw",
-        "l_shoulder_pitch",
-        "neck_pitch",
-        "r_shoulder_pitch",
-        "l_knee",
-        "r_knee",
-        "l_shoulder_roll",
-        "neck_roll",
-        "r_shoulder_roll",
-        "l_ankle_pitch",
-        "r_ankle_pitch",
-        "neck_yaw",
-        "l_ankle_roll",
-        "r_ankle_roll",
-        "l_shoulder_yaw",
-        "r_shoulder_yaw",
-        "l_elbow",
-        "r_elbow",
-    ]
+@pytest.mark.parametrize("robot_name", ["ergoCub", "ANYmal"])
+def test_load_urdf_into_mjcf(robot_name):
+    robot = ROBOTS[robot_name]
+    controlled_joints = robot["controlled_joints"]
+    urdf_path = robot["urdf_path"]
+    mesh_path = robot["mesh_path"]
 
     control_modes = [ControlMode.TORQUE] * len(controlled_joints)
     stiffness = [0.0] * len(controlled_joints)
     damping = [0.0] * len(controlled_joints)
-
     cfg = URDFtoMuJoCoLoaderCfg(controlled_joints, control_modes, stiffness, damping)
 
-    urdf_string = str(
-        rru.resolve_robotics_uri("package://ergoCub/robots/ergoCubSN002/model.urdf")
-    )
-    mesh_path = get_mesh_path(ET.parse(urdf_string).getroot())
+    loader = URDFtoMuJoCoLoader.load_urdf(urdf_path, mesh_path, cfg)
+    mjcf = loader.get_mjcf_string()
+
+    assert mjcf is not None
 
-    loader = URDFtoMuJoCoLoader.load_urdf(urdf_string, mesh_path, cfg)
 
+@pytest.mark.parametrize("robot_name", ["ergoCub", "ANYmal"])
+def test_total_mass(robot_name):
+    robot = ROBOTS[robot_name]
+    controlled_joints = robot["controlled_joints"]
+    urdf_path = robot["urdf_path"]
+    mesh_path = robot["mesh_path"]
+
+    control_modes = [ControlMode.TORQUE] * len(controlled_joints)
+    stiffness = [0.0] * len(controlled_joints)
+    damping = [0.0] * len(controlled_joints)
+    cfg = URDFtoMuJoCoLoaderCfg(controlled_joints, control_modes, stiffness, damping)
+
+    loader = URDFtoMuJoCoLoader.load_urdf(urdf_path, mesh_path, cfg)
     mjcf = loader.get_mjcf_string()
 
-    assert mjcf is not None
+    model = mujoco.MjModel.from_xml_string(mjcf)
+    data = mujoco.MjData(model)
+
+    total_mass_mj = sum(model.body_mass)
+
+    # Test with iDynTree
+    model_loader = idyntree.ModelLoader()
+    model_loader.loadReducedModelFromFile(urdf_path, controlled_joints)
+    idt_model = model_loader.model()
+    total_mass_idt = idt_model.getTotalMass()
+
+    assert total_mass_mj == pytest.approx(total_mass_idt, abs=1e-4)