code/snow_rl/terrain.py

"""Per-environment correlated snow-depth fields (Step 2, part 1).

Ported from hackathon-everest's `terrain.py`, with two deliberate changes:

* **Snow is not geometry.** everest bakes terrain into a spatial grid it samples on the CPU;
  Playground's rough-terrain env bakes it into a MuJoCo heightfield. Neither works here.
  A heightfield models bumpy *solid* ground, but snow depth is a surface you sink *into* --
  different physics. So the depth field lives as a plain JAX array and never becomes a geom.
  The MuJoCo ground stays flat and represents the hard substrate under the snow.
* **Fixed correlation length.** everest draws a per-field correlation length. Doing that per
  environment would need a per-env (H,H) blur matrix (~419 MB at 4096 envs x 160 cells).
  The blur matrices here are shared constants and per-env variety comes from the noise,
  the regime, and the amplitude instead. Contiguity -- the property the brief actually
  requires -- is preserved.

Memory: a 160x40 corridor at 4096 envs is ~105 MB, versus ~1.07 GB for per-env heightfields.
"""
from __future__ import annotations

from typing import NamedTuple

import jax
import jax.numpy as jnp

from snow_rl import sinkage

# Regime mixture from the brief: 30% flat, 40% shallow, 30% deep, mixed WITHIN each batch
# so the policy learns a conditional response rather than an averaged gait.
REGIME_EDGES = (0.30, 0.70)
REGIME_DEPTH_RANGE_M = (
    (0.000, 0.010),  # flat
    (0.000, 0.100),  # shallow
    (0.100, 0.250),  # deep
)
MAX_DEPTH_M = 0.25


class TerrainConfig(NamedTuple):
    length_m: float = 8.0        # along the walking corridor (x)
    width_m: float = 2.0         # across it (y)
    cell_size_m: float = 0.05
    correlation_length_m: float = 0.6

    @property
    def n_rows(self) -> int:
        return round(self.length_m / self.cell_size_m)

    @property
    def n_cols(self) -> int:
        return round(self.width_m / self.cell_size_m)


DEFAULT_CONFIG = TerrainConfig()


class SnowField(NamedTuple):
    """One environment's snow. Leading axis is the env batch under vmap."""
    depth_m: jax.Array               # (H, W) snow thickness above the hard substrate
    bearing_n: jax.Array             # (H, W) load the patch supports before failing
    stiffness_n_per_m: jax.Array     # (H, W)
    damping_ns_per_m: jax.Array      # (H, W)
    void: jax.Array                  # (H, W) bool, hollow under a crust
    regime: jax.Array                # () 0 flat, 1 shallow, 2 deep -- diagnostics only


# --- bearing capacity, derived from the sinkage law rather than generated independently ---
#
# An earlier version drew `bearing_n` from its own material noise, which made it statistically
# INDEPENDENT of snow depth (measured corr = -0.01). That contradicted this module's own
# physics: snow_rl.sinkage already gives shallow snow far more support, because the foot
# bottoms out on the substrate (the `3k(d - d_snow)` term). The label did not describe the
# force law it was labelling, and a support estimate therefore carried no information about
# how high to lift a foot.
#
# Bearing is now defined as the resisting force at a reference sinkage -- "the load this
# foothold carries before the foot has sunk 3 cm". That is consistent with the physics by
# construction, and reproduces the real coupling: thin snow over a hard substrate is strong,
# deep powder is weak.
#
# The log map is fitted so the distribution stays commensurate with the bearing_capacity_n
# everest's estimator was calibrated against (median 363 N vs everest's 383 N, no clipping).
FAILURE_SINKAGE_M = 0.03
BEARING_LOG_GAIN = 175.0
BEARING_LOG_OFFSET_N = 50.0
BEARING_RANGE_N = (100.0, 860.0)


def bearing_capacity_n(depth_m: jax.Array, stiffness_n_per_m: jax.Array) -> jax.Array:
    """Load a foothold carries before sinking past FAILURE_SINKAGE_M."""
    raw = sinkage.resistance_n(
        jnp.full_like(depth_m, FAILURE_SINKAGE_M), jnp.zeros_like(depth_m),
        depth_m, stiffness_n_per_m, jnp.zeros_like(depth_m),
        jnp.zeros_like(depth_m, dtype=bool),
    )
    scaled = BEARING_LOG_GAIN * jnp.log1p(raw / BEARING_LOG_OFFSET_N)
    return jnp.clip(scaled, *BEARING_RANGE_N)


def _blur_matrix(n: int, sigma_cells: float) -> jax.Array:
    """Row-normalised Gaussian smoothing operator. Shared across envs, so it costs nothing."""
    i = jnp.arange(n)
    w = jnp.exp(-0.5 * ((i[:, None] - i[None, :]) / max(sigma_cells, 0.5)) ** 2)
    return w / jnp.sum(w, axis=1, keepdims=True)


def _correlated_unit_field(key: jax.Array, cfg: TerrainConfig) -> jax.Array:
    """Spatially correlated noise mapped to [0, 1] by quantile stretch (as everest does)."""
    sigma = cfg.correlation_length_m / cfg.cell_size_m
    raw = jax.random.normal(key, (cfg.n_rows, cfg.n_cols))
    smooth = _blur_matrix(cfg.n_rows, sigma) @ raw @ _blur_matrix(cfg.n_cols, sigma).T
    lo, hi = jnp.quantile(smooth, jnp.array([0.02, 0.98]))
    return jnp.clip((smooth - lo) / jnp.maximum(hi - lo, 1e-9), 0.0, 1.0)


def generate(key: jax.Array, cfg: TerrainConfig = DEFAULT_CONFIG,
             depth_scale: float = 1.0) -> SnowField:
    """Generate one environment's snow field. vmap over keys for a batch.

    depth_scale shrinks every regime's depth range, which is how the training curriculum
    works. Measured, a policy warm-started from flat ground handles snow up to ~70 mm well
    (60% velocity tracking, no falls) and collapses at ~200 mm (-13% tracking, falls in 85
    steps). Starting at full difficulty it converges on the easy 70% of terrain and treats
    the deep third as unsolvable noise. Ramping the ceiling gives it a reachable next step.
    """
    k_reg, k_depth, k_stiff, k_dens, k_wet, k_void = jax.random.split(key, 6)

    u = jax.random.uniform(k_reg, ())
    regime = (u >= REGIME_EDGES[0]).astype(jnp.int32) + (u >= REGIME_EDGES[1]).astype(jnp.int32)
    lows = jnp.array([r[0] for r in REGIME_DEPTH_RANGE_M])
    highs = jnp.array([r[1] for r in REGIME_DEPTH_RANGE_M])
    depth = depth_scale * (
        lows[regime] + _correlated_unit_field(k_depth, cfg) * (highs[regime] - lows[regime])
    )


    # Stiffness is COUPLED TO DEPTH, not independent of it. An earlier version drew the two
    # from separate correlated fields, so training terrain paired deep snow with any stiffness
    # while the evaluation strip deliberately pairs deep-with-soft -- the policy was tested on
    # a case it had only seen by chance. Physically the coupling is the realistic option too:
    # a thin wind-packed layer over substrate is firm, deep fresh snow is not.
    #
    # The material field still supplies most of the variation, so depth does not determine
    # stiffness; it shifts the distribution.
    depth_factor = jnp.clip(depth / MAX_DEPTH_M, 0.0, 1.0)
    log_lo = jnp.log(2_500.0)
    log_hi = jnp.log(45_000.0) + (jnp.log(6_000.0) - jnp.log(45_000.0)) * depth_factor
    stiffness = jnp.exp(log_lo + _correlated_unit_field(k_stiff, cfg) * (log_hi - log_lo))
    density = 180.0 + _correlated_unit_field(k_dens, cfg) * (650.0 - 180.0)
    wetness = _correlated_unit_field(k_wet, cfg) * 0.55
    damping = jnp.clip(18.0 + 0.07 * density + 85.0 * wetness, 20.0, 140.0)
    bearing = bearing_capacity_n(depth, stiffness)

    # Voids only exist under snow deep enough to bridge one.
    void = (_correlated_unit_field(k_void, cfg) > 0.88) & (depth > 0.08)

    return SnowField(depth, bearing, stiffness, damping, void, regime)


def crossing_strip(cfg: TerrainConfig = DEFAULT_CONFIG,
                   max_depth_m: float = 0.190) -> SnowField:
    """Deterministic flat -> shallow -> deep corridor for Step 6 evaluation.

    Fixed and seedless so every eval run crosses identical ground.

    **Stiffness varies with depth, not just depth alone.** An earlier version held stiffness
    at a uniform 12 kN/m, which made the deep section nearly as easy as the shallow one: at
    that stiffness the snow carries body weight after ~4 cm, so a foot sinks 42 mm in both
    70 mm and 190 mm of snow. The deep end was deep but firm -- visually undramatic, and more
    importantly it gave the policy almost nothing extra to learn. Deep snow in the world is
    usually softer, so stiffness now falls from firm windpack to soft powder as depth grows:
    the foot sinks ~12 mm on the flat and ~144 mm in the deep, which is the behaviour the
    clearance reward is supposed to be teaching a policy to anticipate.
    """
    x = jnp.linspace(0.0, 1.0, cfg.n_rows)[:, None]

    def smoothstep(edge0, edge1, value):
        t = jnp.clip((value - edge0) / (edge1 - edge0), 0.0, 1.0)
        return t * t * (3.0 - 2.0 * t)

    # Smooth transitions, not steps. Hard `where` boundaries put a vertical wall of snow in
    # the corridor and a discontinuous jump in the sinkage force.
    shallow = smoothstep(0.28, 0.38, x)
    deep = smoothstep(0.61, 0.71, x)
    mid = 0.005 + (max_depth_m - 0.005) * 0.37
    depth = 0.005 + (mid - 0.005) * shallow + (max_depth_m - mid) * deep
    # stiffness falls with depth, as in the generated terrain
    frac = jnp.clip(max_depth_m / 0.190, 0.0, 1.0)
    deep_k = 28_000.0 + (3_000.0 - 28_000.0) * frac
    mid_k = 28_000.0 + (12_000.0 - 28_000.0) * frac
    stiffness = jnp.broadcast_to(
        28_000.0 + (mid_k - 28_000.0) * shallow + (deep_k - mid_k) * deep,
        (cfg.n_rows, cfg.n_cols),
    )
    depth = jnp.broadcast_to(depth, (cfg.n_rows, cfg.n_cols))
    damping = jnp.full_like(depth, 60.0)
    return SnowField(depth, bearing_capacity_n(depth, stiffness), stiffness, damping,
                     jnp.zeros_like(depth, bool), jnp.int32(1))


def _reflect(index: jax.Array, n: int) -> jax.Array:
    """Fold an unbounded cell index into [0, n-1] by reflection, so the field tiles seamlessly."""
    period = 2.0 * (n - 1)
    folded = jnp.mod(index, period)
    return jnp.where(folded > n - 1, period - folded, folded)


def sample(field: SnowField, cfg: TerrainConfig, xy_m: jax.Array) -> tuple[jax.Array, ...]:
    """Bilinearly sample the field at world positions.

    xy_m: (..., 2) with x along the corridor, y across it, origin at the corridor centre.
    Returns (depth_m, bearing_n, stiffness, damping, void_prob) each shaped (...,).
    """
    # Reflect-tile rather than clamp. A 20 s episode at up to 1.0 m/s covers ~20 m, but the
    # corridor is 8 m -- clamping pinned a walking policy to the boundary row for most of every
    # episode, feeding it one frozen depth value. Reflection (triangle wave) makes the field
    # effectively infinite with no seam discontinuity, which plain modulo tiling would create.
    fx = _reflect((xy_m[..., 0] + cfg.length_m / 2) / cfg.cell_size_m - 0.5, cfg.n_rows)
    fy = _reflect((xy_m[..., 1] + cfg.width_m / 2) / cfg.cell_size_m - 0.5, cfg.n_cols)
    i0, j0 = jnp.floor(fx).astype(int), jnp.floor(fy).astype(int)
    i1 = jnp.minimum(i0 + 1, cfg.n_rows - 1)
    j1 = jnp.minimum(j0 + 1, cfg.n_cols - 1)
    tx, ty = fx - i0, fy - j0

    def bilinear(grid: jax.Array) -> jax.Array:
        g = grid.astype(float)
        return (
            g[i0, j0] * (1 - tx) * (1 - ty) + g[i1, j0] * tx * (1 - ty)
            + g[i0, j1] * (1 - tx) * ty + g[i1, j1] * tx * ty
        )

    return tuple(bilinear(g) for g in
                 (field.depth_m, field.bearing_n, field.stiffness_n_per_m,
                  field.damping_ns_per_m, field.void))


batched_generate = jax.vmap(generate, in_axes=(0, None, None))
batched_sample = jax.vmap(sample, in_axes=(0, None, 0))