"""YAM stack task: two green blocks and a flat TARGET REGION (a paper-thin marker on the table). Both blocks must end up stacked on the target -- first block placed on the marker, second block placed on top of the first. Stacking is the precision case: the second block has to land within ~1.5 cm of the first one's centre or it topples off, so the placement uses the same closed-loop descend as the pick rather than an open-loop drop. python scripts/yam_stack_target.py --headless --video outputs/tasks/stack.mp4 """ import argparse, sys, os from isaaclab.app import AppLauncher parser = argparse.ArgumentParser() parser.add_argument("--blocks", default="cube_g,cube_g2", help="the blocks to stack, bottom first") parser.add_argument("--block_xy", default="-0.03,0.11;0.04,0.10", help="';'-separated start x,y per block") parser.add_argument("--target_xy", default="0.00,-0.20", help="target region centre x,y (env-local)") parser.add_argument("--target_size", type=float, default=0.11, help="target region side length (m)") parser.add_argument("--block_rpy", default="0,0,0", help="roll,pitch,yaw in DEGREES applied to every block when seating it. RoboTwin " "cups spawn upside-down, and an inverted cup has its rim at the bottom, so " "the jaws grip a narrowing taper and it slips out.") parser.add_argument("--episode", type=int, default=-1) parser.add_argument("--video", default="outputs/tasks/yam_stack_target.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 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 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) _cfg.episode_length_s = 1.0e6 try: _cfg.terminations.time_out = None except Exception as _e: print("[st] 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) rroot = R.data.root_pos_w[0].cpu().numpy()-origin; rrootq = R.data.root_quat_w[0].cpu().numpy() 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]); TABLE_TOP = 0.45 OPEN, CLOSE = 1.0, -1.0 BLOCKS = [b for b in args.blocks.split(",") if b] TXY = [float(v) for v in args.target_xy.split(",")] TARGET = np.array([TXY[0], TXY[1], TABLE_TOP], np.float32) # ---- the TARGET REGION: a paper-thin plate on the table (visual marker, collidable but flat) ---- import isaaclab.sim as sim_utils _mark = sim_utils.CuboidCfg(size=(args.target_size, args.target_size, 0.002), visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.95, 0.95, 0.92)), collision_props=sim_utils.CollisionPropertiesCfg()) _mark.func("/World/envs/env_0/target_region", _mark, translation=tuple((origin+TARGET+np.array([0, 0, 0.001], np.float32)).astype(float).tolist())) _edge = sim_utils.CuboidCfg(size=(args.target_size*1.06, args.target_size*1.06, 0.001), visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.15, 0.55, 0.25))) _edge.func("/World/envs/env_0/target_edge", _edge, translation=tuple((origin+TARGET).astype(float).tolist())) print(f"[st] target region at world={np.round(origin+TARGET,3)} size={args.target_size}", 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) rp_home, rq_home = eef_root(R, Rbn, rroot, rrootq) gq = qR(np.stack([np.array([0., 1., 0.]), np.array([1., 0., 0.]), np.array([0., 0., -1.])], axis=1)) def act(rp, rg): return torch.tensor(np.concatenate([lp0, lq0, [1.0], rp, gq, [rg]]), dtype=torch.float32, device=dev).view(1, -1) missing = [b for b in BLOCKS if b not in u.scene.rigid_objects] if missing: raise SystemExit(f"[st] blocks not in scene: {missing}") START_XY = {} for name, xy in zip(BLOCKS, args.block_xy.split(";")): x, y = [float(v) for v in xy.split(",")] START_XY[name] = (x, y) ro = u.scene.rigid_objects[name] ro.write_root_pose_to_sim(torch.tensor(np.concatenate([origin+np.array([x, y, 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)) print(f"[st] placed {name} at ({x},{y})", flush=True) for _ in range(60): env.step(act(rp_home, 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, tag, 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(f"[st] friction set failed on {tag}:", e, flush=True) _boost(R.root_physx_view, "right_robot") for b in BLOCKS: _boost(u.scene.rigid_objects[b].root_physx_view, b) def objw(name): return u.scene.rigid_objects[name].data.root_pos_w[0].cpu().numpy()-origin def r_eef(): p, _ = eef_root(R, Rbn, rroot, rrootq); return p def fsep(): jn = list(R.data.joint_names) return (float(R.data.joint_pos[0, jn.index("left_finger")].item()) + float(R.data.joint_pos[0, jn.index("right_finger")].item()))/2 import omni.usd from pxr import UsdGeom, Usd stage = omni.usd.get_context().get_stage() bbc = UsdGeom.BBoxCache(Usd.TimeCode.Default(), [UsdGeom.Tokens.default_, UsdGeom.Tokens.render]) def obj_size(name): try: pp = u.scene.rigid_objects[name].root_physx_view.prim_paths[0] rng = bbc.ComputeWorldBound(stage.GetPrimAtPath(pp)).ComputeAlignedRange() return np.array(rng.GetMax())-np.array(rng.GetMin()) except Exception as e: print(f"[st] size probe failed for {name}:", e, flush=True); return np.array([0.05, 0.05, 0.05]) # Seat every object ON the table before starting: they are written in 10 cm above it, and a # cup-shaped object that tall lands on its side, which no top-down grasp can recover. _br, _bp, _by = [np.radians(float(v)) for v in args.block_rpy.split(",")] _BQ = np.array([ np.cos(_br/2)*np.cos(_bp/2)*np.cos(_by/2)+np.sin(_br/2)*np.sin(_bp/2)*np.sin(_by/2), np.sin(_br/2)*np.cos(_bp/2)*np.cos(_by/2)-np.cos(_br/2)*np.sin(_bp/2)*np.sin(_by/2), np.cos(_br/2)*np.sin(_bp/2)*np.cos(_by/2)+np.sin(_br/2)*np.cos(_bp/2)*np.sin(_by/2), np.cos(_br/2)*np.cos(_bp/2)*np.sin(_by/2)-np.sin(_br/2)*np.sin(_bp/2)*np.cos(_by/2)]) for _n in BLOCKS: _e = obj_size(_n) _x, _y = START_XY[_n] _ro = u.scene.rigid_objects[_n] _ro.write_root_pose_to_sim(torch.tensor( np.concatenate([origin+np.array([_x, _y, TABLE_TOP+float(_e[2])/2.0+0.004]), _BQ]), dtype=torch.float32, device=dev).view(1, 7)) _ro.write_root_velocity_to_sim(torch.zeros((1, 6), device=dev)) print(f"[st] seated {_n} upright: size={np.round(_e,3)}", flush=True) for _ in range(80): env.step(act(rp_home, OPEN)); _freeze_left() frames = []; _phase = {"v": "approach"}; _CUR = {"tgt": None, "grip": "OPEN", "obj": ""}; _RESULT = {"v": ""} SEM = {"approach": "1. PRM APPROACH", "descend": "2. DESCEND onto block", "close": "3. CLOSE-GRASP", "carry": "4. LIFT + CARRY to target", "place": "5. PLACE (closed-loop)", "release": "6. RELEASE", "retreat": "7. RETREAT", "done": "DONE"} def capture(): img = env.render() if img is None: return im = Image.fromarray(np.asarray(img)[..., :3].copy()); d = ImageDraw.Draw(im) lines = ["=== STACK BLOCKS ON TARGET REGION ===" + (f" EP {args.episode}" if args.episode >= 0 else "")] if _RESULT["v"]: lines.append(f"RESULT: {_RESULT['v']}") lines += [f"ACTION: {SEM.get(_phase['v'], _phase['v'])}", f"block={_CUR['obj']} gripper={_CUR['grip']}", f"target=({TARGET[0]:+.2f},{TARGET[1]:+.2f}) size={args.target_size:.2f}m"] e = r_eef(); t = _CUR["tgt"] if t is not None: lines.append(f"eef(root) [{e[0]:+.2f} {e[1]:+.2f} {e[2]:+.2f}] err={np.linalg.norm(e-t):.3f}m") d.rectangle([0, 0, 430, 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)) _CORR = {"v": np.zeros(3, np.float32)}; _CMD = {"p": None} def _seg_start(): return _CMD["p"].copy() if _CMD["p"] is not None else r_eef().astype(np.float32) def _drive(cp, rg, corr): _CMD["p"] = np.asarray(cp, np.float32) env.step(act((cp+corr).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): return float(0.5-0.5*np.cos(np.pi*min(max(a, 0.0), 1.0))) def fillet(pts, r=0.06, n=6): 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 _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 flow(pts, rg, speed=0.008, settle=16, 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") 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))) cp = _point_at(poly, seglens, min(total, a*total)); _CUR["tgt"] = cp if on_step is not None: on_step(min(1.0, (k+1)/float(nsteps))) corr = _drive(cp, rg, corr) if k % 2 == 0: capture() return float(np.linalg.norm(poly[-1]-r_eef())) def go_converge(target, rg, tol=0.006, max_n=150): 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))); corr = _drive((1-a)*sp+a*tp, rg, corr) if k % 3 == 0: capture() if a >= 1.0 and np.linalg.norm(r_eef()-tp) < tol: break return float(np.linalg.norm(r_eef()-tp)) def go(target, rg, n): sp = _seg_start(); tp = np.asarray(target, 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)); corr = _drive((1-a)*sp+a*tp, rg, corr) if k % 3 == 0: capture() TABLE_ROOT_Z = float(TABLE_TOP-rroot[2]) obstacles = [yam_prm.Box(center=(np.array([-0.38, -0.15, 0.70])-rroot), half=np.array([0.02, 0.45, 0.28])), yam_prm.Box(center=(np.array([-0.10, -0.45, 0.70])-rroot), half=np.array([0.45, 0.02, 0.28]))] def pick_and_stack(name, layer): """layer 0 = on the target marker, layer 1 = on top of the block below.""" ext = obj_size(name) w = objw(name) o = Rq(rrootq).T@(np.array([w[0], w[1], TABLE_TOP+float(ext[2])/2.0], np.float32)-rroot) grasp = np.array([o[0], o[1], max(float(o[2])-0.005, TABLE_ROOT_Z+0.010)], np.float32) pre = np.array([o[0], o[1], float(o[2])+0.12], np.float32) _CUR["obj"] = name print(f"[st] --- block {layer}: {name} size={np.round(ext,3)} grasp={np.round(grasp,3)} ---", flush=True) _phase["v"] = "approach" prm = yam_prm.PRM(bounds_lo=np.array([-0.05, -0.35, -0.03]), bounds_hi=np.array([0.55, 0.45, 0.40]), obstacles=obstacles, clearance=0.05, num_samples=300, k=12, 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, obstacles, 0.05, iters=200, seed=2), 12).astype(np.float32) flow(list(wps[1:]), OPEN) _phase["v"] = "descend" print(f"[st] descend err={go_converge(grasp, OPEN, tol=0.008, max_n=140):.3f}", flush=True) for _ in range(10): _drive(grasp, OPEN, _CORR["v"]) _phase["v"] = "close" hold = r_eef().astype(np.float32) prev = fsep(); stall = 0 for k in range(160): _drive(hold, CLOSE, _CORR["v"]) if k % 6 == 0: capture() cur = fsep() stall = stall+1 if abs(cur-prev) < 0.0002 else 0 prev = cur if stall >= 8 and cur < -0.002: print(f"[st] jaws stalled at fsep={cur:.4f}", flush=True); break # Where this block must land: on the marker, or centred on top of the block below. stack_z = TABLE_TOP+0.002+float(ext[2])*(layer+0.5) place_local = np.array([TARGET[0], TARGET[1], stack_z], np.float32) place = (Rq(rrootq).T@(place_local-rroot)).astype(np.float32) above = place+np.array([0, 0, 0.12], np.float32) lift = hold+np.array([0, 0, 0.16], np.float32) z0 = objw(name)[2]; zmax = {"v": z0} def _step(frac): zmax["v"] = max(zmax["v"], objw(name)[2]) _phase["v"] = "carry" _phase["v"] = "carry" flow(fillet([_seg_start(), lift, above], r=0.06)[1:], CLOSE, on_step=_step) # closed-loop descent onto the stack: an open-loop drop misses by more than a block width _phase["v"] = "place" perr = go_converge(place, CLOSE, tol=0.005, max_n=170) print(f"[st] placed: target={np.round(place,3)} err={perr:.3f} peak_z={zmax['v']:.3f}", flush=True) _phase["v"] = "release"; go(place, OPEN, 40) _phase["v"] = "retreat"; flow([above], OPEN) return zmax["v"]-z0 > 0.04 for i, b in enumerate(BLOCKS): pick_and_stack(b, i) _phase["v"] = "retreat"; flow([rp_home], OPEN) for _ in range(60): env.step(act(rp_home, OPEN)); _freeze_left() # ---- success: both blocks on the target footprint, and the second one ON TOP of the first ---- half = args.target_size/2.0+0.02 pos = {b: objw(b) for b in BLOCKS} on_target = {b: (abs(pos[b][0]-TARGET[0]) < half and abs(pos[b][1]-TARGET[1]) < half) for b in BLOCKS} zs = sorted((pos[b][2], b) for b in BLOCKS) stacked = len(BLOCKS) < 2 or (zs[-1][0]-zs[0][0]) > 0.03 for b in BLOCKS: print(f"[st] {b}: final=({pos[b][0]:.3f},{pos[b][1]:.3f},{pos[b][2]:.3f}) on_target={on_target[b]}", flush=True) _RESULT["v"] = "SUCCESS" if (all(on_target.values()) and stacked) else "FAIL" _phase["v"] = "done" for _ in range(16): capture() print(f"[st] EPISODE_RESULT: {_RESULT['v']} on_target={sum(on_target.values())}/{len(BLOCKS)} " f"stacked={stacked} dz={zs[-1][0]-zs[0][0]:.3f}", flush=True) os.makedirs(os.path.dirname(args.video), exist_ok=True) # Drop the warm-up frames: before the renderer settles they come out with the wrong camera # pose, unresolved textures and missing geometry. if len(frames) > 6: frames = frames[2:] if frames: imageio.mimsave(args.video, frames, fps=14) print(f"[st] video -> {args.video} ({len(frames)} frames)", flush=True) env.close(); app.close(); print("YAM_STACK_OK", flush=True)