Fixes a bug in terrain generation when `make_plane` is used (#53)

# Description

<!--
Thank you for your interest in sending a pull request. Please make sure
to check the contribution guidelines.

Link: https://isaac-orbit.github.io/orbit/source/refs/contributing.html
-->

Previously the `make_plane` function was always offsetting the terrain
origin such that the origin of the plane coincides with (0, 0). This was
done because when adding a ground plane, we wanted the origin to match
the simulation origin. However, it introduced a bug in sub-terrain
generator where we assume that all sub-terrains are origin at (size_x *
0.5, size_y * 0.5).

The PR fixes this issue by making an explicit argument to `make_plane`
called `centered`. Also adds comments to make it clear what the
sub-terrain generator and terrain importer does.

<!-- As a practice, it is recommended to open an issue to have
discussions on the proposed pull request.
This makes it easier for the community to keep track of what is being
developed or added, and if a given feature
is demanded by more than one party. -->

## Type of change

<!-- As you go through the list, delete the ones that are not
applicable. -->

- Bug fix (non-breaking change which fixes an issue)

## Screenshots

| Before | After |
| ------ | ----- |
|
![bug](https://github.com/isaac-orbit/orbit/assets/12863862/1ce9db23-b97b-44d9-8a4e-1c543eab914b)|
![fix](https://github.com/isaac-orbit/orbit/assets/12863862/b347077b-daa6-4d76-a307-4e97450ccce9)
|


## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./orbit.sh --format`
- [x] 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

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->

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

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

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

 Changelog
 ---------
 
+0.4.4 (2023-07-05)
+~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added the :attr:`omni.isaac.orbit.terrains.TerrainGeneratorCfg.seed` to make generation of terrains reproducible.
+  The default value is ``None`` which means that the seed is not set.
+
+Fixed
+^^^^^
+
+* Fixed the :meth:`omni.isaac.orbit.terrains.trimesh.utils.make_plane` method to handle the case when the
+  plane origin does not need to be centered.
+
+Changed
+^^^^^^^
+
+* Changed the saving of ``origins`` in :class:`omni.isaac.orbit.terrains.TerrainGenerator` class to be in CSV format
+  instead of NPY format.
+
+
 0.4.3 (2023-06-28)
 ~~~~~~~~~~~~~~~~~~
 

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

@@ -22,7 +22,13 @@ from omni.isaac.orbit.utils import configclass
 
 @configclass
 class SubTerrainBaseCfg:
-    """Base class for terrain configurations."""
+    """Base class for terrain configurations.
+
+    All the sub-terrain configurations must inherit from this class.
+
+    The :attr:`size` attribute is the size of the generated sub-terrain. Based on this, the terrain must
+    extend from :math:`(0, 0)` to :math:`(size[0], size[1])`.
+    """
 
     function: Callable[[float, "SubTerrainBaseCfg"], Tuple[List[trimesh.Trimesh], np.ndarray]] = MISSING
     """Function to generate the terrain.
@@ -48,6 +54,12 @@ class SubTerrainBaseCfg:
 class TerrainGeneratorCfg:
     """Configuration for the terrain generator."""
 
+    seed: Optional[int] = None
+    """The seed for the random number generator. Defaults to :obj:`None`.
+
+    If :obj:`None`, the seed is not set.
+    """
+
     size: Tuple[float, float] = MISSING
     """The width (along x) and length (along y) of each sub-terrain (in m).
 

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

@@ -5,6 +5,7 @@
 
 import numpy as np
 import os
+import torch
 import trimesh
 from typing import List, Tuple
 
@@ -74,6 +75,10 @@ class TerrainGenerator:
                 sub_cfg.vertical_scale = self.cfg.vertical_scale
                 sub_cfg.slope_threshold = self.cfg.slope_threshold
 
+        # set the seed for reproducibility
+        if self.cfg.seed is not None:
+            torch.manual_seed(self.cfg.seed)
+            np.random.seed(self.cfg.seed)
         # create a list of all sub-terrains
         self.terrain_meshes = list()
         self.terrain_origins = np.zeros((self.cfg.num_rows, self.cfg.num_cols, 3))
@@ -183,9 +188,7 @@ class TerrainGenerator:
         # add mesh to the list
         self.terrain_meshes.append(mesh)
         # add origin to the list
-        self.terrain_origins[row, col, 0] = origin[0] + row * self.cfg.size[0]
-        self.terrain_origins[row, col, 1] = origin[1] + col * self.cfg.size[1]
-        self.terrain_origins[row, col, 2] = origin[2]
+        self.terrain_origins[row, col] = origin + transform[:3, -1]
 
     def _get_terrain_mesh(self, difficulty: float, cfg: SubTerrainBaseCfg) -> Tuple[trimesh.Trimesh, np.ndarray]:
         """Generate a sub-terrain mesh based on the input difficulty parameter.
@@ -193,6 +196,10 @@ class TerrainGenerator:
         If caching is enabled, the sub-terrain is cached and loaded from the cache if it exists.
         The cache is stored in the cache directory specified in the configuration.
 
+        .. Note:
+            This function centers the 2D center of the mesh and its specified origin such that the
+            2D center becomes :math:`(0, 0)` instead of :math:`(size[0] / 2, size[1] / 2).
+
         Args:
             difficulty (float): The difficulty parameter.
             cfg (SubTerrainBaseCfg): The configuration of the sub-terrain.
@@ -202,19 +209,20 @@ class TerrainGenerator:
         """
         # add other parameters to the sub-terrain configuration
         cfg.difficulty = float(difficulty)
+        cfg.seed = self.cfg.seed
         # generate hash for the sub-terrain
         sub_terrain_hash = dict_to_md5_hash(cfg.to_dict())
         # generate the file name
         sub_terrain_cache_dir = os.path.join(self.cfg.cache_dir, sub_terrain_hash)
-        sub_terrain_obj_filename = os.path.join(sub_terrain_cache_dir, "mesh.stl")
-        sub_terrain_np_filename = os.path.join(sub_terrain_cache_dir, "origin.npy")
+        sub_terrain_stl_filename = os.path.join(sub_terrain_cache_dir, "mesh.stl")
+        sub_terrain_csv_filename = os.path.join(sub_terrain_cache_dir, "origin.csv")
         sub_terrain_meta_filename = os.path.join(sub_terrain_cache_dir, "cfg.yaml")
 
         # check if hash exists - if true, load the mesh and origin and return
-        if self.cfg.use_cache and os.path.exists(sub_terrain_obj_filename):
+        if self.cfg.use_cache and os.path.exists(sub_terrain_stl_filename):
             # load existing mesh
-            mesh = trimesh.load_mesh(sub_terrain_obj_filename)
-            origin = np.load(sub_terrain_np_filename)
+            mesh = trimesh.load_mesh(sub_terrain_stl_filename)
+            origin = np.loadtxt(sub_terrain_csv_filename, delimiter=",")
             # return the generated mesh
             return mesh, origin
 
@@ -223,15 +231,18 @@ class TerrainGenerator:
         mesh = trimesh.util.concatenate(meshes)
         # offset mesh such that they are in their center
         transform = np.eye(4)
-        transform[0:2, -1] = -self.cfg.size[0] * 0.5, -self.cfg.size[1] * 0.5
+        transform[0:2, -1] = -cfg.size[0] * 0.5, -cfg.size[1] * 0.5
         mesh.apply_transform(transform)
+        # change origin to be in the center of the sub-terrain
+        origin += transform[0:3, -1]
+
         # if caching is enabled, save the mesh and origin
         if self.cfg.use_cache:
             # create the cache directory
             os.makedirs(sub_terrain_cache_dir, exist_ok=True)
             # save the data
-            mesh.export(sub_terrain_obj_filename)
-            np.save(sub_terrain_np_filename, origin)
+            mesh.export(sub_terrain_stl_filename)
+            np.savetxt(sub_terrain_csv_filename, origin, delimiter=",", header="x,y,z")
             dump_yaml(sub_terrain_meta_filename, cfg)
         # return the generated mesh
         return mesh, origin

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

@@ -77,7 +77,7 @@ class TerrainImporter:
             ValueError: If a terrain with the same key already exists.
         """
         # create a plane
-        mesh = make_plane(size, height=0.0)
+        mesh = make_plane(size, height=0.0, centered=True)
         # store the mesh
         self.meshes[key] = mesh
         # create a warp mesh

source/extensions/omni.isaac.orbit/omni/isaac/orbit/terrains/trimesh/mesh_terrains.py

@@ -16,7 +16,7 @@ from typing import TYPE_CHECKING
 from .utils import make_border, make_plane
 
 if TYPE_CHECKING:
-    import omni.isaac.orbit.terrains.trimesh.mesh_terrains_cfg as mesh_terrains_cfg
+    from . import mesh_terrains_cfg
 
 
 def flat_terrain(
@@ -40,11 +40,11 @@ def flat_terrain(
             of the terrain (in m).
     """
     # compute the position of the terrain
-    position = (cfg.size[0] / 2.0, cfg.size[1] / 2.0, 0.0)
+    origin = (cfg.size[0] / 2.0, cfg.size[1] / 2.0, 0.0)
     # compute the vertices of the terrain
-    plane_mesh = make_plane(cfg.size, 0.0)
+    plane_mesh = make_plane(cfg.size, 0.0, center_zero=False)
     # return the tri-mesh and the position
-    return [plane_mesh], np.array(position)
+    return [plane_mesh], np.array(origin)
 
 
 def pyramid_stairs_terrain(
@@ -719,7 +719,7 @@ def star_terrain(
     inner_size = (cfg.size[0] - 2 * bar_width, cfg.size[1] - 2 * bar_width)
     meshes_list += make_border(cfg.size, inner_size, bar_height, platform_center)
     # Generate the ground
-    ground = make_plane(cfg.size, -bar_height)
+    ground = make_plane(cfg.size, -bar_height, center_zero=False)
     meshes_list.append(ground)
     # specify the origin of the terrain
     origin = np.asarray([0.5 * cfg.size[0], 0.5 * cfg.size[1], 0.0])
@@ -850,7 +850,7 @@ def repeated_objects_terrain(
             meshes_list.append(object_mesh)
 
     # generate a ground plane for the terrain
-    ground_plane = make_plane(cfg.size, height=0.0)
+    ground_plane = make_plane(cfg.size, height=0.0, center_zero=False)
     meshes_list.append(ground_plane)
     # generate a platform in the middle
     dim = (cfg.platform_width, cfg.platform_width, 0.5 * height)

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

@@ -3,23 +3,30 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
+from __future__ import annotations
 
 import numpy as np
 import scipy.spatial.transform as tf
 import trimesh
-from typing import List, Tuple
 
 """
 Primitive functions to generate meshes.
 """
 
 
-def make_plane(size: Tuple[float, float], height: float) -> trimesh.Trimesh:
+def make_plane(size: tuple[float, float], height: float, center_zero: bool = True) -> trimesh.Trimesh:
     """Generate a plane mesh.
 
+    If :obj:`center_zero` is True, the origin is at center of the plane mesh i.e. the mesh extends from
+    :math:`(-size[0] / 2, -size[1] / 2, 0)` to :math:`(size[0] / 2, size[1] / 2, height)`.
+    Otherwise, the origin is :math:`(size[0] / 2, size[1] / 2)` and the mesh extends from
+    :math:`(0, 0, 0)` to :math:`(size[0], size[1], height)`.
+
     Args:
         size (Tuple[float, float]): The length (along x) and width (along y) of the terrain (in m).
         height (float): The height of the plane (in m).
+        center_zero (bool, optional): Whether the 2D origin of the plane is set to the center of mesh.
+            Defaults to True.
 
     Returns:
         trimesh.Trimesh: A trimesh.Trimesh objects for the plane.
@@ -33,14 +40,16 @@ def make_plane(size: Tuple[float, float], height: float) -> trimesh.Trimesh:
     vertices = np.array([x0, x1, x2, x3])
     faces = np.array([[1, 0, 2], [2, 3, 1]])
     plane_mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
-    plane_mesh.apply_translation(-np.array([size[0] / 2.0, size[1] / 2.0, 0.0]))
+    # center the plane at the origin
+    if center_zero:
+        plane_mesh.apply_translation(-np.array([size[0] / 2.0, size[1] / 2.0, 0.0]))
     # return the tri-mesh and the position
     return plane_mesh
 
 
 def make_border(
-    size: Tuple[float, float], inner_size: Tuple[float, float], height: float, position: Tuple[float, float, float]
-) -> List[trimesh.Trimesh]:
+    size: tuple[float, float], inner_size: tuple[float, float], height: float, position: tuple[float, float, float]
+) -> list[trimesh.Trimesh]:
     """Generate meshes for a rectangular border with a hole in the middle.
 
     .. code:: text
@@ -94,7 +103,7 @@ def make_box(
     length: float,
     width: float,
     height: float,
-    center: Tuple[float, float, float],
+    center: tuple[float, float, float],
     max_yx_angle: float = 0,
     degrees: bool = True,
 ) -> trimesh.Trimesh:
@@ -128,7 +137,7 @@ def make_box(
 
 
 def make_cylinder(
-    radius: float, height: float, center: Tuple[float, float, float], max_yx_angle: float = 0, degrees: bool = True
+    radius: float, height: float, center: tuple[float, float, float], max_yx_angle: float = 0, degrees: bool = True
 ) -> trimesh.Trimesh:
     """Generate a cylinder mesh with a random orientation.
 
@@ -158,7 +167,7 @@ def make_cylinder(
 
 
 def make_cone(
-    radius: float, height: float, center: Tuple[float, float, float], max_yx_angle: float = 0, degrees: bool = True
+    radius: float, height: float, center: tuple[float, float, float], max_yx_angle: float = 0, degrees: bool = True
 ) -> trimesh.Trimesh:
     """Generate a cone mesh with a random orientation.