scripts/yam_tool_tasks.py

"""YAM contact-rich / tool-use tasks -- the non-pick-and-place skills.

    --task push   : slide a block to a target region WITHOUT grasping it (closed jaw as a finger)
    --task wipe   : grasp the brush and wipe a target region along a lawn-mower path
    --task sweep  : grasp the brush and sweep loose beads into a target zone
    --task pour   : grasp a cup of beads, carry it over a bowl and ROTATE the wrist to pour

These need a trajectory, not just waypoints: the end-effector follows a parametric path, and
`pour` also drives the wrist orientation (slerp from top-down to tipped) while it moves.

    python scripts/yam_tool_tasks.py --headless --task push --video outputs/tasks/push.mp4
"""
import argparse, sys, os
from isaaclab.app import AppLauncher

parser = argparse.ArgumentParser()
parser.add_argument("--task", default="push", choices=["push", "wipe", "sweep", "pour"])
parser.add_argument("--obj", default="cube_r", help="pushed block (push) / cup (pour)")
parser.add_argument("--tool", default="rw_brush", help="tool object for wipe/sweep")
parser.add_argument("--obj_xy", default="0.02,0.12", help="object start x,y (env-local)")
parser.add_argument("--target_xy", default="0.02,-0.10", help="target region centre x,y")
parser.add_argument("--target_size", type=float, default=0.10)
parser.add_argument("--beads", type=int, default=6, help="beads used by sweep/pour")
parser.add_argument("--pour_deg", type=float, default=105.0, help="wrist tilt angle for pouring")
parser.add_argument("--episode", type=int, default=-1)
parser.add_argument("--video", default="outputs/tasks/yam_tool.mp4")
AppLauncher.add_app_launcher_args(parser)
args = parser.parse_args(); args.headless = True; args.enable_cameras = True
app = AppLauncher(args).app

import numpy as np, torch, gymnasium as gym
import imageio.v2 as imageio
from PIL import Image, ImageDraw
REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(REPO, "source")); sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import yam_prm
import bimanual.tasks.manager_based.yam  # noqa
from isaaclab_tasks.utils import parse_env_cfg

TASK = "Template-YAM-Play-v0"; dev = "cuda:0"
_cfg = parse_env_cfg(TASK, device=dev, num_envs=1)
_cfg.episode_length_s = 1.0e6
try:
    _cfg.terminations.time_out = None
except Exception as _e:
    print("[tt] time_out disable failed:", _e)
try:
    _cfg.viewer.eye = (0.9, -0.9, 1.15); _cfg.viewer.lookat = (0.05, 0.0, 0.5); _cfg.viewer.resolution = (720, 540)
except Exception as _e:
    print("viewer cfg:", _e)
env = gym.make(TASK, cfg=_cfg, render_mode="rgb_array"); u = env.unwrapped; env.reset()


def Rq(q):
    w, x, y, z = 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 qR(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 slerp(q0, q1, t):
    q0 = q0/(np.linalg.norm(q0)+1e-9); q1 = q1/(np.linalg.norm(q1)+1e-9)
    d = float(np.dot(q0, q1))
    if d < 0:
        q1 = -q1; d = -d
    if d > 0.9995:
        q = q0+t*(q1-q0); return q/(np.linalg.norm(q)+1e-9)
    th = np.arccos(d); q2 = q1-q0*d; q2 /= (np.linalg.norm(q2)+1e-9)
    return q0*np.cos(th*t)+q2*np.sin(th*t)


origin = u.scene.env_origins[0].cpu().numpy()
R = u.scene["right_robot"]; Rbn = list(R.data.body_names)
L = u.scene["left_robot"]; Lbn = list(L.data.body_names)
rroot = R.data.root_pos_w[0].cpu().numpy()-origin; rrootq = R.data.root_quat_w[0].cpu().numpy()
lroot = L.data.root_pos_w[0].cpu().numpy()-origin; lrootq = L.data.root_quat_w[0].cpu().numpy()
OFF = np.array([0, 0, 0.13]); TABLE_TOP = 0.45
OPEN, CLOSE = 1.0, -1.0
TXY = [float(v) for v in args.target_xy.split(",")]
TARGET = np.array([TXY[0], TXY[1], TABLE_TOP], np.float32)
OXY = [float(v) for v in args.obj_xy.split(",")]

import isaaclab.sim as sim_utils
# target region marker (paper-thin, so it never blocks the tool)
_m = sim_utils.CuboidCfg(size=(args.target_size, args.target_size, 0.0015),
                         visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.20, 0.65, 0.30)))
_m.func("/World/envs/env_0/target_region", _m,
        translation=tuple((origin+TARGET+np.array([0, 0, 0.0008], np.float32)).astype(float).tolist()))
print(f"[tt] target region at {np.round(origin+TARGET,3)} size={args.target_size}", flush=True)


def eef_root(a, bn, root, rootq):
    i = bn.index("link_6"); p = a.data.body_pos_w[0, i].cpu().numpy()-origin
    q = a.data.body_quat_w[0, i].cpu().numpy()
    return Rq(rootq).T@((p+Rq(q)@OFF)-root), q


lp0, lq0 = eef_root(L, Lbn, lroot, lrootq)
rp_home, rq_home = eef_root(R, Rbn, rroot, rrootq)
# top-down grasp: jaw closes along world y; the tipped pour pose rotates about the tool's x axis
GQ = qR(np.stack([np.array([0., 1., 0.]), np.array([1., 0., 0.]), np.array([0., 0., -1.])], axis=1))


def tipped_quat(deg):
    """Top-down orientation rolled by `deg` about world x -- the pouring motion."""
    a = np.radians(deg); ca, sa = np.cos(a), np.sin(a)
    Rx = np.array([[1, 0, 0], [0, ca, -sa], [0, sa, ca]])
    return qR(Rx@np.stack([np.array([0., 1., 0.]), np.array([1., 0., 0.]), np.array([0., 0., -1.])], axis=1))


def act(rp, rq, rg):
    return torch.tensor(np.concatenate([lp0, lq0, [1.0], rp, rq, [rg]]),
                        dtype=torch.float32, device=dev).view(1, -1)


def place(name, xy, z=0.55, quat=(1, 0, 0, 0)):
    if name not in u.scene.rigid_objects:
        print(f"[tt] MISSING {name}", flush=True); return False
    ro = u.scene.rigid_objects[name]
    ro.write_root_pose_to_sim(torch.tensor(np.concatenate([origin+np.array([xy[0], xy[1], z]), quat]),
                                           dtype=torch.float32, device=dev).view(1, 7))
    ro.write_root_velocity_to_sim(torch.zeros((1, 6), device=dev))
    return True


BEADS = [f"bead{i}" for i in range(args.beads)]
TOOLNAME = args.tool if args.task in ("wipe", "sweep") else None
if args.task == "push":
    place(args.obj, OXY)
elif args.task == "pour":
    place(args.obj, OXY)
elif args.task in ("wipe", "sweep"):
    place(TOOLNAME, OXY)
for _ in range(60):
    env.step(act(rp_home, GQ, OPEN))

import omni.usd
from pxr import UsdGeom, Usd
stage = omni.usd.get_context().get_stage()
bbc = UsdGeom.BBoxCache(Usd.TimeCode.Default(), [UsdGeom.Tokens.default_, UsdGeom.Tokens.render])


def size_of(name):
    try:
        pp = u.scene.rigid_objects[name].root_physx_view.prim_paths[0]
        rng = bbc.ComputeWorldBound(stage.GetPrimAtPath(pp)).ComputeAlignedRange()
        return np.array(rng.GetMax())-np.array(rng.GetMin())
    except Exception:
        return np.array([0.05, 0.05, 0.05])


# beads go INSIDE the cup for pour, or scattered on the table for sweep
if args.task == "pour":
    cs = size_of(args.obj)
    for i, b in enumerate(BEADS):
        ang = 2*np.pi*i/max(len(BEADS), 1)
        place(b, [OXY[0]+0.008*np.cos(ang), OXY[1]+0.008*np.sin(ang)],
              z=TABLE_TOP+float(cs[2])*0.55+0.012*(i//3))
elif args.task == "sweep":
    for i, b in enumerate(BEADS):
        place(b, [OXY[0]-0.06+0.028*(i % 3), OXY[1]-0.10-0.03*(i//3)], z=TABLE_TOP+0.012)
for _ in range(80):
    env.step(act(rp_home, GQ, OPEN))

lhome_q = L.data.joint_pos[0].clone(); _lz = torch.zeros((1, lhome_q.shape[0]), device=dev)


def _freeze_left():
    L.write_joint_state_to_sim(lhome_q.view(1, -1), _lz)


def _boost(view, tag, s=1.6, d=1.4):
    try:
        m = view.get_material_properties().clone(); m[..., 0] = s; m[..., 1] = d
        view.set_material_properties(m, torch.arange(m.shape[0], dtype=torch.int32, device=m.device))
    except Exception as e:
        print(f"[tt] friction failed {tag}:", e, flush=True)


_boost(R.root_physx_view, "right")
for n in ([args.obj] if args.task in ("push", "pour") else []) + ([TOOLNAME] if TOOLNAME else []):
    if n in u.scene.rigid_objects:
        _boost(u.scene.rigid_objects[n].root_physx_view, n)


def objw(name):
    return u.scene.rigid_objects[name].data.root_pos_w[0].cpu().numpy()-origin


def r_eef():
    p, _ = eef_root(R, Rbn, rroot, rrootq); return p


def fsep():
    jn = list(R.data.joint_names)
    return (float(R.data.joint_pos[0, jn.index("left_finger")].item())
            + float(R.data.joint_pos[0, jn.index("right_finger")].item()))/2


frames = []; _phase = {"v": "start"}; _CUR = {"tgt": None, "grip": "OPEN"}; _RESULT = {"v": ""}
TITLE = {"push": "PUSH BLOCK TO TARGET (no grasp)", "wipe": "WIPE TARGET REGION WITH BRUSH",
         "sweep": "SWEEP DEBRIS INTO TARGET ZONE", "pour": "POUR CUP INTO BOWL (wrist rotation)"}[args.task]


def capture():
    img = env.render()
    if img is None:
        return
    im = Image.fromarray(np.asarray(img)[..., :3].copy()); d = ImageDraw.Draw(im)
    lines = [f"=== {TITLE} ===" + (f"  EP {args.episode}" if args.episode >= 0 else "")]
    if _RESULT["v"]:
        lines.append(f"RESULT: {_RESULT['v']}")
    lines += [f"ACTION: {_phase['v']}", f"gripper={_CUR['grip']}"]
    e = r_eef(); t = _CUR["tgt"]
    if t is not None:
        lines.append(f"eef(root) [{e[0]:+.2f} {e[1]:+.2f} {e[2]:+.2f}]  err={np.linalg.norm(e-t):.3f}m")
    d.rectangle([0, 0, 430, 18*len(lines)+6], fill=(0, 0, 0))
    y = 3
    for ln in lines:
        d.text((6, y), ln, fill=(255, 235, 60)); y += 18
    frames.append(np.array(im))


_CORR = {"v": np.zeros(3, np.float32)}; _CMD = {"p": None, "q": GQ.copy()}


def _seg_start():
    return _CMD["p"].copy() if _CMD["p"] is not None else r_eef().astype(np.float32)


def _drive(cp, cq, rg, corr):
    _CMD["p"] = np.asarray(cp, np.float32); _CMD["q"] = np.asarray(cq, np.float32)
    env.step(act((cp+corr).astype(np.float32), np.asarray(cq, np.float32), rg)); _freeze_left()
    e = cp-r_eef(); e = np.where(np.abs(e) > 0.008, e, 0.0)
    corr = np.clip(corr+0.08*e, -0.10, 0.10); corr[2] = max(float(corr[2]), -0.07)
    _CORR["v"] = corr
    return corr


def _ease(a):
    return float(0.5-0.5*np.cos(np.pi*min(max(a, 0.0), 1.0)))


def _point_at(poly, seglens, s):
    acc = 0.0
    for i, Lg in enumerate(seglens):
        if acc+Lg >= s or i == len(seglens)-1:
            t = min(max((s-acc)/max(Lg, 1e-6), 0.0), 1.0)
            return poly[i]+(poly[i+1]-poly[i])*t
        acc += Lg
    return poly[-1]


def flow(pts, rg, speed=0.008, settle=14, quat_fn=None, ease=True):
    """Glide along a polyline; quat_fn(frac) drives the wrist orientation along the way."""
    sp = _seg_start(); poly = [sp]+[np.asarray(p, np.float32) for p in pts]
    seglens = [float(np.linalg.norm(poly[i+1]-poly[i])) for i in range(len(poly)-1)]
    total = float(sum(seglens)); _CUR["grip"] = ("CLOSE" if rg < 0 else "OPEN")
    nsteps = max(int(total/speed*(np.pi/2 if ease else 1.0)), 1); corr = _CORR["v"]
    for k in range(nsteps+settle):
        f = min(1.0, (k+1)/float(nsteps))
        a = _ease(f) if ease else f
        cp = _point_at(poly, seglens, min(total, a*total)); _CUR["tgt"] = cp
        cq = quat_fn(f) if quat_fn is not None else _CMD["q"]
        corr = _drive(cp, cq, rg, corr)
        if k % 2 == 0:
            capture()
    return float(np.linalg.norm(poly[-1]-r_eef()))


def go_converge(target, rg, tol=0.008, max_n=140, quat=None):
    sp = _seg_start(); tp = np.asarray(target, np.float32)
    _CUR["tgt"] = tp; _CUR["grip"] = ("CLOSE" if rg < 0 else "OPEN")
    q = _CMD["q"] if quat is None else np.asarray(quat, np.float32)
    corr = _CORR["v"]; ramp = max(int(max_n*0.6), 18)
    for k in range(max_n):
        a = _ease(min(1.0, (k+1)/float(ramp))); corr = _drive((1-a)*sp+a*tp, q, rg, corr)
        if k % 3 == 0:
            capture()
        if a >= 1.0 and np.linalg.norm(r_eef()-tp) < tol:
            break
    return float(np.linalg.norm(r_eef()-tp))


def hold(n, rg):
    corr = _CORR["v"]
    for k in range(n):
        corr = _drive(_CMD["p"], _CMD["q"], rg, corr)
        if k % 3 == 0:
            capture()


def to_root(w):
    return (Rq(rrootq).T@(np.asarray(w, np.float32)-rroot)).astype(np.float32)


TABLE_ROOT_Z = float(TABLE_TOP-rroot[2])
obstacles = [yam_prm.Box(center=(np.array([-0.38, -0.15, 0.70])-rroot), half=np.array([0.02, 0.45, 0.28])),
             yam_prm.Box(center=(np.array([-0.10, -0.45, 0.70])-rroot), half=np.array([0.45, 0.02, 0.28]))]


def prm_to(goal, rg):
    prm = yam_prm.PRM(bounds_lo=np.array([-0.05, -0.35, -0.03]), bounds_hi=np.array([0.55, 0.45, 0.40]),
                      obstacles=obstacles, clearance=0.05, num_samples=300, k=12, seed=1)
    p = prm.plan(r_eef().astype(np.float64), np.asarray(goal, np.float64))
    if p is None:
        p = np.stack([r_eef(), goal]).astype(np.float32)
    wps = yam_prm.resample_polyline(yam_prm.shortcut(p, obstacles, 0.05, iters=200, seed=2), 12).astype(np.float32)
    flow(list(wps[1:]), rg)


def grasp_object(name, top_offset=None):
    """Standard closed-loop top-down grasp, reused by the tool tasks."""
    ext = size_of(name); w = objw(name)
    # Grasp height comes from the object's LIVE pose, not from the authored bbox: the bbox is the
    # pre-scale, pre-settle one, and for the brush it put the grasp 27 cm above the table, so the
    # jaw closed on air and the whole sweep ran with an empty hand.
    gz = float(w[2]) if top_offset is None else float(w[2])+max(float(ext[2])/2.0-top_offset, 0.0)
    o = to_root(np.array([w[0], w[1], gz], np.float32))
    print(f"[tt] {name} live pose=({w[0]:.3f},{w[1]:.3f},{w[2]:.3f}) -> grasp z={gz:.3f}", flush=True)
    grasp = np.array([o[0], o[1], max(float(o[2]), TABLE_ROOT_Z+0.010)], np.float32)
    _phase["v"] = "APPROACH"; prm_to(grasp+np.array([0, 0, 0.12], np.float32), OPEN)
    _phase["v"] = "DESCEND"; print(f"[tt] descend err={go_converge(grasp, OPEN):.3f}", flush=True)
    _phase["v"] = "CLOSE-GRASP"
    h = r_eef().astype(np.float32); prev = fsep(); stall = 0
    for k in range(160):
        _drive(h, _CMD["q"], CLOSE, _CORR["v"])
        if k % 6 == 0:
            capture()
        cur = fsep()
        stall = stall+1 if abs(cur-prev) < 0.0002 else 0
        prev = cur
        if stall >= 8 and cur < -0.002:
            print(f"[tt] jaws stalled fsep={cur:.4f}", flush=True); break
    # POST-GRASP GATE: lift a little and check the object actually comes along. Without this the
    # script happily performs the entire sweep/pour with an empty hand and only reports failure at
    # the very end, which hides WHERE it went wrong.
    z_before = objw(name)[2]
    corr = _CORR["v"]
    for _ in range(45):
        corr = _drive(h+np.array([0, 0, 0.05], np.float32), _CMD["q"], CLOSE, corr)
    z_after = objw(name)[2]
    if z_after-z_before < 0.02:
        print(f"[tt] GRASP FAILED: {name} did not rise with the gripper "
              f"(z {z_before:.3f} -> {z_after:.3f}, fsep={fsep():.4f}). "
              f"The jaw closed beside/above the object rather than around it.", flush=True)
        _RESULT["v"] = "FAIL (no grasp)"
        return None, ext
    print(f"[tt] grasp confirmed: {name} rose {z_after-z_before:+.3f} m with the gripper", flush=True)
    return h, ext


# =============================== the tasks ================================================
if args.task == "push":
    # No grasp at all: close the empty jaw and use it as a finger to slide the block.
    ext = size_of(args.obj); w0 = objw(args.obj)
    push_z = TABLE_TOP+float(ext[2])*0.45
    d = np.array([TARGET[0]-w0[0], TARGET[1]-w0[1]], np.float32)
    n = d/(np.linalg.norm(d)+1e-9)
    behind = np.array([w0[0]-n[0]*(float(ext[0])/2+0.045), w0[1]-n[1]*(float(ext[1])/2+0.045), push_z], np.float32)
    contact = np.array([w0[0]-n[0]*(float(ext[0])/2+0.012), w0[1]-n[1]*(float(ext[1])/2+0.012), push_z], np.float32)
    end = np.array([TARGET[0]-n[0]*(float(ext[0])/2+0.012), TARGET[1]-n[1]*(float(ext[1])/2+0.012), push_z], np.float32)
    print(f"[tt] push {args.obj} from ({w0[0]:.3f},{w0[1]:.3f}) to target ({TARGET[0]:.3f},{TARGET[1]:.3f})", flush=True)
    _phase["v"] = "CLOSE JAW (finger)"
    for _ in range(40):
        _drive(_seg_start(), GQ, CLOSE, _CORR["v"]); capture()
    _phase["v"] = "APPROACH behind block"; prm_to(to_root(behind+np.array([0, 0, 0.10], np.float32)), CLOSE)
    _phase["v"] = "LOWER to push height"; go_converge(to_root(behind), CLOSE, tol=0.008, max_n=130)
    _phase["v"] = "MAKE CONTACT"; go_converge(to_root(contact), CLOSE, tol=0.006, max_n=90)
    _phase["v"] = "PUSH (closed-loop, re-centred on the block)"
    # Open-loop pushing fails the moment the block skids off the contact normal: the pusher keeps
    # driving down its planned line while the block squirts off sideways. Instead, re-derive the
    # pusher pose from the block's LIVE position every step -- stay on the line from the block to
    # the target, a fixed gap behind the block, and advance the goal gradually.
    steps = 220
    corr = _CORR["v"]
    for k in range(steps):
        w = objw(args.obj)
        to_t = np.array([TARGET[0]-w[0], TARGET[1]-w[1]], np.float32)
        dist = float(np.linalg.norm(to_t))
        if dist < 0.015:
            break
        nn = to_t/(dist+1e-9)
        back = float(ext[0])/2+0.012
        # advance a little past the block each step so contact is maintained, not just touched
        adv = min(0.02, dist)
        cp_w = np.array([w[0]-nn[0]*back+nn[0]*adv, w[1]-nn[1]*back+nn[1]*adv, push_z], np.float32)
        cp = to_root(cp_w); _CUR["tgt"] = cp
        corr = _drive(cp, GQ, CLOSE, corr)
        if k % 2 == 0:
            capture()
    print(f"[tt] push servo done: block=({objw(args.obj)[0]:.3f},{objw(args.obj)[1]:.3f}) "
          f"target=({TARGET[0]:.3f},{TARGET[1]:.3f})", flush=True)
    hold(30, CLOSE)
    _phase["v"] = "RETREAT"; flow([to_root(end+np.array([0, 0, 0.14], np.float32))], CLOSE)
    wf = objw(args.obj)
    err = float(np.hypot(wf[0]-TARGET[0], wf[1]-TARGET[1]))
    _RESULT["v"] = "SUCCESS" if err < args.target_size/2+0.02 else "FAIL"
    print(f"[tt] EPISODE_RESULT: {_RESULT['v']} block=({wf[0]:.3f},{wf[1]:.3f}) target=({TARGET[0]:.3f},{TARGET[1]:.3f}) err={err:.3f}", flush=True)

elif args.task in ("wipe", "sweep"):
    # No top_offset: the brush lies FLAT, so its live centre height is the grasp height. Deriving
    # it from the authored bbox (which reports an 18.8 cm tall brush) aims 6 cm above the handle.
    hold_pose, ext = grasp_object(TOOLNAME)
    if hold_pose is None:
        raise SystemExit("[tt] aborting: the brush was never actually grasped")
    tip_drop = float(ext[2])*0.55            # how far the tool tip reaches below the grip point
    work_z = TABLE_TOP+tip_drop+0.004
    _phase["v"] = "LIFT tool"
    flow([hold_pose+np.array([0, 0, 0.12], np.float32)], CLOSE)
    if args.task == "wipe":
        # lawn-mower path over the target square, keeping the tool on the surface
        half = args.target_size/2
        lanes = []
        for i, yy in enumerate(np.linspace(-half, half, 4)):
            xs = [-half, half] if i % 2 == 0 else [half, -half]
            lanes += [np.array([TARGET[0]+xs[0], TARGET[1]+yy, work_z], np.float32),
                      np.array([TARGET[0]+xs[1], TARGET[1]+yy, work_z], np.float32)]
        _phase["v"] = "MOVE over region"
        go_converge(to_root(lanes[0]+np.array([0, 0, 0.10], np.float32)), CLOSE, tol=0.01, max_n=130)
        _phase["v"] = "PRESS DOWN"; go_converge(to_root(lanes[0]), CLOSE, tol=0.008, max_n=110)
        _phase["v"] = "WIPE (lawn-mower path)"
        flow([to_root(p) for p in lanes[1:]], CLOSE, speed=0.006)
        covered = True
    else:
        # sweep strokes: drag the beads toward the target zone in three passes
        strokes = []
        for i, xoff in enumerate((-0.03, 0.0, 0.03)):
            start = np.array([OXY[0]+xoff, OXY[1]-0.14, work_z], np.float32)
            end = np.array([TARGET[0]+xoff*0.5, TARGET[1]+0.02, work_z], np.float32)
            strokes.append((start, end))
        _phase["v"] = "MOVE behind debris"
        go_converge(to_root(strokes[0][0]+np.array([0, 0, 0.10], np.float32)), CLOSE, tol=0.01, max_n=130)
        for i, (s, e) in enumerate(strokes):
            _phase["v"] = f"SWEEP stroke {i+1}/3"
            go_converge(to_root(s+np.array([0, 0, 0.06], np.float32)), CLOSE, tol=0.01, max_n=80)
            go_converge(to_root(s), CLOSE, tol=0.008, max_n=90)
            flow([to_root(e)], CLOSE, speed=0.005)
            flow([to_root(e+np.array([0, 0, 0.08], np.float32))], CLOSE)
        covered = True
    _phase["v"] = "LIFT tool away"
    flow([_seg_start()+np.array([0, 0, 0.12], np.float32)], CLOSE)
    if args.task == "sweep":
        inzone = sum(1 for b in BEADS if b in u.scene.rigid_objects
                     and abs(objw(b)[0]-TARGET[0]) < args.target_size/2+0.03
                     and abs(objw(b)[1]-TARGET[1]) < args.target_size/2+0.03)
        _RESULT["v"] = "SUCCESS" if inzone >= max(1, len(BEADS)//2) else "FAIL"
        print(f"[tt] EPISODE_RESULT: {_RESULT['v']} beads_in_zone={inzone}/{len(BEADS)}", flush=True)
    else:
        e = np.linalg.norm(r_eef()-to_root(np.array([TARGET[0], TARGET[1], work_z], np.float32)))
        _RESULT["v"] = "SUCCESS" if covered else "FAIL"
        print(f"[tt] EPISODE_RESULT: {_RESULT['v']} wiped the {args.target_size:.2f} m region, "
              f"final tool err={e:.3f}", flush=True)

elif args.task == "pour":
    hold_pose, ext = grasp_object(args.obj)
    z_before = [objw(b)[2] for b in BEADS if b in u.scene.rigid_objects]
    _phase["v"] = "LIFT cup"
    flow([hold_pose+np.array([0, 0, 0.16], np.float32)], CLOSE)
    over = to_root(np.array([TARGET[0], TARGET[1]+0.03, TABLE_TOP+0.20], np.float32))
    _phase["v"] = "CARRY over bowl"
    flow([over], CLOSE)
    _phase["v"] = "POUR (rotate wrist)"
    tip = tipped_quat(args.pour_deg)
    n = 130
    corr = _CORR["v"]
    for k in range(n):                       # rotate in place: position held, orientation slerped
        f = _ease((k+1)/float(n))
        corr = _drive(_CMD["p"], slerp(GQ, tip, f), CLOSE, corr)
        if k % 2 == 0:
            capture()
    _phase["v"] = "HOLD (contents fall out)"
    hold(70, CLOSE)
    _phase["v"] = "RETURN upright"
    corr = _CORR["v"]
    for k in range(90):
        f = _ease((k+1)/90.0)
        corr = _drive(_CMD["p"], slerp(tip, GQ, f), CLOSE, corr)
        if k % 2 == 0:
            capture()
    _phase["v"] = "RETREAT"
    flow([over+np.array([0, 0, 0.06], np.float32)], CLOSE)
    poured = sum(1 for b in BEADS if b in u.scene.rigid_objects
                 and abs(objw(b)[0]-TARGET[0]) < 0.12 and abs(objw(b)[1]-TARGET[1]) < 0.12
                 and objw(b)[2] < TABLE_TOP+0.10)
    _RESULT["v"] = "SUCCESS" if poured >= max(1, len(BEADS)//2) else "FAIL"
    print(f"[tt] EPISODE_RESULT: {_RESULT['v']} poured={poured}/{len(BEADS)} beads into the target area", flush=True)

for _ in range(16):
    capture()
os.makedirs(os.path.dirname(args.video), exist_ok=True)
# Drop the warm-up frames: before the renderer settles they come out with the wrong camera
# pose, unresolved textures and missing geometry.
if len(frames) > 6:
    frames = frames[2:]
if frames:
    imageio.mimsave(args.video, frames, fps=14)
print(f"[tt] video -> {args.video} ({len(frames)} frames)", flush=True)
env.close(); app.close(); print("YAM_TOOL_OK", flush=True)