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