[wip] Add support for heightmap visualization

src/jaxsim/mujoco/loaders.py

@@ -1,7 +1,7 @@
 import pathlib
 import tempfile
 import warnings
-from typing import Any
+from typing import Any, Callable
 
 import mujoco as mj
 import rod.urdf.exporter
@@ -43,6 +43,36 @@ def load_rod_model(
     return models[model_name]
 
 
+def generate_hfield(heightmap: Callable, size: tuple[int, int] = (10, 10)) -> str:
+    """"""
+
+    import numpy as np
+
+    # Generate the heightmap.
+    heightmap = heightmap(size)
+
+    # Check the heightmap dimensions.
+    if heightmap.shape != size:
+        raise ValueError(
+            f"Heightmap dimensions {heightmap.shape} do not match the size {size}"
+        )
+
+    # Create the hfield file.
+    with tempfile.NamedTemporaryFile(mode="w+", suffix=".hfield") as hfield_file:
+
+        # Write the header.
+        hfield_file.write(f"{size[0]} {size[1]}\n")
+
+        # Write the heightmap.
+        for row in heightmap:
+            hfield_file.write(" ".join(map(str, row)) + "\n")
+
+        # Move the current position to the beginning.
+        hfield_file.seek(0)
+
+        return hfield_file.name
+
+
 class RodModelToMjcf:
     """"""
 
@@ -129,6 +159,8 @@ def add_floating_joint(
     def convert(
         rod_model: rod.Model,
         considered_joints: list[str] | None = None,
+        plane_normal: tuple[float, float, float] = (0, 0, 1),
+        heightmap: pathlib.Path | Callable | None = None,
     ) -> tuple[str, dict[str, Any]]:
         """"""
 
@@ -364,18 +396,37 @@ def convert(
 
         worldbody_scene_element = ET.SubElement(mujoco_element, "worldbody")
 
-        _ = ET.SubElement(
-            worldbody_scene_element,
-            "geom",
-            name="floor",
-            type="plane",
-            size="0 0 0.05",
-            material="plane_material",
-            condim="3",
-            contype="1",
-            conaffinity="1",
-            zaxis=" ".join(map(str, plane_normal)),
-        )
+        if heightmap:
+            _ = ET.SubElement(
+                worldbody_scene_element,
+                "geom",
+                name="floor",
+                type="hfield",
+                size="10 10 0.05",
+                material="plane_material",
+                condim="3",
+                contype="1",
+                conaffinity="1",
+                zaxis=" ".join(map(str, plane_normal)),
+                file=(
+                    heightmap
+                    if isinstance(heightmap, pathlib.Path)
+                    else generate_hfield(heightmap)
+                ),
+            )
+        else:
+            _ = ET.SubElement(
+                worldbody_scene_element,
+                "geom",
+                name="floor",
+                type="plane",
+                size="0 0 0.05",
+                material="plane_material",
+                condim="3",
+                contype="1",
+                conaffinity="1",
+                zaxis=" ".join(map(str, plane_normal)),
+            )
 
         _ = ET.SubElement(
             worldbody_scene_element,
@@ -451,6 +502,7 @@ def convert(
         considered_joints: list[str] | None = None,
         model_name: str | None = None,
         plane_normal: tuple[float, float, float] = (0, 0, 1),
+        heightmap: str | Callable | None = None,
     ) -> tuple[str, dict[str, Any]]:
         """"""
 
@@ -466,6 +518,7 @@ def convert(
             rod_model=rod_model,
             considered_joints=considered_joints,
             zaxis=plane_normal,
+            heightmap=heightmap,
         )
 
 
@@ -476,6 +529,7 @@ def convert(
         considered_joints: list[str] | None = None,
         model_name: str | None = None,
         plane_normal: tuple[float, float, float] = (0, 0, 1),
+        heightmap: str | Callable | None = None,
     ) -> tuple[str, dict[str, Any]]:
         """"""
 
@@ -491,4 +545,5 @@ def convert(
             rod_model=rod_model,
             considered_joints=considered_joints,
             zaxis=plane_normal,
+            heightmap=heightmap,
         )