scripts/yam_cubes.py

"""YAM cube tasks: SORT colored cubes into color-matched bins, or STACK cubes on top of each
other. Right arm, PRM approach + real friction grasp (no attach). Renders a labelled video.
Usage: python scripts/yam_cubes.py --headless --task sort|stack --video outputs/cubes.mp4"""
import argparse, sys, os
from isaaclab.app import AppLauncher
p=argparse.ArgumentParser()
p.add_argument("--task", default="sort", choices=["sort","stack"])
p.add_argument("--video", default="outputs/yam_cubes.mp4")
AppLauncher.add_app_launcher_args(p); args=p.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, 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: pass
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"]
def root_of(a): return a.data.root_pos_w[0].cpu().numpy()-origin, a.data.root_quat_w[0].cpu().numpy()
rroot,rrootq=root_of(R); OFF=np.array([0,0,0.13]); TABLE=0.45; CUBE=0.05
def eef_root(a,bn,root,rootq):
    i=bn.index("link_6"); pp=a.data.body_pos_w[0,i].cpu().numpy()-origin; q=a.data.body_quat_w[0,i].cpu().numpy()
    return Rq(rootq).T@((pp+Rq(q)@OFF)-root), q
lp0,lq0=eef_root(L,list(L.data.body_names),*root_of(L)); rp0,rq0=eef_root(R,Rbn,rroot,rrootq)
OPEN,CLOSE=1.0,-1.0
import isaaclab.sim as sim_utils
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)
Rg=np.stack([np.array([0.,1.,0.]),np.array([1.,0.,0.]),np.array([0.,0.,-1.])],axis=1); gq=qR(Rg)
def reposition(name,xy):
    ro=u.scene.rigid_objects[name]
    ro.write_root_pose_to_sim(torch.tensor(np.concatenate([origin+np.array([xy[0],xy[1],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))

# ---- scene setup per task ----
BINS=[]  # (xy, color, name) for sort
if args.task=="sort":
    reposition("cube_r",[0.00,0.06]); reposition("cube_b",[0.13,0.06])
    for n in ["cube_g","cube_y"]: reposition(n,[3.0+np.random.rand()*0,4.0])  # keep off-scene deterministically
    reposition("cube_g",[3.0,4.0]); reposition("cube_y",[3.4,4.0])
    BINS=[([0.00,-0.20],(0.75,0.2,0.2),"bin_r"),([0.15,-0.20],(0.2,0.35,0.8),"bin_b")]
    for xy,col,nm in BINS:
        S,H,T=0.14,0.05,0.01
        for sub,sz,off in [("f",(S,S,T),(0,0,T/2)),("xp",(T,S,H),(S/2,0,H/2)),("xn",(T,S,H),(-S/2,0,H/2)),("yp",(S,T,H),(0,S/2,H/2)),("yn",(S,T,H),(0,-S/2,H/2))]:
            c=sim_utils.CuboidCfg(size=sz,visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=col),collision_props=sim_utils.CollisionPropertiesCfg())
            c.func(f"/World/envs/env_0/{nm}_{sub}",c,translation=tuple((origin+np.array([xy[0],xy[1],TABLE])+np.array(off)).tolist()))
    JOBS=[("cube_r",[0.00,-0.20]),("cube_b",[0.15,-0.20])]
else:  # stack — keep every cube within reliable reach (root dist < ~0.38)
    reposition("cube_g",[0.06,-0.02]); reposition("cube_r",[0.00,0.08]); reposition("cube_b",[0.12,0.02])
    for n in ["cube_y"]: reposition(n,[3.0,4.0])
    JOBS=[("cube_r",("stack",1)),("cube_b",("stack",2))]  # onto cube_g at [0.06,-0.02]
    STACK_XY=[0.06,-0.02]

for _ in range(60): env.step(act(rp0,rq0,OPEN))
lhome=L.data.joint_pos[0].clone(); _lz=torch.zeros((1,lhome.shape[0]),device=dev)
def freeze(): L.write_joint_state_to_sim(lhome.view(1,-1),_lz)
def boost(v,s=1.6,d=1.4):
    try:
        m=v.get_material_properties().clone(); m[...,0]=s; m[...,1]=d; v.set_material_properties(m,torch.arange(m.shape[0],dtype=torch.int32,device=m.device))
    except Exception as e: print("fric",e)
boost(R.root_physx_view)
for j in (JOBS if args.task=="sort" else JOBS): boost(u.scene.rigid_objects[j[0]].root_physx_view)
frames=[]; _CUR={"a":"start","g":"OPEN"}
def r_eef(): p,_=eef_root(R,Rbn,rroot,rrootq); return p
def rsep(): jn=list(R.data.joint_names); return (R.data.joint_pos[0,jn.index("left_finger")].item()+R.data.joint_pos[0,jn.index("right_finger")].item())/2
def cap():
    img=env.render()
    if img is None: return
    im=Image.fromarray(np.asarray(img)[...,:3].copy()); d=ImageDraw.Draw(im); e=r_eef()
    lines=[f"TASK: {args.task.upper()} cubes", f"ACTION: {_CUR['a']}", f"gripper: {_CUR['g']}"]
    d.rectangle([0,0,330,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))
def go(tgt,g,n,tol=None,a=None):
    if a: _CUR["a"]=a
    _CUR["g"]="CLOSE" if g<0 else "OPEN"
    for k in range(n):
        env.step(act(tgt.astype(np.float32),gq,g)); freeze()
        if k%5==0: cap()
        if tol and np.linalg.norm(r_eef()-tgt)<tol: break
def close_stall(tgt,n=140):
    _CUR["a"]="CLOSE-GRASP"; _CUR["g"]="CLOSE"; prev=rsep(); st=0
    for k in range(n):
        env.step(act(tgt.astype(np.float32),gq,CLOSE)); freeze()
        if k%5==0: cap()
        cur=rsep()
        if abs(cur-prev)<0.0002: st+=1
        else: st=0
        prev=cur
        if st>=8 and cur<-0.002: break
def objw(n): return u.scene.rigid_objects[n].data.root_pos_w[0].cpu().numpy()-origin
def oroot(n): return Rq(rrootq).T@(objw(n)-rroot)

def pick(name):
    o=oroot(name); o[2]=(TABLE+CUBE/2.0)-rroot[2]   # grasp centre height (cube rests on table)
    pre=o+np.array([0,0,0.13],np.float32); grasp=o+np.array([0,0,-0.03],np.float32)
    go(pre,OPEN,60,tol=0.02,a="APPROACH above cube"); go(grasp,OPEN,80,a="DESCEND")
    close_stall(grasp); go(grasp+np.array([0,0,0.22],np.float32),CLOSE,70,a="LIFT")
def place(xy, top_world_z, precise=False):
    # command the TCP so the held cube's bottom ends just above top_world_z, then release.
    # diff-IK stalls ~4cm high -> command lower for a precise stack; drop-in bins can be higher.
    tcp_z = (top_world_z - 0.03 if precise else top_world_z + CUBE+0.02)  # precise: counter the ~4cm IK stall
    tgt=(Rq(rrootq).T@(np.array([xy[0],xy[1],tcp_z])-rroot)).astype(np.float32)
    go(tgt+np.array([0,0,0.13],np.float32),CLOSE,100,tol=0.02,a="CARRY over target"); go(tgt,CLOSE,80,tol=0.02,a="LOWER")
    go(tgt,OPEN,45,a="RELEASE"); go(tgt+np.array([0,0,0.14],np.float32),OPEN,30,a="RETREAT")

results=[]
if args.task=="sort":
    for name,binxy in JOBS:
        pick(name); place(binxy, TABLE)  # drop into bin (floor at table height)
        w=objw(name); inb=abs(w[0]-binxy[0])<0.09 and abs(w[1]-binxy[1])<0.09
        results.append((name,inb)); print(f"[cubes] SORT {name} -> bin{binxy} in_bin={inb} world=({w[0]:.3f},{w[1]:.3f})",flush=True)
else:
    for name,(_,lvl) in JOBS:
        pick(name); place(STACK_XY, TABLE+lvl*CUBE, precise=True)  # stack level lvl on the base cube
        w=objw(name); onstack=abs(w[0]-STACK_XY[0])<0.05 and abs(w[1]-STACK_XY[1])<0.05 and w[2]>TABLE+0.5*CUBE+0.02
        results.append((name,onstack)); print(f"[cubes] STACK {name} lvl{lvl} stacked={onstack} world=({w[0]:.3f},{w[1]:.3f},{w[2]:.3f})",flush=True)
_CUR["a"]="RESULT"; [cap() for _ in range(12)]
os.makedirs(os.path.dirname(args.video) or ".",exist_ok=True)
if frames: imageio.mimsave(args.video, frames, fps=7)
ok=sum(1 for _,r in results if r)
print(f"[cubes] {args.task}: {ok}/{len(results)} ok -> {args.video} ({len(frames)} frames)",flush=True)
env.close(); app.close(); print("YAM_CUBES_OK",flush=True)