cyy/go2_rl_robotlab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Align Go2 terrain generation with gym

Browse Source
This commit is contained in:
wertyuilife
2026-04-30 20:27:45 +08:00
parent b5d03f6b30
commit 0cdd3b9ac1
3 changed files with 51 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
source/robot_lab/robot_lab/tasks/go2/mdp/commands.py
View File

@@ -94,7 +94,9 @@ class Go2RLGymCommand(CommandTerm):
idxs = is_robot_on_terrain(self._env, terrain_type).nonzero().flatten()
if len(idxs) > 0:
self.terrain_idxs[idxs] = self.terrain_type2idx[terrain_type]
self.terrain_length = self._env.scene.terrain.cfg.terrain_generator.size[0]
terrain_cfg = self._env.scene.terrain.cfg.terrain_generator
sub_terrain_border_width = getattr(terrain_cfg, "sub_terrain_border_width", 0.0) or 0.0
self.terrain_length = max(0.0, terrain_cfg.size[0] - 2.0 * sub_terrain_border_width)
@property
def command(self) -> torch.Tensor:

7
source/robot_lab/robot_lab/tasks/go2/mdp/curriculums.py
View File

@@ -209,9 +209,12 @@ def terrain_levels_vel_gym(env: ManagerBasedRLEnv, env_ids: Sequence[int]) -> fl
resampling_time = command.cfg.resampling_time
zero_prob = command.zero_command_prob
move_up = max_move_dist > terrain.cfg.terrain_generator.size[0] / 2
terrain_cfg = terrain.cfg.terrain_generator
sub_terrain_border_width = getattr(terrain_cfg, "sub_terrain_border_width", 0.0) or 0.0
terrain_length = max(0.0, terrain_cfg.size[0] - 2.0 * sub_terrain_border_width)
move_up = max_move_dist > terrain_length / 2
target_dist = torch.norm(cmd_accum, dim=1) * (resampling_time * (1 - zero_prob))
move_down = (max_move_dist < target_dist * 0.5) * ~move_up
terrain.update_env_origins(env_ids, move_up, move_down)
return torch.mean(terrain.terrain_levels.float()) # type: ignore
return torch.mean(terrain.terrain_levels.float()) # type: ignore

56
source/robot_lab/robot_lab/tasks/go2/mdp/terrains.py
View File

@@ -12,7 +12,22 @@ 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 __init__(self, cfg, device: str = "cpu"):
self._apply_sub_terrain_border_width(cfg)
super().__init__(cfg, device)
def _apply_sub_terrain_border_width(self, cfg):
border_width = getattr(cfg, "sub_terrain_border_width", None)
if border_width is None:
return
for sub_cfg in cfg.sub_terrains.values():
if hasattr(sub_cfg, "border_width"):
sub_cfg.border_width = border_width
def _get_terrain_mesh(self, difficulty, cfg):
difficulty = self._maybe_use_gym_difficulty(difficulty)
override = getattr(cfg, "slope_threshold_override", None)
if override is None:
return super()._get_terrain_mesh(difficulty, cfg)
@@ -24,6 +39,28 @@ class PerSubTerrainSlopeThresholdGenerator(TerrainGenerator):
finally:
cfg.slope_threshold = original_slope_threshold
def _maybe_use_gym_difficulty(self, difficulty: float) -> float:
if not getattr(self.cfg, "use_gym_difficulty", False):
return difficulty
lower, upper = self.cfg.difficulty_range
if upper > lower:
normalized_difficulty = (float(difficulty) - lower) / (upper - lower)
else:
normalized_difficulty = float(difficulty)
normalized_difficulty = np.clip(normalized_difficulty, 0.0, 1.0)
num_levels = self.cfg.num_rows
level = min(int(np.floor(normalized_difficulty * num_levels)), num_levels - 1)
return level / num_levels
@configclass
class Go2TerrainGeneratorCfg(terrain_gen.TerrainGeneratorCfg):
class_type: type = PerSubTerrainSlopeThresholdGenerator
sub_terrain_border_width: float | None = None
use_gym_difficulty: bool = False
def with_slope_threshold(sub_terrain_cfg, slope_threshold: float | None):
"""Attach a per-sub-terrain slope-threshold override to a terrain config."""
@@ -130,23 +167,23 @@ class RoughSlopeTerrainCfg(terrain_gen.HfPyramidSlopedTerrainCfg):
downsampled_scale: float = 0.2
TERRAIN_CFG = terrain_gen.TerrainGeneratorCfg(
class_type=PerSubTerrainSlopeThresholdGenerator,
size=(8.0, 8.0),
TERRAIN_CFG = Go2TerrainGeneratorCfg(
size=(9.0, 9.0), # 8.0 terrain + 0.5*2 sub_terrain_border_width
border_width=25.0,
sub_terrain_border_width=0.5,
num_rows=10,
num_cols=20,
horizontal_scale=0.1,
vertical_scale=0.005,
# slope correction = 0.75 ~ 36.9 degrees by default,
# 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,
slope_threshold=0.75,
use_gym_difficulty=True,
use_cache=False,
sub_terrains={
"wave": with_slope_threshold(
WaveTerrainCfg(
proportion=0.05,
border_width=0.5,
),
10.0, # effectively disable slope correction for wave terrain
),
@@ -155,7 +192,6 @@ TERRAIN_CFG = terrain_gen.TerrainGeneratorCfg(
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
),
@@ -164,14 +200,12 @@ TERRAIN_CFG = terrain_gen.TerrainGeneratorCfg(
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": with_slope_threshold(
RoughSlopeTerrainCfg(
proportion=0.05,
border_width=0.5,
),
10.0, # effectively disable slope correction for rough_slope terrain
),
@@ -181,7 +215,6 @@ TERRAIN_CFG = terrain_gen.TerrainGeneratorCfg(
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
),
@@ -191,7 +224,6 @@ TERRAIN_CFG = terrain_gen.TerrainGeneratorCfg(
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
),
@@ -202,7 +234,6 @@ TERRAIN_CFG = terrain_gen.TerrainGeneratorCfg(
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
),
@@ -214,7 +245,6 @@ TERRAIN_CFG = terrain_gen.TerrainGeneratorCfg(
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
),