scripts/yam_multi_pick.py

"""YAM multi-object sorting: pick a chosen SUBSET of the objects on the table into one basket
and leave the rest untouched.

The default is the fruit-basket task: two fruits (grape, apple) and one non-fruit distractor
(a can) sit on the table; only the fruits belong in the basket, so a run that also dumps the can
in -- or knocks it over -- fails.

Motion is the same pipeline as ``yam_grasp_prm.py``: task-space PRM approach, closed-loop
vertical descend onto the object, close-until-contact friction grasp, then ONE filleted
lift->carry->lower arc per object (no stop-and-turn at the top of the lift).

    python scripts/yam_multi_pick.py --headless \
        --pick grape,apple --leave can --basket_xy=0.02,-0.24 \
        --video outputs/tasks/task7_fruit_basket.mp4
"""
import argparse, sys, os
from isaaclab.app import AppLauncher

parser = argparse.ArgumentParser()
parser.add_argument("--pick", default="grape,apple", help="objects that must end up in the basket")
parser.add_argument("--leave", default="can", help="distractors that must STAY on the table (may be empty)")
parser.add_argument("--pick_xy", default="-0.03,0.10;0.02,0.10",
                    help="';'-separated x,y for each --pick object (env-local)")
parser.add_argument("--leave_xy", default="0.14,0.02", help="';'-separated x,y for each --leave object")
parser.add_argument("--basket_xy", default="0.02,-0.24", help="basket centre x,y (env-local)")
parser.add_argument("--container", default="",
                    help="RoboTwin container USD subpath under $ROBOTWIN_USD (e.g. 110_basket/base0.usd). "
                         "Empty = build a primitive box instead.")
parser.add_argument("--container_scale", type=float, default=1.0)
parser.add_argument("--container_rpy", default="0,0,0",
                    help="container roll,pitch,yaw in DEGREES. RoboTwin GLBs are not all authored "
                         "z-up, so a basket/bin often needs a 90 deg roll to stand opening-up.")
parser.add_argument("--on_top", action="store_true",
                    help="place the objects ON TOP of the container (a scale / plate / coaster) "
                         "instead of dropping them inside it")
parser.add_argument("--label", default="FRUIT BASKET", help="title shown in the video overlay")
parser.add_argument("--episode", type=int, default=-1)
parser.add_argument("--video", default="outputs/tasks/multi_pick.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, json as _json
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)
# One long scripted sequence: the 12 s task episode length would auto-reset the env mid-run
# and snap the arm back to its home joints.
_cfg.episode_length_s = 1.0e6
try:
    _cfg.terminations.time_out = None
except Exception as _e:
    print("[cfg] 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


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
PICK = [s for s in args.pick.split(",") if s]
LEAVE = [s for s in args.leave.split(",") if s]
BXY = [float(v) for v in args.basket_xy.split(",")]
BASKET = np.array([BXY[0], BXY[1], TABLE_TOP], np.float32)

# ---- basket: primitive cuboids (license-clean), wide + shallow so objects settle inside ----
import isaaclab.sim as sim_utils
S, H, T = 0.26, 0.05, 0.010


def _cub(name, size, off, color=(0.55, 0.38, 0.22)):
    c = sim_utils.CuboidCfg(size=tuple(size),
                            visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=color),
                            collision_props=sim_utils.CollisionPropertiesCfg())
    c.func(f"/World/envs/env_0/basket_{name}", c,
           translation=tuple((origin+BASKET+np.array(off, np.float32)).astype(float).tolist()))


CONTAINER_TOP = {"z": TABLE_TOP}          # rim height, filled in for a real container asset
if args.container:
    # ---- real RoboTwin basket, kinematic so it cannot be shoved around, standing ON the table ----
    ROBOTWIN_USD = os.environ.get("ROBOTWIN_USD", "/home/horde/xiaotong/robotwin_assets/usd")
    ccfg = sim_utils.UsdFileCfg(usd_path=f"{ROBOTWIN_USD}/{args.container}",
                                rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True),
                                scale=(args.container_scale,)*3)
    _cr, _cp, _cy = [np.radians(float(v)) for v in args.container_rpy.split(",")]
    _sr, _cr_ = np.sin(_cr/2), np.cos(_cr/2)
    _sp, _cp_ = np.sin(_cp/2), np.cos(_cp/2)
    _sy, _cy_ = np.sin(_cy/2), np.cos(_cy/2)
    _cq = (float(_cr_*_cp_*_cy_+_sr*_sp*_sy), float(_sr*_cp_*_cy_-_cr_*_sp*_sy),
           float(_cr_*_sp*_cy_+_sr*_cp_*_sy), float(_cr_*_cp_*_sy-_sr*_sp*_cy_))
    ccfg.func("/World/envs/env_0/rw_basket", ccfg,
              translation=tuple((origin+BASKET).astype(float).tolist()), orientation=_cq)
    import omni.usd as _ou
    from pxr import UsdGeom as _UG, Usd as _U, Gf as _Gf
    _stage = _ou.get_context().get_stage()
    _prim = _stage.GetPrimAtPath("/World/envs/env_0/rw_basket")
    _bbc = _UG.BBoxCache(_U.TimeCode.Default(), [_UG.Tokens.default_, _UG.Tokens.render])
    _rng = _bbc.ComputeWorldBound(_prim).ComputeAlignedRange()
    # the USD origin is wherever the mesh was authored, so lift it until its base meets the table
    _dz = float(origin[2]+TABLE_TOP)-float(_rng.GetMin()[2])
    for _op in _UG.Xformable(_prim).GetOrderedXformOps():
        if _op.GetOpType() == _UG.XformOp.TypeTranslate:
            _t = _op.Get(); _op.Set(_Gf.Vec3d(float(_t[0]), float(_t[1]), float(_t[2])+_dz)); break
    _rng = _bbc.ComputeWorldBound(_prim).ComputeAlignedRange()
    _ext = np.array(_rng.GetMax())-np.array(_rng.GetMin())
    CONTAINER_TOP["z"] = float(_rng.GetMax()[2])-origin[2]
    S = float(min(_ext[0], _ext[1]))          # inner span used by the success check
    print(f"[mp] RoboTwin basket {args.container} raised {_dz:+.3f} m to stand on the table; "
          f"extents={np.round(_ext,3)} rim_z={CONTAINER_TOP['z']:.3f}", flush=True)
else:
    _cub("floor", (S, S, T), (0, 0, T/2))
    _cub("xp", (T, S, H), (S/2, 0, H/2)); _cub("xn", (T, S, H), (-S/2, 0, H/2))
    _cub("yp", (S, T, H), (0, S/2, H/2)); _cub("yn", (S, T, H), (0, -S/2, H/2))
    print(f"[mp] primitive basket at world={np.round(origin+BASKET,3)} span={S}", 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)
OPEN, CLOSE = 1.0, -1.0


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


# top-down grasp orientation (jaw closes along world y)
gq = qR(np.stack([np.array([0., 1., 0.]), np.array([1., 0., 0.]), np.array([0., 0., -1.])], axis=1))


def place_objects():
    """Put every task object at its scripted spot; everything else stays parked off-scene."""
    specs = list(zip(PICK, args.pick_xy.split(";"))) + list(zip(LEAVE, args.leave_xy.split(";")))
    for name, xy in specs:
        if name not in u.scene.rigid_objects:
            print(f"[mp] MISSING object {name} in scene", flush=True); continue
        x, y = [float(v) for v in xy.split(",")]
        ro = u.scene.rigid_objects[name]
        pw = torch.tensor(np.concatenate([origin+np.array([x, y, 0.55]), [1, 0, 0, 0]]),
                          dtype=torch.float32, device=dev).view(1, 7)
        ro.write_root_pose_to_sim(pw); ro.write_root_velocity_to_sim(torch.zeros((1, 6), device=dev))
        print(f"[mp] placed {name} at ({x},{y})", flush=True)


place_objects()
for _ in range(60):
    env.step(act(rp_home, 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"[mp] friction set failed on {tag}:", e, flush=True)


_boost(R.root_physx_view, "right_robot")
for n in PICK+LEAVE:
    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


# ---- object size probe (pose-independent), used for grasp height and placement height ----
import omni.usd
from pxr import UsdGeom, Usd
stage = omni.usd.get_context().get_stage()
bbcache = UsdGeom.BBoxCache(Usd.TimeCode.Default(), [UsdGeom.Tokens.default_, UsdGeom.Tokens.render])


def obj_size(name):
    try:
        pp = u.scene.rigid_objects[name].root_physx_view.prim_paths[0]
        rng = bbcache.ComputeWorldBound(stage.GetPrimAtPath(pp)).ComputeAlignedRange()
        return np.array(rng.GetMax())-np.array(rng.GetMin())
    except Exception as e:
        print(f"[mp] size probe failed for {name}:", e, flush=True); return None


# ---- overlay ----
frames = []; _phase = {"v": "start"}; _CUR = {"tgt": None, "grip": "OPEN", "obj": ""}; _RESULT = {"v": ""}
SEM = {"approach": "1. PRM APPROACH", "descend": "2. DESCEND onto object", "close": "3. CLOSE-GRASP",
       "carry": "4. LIFT + CARRY to basket", "release": "5. RELEASE into basket",
       "retreat": "6. RETREAT", "done": "DONE"}


def capture():
    img = env.render()
    if img is None:
        return
    im = Image.fromarray(np.asarray(img)[..., :3].copy()); d = ImageDraw.Draw(im)
    lines = [f"=== {args.label}" + (f"  EPISODE {args.episode}" if args.episode >= 0 else "") + " ==="]
    if _RESULT["v"]:
        lines.append(f"RESULT: {_RESULT['v']}")
    lines += [f"ACTION: {SEM.get(_phase['v'], _phase['v'])}",
              f"target={_CUR['obj']}  gripper={_CUR['grip']}",
              f"pick={','.join(PICK)}   leave={','.join(LEAVE)}"]
    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, 400, 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))


# ---- smooth executor (shared design with yam_grasp_prm.py): the integral correction and the
#      last COMMANDED point persist across segments, so no segment boundary snaps the pose. ----
_CORR = {"v": np.zeros(3, np.float32)}; _CMD = {"p": None}


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


def _drive(cp, rg, corr):
    _CMD["p"] = np.asarray(cp, np.float32)
    env.step(act((cp+corr).astype(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.06)
    _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 fillet(pts, r=0.06, n=6):
    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 = p1-d0/l0*rr; b = 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, 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=16, on_step=None):
    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)), 1); corr = _CORR["v"]
    for k in range(nsteps+settle):
        a = _ease(min(1.0, (k+1)/float(nsteps)))
        cp = _point_at(poly, seglens, min(total, a*total)); _CUR["tgt"] = cp
        if on_step is not None:
            on_step(min(1.0, (k+1)/float(nsteps)))
        corr = _drive(cp, 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):
    sp = _seg_start(); tp = np.asarray(target, np.float32)
    _CUR["tgt"] = tp; _CUR["grip"] = ("CLOSE" if rg < 0 else "OPEN")
    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))); cp = (1-a)*sp+a*tp
        corr = _drive(cp, 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 go(target, rg, n):
    sp = _seg_start(); tp = np.asarray(target, np.float32)
    _CUR["tgt"] = tp; _CUR["grip"] = ("CLOSE" if rg < 0 else "OPEN"); corr = _CORR["v"]
    for k in range(n):
        a = _ease((k+1)/float(n)); corr = _drive((1-a)*sp+a*tp, rg, corr)
        if k % 3 == 0:
            capture()


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


def pick_into_basket(name, idx):
    ext = obj_size(name)
    w = objw(name)
    if ext is not None:
        w = np.array([w[0], w[1], TABLE_TOP+float(ext[2])/2.0])
    o = Rq(rrootq).T@(w-rroot)
    grasp = np.array([o[0], o[1], max(float(o[2])-0.005, TABLE_ROOT_Z+0.010)], np.float32)
    pre = np.array([o[0], o[1], float(o[2])+0.12], np.float32)
    _CUR["obj"] = name
    print(f"[mp] --- {idx}. {name}: obj_root={np.round(o,3)} grasp={np.round(grasp,3)} ---", flush=True)

    _phase["v"] = "approach"
    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), pre.astype(np.float64))
    if p is None:
        p = np.stack([r_eef(), pre]).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:]), OPEN)

    _phase["v"] = "descend"
    d_err = go_converge(grasp, OPEN, tol=0.008, max_n=140)
    for _ in range(10):
        _drive(grasp, OPEN, _CORR["v"])
    print(f"[mp]   descend err={d_err:.3f}", flush=True)

    _phase["v"] = "close"
    hold = r_eef().astype(np.float32)
    prev = fsep(); stall = 0
    for k in range(160):
        _drive(hold, 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"[mp]   jaws stalled at fsep={cur:.4f}", flush=True); break

    _phase["v"] = "carry"
    z0 = objw(name)[2]
    # The object is rarely centred between the fingers -- a long grape cluster grabbed near one
    # end hangs several cm off the gripper axis. Aim the OBJECT at the basket by subtracting the
    # measured object-vs-eef offset from the placement target.
    _eef_w = rroot+Rq(rrootq)@r_eef()
    _grip_off = objw(name)[:2]-_eef_w[:2]
    print(f"[mp]   object offset in jaws = {np.round(_grip_off,3)} m", flush=True)
    half_h = (float(ext[2])/2.0 if ext is not None else 0.05)
    # release just above the rim for a real basket (it has walls to clear), or just above the
    # floor for the primitive tray
    _rim = CONTAINER_TOP["z"]-TABLE_TOP
    if args.on_top:
        # a scale/plate is a platform, not a bin: set the object down ON its top surface, and
        # spread multiple objects so the second one does not land on the first
        _spread = np.array([0.035*(idx-1), 0.0, 0.0], np.float32)
        drop_local = BASKET+_spread+np.array([0, 0, _rim+half_h+0.004], np.float32)
    else:
        # Spread the drop points along x: dropping every object on the same spot means the second
        # one lands on the first and a round fruit rolls straight back out of a shallow tray.
        _dx = (idx-(len(PICK)+1)/2.0)*0.075
        drop_local = BASKET+np.array([_dx, 0, max(0.010, _rim)+half_h+0.02], np.float32)
    drop_local = drop_local-np.array([_grip_off[0], _grip_off[1], 0.0], np.float32)
    place = (Rq(rrootq).T@(drop_local-rroot)).astype(np.float32)
    above = place+np.array([0, 0, 0.10], np.float32)
    lift = hold+np.array([0, 0, 0.16], np.float32)
    zmax = {"v": z0}

    def _step(frac):
        zmax["v"] = max(zmax["v"], objw(name)[2])
        _phase["v"] = "carry" if frac < 0.80 else "release"

    flow(fillet([_seg_start(), lift, above, place], r=0.06)[1:], CLOSE, on_step=_step)
    print(f"[mp]   carried: peak z={zmax['v']:.3f} (start {z0:.3f}) place_err={np.linalg.norm(r_eef()-place):.3f}", flush=True)

    _phase["v"] = "release"; go(place, OPEN, 40)
    _phase["v"] = "retreat"; flow([above], OPEN)
    return zmax["v"]-z0 > 0.05


lifted = {}
for i, name in enumerate(PICK, 1):
    if name in u.scene.rigid_objects:
        lifted[name] = pick_into_basket(name, i)

_phase["v"] = "retreat"; flow([rp_home], OPEN)

# ---- success: every --pick object inside the basket, every --leave object still outside ----
in_basket = {}
for name in PICK+LEAVE:
    if name not in u.scene.rigid_objects:
        continue
    w = objw(name)
    if args.on_top:
        # resting ON the platform: centred on it and sitting at (or just above) its top surface
        in_basket[name] = (abs(w[0]-BASKET[0]) < S/2+0.03 and abs(w[1]-BASKET[1]) < S/2+0.03
                           and w[2] > CONTAINER_TOP["z"]-0.01)
    else:
        in_basket[name] = (abs(w[0]-BASKET[0]) < S/2+0.02 and abs(w[1]-BASKET[1]) < S/2+0.02
                           and w[2] < BASKET[2]+0.16)
    print(f"[mp] {name}: final=({w[0]:.3f},{w[1]:.3f},{w[2]:.3f}) in_basket={in_basket[name]}", flush=True)
ok_pick = all(in_basket.get(n, False) for n in PICK)
ok_leave = all(not in_basket.get(n, False) for n in LEAVE)
_RESULT["v"] = "SUCCESS" if (ok_pick and ok_leave) else "FAIL"
_phase["v"] = "done"
for _ in range(14):
    capture()
print(f"[mp] EPISODE_RESULT: {_RESULT['v']} picked={sum(in_basket.get(n,False) for n in PICK)}/{len(PICK)} "
      f"distractors_left_out={sum(not in_basket.get(n,False) for n in LEAVE)}/{len(LEAVE)}", flush=True)

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"[mp] video -> {args.video} ({len(frames)} frames)", flush=True)
env.close(); app.close(); print("YAM_MULTI_PICK_OK", flush=True)