Repository: Genesis-Embodied-AI/Genesis
File: genesis/utils/terrain.py (Python)

Code before:
    Returns
    -------
    heights : (nx, ny)  float32 ndarray
        Highest elevation per coarse cell (NaN if no hit).
    xs      : (nx,)     float32 ndarray
    ys      : (ny,)     float32 ndarray
        Coordinates of grid lines (cell centres).

    Notes
    -----
    • Requires `pyembree` or `rtree` for fast ray‑tracing.
    • Memory cost grows as  oversample².  For very large terrains keep
      `oversample` ≤ 4 unless you have ample RAM.
    """
    mesh = trimesh.load(path, force="mesh")

    # -------------------------------- axis handling ---------------------------
    if up_axis.lower() == "y":
        # rotate so Z becomes up (‑90° around X)
        T = trimesh.transformations.rotation_matrix(np.deg2rad(-90), [1, 0, 0])
        mesh.apply_transform(T)

    (minx, miny, _), (maxx, maxy, maxz) = mesh.bounds

    # -------------------------------- coarse grid ----------------------------
    dx, dy = (spacing, spacing) if isinstance(spacing, (int, float)) else spacing
    nx = math.ceil((maxx - minx) / dx) + 1
    ny = math.ceil((maxy - miny) / dy) + 1

    xs = np.linspace(minx, maxx, nx, dtype=np.float32)
    ys = np.linspace(miny, maxy, ny, dtype=np.float32)

    # -------------------------------- fine grid ------------------------------
    fx = nx * oversample
    fy = ny * oversample
    fx_lin = np.linspace(minx, maxx, fx, dtype=np.float32)
    fy_lin = np.linspace(miny, maxy, fy, dtype=np.float32)
    fxx, fyy = np.meshgrid(fx_lin, fy_lin, indexing="ij")

    origins = np.c_[fxx.ravel(), fyy.ravel(), np.full(fxx.size, maxz + 1.0)]
    directions = np.tile([0.0, 0.0, -1.0], (origins.shape[0], 1))

    # -------------------------------- ray cast -------------------------------
    locs, hit_ids, _ = mesh.ray.intersects_location(ray_origins=origins, ray_directions=directions, multiple_hits=False)

    h_fine = np.full(fxx.size, np.nan, dtype=np.float32)
    h_fine[hit_ids] = locs[:, 2]
    h_fine = h_fine.reshape((fx, fy))  # (fx, fy) = (nx*k, ny*k)

    # -------------------------------- down‑sample ----------------------------
    # reshape to (nx, k, ny, k) then take max over the 2 fine axes

Review comment (suggestion):
It would be nice to avoid duplicating every lines for x and y:
```python
from collections.abc import ValuesView, Sequence

if np.isscalar(spacing):
    spacing = (spacing, spacing)

n, s, f, f_lin = [None, None], [None, None], [None, None], [None, None]
for i, (vmin, vmax, delta) in range(2):
    # -------------------------------- coarse grid ----------------------------
    vmin, vmax = mesh.bounds[0][i], mesh.bounds[1][i]
    n[i] = math.ceil((vmax - vmin) / spacing[i]) + 1
    s[i] = np.linspace(vmin, vmax, n[i], dtype=np.float32)

    # -------------------------------- fine grid ------------------------------
    f[i] =  n[i] * oversample
    f_lin[i] = np.linspace(vmin, vmax, f[i], dtype=np.float32)

f_grid = np.meshgrid(*f_lin, indexing="ij")
```

Code after:

    Returns
    -------
    heights : (nx, ny)  float32 ndarray
        Highest elevation per coarse cell (NaN if no hit).
    xs      : (nx,)     float32 ndarray
    ys      : (ny,)     float32 ndarray
        Coordinates of grid lines (cell centres).

    Notes
    -----
    • Memory cost grows as  oversample².
    """
    if np.isscalar(spacing):
        spacing = (spacing, spacing)

    mesh = trimesh.load(path, force="mesh")

    # -------------------------------- axis handling ---------------------------
    if up_axis.lower() == "y":
        # rotate so Z becomes up (‑90° around X)
        T = trimesh.transformations.rotation_matrix(np.deg2rad(-90), [1, 0, 0])
        mesh.apply_transform(T)

    (minx, miny, _), (maxx, maxy, maxz) = mesh.bounds

    # -------------------------------- coarse grid ----------------------------
    dx, dy = spacing
    nx = math.ceil((maxx - minx) / dx) + 1
    ny = math.ceil((maxy - miny) / dy) + 1

    xs = np.linspace(minx, maxx, nx, dtype=np.float32)
    ys = np.linspace(miny, maxy, ny, dtype=np.float32)

    # -------------------------------- fine grid ------------------------------
    fx = nx * oversample
    fy = ny * oversample
    fx_lin = np.linspace(minx, maxx, fx, dtype=np.float32)
    fy_lin = np.linspace(miny, maxy, fy, dtype=np.float32)
    fxx, fyy = np.meshgrid(fx_lin, fy_lin, indexing="ij")

    origins = np.stack((fxx.ravel(), fyy.ravel(), np.full(fxx.size, maxz + 1.0)), axis=-1)
    directions = np.tile([0.0, 0.0, -1.0], (origins.shape[0], 1))

    # -------------------------------- ray cast -------------------------------
    locs, hit_ids, _ = mesh.ray.intersects_location(ray_origins=origins, ray_directions=directions, multiple_hits=False)

    h_fine = np.full(fxx.size, np.nan, dtype=np.float32)
    h_fine[hit_ids] = locs[:, 2]
    h_fine = h_fine.reshape((fx, fy))  # (fx, fy) = (nx*k, ny*k)