Initial commit

s10_supermimic/core/wheel_intrusion.py

@@ -0,0 +1,101 @@
+"""Always-on wheel/body projection constraints for the safe v7 stance.
+
+The body is 0.60 m long in ``base_link`` (x in [-0.30, 0.30]).
+
+``wheel_body_intrusion`` keeps every individual wheel out of the opposite
+three quarters: a front wheel inside the rear three quarters (x < +0.15) or a
+rear wheel inside the front three quarters (x > -0.15) shortens the support
+polygon and makes the robot tip-prone. Its hinge is exactly zero outside the
+outer quarters and grows cubically inside, so a normal extended stance is
+untouched while a wheel tucked under the body is penalised quickly.
+
+``pair_body_intrusion`` adds the wider five-sixths restriction requested for
+v7: both front wheels may not sit inside the rear five sixths (x < +0.20) at
+the same time, and both rear wheels may not sit inside the front five sixths
+(x > -0.20). A single leg crossing the line is allowed; a leg only counts as
+inside after it has stayed inside continuously for ``PAIR_GRACE_S`` (0.5 s).
+The pair penalty is therefore zero during a brief swing-through and applies
+immediately once both same-side legs are armed.
+
+Both terms are kinematic and need no terrain scan. Wheels are ordered
+fl, fr, hl, hr.
+"""
+import torch
+
+from .rolling_rewards import wheel_geometry
+
+BODY_HALF_LENGTH = 0.30
+# Each wheel must stay in the outer quarter of the body length: x > +0.15 for
+# front wheels and x < -0.15 for rear wheels.
+WHEEL_STRIP_BOUNDARY = BODY_HALF_LENGTH / 2.0
+# Rear/front 5/6 of a 0.60 m body => the boundary sits at x = +-0.20.
+PAIR_BOUNDARY = 2.0 * BODY_HALF_LENGTH / 3.0
+PAIR_GRACE_S = 0.5
+
+
+def wheel_body_intrusion(state, default_q, boundary=WHEEL_STRIP_BOUNDARY):
+  point, _ = wheel_geometry(state, default_q)
+  x = point[..., 0]
+  # Front wheels must stay ahead of +L/4, rear wheels behind -L/4.
+  violation = torch.cat(((boundary - x[:, :2]).clamp_min(0.),
+                         (x[:, 2:] + boundary).clamp_min(0.)), -1)
+  normalized = violation / boundary
+  return normalized.pow(3).mean(-1)
+
+
+def pair_region(state, default_q, boundary=PAIR_BOUNDARY):
+  """Per-wheel inside mask and normalised depth for the 5/6 restriction."""
+  point, _ = wheel_geometry(state, default_q)
+  x = point[..., 0]
+  inside = torch.cat((x[:, :2] < boundary, x[:, 2:] > -boundary), -1)
+  depth = torch.cat(((boundary - x[:, :2]).clamp_min(0.),
+                     (x[:, 2:] + boundary).clamp_min(0.)), -1) / boundary
+  return inside, depth
+
+
+def pair_body_intrusion(state, default_q, timers, dt, grace=PAIR_GRACE_S):
+  """Penalise both front or both rear wheels staying inside the 5/6 region.
+
+  ``timers`` holds the continuous inside time of each wheel. Returning it lets
+  the caller keep the 0.5 s grace across environment steps. A wheel outside
+  the boundary resets its own timer, so a single-leg pass-through never
+  contributes.
+  """
+  inside, depth = pair_region(state, default_q)
+  timers = torch.where(inside, timers + dt, torch.zeros_like(timers))
+  armed = timers >= grace - 1e-6
+  front = armed[:, 0] & armed[:, 1]
+  rear = armed[:, 2] & armed[:, 3]
+  cost = (front.float() * depth[:, :2].mean(-1)
+          + rear.float() * depth[:, 2:].mean(-1))
+  return cost, timers
+
+
+def _bridge(env):
+  from .terrain import TerrainBridge, TerrainField
+  bridge = getattr(env, '_s10_posture_bridge', None)
+  if bridge is None:
+    bridge = TerrainBridge(env, TerrainField.from_env(env))
+    env._s10_posture_bridge = bridge
+  return bridge
+
+
+def env_wheel_body_intrusion(env):
+  bridge = _bridge(env)
+  return wheel_body_intrusion(
+    bridge.state(), bridge.robot.data.default_joint_pos[0, bridge.legs])
+
+
+def env_pair_body_intrusion(env):
+  bridge = _bridge(env)
+  timers = getattr(env, '_s10_pair_timer', None)
+  if timers is None or timers.shape[0] != env.num_envs:
+    timers = torch.zeros(env.num_envs, 4, device=env.device)
+  # episode_length_buf is 1 on the first reward step after an auto-reset.
+  timers = timers.clone()
+  timers[env.episode_length_buf <= 1] = 0.
+  cost, timers = pair_body_intrusion(
+    bridge.state(), bridge.robot.data.default_joint_pos[0, bridge.legs],
+    timers, env.step_dt)
+  env._s10_pair_timer = timers
+  return cost