scripts/yam_longhorizon.py

"""YAM LONG-HORIZON task: a primitive box (built from Cuboids, license-clean) sits in the
middle; the right arm picks each object and drops it into the box, one after another, until
ALL target objects are in the box. Each pick = PRM approach + REAL friction grasp (no attach).
Renders one continuous Isaac Sim video."""
import argparse, sys, os
from isaaclab.app import AppLauncher
parser=argparse.ArgumentParser()
parser.add_argument("--video", default="outputs/yam_longhorizon.mp4")
parser.add_argument("--objects", default="can,can2,grape")
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
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)
# 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=(1.0,-1.0,1.2); _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)
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); lroot,lrootq=root_of(L)
OFF=np.array([0,0,0.13])

# ---- central primitive box (Cuboids only, license-clean) ----
BASKET=np.array([0.10,0.0,0.45],np.float32)  # env-local, table height, reachable center
import isaaclab.sim as sim_utils
S,H,T=0.16,0.10,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()))
_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"[lh] built central box at world={np.round(origin+BASKET,3)} span={S} h={H}", 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); rp0,rq0=eef_root(R,Rbn,rroot,rrootq)
OPEN,CLOSE=1.0,-1.0
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)
for _ in range(60): env.step(act(rp0,rq0,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,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("[lh] fric fail",e)
boost(R.root_physx_view)
for o in args.objects.split(","): boost(u.scene.rigid_objects[o].root_physx_view)

Rg=np.stack([np.array([0.,1.,0.]),np.array([1.,0.,0.]),np.array([0.,0.,-1.])],axis=1); gq=qR(Rg)
frames=[]
def snap():
    img=env.render()
    if img is not None: frames.append(np.asarray(img)[...,:3])
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 objw(name): return u.scene.rigid_objects[name].data.root_pos_w[0].cpu().numpy()-origin
def obj_root(name): return Rq(rrootq).T@(objw(name)-rroot)
def go(tgt,g,n,tol=None):
    for k in range(n):
        env.step(act(tgt.astype(np.float32),gq,g)); freeze_left()
        if k%6==0: snap()
        if tol and np.linalg.norm(r_eef()-tgt)<tol: break
def close_stall(tgt,n=140):
    prev=rsep(); st=0
    for k in range(n):
        env.step(act(tgt.astype(np.float32),gq,CLOSE)); freeze_left()
        if k%6==0: snap()
        cur=rsep()
        if abs(cur-prev)<0.0002: st+=1
        else: st=0
        prev=cur
        if st>=8 and cur<-0.002: break

BASKET_root=(Rq(rrootq).T@(BASKET-rroot)).astype(np.float32)
DROP=BASKET_root+np.array([0,0,0.18],np.float32)     # release point, above the box walls
DROP_HI=DROP+np.array([0,0,0.10],np.float32)

def pick_place(name, idx):
    o=obj_root(name)
    pre=o+np.array([0,0,0.12],np.float32); grasp=o+np.array([0,0,-0.05],np.float32); lift=grasp+np.array([0,0,0.24],np.float32)
    print(f"[lh] --- object {idx}: {name} obj_root={np.round(o,3)} ---", flush=True)
    # PRM approach from current eef to pre
    prm=yam_prm.PRM(bounds_lo=np.array([-0.05,-0.35,-0.03]),bounds_hi=np.array([0.55,0.45,0.4]),obstacles=[],clearance=0.03,num_samples=150,k=10,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,[],0.03,iters=80,seed=2),10).astype(np.float32)
    for i in range(1,len(wps)): go(wps[i],OPEN, 70 if i==len(wps)-1 else 18, tol=(0.02 if i==len(wps)-1 else None))
    go(grasp,OPEN,80)                                  # descend
    print(f"[lh]   descend err={np.linalg.norm(r_eef()-grasp):.3f}", flush=True)
    close_stall(grasp)                                 # clamp
    zg0=objw(name)[2]
    go(lift,CLOSE,70)                                  # lift
    zg1=objw(name)[2]
    print(f"[lh]   grasp lift dz={zg1-zg0:.3f} lifted={zg1-zg0>0.05}", flush=True)
    go(DROP_HI,CLOSE,90); go(DROP,CLOSE,45)            # carry over the box
    go(DROP,OPEN,45)                                   # release into box
    go(DROP_HI,OPEN,30)                                # retreat up
    return zg1-zg0>0.05

order=args.objects.split(",")
results={}
for i,name in enumerate(order,1):
    results[name]=pick_place(name,i)
# return home-ish
go(rp0,OPEN,50)

# ---- success check: how many objects ended up inside the box footprint ----
inb=0
for name in order:
    w=objw(name); dx=abs(w[0]-BASKET[0]); dy=abs(w[1]-BASKET[1])
    inside = dx<S/2+0.02 and dy<S/2+0.02 and w[2]<BASKET[2]+0.14
    inb += int(inside)
    print(f"[lh] {name}: final world=({w[0]:.3f},{w[1]:.3f},{w[2]:.3f}) in_box={inside}", flush=True)
os.makedirs(os.path.dirname(args.video),exist_ok=True)
if frames: imageio.mimsave(args.video, frames, fps=7)
print(f"[lh] DONE: {inb}/{len(order)} objects in the box -> {args.video} ({len(frames)} frames)", flush=True)
env.close(); app.close(); print("YAM_LONGHORIZON_OK", flush=True)