Fixes the terrain importer to use the new spawning functions for materials (#118)

# Description

This MR allows coloring a mesh using its vertex colors. Right now it
only demonstrates how to do this based on height but one could use it in
more creative ways. For example, coloring boxes in a different color.

It also updates the `TerrainImporter` to use the visual and physics
material configuration classes.

## Type of change

- New feature (non-breaking change which adds functionality)
- This change requires a documentation update

## Screenshots

| Random Color | Color by Height |
| ------ | ----- |
|
![color](https://github.com/isaac-orbit/orbit/assets/12863862/b2c8fa13-aa6f-44ff-a195-509d422a62b6)
|
![tough](https://github.com/isaac-orbit/orbit/assets/12863862/edc14879-8f14-4ebd-9578-a7eaa26eae26)
|

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./orbit.sh --format`
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file

source/extensions/omni.isaac.orbit/config/extension.toml

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.9.1"
+version = "0.9.2"
 
 # Description
 title = "ORBIT framework for Robot Learning"

source/extensions/omni.isaac.orbit/docs/CHANGELOG.rst

 Changelog
 ---------
 
+0.9.2 (2023-08-22)
+~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added the ability to color meshes in the :class:`omni.isaac.orbit.terrain.TerrainGenerator` class. Currently,
+  it only supports coloring the mesh randomly (``"random"``), based on the terrain height (``"height"``), and
+  no coloring (``"none"``).
+
+Fixed
+^^^^^
+
+* Modified the :class:`omni.isaac.orbit.terrain.TerrainImporter` class to configure visual and physics materials
+  based on the configuration object.
+
+
 0.9.1 (2023-08-18)
 ~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/spawners/from_files/from_files.py

@@ -172,18 +172,18 @@ def spawn_ground_plane(prim_path: str, cfg: from_files_cfg.GroundPlaneCfg, **kwa
 
     # Create physics material
     if cfg.physics_material is not None:
-        cfg.physics_material.func(f"{prim_path}/groundMaterial", cfg.physics_material)
+        cfg.physics_material.func(f"{prim_path}/physicsMaterial", cfg.physics_material)
         # Apply physics material to ground plane
         collision_prim_path = prim_utils.get_prim_path(
             prim_utils.get_first_matching_child_prim(
                 prim_path, predicate=lambda x: prim_utils.get_prim_type_name(x) == "Plane"
             )
         )
-        bind_physics_material(collision_prim_path, f"{prim_path}/groundMaterial")
+        bind_physics_material(collision_prim_path, f"{prim_path}/physicsMaterial")
 
     # Scale only the mesh
     # Warning: This is specific to the default grid plane asset.
-    if cfg.size is not None and prim_utils.is_prim_path_valid(f"{prim_path}/Enviroment"):
+    if prim_utils.is_prim_path_valid(f"{prim_path}/Enviroment"):
         # compute scale from size
         scale = (cfg.size[0] / 100.0, cfg.size[1] / 100.0, 1.0)
         # apply scale to the mesh

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sim/spawners/from_files/from_files_cfg.py

@@ -72,7 +72,7 @@ class GroundPlaneCfg(SpawnerCfg):
 
     If None, then the color remains unchanged.
     """
-    size: tuple[float, float] | None = None
-    """The size of the ground plane. Defaults to None, which is 100 m x 100 m."""
+    size: tuple[float, float] = (100.0, 100.0)
+    """The size of the ground plane. Defaults to 100 m x 100 m."""
     physics_material: materials.RigidBodyMaterialCfg = materials.RigidBodyMaterialCfg()
     """Physics material properties. Defaults to the default rigid body material."""

source/extensions/omni.isaac.orbit/omni/isaac/orbit/terrains/terrain_generator.py

@@ -16,6 +16,7 @@ from omni.isaac.orbit.utils.timer import Timer
 from .height_field import HfTerrainBaseCfg
 from .terrain_generator_cfg import SubTerrainBaseCfg, TerrainGeneratorCfg
 from .trimesh.utils import make_border
+from .utils import color_meshes_by_height
 
 
 class TerrainGenerator:
@@ -92,6 +93,17 @@ class TerrainGenerator:
         self._add_terrain_border()
         # combine all the sub-terrains into a single mesh
         self.terrain_mesh = trimesh.util.concatenate(self.terrain_meshes)
+        # color the terrain mesh
+        if self.cfg.color_scheme == "height":
+            self.terrain_mesh = color_meshes_by_height(self.terrain_mesh)
+        elif self.cfg.color_scheme == "random":
+            self.terrain_mesh.visual.vertex_colors = np.random.choice(
+                range(256), size=(len(self.terrain_mesh.vertices), 4)
+            )
+        elif self.cfg.color_scheme == "none":
+            pass
+        else:
+            raise ValueError(f"Invalid color scheme: {self.cfg.color_scheme}.")
         # offset the entire terrain and origins so that it is centered
         # -- terrain mesh
         transform = np.eye(4)

source/extensions/omni.isaac.orbit/omni/isaac/orbit/terrains/terrain_generator_cfg.py

@@ -18,6 +18,7 @@ import numpy as np
 import trimesh
 from dataclasses import MISSING
 from typing import Callable
+from typing_extensions import Literal
 
 from omni.isaac.orbit.utils import configclass
 
@@ -85,6 +86,16 @@ class TerrainGeneratorCfg:
     num_cols: int = 1
     """Number of columns of sub-terrains to generate. Defaults to 1."""
 
+    color_scheme: Literal["height", "random", "none"] = "none"
+    """Color scheme to use for the terrain. Defaults to "none".
+
+    The available color schemes are:
+
+    - "height": Color based on the height of the terrain.
+    - "random": Random color scheme.
+    - "none": No color scheme.
+    """
+
     horizontal_scale: float = 0.1
     """The discretization of the terrain along the x and y axes (in m). Defaults to 0.1.
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/terrains/terrain_importer.py

@@ -15,9 +15,9 @@ import warp
 from omni.isaac.core.simulation_context import SimulationContext
 from pxr import UsdGeom
 
+import omni.isaac.orbit.sim as sim_utils
 from omni.isaac.orbit.markers import VisualizationMarkers
 from omni.isaac.orbit.markers.config import FRAME_MARKER_CFG
-from omni.isaac.orbit.utils.kit import create_ground_plane
 from omni.isaac.orbit.utils.warp import convert_to_warp_mesh
 
 from .terrain_generator import TerrainGenerator
@@ -110,7 +110,7 @@ class TerrainImporter:
     Operations - Import.
     """
 
-    def import_ground_plane(self, key: str, size: tuple[int, int] = (2.0e6, 2.0e6), **kwargs):
+    def import_ground_plane(self, key: str, size: tuple[int, int] = (2.0e6, 2.0e6)):
         """Add a plane to the terrain importer.
 
         Args:
@@ -131,21 +131,11 @@ class TerrainImporter:
         device = "cuda" if "cuda" in self.device else "cpu"
         self.warp_meshes[key] = convert_to_warp_mesh(mesh.vertices, mesh.faces, device=device)
 
-        # properties for the terrain
-        mesh_props = {
-            "color": self.cfg.color,
-            "static_friction": self.cfg.static_friction,
-            "dynamic_friction": self.cfg.dynamic_friction,
-            "restitution": self.cfg.restitution,
-            "improve_patch_friction": self.cfg.improve_patch_friction,
-            "combine_mode": self.cfg.combine_mode,
-        }
-        # update the properties
-        mesh_props.update(kwargs)
-        # import the grid mesh
-        create_ground_plane(self.cfg.prim_path, **mesh_props)
-
-    def import_mesh(self, key: str, mesh: trimesh.Trimesh, **kwargs):
+        # get the mesh
+        ground_plane_cfg = sim_utils.GroundPlaneCfg(physics_material=self.cfg.physics_material)
+        ground_plane_cfg.func(self.cfg.prim_path, ground_plane_cfg)
+
+    def import_mesh(self, key: str, mesh: trimesh.Trimesh):
         """Import a mesh into the simulator.
 
         The mesh is imported into the simulator under the prim path ``cfg.prim_path/{key}``. The created path
@@ -154,7 +144,6 @@ class TerrainImporter:
         Args:
             key (str): The key to store the mesh.
             mesh (trimesh.Trimesh): The mesh to import.
-            **kwargs: The properties of the mesh. If not provided, the default properties are used.
 
         Raises:
             ValueError: If a terrain with the same key already exists.
@@ -171,19 +160,13 @@ class TerrainImporter:
         # get the mesh
         mesh = self.meshes[key]
         mesh_prim_path = self.cfg.prim_path + f"/{key}"
-        # properties for the terrain
-        mesh_props = {
-            "color": self.cfg.color,
-            "static_friction": self.cfg.static_friction,
-            "dynamic_friction": self.cfg.dynamic_friction,
-            "restitution": self.cfg.restitution,
-            "improve_patch_friction": self.cfg.improve_patch_friction,
-            "combine_mode": self.cfg.combine_mode,
-        }
-        # update the properties
-        mesh_props.update(kwargs)
         # import the mesh
-        create_prim_from_mesh(mesh_prim_path, mesh.vertices, mesh.faces, **mesh_props)
+        create_prim_from_mesh(
+            mesh_prim_path,
+            mesh,
+            visual_material=self.cfg.visual_material,
+            physics_material=self.cfg.physics_material,
+        )
 
     def import_usd(self, key: str, usd_path: str):
         """Import a mesh from a USD file.
@@ -208,6 +191,7 @@ class TerrainImporter:
             raise ValueError(f"Mesh with key {key} already exists. Existing keys: {self.meshes.keys()}.")
         # add the prim path
         prim_utils.create_prim(self.cfg.prim_path + f"/{key}", usd_path=usd_path)
+
         # traverse the prim and get the collision mesh
         # THINK: Should the user specify the collision mesh?
         mesh_prim = prim_utils.get_first_matching_child_prim(

source/extensions/omni.isaac.orbit/omni/isaac/orbit/terrains/terrain_importer_cfg.py

@@ -9,6 +9,7 @@ from dataclasses import MISSING
 from typing import TYPE_CHECKING
 from typing_extensions import Literal
 
+import omni.isaac.orbit.sim as sim_utils
 from omni.isaac.orbit.utils import configclass
 
 from .terrain_importer import TerrainImporter
@@ -24,6 +25,9 @@ class TerrainImporterCfg:
     cls_name: type = TerrainImporter
     """The class name of the terrain importer."""
 
+    collision_group: int = -1
+    """The collision group of the terrain. Defaults to -1."""
+
     prim_path: str = MISSING
     """The absolute path of the USD terrain prim.
 
@@ -58,28 +62,24 @@ class TerrainImporterCfg:
       This parameter is used only when the ``terrain_type`` is ``"plane"`` or ``"usd"``.
     """
 
-    color: tuple[float, float, float] | None = (0.065, 0.0725, 0.080)
-    """The color of the terrain. Defaults to (0.065, 0.0725, 0.080).
-
-    If :obj:`None`, no color is applied to the prim.
-    """
+    visual_material: sim_utils.VisualMaterialCfg | None = sim_utils.PreviewSurfaceCfg(
+        diffuse_color=(0.065, 0.0725, 0.080)
+    )
+    """The visual material of the terrain. Defaults to a dark gray color material.
 
-    static_friction: float = 1.0
-    """The static friction coefficient of the terrain. Defaults to 1.0."""
+    The material is created at the path: ``{prim_path}/visualMaterial``. If `None`, then no material is created.
 
-    dynamic_friction: float = 1.0
-    """The dynamic friction coefficient of the terrain. Defaults to 1.0."""
-
-    restitution: float = 0.0
-    """The restitution coefficient of the terrain. Defaults to 0.0."""
+    .. note::
+        This parameter is used only when the ``terrain_type`` is ``"generator"``.
+    """
 
-    improve_patch_friction: bool = False
-    """Whether to enable patch friction. Defaults to False."""
+    physics_material: sim_utils.RigidBodyMaterialCfg = sim_utils.RigidBodyMaterialCfg()
+    """The physics material of the terrain. Defaults to a default physics material.
 
-    combine_mode: str = "average"
-    """Determines the way physics materials will be combined during collisions. Defaults to `average`.
+    The material is created at the path: ``{prim_path}/physicsMaterial``.
 
-    Available options are `average`, `min`, `multiply`, `multiply`, and `max`.
+    .. note::
+        This parameter is used only when the ``terrain_type`` is ``"generator"`` or ``"plane"``.
     """
 
     max_init_terrain_level: int | None = None

source/extensions/omni.isaac.orbit/omni/isaac/orbit/terrains/utils.py

@@ -20,6 +20,7 @@ def color_meshes_by_height(meshes: List[trimesh.Trimesh], **kwargs) -> trimesh.T
     Keyword Args:
         color (List[int]): A list of 3 integers in the range [0,255] representing the RGB
             color of the mesh. Used when the z-coordinates of all vertices are the same.
+        color_map (str): The name of the color map to be used. Defaults to "turbo".
 
     Returns:
         trimesh.Trimesh: A trimesh object with the vertices colored based on the z-coordinate (height) of each vertex.
@@ -38,15 +39,18 @@ def color_meshes_by_height(meshes: List[trimesh.Trimesh], **kwargs) -> trimesh.T
     else:
         # Normalize the heights to [0,1]
         heights_normalized = (heights - np.min(heights)) / (np.max(heights) - np.min(heights))
+        # clip lower and upper bounds to have better color mapping
+        heights_normalized = np.clip(heights_normalized, 0.1, 0.9)
         # Get the color for each vertex based on the height
-        colors = trimesh.visual.color.interpolate(heights_normalized, color_map="turbo")
+        color_map = kwargs.pop("color_map", "turbo")
+        colors = trimesh.visual.color.interpolate(heights_normalized, color_map=color_map)
         # Set the vertex colors
         mesh.visual.vertex_colors = colors
     # Return the mesh
     return mesh
 
 
-def create_prim_from_mesh(prim_path: str, vertices: np.ndarray, triangles: np.ndarray, **kwargs):
+def create_prim_from_mesh(prim_path: str, mesh: trimesh.Trimesh, **kwargs):
     """Create a USD prim with mesh defined from vertices and triangles.
 
     The function creates a USD prim with a mesh defined from vertices and triangles. It performs the
@@ -59,67 +63,62 @@ def create_prim_from_mesh(prim_path: str, vertices: np.ndarray, triangles: np.nd
 
     Args:
         prim_path (str): The path to the primitive to be created.
-        vertices (np.ndarray): The vertices of the mesh. Shape is :math:`(N, 3)`, where :math:`N`
-            is the number of vertices.
-        triangles (np.ndarray): The triangles of the mesh as references to vertices for each triangle.
-            Shape is :math:`(M, 3)`, where :math:`M` is the number of triangles / faces.
+        mesh (trimesh.Trimesh): The mesh to be used for the primitive.
 
     Keyword Args:
         translation (Optional[Sequence[float]]): The translation of the terrain. Defaults to None.
         orientation (Optional[Sequence[float]]): The orientation of the terrain. Defaults to None.
-        scale (Optional[Sequence[float]]): The scale of the terrain. Defaults to None.
-        color (Optional[tuple]): The color of the terrain. Defaults to (0.065, 0.0725, 0.080).
-        static_friction (float): The static friction of the terrain. Defaults to 1.0.
-        dynamic_friction (float): The dynamic friction of the terrain. Defaults to 1.0.
-        restitution (float): The restitution of the terrain. Defaults to 0.0.
-        improve_patch_friction (bool): Whether to enable patch friction. Defaults to False.
-        combine_mode (str): Determines the way physics materials will be combined during collisions.
-            Available options are `average`, `min`, `multiply`, `multiply`, and `max`. Defaults to `average`.
+        visual_material (Optional[sim_utils.VisualMaterialCfg]): The visual material to apply. Defaults to None.
+        physics_material (Optional[sim_utils.RigidBodyMaterialCfg]): The physics material to apply. Defaults to None.
     """
     # need to import these here to prevent isaacsim launching when importing this module
     import omni.isaac.core.utils.prims as prim_utils
-    from omni.isaac.core.materials import PhysicsMaterial, PreviewSurface
-    from omni.isaac.core.prims import GeometryPrim, XFormPrim
-    from pxr import PhysxSchema
+    from pxr import UsdGeom
+
+    import omni.isaac.orbit.sim as sim_utils
 
     # create parent prim
     prim_utils.create_prim(prim_path, "Xform")
     # create mesh prim
-    prim_utils.create_prim(
+    prim = prim_utils.create_prim(
         f"{prim_path}/mesh",
         "Mesh",
         translation=kwargs.get("translation"),
         orientation=kwargs.get("orientation"),
-        scale=kwargs.get("scale"),
         attributes={
-            "points": vertices,
-            "faceVertexIndices": triangles.flatten(),
-            "faceVertexCounts": np.asarray([3] * len(triangles)),
+            "points": mesh.vertices,
+            "faceVertexIndices": mesh.faces.flatten(),
+            "faceVertexCounts": np.asarray([3] * len(mesh.faces)),
             "subdivisionScheme": "bilinear",
         },
     )
+    # apply collider properties
+    collider_cfg = sim_utils.CollisionPropertiesCfg(collision_enabled=True)
+    sim_utils.define_collision_properties(prim.GetPrimPath(), collider_cfg)
+    # add rgba color to the mesh primvars
+    if mesh.visual.vertex_colors is not None:
+        # obtain color from the mesh
+        rgba_colors = np.asarray(mesh.visual.vertex_colors).astype(np.float32) / 255.0
+        # displayColor is a primvar attribute that is used to color the mesh
+        color_prim_attr = prim.GetAttribute("primvars:displayColor")
+        color_prim_var = UsdGeom.Primvar(color_prim_attr)
+        color_prim_var.SetInterpolation(UsdGeom.Tokens.vertex)
+        color_prim_attr.Set(rgba_colors[:, :3])
+        # displayOpacity is a primvar attribute that is used to set the opacity of the mesh
+        display_prim_attr = prim.GetAttribute("primvars:displayOpacity")
+        display_prim_var = UsdGeom.Primvar(display_prim_attr)
+        display_prim_var.SetInterpolation(UsdGeom.Tokens.vertex)
+        display_prim_var.Set(rgba_colors[:, 3])
 
     # create visual material
-    color = kwargs.get("color", (0.065, 0.0725, 0.080))
-    if color is not None:
-        material = PreviewSurface(f"{prim_path}/visualMaterial", color=np.asarray(color))
-        XFormPrim(f"{prim_path}/mesh").apply_visual_material(material)
-
+    if kwargs.get("visual_material") is not None:
+        visual_material_cfg: sim_utils.VisualMaterialCfg = kwargs.get("visual_material")
+        # spawn the material
+        visual_material_cfg.func(f"{prim_path}/visualMaterial", visual_material_cfg)
+        sim_utils.bind_visual_material(prim.GetPrimPath(), f"{prim_path}/visualMaterial")
     # create physics material
-    material = PhysicsMaterial(
-        f"{prim_path}/physicsMaterial",
-        static_friction=kwargs.get("static_friction", 1.0),
-        dynamic_friction=kwargs.get("dynamic_friction", 1.0),
-        restitution=kwargs.get("restitution", 0.0),
-    )
-    # apply PhysX Rigid Material schema
-    physx_material_api = PhysxSchema.PhysxMaterialAPI.Apply(material.prim)
-    # set patch friction property
-    improve_patch_friction = kwargs.get("improve_patch_friction", False)
-    physx_material_api.CreateImprovePatchFrictionAttr().Set(improve_patch_friction)
-    # set combination mode for coefficients
-    combine_mode = kwargs.get("combine_mode", "multiply")
-    physx_material_api.CreateFrictionCombineModeAttr().Set(combine_mode)
-    physx_material_api.CreateRestitutionCombineModeAttr().Set(combine_mode)
-    # apply physics material to ground plane
-    GeometryPrim(f"{prim_path}/mesh", collision=True).apply_physics_material(material)
+    if kwargs.get("physics_material") is not None:
+        physics_material_cfg: sim_utils.RigidBodyMaterialCfg = kwargs.get("physics_material")
+        # spawn the material
+        physics_material_cfg.func(f"{prim_path}/physicsMaterial", physics_material_cfg)
+        sim_utils.bind_physics_material(prim.GetPrimPath(), f"{prim_path}/physicsMaterial")

source/extensions/omni.isaac.orbit/test/terrains/test_terrain_importer.py

@@ -39,9 +39,17 @@ parser.add_argument("--geom_sphere", action="store_true", default=False, help="W
 parser.add_argument(
     "--terrain_type",
     type=str,
+    choices=["generated", "usd", "plane"],
     default="generator",
     help="Type of terrain to import. Can be 'generated' or 'usd' or 'plane'.",
 )
+parser.add_argument(
+    "--color_scheme",
+    type=str,
+    default="height",
+    choices=["height", "random", "none"],
+    help="The color scheme to use for the generated terrain.",
+)
 args_cli = parser.parse_args()
 
 # launch omniverse app
@@ -69,7 +77,7 @@ from omni.isaac.orbit.terrains.terrain_importer import TerrainImporter
 from omni.isaac.orbit.utils.assets import ISAAC_NUCLEUS_DIR
 
 
-def main(terrain_type: str):
+def main():
     """Generates a terrain from orbit."""
 
     # Load kit helper
@@ -100,8 +108,8 @@ def main(terrain_type: str):
         env_spacing=3.0,
         prim_path="/World/ground",
         max_init_terrain_level=None,
-        terrain_type=terrain_type,
-        terrain_generator=ROUGH_TERRAINS_CFG.replace(curriculum=True),
+        terrain_type=args_cli.terrain_type,
+        terrain_generator=ROUGH_TERRAINS_CFG.replace(curriculum=True, color_scheme=args_cli.color_scheme),
         usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Terrains/rough_plane.usd",
     )
     terrain_importer = TerrainImporter(terrain_importer_cfg)
@@ -191,6 +199,6 @@ def main(terrain_type: str):
 
 if __name__ == "__main__":
     # Runs the main function
-    main(terrain_type=args_cli.terrain_type)
+    main()
     # Close the simulator
     simulation_app.close()