Add support for inclined planes and heightmap visualization

src/jaxsim/api/model.py

@@ -52,6 +52,8 @@ class JaxSimModel(JaxsimDataclass):
     def build_from_model_description(
         model_description: str | pathlib.Path | rod.Model,
         model_name: str | None = None,
+        *,
+        terrain: jaxsim.terrain.Terrain | None = None,
         is_urdf: bool | None = None,
         considered_joints: list[str] | None = None,
     ) -> JaxSimModel:
@@ -65,6 +67,8 @@ def build_from_model_description(
             model_name:
                 The optional name of the model that overrides the one in
                 the description.
+            terrain:
+                The optional terrain to consider.
             is_urdf:
                 Whether the model description is a URDF or an SDF. This is
                 automatically inferred if the model description is a path to a file.
@@ -92,7 +96,9 @@ def build_from_model_description(
 
         # Build the model
         model = JaxSimModel.build(
-            model_description=intermediate_description, model_name=model_name
+            model_description=intermediate_description,
+            model_name=model_name,
+            terrain=terrain,
         )
 
         # Store the origin of the model, in case downstream logic needs it
@@ -105,6 +111,8 @@ def build_from_model_description(
     def build(
         model_description: jaxsim.parsers.descriptions.ModelDescription,
         model_name: str | None = None,
+        *,
+        terrain: jaxsim.terrain.Terrain | None = None,
     ) -> JaxSimModel:
         """
         Build a Model object from an intermediate model description.
@@ -115,6 +123,8 @@ def build(
                 of the model.
             model_name:
                 The optional name of the model overriding the physics model name.
+            terrain:
+                The optional terrain to consider.
 
         Returns:
             The built Model object.
@@ -130,6 +140,7 @@ def build(
             kin_dyn_parameters=js.kin_dyn_parameters.KynDynParameters.build(
                 model_description=model_description
             ),
+            terrain=terrain or JaxSimModel.__dataclass_fields__["terrain"].default,
         )
 
         return model

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,17 +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",
-        )
+        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,
@@ -412,8 +464,8 @@ def convert(
             "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",
         )
 
@@ -449,6 +501,8 @@ 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: str | Callable | None = None,
     ) -> tuple[str, dict[str, Any]]:
         """"""
 
@@ -461,7 +515,10 @@ def convert(
 
         # 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,
         )
 
 
@@ -471,6 +528,8 @@ 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: str | Callable | None = None,
     ) -> tuple[str, dict[str, Any]]:
         """"""
 
@@ -483,5 +542,8 @@ def convert(
 
         # 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,
         )

src/jaxsim/terrain/__init__.py

@@ -1,2 +1,2 @@
 from . import terrain
-from .terrain import FlatTerrain, Terrain
+from .terrain import FlatTerrain, PlaneTerrain, Terrain