source/bimanual/yam/motion/arm.py

"""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)