| """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]) |
|
|
| 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 |
| self.name = name |
| self.quat = grasp_quat("y") |
| self._corr = np.zeros(3, np.float32) |
| self._cmd = None |
|
|
| |
| 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 |
|
|
| |
| 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 |
| |
| |
| 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) |
|
|
| |
| 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: |
| |
| |
| 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) |
|
|
