Adds terrain-based 2D-pose command and reset (#424)

# Description

* Added functionality to sample flat patches on a generated terrain.
Added separate normal and terrain-based position commands. Terrain-based
commands rely on the terrain to sample flat patches.
* Added a terrain-based root reset function to reset the robot to a
random flat patch.

The MR includes the following fixes:

* Fixed command resample termination function.

## Type of change

- Bug fix (non-breaking change which fixes an issue)
- New feature (non-breaking change which adds functionality)
- This change requires a documentation update

## Screenshots

![poses](https://github.com/isaac-orbit/orbit/assets/12863862/678d3cde-63da-4c57-8aae-4c84440f55e3)

## 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 run all the tests with `./orbit.sh --test` and they pass
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

---------
Co-authored-by: Mayank Mittal <mittalma@leggedrobotics.com>

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

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

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

 Changelog
 ---------
 
+0.12.0 (2024-03-08)
+~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added functionality to sample flat patches on a generated terrain. This can be configured using
+  :attr:`omni.isaac.orbit.terrains.SubTerrainBaseCfg.flat_patch_sampling` attribute.
+* Added a randomization function for setting terrain-aware root state. Through this, an asset can be
+  reset to a randomly sampled flat patches.
+
+Fixed
+^^^^^
+
+* Separated normal and terrain-base position commands. The terrain based commands rely on the
+  terrain to sample flat patches for setting the target position.
+* Fixed command resample termination function.
+
+Changed
+^^^^^^^
+
+* Added the attribute :attr:`omni.isaac.orbit.envs.mdp.commands.UniformVelocityCommandCfg.heading_control_stiffness`
+  to control the stiffness of the heading control term in the velocity command term. Earlier, this was
+  hard-coded to 0.5 inside the term.
+
+Removed
+^^^^^^^
+
+* Removed the function :meth:`sample_new_targets` in the terrain importer. Instead the attribute
+  :attr:`omni.isaac.orbit.terrains.TerrainImporter.flat_patches` should be used to sample new targets.
+
+
 0.11.3 (2024-03-04)
 ~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/commands/__init__.py

@@ -8,11 +8,12 @@
 from .commands_cfg import (
     NormalVelocityCommandCfg,
     NullCommandCfg,
-    TerrainBasedPositionCommandCfg,
+    UniformPose2dCommandCfg,
     UniformPoseCommandCfg,
+    UniformTerrainBasedPose2dCommandCfg,
     UniformVelocityCommandCfg,
 )
 from .null_command import NullCommand
+from .pose_2d_command import UniformPose2dCommand, UniformTerrainBasedPose2dCommandCfg
 from .pose_command import UniformPoseCommand
-from .position_command import TerrainBasedPositionCommand
 from .velocity_command import NormalVelocityCommand, UniformVelocityCommand

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/commands/commands_cfg.py

@@ -12,14 +12,10 @@ from omni.isaac.orbit.managers import CommandTermCfg
 from omni.isaac.orbit.utils import configclass
 
 from .null_command import NullCommand
+from .pose_2d_command import UniformPose2dCommand, UniformTerrainBasedPose2dCommand
 from .pose_command import UniformPoseCommand
-from .position_command import TerrainBasedPositionCommand
 from .velocity_command import NormalVelocityCommand, UniformVelocityCommand
 
-"""
-Null-command generator.
-"""
-
 
 @configclass
 class NullCommandCfg(CommandTermCfg):
@@ -33,11 +29,6 @@ class NullCommandCfg(CommandTermCfg):
         self.resampling_time_range = (math.inf, math.inf)
 
 
-"""
-Locomotion-specific command generators.
-"""
-
-
 @configclass
 class UniformVelocityCommandCfg(CommandTermCfg):
     """Configuration for the uniform velocity command generator."""
@@ -53,6 +44,8 @@ class UniformVelocityCommandCfg(CommandTermCfg):
     target heading is sampled uniformly from provided range. Otherwise, the angular velocity
     command is sampled uniformly from provided range.
     """
+    heading_control_stiffness: float = MISSING
+    """Scale factor to convert the heading error to angular velocity command."""
     rel_standing_envs: float = MISSING
     """Probability threshold for environments where the robots that are standing still."""
     rel_heading_envs: float = MISSING
@@ -130,10 +123,10 @@ class UniformPoseCommandCfg(CommandTermCfg):
 
 
 @configclass
-class TerrainBasedPositionCommandCfg(CommandTermCfg):
-    """Configuration for the terrain-based position command generator."""
+class UniformPose2dCommandCfg(CommandTermCfg):
+    """Configuration for the uniform 2D-pose command generator."""
 
-    class_type: type = TerrainBasedPositionCommand
+    class_type: type = UniformPose2dCommand
 
     asset_name: str = MISSING
     """Name of the asset in the environment for which the commands are generated."""
@@ -147,8 +140,12 @@ class TerrainBasedPositionCommandCfg(CommandTermCfg):
 
     @configclass
     class Ranges:
-        """Uniform distribution ranges for the velocity commands."""
+        """Uniform distribution ranges for the position commands."""
 
+        pos_x: tuple[float, float] = MISSING
+        """Range for the x position (in m)."""
+        pos_y: tuple[float, float] = MISSING
+        """Range for the y position (in m)."""
         heading: tuple[float, float] = MISSING
         """Heading range for the position commands (in rad).
 
@@ -157,3 +154,23 @@ class TerrainBasedPositionCommandCfg(CommandTermCfg):
 
     ranges: Ranges = MISSING
     """Distribution ranges for the position commands."""
+
+
+@configclass
+class UniformTerrainBasedPose2dCommandCfg(UniformPose2dCommandCfg):
+    """Configuration for the terrain-based position command generator."""
+
+    class_type = UniformTerrainBasedPose2dCommand
+
+    @configclass
+    class Ranges:
+        """Uniform distribution ranges for the position commands."""
+
+        heading: tuple[float, float] = MISSING
+        """Heading range for the position commands (in rad).
+
+        Used only if :attr:`simple_heading` is False.
+        """
+
+    ranges: Ranges = MISSING
+    """Distribution ranges for the sampled commands."""

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/commands/velocity_command.py

@@ -120,8 +120,9 @@ class UniformVelocityCommand(CommandTerm):
             # resolve indices of heading envs
             env_ids = self.is_heading_env.nonzero(as_tuple=False).flatten()
             # compute angular velocity
+            heading_error = math_utils.wrap_to_pi(self.heading_target[env_ids] - self.robot.data.heading_w[env_ids])
             self.vel_command_b[env_ids, 2] = torch.clip(
-                0.5 * math_utils.wrap_to_pi(self.heading_target[env_ids] - self.robot.data.heading_w[env_ids]),
+                self.cfg.heading_control_stiffness * heading_error,
                 min=self.cfg.ranges.ang_vel_z[0],
                 max=self.cfg.ranges.ang_vel_z[1],
             )

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/randomizations.py

@@ -220,6 +220,76 @@ def reset_root_state_uniform(
     asset.write_root_velocity_to_sim(velocities, env_ids=env_ids)
 
 
+def reset_robot_root_from_terrain(
+    env: BaseEnv,
+    env_ids: torch.Tensor,
+    pose_range: dict[str, tuple[float, float]],
+    velocity_range: dict[str, tuple[float, float]],
+    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
+):
+    """Reset the robot root state by sampling a random valid pose from the terrain.
+
+    This function samples a random valid pose(based on flat patches) from the terrain and sets the root state
+    of the robot to this pose. The function also samples random velocities from the given ranges and sets them
+    into the physics simulation.
+
+    Note:
+        The function expects the terrain to have valid flat patches under the key "init_pos". The flat patches
+        are used to sample the random pose for the robot.
+
+    Raises:
+        ValueError: If the terrain does not have valid flat patches under the key "init_pos".
+    """
+    # access the used quantities (to enable type-hinting)
+    asset: RigidObject | Articulation = env.scene[asset_cfg.name]
+
+    # obtain all flat patches corresponding to the valid poses
+    valid_poses: torch.Tensor = env.scene.terrain.flat_patches.get("init_pos")
+    if valid_poses is None:
+        raise ValueError(
+            "The randomization term 'reset_robot_root_from_terrain' requires valid flat patches under 'init_pos'."
+            f" Found: {list(env.scene.terrain.flat_patches.keys())}"
+        )
+
+    # sample random valid poses
+    ids = torch.randint(0, valid_poses.shape[2], size=(len(env_ids),), device=env.device)
+    positions = valid_poses[env.scene.terrain.terrain_levels[env_ids], env.scene.terrain.terrain_types[env_ids], ids]
+    positions += asset.data.default_root_state[env_ids, :3]
+
+    # sample random orientations
+    ranges = torch.tensor(
+        [
+            pose_range.get("roll", (0.0, 0.0)),
+            pose_range.get("pitch", (0.0, 0.0)),
+            pose_range.get("yaw", (0.0, 0.0)),
+        ],
+        device=env.device,
+    )
+    euler_angles = torch.zeros_like(positions).uniform_()
+    euler_angles = ranges[0] + (ranges[1] - ranges[0]) * euler_angles
+    # convert to quaternions
+    orientations = quat_from_euler_xyz(euler_angles[:, 0], euler_angles[:, 1], euler_angles[:, 2])
+
+    # sample random velocities
+    ranges = torch.tensor(
+        [
+            velocity_range.get("x", (0.0, 0.0)),
+            velocity_range.get("y", (0.0, 0.0)),
+            velocity_range.get("z", (0.0, 0.0)),
+            velocity_range.get("roll", (0.0, 0.0)),
+            velocity_range.get("pitch", (0.0, 0.0)),
+            velocity_range.get("yaw", (0.0, 0.0)),
+        ],
+        device=env.device,
+    )
+    velocities = torch.zeros(len(env_ids), 6, device=asset.device).uniform_()
+    velocities = ranges[:, 0] + (ranges[:, 1] - ranges[:, 0]) * velocities
+
+    # set into the physics simulation
+    asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids)
+    asset.write_root_velocity_to_sim(velocities, env_ids=env_ids)
+
+
 def reset_joints_by_scale(
     env: BaseEnv,
     env_ids: torch.Tensor,

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/terminations.py

@@ -20,6 +20,7 @@ from omni.isaac.orbit.sensors import ContactSensor
 
 if TYPE_CHECKING:
     from omni.isaac.orbit.envs import RLTaskEnv
+    from omni.isaac.orbit.managers.command_manager import CommandTerm
 
 """
 MDP terminations.
@@ -31,15 +32,14 @@ def time_out(env: RLTaskEnv) -> torch.Tensor:
     return env.episode_length_buf >= env.max_episode_length
 
 
-def command_resample(env: RLTaskEnv, num_resamples: int = 1) -> torch.Tensor:
+def command_resample(env: RLTaskEnv, command_name: str, num_resamples: int = 1) -> torch.Tensor:
     """Terminate the episode based on the total number of times commands have been re-sampled.
 
     This makes the maximum episode length fluid in nature as it depends on how the commands are
     sampled. It is useful in situations where delayed rewards are used :cite:`rudin2022advanced`.
     """
-    return torch.logical_and(
-        (env.command_manager.time_left <= env.step_dt), (env.command_manager.command_counter == num_resamples)
-    )
+    command: CommandTerm = env.command_manager.get_term(command_name)
+    return torch.logical_and((command.time_left <= env.step_dt), (command.command_counter == num_resamples))
 
 
 """

source/extensions/omni.isaac.orbit/omni/isaac/orbit/managers/command_manager.py

@@ -354,6 +354,17 @@ class CommandManager(ManagerBase):
         """
         return self._terms[name].command
 
+    def get_term(self, name: str) -> CommandTerm:
+        """Returns the command term with the specified name.
+
+        Args:
+            name: The name of the command term.
+
+        Returns:
+            The command term with the specified name.
+        """
+        return self._terms[name]
+
     """
     Helper functions.
     """

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

@@ -23,7 +23,7 @@ There are two main components in this package:
 
 from .height_field import *  # noqa: F401, F403
 from .terrain_generator import TerrainGenerator
-from .terrain_generator_cfg import SubTerrainBaseCfg, TerrainGeneratorCfg
+from .terrain_generator_cfg import FlatPatchSamplingCfg, SubTerrainBaseCfg, TerrainGeneratorCfg
 from .terrain_importer import TerrainImporter
 from .terrain_importer_cfg import TerrainImporterCfg
 from .trimesh import *  # noqa: F401, F403

source/extensions/omni.isaac.orbit/omni/isaac/orbit/terrains/config/rough.py

@@ -19,7 +19,6 @@ ROUGH_TERRAINS_CFG = TerrainGeneratorCfg(
     horizontal_scale=0.1,
     vertical_scale=0.005,
     slope_threshold=0.75,
-    difficulty_choices=(0.5, 0.75, 0.9),
     use_cache=False,
     sub_terrains={
         "pyramid_stairs": terrain_gen.MeshPyramidStairsTerrainCfg(

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

@@ -10,14 +10,17 @@ import os
 import torch
 import trimesh
 
+import carb
+
 from omni.isaac.orbit.utils.dict import dict_to_md5_hash
 from omni.isaac.orbit.utils.io import dump_yaml
 from omni.isaac.orbit.utils.timer import Timer
+from omni.isaac.orbit.utils.warp import convert_to_warp_mesh
 
 from .height_field import HfTerrainBaseCfg
-from .terrain_generator_cfg import SubTerrainBaseCfg, TerrainGeneratorCfg
+from .terrain_generator_cfg import FlatPatchSamplingCfg, SubTerrainBaseCfg, TerrainGeneratorCfg
 from .trimesh.utils import make_border
-from .utils import color_meshes_by_height
+from .utils import color_meshes_by_height, find_flat_patches
 
 
 class TerrainGenerator:
@@ -41,6 +44,11 @@ class TerrainGenerator:
     The difficulty is varied linearly over the number of rows (i.e. along x). If a curriculum
     is not used, the terrains are generated randomly.
 
+    If the :obj:`cfg.flat_patch_sampling` is specified for a sub-terrain, flat patches are sampled
+    on the terrain. These can be used for spawning robots, targets, etc. The sampled patches are stored
+    in the :obj:`flat_patches` dictionary. The key specifies the intention of the flat patches and the
+    value is a tensor containing the flat patches for each sub-terrain.
+
     If the flag :obj:`cfg.use_cache` is set to True, the terrains are cached based on their
     sub-terrain configurations. This means that if the same sub-terrain configuration is used
     multiple times, the terrain is only generated once and then reused. This is useful when
@@ -53,18 +61,32 @@ class TerrainGenerator:
     """List of trimesh.Trimesh objects for all the generated sub-terrains."""
     terrain_origins: np.ndarray
     """The origin of each sub-terrain. Shape is (num_rows, num_cols, 3)."""
+    flat_patches: dict[str, torch.Tensor]
+    """A dictionary of sampled valid (flat) patches for each sub-terrain.
+
+    The dictionary keys are the names of the flat patch sampling configurations. This maps to a
+    tensor containing the flat patches for each sub-terrain. The shape of the tensor is
+    (num_rows, num_cols, num_patches, 3).
+
+    For instance, the key "root_spawn" maps to a tensor containing the flat patches for spawning an asset.
+    Similarly, the key "target_spawn" maps to a tensor containing the flat patches for setting targets.
+    """
 
-    def __init__(self, cfg: TerrainGeneratorCfg):
+    def __init__(self, cfg: TerrainGeneratorCfg, device: str = "cpu"):
         """Initialize the terrain generator.
 
         Args:
             cfg: Configuration for the terrain generator.
+            device: The device to use for the flat patches tensor.
         """
         # check inputs
         if len(cfg.sub_terrains) == 0:
             raise ValueError("No sub-terrains specified! Please add at least one sub-terrain.")
         # store inputs
         self.cfg = cfg
+        self.device = device
+        # -- valid patches
+        self.flat_patches = {}
         # set common values to all sub-terrains config
         for sub_cfg in self.cfg.sub_terrains.values():
             # size of all terrains
@@ -82,6 +104,7 @@ class TerrainGenerator:
         # 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))
+
         # parse configuration and add sub-terrains
         # create terrains based on curriculum or randomly
         if self.cfg.curriculum:
@@ -94,6 +117,7 @@ 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)
@@ -105,6 +129,7 @@ class TerrainGenerator:
             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)
@@ -112,6 +137,10 @@ class TerrainGenerator:
         self.terrain_mesh.apply_transform(transform)
         # -- terrain origins
         self.terrain_origins += transform[:3, -1]
+        # -- valid patches
+        terrain_origins_torch = torch.tensor(self.terrain_origins, dtype=torch.float, device=self.device).unsqueeze(2)
+        for name, value in self.flat_patches.items():
+            self.flat_patches[name] = value + terrain_origins_torch
 
     """
     Terrain generator functions.
@@ -132,11 +161,11 @@ class TerrainGenerator:
             # randomly sample terrain index
             sub_index = np.random.choice(len(proportions), p=proportions)
             # randomly sample difficulty parameter
-            difficulty = np.random.choice(self.cfg.difficulty_choices)
+            difficulty = np.random.uniform(*self.cfg.difficulty_range)
             # generate terrain
             mesh, origin = self._get_terrain_mesh(difficulty, sub_terrains_cfgs[sub_index])
             # add to sub-terrains
-            self._add_sub_terrain(mesh, origin, sub_row, sub_col)
+            self._add_sub_terrain(mesh, origin, sub_row, sub_col, sub_terrains_cfgs[sub_index])
 
     def _generate_curriculum_terrains(self):
         """Add terrains based on the difficulty parameter."""
@@ -157,12 +186,14 @@ class TerrainGenerator:
         # curriculum-based sub-terrains
         for sub_col in range(self.cfg.num_cols):
             for sub_row in range(self.cfg.num_rows):
-                # vary the difficulty parameter
-                difficulty = sub_row / self.cfg.num_rows
+                # vary the difficulty parameter linearly over the number of rows
+                lower, upper = self.cfg.difficulty_range
+                difficulty = (sub_row + np.random.uniform()) / self.cfg.num_rows
+                difficulty = lower + (upper - lower) * difficulty
                 # generate terrain
                 mesh, origin = self._get_terrain_mesh(difficulty, sub_terrains_cfgs[sub_indices[sub_col]])
                 # add to sub-terrains
-                self._add_sub_terrain(mesh, origin, sub_row, sub_col)
+                self._add_sub_terrain(mesh, origin, sub_row, sub_col, sub_terrains_cfgs[sub_indices[sub_col]])
 
     """
     Internal helper functions.
@@ -186,8 +217,13 @@ class TerrainGenerator:
         # add the border to the list of meshes
         self.terrain_meshes.append(border)
 
-    def _add_sub_terrain(self, mesh: trimesh.Trimesh, origin: np.ndarray, row: int, col: int):
-        """Add input sub-terrain to the list of sub-terrains meshes and origins.
+    def _add_sub_terrain(
+        self, mesh: trimesh.Trimesh, origin: np.ndarray, row: int, col: int, sub_terrain_cfg: SubTerrainBaseCfg
+    ):
+        """Add input sub-terrain to the list of sub-terrains.
+
+        This function adds the input sub-terrain mesh to the list of sub-terrains and updates the origin
+        of the sub-terrain in the list of origins. It also samples flat patches if specified.
 
         Args:
             mesh: The mesh of the sub-terrain.
@@ -195,6 +231,32 @@ class TerrainGenerator:
             row: The row index of the sub-terrain.
             col: The column index of the sub-terrain.
         """
+        # sample flat patches if specified
+        if sub_terrain_cfg.flat_patch_sampling is not None:
+            carb.log_info(f"Sampling flat patches for sub-terrain at (row, col):  ({row}, {col})")
+            # convert the mesh to warp mesh
+            wp_mesh = convert_to_warp_mesh(mesh.vertices, mesh.faces, device=self.device)
+            # sample flat patches based on each patch configuration for that sub-terrain
+            for name, patch_cfg in sub_terrain_cfg.flat_patch_sampling.items():
+                patch_cfg: FlatPatchSamplingCfg
+                # create the flat patches tensor (if not already created)
+                if name not in self.flat_patches:
+                    self.flat_patches[name] = torch.zeros(
+                        (self.cfg.num_rows, self.cfg.num_cols, patch_cfg.num_patches, 3), device=self.device
+                    )
+                # add the flat patches to the tensor
+                self.flat_patches[name][row, col] = find_flat_patches(
+                    wp_mesh=wp_mesh,
+                    origin=origin,
+                    num_patches=patch_cfg.num_patches,
+                    patch_radius=patch_cfg.patch_radius,
+                    x_range=patch_cfg.x_range,
+                    y_range=patch_cfg.y_range,
+                    z_range=patch_cfg.z_range,
+                    max_height_diff=patch_cfg.max_height_diff,
+                )
+
+        # transform the mesh to the correct position
         transform = np.eye(4)
         transform[0:2, -1] = (row + 0.5) * self.cfg.size[0], (col + 0.5) * self.cfg.size[1]
         mesh.apply_transform(transform)

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

@@ -23,6 +23,45 @@ from typing import Literal
 from omni.isaac.orbit.utils import configclass
 
 
+@configclass
+class FlatPatchSamplingCfg:
+    """Configuration for sampling flat patches on the sub-terrain.
+
+    For a given sub-terrain, this configuration specifies how to sample flat patches on the terrain.
+    The sampled flat patches can be used for spawning robots, targets, etc.
+
+    Please check the function :meth:`~omni.isaac.orbit.terrains.utils.find_flat_patches` for more details.
+    """
+
+    num_patches: int = MISSING
+    """Number of patches to sample."""
+
+    patch_radius: float | list[float] = MISSING
+    """Radius of the patches.
+
+    A list of radii can be provided to check for patches of different sizes. This is useful to deal with
+    cases where the terrain may have holes or obstacles in some areas.
+    """
+
+    x_range: tuple[float, float] = (-1e6, 1e6)
+    """The range of x-coordinates to sample from. Defaults to (-1e6, 1e6).
+
+    This range is internally clamped to the size of the terrain mesh.
+    """
+
+    y_range: tuple[float, float] = (-1e6, 1e6)
+    """The range of y-coordinates to sample from. Defaults to (-1e6, 1e6).
+
+    This range is internally clamped to the size of the terrain mesh.
+    """
+
+    z_range: tuple[float, float] = (-1e6, 1e6)
+    """Allowed range of z-coordinates for the sampled patch. Defaults to (-1e6, 1e6)."""
+
+    max_height_diff: float = MISSING
+    """Maximum allowed height difference between the highest and lowest points on the patch."""
+
+
 @configclass
 class SubTerrainBaseCfg:
     """Base class for terrain configurations.
@@ -52,6 +91,14 @@ class SubTerrainBaseCfg:
     size: tuple[float, float] = MISSING
     """The width (along x) and length (along y) of the terrain (in m)."""
 
+    flat_patch_sampling: dict[str, FlatPatchSamplingCfg] | None = None
+    """Dictionary of configurations for sampling flat patches on the sub-terrain. Defaults to None,
+    in which case no flat patch sampling is performed.
+
+    The keys correspond to the name of the flat patch sampling configuration and the values are the
+    corresponding configurations.
+    """
+
 
 @configclass
 class TerrainGeneratorCfg:
@@ -115,16 +162,17 @@ class TerrainGeneratorCfg:
     """
 
     sub_terrains: dict[str, SubTerrainBaseCfg] = MISSING
-    """List of sub-terrain configurations."""
+    """Dictionary of sub-terrain configurations.
 
-    difficulty_choices: list[float] = [0.5, 0.75, 0.9]
-    """List of difficulty choices. Defaults to [0.5, 0.75, 0.9].
+    The keys correspond to the name of the sub-terrain configuration and the values are the corresponding
+    configurations.
+    """
 
-    The difficulty choices are used to sample the difficulty of the generated terrain. The specified
-    choices are randomly sampled with equal probability.
+    difficulty_range: tuple[float, float] = (0.0, 1.0)
+    """The range of difficulty values for the sub-terrains. Defaults to (0.0, 1.0).
 
-    Note:
-      This is used only when curriculum-based generation is disabled.
+    If curriculum is enabled, the terrains will be generated based on this range in ascending order
+    of difficulty. Otherwise, the terrains will be generated based on this range in a random order.
     """
 
     use_cache: bool = False

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

@@ -8,7 +8,6 @@ from __future__ import annotations
 import numpy as np
 import torch
 import trimesh
-from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 import warp
@@ -77,6 +76,8 @@ class TerrainImporter:
         self.warp_meshes = dict()
         self.env_origins = None
         self.terrain_origins = None
+        # private variables
+        self._terrain_flat_patches = dict()
 
         # auto-import the terrain based on the config
         if self.cfg.terrain_type == "generator":
@@ -84,10 +85,12 @@ class TerrainImporter:
             if self.cfg.terrain_generator is None:
                 raise ValueError("Input terrain type is 'generator' but no value provided for 'terrain_generator'.")
             # generate the terrain
-            terrain_generator = TerrainGenerator(cfg=self.cfg.terrain_generator)
+            terrain_generator = TerrainGenerator(cfg=self.cfg.terrain_generator, device=self.device)
             self.import_mesh("terrain", terrain_generator.terrain_mesh)
             # configure the terrain origins based on the terrain generator
             self.configure_env_origins(terrain_generator.terrain_origins)
+            # refer to the flat patches
+            self._terrain_flat_patches = terrain_generator.flat_patches
         elif self.cfg.terrain_type == "usd":
             # check if config is provided
             if self.cfg.usd_path is None:
@@ -119,6 +122,17 @@ class TerrainImporter:
         """
         return True
 
+    @property
+    def flat_patches(self) -> dict[str, torch.Tensor]:
+        """A dictionary containing the sampled valid (flat) patches for the terrain.
+
+        This is only available if the terrain type is 'generator'. For other terrain types, this feature
+        is not available and the function returns an empty dictionary.
+
+        Please refer to the :attr:`TerrainGenerator.flat_patches` for more information.
+        """
+        return self._terrain_flat_patches
+
     """
     Operations - Visibility.
     """
@@ -305,26 +319,6 @@ class TerrainImporter:
         # update the env origins
         self.env_origins[env_ids] = self.terrain_origins[self.terrain_levels[env_ids], self.terrain_types[env_ids]]
 
-    """
-    Operations - Sampling
-    """
-
-    def sample_new_targets(self, env_ids: Sequence[int]) -> torch.Tensor:
-        """Samples terrain-aware locations of flat patches to set spawn or target locations.
-
-        Note:
-            This is a dummy function that returns the environment origins as target locations.
-            Please inherit the class and reimplement the function for specific terrain types
-
-        Args:
-            env_ids: The environment indices to sample targets locations for.
-
-        Returns:
-            The sampled target locations as (x, y, z). Shape is (N, 3).
-        """
-
-        return self.env_origins[env_ids]
-
     """
     Internal helpers.
     """

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

@@ -6,8 +6,13 @@
 from __future__ import annotations
 
 import numpy as np
+import torch
 import trimesh
 
+import warp as wp
+
+from omni.isaac.orbit.utils.warp import raycast_mesh
+
 
 def color_meshes_by_height(meshes: list[trimesh.Trimesh], **kwargs) -> trimesh.Trimesh:
     """
@@ -123,3 +128,146 @@ def create_prim_from_mesh(prim_path: str, mesh: trimesh.Trimesh, **kwargs):
         # 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")
+
+
+def find_flat_patches(
+    wp_mesh: wp.Mesh,
+    num_patches: int,
+    patch_radius: float | list[float],
+    origin: np.ndarray | torch.Tensor | tuple[float, float, float],
+    x_range: tuple[float, float],
+    y_range: tuple[float, float],
+    z_range: tuple[float, float],
+    max_height_diff: float,
+):
+    """Finds flat patches of given radius in the input mesh.
+
+    The function finds flat patches of given radius based on the search space defined by the input ranges.
+    The search space is characterized by origin in the mesh frame, and the x, y, and z ranges. The x and y
+    ranges are used to sample points in the 2D region around the origin, and the z range is used to filter
+    patches based on the height of the points.
+
+    The function performs rejection sampling to find the patches based on the following steps:
+
+    1. Sample patch locations in the 2D region around the origin.
+    2. Define a ring of points around each patch location to query the height of the points using ray-casting.
+    3. Reject patches that are outside the z range or have a height difference that is too large.
+    4. Keep sampling until all patches are valid.
+
+    Args:
+        wp_mesh: The warp mesh to find patches in.
+        num_patches: The desired number of patches to find.
+        patch_radius: The radii used to form patches. If a list is provided, multiple patch sizes are checked.
+            This is useful to deal with holes or other artifacts in the mesh.
+        origin: The origin defining the center of the search space. This is specified in the mesh frame.
+        x_range: The range of X coordinates to sample from.
+        y_range: The range of Y coordinates to sample from.
+        z_range: The range of valid Z coordinates used for filtering patches.
+        max_height_diff: The maximum allowable distance between the lowest and highest points
+            on a patch to consider it as valid. If the difference is greater than this value,
+            the patch is rejected.
+
+    Returns:
+        A tensor of shape (num_patches, 3) containing the flat patches. The patches are defined in the mesh frame.
+
+    Raises:
+        RuntimeError: If the function fails to find valid patches. This can happen if the input parameters
+            are not suitable for finding valid patches and maximum number of iterations is reached.
+    """
+    # set device to warp mesh device
+    device = wp.device_to_torch(wp_mesh.device)
+
+    # resolve inputs to consistent type
+    # -- patch radii
+    if isinstance(patch_radius, float):
+        patch_radius = [patch_radius]
+    # -- origin
+    if isinstance(origin, np.ndarray):
+        origin = torch.from_numpy(origin).to(torch.float).to(device)
+    elif isinstance(origin, torch.Tensor):
+        origin = origin.to(device)
+    else:
+        origin = torch.tensor(origin, dtype=torch.float, device=device)
+
+    # create ranges for the x and y coordinates around the origin.
+    # The provided ranges are bounded by the mesh's bounding box.
+    x_range = (
+        max(x_range[0] + origin[0].item(), wp_mesh.points.numpy()[:, 0].min()),
+        min(x_range[1] + origin[0].item(), wp_mesh.points.numpy()[:, 0].max()),
+    )
+    y_range = (
+        max(y_range[0] + origin[1].item(), wp_mesh.points.numpy()[:, 1].min()),
+        min(y_range[1] + origin[1].item(), wp_mesh.points.numpy()[:, 1].max()),
+    )
+    z_range = (
+        z_range[0] + origin[2].item(),
+        z_range[1] + origin[2].item(),
+    )
+
+    # create a circle of points around (0, 0) to query validity of the patches
+    # the ring of points is uniformly distributed around the circle
+    angle = torch.linspace(0, 2 * np.pi, 10, device=device)
+    query_x = []
+    query_y = []
+    for radius in patch_radius:
+        query_x.append(radius * torch.cos(angle))
+        query_y.append(radius * torch.sin(angle))
+    query_x = torch.cat(query_x).unsqueeze(1)  # dim: (num_radii * 10, 1)
+    query_y = torch.cat(query_y).unsqueeze(1)  # dim: (num_radii * 10, 1)
+    # dim: (num_radii * 10, 3)
+    query_points = torch.cat([query_x, query_y, torch.zeros_like(query_x)], dim=-1)
+
+    # create buffers
+    # -- a buffer to store indices of points that are not valid
+    points_ids = torch.arange(num_patches, device=device)
+    # -- a buffer to store the flat patches locations
+    flat_patches = torch.zeros(num_patches, 3, device=device)
+
+    # sample points and raycast to find the height.
+    # 1. Reject points that are outside the z_range or have a height difference that is too large.
+    # 2. Keep sampling until all points are valid.
+    iter_count = 0
+    while len(points_ids) > 0 and iter_count < 10000:
+        # sample points in the 2D region around the origin
+        pos_x = torch.empty(len(points_ids), device=device).uniform_(*x_range)
+        pos_y = torch.empty(len(points_ids), device=device).uniform_(*y_range)
+        flat_patches[points_ids, :2] = torch.stack([pos_x, pos_y], dim=-1)
+
+        # define the query points to check validity of the patch
+        # dim: (num_patches, num_radii * 10, 3)
+        points = flat_patches[points_ids].unsqueeze(1) + query_points
+        points[..., 2] = 100.0
+        # ray-cast direction is downwards
+        dirs = torch.zeros_like(points)
+        dirs[..., 2] = -1.0
+
+        # ray-cast to find the height of the patches
+        ray_hits = raycast_mesh(points.view(-1, 3), dirs.view(-1, 3), wp_mesh)[0]
+        heights = ray_hits.view(points.shape)[..., 2]
+        # set the height of the patches
+        # note: for invalid patches, they would be overwritten in the next iteration
+        #   so it's safe to set the height to the last value
+        flat_patches[points_ids, 2] = heights[..., -1]
+
+        # check validity
+        # -- height is within the z range
+        not_valid = torch.any(torch.logical_or(heights < z_range[0], heights > z_range[1]), dim=1)
+        # -- height difference is within the max height difference
+        not_valid = torch.logical_or(not_valid, (heights.max(dim=1)[0] - heights.min(dim=1)[0]) > max_height_diff)
+
+        # remove invalid patches indices
+        points_ids = points_ids[not_valid]
+        # increment count
+        iter_count += 1
+
+    # check all patches are valid
+    if len(points_ids) > 0:
+        raise RuntimeError(
+            "Failed to find valid patches! Please check the input parameters."
+            f"\n\tMaximum number of iterations reached: {iter_count}"
+            f"\n\tNumber of invalid patches: {len(points_ids)}"
+            f"\n\tMaximum height difference: {max_height_diff}"
+        )
+
+    # return the flat patches (in the mesh frame)
+    return flat_patches - origin

source/extensions/omni.isaac.orbit_tasks/omni/isaac/orbit_tasks/locomotion/velocity/velocity_env_cfg.py

@@ -98,6 +98,7 @@ class CommandsCfg:
         rel_standing_envs=0.02,
         rel_heading_envs=1.0,
         heading_command=True,
+        heading_control_stiffness=0.5,
         debug_vis=True,
         ranges=mdp.UniformVelocityCommandCfg.Ranges(
             lin_vel_x=(-1.0, 1.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.0, 1.0), heading=(-math.pi, math.pi)

source/standalone/demos/procedural_terrain.py

+# Copyright (c) 2022-2024, The ORBIT Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+This script demonstrates procedural terrains with flat patches.
+
+Example usage:
+
+.. code-block:: bash
+
+    # Generate terrain with height color scheme
+    ./orbit.sh -p source/standalone/demos/procedural_terrain.py --color_scheme height
+
+    # Generate terrain with random color scheme
+    ./orbit.sh -p source/standalone/demos/procedural_terrain.py --color_scheme random
+
+    # Generate terrain with no color scheme
+    ./orbit.sh -p source/standalone/demos/procedural_terrain.py --color_scheme none
+
+    # Generate terrain with curriculum
+    ./orbit.sh -p source/standalone/demos/procedural_terrain.py --use_curriculum
+
+    # Generate terrain with curriculum along with flat patches
+    ./orbit.sh -p source/standalone/demos/procedural_terrain.py --use_curriculum --show_flat_patches
+
+"""
+
+from __future__ import annotations
+
+"""Launch Isaac Sim Simulator first."""
+
+
+import argparse
+
+from omni.isaac.orbit.app import AppLauncher
+
+# add argparse arguments
+parser = argparse.ArgumentParser(description="This script demonstrates procedural terrain generation.")
+parser.add_argument(
+    "--color_scheme",
+    type=str,
+    default="none",
+    choices=["height", "random", "none"],
+    help="Color scheme to use for the terrain generation.",
+)
+parser.add_argument(
+    "--use_curriculum",
+    action="store_true",
+    default=False,
+    help="Whether to use the curriculum for the terrain generation.",
+)
+parser.add_argument(
+    "--show_flat_patches",
+    action="store_true",
+    default=False,
+    help="Whether to show the flat patches computed during the terrain generation.",
+)
+# append AppLauncher cli args
+AppLauncher.add_app_launcher_args(parser)
+# parse the arguments
+args_cli = parser.parse_args()
+
+# launch omniverse app
+app_launcher = AppLauncher(args_cli)
+simulation_app = app_launcher.app
+
+"""Rest everything follows."""
+
+import random
+import torch
+import traceback
+
+import carb
+
+import omni.isaac.orbit.sim as sim_utils
+from omni.isaac.orbit.assets import AssetBase
+from omni.isaac.orbit.markers import VisualizationMarkers, VisualizationMarkersCfg
+from omni.isaac.orbit.terrains import FlatPatchSamplingCfg, TerrainImporter, TerrainImporterCfg
+
+##
+# Pre-defined configs
+##
+from omni.isaac.orbit.terrains.config.rough import ROUGH_TERRAINS_CFG  # isort:skip
+
+
+def design_scene() -> tuple[dict, torch.Tensor]:
+    """Designs the scene."""
+    # Lights
+    cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75))
+    cfg.func("/World/Light", cfg)
+
+    # Parse terrain generation
+    terrain_gen_cfg = ROUGH_TERRAINS_CFG.replace(curriculum=args_cli.use_curriculum, color_scheme=args_cli.color_scheme)
+
+    # Add flat patch configuration
+    # Note: To have separate colors for each sub-terrain type, we set the flat patch sampling configuration name
+    #   to the sub-terrain name. However, this is not how it should be used in practice. The key name should be
+    #   the intention of the flat patch. For instance, "source" or "target" for spawn and command related flat patches.
+    if args_cli.show_flat_patches:
+        for sub_terrain_name, sub_terrain_cfg in terrain_gen_cfg.sub_terrains.items():
+            sub_terrain_cfg.flat_patch_sampling = {
+                sub_terrain_name: FlatPatchSamplingCfg(num_patches=10, patch_radius=0.5, max_height_diff=0.05)
+            }
+
+    # Handler for terrains importing
+    terrain_importer_cfg = TerrainImporterCfg(
+        num_envs=2048,
+        env_spacing=3.0,
+        prim_path="/World/ground",
+        max_init_terrain_level=None,
+        terrain_type="generator",
+        terrain_generator=terrain_gen_cfg,
+        debug_vis=True,
+    )
+    # Remove visual material for height and random color schemes to use the default material
+    if args_cli.color_scheme in ["height", "random"]:
+        terrain_importer_cfg.visual_material = None
+    # Create terrain importer
+    terrain_importer = TerrainImporter(terrain_importer_cfg)
+
+    # Show the flat patches computed
+    if args_cli.show_flat_patches:
+        # Configure the flat patches
+        vis_cfg = VisualizationMarkersCfg(prim_path="/Visuals/TerrainFlatPatches", markers={})
+        for name in terrain_importer.flat_patches:
+            vis_cfg.markers[name] = sim_utils.CylinderCfg(
+                radius=0.5,  # note: manually set to the patch radius for visualization
+                height=0.1,
+                visual_material=sim_utils.GlassMdlCfg(glass_color=(random.random(), random.random(), random.random())),
+            )
+        flat_patches_visualizer = VisualizationMarkers(vis_cfg)
+
+        # Visualize the flat patches
+        all_patch_locations = []
+        all_patch_indices = []
+        for i, patch_locations in enumerate(terrain_importer.flat_patches.values()):
+            num_patch_locations = patch_locations.view(-1, 3).shape[0]
+            # store the patch locations and indices
+            all_patch_locations.append(patch_locations.view(-1, 3))
+            all_patch_indices += [i] * num_patch_locations
+        # combine the patch locations and indices
+        flat_patches_visualizer.visualize(torch.cat(all_patch_locations), marker_indices=all_patch_indices)
+
+    # return the scene information
+    scene_entities = {"terrain": terrain_importer}
+    return scene_entities, terrain_importer.env_origins
+
+
+def run_simulator(sim: sim_utils.SimulationContext, entities: dict[str, AssetBase], origins: torch.Tensor):
+    """Runs the simulation loop."""
+    # Simulate physics
+    while simulation_app.is_running():
+        # perform step
+        sim.step()
+
+
+def main():
+    """Main function."""
+
+    # Initialize the simulation context
+    sim = sim_utils.SimulationContext(sim_utils.SimulationCfg(dt=0.01, substeps=1))
+    # Set main camera
+    sim.set_camera_view(eye=[2.5, 2.5, 2.5], target=[0.0, 0.0, 0.0])
+    # design scene
+    scene_entities, scene_origins = design_scene()
+    # Play the simulator
+    sim.reset()
+    # Now we are ready!
+    print("[INFO]: Setup complete...")
+    # Run the simulator
+    run_simulator(sim, scene_entities, scene_origins)
+
+
+if __name__ == "__main__":
+    try:
+        # run the main execution
+        main()
+    except Exception as err:
+        carb.log_error(err)
+        carb.log_error(traceback.format_exc())
+        raise
+    finally:
+        # close sim app
+        simulation_app.close()