scripts/yam_stack_target.py

"""YAM stack task: two green blocks and a flat TARGET REGION (a paper-thin marker on the table).
Both blocks must end up stacked on the target -- first block placed on the marker, second block
placed on top of the first.

Stacking is the precision case: the second block has to land within ~1.5 cm of the first one's
centre or it topples off, so the placement uses the same closed-loop descend as the pick rather
than an open-loop drop.

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

parser = argparse.ArgumentParser()
parser.add_argument("--blocks", default="cube_g,cube_g2", help="the blocks to stack, bottom first")
parser.add_argument("--block_xy", default="-0.03,0.11;0.04,0.10", help="';'-separated start x,y per block")
parser.add_argument("--target_xy", default="0.00,-0.20", help="target region centre x,y (env-local)")
parser.add_argument("--target_size", type=float, default=0.11, help="target region side length (m)")
parser.add_argument("--block_rpy", default="0,0,0",
                    help="roll,pitch,yaw in DEGREES applied to every block when seating it. RoboTwin "
                         "cups spawn upside-down, and an inverted cup has its rim at the bottom, so "
                         "the jaws grip a narrowing taper and it slips out.")
parser.add_argument("--episode", type=int, default=-1)
parser.add_argument("--video", default="outputs/tasks/yam_stack_target.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("[st] 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
OPEN, CLOSE = 1.0, -1.0
BLOCKS = [b for b in args.blocks.split(",") if b]
TXY = [float(v) for v in args.target_xy.split(",")]
TARGET = np.array([TXY[0], TXY[1], TABLE_TOP], np.float32)

# ---- the TARGET REGION: a paper-thin plate on the table (visual marker, collidable but flat) ----
import isaaclab.sim as sim_utils
_mark = sim_utils.CuboidCfg(size=(args.target_size, args.target_size, 0.002),
                            visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.95, 0.95, 0.92)),
                            collision_props=sim_utils.CollisionPropertiesCfg())
_mark.func("/World/envs/env_0/target_region", _mark,
           translation=tuple((origin+TARGET+np.array([0, 0, 0.001], np.float32)).astype(float).tolist()))
_edge = sim_utils.CuboidCfg(size=(args.target_size*1.06, args.target_size*1.06, 0.001),
                            visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.15, 0.55, 0.25)))
_edge.func("/World/envs/env_0/target_edge", _edge,
           translation=tuple((origin+TARGET).astype(float).tolist()))
print(f"[st] target region at world={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)
gq = qR(np.stack([np.array([0., 1., 0.]), np.array([1., 0., 0.]), np.array([0., 0., -1.])], axis=1))


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


missing = [b for b in BLOCKS if b not in u.scene.rigid_objects]
if missing:
    raise SystemExit(f"[st] blocks not in scene: {missing}")
START_XY = {}
for name, xy in zip(BLOCKS, args.block_xy.split(";")):
    x, y = [float(v) for v in xy.split(",")]
    START_XY[name] = (x, y)
    ro = u.scene.rigid_objects[name]
    ro.write_root_pose_to_sim(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_velocity_to_sim(torch.zeros((1, 6), device=dev))
    print(f"[st] placed {name} at ({x},{y})", flush=True)
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"[st] friction set failed on {tag}:", e, flush=True)


_boost(R.root_physx_view, "right_robot")
for b in BLOCKS:
    _boost(u.scene.rigid_objects[b].root_physx_view, b)


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


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 obj_size(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 as e:
        print(f"[st] size probe failed for {name}:", e, flush=True); return np.array([0.05, 0.05, 0.05])


# Seat every object ON the table before starting: they are written in 10 cm above it, and a
# cup-shaped object that tall lands on its side, which no top-down grasp can recover.
_br, _bp, _by = [np.radians(float(v)) for v in args.block_rpy.split(",")]
_BQ = np.array([
    np.cos(_br/2)*np.cos(_bp/2)*np.cos(_by/2)+np.sin(_br/2)*np.sin(_bp/2)*np.sin(_by/2),
    np.sin(_br/2)*np.cos(_bp/2)*np.cos(_by/2)-np.cos(_br/2)*np.sin(_bp/2)*np.sin(_by/2),
    np.cos(_br/2)*np.sin(_bp/2)*np.cos(_by/2)+np.sin(_br/2)*np.cos(_bp/2)*np.sin(_by/2),
    np.cos(_br/2)*np.cos(_bp/2)*np.sin(_by/2)-np.sin(_br/2)*np.sin(_bp/2)*np.cos(_by/2)])
for _n in BLOCKS:
    _e = obj_size(_n)
    _x, _y = START_XY[_n]
    _ro = u.scene.rigid_objects[_n]
    _ro.write_root_pose_to_sim(torch.tensor(
        np.concatenate([origin+np.array([_x, _y, TABLE_TOP+float(_e[2])/2.0+0.004]), _BQ]),
        dtype=torch.float32, device=dev).view(1, 7))
    _ro.write_root_velocity_to_sim(torch.zeros((1, 6), device=dev))
    print(f"[st] seated {_n} upright: size={np.round(_e,3)}", flush=True)
for _ in range(80):
    env.step(act(rp_home, OPEN)); _freeze_left()


frames = []; _phase = {"v": "approach"}; _CUR = {"tgt": None, "grip": "OPEN", "obj": ""}; _RESULT = {"v": ""}
SEM = {"approach": "1. PRM APPROACH", "descend": "2. DESCEND onto block", "close": "3. CLOSE-GRASP",
       "carry": "4. LIFT + CARRY to target", "place": "5. PLACE (closed-loop)",
       "release": "6. RELEASE", "retreat": "7. 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 = ["=== STACK BLOCKS ON TARGET REGION ===" + (f"  EP {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"block={_CUR['obj']}  gripper={_CUR['grip']}",
              f"target=({TARGET[0]:+.2f},{TARGET[1]:+.2f})  size={args.target_size:.2f}m"]
    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}


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.006, max_n=150):
    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))); corr = _drive((1-a)*sp+a*tp, 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])
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 pick_and_stack(name, layer):
    """layer 0 = on the target marker, layer 1 = on top of the block below."""
    ext = obj_size(name)
    w = objw(name)
    o = Rq(rrootq).T@(np.array([w[0], w[1], TABLE_TOP+float(ext[2])/2.0], np.float32)-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"[st] --- block {layer}: {name} size={np.round(ext,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"
    print(f"[st]   descend err={go_converge(grasp, OPEN, tol=0.008, max_n=140):.3f}", flush=True)
    for _ in range(10):
        _drive(grasp, OPEN, _CORR["v"])

    _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"[st]   jaws stalled at fsep={cur:.4f}", flush=True); break

    # Where this block must land: on the marker, or centred on top of the block below.
    stack_z = TABLE_TOP+0.002+float(ext[2])*(layer+0.5)
    place_local = np.array([TARGET[0], TARGET[1], stack_z], np.float32)
    place = (Rq(rrootq).T@(place_local-rroot)).astype(np.float32)
    above = place+np.array([0, 0, 0.12], np.float32)
    lift = hold+np.array([0, 0, 0.16], np.float32)
    z0 = objw(name)[2]; zmax = {"v": z0}

    def _step(frac):
        zmax["v"] = max(zmax["v"], objw(name)[2])
        _phase["v"] = "carry"

    _phase["v"] = "carry"
    flow(fillet([_seg_start(), lift, above], r=0.06)[1:], CLOSE, on_step=_step)
    # closed-loop descent onto the stack: an open-loop drop misses by more than a block width
    _phase["v"] = "place"
    perr = go_converge(place, CLOSE, tol=0.005, max_n=170)
    print(f"[st]   placed: target={np.round(place,3)} err={perr:.3f} peak_z={zmax['v']:.3f}", flush=True)
    _phase["v"] = "release"; go(place, OPEN, 40)
    _phase["v"] = "retreat"; flow([above], OPEN)
    return zmax["v"]-z0 > 0.04


for i, b in enumerate(BLOCKS):
    pick_and_stack(b, i)
_phase["v"] = "retreat"; flow([rp_home], OPEN)
for _ in range(60):
    env.step(act(rp_home, OPEN)); _freeze_left()

# ---- success: both blocks on the target footprint, and the second one ON TOP of the first ----
half = args.target_size/2.0+0.02
pos = {b: objw(b) for b in BLOCKS}
on_target = {b: (abs(pos[b][0]-TARGET[0]) < half and abs(pos[b][1]-TARGET[1]) < half) for b in BLOCKS}
zs = sorted((pos[b][2], b) for b in BLOCKS)
stacked = len(BLOCKS) < 2 or (zs[-1][0]-zs[0][0]) > 0.03
for b in BLOCKS:
    print(f"[st] {b}: final=({pos[b][0]:.3f},{pos[b][1]:.3f},{pos[b][2]:.3f}) on_target={on_target[b]}", flush=True)
_RESULT["v"] = "SUCCESS" if (all(on_target.values()) and stacked) else "FAIL"
_phase["v"] = "done"
for _ in range(16):
    capture()
print(f"[st] EPISODE_RESULT: {_RESULT['v']} on_target={sum(on_target.values())}/{len(BLOCKS)} "
      f"stacked={stacked} dz={zs[-1][0]-zs[0][0]:.3f}", 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"[st] video -> {args.video} ({len(frames)} frames)", flush=True)
env.close(); app.close(); print("YAM_STACK_OK", flush=True)