diff --git a/source/robot_lab/robot_lab/tasks/go2/mdp/terrains.py b/source/robot_lab/robot_lab/tasks/go2/mdp/terrains.py index 764fec6..b09a4d9 100644 --- a/source/robot_lab/robot_lab/tasks/go2/mdp/terrains.py +++ b/source/robot_lab/robot_lab/tasks/go2/mdp/terrains.py @@ -3,16 +3,40 @@ from __future__ import annotations import numpy as np import isaaclab.terrains as terrain_gen +from isaaclab.terrains.terrain_generator import TerrainGenerator from isaaclab.utils import configclass from isaaclab.terrains.height_field import hf_terrains from isaaclab.terrains.height_field.utils import height_field_to_mesh +class PerSubTerrainSlopeThresholdGenerator(TerrainGenerator): + """Terrain generator that lets each height-field sub-terrain override slope_threshold.""" + + def _get_terrain_mesh(self, difficulty, cfg): + override = getattr(cfg, "slope_threshold_override", None) + if override is None: + return super()._get_terrain_mesh(difficulty, cfg) + + original_slope_threshold = getattr(cfg, "slope_threshold", None) + cfg.slope_threshold = override + try: + return super()._get_terrain_mesh(difficulty, cfg) + finally: + cfg.slope_threshold = original_slope_threshold + + +def with_slope_threshold(sub_terrain_cfg, slope_threshold: float | None): + """Attach a per-sub-terrain slope-threshold override to a terrain config.""" + sub_terrain_cfg.slope_threshold_override = slope_threshold + return sub_terrain_cfg + + # ----------------------------------------------------------------------------- # Default RobotLab terrain setup # ----------------------------------------------------------------------------- DEFAULT_TERRAIN_CFG = terrain_gen.TerrainGeneratorCfg( + class_type=PerSubTerrainSlopeThresholdGenerator, size=(8.0, 8.0), border_width=25.0, num_rows=10, @@ -107,65 +131,92 @@ class RoughSlopeTerrainCfg(terrain_gen.HfPyramidSlopedTerrainCfg): TERRAIN_CFG = terrain_gen.TerrainGeneratorCfg( + class_type=PerSubTerrainSlopeThresholdGenerator, size=(8.0, 8.0), border_width=25.0, num_rows=10, num_cols=20, horizontal_scale=0.1, vertical_scale=0.005, - slope_threshold=2.0, + # slope correction = 0.75 ~ 36.9 degrees by default, + # but recommended to set for each terrain type separately using with_slope_threshold + slope_threshold=0.75, use_cache=False, sub_terrains={ - "wave": WaveTerrainCfg( - proportion=0.05, - border_width=0.5, + "wave": with_slope_threshold( + WaveTerrainCfg( + proportion=0.05, + border_width=0.5, + ), + 10.0, # effectively disable slope correction for wave terrain ), - "slope_up": terrain_gen.HfInvertedPyramidSlopedTerrainCfg( - proportion=0.10, - slope_range=(0.1, 0.568), - platform_width=3.0, - border_width=0.5, + "slope_up": with_slope_threshold( + terrain_gen.HfInvertedPyramidSlopedTerrainCfg( + proportion=0.10, + slope_range=(0.1, 0.568), + platform_width=3.0, + border_width=0.5, + ), + 10.0, # effectively disable slope correction for slope_up terrain ), - "slope_down": terrain_gen.HfPyramidSlopedTerrainCfg( - proportion=0.10, - slope_range=(0.1, 0.568), - platform_width=3.0, - border_width=0.5, + "slope_down": with_slope_threshold( + terrain_gen.HfPyramidSlopedTerrainCfg( + proportion=0.10, + slope_range=(0.1, 0.568), + platform_width=3.0, + border_width=0.5, + ), + 10.0, # effectively disable slope correction for slope_down terrain ), - "rough_slope": RoughSlopeTerrainCfg( - proportion=0.05, - border_width=0.5, + "rough_slope": with_slope_threshold( + RoughSlopeTerrainCfg( + proportion=0.05, + border_width=0.5, + ), + 10.0, # effectively disable slope correction for rough_slope terrain ), - "stairs_up": terrain_gen.HfInvertedPyramidStairsTerrainCfg( - proportion=0.25, - step_height_range=(0.05, 0.257), - step_width=0.31, - platform_width=3.0, - border_width=0.5, + "stairs_up": with_slope_threshold( + terrain_gen.HfInvertedPyramidStairsTerrainCfg( + proportion=0.25, + step_height_range=(0.05, 0.257), + step_width=0.31, + platform_width=3.0, + border_width=0.5, + ), + 0.25, # enable slope correction for rough_slope terrain by 14.0 degrees, which is recommended for stairs terrain ), - "stairs_down": terrain_gen.HfPyramidStairsTerrainCfg( - proportion=0.10, - step_height_range=(0.05, 0.257), - step_width=0.31, - platform_width=3.0, - border_width=0.5, + "stairs_down": with_slope_threshold( + terrain_gen.HfPyramidStairsTerrainCfg( + proportion=0.10, + step_height_range=(0.05, 0.257), + step_width=0.31, + platform_width=3.0, + border_width=0.5, + ), + 0.25, # enable slope correction for rough_slope terrain by 14.0 degrees, which is recommended for stairs terrain ), - "obstacles": terrain_gen.HfDiscreteObstaclesTerrainCfg( - proportion=0.20, - obstacle_width_range=(1.0, 2.0), - obstacle_height_range=(0.05, 0.275), - num_obstacles=20, - platform_width=3.0, - border_width=0.5, + "obstacles": with_slope_threshold( + terrain_gen.HfDiscreteObstaclesTerrainCfg( + proportion=0.20, + obstacle_width_range=(1.0, 2.0), + obstacle_height_range=(0.05, 0.275), + num_obstacles=20, + platform_width=3.0, + border_width=0.5, + ), + 0.25, # enable slope correction for rough_slope terrain by 14.0 degrees ), - "stepping_stones": terrain_gen.HfSteppingStonesTerrainCfg( - proportion=0.0, - stone_height_max=0.0, - stone_width_range=(0.075, 1.575), - stone_distance_range=(0.05, 0.10), - holes_depth=-10.0, - platform_width=4.0, - border_width=0.5, + "stepping_stones": with_slope_threshold( + terrain_gen.HfSteppingStonesTerrainCfg( + proportion=0.0, + stone_height_max=0.0, + stone_width_range=(0.075, 1.575), + stone_distance_range=(0.05, 0.10), + holes_depth=-10.0, + platform_width=4.0, + border_width=0.5, + ), + 0.25, # enable slope correction for rough_slope terrain by 14.0 degrees ), "gap": terrain_gen.MeshGapTerrainCfg( proportion=0.0,