Hugging Face's logo

BAAI
/
PingPongRealWorldRL

Reinforcement Learning
ONNX
English
robotics
unitree-h2
table-tennis
Model card Files
xet
 Community
PingPongRealWorldRL / runtime /ball.py
KurisiTOMA's picture
KurisiTOMA
chore: publish H2 real-world RL system
d8290eb verified
Raw History Blame Contribute Delete
13.3 kB
"""Kinematic table-tennis ball, ported from ``KinematicPingPongCommand`` (training).
Same drag, fitted bounce regression, net/landing logic, incoming-velocity solver
and virtual-paddle contact model as the Isaac Lab environment, in numpy.
"""
from __future__ import annotations
import math
from dataclasses import dataclass, field
import numpy as np
from .mathutil import quat_apply, quat_from_euler_xyz, quat_mul, quat_to_matrix
# ---- constants (tracking_pingpong_online_closd_ball_env_cfg_FB + KinematicPingPongCommandCfg) ----
BALL_RADIUS = 0.02
GRAVITY = 9.81
DRAG_K = 0.09213392087790545 # 1/m, v <- v / (1 + k|v|dt)
BALL_MAX_SPEED = 12.0
TABLE_HEIGHT = 0.76
TABLE_X_RANGE = (0.50, 3.24)
TABLE_HALF_WIDTH = 0.76
NET_X = 1.87
NET_HEIGHT = 0.91
RETURN_X_RANGE = (1.87, 3.24)
KILL_Z = 0.15
KILL_X_RANGE = (-0.8, 3.6)
KILL_HALF_Y = 1.5
PADDLE_CONTACT_DISTANCE = 0.01
PADDLE_COOLDOWN_S = 0.12
PADDLE_RESTITUTION = 0.8
PADDLE_TANGENT_KEEP = 1.0
PADDLE_VELOCITY_GAIN = 1.0
PADDLE_SEPARATION = 0.04
PADDLE_NORMAL_LOCAL = np.array([0.0, 1.0, 0.0])
MIN_CLOSING_SPEED = 0.05
LANDING_QUALITY_STD = 0.45
# fitted bounce regression: vx' = a0 + a1 vx + a2 vz ; vy' = b0 + b1 vy ; vz' = c0 |vz| + c1 |vz|^2
_BOUNCE_A = (-0.74010409, 0.88783120, -0.28152663)
_BOUNCE_B = (-0.01392620, 0.82615466)
_BOUNCE_C = (0.91228877, -0.00550677)
def drag_step(v: np.ndarray, dt: float) -> np.ndarray:
speed = float(np.linalg.norm(v))
v = v / max(1.0 + DRAG_K * speed * dt, 1e-9)
speed = float(np.linalg.norm(v))
if speed > BALL_MAX_SPEED:
v = v * (BALL_MAX_SPEED / speed)
return v
def bounce_velocity(v: np.ndarray) -> np.ndarray:
vz_in = max(-float(v[2]), 0.0)
return np.array(
[
_BOUNCE_A[0] + _BOUNCE_A[1] * v[0] + _BOUNCE_A[2] * v[2],
_BOUNCE_B[0] + _BOUNCE_B[1] * v[1],
_BOUNCE_C[0] * vz_in + _BOUNCE_C[1] * vz_in * vz_in,
],
dtype=np.float64,
)
def _on_table(x: float, y: float) -> bool:
return TABLE_X_RANGE[0] <= x <= TABLE_X_RANGE[1] and abs(y) <= TABLE_HALF_WIDTH
def integrate(p: np.ndarray, v: np.ndarray, dt: float) -> tuple[np.ndarray, np.ndarray, bool, np.ndarray | None]:
"""One ball step: drag, gravity, position, table bounce. Returns (p, v, bounced, contact_xy)."""
last_p = p.copy()
prev_v = v.copy()
v = drag_step(v, dt)
v = v + np.array([0.0, 0.0, -GRAVITY * dt])
p = p + v * dt
z_table = TABLE_HEIGHT + BALL_RADIUS
bounced = False
contact = None
if last_p[2] > z_table >= p[2] and v[2] < 0.0:
alpha = float(np.clip((z_table - last_p[2]) / (p[2] - last_p[2]), 0.0, 1.0))
contact_xy = last_p[:2] + alpha * (p[:2] - last_p[:2])
if _on_table(float(contact_xy[0]), float(contact_xy[1])):
t1 = alpha * dt
vc = drag_step(prev_v, t1) + np.array([0.0, 0.0, -GRAVITY * t1])
pc = np.array([contact_xy[0], contact_xy[1], z_table])
vb = bounce_velocity(vc)
t2 = (1.0 - alpha) * dt
v = drag_step(vb, t2) + np.array([0.0, 0.0, -GRAVITY * t2])
p = pc + v * t2
bounced = True
contact = pc
return p, v, bounced, contact
def simulate(p: np.ndarray, v: np.ndarray, duration: float, dt: float) -> tuple[np.ndarray, np.ndarray, int]:
steps = max(1, int(round(duration / dt)))
h = duration / steps
bounces = 0
for _ in range(steps):
p, v, b, _ = integrate(p, v, h)
bounces += int(b)
return p, v, bounces
@dataclass
class Crossing:
found: bool
time: float = math.nan
pos: np.ndarray = field(default_factory=lambda: np.full(3, np.nan))
vel: np.ndarray = field(default_factory=lambda: np.full(3, np.nan))
bounces: int = 0
def predict_plane_crossing(p: np.ndarray, v: np.ndarray, plane_x: float, dt: float, max_time: float = 3.0) -> Crossing:
"""Roll the ball (travelling toward -x) until it crosses ``x = plane_x``."""
t = 0.0
bounces = 0
while t < max_time:
last_p, last_v = p.copy(), v.copy()
p, v, b, _ = integrate(p, v, dt)
bounces += int(b)
t += dt
if last_p[0] > plane_x >= p[0]:
alpha = float(np.clip((last_p[0] - plane_x) / max(last_p[0] - p[0], 1e-9), 0.0, 1.0))
pos = last_p + alpha * (p - last_p)
vel = last_v + alpha * (v - last_v)
return Crossing(True, t - dt + alpha * dt, pos, vel, bounces)
if p[2] < KILL_Z:
break
return Crossing(False)
def solve_incoming_velocity(start: np.ndarray, target: np.ndarray, t_hit: float, rng: np.random.Generator, dt: float, iters: int = 2) -> np.ndarray:
"""Training ``_solve_incoming_velocity_to_hit``: bounce-waypoint guess + shooting corrections."""
ratio = rng.uniform(0.45, 0.70)
t1 = max(ratio * t_hit, 0.08)
bounce = np.array([rng.uniform(0.75, 1.45), target[1] + rng.uniform(-0.10, 0.10), TABLE_HEIGHT + BALL_RADIUS])
g_vec = np.array([0.0, 0.0, -GRAVITY])
v = (bounce - start - 0.5 * g_vec * t1 * t1) / t1
for _ in range(iters):
pred, _, _ = simulate(start.copy(), v.copy(), t_hit, dt)
v = v + (target - pred) / max(t_hit, 1e-6)
return v
@dataclass
class IncomingBall:
start: np.ndarray
vel: np.ndarray
crossing: Crossing
def sample_incoming_ball(
target: np.ndarray,
t_hit: float,
plane_x: float,
hit_box: tuple[tuple[float, float], tuple[float, float], tuple[float, float]],
rng: np.random.Generator,
dt: float,
attempts: int = 4,
pos_tol: float = 0.04,
time_tol: float = 0.04,
) -> IncomingBall | None:
"""Training ``_reset_online_ball`` (synthetic branch): best of ``attempts`` candidates."""
(x_lo, x_hi), (y_lo, y_hi), (z_lo, z_hi) = hit_box
best: tuple[int, float, IncomingBall] | None = None
for _ in range(attempts):
start = np.array([rng.uniform(1.67, 3.24), float(np.clip(target[1] + rng.uniform(-0.25, 0.25), -0.76, 0.76)), rng.uniform(0.90, 1.50)])
v = solve_incoming_velocity(start, target, t_hit, rng, dt)
cr = predict_plane_crossing(start.copy(), v.copy(), plane_x, dt)
if not cr.found:
continue
in_range = x_lo - 1e-5 <= cr.pos[0] <= x_hi + 1e-5 and y_lo - 1e-5 <= cr.pos[1] <= y_hi + 1e-5 and z_lo - 1e-5 <= cr.pos[2] <= z_hi + 1e-5
one_bounce = cr.bounces == 1
pos_err = float(np.linalg.norm(cr.pos - target))
time_err = abs(cr.time - t_hit)
if in_range and one_bounce and pos_err <= pos_tol and time_err <= time_tol:
tier = 0
elif in_range and one_bounce:
tier = 1
elif one_bounce:
tier = 2
else:
tier = 3
score = pos_err + 0.25 * time_err
cand = IncomingBall(start, v, cr)
if best is None or (tier, score) < (best[0], best[1]):
best = (tier, score, cand)
return best[2] if best is not None else None
@dataclass
class PaddleState:
pos: np.ndarray # virtual paddle centre, world
quat: np.ndarray # wxyz, world
lin_vel: np.ndarray # world
def virtual_paddle(body_pos: np.ndarray, body_quat: np.ndarray, body_lin_vel: np.ndarray, body_ang_vel: np.ndarray,
offset_xyz: np.ndarray, offset_rpy: np.ndarray) -> PaddleState:
"""Training ``_get_virtual_paddle_state_for_ball``."""
r = quat_apply(body_quat, offset_xyz)
pos = body_pos + r
quat = quat_mul(body_quat, quat_from_euler_xyz(*offset_rpy))
vel = body_lin_vel + np.cross(body_ang_vel, r)
return PaddleState(pos, quat, vel)
@dataclass
class ReturnPrediction:
over_net: bool
on_table: bool
landing_xy: np.ndarray | None
landing_error: float
landing_quality: float
def predict_return(p: np.ndarray, v: np.ndarray, target_landing_xy: np.ndarray, dt: float, max_time: float = 3.0) -> ReturnPrediction:
"""Roll the outgoing ball: legal net crossing, first table bounce inside the return box, landing error."""
over_net = False
t = 0.0
while t < max_time:
last_p = p.copy()
p, v, bounced, contact = integrate(p, v, dt)
t += dt
if last_p[0] < NET_X <= p[0]:
alpha = (NET_X - last_p[0]) / max(p[0] - last_p[0], 1e-9)
cross = last_p + alpha * (p - last_p)
if abs(cross[1]) <= TABLE_HALF_WIDTH and cross[2] > NET_HEIGHT:
over_net = True
if bounced and contact is not None:
on_table = over_net and RETURN_X_RANGE[0] <= contact[0] <= RETURN_X_RANGE[1] and abs(contact[1]) <= TABLE_HALF_WIDTH
err = float(np.linalg.norm(contact[:2] - target_landing_xy))
return ReturnPrediction(over_net, on_table, contact[:2].copy(), err, math.exp(-(err / LANDING_QUALITY_STD) ** 2))
if p[2] < KILL_Z or p[0] > KILL_X_RANGE[1] or abs(p[1]) > KILL_HALF_Y:
break
return ReturnPrediction(over_net, False, None, math.nan, 0.0)
class KinematicBall:
"""Live ball state advanced once per control step, with paddle contact detection."""
def __init__(self, half_width: float, half_height: float):
self.half_width = float(half_width)
self.half_height = float(half_height)
self.pos = np.zeros(3)
self.vel = np.zeros(3)
self.active = False
self.hit = False
self.over_net = False
self.landed = False
self.landing_xy: np.ndarray | None = None
self.cooldown = 0.0
self._last_paddle: PaddleState | None = None
self.hit_paddle_speed = 0.0
self.outgoing_vel = np.zeros(3)
def launch(self, pos: np.ndarray, vel: np.ndarray) -> None:
self.pos = np.asarray(pos, dtype=np.float64).copy()
self.vel = np.asarray(vel, dtype=np.float64).copy()
self.active = True
self.hit = False
self.over_net = False
self.landed = False
self.landing_xy = None
self.cooldown = 0.0
self._last_paddle = None
def deactivate(self) -> None:
self.active = False
def step(self, dt: float, paddle: PaddleState | None, forehand: bool) -> dict[str, bool]:
events = {"hit": False, "landed": False}
if not self.active:
self._last_paddle = paddle
return events
self.cooldown = max(0.0, self.cooldown - dt)
last_pos = self.pos.copy()
self.pos, self.vel, bounced, contact = integrate(self.pos, self.vel, dt)
if self.hit and last_pos[0] < NET_X <= self.pos[0]:
alpha = (NET_X - last_pos[0]) / max(self.pos[0] - last_pos[0], 1e-9)
cross = last_pos + alpha * (self.pos - last_pos)
if abs(cross[1]) <= TABLE_HALF_WIDTH and cross[2] > NET_HEIGHT:
self.over_net = True
if bounced and contact is not None and self.hit and not self.landed:
self.landed = True
self.landing_xy = contact[:2].copy()
events["landed"] = True
if paddle is not None and not self.hit and self.cooldown <= 0.0:
if self._try_paddle_contact(last_pos, paddle, forehand):
events["hit"] = True
if self.pos[2] < KILL_Z or not (KILL_X_RANGE[0] <= self.pos[0] <= KILL_X_RANGE[1]) or abs(self.pos[1]) > KILL_HALF_Y:
self.active = False
self._last_paddle = paddle
return events
def _face_ok(self, u: float, w: float) -> bool:
a = self.half_width + BALL_RADIUS
b = self.half_height + BALL_RADIUS
return (u / a) ** 2 + (w / b) ** 2 <= 1.0
def _try_paddle_contact(self, last_pos: np.ndarray, paddle: PaddleState, forehand: bool) -> bool:
rot = quat_to_matrix(paddle.quat)
local = rot.T @ (self.pos - paddle.pos)
d = float(local[1])
u, w = float(local[0]), float(local[2])
contact = abs(d) <= BALL_RADIUS + PADDLE_CONTACT_DISTANCE and self._face_ok(u, w)
if not contact and self._last_paddle is not None:
rot_last = quat_to_matrix(self._last_paddle.quat)
local_last = rot_last.T @ (last_pos - self._last_paddle.pos)
d_last = float(local_last[1])
if (d_last < 0.0) != (d < 0.0) and abs(d - d_last) > 1e-9:
alpha = float(np.clip(-d_last / (d - d_last), 0.0, 1.0))
ui = local_last[0] + alpha * (u - local_last[0])
wi = local_last[2] + alpha * (w - local_last[2])
contact = self._face_ok(float(ui), float(wi))
if not contact:
return False
sign = 1.0 if forehand else -1.0
normal = rot @ (PADDLE_NORMAL_LOCAL * sign)
rel = self.vel - paddle.lin_vel
if float(np.dot(rel, normal)) > 0.0:
normal = -normal
vn = float(np.dot(rel, normal))
if vn > -MIN_CLOSING_SPEED:
return False
vt = rel - vn * normal
self.vel = PADDLE_VELOCITY_GAIN * paddle.lin_vel + (-PADDLE_RESTITUTION * vn * normal + PADDLE_TANGENT_KEEP * vt)
tangent_w = rot @ np.array([u, 0.0, w])
self.pos = paddle.pos + tangent_w + PADDLE_SEPARATION * normal
self.hit = True
self.cooldown = PADDLE_COOLDOWN_S
self.hit_paddle_speed = float(np.linalg.norm(paddle.lin_vel))
self.outgoing_vel = self.vel.copy()
return True