source/bimanual/yam/solvers/base.py

"""Shared grasp/place primitives every solver builds on."""
from __future__ import annotations

import numpy as np

from ..motion.arm import OPEN, CLOSE, grasp_quat, fillet
from ..motion.planner import plan_path
from ..envs.scene import TABLE_TOP


def grasp_pose(env, name, grasp_top=None, offset=(0.0, 0.0), live_z=False):
    """LIVE xy from physics + a height from the measured size. Never the bbox for position.

    `offset` shifts the grip along world xy away from the centroid. A long part only presents a
    graspable stretch at one end -- a bar sitting in a sleeve has its centroid buried inside the
    housing, so gripping the centroid means gripping the housing.

    `live_z` takes the height from physics too. The default assumes the object is standing on the
    TABLE, which is wrong for anything resting on something else: a lid sitting on a 7.5 cm pot
    rim is grasped 7.5 cm too low, so the hand drives into the pot wall and the jaws shut beside
    the lid. Pass live_z=True whenever the object starts on a fixture.
    """
    live = env.scene.object_pos(name)
    ext = env.scene.object_size(name)
    base = float(live[2])-float(ext[2])/2.0 if live_z else TABLE_TOP
    z = base+float(ext[2])-float(grasp_top) if grasp_top is not None else base+float(ext[2])/2.0
    return np.array([live[0]+float(offset[0]), live[1]+float(offset[1]),
                     max(z, TABLE_TOP+0.010)], np.float32), ext


JAW_MAX = 0.094           # measured full opening of the YAM gripper
JAW_SAFE = 0.078          # beyond this the fingers clip the object while descending


def check_graspable(env, name, jaw=None):
    """Warn when the axis the jaws must close across is too wide for them.

    Every "jaw closed beside it" failure in this suite traces back to an object scaled past the
    9.4 cm opening: the fingers strike its sides on the way down, the descent stalls a few
    centimetres short, and the clamp then shuts on air. It is invisible in the numbers unless
    someone divides the bbox by the jaw, so do it here.
    """
    ext = env.scene.object_size(name)
    axis = jaw if jaw in ("x", "y") else auto_jaw(env, name)
    w = float(ext[0] if axis == "x" else ext[1])
    if w > JAW_SAFE:
        print(f"[solver] WARNING: {name} is {w*100:.1f} cm across the '{axis}' axis, but the jaw "
              f"opens {JAW_MAX*100:.1f} cm (safe <= {JAW_SAFE*100:.0f} cm). The fingers will clip "
              f"it on the way down -- scale the asset down or grip a narrower feature.", flush=True)
    return w


def auto_jaw(env, name):
    """Close across the NARROWER horizontal extent -- the jaw only opens 9.4 cm."""
    ext = env.scene.object_size(name)
    return "x" if float(ext[0]) < float(ext[1])*0.9 else "y"


def approach_and_grasp(env, arm, name, jaw="auto", grasp_top=None, max_gap=0.045,
                       grasp_yaw=0.0, grasp_offset=(0.0, 0.0), verify_lift=None, live_z=False,
                       prefix=""):
    """PRM approach -> closed-loop descend -> clamp -> VERIFIED grasp."""
    rec = env.recorder
    target, ext = grasp_pose(env, name, grasp_top, offset=grasp_offset, live_z=live_z)
    check_graspable(env, name, None if jaw == "auto" else jaw)
    arm.quat = grasp_quat(auto_jaw(env, name) if jaw == "auto" else jaw, yaw_deg=grasp_yaw)
    goal = arm.to_root(target)
    rec.phase = f"{prefix}1. PRM APPROACH"
    arm.flow(plan_path(arm.eef(), goal+np.array([0, 0, 0.12], np.float32), arm.root), OPEN)
    rec.phase = f"{prefix}2. DESCEND onto {name}"
    err = arm.move_to(goal, OPEN, tol=0.008, max_steps=140)
    arm.hold(10, OPEN)
    print(f"[solver] descend err={err:.3f} target={np.round(goal,3)}", flush=True)
    rec.phase = f"{prefix}3. CLOSE-GRASP"
    kw = {} if verify_lift is None else {"verify_lift": verify_lift}
    res = arm.grasp(lambda: env.scene.object_pos(name)[2], max_gap=max_gap, **kw)
    print(f"[solver] grasp: {res.reason}", flush=True)
    return res, ext


def place_into(env, arm, region, half_h, approach=(0.0, -1.0), standoff=0.16, tilt_deg=0.0,
               prefix=""):
    """Slide an object in through a fixture's open face, optionally TILTED.

    place_at() carries over the target and lowers straight down, which is impossible for anything
    with a roof: the cupboard's top panel blocks the descent, so the object is released above the
    unit and lands 11 cm away on the table. Here the hand rises to shelf height OUTSIDE the
    opening, drives in, releases, and backs out the way it came.

    `tilt_deg` leans the wrist so the object enters nose-first. A cupboard opening is barely
    taller than the box, and a level box catches its top corner on the shelf above; tilted, the
    leading edge goes under the lip first and the box levels out as it is released.
    """
    rec = env.recorder
    reg = env.scene.regions[region]
    level_quat = arm.quat.copy()
    d = np.asarray(approach, float)
    d = d/(np.linalg.norm(d)+1e-9)
    z = float(reg.get("shelf_z", reg.get("top_z", TABLE_TOP)))+half_h+0.008
    centre = np.array([reg["xy"][0], reg["xy"][1]], float)
    outside = centre-d*standoff

    rec.phase = f"{prefix}4. RISE to shelf height, clear of the opening"
    arm.flow([arm._seg_start()+np.array([0, 0, 0.10], np.float32)], CLOSE)
    arm.move_to(arm.to_root(np.array([outside[0], outside[1], z+0.02], np.float32)), CLOSE,
                tol=0.012, max_steps=150)
    if abs(tilt_deg) > 1e-6:
        rec.phase = f"{prefix}4b. TILT the box nose-down to clear the shelf lip"
        arm.quat = grasp_quat(auto_jaw(env, env._last_placed) if getattr(env, "_last_placed", None)
                              else "y", tilt_deg=abs(tilt_deg),
                              tilt_sign=1.0 if tilt_deg > 0 else -1.0)
        arm.hold(18, CLOSE)
    rec.phase = f"{prefix}5. SLIDE IN through the open face"
    arm.move_to(arm.to_root(np.array([outside[0], outside[1], z], np.float32)), CLOSE,
                tol=0.010, max_steps=90)
    err = arm.move_to(arm.to_root(np.array([centre[0], centre[1], z], np.float32)), CLOSE,
                      tol=0.012, max_steps=170)
    if abs(tilt_deg) > 1e-6:
        rec.phase = f"{prefix}5b. LEVEL OFF so the box sets down flat"
        arm.quat = level_quat
        arm.hold(22, CLOSE)
    rec.phase = f"{prefix}6. RELEASE on the shelf"
    arm.hold(20, OPEN)
    rec.phase = f"{prefix}7. BACK OUT the way it came in"
    arm.move_to(arm.to_root(np.array([outside[0], outside[1], z], np.float32)), OPEN,
                tol=0.02, max_steps=140)
    arm.flow([arm._seg_start()+np.array([0, 0, 0.14], np.float32)], OPEN)
    fin = env.scene.object_pos(env._last_placed) if hasattr(env, "_last_placed") else None
    print(f"[solver] shelf insert tracking err {err:.3f}"
          + (f", object at {np.round(fin, 3)}" if fin is not None else ""), flush=True)
    return err


def place_at(env, arm, region, half_h, on_top=False, lift=0.16, dx=0.0, prefix=""):
    """Carry to a region and set down as ONE filleted arc, then release.

    `dx` spreads multiple objects inside the region. It is applied TANGENTIALLY (perpendicular to
    the arm->region direction) rather than along world x: an x-offset pushes one drop point
    further from the shoulder than the other, and past ~0.35 m the arm simply cannot get there,
    so the second object of a pair would fail while the first succeeded.
    """
    rec = env.recorder
    reg = env.scene.regions[region]
    rim = float(reg.get("top_z", TABLE_TOP))-TABLE_TOP
    drop_z = TABLE_TOP+(rim+half_h+0.006 if on_top else max(0.010, rim)+half_h+0.02)
    centre = np.array([reg["xy"][0], reg["xy"][1]], np.float64)
    radial = centre-np.asarray(arm.root[:2], np.float64)
    n = np.linalg.norm(radial)
    tangent = np.array([-radial[1], radial[0]])/n if n > 1e-6 else np.array([1.0, 0.0])
    spot = centre+tangent*dx
    place = arm.to_root(np.array([spot[0], spot[1], drop_z], np.float32))
    above = place+np.array([0, 0, 0.10], np.float32)
    top = arm._seg_start()+np.array([0, 0, lift], np.float32)
    rec.phase = f"{prefix}4. LIFT + CARRY to {region}"
    arm.flow(fillet([arm._seg_start(), top, above, place])[1:], CLOSE)
    rec.phase = f"{prefix}5. RELEASE"
    arm.move_to(place, OPEN, tol=0.01, max_steps=45)
    # measure BEFORE retreating: measuring after reports how far the arm backed off, not how
    # accurately it placed (that read 0.37 m on an episode that placed the can perfectly)
    err = float(np.linalg.norm(arm.eef()-place))
    rec.phase = f"{prefix}6. RETREAT"
    arm.flow([above], OPEN)
    return err