"""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)=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)