78 lines
3.4 KiB
Python
78 lines
3.4 KiB
Python
"""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]
|