|
| 1 | +"""Terrain generator: dataclass cfg → 2D layout + Genesis kwargs. |
| 2 | +
|
| 3 | +The generator is a pure-Python helper: it builds the 2D ``subterrain_types`` layout |
| 4 | +(random by proportion or explicit), aggregates the per-type parameters into the dict |
| 5 | +shape Genesis expects, and computes the planar centers of each cell so callers can |
| 6 | +place envs on them later (curriculum reset, locomotion goals). |
| 7 | +
|
| 8 | +Note: Genesis keys ``subterrain_parameters`` by the *type string*, not per cell. If |
| 9 | +two ``SubTerrainCfg`` instances map to the same ``genesis_type`` (e.g. two |
| 10 | +``PyramidStairsCfg`` with different ``step_width``), only one parameter set survives. |
| 11 | +Use distinct keys in ``TerrainGeneratorCfg.sub_terrains`` for distinct geometries; for |
| 12 | +in-type variation use ``randomize=True`` on the importer. |
| 13 | +""" |
| 14 | + |
| 15 | +import random |
| 16 | +from dataclasses import dataclass, field |
| 17 | +from typing import Any |
| 18 | + |
| 19 | +import torch |
| 20 | + |
| 21 | +from genelab.terrains.sub_terrain import SubTerrainCfg |
| 22 | + |
| 23 | + |
| 24 | +@dataclass |
| 25 | +class TerrainGeneratorCfg: |
| 26 | + """Composition of a heightfield terrain grid. |
| 27 | +
|
| 28 | + The grid is ``num_rows`` × ``num_cols`` cells; each cell has size |
| 29 | + ``subterrain_size`` (in meters). When ``layout`` is ``None``, cell types are |
| 30 | + sampled by ``SubTerrainCfg.proportion`` using a deterministic ``seed``. |
| 31 | + """ |
| 32 | + |
| 33 | + num_rows: int = 4 |
| 34 | + num_cols: int = 4 |
| 35 | + subterrain_size: tuple[float, float] = (8.0, 8.0) |
| 36 | + horizontal_scale: float = 0.1 |
| 37 | + vertical_scale: float = 0.005 |
| 38 | + pos: tuple[float, float, float] = (0.0, 0.0, 0.0) |
| 39 | + sub_terrains: dict[str, SubTerrainCfg] = field(default_factory=dict) |
| 40 | + layout: tuple[tuple[str, ...], ...] | None = None |
| 41 | + curriculum: bool = False |
| 42 | + seed: int = 0 |
| 43 | + |
| 44 | + |
| 45 | +class TerrainGenerator: |
| 46 | + """Translate :class:`TerrainGeneratorCfg` into Genesis ``Terrain`` kwargs.""" |
| 47 | + |
| 48 | + def __init__(self, cfg: TerrainGeneratorCfg) -> None: |
| 49 | + if not cfg.sub_terrains: |
| 50 | + raise ValueError("TerrainGeneratorCfg.sub_terrains must contain at least one entry") |
| 51 | + self.cfg = cfg |
| 52 | + self._layout: list[list[str]] = self._build_layout() |
| 53 | + self._env_origins: torch.Tensor = self._compute_origins() |
| 54 | + |
| 55 | + # ------------------------------------------------------------------ layout |
| 56 | + |
| 57 | + def _build_layout(self) -> list[list[str]]: |
| 58 | + cfg = self.cfg |
| 59 | + if cfg.layout is not None: |
| 60 | + rows = [list(row) for row in cfg.layout] |
| 61 | + if len(rows) != cfg.num_rows or any(len(r) != cfg.num_cols for r in rows): |
| 62 | + raise ValueError( |
| 63 | + f"layout shape {len(rows)}x{len(rows[0]) if rows else 0} " |
| 64 | + f"does not match num_rows x num_cols = {cfg.num_rows}x{cfg.num_cols}" |
| 65 | + ) |
| 66 | + unknown = {name for row in rows for name in row} - set(cfg.sub_terrains.keys()) |
| 67 | + if unknown: |
| 68 | + raise ValueError(f"layout references unknown sub_terrains: {sorted(unknown)}") |
| 69 | + return rows |
| 70 | + # Random tiling weighted by proportion. |
| 71 | + names = list(cfg.sub_terrains.keys()) |
| 72 | + weights = [cfg.sub_terrains[n].proportion for n in names] |
| 73 | + if sum(weights) <= 0: |
| 74 | + raise ValueError("sub_terrains proportions sum to zero") |
| 75 | + rng = random.Random(cfg.seed) |
| 76 | + # PR3 (curriculum) will refine row ordering by difficulty; for PR1 every row uses |
| 77 | + # the same proportion distribution regardless of ``cfg.curriculum``. |
| 78 | + return [rng.choices(names, weights=weights, k=cfg.num_cols) for _ in range(cfg.num_rows)] |
| 79 | + |
| 80 | + # ------------------------------------------------------------------ origins |
| 81 | + |
| 82 | + def _compute_origins(self) -> torch.Tensor: |
| 83 | + cfg = self.cfg |
| 84 | + sx, sy = cfg.subterrain_size |
| 85 | + px, py, pz = cfg.pos |
| 86 | + origins = torch.zeros((cfg.num_rows, cfg.num_cols, 3)) |
| 87 | + for i in range(cfg.num_rows): |
| 88 | + for j in range(cfg.num_cols): |
| 89 | + origins[i, j, 0] = px + (i + 0.5) * sx |
| 90 | + origins[i, j, 1] = py + (j + 0.5) * sy |
| 91 | + origins[i, j, 2] = pz |
| 92 | + return origins |
| 93 | + |
| 94 | + # ------------------------------------------------------------------ outputs |
| 95 | + |
| 96 | + @property |
| 97 | + def layout(self) -> list[list[str]]: |
| 98 | + """2D list of sub-terrain cfg keys at each (row, col).""" |
| 99 | + return self._layout |
| 100 | + |
| 101 | + @property |
| 102 | + def env_origins(self) -> torch.Tensor: |
| 103 | + """Per-cell planar centers, shape ``(num_rows, num_cols, 3)``. |
| 104 | +
|
| 105 | + ``z`` is set to ``cfg.pos[2]`` (the terrain root height); callers needing the |
| 106 | + actual surface elevation should sample :attr:`TerrainImporter.heightfield` |
| 107 | + after spawn. |
| 108 | + """ |
| 109 | + return self._env_origins |
| 110 | + |
| 111 | + @property |
| 112 | + def genesis_kwargs(self) -> dict[str, Any]: |
| 113 | + """Kwargs ready to splat into ``gs.morphs.Terrain(...)``.""" |
| 114 | + cfg = self.cfg |
| 115 | + types_2d = [[cfg.sub_terrains[name].genesis_type() for name in row] for row in self._layout] |
| 116 | + params: dict[str, dict[str, Any]] = {} |
| 117 | + for sub_cfg in cfg.sub_terrains.values(): |
| 118 | + gtype = sub_cfg.genesis_type() |
| 119 | + p = sub_cfg.to_genesis_params() |
| 120 | + if p: |
| 121 | + params[gtype] = p |
| 122 | + kwargs: dict[str, Any] = { |
| 123 | + "n_subterrains": (cfg.num_rows, cfg.num_cols), |
| 124 | + "subterrain_size": cfg.subterrain_size, |
| 125 | + "horizontal_scale": cfg.horizontal_scale, |
| 126 | + "vertical_scale": cfg.vertical_scale, |
| 127 | + "subterrain_types": types_2d, |
| 128 | + "pos": cfg.pos, |
| 129 | + } |
| 130 | + if params: |
| 131 | + kwargs["subterrain_parameters"] = params |
| 132 | + return kwargs |
0 commit comments