Blame for code/snow_rl/belief.py

"""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])