source/bimanual/yam/solvers/tool_use.py

"""Tool skills: grasp a tool, then drive its tip along a work path (wipe / sweep)."""
import numpy as np

from ..motion.arm import OPEN, CLOSE
from ..envs.scene import TABLE_TOP
from .base import approach_and_grasp


def _work_path(kind, reg, work_z, size):
    half = size/2
    if kind == "wipe":                      # lawn-mower lanes over the region
        pts = []
        for i, yy in enumerate(np.linspace(-half, half, 4)):
            xs = [-half, half] if i % 2 == 0 else [half, -half]
            pts += [np.array([reg["xy"][0]+xs[0], reg["xy"][1]+yy, work_z], np.float32),
                    np.array([reg["xy"][0]+xs[1], reg["xy"][1]+yy, work_z], np.float32)]
        return pts
    return [np.array([reg["xy"][0], reg["xy"][1]+0.02, work_z], np.float32)]


def solve(env, tool, target, kind="wipe", arm="right", debris=None, max_gap=0.045,
          grasp_top=None, press=0.004, grasp_offset=(0.0, 0.0), jaw="auto",
          live_z=False, work_h=None):
    a = env.arms[arm]
    rec = env.recorder
    res, ext = approach_and_grasp(env, a, tool, grasp_top=grasp_top, max_gap=max_gap,
                                  grasp_offset=grasp_offset, jaw=jaw, live_z=live_z)
    if not res.ok:
        rec.phase = "ABORT: the tool was never grasped"
        return {"grasped": False, "reason": res.reason}
    reg = env.scene.regions[target]
    # Working height = how far the grip is above the tool's own BASE. A long-handled tool is held
    # near the top, so assuming mid-height would drive its head straight through the table.
    # `work_h` states the grip height directly. The bbox-derived guess assumes the tool STANDS
    # UP; a brush lying flat on the table is gripped ~2 cm up, not at 55% of its 19 cm length,
    # and the guess would hold it a hand's width above the work.
    if work_h is not None:
        grip_h = float(work_h)
    else:
        grip_h = float(ext[2])-float(grasp_top) if grasp_top is not None else float(ext[2])*0.55
    # ride on the region's own top surface: wiping a board that sits 1 cm proud of the table at
    # table height means pressing the tool 1 cm into the board
    surface = float(reg.get("top_z", TABLE_TOP))
    work_z = surface+grip_h-float(press)
    print(f"[solver] tool grip {grip_h:.3f} m above its base -> work_z {work_z-TABLE_TOP:.3f} m",
          flush=True)
    rec.phase = "4. LIFT the tool"
    a.flow([res.hold_pose+np.array([0, 0, 0.12], np.float32)], CLOSE)
    path = _work_path(kind, reg, work_z, float(reg.get("half", 0.06))*2)
    rec.phase = "5. MOVE over the work area"
    a.move_to(a.to_root(path[0]+np.array([0, 0, 0.10], np.float32)), CLOSE, tol=0.01, max_steps=130)
    rec.phase = "6. PRESS DOWN"
    a.move_to(a.to_root(path[0]), CLOSE, tol=0.008, max_steps=110)
    rec.phase = f"7. {kind.upper()} along the work path"
    a.flow([a.to_root(p) for p in path[1:]], CLOSE, speed=0.006)
    rec.phase = "8. SET THE TOOL DOWN and release"
    # Ending the episode with the tool still clamped in mid-air leaves every final-pose check
    # reading the gripper's height, not the work: put it down where the job finished.
    a.move_to(a.to_root(path[-1]+np.array([0, 0, 0.004], np.float32)), CLOSE,
              tol=0.010, max_steps=90)
    a.hold(15, OPEN)
    rec.phase = "9. RETREAT"
    a.flow([a._seg_start()+np.array([0, 0, 0.14], np.float32)], OPEN)
    out = {"grasped": True}
    if debris:
        half = float(reg.get("half", 0.06))+0.03
        inzone = sum(1 for d in debris if d in env.scene.objects
                     and abs(env.scene.object_pos(d)[0]-reg["xy"][0]) < half
                     and abs(env.scene.object_pos(d)[1]-reg["xy"][1]) < half)
        out["in_zone"] = inzone
        print(f"[solver] {inzone}/{len(debris)} debris in the zone", flush=True)
    return out