102 lines
4.0 KiB
Python
102 lines
4.0 KiB
Python
"""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
|