jaxsim 0.2.dev191__py3-none-any.whl → 0.6.1.dev5__py3-none-any.whl

jaxsim/mujoco/loaders.py

@@ -1,19 +1,39 @@
 import pathlib
 import tempfile
 import warnings
+from collections.abc import Sequence
 from typing import Any
 
 import mujoco as mj
+import numpy as np
 import rod.urdf.exporter
 from lxml import etree as ET
 
+from jaxsim import logging
+
+from .utils import MujocoCamera
+
+MujocoCameraType = (
+    MujocoCamera | Sequence[MujocoCamera] | dict[str, str] | Sequence[dict[str, str]]
+)
+
 
 def load_rod_model(
     model_description: str | pathlib.Path | rod.Model,
     is_urdf: bool | None = None,
     model_name: str | None = None,
 ) -> rod.Model:
-    """"""
+    """
+    Load a ROD model from a URDF/SDF file or a ROD model.
+
+    Args:
+        model_description: The URDF/SDF file or ROD model to load.
+        is_urdf: Whether to force parsing the model description as a URDF file.
+        model_name: The name of the model to load from the resource.
+
+    Returns:
+        rod.Model: The loaded ROD model.
+    """
 
     # Parse the SDF resource.
     sdf_element = rod.Sdf.load(sdf=model_description, is_urdf=is_urdf)
@@ -43,14 +63,77 @@ def load_rod_model(
     return models[model_name]
 
 
+class ModelToMjcf:
+    """
+    Class to convert a URDF/SDF file or a ROD model to a Mujoco MJCF string.
+    """
+
+    @staticmethod
+    def convert(
+        model: str | pathlib.Path | rod.Model,
+        considered_joints: list[str] | None = None,
+        plane_normal: tuple[float, float, float] = (0, 0, 1),
+        heightmap: bool | None = None,
+        heightmap_samples_xy: tuple[int, int] = (101, 101),
+        cameras: MujocoCameraType = (),
+    ) -> tuple[str, dict[str, Any]]:
+        """
+        Convert a model to a Mujoco MJCF string.
+
+        Args:
+            model: The URDF/SDF file or ROD model to convert.
+            considered_joints: The list of joint names to consider in the conversion.
+            plane_normal: The normal vector of the plane.
+            heightmap: Whether to generate a heightmap.
+            heightmap_samples_xy: The number of points in the heightmap grid.
+            cameras: The custom cameras to add to the scene.
+
+        Returns:
+            A tuple containing the MJCF string and the dictionary of assets.
+        """
+
+        match model:
+            case rod.Model():
+                rod_model = model
+            case str() | pathlib.Path():
+                # Convert the JaxSim model to a ROD model.
+                rod_model = load_rod_model(
+                    model_description=model,
+                    is_urdf=None,
+                    model_name=None,
+                )
+            case _:
+                raise TypeError(f"Unsupported type for 'model': {type(model)}")
+
+        # Convert the ROD model to MJCF.
+        return RodModelToMjcf.convert(
+            rod_model=rod_model,
+            considered_joints=considered_joints,
+            plane_normal=plane_normal,
+            heightmap=heightmap,
+            heightmap_samples_xy=heightmap_samples_xy,
+            cameras=cameras,
+        )
+
+
 class RodModelToMjcf:
-    """"""
+    """
+    Class to convert a ROD model to a Mujoco MJCF string.
+    """
 
     @staticmethod
     def assets_from_rod_model(
         rod_model: rod.Model,
     ) -> dict[str, bytes]:
-        """"""
+        """
+        Generate a dictionary of assets from a ROD model.
+
+        Args:
+            rod_model: The ROD model to extract the assets from.
+
+        Returns:
+            dict: A dictionary of assets.
+        """
 
         import resolve_robotics_uri_py
 
@@ -85,7 +168,17 @@ class RodModelToMjcf:
         base_link_name: str,
         floating_joint_name: str = "world_to_base",
     ) -> str:
-        """"""
+        """
+        Add a floating joint to a URDF string.
+
+        Args:
+            urdf_string: The URDF string to modify.
+            base_link_name: The name of the base link to attach the floating joint.
+            floating_joint_name: The name of the floating joint to add.
+
+        Returns:
+            str: The modified URDF string.
+        """
 
         with tempfile.NamedTemporaryFile(mode="w+", suffix=".urdf") as urdf_file:
 
@@ -102,7 +195,7 @@ class RodModelToMjcf:
 
         if root.find(f".//joint[@name='{floating_joint_name}']") is not None:
             msg = f"The URDF already has a floating joint '{floating_joint_name}'"
-            warnings.warn(msg)
+            warnings.warn(msg, stacklevel=2)
             return ET.tostring(root, pretty_print=True).decode()
 
         # Create the "world" link if it doesn't exist.
@@ -129,8 +222,25 @@ class RodModelToMjcf:
     def convert(
         rod_model: rod.Model,
         considered_joints: list[str] | None = None,
+        plane_normal: tuple[float, float, float] = (0, 0, 1),
+        heightmap: bool | None = None,
+        heightmap_samples_xy: tuple[int, int] = (101, 101),
+        cameras: MujocoCameraType = (),
     ) -> tuple[str, dict[str, Any]]:
-        """"""
+        """
+        Convert a ROD model to a Mujoco MJCF string.
+
+        Args:
+            rod_model: The ROD model to convert.
+            considered_joints: The list of joint names to consider in the conversion.
+            plane_normal: The normal vector of the plane.
+            heightmap: Whether to generate a heightmap.
+            heightmap_samples_xy: The number of points in the heightmap grid.
+            cameras: The custom cameras to add to the scene.
+
+        Returns:
+            A tuple containing the MJCF string and the dictionary of assets.
+        """
 
         # -------------------------------------
         # Convert the model description to URDF
@@ -144,10 +254,9 @@ class RodModelToMjcf:
         )
 
         # If considered joints are passed, make sure that they are all part of the model.
-        if considered_joints - set([j.name for j in rod_model.joints()]):
-            extra_joints = set(considered_joints) - set(
-                [j.name for j in rod_model.joints()]
-            )
+        if considered_joints - {j.name for j in rod_model.joints()}:
+            extra_joints = set(considered_joints) - {j.name for j in rod_model.joints()}
+
             msg = f"Couldn't find the following joints in the model: '{extra_joints}'"
             raise ValueError(msg)
 
@@ -155,14 +264,14 @@ class RodModelToMjcf:
         joints_dict = {j.name: j for j in rod_model.joints()}
 
         # Convert all the joints not considered to fixed joints.
-        for joint_name in set(j.name for j in rod_model.joints()) - considered_joints:
+        for joint_name in {j.name for j in rod_model.joints()} - considered_joints:
             joints_dict[joint_name].type = "fixed"
 
         # Convert the ROD model to URDF.
-        urdf_string = rod.urdf.exporter.UrdfExporter.sdf_to_urdf_string(
+        urdf_string = rod.urdf.exporter.UrdfExporter(
+            gazebo_preserve_fixed_joints=False, pretty=True
+        ).to_urdf_string(
             sdf=rod.Sdf(model=rod_model, version="1.7"),
-            gazebo_preserve_fixed_joints=False,
-            pretty=True,
         )
 
         # -------------------------------------
@@ -198,8 +307,6 @@ class RodModelToMjcf:
         )
 
         urdf_string = ET.tostring(root, pretty_print=True).decode()
-        # print(urdf_string)
-        # raise
 
         # ------------------------------
         # Post-process all dummy visuals
@@ -207,7 +314,6 @@ class RodModelToMjcf:
 
         parser = ET.XMLParser(remove_blank_text=True)
         root: ET._Element = ET.fromstring(text=urdf_string.encode(), parser=parser)
-        import numpy as np
 
         # Give a tiny radius to all dummy spheres
         for geometry in root.findall(".//visual/geometry[sphere]"):
@@ -233,13 +339,13 @@ class RodModelToMjcf:
 
         # Load the URDF model into Mujoco.
         assets = RodModelToMjcf.assets_from_rod_model(rod_model=rod_model)
-        mj_model = mj.MjModel.from_xml_string(xml=urdf_string, assets=assets)  # noqa
+        mj_model = mj.MjModel.from_xml_string(xml=urdf_string, assets=assets)
 
         # Get the joint names.
-        mj_joint_names = set(
+        mj_joint_names = {
             mj.mj_id2name(mj_model, mj.mjtObj.mjOBJ_JOINT, idx)
             for idx in range(mj_model.njnt)
-        )
+        }
 
         # Check that the Mujoco model only has the considered joints.
         if mj_joint_names != considered_joints:
@@ -265,7 +371,7 @@ class RodModelToMjcf:
         root: ET._Element = tree.getroot()
 
         # Find the <mujoco> element (might be the root itself).
-        mujoco_element: ET._Element = list(root.iter("mujoco"))[0]
+        mujoco_element: ET._Element = next(iter(root.iter("mujoco")))
 
         # --------------
         # Add the motors
@@ -310,7 +416,7 @@ class RodModelToMjcf:
         # Set alpha=0 to the color of all collision elements
         for geometry_element in mujoco_element.findall(".//geom[@rgba]"):
             if geometry_element.attrib.get("name") in collision_names:
-                r, g, b, a = geometry_element.attrib["rgba"].split(" ")
+                r, g, b, _ = geometry_element.attrib["rgba"].split(" ")
                 geometry_element.set("rgba", f"{r} {g} {b} 0")
 
         # -----------------------
@@ -358,6 +464,21 @@ class RodModelToMjcf:
             texuniform="true",
         )
 
+        _ = (
+            ET.SubElement(
+                asset_element,
+                "hfield",
+                name="terrain",
+                nrow=f"{int(heightmap_samples_xy[0])}",
+                ncol=f"{int(heightmap_samples_xy[1])}",
+                # The following 'size' is a placeholder, it is updated dynamically
+                # when a hfield/heightmap is stored into MjData.
+                size="1 1 1 1",
+            )
+            if heightmap
+            else None
+        )
+
         # ----------------------------------
         # Populate the scene with the assets
         # ----------------------------------
@@ -368,12 +489,14 @@ class RodModelToMjcf:
             worldbody_scene_element,
             "geom",
             name="floor",
-            type="plane",
+            type="plane" if not heightmap else "hfield",
             size="0 0 0.05",
             material="plane_material",
             condim="3",
             contype="1",
             conaffinity="1",
+            zaxis=" ".join(map(str, plane_normal)),
+            **({"hfield": "terrain"} if heightmap else {}),
         )
 
         _ = ET.SubElement(
@@ -407,16 +530,28 @@ class RodModelToMjcf:
             raise RuntimeError("Failed to find the <worldbody> element of the model")
 
         # Camera attached to the model
+        # It can be manually copied from `python -m mujoco.viewer --mjcf=<URDF_PATH>`
         _ = ET.SubElement(
             worldbody_element,
             "camera",
             name="track",
             mode="trackcom",
-            pos="1 0 5",
-            zaxis="0 0 1",
+            pos="1.930 -2.279 0.556",
+            xyaxes="0.771 0.637 0.000 -0.116 0.140 0.983",
             fovy="60",
         )
 
+        # Add user-defined camera.
+        for camera in cameras if isinstance(cameras, Sequence) else [cameras]:
+
+            mj_camera = (
+                camera
+                if isinstance(camera, MujocoCamera)
+                else MujocoCamera.build(**camera)
+            )
+
+            _ = ET.SubElement(worldbody_element, "camera", mj_camera.asdict())
+
         # ------------------------------------------------
         # Add a light following the  CoM of the first link
         # ------------------------------------------------
@@ -444,13 +579,35 @@ class RodModelToMjcf:
 
 
 class UrdfToMjcf:
+    """
+    Class to convert a URDF file to a Mujoco MJCF string.
+    """
+
     @staticmethod
     def convert(
         urdf: str | pathlib.Path,
         considered_joints: list[str] | None = None,
         model_name: str | None = None,
+        plane_normal: tuple[float, float, float] = (0, 0, 1),
+        heightmap: bool | None = None,
+        cameras: MujocoCameraType = (),
     ) -> tuple[str, dict[str, Any]]:
-        """"""
+        """
+        Convert a URDF file to a Mujoco MJCF string.
+
+        Args:
+            urdf: The URDF file to convert.
+            considered_joints: The list of joint names to consider in the conversion.
+            model_name: The name of the model to convert.
+            plane_normal: The normal vector of the plane.
+            heightmap: Whether to generate a heightmap.
+            cameras: The list of cameras to add to the scene.
+
+        Returns:
+            tuple: A tuple containing the MJCF string and the assets dictionary.
+        """
+
+        logging.warning("This method is deprecated. Use 'ModelToMjcf.convert' instead.")
 
         # Get the ROD model.
         rod_model = load_rod_model(
@@ -461,18 +618,44 @@ class UrdfToMjcf:
 
         # Convert the ROD model to MJCF.
         return RodModelToMjcf.convert(
-            rod_model=rod_model, considered_joints=considered_joints
+            rod_model=rod_model,
+            considered_joints=considered_joints,
+            plane_normal=plane_normal,
+            heightmap=heightmap,
+            cameras=cameras,
         )
 
 
 class SdfToMjcf:
+    """
+    Class to convert a SDF file to a Mujoco MJCF string.
+    """
+
     @staticmethod
     def convert(
         sdf: str | pathlib.Path,
         considered_joints: list[str] | None = None,
         model_name: str | None = None,
+        plane_normal: tuple[float, float, float] = (0, 0, 1),
+        heightmap: bool | None = None,
+        cameras: MujocoCameraType = (),
     ) -> tuple[str, dict[str, Any]]:
-        """"""
+        """
+        Convert a SDF file to a Mujoco MJCF string.
+
+        Args:
+            sdf: The SDF file to convert.
+            considered_joints: The list of joint names to consider in the conversion.
+            model_name: The name of the model to convert.
+            plane_normal: The normal vector of the plane.
+            heightmap: Whether to generate a heightmap.
+            cameras: The list of cameras to add to the scene.
+
+        Returns:
+            tuple: A tuple containing the MJCF string and the assets dictionary.
+        """
+
+        logging.warning("This method is deprecated. Use 'ModelToMjcf.convert' instead.")
 
         # Get the ROD model.
         rod_model = load_rod_model(
@@ -483,5 +666,9 @@ class SdfToMjcf:
 
         # Convert the ROD model to MJCF.
         return RodModelToMjcf.convert(
-            rod_model=rod_model, considered_joints=considered_joints
+            rod_model=rod_model,
+            considered_joints=considered_joints,
+            plane_normal=plane_normal,
+            heightmap=heightmap,
+            cameras=cameras,
         )