|
| 1 | +"""Interactive calibration grid recorder for DecRAS. |
| 2 | +
|
| 3 | +Connects to leader + follower SO-101 arms. The user teleoperates the arm to |
| 4 | +positions across the workspace, measures (x, y, z) with a ruler, and presses |
| 5 | +ENTER to record the joint angles together with the Cartesian position. |
| 6 | +
|
| 7 | +Saves to calibration/calibration_data.json. Supports resuming (appends to |
| 8 | +existing file) and displays workspace coverage summary after each point. |
| 9 | +
|
| 10 | +Usage: |
| 11 | + uv run python -m calibration.record_grid |
| 12 | +
|
| 13 | + # Custom ports: |
| 14 | + uv run python -m calibration.record_grid \\ |
| 15 | + --follower-port /dev/ttyACM0 \\ |
| 16 | + --leader-port /dev/ttyACM1 |
| 17 | +
|
| 18 | + # Custom output file: |
| 19 | + uv run python -m calibration.record_grid --output calibration/my_data.json |
| 20 | +""" |
| 21 | + |
| 22 | +import argparse |
| 23 | +import json |
| 24 | +import logging |
| 25 | +import os |
| 26 | +import sys |
| 27 | +import time |
| 28 | +from datetime import datetime, timezone |
| 29 | +from pathlib import Path |
| 30 | + |
| 31 | +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") |
| 32 | +logger = logging.getLogger(__name__) |
| 33 | + |
| 34 | +# Joint names matching the SO-101 arm (same as mcp_server/robot/lerobot.py) |
| 35 | +JOINT_NAMES = [ |
| 36 | + "shoulder_pan", |
| 37 | + "shoulder_lift", |
| 38 | + "elbow_flex", |
| 39 | + "wrist_flex", |
| 40 | + "wrist_roll", |
| 41 | + "gripper", |
| 42 | +] |
| 43 | + |
| 44 | +ARM_JOINT_NAMES = JOINT_NAMES[:-1] |
| 45 | + |
| 46 | +# Calibration grid spec from ARCHITECTURE.md |
| 47 | +GRID_SPEC = { |
| 48 | + "x_range": (0.15, 0.35), # meters from base |
| 49 | + "y_range": (-0.15, 0.15), |
| 50 | + "z_range": (0.00, 0.15), # table to +15cm |
| 51 | + "target_points": 75, |
| 52 | +} |
| 53 | + |
| 54 | +DEFAULT_OUTPUT = Path(__file__).parent / "calibration_data.json" |
| 55 | + |
| 56 | + |
| 57 | +def load_existing_data(path: Path) -> list[dict]: |
| 58 | + """Load existing calibration data for resume support.""" |
| 59 | + if path.exists(): |
| 60 | + with open(path) as f: |
| 61 | + data = json.load(f) |
| 62 | + return data.get("points", []) |
| 63 | + return [] |
| 64 | + |
| 65 | + |
| 66 | +def save_data(path: Path, points: list[dict]) -> None: |
| 67 | + """Save calibration data to JSON.""" |
| 68 | + path.parent.mkdir(parents=True, exist_ok=True) |
| 69 | + data = { |
| 70 | + "format_version": 1, |
| 71 | + "grid_spec": GRID_SPEC, |
| 72 | + "num_points": len(points), |
| 73 | + "last_updated": datetime.now(timezone.utc).isoformat(), |
| 74 | + "points": points, |
| 75 | + } |
| 76 | + with open(path, "w") as f: |
| 77 | + json.dump(data, f, indent=2) |
| 78 | + |
| 79 | + |
| 80 | +def compute_coverage(points: list[dict]) -> dict: |
| 81 | + """Compute workspace coverage stats from recorded points.""" |
| 82 | + if not points: |
| 83 | + return {"count": 0, "x_range": None, "y_range": None, "z_range": None} |
| 84 | + |
| 85 | + xs = [p["position"]["x"] for p in points] |
| 86 | + ys = [p["position"]["y"] for p in points] |
| 87 | + zs = [p["position"]["z"] for p in points] |
| 88 | + |
| 89 | + return { |
| 90 | + "count": len(points), |
| 91 | + "x_range": (min(xs), max(xs)), |
| 92 | + "y_range": (min(ys), max(ys)), |
| 93 | + "z_range": (min(zs), max(zs)), |
| 94 | + } |
| 95 | + |
| 96 | + |
| 97 | +def print_coverage(points: list[dict]) -> None: |
| 98 | + """Display workspace coverage summary.""" |
| 99 | + cov = compute_coverage(points) |
| 100 | + print(f"\n Recorded points: {cov['count']} / {GRID_SPEC['target_points']}") |
| 101 | + if cov["count"] > 0: |
| 102 | + print(f" X coverage: {cov['x_range'][0]:.3f} – {cov['x_range'][1]:.3f} m " |
| 103 | + f"(target: {GRID_SPEC['x_range'][0]:.2f} – {GRID_SPEC['x_range'][1]:.2f})") |
| 104 | + print(f" Y coverage: {cov['y_range'][0]:.3f} – {cov['y_range'][1]:.3f} m " |
| 105 | + f"(target: {GRID_SPEC['y_range'][0]:.2f} – {GRID_SPEC['y_range'][1]:.2f})") |
| 106 | + print(f" Z coverage: {cov['z_range'][0]:.3f} – {cov['z_range'][1]:.3f} m " |
| 107 | + f"(target: {GRID_SPEC['z_range'][0]:.2f} – {GRID_SPEC['z_range'][1]:.2f})") |
| 108 | + print() |
| 109 | + |
| 110 | + |
| 111 | +def parse_position(prompt: str) -> tuple[float, float, float] | None: |
| 112 | + """Prompt user for x, y, z position. Returns None on empty/quit input.""" |
| 113 | + try: |
| 114 | + raw = input(prompt).strip() |
| 115 | + except EOFError: |
| 116 | + return None |
| 117 | + |
| 118 | + if not raw or raw.lower() in ("q", "quit", "exit"): |
| 119 | + return None |
| 120 | + |
| 121 | + parts = raw.replace(",", " ").split() |
| 122 | + if len(parts) != 3: |
| 123 | + print(" Expected 3 values: x y z (meters). Example: 0.25 0.05 0.10") |
| 124 | + return parse_position(prompt) |
| 125 | + |
| 126 | + try: |
| 127 | + x, y, z = float(parts[0]), float(parts[1]), float(parts[2]) |
| 128 | + except ValueError: |
| 129 | + print(" Could not parse numbers. Example: 0.25 0.05 0.10") |
| 130 | + return parse_position(prompt) |
| 131 | + |
| 132 | + return (x, y, z) |
| 133 | + |
| 134 | + |
| 135 | +def read_joint_angles(robot) -> dict[str, float]: |
| 136 | + """Read current joint angles from the follower arm.""" |
| 137 | + obs = robot.get_observation() |
| 138 | + positions = {} |
| 139 | + for name in JOINT_NAMES: |
| 140 | + key = f"{name}.pos" |
| 141 | + positions[name] = float(obs.get(key, 0.0)) |
| 142 | + return positions |
| 143 | + |
| 144 | + |
| 145 | +def connect_hardware(follower_port: str, leader_port: str): |
| 146 | + """Connect to leader and follower arms. Returns (follower, leader, processors).""" |
| 147 | + from lerobot.robots.so_follower import SOFollower |
| 148 | + from lerobot.robots.so_follower.config_so_follower import SOFollowerRobotConfig |
| 149 | + from lerobot.teleoperators.so_leader import SOLeader |
| 150 | + from lerobot.teleoperators.so_leader.config_so_leader import SOLeaderTeleopConfig |
| 151 | + from lerobot.processor import make_default_processors |
| 152 | + |
| 153 | + follower = SOFollower(SOFollowerRobotConfig( |
| 154 | + port=follower_port, |
| 155 | + id="decras_follower", |
| 156 | + use_degrees=True, |
| 157 | + )) |
| 158 | + leader = SOLeader(SOLeaderTeleopConfig( |
| 159 | + port=leader_port, |
| 160 | + id="decras_leader", |
| 161 | + use_degrees=True, |
| 162 | + )) |
| 163 | + |
| 164 | + follower.connect() |
| 165 | + leader.connect() |
| 166 | + |
| 167 | + teleop_action_proc, robot_action_proc, robot_obs_proc = make_default_processors() |
| 168 | + |
| 169 | + return follower, leader, teleop_action_proc, robot_action_proc, robot_obs_proc |
| 170 | + |
| 171 | + |
| 172 | +def teleop_step(follower, leader, teleop_action_proc, robot_action_proc, robot_obs_proc): |
| 173 | + """Execute one teleoperation step: read leader, command follower.""" |
| 174 | + obs = follower.get_observation() |
| 175 | + obs_processed = robot_obs_proc(obs) |
| 176 | + act = leader.get_action() |
| 177 | + act_processed = teleop_action_proc((act, obs)) |
| 178 | + robot_action = robot_action_proc((act_processed, obs)) |
| 179 | + follower.send_action(robot_action) |
| 180 | + |
| 181 | + |
| 182 | +def parse_args() -> argparse.Namespace: |
| 183 | + parser = argparse.ArgumentParser( |
| 184 | + description="Record calibration grid points for joint-space lookup" |
| 185 | + ) |
| 186 | + parser.add_argument( |
| 187 | + "--follower-port", |
| 188 | + default=os.environ.get("LEROBOT_FOLLOWER_PORT", "/dev/decras_follower"), |
| 189 | + help="Follower arm serial port", |
| 190 | + ) |
| 191 | + parser.add_argument( |
| 192 | + "--leader-port", |
| 193 | + default=os.environ.get("LEROBOT_LEADER_PORT", "/dev/decras_leader"), |
| 194 | + help="Leader arm serial port", |
| 195 | + ) |
| 196 | + parser.add_argument( |
| 197 | + "--output", "-o", |
| 198 | + type=Path, |
| 199 | + default=DEFAULT_OUTPUT, |
| 200 | + help=f"Output JSON file (default: {DEFAULT_OUTPUT})", |
| 201 | + ) |
| 202 | + parser.add_argument( |
| 203 | + "--fps", type=int, default=50, |
| 204 | + help="Teleoperation control frequency in Hz (default: 50)", |
| 205 | + ) |
| 206 | + return parser.parse_args() |
| 207 | + |
| 208 | + |
| 209 | +def main(): |
| 210 | + args = parse_args() |
| 211 | + |
| 212 | + # Load existing data for resume |
| 213 | + points = load_existing_data(args.output) |
| 214 | + if points: |
| 215 | + logger.info(f"Resuming: loaded {len(points)} existing points from {args.output}") |
| 216 | + else: |
| 217 | + logger.info(f"Starting fresh. Output: {args.output}") |
| 218 | + |
| 219 | + # Connect hardware |
| 220 | + logger.info("Connecting to leader + follower arms...") |
| 221 | + try: |
| 222 | + follower, leader, tap, rap, rop = connect_hardware( |
| 223 | + args.follower_port, args.leader_port |
| 224 | + ) |
| 225 | + except Exception as e: |
| 226 | + logger.error(f"Failed to connect: {e}") |
| 227 | + sys.exit(1) |
| 228 | + |
| 229 | + logger.info("Connected. Starting teleoperation loop.") |
| 230 | + print("=" * 60) |
| 231 | + print(" DecRAS — Calibration Grid Recorder") |
| 232 | + print("=" * 60) |
| 233 | + print() |
| 234 | + print(" Move the arm using the leader, then press ENTER to record.") |
| 235 | + print(" You will be prompted to enter the measured x, y, z position.") |
| 236 | + print(" Type 'q' or Ctrl+C to stop and save.") |
| 237 | + print(" Type 'd' + ENTER to delete the last recorded point.") |
| 238 | + print() |
| 239 | + print_coverage(points) |
| 240 | + |
| 241 | + # Teleoperation runs in background; recording happens on ENTER |
| 242 | + import threading |
| 243 | + import select |
| 244 | + |
| 245 | + stop_event = threading.Event() |
| 246 | + |
| 247 | + def _teleop_loop(): |
| 248 | + """Run teleoperation at target FPS until stopped.""" |
| 249 | + dt = 1.0 / args.fps |
| 250 | + while not stop_event.is_set(): |
| 251 | + t0 = time.perf_counter() |
| 252 | + try: |
| 253 | + teleop_step(follower, leader, tap, rap, rop) |
| 254 | + except Exception as e: |
| 255 | + logger.warning(f"Teleop step error: {e}") |
| 256 | + elapsed = time.perf_counter() - t0 |
| 257 | + remaining = dt - elapsed |
| 258 | + if remaining > 0: |
| 259 | + time.sleep(remaining) |
| 260 | + |
| 261 | + teleop_thread = threading.Thread(target=_teleop_loop, daemon=True) |
| 262 | + teleop_thread.start() |
| 263 | + |
| 264 | + try: |
| 265 | + point_num = len(points) + 1 |
| 266 | + while True: |
| 267 | + try: |
| 268 | + cmd = input(f" [{point_num}] Press ENTER to record (or 'q' to quit, 'd' to delete last): ").strip().lower() |
| 269 | + except EOFError: |
| 270 | + break |
| 271 | + |
| 272 | + if cmd in ("q", "quit", "exit"): |
| 273 | + break |
| 274 | + |
| 275 | + if cmd == "d": |
| 276 | + if points: |
| 277 | + removed = points.pop() |
| 278 | + save_data(args.output, points) |
| 279 | + print(f" Deleted point at ({removed['position']['x']:.3f}, " |
| 280 | + f"{removed['position']['y']:.3f}, {removed['position']['z']:.3f})") |
| 281 | + point_num = len(points) + 1 |
| 282 | + print_coverage(points) |
| 283 | + else: |
| 284 | + print(" No points to delete.") |
| 285 | + continue |
| 286 | + |
| 287 | + if cmd and cmd not in ("", "y", "yes"): |
| 288 | + print(f" Unknown command: '{cmd}'. Press ENTER to record or 'q' to quit.") |
| 289 | + continue |
| 290 | + |
| 291 | + # Read joint angles from follower |
| 292 | + joints = read_joint_angles(follower) |
| 293 | + arm_joints = {k: v for k, v in joints.items() if k != "gripper"} |
| 294 | + |
| 295 | + print(f" Joint angles: {', '.join(f'{k}={v:.1f}' for k, v in arm_joints.items())}") |
| 296 | + |
| 297 | + # Get position measurement from user |
| 298 | + pos = parse_position(" Enter position (x y z in meters): ") |
| 299 | + if pos is None: |
| 300 | + break |
| 301 | + |
| 302 | + x, y, z = pos |
| 303 | + |
| 304 | + point = { |
| 305 | + "index": point_num, |
| 306 | + "position": {"x": x, "y": y, "z": z}, |
| 307 | + "joints": joints, |
| 308 | + "timestamp": datetime.now(timezone.utc).isoformat(), |
| 309 | + } |
| 310 | + points.append(point) |
| 311 | + save_data(args.output, points) |
| 312 | + |
| 313 | + print(f" Saved point {point_num}: pos=({x:.3f}, {y:.3f}, {z:.3f})") |
| 314 | + print_coverage(points) |
| 315 | + point_num += 1 |
| 316 | + |
| 317 | + except KeyboardInterrupt: |
| 318 | + print("\n Interrupted.") |
| 319 | + |
| 320 | + finally: |
| 321 | + stop_event.set() |
| 322 | + teleop_thread.join(timeout=2.0) |
| 323 | + save_data(args.output, points) |
| 324 | + logger.info(f"Saved {len(points)} points to {args.output}") |
| 325 | + |
| 326 | + try: |
| 327 | + follower.disconnect() |
| 328 | + except Exception: |
| 329 | + pass |
| 330 | + try: |
| 331 | + leader.disconnect() |
| 332 | + except Exception: |
| 333 | + pass |
| 334 | + |
| 335 | + print_coverage(points) |
| 336 | + print("Done.") |
| 337 | + |
| 338 | + |
| 339 | +if __name__ == "__main__": |
| 340 | + main() |
0 commit comments