code/snow_rl/snow_env.py

"""Snow locomotion environment (Step 4).

Wraps MuJoCo Playground's G1 joystick env. Three changes, nothing else:

1. **Snow forces.** Each step, the per-environment depth field is sampled at both feet and a
   sinkage + ploughing force is written into `data.xfrc_applied`. The MuJoCo ground stays a
   flat plane representing the hard substrate -- snow is never geometry (see terrain.py).
2. **Observations.** The 8 estimator values plus a 3-step history are appended to the actor's
   `state` observation. The critic's `privileged_state` additionally receives true depth and
   true bearing per foot -- training only, discarded at deployment.
3. **Rewards.** Three terms on top of Playground's stock G1 terms, all scaled by confidence.

The design rule, from the brief: **the reward reads the estimate, never the truth.** Truth is
permitted only as the corruption function's input, in the critic, and in physics-consequence
terms (stumble). If the clearance target read true depth, the policy could succeed while
ignoring the sensor entirely and the experiment would prove nothing.

Terrain is stored as a shared BANK of pre-generated fields with a per-environment index,
rather than a field per environment. A field in `state.info` would be carried through every
rollout transition -- 160x40 floats x 8192 envs x unroll 20 is several GB. The bank costs
~1.6 MB total and an int per environment.
"""
from __future__ import annotations

from typing import Any

import jax
import jax.numpy as jnp
import numpy as np
from ml_collections import config_dict
from mujoco import mjx
from mujoco_playground._src import mjx_env
from mujoco_playground._src.locomotion.g1 import joystick

from snow_rl import belief, estimator, sinkage, support
from snow_rl.terrain import DEFAULT_CONFIG as TERRAIN_CFG
from snow_rl.terrain import SnowField, batched_generate, sample

TERRAIN_BANK_SIZE = 64
HISTORY_STEPS = 3
CLEARANCE_MARGIN_M = 0.05
CLEARANCE_TOLERANCE = 0.01      # Gaussian width on the clearance target
FLAT_DEPTH_THRESHOLD_M = 0.03   # below this the estimate is calling the ground firm
FLAT_NOMINAL_CLEARANCE_M = 0.07
# Absolute plough threshold, not a ratio. The brief specified "horizontal force > 2x
# vertical", which is right for CONTACT forces where a dragging foot can exceed vertical
# load. In this model vertical support is a stiff spring and ploughing is a small tangential
# drag: swept over the whole parameter space the ratio maxes out at 0.076 against a threshold
# of 2.0, so the term could never fire and was identically 0.000 across a 60M-step run.
# ~15 N of drag on a swinging foot means it is meaningfully submerged.
STUMBLE_PLOUGH_N = 15.0
# A foot is "in contact" when snow carries this much load -- ~6% of body weight. Floor
# contact alone is not enough once snow is a force rather than geometry (see
# _stepped_with_snow_contact).
SNOW_CONTACT_FORCE_N = 20.0
# MJX constraint-row budget. mujoco_warp's default heuristic is
#   njmax = max(nv*2.26*(nhfield>0)*18 + 53, ...)
# which collapses to 53 -> rounded to 64 for any model without a heightfield. The G1 needs
# 93-94 rows once the gait is dynamic (standing is only 61, which is why this appears mid-run
# and not at reset). Past the limit MuJoCo SILENTLY DROPS contact constraints -- feet lose
# grip or sink through -- while printing "nefc overflow" and continuing to train.
# Playground's stock value is 90. 128 is the next valid size above the observed 93-94 peak.
NJMAX = 128
# Slope. Everything before this was level-ground walking, which is a real gap for a
# mountaineering project: on the flat, losing a foothold means falling over; on a slope it
# means sliding. The source project's terrain generator had a slope term (+/-0.25) that was
# dropped when the field was ported to JAX, because our snow is a force field over a flat
# MuJoCo plane rather than geometry.
#
# Rather than tilt the ground mesh, tilt GRAVITY. A uniform slope is exactly equivalent to
# rotating the gravity vector, it costs nothing in the physics, and it keeps the snow force
# field (which is defined against a flat substrate) valid without reworking terrain.py.
SLOPE_RAD = 0.0
# Uncertainty the belief map attaches to one contact observation. The estimator's own
# ensemble spread is only weakly informative (corr with true error +0.35), so a fixed
# conservative value is more honest than pretending the spread is calibrated.
MAP_OBS_STD_N = 70.0
MAP_OBS_STD_M = 0.05


def default_config() -> config_dict.ConfigDict:
    """Playground's G1 config plus the snow terms. Stock values are untouched."""
    cfg = joystick.default_config()
    # Measured on the first 60M-step run, the snow terms were 6.1% of total episode reward
    # (clearance 6.8%, flat_efficiency 0.1%, overload 0.6%) -- Playground's stock objective
    # accounted for the other 94%, so the policy had almost no gradient pressure to change
    # its gait for snow. Scaled up ~3.5x to reach roughly a quarter of total reward: enough
    # to steer, short of drowning velocity tracking (whose failure mode is a policy that
    # stands still or creeps).
    # Back to the brief's original weights. Five tuning cycles (0.8 -> 3 -> 5 on clearance,
    # 0.4 -> 1.5 -> 3.0 on overload) never beat them: measured velocity tracking went
    # 18% -> -5% -> -8% while survival rose, because not walking maximises survival. The
    # bottleneck was never the reward weights -- it is that a flat-ground policy is thrown
    # into 200 mm snow from step 0, which the curriculum now addresses instead.
    cfg.reward_config.scales.snow_clearance = 0.8
    cfg.reward_config.scales.snow_stumble = -1.0
    cfg.reward_config.scales.snow_flat_efficiency = -0.5
    cfg.reward_config.scales.snow_overload = -0.4
    # Playground's stock njmax is 90 constraint rows; a dynamic G1 gait needs 93-94 and
    # MuJoCo SILENTLY DROPS contacts past the limit. This line has been deleted twice by
    # index-based edits to the reward block above -- if you rewrite those weights, check it
    # survived. `make_snow_env()._config.njmax` should read 128, not 90.
    cfg.njmax = NJMAX
    cfg.snow = config_dict.create(
        enable_forces=True,
        # Zeroing this blanks the estimator block in the observation while leaving snow
        # physics and rewards intact. This is the Step 6 ablation that isolates the sensor
        # channel -- see docs/PIPELINE.md section 8.
        estimator_channel=1.0,
        # Separate flags so each half's contribution is measurable on its own. everest's
        # benchmark could not separate its map from its support gate; this can.
        map_channel=1.0,
        support_channel=1.0,
        false_safe_rate=-1.0,   # <0 means "use the calibrated point-estimate rate"
        terrain_bank_size=TERRAIN_BANK_SIZE,
        terrain_seed=0,
        # Uphill slope in radians. 0.20 rad ~ 11.5 deg, a realistic snow slope; avalanche
        # terrain starts around 0.5 rad (30 deg). Applied by rotating gravity, so positive
        # values mean the robot walks uphill along +x.
        slope_rad=SLOPE_RAD,
        # Curriculum: 1.0 is full difficulty (0-25 cm). Lower values shrink every regime's
        # depth range so an early stage trains on snow the policy can actually cope with.
        depth_scale=1.0,
    )
    return cfg


class SnowJoystick(joystick.Joystick):
    """G1 joystick locomotion over variable-depth snow."""

    def __init__(
        self,
        task: str = "flat_terrain",
        config: config_dict.ConfigDict | None = None,
        config_overrides: dict[str, Any] | None = None,
    ):
        # Playground downloads mujoco_menagerie inside `locomotion.load()`. Constructing
        # Joystick directly bypasses that, so a fresh machine has no robot meshes and model
        # compilation fails on a missing STL. This does not reproduce anywhere menagerie is
        # already cached -- which is every development machine, and no fresh container.
        mjx_env.ensure_menagerie_exists()
        super().__init__(task=task, config=config or default_config(),
                         config_overrides=config_overrides)

    def _post_init(self) -> None:
        super()._post_init()
        slope = float(self._config.snow.slope_rad)
        if slope != 0.0:
            # Rotate gravity about +y: the robot walking along +x is then walking uphill.
            g = float(np.linalg.norm(self._mj_model.opt.gravity))
            # NEGATIVE x component: gravity pulls the robot BACK along -x while the
            # walking command drives +x, so positive slope_rad means walking UPHILL.
            # The sign was inverted first time round, which made gravity assist the command
            # -- passive drift measured +1.02 m downhill at 17 deg, and "slope performance"
            # improved with steepness because the robot was being pushed along.
            tilted = np.array([-g * np.sin(slope), 0.0, -g * np.cos(slope)])
            self._mj_model.opt.gravity[:] = tilted
            self._mjx_model = self._mjx_model.replace(opt=self._mjx_model.opt.replace(
                gravity=jnp.asarray(tilted)))
        self._terrain_cfg = TERRAIN_CFG
        rate = self._config.snow.false_safe_rate
        self._est_params = estimator.load_params(
            false_safe_rate=None if rate < 0 else rate
        )
        # Feet bodies carry the snow force; site ids come from the parent env.
        self._feet_body_id = jnp.asarray(self._mj_model.site_bodyid[self._feet_site_id])
        self._terrain_bank: SnowField = batched_generate(
            jax.random.split(jax.random.key(self._config.snow.terrain_seed),
                             self._config.snow.terrain_bank_size),
            self._terrain_cfg,
            float(self._config.snow.depth_scale),
        )

    # ---------------- terrain ----------------

    def _field_at(self, index: jax.Array) -> SnowField:
        return jax.tree_util.tree_map(lambda leaf: leaf[index], self._terrain_bank)

    def _sample_feet(self, field: SnowField, feet_xyz: jax.Array):
        """Snow state under each foot. Returns (depth, bearing, stiffness, damping, void)."""
        return sample(field, self._terrain_cfg, feet_xyz[:, :2])

    # ---------------- episode ----------------

    def reset(self, rng: jax.Array) -> mjx_env.State:
        rng, terrain_rng, est_rng = jax.random.split(rng, 3)
        state = super().reset(rng)

        index = jax.random.randint(terrain_rng, (), 0, self._config.snow.terrain_bank_size)
        est_state = estimator.init_state(est_rng, self._est_params)
        feet_xyz = state.data.site_xpos[self._feet_site_id]
        field = self._field_at(index)
        depth, bearing, _, _, void = self._sample_feet(field, feet_xyz)

        obs8, est_state = estimator.estimate(
            est_rng, est_state, depth, bearing, void > 0.5,
            jnp.ones((2,), dtype=bool), self._est_params,
        )
        belief_map = belief.init()
        state.info["snow_belief"] = belief_map
        state.info["snow_map_readout"] = belief.readout(belief_map, feet_xyz[:, :2])
        state.info["snow_support_obs"] = support.observation(
            support.evaluate(bearing * 0.0 + support.G1_WEIGHT_N,
                             jnp.zeros(2), jnp.zeros(2))
        )
        state.info["snow_overload"] = jnp.zeros(())
        state.info["snow_terrain_index"] = index
        state.info["snow_obs"] = obs8
        state.info["snow_history"] = jnp.tile(obs8, HISTORY_STEPS)
        state.info["snow_truth"] = jnp.concatenate([depth, bearing])
        state.info["snow_force"] = jnp.zeros((2, 3))
        state.info["snow_last_feet_pos"] = feet_xyz
        state.info["snow_est_last_obs"] = est_state.last_obs
        state.info["snow_est_false_safe"] = est_state.false_safe
        state.info["snow_est_noise_scale"] = est_state.noise_scale

        state = state.replace(obs=self._get_obs(state.data, state.info,
                                                state.info["last_contact"]))
        for key in ("snow_clearance", "snow_stumble", "snow_flat_efficiency",
                    "snow_overload"):
            state.metrics[f"reward/{key}"] = jnp.zeros(())
        state.metrics["snow_depth_under_feet"] = jnp.mean(depth)
        state.metrics["snow_contact_fraction"] = jnp.zeros(())
        return state

    def step(self, state: mjx_env.State, action: jax.Array) -> mjx_env.State:
        field = self._field_at(state.info["snow_terrain_index"])
        feet_xyz = state.data.site_xpos[self._feet_site_id]
        # Finite-difference foot velocity; MJX site velocities are not directly exposed.
        feet_vel = (feet_xyz - state.info["snow_last_feet_pos"]) / self.dt
        depth, bearing, stiffness, damping, void = self._sample_feet(field, feet_xyz)
        void_b = void > 0.5

        # Force for the estimator/reward at this control instant. The force actually applied
        # to the physics is recomputed each substep inside _physics_with_snow.
        force = sinkage.foot_force_world_n(
            feet_xyz[:, 2], feet_vel, jnp.zeros((2,)), depth, stiffness, damping, void_b
        )
        force = force * self._config.snow.enable_forces

        # A foot that has just landed starts a new stance -- the false-safe latch resamples
        # here rather than every tick, which would inflate the per-footstep rate ~20x.
        state.info["rng"], est_rng = jax.random.split(state.info["rng"])
        new_stance = self._bearing_load(force) & ~state.info["last_contact"].astype(bool)
        est_state = estimator.EstimatorState(
            last_obs=state.info["snow_est_last_obs"],
            false_safe=state.info["snow_est_false_safe"],
            noise_scale=state.info["snow_est_noise_scale"],
        )
        obs8, est_state = estimator.estimate(
            est_rng, est_state, depth, bearing, void_b, new_stance, self._est_params
        )

        # Fuse each foot's estimate into the spatial memory. Reads the ESTIMATE, not truth.
        est = obs8.reshape(estimator.N_FEET, estimator.N_CHANNELS)
        support_est_n = est[:, 1] * estimator.SUPPORT_NORM_N
        implied = estimator.implied_depth_m(est, self._est_params)
        belief_map = state.info["snow_belief"]
        for foot in range(estimator.N_FEET):
            belief_map = belief.update(
                belief_map, feet_xyz[foot, :2], support_est_n[foot],
                jnp.asarray(MAP_OBS_STD_N), implied[foot], jnp.asarray(MAP_OBS_STD_M),
            )
        state.info["snow_belief"] = belief_map
        state.info["snow_map_readout"] = belief.readout(belief_map, feet_xyz[:, :2])

        sup = support.evaluate(support_est_n, jnp.full((2,), MAP_OBS_STD_N), force[:, 2])
        state.info["snow_support_obs"] = support.observation(sup)
        state.info["snow_overload"] = support.overload_cost(
            sup, self._bearing_load(force)
        )

        state.info["snow_obs"] = obs8
        state.info["snow_history"] = jnp.concatenate(
            [obs8, state.info["snow_history"][: estimator.OBS_SIZE * (HISTORY_STEPS - 1)]]
        )
        state.info["snow_truth"] = jnp.concatenate([depth, bearing])
        state.info["snow_force"] = force
        state.info["snow_last_feet_pos"] = feet_xyz
        state.info["snow_est_last_obs"] = est_state.last_obs
        state.info["snow_est_false_safe"] = est_state.false_safe
        state.info["snow_est_noise_scale"] = est_state.noise_scale

        state = self._stepped_with_snow_contact(state, action, field)
        state.metrics["snow_depth_under_feet"] = jnp.mean(depth)
        return state

    def _snow_force_at(self, data, prev_feet_xyz, field, dt: float):
        """Snow force on both feet given the current physics state."""
        feet_xyz = data.site_xpos[self._feet_site_id]
        feet_vel = (feet_xyz - prev_feet_xyz) / dt
        depth, _, stiffness, damping, void = self._sample_feet(field, feet_xyz)
        force = sinkage.foot_force_world_n(
            feet_xyz[:, 2], feet_vel, jnp.zeros((2,)), depth, stiffness, damping, void > 0.5
        )
        return force * self._config.snow.enable_forces, feet_xyz

    def _physics_with_snow(self, data, motor_targets: jax.Array, field):
        """Substep the physics, recomputing the snow force at EVERY substep.

        Playground's `mjx_env.step` holds `xfrc_applied` constant across all 10 substeps.
        For a stiff spring applied as an external force that zero-order hold is an explicit
        integrator, and it is unstable above k ~ m/ctrl_dt^2 ~ 5000 N/m with a ~2 kg effective
        foot mass. The terrain generates stiffness up to 45000 N/m, so most of the range
        diverged: measured growth 4.0x at k=12000 and 9.1x at k=45000, which showed up in
        rollouts as the force pinned at its 2000 N clamp and swing feet launched to 0.67 m.

        Recomputing inside the substep loop moves the limit to m/sim_dt^2 ~ 500000 N/m,
        comfortably above the terrain's maximum. Measured growth is 1.0x at every stiffness.

        The estimator, belief map and reward stay at the 50 Hz control rate -- those are
        sensor-rate quantities and nothing about them is unstable.
        """
        def substep(carry, _):
            d, prev = carry
            force, feet_xyz = self._snow_force_at(d, prev, field, self.sim_dt)
            xfrc = jnp.zeros_like(d.xfrc_applied).at[self._feet_body_id, :3].set(force)
            d = d.replace(ctrl=motor_targets, xfrc_applied=xfrc)
            d = mjx.step(self.mjx_model, d)
            return (d, feet_xyz), force

        feet0 = data.site_xpos[self._feet_site_id]
        (data, _), forces = jax.lax.scan(
            substep, (data, feet0), None, length=self.n_substeps
        )
        # Report the last substep's force: what the feet are feeling as control resumes.
        return data, forces[-1]

    def _bearing_load(self, snow_force: jax.Array) -> jax.Array:
        """Feet whose vertical snow force is carrying meaningful load."""
        return snow_force[:, 2] > SNOW_CONTACT_FORCE_N

    def _stepped_with_snow_contact(self, state: mjx_env.State, action: jax.Array,
                                   field) -> mjx_env.State:
        """Playground's Joystick.step with ONE change: what counts as foot contact.

        Playground reads a floor-contact sensor at z=0. Once snow is a force rather than
        geometry that sensor stops meaning "the foot is standing on something": measured on
        deep terrain it fired on 2 of 35 steps, because the snow held the foot at z~0.156 m
        and it never reached the plane. Everything keyed to it broke -- `first_contact` never
        fired so the clearance reward was dead exactly where the signal lives, `swing_peak`
        never reset, and `feet_air_time` grew without bound.

        Contact here means BEARING LOAD: the floor sensor OR snow carrying real weight.

        This mirrors the upstream step body; re-check it when bumping mujoco_playground.
        """
        state.info["rng"], push1_rng, push2_rng = jax.random.split(state.info["rng"], 3)
        push_theta = jax.random.uniform(push1_rng, maxval=2 * jnp.pi)
        push_magnitude = jax.random.uniform(
            push2_rng,
            minval=self._config.push_config.magnitude_range[0],
            maxval=self._config.push_config.magnitude_range[1],
        )
        push = jnp.array([jnp.cos(push_theta), jnp.sin(push_theta)])
        push *= (
            jnp.mod(state.info["push_step"] + 1, state.info["push_interval_steps"]) == 0
        )
        push *= self._config.push_config.enable
        qvel = state.data.qvel
        qvel = qvel.at[:2].set(push * push_magnitude + qvel[:2])
        state = state.replace(data=state.data.replace(qvel=qvel))

        motor_targets = self._default_pose + action * self._config.action_scale
        data, applied_force = self._physics_with_snow(state.data, motor_targets, field)
        state.info["motor_targets"] = motor_targets
        state.info["snow_force"] = applied_force

        floor_contact = jnp.array([
            data.sensordata[self._mj_model.sensor_adr[sensorid]] > 0
            for sensorid in self._feet_floor_found_sensor
        ])
        # >>> the one deviation from upstream <<<
        contact = floor_contact | self._bearing_load(state.info["snow_force"])

        contact_filt = contact | state.info["last_contact"]
        first_contact = (state.info["feet_air_time"] > 0.0) * contact_filt
        state.info["feet_air_time"] += self.dt
        p_fz = data.site_xpos[self._feet_site_id][..., -1]
        state.info["swing_peak"] = jnp.maximum(state.info["swing_peak"], p_fz)

        obs = self._get_obs(data, state.info, contact)
        done = self._get_termination(data)

        rewards = self._get_reward(
            data, action, state.info, state.metrics, done, first_contact, contact
        )
        rewards = {k: v * self._config.reward_config.scales[k] for k, v in rewards.items()}
        reward = sum(rewards.values()) * self.dt

        state.info["push"] = push
        state.info["step"] += 1
        state.info["push_step"] += 1
        phase_tp1 = state.info["phase"] + state.info["phase_dt"]
        state.info["phase"] = jnp.fmod(phase_tp1 + jnp.pi, 2 * jnp.pi) - jnp.pi
        state.info["last_last_act"] = state.info["last_act"]
        state.info["last_act"] = action
        state.info["rng"], cmd_rng = jax.random.split(state.info["rng"])
        state.info["command"] = jnp.where(
            state.info["step"] > 500, self.sample_command(cmd_rng), state.info["command"]
        )
        state.info["step"] = jnp.where(done | (state.info["step"] > 500), 0,
                                       state.info["step"])
        state.info["feet_air_time"] *= ~contact
        state.info["last_contact"] = contact
        state.info["swing_peak"] *= ~contact
        for k, v in rewards.items():
            state.metrics[f"reward/{k}"] = v
        state.metrics["swing_peak"] = jnp.mean(state.info["swing_peak"])
        state.metrics["snow_contact_fraction"] = jnp.mean(contact.astype(float))

        done = done.astype(reward.dtype)
        return state.replace(data=data, obs=obs, reward=reward, done=done)

    # ---------------- observations ----------------

    def _snow_block(self, info: dict[str, Any]) -> jax.Array:
        # The parent's reset() calls _get_obs before this env has populated its info keys.
        # reset() recomputes the observation once they exist, so zeros here are transient.
        history = info.get("snow_history")
        if history is None:
            return jnp.zeros(estimator.OBS_SIZE * HISTORY_STEPS)
        return history * self._config.snow.estimator_channel

    def _assurance_block(self, info: dict[str, Any]) -> jax.Array:
        """everest's half: spatial memory plus bilateral support reserve."""
        readout = info.get("snow_map_readout")
        if readout is None:
            readout = jnp.zeros(belief.READOUT_SIZE)
        sup = info.get("snow_support_obs")
        if sup is None:
            sup = jnp.zeros(5)
        return jnp.hstack([
            readout * self._config.snow.map_channel,
            sup * self._config.snow.support_channel,
        ])

    def _get_obs(self, data, info: dict[str, Any], contact: jax.Array):
        obs = super()._get_obs(data, info, contact)
        block = jnp.hstack([self._snow_block(info), self._assurance_block(info)])
        return {
            "state": jnp.hstack([obs["state"], block]),
            # Privileged: the critic sees ground truth, the actor never does.
            "privileged_state": jnp.hstack([
                obs["privileged_state"], block,
                info.get("snow_truth", jnp.zeros(2 * estimator.N_FEET)),
            ]),
        }

    # ---------------- rewards ----------------

    def _get_reward(self, data, action, info, metrics, done, first_contact, contact):
        rewards = super()._get_reward(data, action, info, metrics, done,
                                      first_contact, contact)
        obs8 = info["snow_obs"].reshape(estimator.N_FEET, estimator.N_CHANNELS)
        implied_depth = estimator.implied_depth_m(obs8, self._est_params)
        confidence = obs8[:, 2]

        rewards["snow_clearance"] = self._reward_snow_clearance(
            info["swing_peak"], implied_depth, confidence, first_contact
        )
        rewards["snow_stumble"] = self._cost_snow_stumble(info["snow_force"], contact)
        rewards["snow_flat_efficiency"] = self._cost_snow_flat_efficiency(
            info["swing_peak"], implied_depth, confidence, first_contact
        )
        # everest's bilateral gate. It cannot halt a locomotion policy, so it becomes a cost
        # for standing on a foothold the estimate says will not hold.
        rewards["snow_overload"] = info["snow_overload"]
        return rewards

    def _reward_snow_clearance(self, swing_peak, implied_depth, confidence, first_contact):
        """Score the swing foot's peak height against the depth the ESTIMATE implies.

        Targets a value rather than rewarding 'higher is better' -- unbounded lifting is its
        own failure mode. Scored once per step, when the swing ends.
        """
        target = implied_depth + CLEARANCE_MARGIN_M
        score = jnp.exp(-((swing_peak - target) ** 2) / CLEARANCE_TOLERANCE)
        return jnp.sum(score * confidence * first_contact)

    def _cost_snow_stumble(self, snow_force, contact):
        """Ploughing drag on a foot that should be swinging freely.

        A physics consequence, so it may read the real force rather than the estimate.
        Graded rather than binary: a step function gives the optimiser no gradient about
        *how much* it is ploughing, only whether it crossed a line.
        """
        horizontal = jnp.linalg.norm(snow_force[:, :2], axis=-1)
        excess = jnp.maximum(horizontal - STUMBLE_PLOUGH_N, 0.0) / STUMBLE_PLOUGH_N
        return jnp.sum(excess * ~contact.astype(bool))

    def _cost_snow_flat_efficiency(self, swing_peak, implied_depth, confidence, first_contact):
        """Penalise lifting high when the estimate says the ground is firm.

        Without this the policy prances everywhere, which technically satisfies the clearance
        term and is obviously wrong.
        """
        on_flat = implied_depth < FLAT_DEPTH_THRESHOLD_M
        excess = jnp.maximum(swing_peak - FLAT_NOMINAL_CLEARANCE_M, 0.0)
        return jnp.sum(excess * on_flat * confidence * first_contact)


def make_snow_env(config_overrides: dict[str, Any] | None = None) -> SnowJoystick:
    return SnowJoystick(config_overrides=config_overrides)


def make_no_map_env(config_overrides: dict[str, Any] | None = None) -> SnowJoystick:
    """Spatial memory blanked, everything else intact.

    This is the run everest could never do. Its own benchmark conflated the map with the
    bilateral gate, and the decomposition showed the gate did all the work -- so the map's
    contribution has never actually been measured. Against the full env, this isolates it.
    """
    overrides = dict(config_overrides or {})
    overrides["snow.map_channel"] = 0.0
    return SnowJoystick(config_overrides=overrides)


def make_no_support_env(config_overrides: dict[str, Any] | None = None) -> SnowJoystick:
    """Bilateral reserve blanked, everything else intact."""
    overrides = dict(config_overrides or {})
    overrides["snow.support_channel"] = 0.0
    return SnowJoystick(config_overrides=overrides)


def make_ablation_env(config_overrides: dict[str, Any] | None = None) -> SnowJoystick:
    """Identical environment with the estimator block zeroed.

    This is the comparison that isolates the sensor channel. Baseline-vs-trained differs in
    three ways -- channel, rewards, AND snow training experience -- so only this run can
    attribute a gap to the sensor.
    """
    overrides = dict(config_overrides or {})
    overrides["snow.estimator_channel"] = 0.0
    overrides["snow.map_channel"] = 0.0
    overrides["snow.support_channel"] = 0.0
    return SnowJoystick(config_overrides=overrides)