Initial commit

s10_supermimic/core/wheel_terrain_rewards.py

@@ -0,0 +1,81 @@
+"""V6 wheel-local terrain constraints; privileged scans never enter the actor.
+
+Clearance is a geometric surrogate, not a measured persistent-contact detector.
+S10 joint geometry comes from ``rolling_rewards.wheel_geometry``.
+"""
+import torch
+from .world_model import GRAVITY, LEG_POS, LEG_VEL
+from .rolling_rewards import wheel_geometry
+
+
+def yaw_frame_points(state, points):
+  # Remove roll/pitch, retaining the heading frame used by TerrainField.scan.
+  up = -torch.nn.functional.normalize(state[:, GRAVITY], dim=-1, eps=1e-6)
+  cp = (up[:, 1:].square().sum(-1)+1e-12).sqrt()
+  sr, cr, sp = up[:, 1] / cp, up[:, 2] / cp, -up[:, 0]
+  x, y, z = points.unbind(-1)
+  yr = cr[:, None]*y - sr[:, None]*z
+  zr = sr[:, None]*y + cr[:, None]*z
+  return torch.stack((cp[:, None]*x + sp[:, None]*zr, yr,
+                      -sp[:, None]*x + cp[:, None]*zr), -1)
+
+
+def sample_scan(scan, xy):
+  """Bilinear samples on the existing 21x13, 10cm heading-frame grid."""
+  shape = xy.shape[:-1]
+  x = ((xy[..., 0]+1.)*10).clamp(0,20).reshape(len(scan),-1)
+  y = ((xy[..., 1]+.6)*10).clamp(0,12).reshape(len(scan),-1)
+  ix, iy = x.floor().long().clamp(0,19), y.floor().long().clamp(0,11)
+  ax, ay = x-ix, y-iy
+  def at(dx,dy): return scan.gather(1,(ix+dx)*13+iy+dy)
+  return (((1-ax)*at(0,0)+ax*at(1,0))*(1-ay)
+          + ((1-ax)*at(0,1)+ax*at(1,1))*ay).reshape(shape)
+
+
+def wheel_terms(state, command, scan, default_q):
+  point, radius = wheel_geometry(state, default_q)
+  point = yaw_frame_points(state, point)
+  xy = point[..., :2]
+  offsets = state.new_tensor([[x,y] for x in (-.15,0.,.15) for y in (-.10,0.,.10)])
+  patch_xy = xy[:,:,None,:]+offsets
+  samples = sample_scan(scan, patch_xy)
+  # Index-based extrema avoid division by the number of equal extrema in
+  # compiled backward through interpolated, nearly-flat patches.
+  relief = samples.max(dim=-1).values-samples.min(dim=-1).values
+  # Outside the scan, never assume flatness and penalize a recovering leg.
+  inside = ((patch_xy[...,0].abs()<=1.) & (patch_xy[...,1].abs()<=.6)).all(-1)
+  flat = (1.-(relief-.02)/.04).clamp(0,1)*inside
+  straight = torch.exp(-(command[:,1]/.25).square()-(command[:,2]).square())
+  gate = flat*straight[:,None]
+  below = sample_scan(scan,xy)
+  bottom = point[...,2]+below-radius
+  lift = ((bottom-.005).clamp_min(0)/.05).square()
+  tilt = ((state[:,GRAVITY][:,:2].square().sum(-1)+1e-12).sqrt()-.08).div(.22).clamp(0,1)
+  motion = state[:,LEG_VEL].reshape(-1,4,3).square().mean(-1)*(1.-.8*tilt[:,None])
+  q = state[:,LEG_POS].reshape(-1,4,3).square()
+  posture_multiplier = torch.where(inside, 1.-.875*(relief/.12).clamp(0,1), .125)
+  posture = 2.*q[:,:,0].sum(-1)+(posture_multiplier*q[:,:,1:].sum(-1)).sum(-1)
+  # Only approaching a higher surface: no reward for lifting on flat ground.
+  speed = command[:,:2].norm(dim=-1)
+  heading_velocity = yaw_frame_points(state, torch.cat((command[:,:2], torch.zeros_like(command[:,:1])), -1)[:,None,:])[:,0,:2]
+  direction = heading_velocity/(heading_velocity.square().sum(-1,keepdim=True)+1e-12).sqrt().clamp_min(.1)
+  ahead_xy = xy+.16*direction[:,None,:]
+  ahead = sample_scan(scan,ahead_xy)
+  rise = below-ahead
+  approaching = ((rise-.04)/.08).clamp(0,1)*(speed>.1)[:,None]*inside
+  deficit = (radius+.015-point[...,2]-ahead).clamp_min(0)/.20
+  clearance = (approaching*deficit/(1.+deficit)).mean(-1)
+  return dict(unnecessary_wheel_lift=(gate*lift).sum(-1),
+              rolling_leg_motion=(gate*motion).mean(-1), leg_posture=posture,
+              wheel_obstacle_clearance=clearance)
+
+
+def env_wheel_term(env, term):
+  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
+  state=bridge.state()
+  return wheel_terms(state,env.command_manager.get_command('twist'),
+                     bridge.field.scan(state,bridge.pose()),
+                     bridge.robot.data.default_joint_pos[0,bridge.legs])[term]