scripts/yam_dualarm.py

"""YAM DUAL-ARM parallel pick(-and-place). Task allocation (which arm -> which object) comes
from the LLM/nearest-arm layer; each arm then runs PRM (planning-only walls) + the REAL
friction grasp IN PARALLEL. Renders a real Isaac Sim video. No kinematic attach anywhere."""
import argparse, sys, os
from isaaclab.app import AppLauncher
parser = argparse.ArgumentParser()
parser.add_argument("--video", default="outputs/yam_dualarm.mp4")
parser.add_argument("--place", action="store_true", help="also carry each object aside and release")
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=(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)
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])
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

# ---- ALLOCATION: each arm picks its own cup (parallel), then places it aside ----
ASSIGN={"left_robot":"rw_cup2","right_robot":"rw_cup"}
print(f"[da] allocation: left->{ASSIGN['left_robot']} right->{ASSIGN['right_robot']} (parallel)", flush=True)
lp0,lq0=eef_root(L,Lbn,lroot,lrootq); rp0,rq0=eef_root(R,Rbn,rroot,rrootq)
OPEN,CLOSE=1.0,-1.0
def act2(lp,lq,lg,rp,rq,rg):
    return torch.tensor(np.concatenate([lp,lq,[lg],rp,rq,[rg]]),dtype=torch.float32,device=dev).view(1,-1)
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))
reposition("rw_cup",[0.05,-0.14]); reposition("rw_cup2",[0.05,0.14])   # symmetric reachable spots
for _ in range(60): env.step(act2(lp0,lq0,OPEN,rp0,rq0,OPEN))
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)); return True
    except Exception as e: print("[da] fric fail",e); return False
boost(R.root_physx_view); boost(L.root_physx_view)
boost(u.scene.rigid_objects["rw_cup"].root_physx_view); boost(u.scene.rigid_objects["rw_cup2"].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)
TABLE=0.45; CENTER_Z=0.47   # cup geometric-centre world height (rests on table, ~4.5cm tall)
def obj_root(name,root,rootq):
    w=u.scene.rigid_objects[name].data.root_pos_w[0].cpu().numpy()-origin
    return Rq(rootq).T@(np.array([w[0],w[1],CENTER_Z],np.float32)-root)
def objz(name): return float(u.scene.rigid_objects[name].data.root_pos_w[0,2].item())
L_obj=obj_root(ASSIGN["left_robot"],lroot,lrootq); R_obj=obj_root(ASSIGN["right_robot"],rroot,rrootq)
L_pre=L_obj+np.array([0,0,0.12],np.float32); L_grasp=L_obj+np.array([0,0,-0.05],np.float32); L_lift=L_grasp+np.array([0,0,0.22],np.float32)
R_pre=R_obj+np.array([0,0,0.12],np.float32); R_grasp=R_obj+np.array([0,0,-0.05],np.float32); R_lift=R_grasp+np.array([0,0,0.22],np.float32)
print(f"[da] left obj_root={np.round(L_obj,3)}  right obj_root={np.round(R_obj,3)}", flush=True)

def plan(start,goal):
    prm=yam_prm.PRM(bounds_lo=np.array([-0.1,-0.45,-0.05]),bounds_hi=np.array([0.55,0.45,0.35]),
                    obstacles=[],clearance=0.04,num_samples=200,k=10,seed=1)
    p=prm.plan(start.astype(np.float64),goal.astype(np.float64))
    if p is None: p=np.stack([start,goal]).astype(np.float32)
    return yam_prm.resample_polyline(yam_prm.shortcut(p,[],0.04,iters=100,seed=2),12).astype(np.float32)
L_path=plan(lp0,L_pre); R_path=plan(rp0,R_pre)

frames=[]
def snap():
    img=env.render()
    if img is not None: frames.append(np.asarray(img)[...,:3])
def eefL(): p,_=eef_root(L,Lbn,lroot,lrootq); return p
def eefR(): p,_=eef_root(R,Rbn,rroot,rrootq); return p
def drive(lt,lg,rt,rg,n,tol=None):
    ls=eefL(); rs=eefR(); lt=np.asarray(lt,np.float32); rt=np.asarray(rt,np.float32)  # smooth LERP (no jerk)
    for k in range(n):
        a=(k+1)/float(n); lc=(1-a)*ls+a*lt; rc=(1-a)*rs+a*rt
        env.step(act2(lc,gq,lg,rc,gq,rg))
        if k%4==0: snap()
        if tol and a>=1.0 and np.linalg.norm(eefL()-lt)<tol and np.linalg.norm(eefR()-rt)<tol: break

# ---- PHASE 1: parallel approach along each arm's PRM path ----
m=max(len(L_path),len(R_path))
for i in range(1,m):
    lt=L_path[min(i,len(L_path)-1)]; rt=R_path[min(i,len(R_path)-1)]; last=(i==m-1)
    drive(lt,OPEN,rt,OPEN, 80 if last else 20, tol=(0.02 if last else None))
print(f"[da] approached: L_err={np.linalg.norm(eefL()-L_pre):.3f} R_err={np.linalg.norm(eefR()-R_pre):.3f}", flush=True)
# ---- PHASE 2: parallel descend ----
drive(L_grasp,OPEN,R_grasp,OPEN,90)
print(f"[da] descended: Leef={np.round(eefL(),3)} Reef={np.round(eefR(),3)}", flush=True)
# ---- PHASE 3: parallel close (both clamp) ----
for k in range(120):
    env.step(act2(L_grasp,gq,CLOSE,R_grasp,gq,CLOSE))
    if k%6==0: snap()
# ---- PHASE 4: parallel lift (verify hold) ----
zL0=objz(ASSIGN["left_robot"]); zR0=objz(ASSIGN["right_robot"])
drive(L_lift,CLOSE,R_lift,CLOSE,80)
zL1=objz(ASSIGN["left_robot"]); zR1=objz(ASSIGN["right_robot"])
# ---- optional PHASE 5: place aside + release ----
if args.place:
    L_drop=np.array([0.10,-0.05,0.10],np.float32); R_drop=np.array([0.10,0.05,0.10],np.float32)  # swap sides (cross place)
    drive(L_lift+np.array([0,0,0.02],np.float32),CLOSE, R_lift+np.array([0,0,0.02],np.float32),CLOSE, 10)
    drive(np.array([L_drop[0],L_drop[1],L_lift[2]],np.float32),CLOSE, np.array([R_drop[0],R_drop[1],R_lift[2]],np.float32),CLOSE, 70)
    drive(L_drop,CLOSE,R_drop,CLOSE,60)
    for k in range(40):
        env.step(act2(L_drop,gq,OPEN,R_drop,gq,OPEN))
        if k%6==0: snap()

os.makedirs(os.path.dirname(args.video),exist_ok=True)
if frames: imageio.mimsave(args.video, frames, fps=6)
print(f"[da] LEFT  {ASSIGN['left_robot']}: dz={zL1-zL0:.3f} lifted={zL1-zL0>0.05}", flush=True)
print(f"[da] RIGHT {ASSIGN['right_robot']}: dz={zR1-zR0:.3f} lifted={zR1-zR0>0.05}", flush=True)
print(f"[da] -> {args.video} ({len(frames)} frames)", flush=True)
env.close(); app.close(); print("YAM_DUALARM_OK", flush=True)