scripts/yam_grasp_prm.py

"""YAM bimanual: right arm grasps the apple, but the HOME->pre-grasp motion is planned
by a task-space PRM (probabilistic roadmap) that routes the end-effector around a box
obstacle placed on the table. The planned polyline is executed via absolute diff-IK.
Renders a real Isaac Sim video."""
import argparse, sys, os
from isaaclab.app import AppLauncher
parser = argparse.ArgumentParser()
parser.add_argument("--obj", default="apple")
parser.add_argument("--cart", action="store_true", help="place the object INTO the cart (else drop on table)")
parser.add_argument("--basket", action="store_true", help="build a primitive box on the table and place the object into it")
parser.add_argument("--obj_xy", default="", help="runtime override of the object's x,y (env-local)")
parser.add_argument("--box_xy", default="0.06,-0.26", help="box x,y (env-local)")
parser.add_argument("--episode", type=int, default=-1, help="episode index (overlay label)")
parser.add_argument("--container", default="", help="RobotWin container USD subpath (e.g. 002_bowl/base1.usd) placed at box_xy")
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 Y-up, so a rack "
                         "or bin usually needs a 90 deg roll to stand upright")
parser.add_argument("--insert", action="store_true",
                    help="ManiSkill-style peg-in-hole: build a socket with a square hole at box_xy "
                         "and insert the object into it instead of dropping it in a box")
parser.add_argument("--hole", type=float, default=0.038, help="socket hole width (m)")
parser.add_argument("--jaw", default="auto", choices=["auto", "x", "y"],
                    help="world axis the jaw closes along; 'auto' picks the object's narrower "
                         "horizontal extent (needed for objects lying on their side)")
parser.add_argument("--grasp_top", type=float, default=-1.0,
                    help="grasp this far below the object's TOP instead of at its mid-height "
                         "(use for tall objects so the shaft below the fingers can enter a hole)")
parser.add_argument("--video", default="outputs/yam_grasp_prm.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
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__)))   # for yam_prm
import yam_prm  # task-space PRM planner (Box, PRM, shortcut, resample_polyline)
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)
# These demos script one long manipulation sequence; the 12 s task episode length would
# auto-reset the env mid-run and snap the arm back to its home joints (a visible pose jump).
_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
obs,_=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)
rroot=R.data.root_pos_w[0].cpu().numpy()-origin; rrootq=R.data.root_quat_w[0].cpu().numpy()
L=u.scene["left_robot"]; Lbn=list(L.data.body_names)
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])
# ---- Walls behind the robot, AUTO-loaded as PRM blockers ----
# The arm faces +x,+y (toward the apple); everything behind the base is walled off.
# Each wall is spawned as visible+collidable geometry AND registered into the PRM
# collision set (BLOCKERS). Coords: env-local center 'c', half-extents 'h'.
# No walls when placing INTO the cart (the cart is BEHIND the robot, so back-walls would block it).
WALLS=[] if args.cart else [
    {"name":"wall_back", "c":np.array([-0.38,-0.15,0.70],np.float32), "h":np.array([0.02,0.45,0.28],np.float32)},  # -x back wall
    {"name":"wall_side", "c":np.array([-0.10,-0.45,0.70],np.float32), "h":np.array([0.45,0.02,0.28],np.float32)},  # -y side wall
]
# Walls are PLANNING-ONLY obstacles: added to the PRM collision set but NOT spawned/rendered
# into the scene (no visible or physical wall). Flip SPAWN_WALLS=True to also render them.
SPAWN_WALLS=False
BLOCKERS=[]; _wall_json=[]
for w in WALLS:
    BLOCKERS.append(((w["c"]-rroot).astype(np.float64), w["h"].astype(np.float64)))   # root-frame for PRM
    _wall_json.append({"center_world":(origin+w["c"]).astype(float).tolist(),"size":(w["h"]*2).astype(float).tolist()})
    if SPAWN_WALLS:
        try:
            import isaaclab.sim as sim_utils
            cub=sim_utils.CuboidCfg(size=tuple((w["h"]*2).astype(float).tolist()),
                visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.72,0.72,0.80)),
                collision_props=sim_utils.CollisionPropertiesCfg())
            cub.func(f"/World/envs/env_0/{w['name']}", cub, translation=tuple((origin+w["c"]).astype(float).tolist()))
        except Exception as _e:
            print("[g] wall spawn failed:", _e, flush=True)
    print(f"[g] PRM blocker {w['name']} (planning_only={not SPAWN_WALLS}) world={np.round(origin+w['c'],3)} size={np.round(w['h']*2,3)}", flush=True)
# ---- optional PRIMITIVE box container (built from Cuboids -> no third-party asset / license-clean) ----
_bxy=[float(v) for v in args.box_xy.split(",")]
BASKET_W=np.array([_bxy[0],_bxy[1],0.45],np.float32)   # env-local, on the table
if args.container:
    # ---- REAL RobotWin container (bowl / tray / plate) placed at box_xy, kinematic (static) ----
    try:
        import isaaclab.sim as sim_utils
        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,args.container_scale,args.container_scale))
        _r,_p,_y=[np.radians(float(v)) for v in args.container_rpy.split(",")]
        _cq=(float(np.cos(_r/2)*np.cos(_p/2)*np.cos(_y/2)+np.sin(_r/2)*np.sin(_p/2)*np.sin(_y/2)),
             float(np.sin(_r/2)*np.cos(_p/2)*np.cos(_y/2)-np.cos(_r/2)*np.sin(_p/2)*np.sin(_y/2)),
             float(np.cos(_r/2)*np.sin(_p/2)*np.cos(_y/2)+np.sin(_r/2)*np.cos(_p/2)*np.sin(_y/2)),
             float(np.cos(_r/2)*np.cos(_p/2)*np.sin(_y/2)-np.sin(_r/2)*np.sin(_p/2)*np.cos(_y/2)))
        ccfg.func("/World/envs/env_0/rw_container", ccfg,
                  translation=tuple((origin+BASKET_W).astype(float).tolist()), orientation=_cq)
        print(f"[g] spawned RobotWin container {args.container} scale={args.container_scale} at world={np.round(origin+BASKET_W,3)}", flush=True)
        # Stand it ON the table: the USD origin is wherever the mesh was authored (often the
        # centre), so spawning at table height buries the lower half. Shift up by the gap
        # between the measured bbox bottom and the table surface.
        try:
            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_container")
            _bbc=_UG.BBoxCache(_U.TimeCode.Default(),[_UG.Tokens.default_,_UG.Tokens.render])
            _rng=_bbc.ComputeWorldBound(_prim).ComputeAlignedRange()
            _dz=float(origin[2]+0.45)-float(_rng.GetMin()[2])       # 0.45 = table top, env-local
            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
            print(f"[g] container stood on table: raised by {_dz:+.3f} m "
                  f"(bbox z was [{float(_rng.GetMin()[2]):.3f},{float(_rng.GetMax()[2]):.3f}])", flush=True)
        except Exception as _e:
            print("[g] container stand-on-table correction failed:", _e, flush=True)
    except Exception as _e:
        print("[g] container spawn failed:", _e, flush=True)
if args.insert:
    # ---- SOCKET with a square hole (primitives, license-clean), ManiSkill peg-insertion style.
    # Four walls leave a `--hole` wide square gap whose floor is the table, so a peg dropped in
    # stands with its base on the table; landing on the walls instead is a clear miss. ----
    try:
        import isaaclab.sim as sim_utils
        HW=float(args.hole)/2.0; WT=0.045; WH=0.05      # half-hole, wall thickness, wall height
        SPAN=args.hole+2*WT
        def _wall(name,size,off,color=(0.35,0.38,0.45)):
            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/socket_{name}", c,
                   translation=tuple((origin+BASKET_W+np.array(off,np.float32)).astype(float).tolist()))
        _wall("xp",(WT,SPAN,WH),( HW+WT/2,0,WH/2)); _wall("xn",(WT,SPAN,WH),(-HW-WT/2,0,WH/2))
        _wall("yp",(SPAN,WT,WH),(0, HW+WT/2,WH/2)); _wall("yn",(SPAN,WT,WH),(0,-HW-WT/2,WH/2))
        print(f"[g] built socket at world={np.round(origin+BASKET_W,3)} hole={args.hole} wall_h={WH}", flush=True)
    except Exception as _e:
        print("[g] socket build failed:", _e, flush=True)
if args.basket and not args.container and not args.insert:
    try:
        import isaaclab.sim as sim_utils
        S,H,T=0.26,0.05,0.010   # wide + shallow tray: tall objects rest in it without clipping/ejection
        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_W+np.array(off,np.float32)).astype(float).tolist()))
        _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"[g] built primitive box at world={np.round(origin+BASKET_W,3)} span={S} wall_h={H}", flush=True)
    except Exception as _e:
        print("[g] basket build failed:", _e, flush=True)
def eef_root(art, bn, root, rootq):
    i=bn.index("link_6"); p=art.data.body_pos_w[0,i].cpu().numpy()-origin; q=art.data.body_quat_w[0,i].cpu().numpy()
    eef_w=p+Rq(q)@OFF; return Rq(rootq).T@(eef_w-root), q
# arm home EEF poses (root frame)
lp0, lq0 = eef_root(L, Lbn, lroot, lrootq)
rp_home, rq_home = eef_root(R, Rbn, rroot, rrootq)

# Let the (placeholder-floating) objects settle onto the table, arms holding home.
def _settle_act():
    import numpy as _np
    return torch.tensor(_np.concatenate([lp0,lq0,[1.0], rp_home,rq_home,[1.0]]),dtype=torch.float32,device=dev).view(1,-1)
# runtime override of the object's x,y (so we can vary grape position per episode)
if args.obj_xy:
    _ox,_oy=[float(v) for v in args.obj_xy.split(",")]
    _ro=u.scene.rigid_objects[args.obj]
    _pw=torch.tensor(np.concatenate([origin+np.array([_ox,_oy,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"[g] object {args.obj} repositioned to env-local ({_ox},{_oy})", flush=True)
z_before=float(u.scene.rigid_objects[args.obj].data.root_pos_w[0,2].item())
for _ in range(60):
    env.step(_settle_act())
z_after=float(u.scene.rigid_objects[args.obj].data.root_pos_w[0,2].item())
print(f"[g] settle: {args.obj} z {z_before:.3f} -> {z_after:.3f}", flush=True)
# ---- freeze the UNUSED left arm rigidly (kinematic hold) so it doesn't drift under its
#      soft diff-IK hold while the right arm works ----
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)
# ---- boost contact friction at RUNTIME (PhysX tensor view) so a REAL friction grasp holds ----
def _boost_friction(view, sfric=1.6, dfric=1.4, tag=""):
    try:
        m=view.get_material_properties().clone()   # (num_envs, num_shapes, 3): static,dynamic,restitution
        m[...,0]=sfric; m[...,1]=dfric
        idx=torch.arange(m.shape[0],dtype=torch.int32,device=m.device)
        view.set_material_properties(m, idx)
        print(f"[g] friction set on {tag}: shapes={m.shape[1]} static={sfric} dynamic={dfric}", flush=True)
        return True
    except Exception as e:
        print(f"[g] friction set FAILED on {tag}:", e, flush=True); return False
_boost_friction(R.root_physx_view, tag="right_robot(gripper)")
_boost_friction(u.scene.rigid_objects[args.obj].root_physx_view, tag=args.obj)
def fsep():
    jn=list(R.data.joint_names); i=jn.index("left_finger"); j=jn.index("right_finger")
    return (float(R.data.joint_pos[0,i].item())+float(R.data.joint_pos[0,j].item()))/2
# grasp reference: LIVE x,y from physics (root_pos_w), and a geometric-centre HEIGHT computed as
# table_top + height/2 (the object rests on the table, so this is robust to USD origin placement).
# bbox SIZE is pose-independent, so it's fine even though ComputeWorldBound reads the authored pose.
TABLE_TOP=0.45
apple_w=u.scene.rigid_objects[args.obj].data.root_pos_w[0].cpu().numpy()-origin  # live x,y (+ origin z)
obj_ext=None
try:
    import omni.usd; from pxr import UsdGeom, Usd
    stage=omni.usd.get_context().get_stage()
    pp=u.scene.rigid_objects[args.obj].root_physx_view.prim_paths[0]
    bb=UsdGeom.BBoxCache(Usd.TimeCode.Default(),[UsdGeom.Tokens.default_,UsdGeom.Tokens.render])
    rng=bb.ComputeWorldBound(stage.GetPrimAtPath(pp)).ComputeAlignedRange()
    import numpy as _np; obj_ext=_np.array(rng.GetMax())-_np.array(rng.GetMin())     # SIZE (pose-independent)
    apple_w[2]=TABLE_TOP+float(obj_ext[2])/2.0
    print(f"[g] {args.obj} size(x,y,z)={_np.round(obj_ext,3)} -> grasp center z={apple_w[2]:.3f} live_xy={_np.round(apple_w[:2],3)}", flush=True)
except Exception as e:
    print("[g] size probe failed, using root_pos z:", e, flush=True)
# Re-seat the object ON the table now that its height is known: --obj_xy drops it from z=0.55,
# which is a 5-10 cm fall for a tall object (a mug lands on its side, and then the top-down grasp
# closes on nothing). Place it so it starts resting upright, then let it settle again.
if obj_ext is not None and args.obj_xy:
    _seat_z=TABLE_TOP+float(obj_ext[2])/2.0+0.004
    _live=u.scene.rigid_objects[args.obj].data.root_pos_w[0].cpu().numpy()-origin
    _ro=u.scene.rigid_objects[args.obj]
    _ro.write_root_pose_to_sim(torch.tensor(
        np.concatenate([origin+np.array([_ox,_oy,_seat_z]),[1,0,0,0]]),
        dtype=torch.float32,device=dev).view(1,7))
    _ro.write_root_velocity_to_sim(torch.zeros((1,6),device=dev))
    for _ in range(70):
        env.step(_settle_act())
    apple_w=u.scene.rigid_objects[args.obj].data.root_pos_w[0].cpu().numpy()-origin
    apple_w[2]=TABLE_TOP+float(obj_ext[2])/2.0
    print(f"[g] re-seated {args.obj} upright at z={_seat_z:.3f}; settled xy="
          f"{np.round(apple_w[:2],3)}", flush=True)
apple_root=Rq(rrootq).T@(apple_w-rroot)
print(f"[g] right_root={np.round(rroot,3)} apple_world={np.round(apple_w,3)} apple_root={np.round(apple_root,3)}", flush=True)
# TOP-DOWN grasp quat (root=identity): approach straight down, jaw closes in world y.
# link6 axes in world: X(jaw)=+y, Y=+x, Z(approach)=-z(down). Reaches accurately (WP0 err~0).
# The jaw axis is chosen from the object's LIVE bbox: a bottle lying on its side is ~9.5 cm
# along its length but only ~2.7 cm across, and the jaw only opens 9.4 cm -- closing along the
# long axis simply cannot grip it. So close across whichever horizontal extent is smaller.
_JAW_Y=np.stack([np.array([0.,1.,0.]), np.array([1.,0.,0.]), np.array([0.,0.,-1.])],axis=1)
_JAW_X=np.stack([np.array([1.,0.,0.]), np.array([0.,-1.,0.]), np.array([0.,0.,-1.])],axis=1)
_jaw=args.jaw
if _jaw=="auto":
    if obj_ext is not None and float(obj_ext[0])<float(obj_ext[1])*0.9: _jaw="x"
    else: _jaw="y"
Rg=_JAW_X if _jaw=="x" else _JAW_Y
gq=qR(Rg)
if obj_ext is not None:
    print(f"[g] jaw axis = {_jaw} (extents x={float(obj_ext[0]):.3f} y={float(obj_ext[1]):.3f}, "
          f"jaw opens {2*0.04695:.3f} m)", flush=True)
frames=[]
rec=[]; _phase={'v':'start'}
import json as _json
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 r_eef():
    p,_=eef_root(R,Rbn,rroot,rrootq); return p
def _record(rg):
    q=R.data.joint_pos[0].cpu().numpy()
    e,_=eef_root(R,Rbn,rroot,rrootq)
    rec.append({"phase":_phase["v"],"frame":len(frames),
                "joints":[float(x) for x in q[:6]],
                "gripper":float((q[6]+q[7])/2),"grip_cmd":("close" if rg<0 else "open"),
                "eef":[float(x) for x in e],"apple_z":appz()})
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)
    th0=np.arccos(d); q2=q1-q0*d; q2/=(np.linalg.norm(q2)+1e-9)
    return q0*np.cos(th0*t)+q2*np.sin(th0*t)
def eef_full():
    p,q=eef_root(R,Rbn,rroot,rrootq); return p.astype(np.float32), q.astype(np.float32)  # pos(root), link6 quat
# ---- SMOOTH executor: per-step LERP(pos)+SLERP(quat) toward the target + integral correction
#      for the diff-IK steady-state stall (adopted from RoboLab's CartesianIKPlanner). ----
# The integral correction and the last COMMANDED point are shared by every executor below.
# Resetting either one between segments makes the commanded pose jump by the whole tracking
# error (up to ~10 cm), which the arm then chases in a few steps -- that is the visible
# "pause, then snap into a new pose" at each phase boundary. Carrying both across segments
# keeps the command continuous from phase to phase.
_CORR={"v":np.zeros(3,np.float32)}
_CMD={"p":None}
def _seg_start():
    return _CMD["p"].copy() if _CMD["p"] is not None else eef_full()[0].astype(np.float32)
def _drive(cp,rg,corr):
    """One step: command cp(+corr), then update the integral correction."""
    _CMD["p"]=np.asarray(cp,np.float32)
    env.step(act((cp+corr).astype(np.float32), gq.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):
    """Cosine ease-in/out so segments start and end at zero velocity (no start jerk)."""
    return float(0.5-0.5*np.cos(np.pi*min(max(a,0.0),1.0)))
def go(rp,rg,n,render=True):
    sp=_seg_start(); tp=np.asarray(rp,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)); cp=(1-a)*sp+a*tp
        corr=_drive(cp,rg,corr)
        if render and k % 3 == 0: capture()
        if k % 3 == 0: _record(rg)
def appz(): return float(u.scene.rigid_objects[args.obj].data.root_pos_w[0,2].item())
# ---- per-frame DEBUG OVERLAY: label each segment with its semantic action + target/eef/err ----
from PIL import Image, ImageDraw
SEM={"WP0_home":"1. APPROACH (plan from home)","WP1_PRM_path":"2. PRM APPROACH (planned path)",
     "WP2_descend":"3. DESCEND onto object","grasp_close":"4. CLOSE-GRASP (clamp)","WP3_lift":"5. LIFT (verify hold)",
     "WP4_to_box":"6. CARRY to box","WP4_to_cart":"6. CARRY to cart","WP4_carry":"6. CARRY",
     "WP4_align":"6. CARRY + ALIGN over hole","WP5_insert":"7. INSERT into hole",
     "WP5_lower_box":"7. LOWER into box","WP5_into_cart":"7. LOWER into cart","WP5_lower":"7. LOWER",
     "WP6_release":"8. RELEASE (open gripper)","WP7_retreat":"9. RETREAT"}
_CUR={"tgt":None,"grip":"OPEN"}; _RESULT={"v":""}
def capture():
    img=env.render()
    if img is None: return
    im=Image.fromarray(np.asarray(img)[...,:3].copy()); d=ImageDraw.Draw(im)
    e=r_eef(); tgt=_CUR["tgt"]; sem=SEM.get(_phase["v"], _phase["v"])
    lines=[]
    if args.episode>=0: lines.append(f"=== EPISODE {args.episode} ===")
    if _RESULT["v"]: lines.append(f"RESULT: {_RESULT['v']}")
    lines += [f"ACTION: {sem}", f"obj={args.obj}  gripper={_CUR['grip']}"]
    if tgt is not None:
        lines.append(f"target(root) [{tgt[0]:+.2f} {tgt[1]:+.2f} {tgt[2]:+.2f}]")
        lines.append(f"eef(root)    [{e[0]:+.2f} {e[1]:+.2f} {e[2]:+.2f}]  err={np.linalg.norm(e-tgt):.3f}m")
    d.rectangle([0,0,372,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))
z0=appz()
OPEN,CLOSE=1.0,-1.0
# ---- LEVEL approach waypoints (right EEF, root frame; top-down grip orientation) ----
# No raise phase: the arm starts at its raised home (t=0). It moves to the object's
# grasp height and comes in FORWARD (horizontal, constant z) — no up/down diving.
# TOP-DOWN pick: pre-grasp ABOVE the apple -> descend so fingertips seat below the equator
# (caging the lower hemisphere) -> clamp -> lift straight up.
# The descend target is the REAL grasp pose: fingertips at the object's mid-height, clamped so
# they never go below the table. Earlier this commanded a point 5 cm INSIDE the table and let the
# diff-IK stall its way up to the object -- which made the hand visibly press into the tabletop and
# pinch the object at the very fingertips. The descend below is closed-loop instead, so it arrives
# at this height and stops.
TABLE_ROOT_Z = float(TABLE_TOP - rroot[2])                    # table surface, root frame
if args.grasp_top >= 0.0 and obj_ext is not None:
    # grasp near the TOP of a tall object: the shaft below the fingers is what enters the hole
    grasp_z = TABLE_ROOT_Z + float(obj_ext[2]) - float(args.grasp_top)
else:
    grasp_z = float(apple_root[2]) - 0.005
grasp_z = max(grasp_z, TABLE_ROOT_Z + 0.010)
grasp = np.array([apple_root[0], apple_root[1], grasp_z], np.float32)
pre   = apple_root + np.array([0.0, 0.0,  0.12], np.float32)  # above the apple
lift  = apple_root + np.array([0.0, 0.0,  0.25], np.float32)  # straight up after clamping
# ---- First step = motion planning: PRM plans a collision-free path straight from the
#      arm's actual HOME end-effector pose to the pre-grasp, with the walls as blockers.
#      (No separate hand-tuned "ready" servo — the whole approach is one planned motion.)
start_eef = rp_home.astype(np.float32)
_phase["v"]="WP0_home"; print(f"[g] start HOME eef_root={np.round(start_eef,3)} -> pre-grasp {np.round(pre,3)}", flush=True)
obstacles=[yam_prm.Box(center=c, half=h) for (c,h) in BLOCKERS]
prm=yam_prm.PRM(bounds_lo=np.array([-0.05,-0.15,-0.03]), bounds_hi=np.array([0.50,0.35,0.30]),
                obstacles=obstacles, clearance=0.05, num_samples=400, k=12, seed=1)
straight_blocked = any(o.segment_hits(start_eef.astype(np.float64), pre.astype(np.float64), pad=0.05) for o in obstacles)
_path=prm.plan(start_eef.astype(np.float64), pre.astype(np.float64))
if _path is None:
    print("[g] PRM: no path found -> straight fallback", flush=True); _path=np.stack([start_eef,pre]).astype(np.float32)
# Dijkstra already returns the shortest roadmap path; shortcut() collapses it to the
# minimal set of corners (fewest waypoints / "走最小步"). We execute THOSE corners,
# converging at each so the diff-IK actually tracks the minimal detour.
_sc=yam_prm.shortcut(_path,obstacles,0.05,iters=300,seed=2)
prm_wps=_sc.astype(np.float32)                       # minimal-corner polyline
prm_world=[(rroot+w).astype(float).tolist() for w in prm_wps]
def go_converge(target, rg, tol=0.02, max_n=100):
    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 k%3==0: _record(rg)
        if a>=1.0 and np.linalg.norm(r_eef()-tp) < tol: break
    return np.linalg.norm(r_eef()-tp)
plen=float(np.sum(np.linalg.norm(np.diff(prm_wps,axis=0),axis=1)))
# Execute ALONG the planned straight line: resample into dense intermediate waypoints and
# track them one by one, so the end-effector stays pinned to the line (no IK "head-swing").
prm_exec=yam_prm.resample_polyline(prm_wps, 14).astype(np.float32)
print(f"[g] PRM: straight_blocked={straight_blocked} raw_verts={len(_path)} min_corners={len(prm_wps)} exec_pts={len(prm_exec)} path_len={plen:.3f}m", flush=True)
# ---- SMOOTH continuous executor: the COMMANDED point glides at constant speed along the
#      polyline (arc-length param) from the current pose -> no start-jerk, no corner-cut, no
#      inter-segment stop. Orientation held at gq (no twist). + integral correction. ----
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 fillet(pts, r=0.05, n=6):
    """Round the interior corners of a polyline with quadratic Beziers.

    A lift-then-carry path has a 90-degree corner at the top of the lift; tracked literally the
    arm stops dead and turns, which is the "hangs in the air then jerks sideways" look. Rounding
    the corner turns it into one continuous arc.
    """
    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 flow(pts, rg, speed=0.008, settle=18, render=True, 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")
    # Ease the arc-length rate in and out so the glide starts/ends at rest. A cosine ease peaks
    # at pi/2 x the mean rate, so stretch the duration by the same factor -- otherwise the
    # mid-path speed jumps ~57% and the carried object gets flung out of the jaws.
    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)))
        s=min(total,a*total); cp=_point_at(poly,seglens,s); _CUR["tgt"]=cp
        if on_step is not None: on_step(min(1.0,(k+1)/float(nsteps)), cp)
        corr=_drive(cp,rg,corr)
        if render and k%2==0: capture()
        if k%2==0: _record(rg)
    return float(np.linalg.norm(poly[-1]-r_eef()))
_phase["v"]="WP1_PRM_path"
flow(list(prm_exec[1:]), OPEN)                  # home -> PRM path -> hover above the object
# ---- straight vertical DESCEND, closed-loop: converge onto the grasp height and stop there ----
_phase["v"]="WP2_descend"
d_err=go_converge(grasp, OPEN, tol=0.008, max_n=140)
for _ in range(12):                             # settle: come to rest before the jaws move
    env.step(act(grasp.astype(np.float32),gq.astype(np.float32),OPEN)); _freeze_left()
capture(); _record(OPEN)
print(f"[g] approach+descend done, eef_root={np.round(r_eef(),3)} grasp={np.round(grasp,3)} err={d_err:.3f} apple_z={appz():.3f}", flush=True)
# ---- REAL physical grasp: close until the jaws STALL on the object (contact), then hold, then LIFT to verify ----
# Clamp while HOLDING the pose the arm actually reached, so the wrist does not keep driving
# downward into the table while the fingers close.
hold=r_eef().astype(np.float32)
_phase["v"]="grasp_close"; print(f"[g] close: fsep_before={fsep():.4f} hold={np.round(hold,3)}", flush=True)
_CUR["tgt"]=hold; _CUR["grip"]="CLOSE"
_prev=fsep(); _stall=0
for kc in range(160):
    env.step(act(hold,gq.astype(np.float32),CLOSE)); _freeze_left()
    if kc%6==0: capture()
    if kc%3==0: _record(CLOSE)
    cur=fsep()
    if abs(cur-_prev)<0.0002: _stall+=1
    else: _stall=0
    _prev=cur
    if _stall>=8 and cur<-0.002:   # stopped moving while still open enough => clamped on the object
        print(f"[g] jaws STALLED (contact) at fsep={cur:.4f} after {kc} steps", flush=True); break
print(f"[g] fsep_after={fsep():.4f}", flush=True)
# ---- LIFT to verify the grasp actually holds (physics only, no attach) ----
z_pre=appz()
lift = hold + np.array([0,0,0.16], np.float32)
_phase["v"]="WP3_lift"; print(f"[g] LIFT to verify hold, appz={z_pre:.3f}", flush=True)
if args.insert:
    # ---- PEG-IN-HOLE: carry over the socket, align, then insert straight down slowly.
    # The object was grasped while standing on the table, so returning the fingers to the same
    # height above the table puts the peg's base back at table level -- i.e. fully seated in the
    # hole, whose floor IS the table. Alignment first, then a pure vertical insert. ----
    hole_local=BASKET_W.copy()
    _above_local=hole_local+np.array([0,0,0.20],np.float32)
    above=(Rq(rrootq).T@(_above_local-rroot)).astype(np.float32)
    seat=np.array([above[0],above[1],float(hold[2])],np.float32)     # same grip height as the pick
    _zmax={"v":appz()}
    def _carry_step(frac, cp):
        _zmax["v"]=max(_zmax["v"], appz())
        _phase["v"]=("WP3_lift" if frac<0.30 else "WP4_align")
    _arc=fillet([_seg_start(), lift, above], r=0.06)[1:]
    flow(_arc, CLOSE, on_step=_carry_step)
    z_lift=_zmax["v"]
    print(f"[g] aligned over hole: peak appz={z_lift:.3f} eef={np.round(r_eef(),3)} "
          f"target={np.round(above,3)} err={np.linalg.norm(r_eef()-above):.3f}", flush=True)
    _phase["v"]="WP5_insert"
    ins_err=go_converge(seat, CLOSE, tol=0.006, max_n=170)           # slow vertical insertion
    print(f"[g] INSERT: seat={np.round(seat,3)} eef={np.round(r_eef(),3)} err={ins_err:.3f} "
          f"obj_z={appz():.3f}", flush=True)
    _phase["v"]="WP6_release"; go(seat, OPEN, 40)
    _phase["v"]="WP7_retreat"; flow([seat+np.array([0,0,0.16],np.float32)], OPEN)
    _placed=True
elif args.basket and not args.container:
    # Lift + carry + lower as ONE filleted arc. Done as three separate segments the arm rose,
    # stopped dead at the top, then set off sideways -- the motion read as three disjoint moves
    # instead of one reach. Corners rounded, and the object's peak height is sampled along the
    # way so the "did the grasp hold" check still works.
    _box_floor=0.010
    _half_h=(float(obj_ext[2])/2.0 if obj_ext is not None else 0.05)
    drop_local = BASKET_W + np.array([0,0,_box_floor+_half_h+0.02], np.float32)
    place=(Rq(rrootq).T@(drop_local-rroot)).astype(np.float32)
    above=place+np.array([0,0,0.10],np.float32)
    _zmax={"v":appz()}
    def _carry_step(frac, cp):
        _zmax["v"]=max(_zmax["v"], appz())
        _phase["v"]=("WP3_lift" if frac<0.25 else ("WP4_to_box" if frac<0.80 else "WP5_lower_box"))
    _arc=fillet([_seg_start(), lift, above, place], r=0.06)[1:]
    flow(_arc, CLOSE, on_step=_carry_step)
    z_lift=_zmax["v"]
    print(f"[g] lift+carry arc: peak appz={z_lift:.3f} fsep={fsep():.4f} "
          f"place={np.round(place,3)} eef={np.round(r_eef(),3)} err={np.linalg.norm(r_eef()-place):.3f}", flush=True)
    _phase["v"]="WP6_release"; print("[g] WP6 release into box", flush=True); go(place, OPEN, 40)
    _phase["v"]="WP7_retreat"; flow([above], OPEN)
    _placed=True
else:
    # every other placement target still lifts as its own segment first
    _placed=False
    flow([lift], CLOSE)
    z_lift=appz(); print(f"[g] lifted appz={z_lift:.3f} fsep={fsep():.4f}", flush=True)
if _placed:
    pass
elif args.container:
    # ---- PLACE INTO the real RobotWin container: probe its world bbox, drop above the rim ----
    try:
        crng=bb.ComputeWorldBound(stage.GetPrimAtPath("/World/envs/env_0/rw_container")).ComputeAlignedRange()
        cmin=np.array(crng.GetMin())-origin; cmax=np.array(crng.GetMax())-origin
        # Set the object DOWN on the container's top surface instead of releasing 6 cm above it:
        # dropped from that height a cup topples and ends up on its side next to the rack, while
        # the xy-only success check still passes.
        _half_h=(float(obj_ext[2])/2.0 if obj_ext is not None else 0.03)
        drop_local=np.array([(cmin[0]+cmax[0])/2.0,(cmin[1]+cmax[1])/2.0,
                             cmax[2]+_half_h+0.006], np.float32)
        print(f"[g] container bbox top_z={cmax[2]:.3f} drop_local={np.round(drop_local,3)}", flush=True)
    except Exception as e:
        print("[g] container probe failed:", e, flush=True); drop_local=BASKET_W+np.array([0,0,0.16],np.float32)
    place=(Rq(rrootq).T@(drop_local-rroot)).astype(np.float32); above=place+np.array([0,0,0.10],np.float32)
    _phase["v"]="WP4_to_box"; print("[g] carry+lower to container (continuous)", flush=True)
    flow([above, place], CLOSE)
    print(f"[g] CONTAINER-REACH: place={np.round(place,3)} eef={np.round(r_eef(),3)} err={np.linalg.norm(r_eef()-place):.3f}", flush=True)
    _phase["v"]="WP6_release"; go(place, OPEN, 40)
    _phase["v"]="WP7_retreat"; flow([above], OPEN)
elif args.cart:
    # ---- PLACE INTO CART: probe the cart's world bbox, aim for above its basket opening ----
    try:
        cpp=u.scene.rigid_objects["cart"].root_physx_view.prim_paths[0]
        crng=bb.ComputeWorldBound(stage.GetPrimAtPath(cpp)).ComputeAlignedRange()
        cmin=np.array(crng.GetMin())-origin; cmax=np.array(crng.GetMax())-origin
        drop_w=np.array([(cmin[0]+cmax[0])/2.0,(cmin[1]+cmax[1])/2.0, cmax[2]-0.02], np.float32)  # just above the top rim
        place=(Rq(rrootq).T@(drop_w-rroot)).astype(np.float32)
        print(f"[g] cart bbox(world-local) x[{cmin[0]:.2f},{cmax[0]:.2f}] y[{cmin[1]:.2f},{cmax[1]:.2f}] top_z={cmax[2]:.2f} -> drop_root={np.round(place,3)}", flush=True)
    except Exception as e:
        print("[g] cart probe failed:", e, flush=True); place=grasp+np.array([0.1,-0.1,0.0],np.float32)
    above=place+np.array([0,0,0.15],np.float32)
    _phase["v"]="WP4_to_cart"; print(f"[g] WP4 carry toward cart above={np.round(above,3)}", flush=True); go(above, CLOSE, 110)
    print(f"[g] CART-REACH: target_above={np.round(above,3)} achieved_eef={np.round(r_eef(),3)} err={np.linalg.norm(r_eef()-above):.3f}", flush=True)
    _phase["v"]="WP5_into_cart"; go(place, CLOSE, 60)
    _phase["v"]="WP6_release"; print(f"[g] WP6 release into cart", flush=True); go(place, OPEN, 45)
    _phase["v"]="WP7_retreat"; go(above, OPEN, 35)
elif args.basket:
    # ---- PLACE INTO the primitive box: carry above box center, lower, release ----
    # Lower the object until it nearly touches the box floor before opening, instead of dropping
    # it from a fixed 14 cm (tall objects bounced back out of the shallow tray).
    _box_floor=0.010                                            # tray floor thickness
    _half_h=(float(obj_ext[2])/2.0 if obj_ext is not None else 0.05)
    drop_local = BASKET_W + np.array([0,0,_box_floor+_half_h+0.02], np.float32)
    place=(Rq(rrootq).T@(drop_local-rroot)).astype(np.float32)
    above=place+np.array([0,0,0.10],np.float32)
    _phase["v"]="WP4_to_box"; print(f"[g] WP4 carry+lower to box (continuous)", flush=True)
    flow([above, place], CLOSE)   # carry over the box + lower, one continuous motion
    print(f"[g] BOX-REACH: place={np.round(place,3)} achieved_eef={np.round(r_eef(),3)} err={np.linalg.norm(r_eef()-place):.3f}", flush=True)
    _phase["v"]="WP6_release"; print(f"[g] WP6 release into box", flush=True); go(place, OPEN, 40)   # stop only to open
    _phase["v"]="WP7_retreat"; flow([above], OPEN)
else:
    # ---- PLACE on the table (default) ----
    place = grasp + np.array([0.10, -0.10, 0.0], np.float32)
    _phase["v"]="WP4_carry";   print(f"[g] WP4 carry (lifted) to place xy={np.round(place[:2],3)}", flush=True); go(place+np.array([0,0,0.22],np.float32), CLOSE, 60)
    _phase["v"]="WP5_lower";   print(f"[g] WP5 lower to table appz={appz():.3f}", flush=True);                  go(place, CLOSE, 55)
    _phase["v"]="WP6_release"; print(f"[g] WP6 open / release", flush=True);                                    go(place, OPEN, 45)
    _phase["v"]="WP7_retreat"; print(f"[g] WP7 retreat up", flush=True);                                        go(place+np.array([0,0,0.20],np.float32), OPEN, 35)
z1=appz()
# ---- episode success ----
if args.insert:
    # Seated means: centred on the hole AND resting at table level inside it. A peg left standing
    # on the socket walls sits a full wall-height higher, so the z test separates the two.
    _w=u.scene.rigid_objects[args.obj].data.root_pos_w[0].cpu().numpy()-origin
    _half=(float(obj_ext[2])/2.0 if obj_ext is not None else 0.065)
    _seated=(abs(_w[0]-BASKET_W[0])<0.02 and abs(_w[1]-BASKET_W[1])<0.02
             and _w[2] < TABLE_TOP+_half+0.02)
    _RESULT["v"]="SUCCESS" if _seated else "FAIL"
    _phase["v"]="RESULT"
    print(f"[g] EPISODE_RESULT: {_RESULT['v']} obj_world=({_w[0]:.3f},{_w[1]:.3f},{_w[2]:.3f}) "
          f"hole=({BASKET_W[0]:.2f},{BASKET_W[1]:.2f}) seated_z<{TABLE_TOP+_half+0.02:.3f}", flush=True)
    for _ in range(14): capture()
elif args.basket:
    _w=u.scene.rigid_objects[args.obj].data.root_pos_w[0].cpu().numpy()-origin
    if args.container:
        # ON a stand/rack: the object must be resting on the container's TOP surface. An xy-only
        # test passes an object that fell off and is lying on the table beside it, which is
        # exactly how a "cup on the rack" episode reported success with the cup on its side.
        try:
            _crng=bb.ComputeWorldBound(stage.GetPrimAtPath("/World/envs/env_0/rw_container")).ComputeAlignedRange()
            _ctop=float(_crng.GetMax()[2])-origin[2]
        except Exception:
            _ctop=TABLE_TOP
        _inbox = (abs(_w[0]-BASKET_W[0])<0.12 and abs(_w[1]-BASKET_W[1])<0.12
                  and _w[2] > _ctop-0.01)
        print(f"[g] on-container check: obj_z={_w[2]:.3f} must exceed container_top-0.01={_ctop-0.01:.3f}", flush=True)
    else:
        _inbox = abs(_w[0]-BASKET_W[0])<0.15 and abs(_w[1]-BASKET_W[1])<0.15 and _w[2]<BASKET_W[2]+0.16
    _RESULT["v"]="SUCCESS" if _inbox else "FAIL"
    _phase["v"]="RESULT"
    print(f"[g] EPISODE_RESULT: {_RESULT['v']} obj_world=({_w[0]:.3f},{_w[1]:.3f},{_w[2]:.3f}) box=({BASKET_W[0]:.2f},{BASKET_W[1]:.2f}) dz={z_lift-z0:.3f}", flush=True)
    for _ in range(14): capture()   # hold the SUCCESS/FAIL result on screen
os.makedirs(os.path.dirname(args.video),exist_ok=True)
# Drop the first captured frames: the renderer has not settled at that point, so they come out
# with the wrong camera pose, unresolved textures and missing geometry (a reviewer reading the
# contact sheet sees a "container missing" scene that never actually existed).
WARMUP_FRAMES=2
if len(frames) > WARMUP_FRAMES+4: frames = frames[WARMUP_FRAMES:]
if frames: imageio.mimsave(args.video, frames, fps=14)
_json.dump({"dt":1.0/30,"joint_names":["joint1","joint2","joint3","joint4","joint5","joint6"],
            "home":[-0.017453,1.640610,1.483530,-1.466077,-0.087266,0.0],"steps":rec,
            "planner":"task-space PRM (roadmap + Dijkstra + shortcut)",
            "prm_path_world":prm_world,
            "obstacles":_wall_json},
           open(os.path.splitext(args.video)[0]+"_pose.json","w"))
print(f"[g] pose json -> {os.path.splitext(args.video)[0]}_pose.json ({len(rec)} steps)", flush=True)
print(f"[g] LIFT dz={z_lift-z0:.3f} lifted={z_lift-z0>0.05} | PLACED final_z={z1:.3f} (task=pick+carry+release) -> {args.video}", flush=True)
env.close(); app.close(); print("YAM_GRASP_OK", flush=True)