Initial commit

2026-11-02 17:10:39 +08:00
commit 102fdbaf47
79 changed files with 56509 additions and 0 deletions
@@ -0,0 +1,77 @@
+"""Terrain-gated rolling preference, shared by real and imagined rewards."""
+import torch
+from .world_model import GRAVITY, LEG_POS, LEG_VEL
+
+
+def rolling_gate(scan, command):
+  local = scan.reshape(-1, 21, 13)[:, 5:16, 3:10].flatten(1)
+  relief = local.amax(-1) - local.amin(-1)
+  # Flat (<2 cm): favor rolling; near obstacles (>=6 cm): allow lifting.
+  flat = (1. - (relief - .02) / .04).clamp(0., 1.)
+  # Non-steerable wheels may require stepping for lateral/yaw maneuvers.
+  straight = torch.exp(-(command[:,1]/.25).square() - (command[:,2]/1.).square())
+  return flat * straight
+
+
+def wheel_geometry(state, default_q):
+  """S10 wheel centres in body coordinates and vertical cylinder half-extents.
+
+  Joint order is hipx (abduction), hipy (flexion), knee. Geometry follows
+  ``S10_description/s10_mjcf/S10.xml``: hipx at x=±0.2277/y=±0.0574, hipy
+  offset y=±0.0779, knee at z=-0.18, wheel at y=±0.045891/z=-0.18,
+  wheel radius 0.081 m and half-width 0.0185 m. The analytic result was
+  checked against MuJoCo ``mj_forward`` to machine precision.
+
+  The cylinder is symmetric about its axle, so wheel phase is unnecessary.
+  """
+  q = (state[:,LEG_POS] + default_q).reshape(-1,4,3)
+  hipx, hipy, knee = q.unbind(-1)
+  side = state.new_tensor([1.,-1.,1.,-1.])
+  # Leg sagittal plane: knee link (0.18 m) + wheel link (0.18 m).
+  x = .18*knee.sin()
+  z = -.18 - .18*knee.cos()
+  x3 = hipy.cos()*x - hipy.sin()*z
+  z3 = hipy.sin()*x + hipy.cos()*z
+  y3 = side*(.0779+.045891)
+  y4 = hipx.cos()*y3 + hipx.sin()*z3
+  z4 = -hipx.sin()*y3 + hipx.cos()*z3
+  point = torch.stack((state.new_tensor([.2277,.2277,-.2277,-.2277]) + x3,
+                       state.new_tensor([.0574,-.0574,.0574,-.0574]) + y4,
+                       z4),-1)
+  up = -torch.nn.functional.normalize(state[:,GRAVITY],dim=-1,eps=1e-6)
+  # Wheel axle is local -y, rotated only by hipx; sign is irrelevant after abs.
+  axle_z = -up[:,1:2]*hipx.cos()+up[:,2:3]*hipx.sin()
+  radius_z = .081*(1.-axle_z.square()).clamp_min(1e-8).sqrt()+.0185*axle_z.abs()
+  return point, radius_z
+
+
+def wheel_bottom_height(state, default_q):
+  """Wheel bottoms relative to the ground below the base.
+
+  Only used on near-flat ground by rolling_gate, not as obstacle clearance.
+  """
+  point, radius_z = wheel_geometry(state, default_q)
+  up = -torch.nn.functional.normalize(state[:,GRAVITY],dim=-1,eps=1e-6)
+  return state[:,:1] + (point*up[:,None,:]).sum(-1) - radius_z
+
+
+def rolling_costs(state, command, scan, default_q):
+  gate = rolling_gate(scan, command)
+  # Tolerate 5 mm of clearance; no positive reward for lifting anywhere.
+  lift = ((wheel_bottom_height(state,default_q)-.005).clamp_min(0)/.05).square().sum(-1)
+  # Reference mjlab task uses mean squared speed and relaxes during recovery.
+  tilt = (state[:,GRAVITY][:,:2].norm(dim=-1)-.08).div(.22).clamp(0.,1.)
+  motion = state[:,LEG_VEL].square().mean(-1) * (1.-.8*tilt)
+  return {'unnecessary_wheel_lift':gate*lift, 'rolling_leg_motion':gate*motion}
+
+
+def env_rolling_cost(env, term):
+  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
+  state = bridge.state()
+  return rolling_costs(state, env.command_manager.get_command('twist'),
+                       bridge.field.scan(state,bridge.pose()),
+                       bridge.robot.data.default_joint_pos[0,bridge.legs])[term]