Upload folder using huggingface_hub

2eeee5a verified 2 months ago

6.27 kB

"""Snow sinkage force (Step 2, part 2).
Follows the resistance law in hackathon-everest's `physics.py::_resistance_n`:
    F = k_eff * d^1.12  +  c * max(v_down, 0)  +  3 * k_eff * max(d - d_support, 0)
with a void band that collapses resistance to 8%. Term by term:
* `k_eff * d^1.12`   progressive compaction -- snow stiffens as it packs.
* `c * v_down`       rate dependence -- snow resists a fast foot harder than a slow one.
                     Only the downward half; snow does not pull a lifting foot back.
* `3 * k_eff * (...)`  bottoming out on the hard substrate once the foot passes through
                     the snow layer. The factor 3 makes that transition sharp.
* void               a hollow under a crust: the foot meets almost nothing, then the
                     substrate term catches it. This is the breakthrough the policy must survive.
Applied to the feet as `xfrc_applied`, not as contact geometry -- see terrain.py for why.
The brief calls for depth-and-velocity dependence and explicitly not granular physics.
"""
from __future__ import annotations

import jax
import jax.numpy as jnp

COMPACTION_EXPONENT = 1.12
SUBSTRATE_STIFFNESS_GAIN = 3.0
VOID_RESISTANCE_FACTOR = 0.08
MAX_FORCE_N = 2_000.0  # numerical guard; ~6x G1 body weight

# Ploughing. Vertical support alone is NOT enough to make deep snow hard: under uniform
# stiffness a foot reaches the same equilibrium sinkage in 0.07 m and 0.19 m of snow, so
# the two regimes would feel identical and the policy would have nothing to learn.
# What actually punishes deep snow is dragging the swing foot through it. Resistance scales
# with submerged frontal area (~ depth x foot width) and snow strength, saturating in speed.
# Sized so a foot swung through deep snow at walking speed meets ~30 N -- roughly 9% of body
# weight, enough to alter gait without immobilising the robot. The learning signal comes from
# clearance: a foot lifted clear of the surface ploughs nothing, so deeper snow demands a
# higher swing to escape the same drag.
PLOUGH_COEFF = 0.25
FOOT_WIDTH_M = 0.09           # G1 foot, approximate
PLOUGH_REFERENCE_SPEED_MPS = 0.4


def resistance_n(
    penetration_m: jax.Array,
    vertical_velocity_mps: jax.Array,
    snow_depth_m: jax.Array,
    stiffness_n_per_m: jax.Array,
    damping_ns_per_m: jax.Array,
    void: jax.Array,
    include_substrate: bool = True,
) -> jax.Array:
    """Upward resisting force on a foot. Zero when the foot is above the snow surface.
    penetration_m: depth below the snow surface, negative when clear of it.
    vertical_velocity_mps: foot velocity, negative downward.
    include_substrate: whether to add the bottoming-out term. TRUE when deriving bearing
    capacity as a label (everest has no geometry, so the term is the only hard ground there).
    FALSE when computing a force to APPLY in MuJoCo, where the substrate is a real plane with
    a real contact solver -- adding it there double-counts the floor and fights the solver,
    which produced spikes to 1530 N (~4.5x body weight) against the 2000 N clamp.
    """
    d = jnp.maximum(penetration_m, 0.0)
    v_down = jnp.maximum(-vertical_velocity_mps, 0.0)

    force = stiffness_n_per_m * d**COMPACTION_EXPONENT + damping_ns_per_m * v_down
    # Void collapses the snow-column term only; the substrate below still stops the foot.
    force = jnp.where(void, force * VOID_RESISTANCE_FACTOR, force)
    if include_substrate:
        force = force + SUBSTRATE_STIFFNESS_GAIN * stiffness_n_per_m * jnp.maximum(
            d - snow_depth_m, 0.0
        )
    return jnp.clip(jnp.where(d > 0.0, force, 0.0), 0.0, MAX_FORCE_N)


def plough_force_n(
    penetration_m: jax.Array,
    horizontal_speed_mps: jax.Array,
    stiffness_n_per_m: jax.Array,
    void: jax.Array,
) -> jax.Array:
    """Magnitude of the horizontal drag on a foot moving through snow.
    Opposes horizontal motion. Zero above the surface; a void offers almost nothing to plough.
    """
    submerged = jnp.maximum(penetration_m, 0.0)
    strength = jnp.where(void, stiffness_n_per_m * VOID_RESISTANCE_FACTOR, stiffness_n_per_m)
    force = (
        PLOUGH_COEFF * strength * submerged * FOOT_WIDTH_M
        * jnp.tanh(horizontal_speed_mps / PLOUGH_REFERENCE_SPEED_MPS)
    )
    return jnp.clip(force, 0.0, MAX_FORCE_N)


def foot_force_world_n(
    foot_pos_z_m: jax.Array,
    foot_vel_xyz_mps: jax.Array,
    ground_z_m: jax.Array,
    snow_depth_m: jax.Array,
    stiffness_n_per_m: jax.Array,
    damping_ns_per_m: jax.Array,
    void: jax.Array,
) -> jax.Array:
    """World-frame (fx, fy, fz) per foot, ready for `data.xfrc_applied[body, :3]`.
    The snow surface sits at `ground_z_m + snow_depth_m`; the hard substrate at `ground_z_m`.
    Vertical term supports the foot; horizontal term ploughs against its motion.
    """
    penetration = ground_z_m + snow_depth_m - foot_pos_z_m
    # include_substrate=False: MuJoCo's floor plane already stops the foot.
    fz = resistance_n(penetration, foot_vel_xyz_mps[..., 2], snow_depth_m,
                      stiffness_n_per_m, damping_ns_per_m, void, include_substrate=False)
    v_h = foot_vel_xyz_mps[..., :2]
    speed = jnp.linalg.norm(v_h, axis=-1)
    drag = plough_force_n(penetration, speed, stiffness_n_per_m, void)
    direction = v_h / jnp.maximum(speed, 1e-6)[..., None]
    return jnp.concatenate([-drag[..., None] * direction, fz[..., None]], axis=-1)


def equilibrium_sinkage_m(
    load_n: jax.Array,
    snow_depth_m: jax.Array,
    stiffness_n_per_m: jax.Array,
    void: jax.Array,
    *,
    iterations: int = 40,
) -> jax.Array:
    """Static sinkage where resistance balances `load_n`, by bisection.
    Diagnostic and test helper -- the training loop integrates the force instead.
    """
    def body(bounds, _):
        lo, hi = bounds
        mid = 0.5 * (lo + hi)
        f = resistance_n(mid, jnp.zeros_like(mid), snow_depth_m, stiffness_n_per_m,
                         jnp.zeros_like(mid), void)
        too_soft = f < load_n
        return (jnp.where(too_soft, mid, lo), jnp.where(too_soft, hi, mid)), None

    lo = jnp.zeros_like(snow_depth_m)
    hi = jnp.full_like(snow_depth_m, 1.0)
    (lo, hi), _ = jax.lax.scan(body, (lo, hi), None, length=iterations)
    return 0.5 * (lo + hi)

	"""Snow sinkage force (Step 2, part 2).

	Follows the resistance law in hackathon-everest's `physics.py::_resistance_n`:

	F = k_eff * d^1.12 + c * max(v_down, 0) + 3 * k_eff * max(d - d_support, 0)

	with a void band that collapses resistance to 8%. Term by term:

	* `k_eff * d^1.12` progressive compaction -- snow stiffens as it packs.
	* `c * v_down` rate dependence -- snow resists a fast foot harder than a slow one.
	Only the downward half; snow does not pull a lifting foot back.
	* `3 * k_eff * (...)` bottoming out on the hard substrate once the foot passes through
	the snow layer. The factor 3 makes that transition sharp.
	* void a hollow under a crust: the foot meets almost nothing, then the
	substrate term catches it. This is the breakthrough the policy must survive.

	Applied to the feet as `xfrc_applied`, not as contact geometry -- see terrain.py for why.
	The brief calls for depth-and-velocity dependence and explicitly not granular physics.
	"""
	from __future__ import annotations

	import jax
	import jax.numpy as jnp

	COMPACTION_EXPONENT = 1.12
	SUBSTRATE_STIFFNESS_GAIN = 3.0
	VOID_RESISTANCE_FACTOR = 0.08
	MAX_FORCE_N = 2_000.0 # numerical guard; ~6x G1 body weight

	# Ploughing. Vertical support alone is NOT enough to make deep snow hard: under uniform
	# stiffness a foot reaches the same equilibrium sinkage in 0.07 m and 0.19 m of snow, so
	# the two regimes would feel identical and the policy would have nothing to learn.
	# What actually punishes deep snow is dragging the swing foot through it. Resistance scales
	# with submerged frontal area (~ depth x foot width) and snow strength, saturating in speed.
	# Sized so a foot swung through deep snow at walking speed meets ~30 N -- roughly 9% of body
	# weight, enough to alter gait without immobilising the robot. The learning signal comes from
	# clearance: a foot lifted clear of the surface ploughs nothing, so deeper snow demands a
	# higher swing to escape the same drag.
	PLOUGH_COEFF = 0.25
	FOOT_WIDTH_M = 0.09 # G1 foot, approximate
	PLOUGH_REFERENCE_SPEED_MPS = 0.4


	def resistance_n(
	penetration_m: jax.Array,
	vertical_velocity_mps: jax.Array,
	snow_depth_m: jax.Array,
	stiffness_n_per_m: jax.Array,
	damping_ns_per_m: jax.Array,
	void: jax.Array,
	include_substrate: bool = True,
	) -> jax.Array:
	"""Upward resisting force on a foot. Zero when the foot is above the snow surface.

	penetration_m: depth below the snow surface, negative when clear of it.
	vertical_velocity_mps: foot velocity, negative downward.

	include_substrate: whether to add the bottoming-out term. TRUE when deriving bearing
	capacity as a label (everest has no geometry, so the term is the only hard ground there).
	FALSE when computing a force to APPLY in MuJoCo, where the substrate is a real plane with
	a real contact solver -- adding it there double-counts the floor and fights the solver,
	which produced spikes to 1530 N (~4.5x body weight) against the 2000 N clamp.
	"""
	d = jnp.maximum(penetration_m, 0.0)
	v_down = jnp.maximum(-vertical_velocity_mps, 0.0)

	force = stiffness_n_per_m * d*COMPACTION_EXPONENT + damping_ns_per_m v_down
	# Void collapses the snow-column term only; the substrate below still stops the foot.
	force = jnp.where(void, force * VOID_RESISTANCE_FACTOR, force)
	if include_substrate:
	force = force + SUBSTRATE_STIFFNESS_GAIN * stiffness_n_per_m * jnp.maximum(
	d - snow_depth_m, 0.0
	)
	return jnp.clip(jnp.where(d > 0.0, force, 0.0), 0.0, MAX_FORCE_N)


	def plough_force_n(
	penetration_m: jax.Array,
	horizontal_speed_mps: jax.Array,
	stiffness_n_per_m: jax.Array,
	void: jax.Array,
	) -> jax.Array:
	"""Magnitude of the horizontal drag on a foot moving through snow.

	Opposes horizontal motion. Zero above the surface; a void offers almost nothing to plough.
	"""
	submerged = jnp.maximum(penetration_m, 0.0)
	strength = jnp.where(void, stiffness_n_per_m * VOID_RESISTANCE_FACTOR, stiffness_n_per_m)
	force = (
	PLOUGH_COEFF * strength * submerged * FOOT_WIDTH_M
	* jnp.tanh(horizontal_speed_mps / PLOUGH_REFERENCE_SPEED_MPS)
	)
	return jnp.clip(force, 0.0, MAX_FORCE_N)


	def foot_force_world_n(
	foot_pos_z_m: jax.Array,
	foot_vel_xyz_mps: jax.Array,
	ground_z_m: jax.Array,
	snow_depth_m: jax.Array,
	stiffness_n_per_m: jax.Array,
	damping_ns_per_m: jax.Array,
	void: jax.Array,
	) -> jax.Array:
	"""World-frame (fx, fy, fz) per foot, ready for `data.xfrc_applied[body, :3]`.

	The snow surface sits at `ground_z_m + snow_depth_m`; the hard substrate at `ground_z_m`.
	Vertical term supports the foot; horizontal term ploughs against its motion.
	"""
	penetration = ground_z_m + snow_depth_m - foot_pos_z_m
	# include_substrate=False: MuJoCo's floor plane already stops the foot.
	fz = resistance_n(penetration, foot_vel_xyz_mps[..., 2], snow_depth_m,
	stiffness_n_per_m, damping_ns_per_m, void, include_substrate=False)
	v_h = foot_vel_xyz_mps[..., :2]
	speed = jnp.linalg.norm(v_h, axis=-1)
	drag = plough_force_n(penetration, speed, stiffness_n_per_m, void)
	direction = v_h / jnp.maximum(speed, 1e-6)[..., None]
	return jnp.concatenate([-drag[..., None] * direction, fz[..., None]], axis=-1)


	def equilibrium_sinkage_m(
	load_n: jax.Array,
	snow_depth_m: jax.Array,
	stiffness_n_per_m: jax.Array,
	void: jax.Array,
	*,
	iterations: int = 40,
	) -> jax.Array:
	"""Static sinkage where resistance balances `load_n`, by bisection.

	Diagnostic and test helper -- the training loop integrates the force instead.
	"""
	def body(bounds, _):
	lo, hi = bounds
	mid = 0.5 * (lo + hi)
	f = resistance_n(mid, jnp.zeros_like(mid), snow_depth_m, stiffness_n_per_m,
	jnp.zeros_like(mid), void)
	too_soft = f < load_n
	return (jnp.where(too_soft, mid, lo), jnp.where(too_soft, hi, mid)), None

	lo = jnp.zeros_like(snow_depth_m)
	hi = jnp.full_like(snow_depth_m, 1.0)
	(lo, hi), _ = jax.lax.scan(body, (lo, hi), None, length=iterations)
	return 0.5 * (lo + hi)