Initial commit

s10_supermimic/core/terrain_posture.py

@@ -0,0 +1,33 @@
+"""Shared real/imagined posture objective for the opt-in terrain v2 task.
+
+S10 joint order is hipx (abduction), hipy (flexion), knee. Hipx keeps a 2x
+weight while hipy/knee relax on rough terrain. v7 sets
+``relief_floor=0.02`` so stairs leave knee/hipy deviation almost unpenalised
+and the policy can flex instead of holding a stiff leg.
+"""
+import torch
+
+LEG_SCALES = (0.25, 0.6, 0.8) * 4
+
+
+def posture_cost(q, scan, relief_floor=0.125):
+  # Only the supporting/approaching ground near the body should relax posture.
+  # scan is body height minus ground; its range is independent of body height.
+  local = scan.reshape(-1, 21, 13)[:, 5:16, 3:10].flatten(1)
+  relief = ((local.amax(-1) - local.amin(-1)) / 0.12).clamp(0., 1.)
+  multiplier = (1. - relief).clamp_min(relief_floor)
+  legs = q.reshape(-1, 4, 3).square()
+  # Hip abduction stays constrained even on steps; flexion is freer there.
+  return 2. * legs[:, :, 0].sum(-1) + multiplier * legs[:, :, 1:].sum((1, 2))
+
+
+def env_posture_cost(env, asset_cfg, relief_floor=0.125):
+  from .terrain import TerrainField, TerrainBridge
+  bridge = getattr(env, "_s10_posture_bridge", None)
+  if bridge is None:
+    bridge = TerrainBridge(env, TerrainField.from_env(env))
+    env._s10_posture_bridge = bridge
+  d = env.scene[asset_cfg.name].data
+  q = d.joint_pos[:, asset_cfg.joint_ids] - d.default_joint_pos[:, asset_cfg.joint_ids]
+  return posture_cost(q, bridge.field.scan(bridge.state(), bridge.pose()),
+                      relief_floor)