"""One YAM arm: kinematics, the smooth Cartesian executor, and the grasp primitive. Every hard-won fix from the demo scripts lives here once, instead of being copy-pasted per task: * the integral correction and the last COMMANDED point persist ACROSS segments, so a phase boundary does not snap the pose (that snap was the visible "pause then jump" in early videos); * segment rates are eased in/out, with the duration stretched by pi/2 so the mid-path speed does not rise and fling the carried object; * polyline corners are filleted, so lift->carry is one arc instead of stop-and-turn; * a clamp only counts as a grasp if it stalls at a plausible object-sized gap, and the object is then verified to actually rise with the gripper. """ from __future__ import annotations from dataclasses import dataclass import numpy as np OPEN, CLOSE = 1.0, -1.0 def quat_to_mat(q): w, x, y, z = [float(v) for v in q] return np.array([[1-2*(y*y+z*z), 2*(x*y-z*w), 2*(x*z+y*w)], [2*(x*y+z*w), 1-2*(x*x+z*z), 2*(y*z-x*w)], [2*(x*z-y*w), 2*(y*z+x*w), 1-2*(x*x+y*y)]]) def mat_to_quat(m): t = m[0, 0]+m[1, 1]+m[2, 2] if t > 0: s = np.sqrt(t+1)*2; w = .25*s; x = (m[2, 1]-m[1, 2])/s; y = (m[0, 2]-m[2, 0])/s; z = (m[1, 0]-m[0, 1])/s elif m[0, 0] > m[1, 1] and m[0, 0] > m[2, 2]: s = np.sqrt(1+m[0, 0]-m[1, 1]-m[2, 2])*2; w = (m[2, 1]-m[1, 2])/s; x = .25*s; y = (m[0, 1]+m[1, 0])/s; z = (m[0, 2]+m[2, 0])/s elif m[1, 1] > m[2, 2]: s = np.sqrt(1+m[1, 1]-m[0, 0]-m[2, 2])*2; w = (m[0, 2]-m[2, 0])/s; x = (m[0, 1]+m[1, 0])/s; y = .25*s; z = (m[1, 2]+m[2, 1])/s else: s = np.sqrt(1+m[2, 2]-m[0, 0]-m[1, 1])*2; w = (m[1, 0]-m[0, 1])/s; x = (m[0, 2]+m[2, 0])/s; y = (m[1, 2]+m[2, 1])/s; z = .25*s q = np.array([w, x, y, z]); q /= np.linalg.norm(q)+1e-9 return q if q[0] >= 0 else -q def grasp_quat(jaw="y", yaw_deg=0.0, tilt_deg=0.0, tilt_sign=1.0): """Top-down grasp orientation. jaw which world axis the JAW CLOSES ALONG. It must be perpendicular to the thing being gripped: closing along a protruding handle pinches its length and holds nothing. yaw_deg rotate the whole grasp about z (for an object lying at an angle). tilt_deg lean off vertical, to hook around a handle from outside rather than press on it. """ if jaw == "x": base = np.stack([np.array([1., 0., 0.]), np.array([0., -1., 0.]), np.array([0., 0., -1.])], axis=1) else: base = np.stack([np.array([0., 1., 0.]), np.array([1., 0., 0.]), np.array([0., 0., -1.])], axis=1) if tilt_deg: t = np.radians(tilt_deg); c, s = np.cos(t), np.sin(t) X = base[:, 0] Z = np.array([0., -tilt_sign*s, -c]) base = np.stack([X, np.cross(Z, X), Z], axis=1) if yaw_deg: a = np.radians(yaw_deg); ca, sa = np.cos(a), np.sin(a) base = np.array([[ca, -sa, 0], [sa, ca, 0], [0, 0, 1]])@base return mat_to_quat(base) def ease(a: float) -> float: """Cosine ease so a segment starts and ends at zero velocity.""" return float(0.5-0.5*np.cos(np.pi*min(max(a, 0.0), 1.0))) def fillet(pts, r=0.06, n=6): """Round the interior corners of a polyline with quadratic Beziers.""" pts = [np.asarray(p, np.float32) for p in pts] if len(pts) < 3: return pts out = [pts[0]] for i in range(1, len(pts)-1): p0, p1, p2 = pts[i-1], pts[i], pts[i+1] d0, d2 = p1-p0, p2-p1 l0, l2 = float(np.linalg.norm(d0)), float(np.linalg.norm(d2)) rr = min(r, 0.45*l0, 0.45*l2) if rr < 1e-4 or l0 < 1e-6 or l2 < 1e-6: out.append(p1); continue a, b = p1-d0/l0*rr, p1+d2/l2*rr out.append(a) for k in range(1, n): t = k/float(n); out.append(((1-t)**2)*a + (2*(1-t)*t)*p1 + (t*t)*b) out.append(b) out.append(pts[-1]) return out def _point_at(poly, seglens, s): acc = 0.0 for i, L in enumerate(seglens): if acc+L >= s or i == len(seglens)-1: t = min(max((s-acc)/max(L, 1e-6), 0.0), 1.0) return poly[i]+(poly[i+1]-poly[i])*t acc += L return poly[-1] @dataclass class GraspResult: """Outcome of a grasp attempt -- `ok` is False unless the object VERIFIABLY came along.""" ok: bool reason: str hold_pose: np.ndarray | None = None finger_sep: float = 0.0 rise: float = 0.0 class ArmController: """Drives one arm in its own root frame. `step_fn(cmd_pos, cmd_quat, grip)` is supplied by the task env: it packs this arm's command into whatever action vector the environment expects (the other arm may be held, frozen, or driven by its own ArmController) and steps the sim once. """ EEF_OFFSET = np.array([0.0, 0.0, 0.13]) # link_6 -> grasp reference point def __init__(self, articulation, body_names, root_pos, root_quat, origin, step_fn, on_step=None, name="arm"): self.art = articulation self.bn = list(body_names) self.root = np.asarray(root_pos, np.float64) self.rootq = np.asarray(root_quat, np.float64) self.origin = np.asarray(origin, np.float64) self.step_fn = step_fn self.on_step = on_step # called once per sim step (recording, probes) self.name = name self.quat = grasp_quat("y") self._corr = np.zeros(3, np.float32) # integral correction, PERSISTENT self._cmd = None # last COMMANDED point, PERSISTENT # ---------------- kinematics ---------------- def eef(self): i = self.bn.index("link_6") p = self.art.data.body_pos_w[0, i].cpu().numpy()-self.origin q = self.art.data.body_quat_w[0, i].cpu().numpy() world = p+quat_to_mat(q)@self.EEF_OFFSET return (quat_to_mat(self.rootq).T@(world-self.root)).astype(np.float32) def to_root(self, world_xyz): return (quat_to_mat(self.rootq).T@(np.asarray(world_xyz, np.float64)-self.root)).astype(np.float32) def finger_sep(self): jn = list(self.art.data.joint_names) return (float(self.art.data.joint_pos[0, jn.index("left_finger")].item()) + float(self.art.data.joint_pos[0, jn.index("right_finger")].item()))/2 # ---------------- executor ---------------- def _seg_start(self): return self._cmd.copy() if self._cmd is not None else self.eef() def _drive(self, cp, grip): """One step: command cp(+correction), then update the integral correction.""" self._cmd = np.asarray(cp, np.float32) self.step_fn(self, (cp+self._corr).astype(np.float32), self.quat, grip) e = cp-self.eef() e = np.where(np.abs(e) > 0.008, e, 0.0) self._corr = np.clip(self._corr+0.08*e, -0.10, 0.10) self._corr[2] = max(float(self._corr[2]), -0.06) if self.on_step is not None: self.on_step() def flow(self, pts, grip, speed=0.008, settle=16, round_corners=True): """Glide along a polyline at an eased constant rate. Corners are rounded by default.""" start = self._seg_start() poly = [start]+[np.asarray(p, np.float32) for p in pts] if round_corners: poly = fillet(poly) seglens = [float(np.linalg.norm(poly[i+1]-poly[i])) for i in range(len(poly)-1)] total = float(sum(seglens)) if total < 1e-6: return 0.0 # cosine ease peaks at pi/2 x the mean rate -- stretch the duration to match, or the # mid-path speed rises ~57% and flings whatever is in the jaws nsteps = max(int(total/speed*(np.pi/2)), 1) for k in range(nsteps+settle): a = ease(min(1.0, (k+1)/float(nsteps))) self._drive(_point_at(poly, seglens, min(total, a*total)), grip) return float(np.linalg.norm(poly[-1]-self.eef())) def move_to(self, target, grip, tol=0.008, max_steps=140): """Closed-loop move: converge onto the target and STOP, instead of over-commanding.""" start = self._seg_start(); tgt = np.asarray(target, np.float32) ramp = max(int(max_steps*0.6), 18) for k in range(max_steps): a = ease(min(1.0, (k+1)/float(ramp))) self._drive((1-a)*start+a*tgt, grip) if a >= 1.0 and np.linalg.norm(self.eef()-tgt) < tol: break return float(np.linalg.norm(self.eef()-tgt)) def hold(self, n, grip): for _ in range(n): self._drive(self._seg_start(), grip) # ---------------- grasp ---------------- def grasp(self, object_z_fn, max_gap=0.045, steps=170, verify_lift=0.05): """Close until the jaws stall, then VERIFY the object rises with the gripper. `max_gap` rejects a stall that happened at an implausibly wide opening -- that is a jaw resting on the object's body (or on the table), not a grip, and accepting it is how an episode ends up miming the whole sequence with an empty hand. """ pose = self.eef().astype(np.float32) self._cmd = pose.copy() prev, stall = self.finger_sep(), 0 stalled = False for _ in range(steps): self._drive(pose, CLOSE) cur = self.finger_sep() stall = stall+1 if abs(cur-prev) < 0.0002 else 0 prev = cur if stall >= 8 and -max_gap < cur < -0.002: stalled = True break sep = self.finger_sep() if not stalled: return GraspResult(False, f"jaws never stalled on an object-sized gap (sep={sep:.4f}, " f"gap={2*abs(sep)*100:.1f} cm)", finger_sep=sep) if verify_lift <= 0.0: # No solo lift check: for a two-arm grip the object cannot rise until BOTH hands # hold it, so the shared lift is the verification and a per-arm one always "fails". return GraspResult(True, f"clamped (gap {2*abs(sep)*100:.1f} cm, lift check deferred)", hold_pose=pose, finger_sep=sep) z0 = float(object_z_fn()) for _ in range(45): self._drive(pose+np.array([0, 0, verify_lift], np.float32), CLOSE) rise = float(object_z_fn())-z0 if rise < 0.02: return GraspResult(False, f"object did not rise with the gripper ({rise:+.3f} m): the " f"jaw closed beside or above it", finger_sep=sep, rise=rise) return GraspResult(True, f"grasped (gap {2*abs(sep)*100:.1f} cm, rose {rise:+.3f} m)", hold_pose=pose, finger_sep=sep, rise=rise)