Upload folder using huggingface_hub

code/snow_rl/snow_env.py

@@ -26,6 +26,7 @@ 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
@@ -60,6 +61,16 @@ SNOW_CONTACT_FORCE_N = 20.0
 # 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.
@@ -85,6 +96,11 @@ def default_config() -> config_dict.ConfigDict:
     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
@@ -98,6 +114,10 @@ def default_config() -> config_dict.ConfigDict:
         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,
@@ -124,6 +144,14 @@ class SnowJoystick(joystick.Joystick):
 
     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))
+            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(

code/snow_rl/terrain.py

@@ -156,7 +156,8 @@ def generate(key: jax.Array, cfg: TerrainConfig = DEFAULT_CONFIG,
     return SnowField(depth, bearing, stiffness, damping, void, regime)
 
 
-def crossing_strip(cfg: TerrainConfig = DEFAULT_CONFIG) -> SnowField:
+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.
@@ -180,9 +181,14 @@ def crossing_strip(cfg: TerrainConfig = DEFAULT_CONFIG) -> SnowField:
     # the corridor and a discontinuous jump in the sinkage force.
     shallow = smoothstep(0.28, 0.38, x)
     deep = smoothstep(0.61, 0.71, x)
-    depth = 0.005 + (0.070 - 0.005) * shallow + (0.190 - 0.070) * deep
+    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 + (12_000.0 - 28_000.0) * shallow + (3_000.0 - 12_000.0) * deep,
+        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))