From 059f9cbcdcc8e81f9cbd2b3ee4947d5a00eee480 Mon Sep 17 00:00:00 2001
From: Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
Date: Fri, 10 Jan 2025 16:04:09 +0100
Subject: [PATCH] Add `ModelToMjcf` class for converting generic models to MJCF
 format

---
 src/jaxsim/mujoco/__init__.py |  2 +-
 src/jaxsim/mujoco/loaders.py  | 59 +++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/src/jaxsim/mujoco/__init__.py b/src/jaxsim/mujoco/__init__.py
index a2366ae21..6019903fc 100644
--- a/src/jaxsim/mujoco/__init__.py
+++ b/src/jaxsim/mujoco/__init__.py
@@ -1,4 +1,4 @@
-from .loaders import RodModelToMjcf, SdfToMjcf, UrdfToMjcf
+from .loaders import ModelToMjcf, RodModelToMjcf, SdfToMjcf, UrdfToMjcf
 from .model import MujocoModelHelper
 from .utils import mujoco_data_from_jaxsim
 from .visualizer import MujocoVideoRecorder, MujocoVisualizer
diff --git a/src/jaxsim/mujoco/loaders.py b/src/jaxsim/mujoco/loaders.py
index ba476299b..0b14f1cad 100644
--- a/src/jaxsim/mujoco/loaders.py
+++ b/src/jaxsim/mujoco/loaders.py
@@ -9,6 +9,8 @@
 import rod.urdf.exporter
 from lxml import etree as ET
 
+from jaxsim import logging
+
 from .utils import MujocoCamera
 
 MujocoCameraType = (
@@ -61,6 +63,59 @@ 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.
@@ -552,6 +607,8 @@ def convert(
             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(
             model_description=urdf,
@@ -598,6 +655,8 @@ def convert(
             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(
             model_description=sdf,