"""YAM contact-rich / tool-use tasks -- the non-pick-and-place skills. --task push : slide a block to a target region WITHOUT grasping it (closed jaw as a finger) --task wipe : grasp the brush and wipe a target region along a lawn-mower path --task sweep : grasp the brush and sweep loose beads into a target zone --task pour : grasp a cup of beads, carry it over a bowl and ROTATE the wrist to pour These need a trajectory, not just waypoints: the end-effector follows a parametric path, and `pour` also drives the wrist orientation (slerp from top-down to tipped) while it moves. python scripts/yam_tool_tasks.py --headless --task push --video outputs/tasks/push.mp4 """ import argparse, sys, os from isaaclab.app import AppLauncher parser = argparse.ArgumentParser() parser.add_argument("--task", default="push", choices=["push", "wipe", "sweep", "pour"]) parser.add_argument("--obj", default="cube_r", help="pushed block (push) / cup (pour)") parser.add_argument("--tool", default="rw_brush", help="tool object for wipe/sweep") parser.add_argument("--obj_xy", default="0.02,0.12", help="object start x,y (env-local)") parser.add_argument("--target_xy", default="0.02,-0.10", help="target region centre x,y") parser.add_argument("--target_size", type=float, default=0.10) parser.add_argument("--beads", type=int, default=6, help="beads used by sweep/pour") parser.add_argument("--pour_deg", type=float, default=105.0, help="wrist tilt angle for pouring") parser.add_argument("--episode", type=int, default=-1) parser.add_argument("--video", default="outputs/tasks/yam_tool.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("[tt] 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 def slerp(q0, q1, t): q0 = q0/(np.linalg.norm(q0)+1e-9); q1 = q1/(np.linalg.norm(q1)+1e-9) d = float(np.dot(q0, q1)) if d < 0: q1 = -q1; d = -d if d > 0.9995: q = q0+t*(q1-q0); return q/(np.linalg.norm(q)+1e-9) th = np.arccos(d); q2 = q1-q0*d; q2 /= (np.linalg.norm(q2)+1e-9) return q0*np.cos(th*t)+q2*np.sin(th*t) 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 TXY = [float(v) for v in args.target_xy.split(",")] TARGET = np.array([TXY[0], TXY[1], TABLE_TOP], np.float32) OXY = [float(v) for v in args.obj_xy.split(",")] import isaaclab.sim as sim_utils # target region marker (paper-thin, so it never blocks the tool) _m = sim_utils.CuboidCfg(size=(args.target_size, args.target_size, 0.0015), visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(0.20, 0.65, 0.30))) _m.func("/World/envs/env_0/target_region", _m, translation=tuple((origin+TARGET+np.array([0, 0, 0.0008], np.float32)).astype(float).tolist())) print(f"[tt] target region at {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) # top-down grasp: jaw closes along world y; the tipped pour pose rotates about the tool's x axis GQ = qR(np.stack([np.array([0., 1., 0.]), np.array([1., 0., 0.]), np.array([0., 0., -1.])], axis=1)) def tipped_quat(deg): """Top-down orientation rolled by `deg` about world x -- the pouring motion.""" a = np.radians(deg); ca, sa = np.cos(a), np.sin(a) Rx = np.array([[1, 0, 0], [0, ca, -sa], [0, sa, ca]]) return qR(Rx@np.stack([np.array([0., 1., 0.]), np.array([1., 0., 0.]), np.array([0., 0., -1.])], axis=1)) 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) def place(name, xy, z=0.55, quat=(1, 0, 0, 0)): if name not in u.scene.rigid_objects: print(f"[tt] MISSING {name}", flush=True); return False ro = u.scene.rigid_objects[name] ro.write_root_pose_to_sim(torch.tensor(np.concatenate([origin+np.array([xy[0], xy[1], z]), quat]), dtype=torch.float32, device=dev).view(1, 7)) ro.write_root_velocity_to_sim(torch.zeros((1, 6), device=dev)) return True BEADS = [f"bead{i}" for i in range(args.beads)] TOOLNAME = args.tool if args.task in ("wipe", "sweep") else None if args.task == "push": place(args.obj, OXY) elif args.task == "pour": place(args.obj, OXY) elif args.task in ("wipe", "sweep"): place(TOOLNAME, OXY) for _ in range(60): env.step(act(rp_home, GQ, OPEN)) 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 size_of(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: return np.array([0.05, 0.05, 0.05]) # beads go INSIDE the cup for pour, or scattered on the table for sweep if args.task == "pour": cs = size_of(args.obj) for i, b in enumerate(BEADS): ang = 2*np.pi*i/max(len(BEADS), 1) place(b, [OXY[0]+0.008*np.cos(ang), OXY[1]+0.008*np.sin(ang)], z=TABLE_TOP+float(cs[2])*0.55+0.012*(i//3)) elif args.task == "sweep": for i, b in enumerate(BEADS): place(b, [OXY[0]-0.06+0.028*(i % 3), OXY[1]-0.10-0.03*(i//3)], z=TABLE_TOP+0.012) for _ in range(80): env.step(act(rp_home, GQ, 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"[tt] friction failed {tag}:", e, flush=True) _boost(R.root_physx_view, "right") for n in ([args.obj] if args.task in ("push", "pour") else []) + ([TOOLNAME] if TOOLNAME else []): if n in u.scene.rigid_objects: _boost(u.scene.rigid_objects[n].root_physx_view, n) 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 frames = []; _phase = {"v": "start"}; _CUR = {"tgt": None, "grip": "OPEN"}; _RESULT = {"v": ""} TITLE = {"push": "PUSH BLOCK TO TARGET (no grasp)", "wipe": "WIPE TARGET REGION WITH BRUSH", "sweep": "SWEEP DEBRIS INTO TARGET ZONE", "pour": "POUR CUP INTO BOWL (wrist rotation)"}[args.task] def capture(): img = env.render() if img is None: return im = Image.fromarray(np.asarray(img)[..., :3].copy()); d = ImageDraw.Draw(im) lines = [f"=== {TITLE} ===" + (f" EP {args.episode}" if args.episode >= 0 else "")] if _RESULT["v"]: lines.append(f"RESULT: {_RESULT['v']}") lines += [f"ACTION: {_phase['v']}", f"gripper={_CUR['grip']}"] 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, "q": GQ.copy()} def _seg_start(): return _CMD["p"].copy() if _CMD["p"] is not None else r_eef().astype(np.float32) def _drive(cp, cq, rg, corr): _CMD["p"] = np.asarray(cp, np.float32); _CMD["q"] = np.asarray(cq, np.float32) env.step(act((cp+corr).astype(np.float32), np.asarray(cq, 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.07) _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 _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=14, quat_fn=None, ease=True): """Glide along a polyline; quat_fn(frac) drives the wrist orientation along the way.""" 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 if ease else 1.0)), 1); corr = _CORR["v"] for k in range(nsteps+settle): f = min(1.0, (k+1)/float(nsteps)) a = _ease(f) if ease else f cp = _point_at(poly, seglens, min(total, a*total)); _CUR["tgt"] = cp cq = quat_fn(f) if quat_fn is not None else _CMD["q"] corr = _drive(cp, cq, rg, corr) if k % 2 == 0: capture() return float(np.linalg.norm(poly[-1]-r_eef())) def go_converge(target, rg, tol=0.008, max_n=140, quat=None): sp = _seg_start(); tp = np.asarray(target, np.float32) _CUR["tgt"] = tp; _CUR["grip"] = ("CLOSE" if rg < 0 else "OPEN") q = _CMD["q"] if quat is None else np.asarray(quat, np.float32) 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, q, 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 hold(n, rg): corr = _CORR["v"] for k in range(n): corr = _drive(_CMD["p"], _CMD["q"], rg, corr) if k % 3 == 0: capture() def to_root(w): return (Rq(rrootq).T@(np.asarray(w, np.float32)-rroot)).astype(np.float32) 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 prm_to(goal, rg): 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), np.asarray(goal, np.float64)) if p is None: p = np.stack([r_eef(), goal]).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:]), rg) def grasp_object(name, top_offset=None): """Standard closed-loop top-down grasp, reused by the tool tasks.""" ext = size_of(name); w = objw(name) # Grasp height comes from the object's LIVE pose, not from the authored bbox: the bbox is the # pre-scale, pre-settle one, and for the brush it put the grasp 27 cm above the table, so the # jaw closed on air and the whole sweep ran with an empty hand. gz = float(w[2]) if top_offset is None else float(w[2])+max(float(ext[2])/2.0-top_offset, 0.0) o = to_root(np.array([w[0], w[1], gz], np.float32)) print(f"[tt] {name} live pose=({w[0]:.3f},{w[1]:.3f},{w[2]:.3f}) -> grasp z={gz:.3f}", flush=True) grasp = np.array([o[0], o[1], max(float(o[2]), TABLE_ROOT_Z+0.010)], np.float32) _phase["v"] = "APPROACH"; prm_to(grasp+np.array([0, 0, 0.12], np.float32), OPEN) _phase["v"] = "DESCEND"; print(f"[tt] descend err={go_converge(grasp, OPEN):.3f}", flush=True) _phase["v"] = "CLOSE-GRASP" h = r_eef().astype(np.float32); prev = fsep(); stall = 0 for k in range(160): _drive(h, _CMD["q"], 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"[tt] jaws stalled fsep={cur:.4f}", flush=True); break # POST-GRASP GATE: lift a little and check the object actually comes along. Without this the # script happily performs the entire sweep/pour with an empty hand and only reports failure at # the very end, which hides WHERE it went wrong. z_before = objw(name)[2] corr = _CORR["v"] for _ in range(45): corr = _drive(h+np.array([0, 0, 0.05], np.float32), _CMD["q"], CLOSE, corr) z_after = objw(name)[2] if z_after-z_before < 0.02: print(f"[tt] GRASP FAILED: {name} did not rise with the gripper " f"(z {z_before:.3f} -> {z_after:.3f}, fsep={fsep():.4f}). " f"The jaw closed beside/above the object rather than around it.", flush=True) _RESULT["v"] = "FAIL (no grasp)" return None, ext print(f"[tt] grasp confirmed: {name} rose {z_after-z_before:+.3f} m with the gripper", flush=True) return h, ext # =============================== the tasks ================================================ if args.task == "push": # No grasp at all: close the empty jaw and use it as a finger to slide the block. ext = size_of(args.obj); w0 = objw(args.obj) push_z = TABLE_TOP+float(ext[2])*0.45 d = np.array([TARGET[0]-w0[0], TARGET[1]-w0[1]], np.float32) n = d/(np.linalg.norm(d)+1e-9) behind = np.array([w0[0]-n[0]*(float(ext[0])/2+0.045), w0[1]-n[1]*(float(ext[1])/2+0.045), push_z], np.float32) contact = np.array([w0[0]-n[0]*(float(ext[0])/2+0.012), w0[1]-n[1]*(float(ext[1])/2+0.012), push_z], np.float32) end = np.array([TARGET[0]-n[0]*(float(ext[0])/2+0.012), TARGET[1]-n[1]*(float(ext[1])/2+0.012), push_z], np.float32) print(f"[tt] push {args.obj} from ({w0[0]:.3f},{w0[1]:.3f}) to target ({TARGET[0]:.3f},{TARGET[1]:.3f})", flush=True) _phase["v"] = "CLOSE JAW (finger)" for _ in range(40): _drive(_seg_start(), GQ, CLOSE, _CORR["v"]); capture() _phase["v"] = "APPROACH behind block"; prm_to(to_root(behind+np.array([0, 0, 0.10], np.float32)), CLOSE) _phase["v"] = "LOWER to push height"; go_converge(to_root(behind), CLOSE, tol=0.008, max_n=130) _phase["v"] = "MAKE CONTACT"; go_converge(to_root(contact), CLOSE, tol=0.006, max_n=90) _phase["v"] = "PUSH (closed-loop, re-centred on the block)" # Open-loop pushing fails the moment the block skids off the contact normal: the pusher keeps # driving down its planned line while the block squirts off sideways. Instead, re-derive the # pusher pose from the block's LIVE position every step -- stay on the line from the block to # the target, a fixed gap behind the block, and advance the goal gradually. steps = 220 corr = _CORR["v"] for k in range(steps): w = objw(args.obj) to_t = np.array([TARGET[0]-w[0], TARGET[1]-w[1]], np.float32) dist = float(np.linalg.norm(to_t)) if dist < 0.015: break nn = to_t/(dist+1e-9) back = float(ext[0])/2+0.012 # advance a little past the block each step so contact is maintained, not just touched adv = min(0.02, dist) cp_w = np.array([w[0]-nn[0]*back+nn[0]*adv, w[1]-nn[1]*back+nn[1]*adv, push_z], np.float32) cp = to_root(cp_w); _CUR["tgt"] = cp corr = _drive(cp, GQ, CLOSE, corr) if k % 2 == 0: capture() print(f"[tt] push servo done: block=({objw(args.obj)[0]:.3f},{objw(args.obj)[1]:.3f}) " f"target=({TARGET[0]:.3f},{TARGET[1]:.3f})", flush=True) hold(30, CLOSE) _phase["v"] = "RETREAT"; flow([to_root(end+np.array([0, 0, 0.14], np.float32))], CLOSE) wf = objw(args.obj) err = float(np.hypot(wf[0]-TARGET[0], wf[1]-TARGET[1])) _RESULT["v"] = "SUCCESS" if err < args.target_size/2+0.02 else "FAIL" print(f"[tt] EPISODE_RESULT: {_RESULT['v']} block=({wf[0]:.3f},{wf[1]:.3f}) target=({TARGET[0]:.3f},{TARGET[1]:.3f}) err={err:.3f}", flush=True) elif args.task in ("wipe", "sweep"): # No top_offset: the brush lies FLAT, so its live centre height is the grasp height. Deriving # it from the authored bbox (which reports an 18.8 cm tall brush) aims 6 cm above the handle. hold_pose, ext = grasp_object(TOOLNAME) if hold_pose is None: raise SystemExit("[tt] aborting: the brush was never actually grasped") tip_drop = float(ext[2])*0.55 # how far the tool tip reaches below the grip point work_z = TABLE_TOP+tip_drop+0.004 _phase["v"] = "LIFT tool" flow([hold_pose+np.array([0, 0, 0.12], np.float32)], CLOSE) if args.task == "wipe": # lawn-mower path over the target square, keeping the tool on the surface half = args.target_size/2 lanes = [] for i, yy in enumerate(np.linspace(-half, half, 4)): xs = [-half, half] if i % 2 == 0 else [half, -half] lanes += [np.array([TARGET[0]+xs[0], TARGET[1]+yy, work_z], np.float32), np.array([TARGET[0]+xs[1], TARGET[1]+yy, work_z], np.float32)] _phase["v"] = "MOVE over region" go_converge(to_root(lanes[0]+np.array([0, 0, 0.10], np.float32)), CLOSE, tol=0.01, max_n=130) _phase["v"] = "PRESS DOWN"; go_converge(to_root(lanes[0]), CLOSE, tol=0.008, max_n=110) _phase["v"] = "WIPE (lawn-mower path)" flow([to_root(p) for p in lanes[1:]], CLOSE, speed=0.006) covered = True else: # sweep strokes: drag the beads toward the target zone in three passes strokes = [] for i, xoff in enumerate((-0.03, 0.0, 0.03)): start = np.array([OXY[0]+xoff, OXY[1]-0.14, work_z], np.float32) end = np.array([TARGET[0]+xoff*0.5, TARGET[1]+0.02, work_z], np.float32) strokes.append((start, end)) _phase["v"] = "MOVE behind debris" go_converge(to_root(strokes[0][0]+np.array([0, 0, 0.10], np.float32)), CLOSE, tol=0.01, max_n=130) for i, (s, e) in enumerate(strokes): _phase["v"] = f"SWEEP stroke {i+1}/3" go_converge(to_root(s+np.array([0, 0, 0.06], np.float32)), CLOSE, tol=0.01, max_n=80) go_converge(to_root(s), CLOSE, tol=0.008, max_n=90) flow([to_root(e)], CLOSE, speed=0.005) flow([to_root(e+np.array([0, 0, 0.08], np.float32))], CLOSE) covered = True _phase["v"] = "LIFT tool away" flow([_seg_start()+np.array([0, 0, 0.12], np.float32)], CLOSE) if args.task == "sweep": inzone = sum(1 for b in BEADS if b in u.scene.rigid_objects and abs(objw(b)[0]-TARGET[0]) < args.target_size/2+0.03 and abs(objw(b)[1]-TARGET[1]) < args.target_size/2+0.03) _RESULT["v"] = "SUCCESS" if inzone >= max(1, len(BEADS)//2) else "FAIL" print(f"[tt] EPISODE_RESULT: {_RESULT['v']} beads_in_zone={inzone}/{len(BEADS)}", flush=True) else: e = np.linalg.norm(r_eef()-to_root(np.array([TARGET[0], TARGET[1], work_z], np.float32))) _RESULT["v"] = "SUCCESS" if covered else "FAIL" print(f"[tt] EPISODE_RESULT: {_RESULT['v']} wiped the {args.target_size:.2f} m region, " f"final tool err={e:.3f}", flush=True) elif args.task == "pour": hold_pose, ext = grasp_object(args.obj) z_before = [objw(b)[2] for b in BEADS if b in u.scene.rigid_objects] _phase["v"] = "LIFT cup" flow([hold_pose+np.array([0, 0, 0.16], np.float32)], CLOSE) over = to_root(np.array([TARGET[0], TARGET[1]+0.03, TABLE_TOP+0.20], np.float32)) _phase["v"] = "CARRY over bowl" flow([over], CLOSE) _phase["v"] = "POUR (rotate wrist)" tip = tipped_quat(args.pour_deg) n = 130 corr = _CORR["v"] for k in range(n): # rotate in place: position held, orientation slerped f = _ease((k+1)/float(n)) corr = _drive(_CMD["p"], slerp(GQ, tip, f), CLOSE, corr) if k % 2 == 0: capture() _phase["v"] = "HOLD (contents fall out)" hold(70, CLOSE) _phase["v"] = "RETURN upright" corr = _CORR["v"] for k in range(90): f = _ease((k+1)/90.0) corr = _drive(_CMD["p"], slerp(tip, GQ, f), CLOSE, corr) if k % 2 == 0: capture() _phase["v"] = "RETREAT" flow([over+np.array([0, 0, 0.06], np.float32)], CLOSE) poured = sum(1 for b in BEADS if b in u.scene.rigid_objects and abs(objw(b)[0]-TARGET[0]) < 0.12 and abs(objw(b)[1]-TARGET[1]) < 0.12 and objw(b)[2] < TABLE_TOP+0.10) _RESULT["v"] = "SUCCESS" if poured >= max(1, len(BEADS)//2) else "FAIL" print(f"[tt] EPISODE_RESULT: {_RESULT['v']} poured={poured}/{len(BEADS)} beads into the target area", flush=True) for _ in range(16): capture() 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"[tt] video -> {args.video} ({len(frames)} frames)", flush=True) env.close(); app.close(); print("YAM_TOOL_OK", flush=True)