"""Touch the red sphere with either hand (Unitree G1, scripted).
Pipeline
1. Locate the sphere in the head camera: fit a circle (in ray space) to the
red blob boundary -> direction + angular radius -> 3D point, using a
pinhole model of the head camera calibrated from the floor checkerboard.
2. Store the target in world coordinates (pelvis odometry) and walk the base
so the target sits at a comfortable spot in front of one shoulder.
3. Re-observe, then reach with an analytic G1 arm model + numerical IK:
sweep the hand forward through the estimate; the sphere lights up while
touched, so the lit stretch of the sweep gives the touch-zone centre
(then z / y sweeps if needed), and the hand holds there.
The arms hang down (out of the head camera view) until the reach starts.
"""
import cv2
import numpy as np
from scipy.optimize import least_squares
# ---------------------------------------------------------------- camera model
F = 72.746 # focal length in px (fovy 60 deg, 84 px)
CX = CY = 41.5
CAM_PITCH = 1.093 # rad below horizontal
CAM_OFF = np.array([0.07, 0.0, 0.40]) # head camera position in pelvis frame
R_SPHERE = 0.047 # effective sphere radius for the boundary detector
def cam_R(th, psi=0.0):
"""Camera->frame rotation (camera x right, y down, z forward)."""
z = np.array([np.cos(th) * np.cos(psi), np.cos(th) * np.sin(psi), -np.sin(th)])
x = np.array([np.sin(psi), -np.cos(psi), 0.0])
y = np.cross(z, x)
return np.stack([x, y, z], 1)
R_CAM = cam_R(CAM_PITCH)
def red_mask(im):
im = im.astype(np.int32)
r, g, b = im[..., 0], im[..., 1], im[..., 2]
return (r > 40) & (r > 2 * g + 8) & (r > 2 * b + 8)
def bright_count(im):
im = im.astype(np.int32)
r, g, b = im[..., 0], im[..., 1], im[..., 2]
return int(((r > 215) & (g < 110) & (b < 110)).sum())
def fit_sphere(im):
"""Fit the red blob. Returns (unit ray in camera frame, angular radius, n boundary pts)."""
m = red_mask(im)
if m.sum() < 4:
return None
n, lab, stats, _ = cv2.connectedComponentsWithStats(m.astype(np.uint8), connectivity=8)
k = 1 + int(np.argmax(stats[1:, cv2.CC_STAT_AREA]))
mk = (lab == k).astype(np.uint8)
cs, _ = cv2.findContours(mk, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
pts = np.concatenate([c.reshape(-1, 2) for c in cs], 0).astype(float)
h, w = m.shape
keep = (pts[:, 0] > 0) & (pts[:, 0] < w - 1) & (pts[:, 1] > 0) & (pts[:, 1] < h - 1)
pts = pts[keep]
if len(pts) < 6:
return None
ys, xs = np.nonzero(mk)
rays = np.stack([(pts[:, 0] - CX) / F, (pts[:, 1] - CY) / F, np.ones(len(pts))], 1)
rays /= np.linalg.norm(rays, axis=1, keepdims=True)
def dirf(p):
d = np.array([p[0], p[1], 1.0])
return d / np.linalg.norm(d)
x0 = [(xs.mean() - CX) / F, (ys.mean() - CY) / F, np.sqrt(mk.sum() / np.pi) / F]
def res(p):
return np.arccos(np.clip(rays @ dirf(p[:2]), -1, 1)) - p[2]
r = least_squares(res, x0)
return dirf(r.x[:2]), r.x[2] + 0.5 / F, len(pts)
def locate(im):
"""Sphere centre in the pelvis frame (x fwd, y left, z up rel. pelvis)."""
s = fit_sphere(im)
if s is None:
return None
d, alpha, npts = s
if alpha <= 0.01:
return None
D = R_SPHERE / np.sin(alpha)
return R_CAM @ (D * d) + CAM_OFF, npts
# ------------------------------------------------------------------ arm model
def rx(a):
c, s = np.cos(a), np.sin(a)
return np.array([[1, 0, 0], [0, c, -s], [0, s, c]])
def ry(a):
c, s = np.cos(a), np.sin(a)
return np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]])
def rz(a):
c, s = np.cos(a), np.sin(a)
return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])
TORSO = np.array([-0.0039635, 0, 0.054])
Q16 = 2 * np.arcsin(0.139201)
TOOL = np.array([0.10, 0.0, 0.0])
def fk(q, side):
"""Tool point and wrist rotation (pelvis frame) for arm joints q (7)."""
s = side
R = np.eye(3)
p = TORSO + np.array([0.0039563, s * 0.10022, 0.23778])
R = rx(s * Q16) @ ry(q[0])
p = p + R @ np.array([0, s * 0.038, -0.013831])
R = R @ rx(-s * Q16) @ rx(q[1])
p = p + R @ np.array([0, s * 0.00624, -0.1032])
R = R @ rz(q[2])
p = p + R @ np.array([0.015783, 0, -0.080518])
R = R @ ry(q[3])
p = p + R @ np.array([0.1, s * 0.00188791, -0.01])
R = R @ rx(q[4])
p = p + R @ np.array([0.038, 0, 0])
R = R @ ry(q[5])
p = p + R @ np.array([0.046, 0, 0])
R = R @ rz(q[6])
return p + R @ TOOL, R
IDX = [0, 1, 2, 3, 5] # joints used by the IK: shoulder pitch/roll/yaw, elbow, wrist pitch
# sane working ranges (left arm) for sp, sr, sy, el, wp; roll/yaw mirrored for the right arm
W_LO = np.array([-1.9, 0.0, -0.9, -0.9, -1.0])
W_HI = np.array([0.5, 0.9, 0.9, 1.6, 1.0])
def limits(side):
lo, hi = W_LO.copy(), W_HI.copy()
if side < 0:
for i in (1, 2):
lo[i], hi[i] = -W_HI[i], -W_LO[i]
return lo, hi
def ik(p, side, q0=None, approach=np.array([1.0, 0, 0]), w_ori=0.05):
lo, hi = limits(side)
if q0 is None:
starts = [np.clip(np.array(s, float) * np.array([1, side, side, 1, 1]), lo + 1e-3, hi - 1e-3)
for s in [(-0.3, 0.1, 0, -0.5, 0), (-0.8, 0.1, 0, -0.3, 0), (-0.5, 0.2, 0, 0.3, 0),
(-1.2, 0.1, 0, 0.3, 0), (0, 0, 0, 0, 0)]]
else:
starts = [np.clip(np.asarray(q0)[IDX], lo + 1e-3, hi - 1e-3)]
def res(x):
q = np.zeros(7)
q[IDX] = x
tip, R = fk(q, side)
return np.r_[(tip - p) * 30, w_ori * (R[:, 0] - approach), 0.003 * x]
best = None
for s0 in starts:
r = least_squares(res, s0, bounds=(lo, hi))
if best is None or r.cost < best.cost:
best = r
q = np.zeros(7)
q[IDX] = best.x
return q, float(np.linalg.norm(fk(q, side)[0] - p))
# --------------------------------------------------------------------- policy
def to_world(p_rel, pel):
x, y, z, yaw = pel
c, s = np.cos(yaw), np.sin(yaw)
return np.array([x + c * p_rel[0] - s * p_rel[1], y + s * p_rel[0] + c * p_rel[1], z + p_rel[2]])
def to_rel(p_w, pel):
x, y, z, yaw = pel
c, s = np.cos(yaw), np.sin(yaw)
dx, dy = p_w[0] - x, p_w[1] - y
return np.array([c * dx + s * dy, -s * dx + c * dy, p_w[2] - z])
X_DES = 0.32 # desired target x in pelvis frame before reaching
Y_DES = 0.12 # desired |y| of target (in front of the reaching shoulder)
SX0, SX1 = -0.06, 0.08 # x sweep range (relative to the estimate)
SW = 0.05 # half range of the lateral (z, y) sweeps
S_SPEED = 0.0015 # sweep speed, m per control step
V_MOVE = 0.003 # repositioning speed, m per control step
KI = 0.05 # joint-space integral gain (gravity sag)
# x-sweep lines (dy, dz) tried in turn around the estimate
OFFSETS = [(0, 0), (0, 0.025), (0, -0.025), (0.025, 0), (-0.025, 0), (0.025, 0.025), (-0.025, 0.025),
(0.025, -0.025), (-0.025, -0.025), (0, 0.05), (0, -0.05), (0.05, 0), (-0.05, 0)]
AXES = {"x": 0, "y": 1, "z": 2}
BASE0 = np.array([0.01, 0.0, 0.02]) # systematic offset of the touch zone from the estimate
# arms hang down while observing/walking so the grippers never hide the sphere
CLEAR_L = np.array([0.0, 0.15, 0.0, 1.3, 0.0, 0.0, 0.0])
CLEAR_R = np.array([0.0, -0.15, 0.0, 1.3, 0.0, 0.0, 0.0])
T_CLEAR = 30 # steps for the arms to get out of the view
SETTLE = 70 # steps standing still after the walk
SETTLE_OBS = 40 # observe from this step of the settle phase on
class Policy:
"""The control policy that is scored on the hidden seeds."""
def reset(self, obs, tools):
self.hold = np.asarray(tools.hold_action(), dtype=np.float64)
self.phase = "observe"
self.t0 = 0
self.obs_buf = []
self.target_w = None
self.side = 1
self.q = np.zeros(7)
self.qi = np.zeros(7)
self.log = []
self.line = 0
self.base = np.zeros(3) # centre of the current search (offset from estimate)
self.g = np.zeros(3) # commanded tip offset from the estimate
self.mode = None
self.unlit = 0
def pel(self, obs):
return np.asarray(obs["low_dim_obs"][18:22], dtype=np.float64)
def observe(self, obs, tools):
im = tools.image("head")
r = locate(im)
if r is None:
return None
p, npts = r
return to_world(p, self.pel(obs)), npts
# ------------------------------------------------------------ search program
def start_xsweep(self):
dy, dz = OFFSETS[self.line % len(OFFSETS)]
c = self.base + np.array([0.0, dy, dz])
self.go(c + np.array([SX0, 0, 0]), ("sweep", "x", c[0] + SX1))
def go(self, p, then):
self.g_to = np.asarray(p, float)
self.mode = "move"
self.after = then
def begin_sweep(self, axis, end):
self.mode = "sweep"
self.axis = axis
self.end = end
self.end0 = end
self.lit_vals = []
self.unlit = 0
def sweep_done(self):
ax = AXES[self.axis]
if self.lit_vals:
v = np.array(self.lit_vals)
centre = 0.5 * (v.min() + v.max())
else:
centre = None
self.log.append(("sweep", self.axis, None if centre is None else round(float(centre), 4),
len(self.lit_vals), self.base.round(3).tolist()))
if self.axis == "x":
if centre is None:
self.line += 1
self.start_xsweep()
return
dy, dz = OFFSETS[self.line % len(OFFSETS)]
self.cur = np.array([centre, self.base[1] + dy, self.base[2] + dz])
self.go(self.cur, ("hold", "x"))
elif self.axis == "z":
if centre is not None:
self.cur[2] = centre
self.go(self.cur + np.array([0, -SW, 0]), ("sweep", "y", self.cur[1] + SW))
else:
if centre is not None:
self.cur[1] = centre
self.go(self.cur, ("hold",))
def step_search(self, lit, tip_off):
m = self.mode
if m == "move":
d = self.g_to - self.g
n = np.linalg.norm(d)
if n <= V_MOVE:
self.g = self.g_to.copy()
if self.after[0] == "sweep":
self.begin_sweep(self.after[1], self.after[2])
else:
self.mode = "hold"
self.hold_kind = self.after[1] if len(self.after) > 1 else "final"
self.unlit = 0
else:
self.g = self.g + d / n * V_MOVE
elif m == "sweep":
ax = AXES[self.axis]
if lit:
self.lit_vals.append(tip_off[ax])
self.unlit = 0
else:
self.unlit += 1
if lit and self.g[ax] >= self.end and self.end < self.end0 + 0.04:
self.end += S_SPEED # still lit at the end of the range: keep going
if (self.lit_vals and self.unlit >= 3) or self.g[ax] >= self.end:
self.sweep_done()
else:
self.g[ax] += S_SPEED
elif m == "hold":
self.unlit = 0 if lit else self.unlit + 1
if self.hold_kind == "x" and self.unlit > 20:
# x centre alone is not enough: centre z, then y
self.go(self.cur + np.array([0, 0, -SW]), ("sweep", "z", self.cur[2] + SW))
elif self.unlit > 40:
# lost it: search again around the last centre
self.base = self.cur.copy()
self.line = 0
self.start_xsweep()
# ------------------------------------------------------------------- act
def act(self, obs, tools):
t = int(obs["t"])
a = self.hold.copy()
pel = self.pel(obs)
low = np.asarray(obs["low_dim_obs"], dtype=np.float64)
# both arms hang clear; the reaching arm is overwritten in the reach phase
a[4:11] = CLEAR_L
a[11:18] = CLEAR_R
if self.phase == "observe":
r = self.observe(obs, tools) if t >= T_CLEAR else None
if r is not None and r[1] >= 20:
self.obs_buf.append(r[0])
self.log.append(("obs", t, pel.round(3).tolist(), None if r is None else (r[0].round(3).tolist(), r[1])))
if len(self.obs_buf) >= 8:
self.target_w = np.median(np.array(self.obs_buf), 0)
self.first_w = self.target_w.copy()
rel = to_rel(self.target_w, pel)
self.side = 1 if rel[1] >= 0 else -1
self.phase = "walk"
self.t0 = t
elif T_CLEAR + 10 < t < T_CLEAR + 160:
a[0] = 0.12 # sphere not well visible: creep forward
return a
rel = to_rel(self.target_w, pel)
if self.phase == "walk":
ex = rel[0] - X_DES
ey = rel[1] - self.side * Y_DES
if abs(ex) < 0.02 and abs(ey) < 0.02 or t - self.t0 > 250:
self.phase = "settle"
self.t0 = t
self.obs_buf = []
self.log.append(("walked", t, pel.round(3).tolist(), rel.round(3).tolist()))
else:
vx = np.clip(1.5 * ex, -0.25, 0.25)
vy = np.clip(1.5 * ey, -0.25, 0.25)
sp = np.hypot(vx, vy)
if sp < 0.07:
k = 0.07 / max(sp, 1e-6)
vx, vy = vx * k, vy * k
a[0], a[1] = vx, vy
return a
if self.phase == "settle":
if t - self.t0 >= SETTLE_OBS:
r = self.observe(obs, tools)
if r is not None and r[1] >= 20:
self.obs_buf.append(r[0])
self.log.append(("settle", t, pel.round(3).tolist(), None if r is None else (r[0].round(3).tolist(), r[1])))
if t - self.t0 >= SETTLE:
if len(self.obs_buf) >= 5:
est = np.median(np.array(self.obs_buf), 0)
if np.linalg.norm(est - self.first_w) < 0.08:
self.target_w = est
self.log.append(("target", t, self.target_w.round(3).tolist(), self.first_w.round(3).tolist()))
self.phase = "reach"
self.t0 = t
self.line = 0
self.base = BASE0.copy()
dy, dz = OFFSETS[0]
self.g = self.base + np.array([SX0, dy, dz])
self.start_xsweep()
self.q_from = (CLEAR_L if self.side > 0 else CLEAR_R).copy()
rel = to_rel(self.target_w, pel)
else:
return a
# ---------------- reach
side = self.side
sl = slice(4, 11) if side > 0 else slice(11, 18)
qm = low[0:7] if side > 0 else low[9:16]
tip_off = fk(qm, side)[0] - rel
lit = bright_count(tools.image("head")) >= 3
tau = t - self.t0
if tau >= 40:
self.step_search(lit, tip_off)
q, err = ik(rel + self.g, side, None if tau == 0 else self.q)
self.q = q
if tau < 40:
w = (tau + 1) / 40.0
cmd = (1 - w) * self.q_from + w * q
else:
self.qi[IDX] = np.clip(self.qi[IDX] + KI * (q - qm)[IDX], -0.15, 0.15)
cmd = q + self.qi
a[sl] = cmd
self.log.append(("reach", t, self.mode, self.line, self.g.round(4).tolist(), tip_off.round(4).tolist(),
rel.round(3).tolist(), round(err, 4), int(lit)))
return a