#!/usr/bin/env python3 """ stack_demo.py — stack the four cubes into a standing tower (privileged demo). Builds a red → yellow → green → blue tower on the table using the same privileged weld-grasp regime as the v2 demo collector (scripts/collect_demos.py), made reliable with three closed-loop tricks the open-loop expert doesn't have: - **weld-only grip**: the palm hovers just above the cube and the weld does the holding — the fingers never close around it. A finger grasp released at tower height reliably flicks the cube off (the opening fingers sweep through it); with the hand clear, release is disturbance-free. - **cube servoing**: the weld makes the cube track the hand rigidly, so the placement loop drives the hand to (target − live hand→cube offset) until the cube centre is within 5 mm — a 4.4 cm cube only stays on the one below if the centres line up within ~2 cm. - **verify + retry**: after each release the cube's settled pose is checked against the tower and re-grabbed if it slipped (up to MAX_ATTEMPTS). This is a scripted, privileged-state demonstration (the same honesty class as the demo collector), NOT the learned v2 policy — see README_v2.md for why tight stacking is beyond the current learned checkpoint's precision. python scripts/stack_demo.py # headless, prints the report python scripts/stack_demo.py --render # MuJoCo viewer, real-time paced """ from __future__ import annotations import argparse import sys import time from pathlib import Path import numpy as np import mujoco sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from env.rx1_env import RX1Env # noqa: E402 from models.rx1_ik import RX1IKSolver, CUBE_SIZE # noqa: E402 ROBOT_CFG = { "simulation": {"physics_timestep": 0.002, "n_substeps": 4}, "vision": {"enabled": False}, # privileged exact state; no GL needed "actuator": {"max_joint_speed": 3.0, "control_clip": 2.5}, } # The tower is built directly on the red cube: its x=0.31 column is well # inside the arm's reliable envelope (table-height placements much past # x≈0.35 don't converge — the servo stalls at several cm). Picks run # farthest-from-red first so the hand never closes next to the tower. BASE = "red_cube" ORDER = ["yellow_cube", "green_cube", "blue_cube"] # bottom level → top GRIP_DZ = 0.025 # palm hover above the cube centre at weld time APPROACH_STEPS = 60 LOWER_STEPS = 50 WELD_STEPS = 5 LIFT_STEPS = 40 CARRY_BUDGET = 200 # coarse servo to 10 cm above the place target MAX_CORRECT_STEPS = 500 # fine servo budget per placement CORRECT_TOL = 0.005 # cube-centre-to-target distance to accept (m) RELEASE_STEPS = 2 MICRO_LIFT_STEPS = 12 # straight-up escape before the clearing retreat MICRO_LIFT_M = 0.05 RETREAT_STEPS = 80 SETTLE_STEPS = 120 MAX_ATTEMPTS = 3 # re-grab and re-place a cube that didn't stay put # A cube rests on the one below when centres are ~one edge apart vertically # and overlap at least half an edge horizontally. XY_TOL = CUBE_SIZE / 2 Z_TOL = CUBE_SIZE / 2 class _Viewer: """No-op-able wrapper so the control loops read the same either way.""" def __init__(self, env, render, pace): self._pace = pace if render else 0.0 self._v = None if render: from mujoco import viewer as mj_viewer self._v = mj_viewer.launch_passive(env.model, env.data) def tick(self): if self._v is not None: self._v.sync() if self._pace: time.sleep(self._pace) def hold_open(self, seconds): if self._v is None: return t0 = time.time() while self._v.is_running() and time.time() - t0 < seconds: self._v.sync() time.sleep(0.05) def close(self): if self._v is not None: self._v.close() def _apply(env, ctrl_nu, weld_bit, weld_on, view): full = np.zeros(env.N_ACTUATORS, dtype=np.float32) full[:env._nu] = ctrl_nu full[weld_bit] = 1.0 if weld_on else 0.0 env.step(full) view.tick() def _reach(env, ik, site_goal, steps, bit, weld_on, view): """Coarse open-loop reach of the grasp site toward a fixed goal.""" for _ in range(steps): q = ik.solve(site_goal, env.data.qpos[:env._nu].copy(), side="right") _apply(env, ik.set_fingers(q, closed=False, side="right"), bit, weld_on, view) def _servo_cube(env, ik, cube, goal, tol, budget, bit, view): """Closed-loop: drive the WELDED cube's centre onto `goal`. The weld makes the cube track the hand rigidly, so command the hand each step to (goal − live hand→cube offset). Returns the final error.""" err = float("inf") for _ in range(budget): live = env.perceive_objects()[cube]["pos"] err = float(np.linalg.norm(live - goal)) if err < tol: break offset = live - env.data.site_xpos[ik._rsite] q = ik.solve(goal - offset, env.data.qpos[:env._nu].copy(), side="right") _apply(env, ik.set_fingers(q, closed=False, side="right"), bit, True, view) return err def stack_one(env, ik, cube, place, view): """Pick `cube` with the privileged weld gripper and servo it onto `place`. The grip is weld-only with the palm held GRIP_DZ above the cube — the fingers never close around it. A finger-grasp release at tower height reliably flicks the cube off (the opening fingers sweep through it), and since the weld is doing the real holding anyway (exactly as in collect_demos.py), keeping the hand clear makes the release disturbance- free. Returns (settled error to place, servo error before release).""" nu = env._nu bit = env._weld_bit[cube] pick0 = env.perceive_objects()[cube]["pos"].copy() # 1. approach above the cube, then hover the palm just over its top face _reach(env, ik, pick0 + [0, 0, 0.12], APPROACH_STEPS, bit, False, view) for _ in range(LOWER_STEPS): live = env.perceive_objects()[cube]["pos"] q = ik.solve(live + [0, 0, GRIP_DZ], env.data.qpos[:nu].copy(), side="right") _apply(env, ik.set_fingers(q, closed=False, side="right"), bit, False, view) # 2. weld on (captures the current palm→cube offset), lift, carry high for _ in range(WELD_STEPS): live = env.perceive_objects()[cube]["pos"] q = ik.solve(live + [0, 0, GRIP_DZ], env.data.qpos[:nu].copy(), side="right") _apply(env, ik.set_fingers(q, closed=False, side="right"), bit, True, view) _reach(env, ik, pick0 + [0, 0, 0.20], LIFT_STEPS, bit, True, view) _servo_cube(env, ik, cube, place + [0, 0, 0.10], 0.02, CARRY_BUDGET, bit, view) # 3. fine placement: servo the cube flush onto the target err = _servo_cube(env, ik, cube, place + [0, 0, 0.001], CORRECT_TOL, MAX_CORRECT_STEPS, bit, view) # 4. release = weld off; nothing touches the cube, so just lift away hold = env.data.ctrl[:nu].copy() for _ in range(RELEASE_STEPS): _apply(env, hold, bit, False, view) lift_from = env.data.site_xpos[ik._rsite].copy() for k in range(MICRO_LIFT_STEPS): goal = lift_from + np.array([0.0, 0.0, MICRO_LIFT_M * (k + 1) / MICRO_LIFT_STEPS]) q = ik.solve(goal, env.data.qpos[:nu].copy(), side="right") _apply(env, ik.set_fingers(q, closed=False, side="right"), bit, False, view) _reach(env, ik, lift_from + [0, 0, 0.15], RETREAT_STEPS, bit, False, view) for _ in range(SETTLE_STEPS): _apply(env, env.data.ctrl[:nu].copy(), bit, False, view) final = env.perceive_objects()[cube]["pos"] return float(np.linalg.norm(final - place)), err def verify_tower(env, names): """Check each cube rests on the one below; return (ok, report_lines).""" pos = {n: env.perceive_objects()[n]["pos"] for n in names} lines, ok = [], True for below, above in zip(names, names[1:]): dxy = float(np.linalg.norm(pos[above][:2] - pos[below][:2])) dz = float(pos[above][2] - pos[below][2]) level_ok = dxy <= XY_TOL and abs(dz - CUBE_SIZE) <= Z_TOL ok &= level_ok lines.append(f" {above:<12} on {below:<12} xy-offset={dxy * 100:5.1f} cm " f"dz={dz * 100:5.1f} cm {'OK' if level_ok else 'OFF'}") return ok, lines def main(): p = argparse.ArgumentParser(description="RX1 expert cube-stacking demo") p.add_argument("--render", action="store_true", help="open the MuJoCo viewer") p.add_argument("--pace", type=float, default=0.008, help="seconds slept per control step when rendering") args = p.parse_args() env = RX1Env(ROBOT_CFG) ik = RX1IKSolver(env.model) mujoco.mj_forward(env.model, env.data) env.data.ctrl[:env._nu] = env.data.qpos[:env._nu] view = _Viewer(env, args.render, args.pace) try: placed = [BASE] for level, cube in enumerate(ORDER, start=1): for attempt in range(1, MAX_ATTEMPTS + 1): top = env.perceive_objects()[placed[-1]]["pos"] place = np.array([top[0], top[1], top[2] + CUBE_SIZE]) print(f"\n▸ level {level}: {cube} → {np.round(place, 3).tolist()}" f"{f' (attempt {attempt})' if attempt > 1 else ''}") final_err, servo_err = stack_one(env, ik, cube, place, view) print(f" servo err {servo_err * 100:.1f} cm → " f"settled err {final_err * 100:.1f} cm") settled = env.perceive_objects()[cube]["pos"] if (np.linalg.norm(settled[:2] - place[:2]) <= XY_TOL and abs(settled[2] - place[2]) <= Z_TOL): break print(" cube did not stay on the tower — re-grabbing") placed.append(cube) ok, lines = verify_tower(env, placed) print(f"\n{'=' * 56}\nTower ({' → '.join(t.split('_')[0] for t in placed)}):") print("\n".join(lines)) print(f"\nRESULT: {'STACK COMPLETE — tower standing' if ok else 'STACK INCOMPLETE'}") if ok: view.hold_open(30) return 0 if ok else 1 finally: view.close() if __name__ == "__main__": sys.exit(main())