"""Local terrain belief map (ported from hackathon-everest's `mapping.py`). This is the half of the problem the RL repo was missing. Without it the policy sees only what is under its feet *right now*; with it, each foot's observation accumulates into a spatial memory that the OTHER foot can read before it lands. That cross-foot transfer is everest's central claim. Notably, everest's own benchmark could not measure whether the map helped: decomposing its 7 unsafe transfers showed all 7 were stance-capacity failures, which the bilateral gate catches and the map has nothing to do with. Here the map is a separate observation channel with its own ablation flag, so its contribution is directly measurable for the first time. Differences from everest's version, all forced by running inside a 50 Hz JAX training loop: * **Two channels, not six.** everest maps six quantities; only support and depth reach the policy, so mapping the rest would cost memory for nothing. * **Coarser.** 40x10 cells at 0.2 m over the corridor, versus everest's 40x40 at 0.05 m. The map lives per-environment in `state.info`, so at 8192 envs this is ~52 MB; everest's resolution would be ~26x that. * **Reflect-tiled**, matching terrain.sample, so map and terrain stay in the same frame. * Update is the same Gaussian-kernel Kalman step: a measurement at one point updates neighbours with weight falling off over a correlation length. """ from __future__ import annotations from typing import NamedTuple import jax import jax.numpy as jnp from snow_rl.terrain import DEFAULT_CONFIG, TerrainConfig, _reflect MAP_ROWS = 40 # along the corridor MAP_COLS = 10 # across it CORRELATION_LENGTH_M = 0.35 MIN_WEIGHT = 0.01 # below this a cell is left untouched # Priors, matching everest's mapping.py defaults where they overlap. PRIOR_SUPPORT_N = 330.0 PRIOR_SUPPORT_STD_N = 190.0 PRIOR_DEPTH_M = 0.08 PRIOR_DEPTH_STD_M = 0.07 SUPPORT_NORM_N = 900.0 DEPTH_NORM_M = 0.25 # Where the policy reads the map: offsets ahead of each foot, along the corridor. LOOKAHEAD_M = (0.0, 0.25, 0.50) READOUT_PER_FOOT = len(LOOKAHEAD_M) * 4 # mean+std for support and depth READOUT_SIZE = 2 * READOUT_PER_FOOT # both feet class BeliefMap(NamedTuple): """Per-environment spatial memory. Leading axis is the env batch under vmap.""" support_mean: jax.Array # (MAP_ROWS, MAP_COLS) N support_var: jax.Array depth_mean: jax.Array # m depth_var: jax.Array def init() -> BeliefMap: shape = (MAP_ROWS, MAP_COLS) return BeliefMap( support_mean=jnp.full(shape, PRIOR_SUPPORT_N), support_var=jnp.full(shape, PRIOR_SUPPORT_STD_N**2), depth_mean=jnp.full(shape, PRIOR_DEPTH_M), depth_var=jnp.full(shape, PRIOR_DEPTH_STD_M**2), ) def _fractional_cell(xy_m: jax.Array, cfg: TerrainConfig): """World position -> fractional cell index, reflect-tiled. Everything in this module works in CELL space and goes through this one function, so `update` and `readout` cannot disagree. They previously did: update measured distances with a periodic wrap while readout reflected, and the two used different half-cell offsets. Beyond the corridor edge the map was written to one cell and read from another, so a readout returned the untouched prior. """ fx = _reflect((xy_m[..., 0] + cfg.length_m / 2) / (cfg.length_m / MAP_ROWS) - 0.5, MAP_ROWS) fy = _reflect((xy_m[..., 1] + cfg.width_m / 2) / (cfg.width_m / MAP_COLS) - 0.5, MAP_COLS) return fx, fy def _kalman(mean, var, weight, observed, observed_std): """Gaussian-kernel Kalman update: distant cells see a larger effective observation variance, so a measurement informs its neighbourhood and not the whole map.""" effective_var = observed_std**2 / jnp.maximum(weight, 1e-6) gain = jnp.where(weight >= MIN_WEIGHT, var / (var + effective_var), 0.0) return mean + gain * (observed - mean), (1.0 - gain) * var def update( belief: BeliefMap, xy_m: jax.Array, # (2,) where the measurement was taken support_n: jax.Array, # () support_std_n: jax.Array, depth_m: jax.Array, depth_std_m: jax.Array, cfg: TerrainConfig = DEFAULT_CONFIG, ) -> BeliefMap: """Fuse one foot's estimate into the map.""" fx, fy = _fractional_cell(xy_m, cfg) rows = jnp.arange(MAP_ROWS)[:, None] cols = jnp.arange(MAP_COLS)[None, :] dx = (rows - fx) * (cfg.length_m / MAP_ROWS) dy = (cols - fy) * (cfg.width_m / MAP_COLS) weight = jnp.exp(-(dx**2 + dy**2) / (2.0 * CORRELATION_LENGTH_M**2)) sm, sv = _kalman(belief.support_mean, belief.support_var, weight, support_n, jnp.maximum(support_std_n, 1.0)) dm, dv = _kalman(belief.depth_mean, belief.depth_var, weight, depth_m, jnp.maximum(depth_std_m, 1e-3)) return BeliefMap(sm, sv, jnp.clip(dm, 0.0, DEPTH_NORM_M), dv) def _index(xy_m: jax.Array, cfg: TerrainConfig): """Nearest cell to a world position. Shares _fractional_cell with `update`.""" fx, fy = _fractional_cell(xy_m, cfg) return (jnp.clip(jnp.round(fx), 0, MAP_ROWS - 1).astype(int), jnp.clip(jnp.round(fy), 0, MAP_COLS - 1).astype(int)) def readout(belief: BeliefMap, feet_xy_m: jax.Array, cfg: TerrainConfig = DEFAULT_CONFIG) -> jax.Array: """What the policy sees: belief ahead of each foot, normalised. Returns (READOUT_SIZE,) laid out per foot, per lookahead offset, as [support_mean, support_std, depth_mean, depth_std]. """ out = [] for foot in range(feet_xy_m.shape[0]): for ahead in LOOKAHEAD_M: probe = feet_xy_m[foot] + jnp.array([ahead, 0.0]) i, j = _index(probe, cfg) out.extend([ belief.support_mean[i, j] / SUPPORT_NORM_N, jnp.sqrt(belief.support_var[i, j]) / SUPPORT_NORM_N, belief.depth_mean[i, j] / DEPTH_NORM_M, jnp.sqrt(belief.depth_var[i, j]) / DEPTH_NORM_M, ]) return jnp.stack(out) def lower_confidence_support_n(belief: BeliefMap, xy_m: jax.Array, cfg: TerrainConfig = DEFAULT_CONFIG, sigmas: float = 2.0) -> jax.Array: """Conservative support at a point: mean minus `sigmas` standard deviations. everest ranks candidate footholds on this rather than the mean, so uncertainty costs a candidate rather than being averaged away. """ i, j = _index(xy_m, cfg) return belief.support_mean[i, j] - sigmas * jnp.sqrt(belief.support_var[i, j])