diff --git a/mujoco_urdf_loader/loader.py b/mujoco_urdf_loader/loader.py index 903cc49..94f9817 100644 --- a/mujoco_urdf_loader/loader.py +++ b/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): """ diff --git a/tests/test_urdf_to_mujoco_loader_class.py b/tests/test_urdf_to_mujoco_loader_class.py index ffe1f94..942435f 100644 --- a/tests/test_urdf_to_mujoco_loader_class.py +++ b/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)