From 8879b8ad6a0ef309a6babb10aeb3f406d306378b Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Fri, 10 Jul 2026 19:41:16 +0800 Subject: [PATCH 01/22] feat: initial lerobot env implementation Ported from jx-qiu/feature/lerobot (d3ee39c) onto the rpent/robots layout: physical_agent.* -> rpent.*, env package moved to robots/lerobot/ with env_client.py, MCP namespace mcp__rpent__. Signed-off-by: Jiaxing Qiu --- deployment/franka/auto_calibrate_cameras.py | 631 ++++++++ deployment/franka/calibrate_charuco_wrist.py | 448 ++++++ deployment/franka/calibration.py | 120 ++ deployment/franka/env_server.py | 954 ++++++++++++ deployment/franka/generate_fiducial_board.py | 190 +++ deployment/franka/run_env_server.sh | 41 + .../lerobot/auto_calibrate_scene_cam.py | 119 ++ deployment/lerobot/calibrate_scene_cam.py | 189 +++ deployment/lerobot/calibration.py | 78 + deployment/lerobot/diagnose_motors.py | 173 +++ deployment/lerobot/env_server.py | 1311 +++++++++++++++++ deployment/lerobot/geometry.py | 256 ++++ deployment/lerobot/kinematics.py | 100 ++ deployment/lerobot/scene_camera.py | 108 ++ .../franka_charuco_7x5_25mm.json | 18 + .../franka_charuco_7x5_25mm.pdf | Bin 0 -> 139890 bytes .../franka_charuco_7x5_25mm.png | Bin 0 -> 8396 bytes resources/lerobot/memory/MEMORY.md | 1 + .../memory/plate_place_low_standoff.md | 7 + robots/lerobot/__init__.py | 73 + robots/lerobot/env_client.py | 141 ++ robots/lerobot/prompt.py | 163 ++ robots/lerobot/toolkit.py | 118 ++ robots/lerobot/tools.py | 578 ++++++++ 24 files changed, 5817 insertions(+) create mode 100644 deployment/franka/auto_calibrate_cameras.py create mode 100644 deployment/franka/calibrate_charuco_wrist.py create mode 100644 deployment/franka/calibration.py create mode 100644 deployment/franka/env_server.py create mode 100644 deployment/franka/generate_fiducial_board.py create mode 100755 deployment/franka/run_env_server.sh create mode 100644 deployment/lerobot/auto_calibrate_scene_cam.py create mode 100644 deployment/lerobot/calibrate_scene_cam.py create mode 100644 deployment/lerobot/calibration.py create mode 100644 deployment/lerobot/diagnose_motors.py create mode 100644 deployment/lerobot/env_server.py create mode 100644 deployment/lerobot/geometry.py create mode 100644 deployment/lerobot/kinematics.py create mode 100644 deployment/lerobot/scene_camera.py create mode 100644 resources/franka/calibration_boards/franka_charuco_7x5_25mm.json create mode 100644 resources/franka/calibration_boards/franka_charuco_7x5_25mm.pdf create mode 100644 resources/franka/calibration_boards/franka_charuco_7x5_25mm.png create mode 100644 resources/lerobot/memory/MEMORY.md create mode 100644 resources/lerobot/memory/plate_place_low_standoff.md create mode 100644 robots/lerobot/__init__.py create mode 100644 robots/lerobot/env_client.py create mode 100644 robots/lerobot/prompt.py create mode 100644 robots/lerobot/toolkit.py create mode 100644 robots/lerobot/tools.py diff --git a/deployment/franka/auto_calibrate_cameras.py b/deployment/franka/auto_calibrate_cameras.py new file mode 100644 index 00000000..3c12a2ce --- /dev/null +++ b/deployment/franka/auto_calibrate_cameras.py @@ -0,0 +1,631 @@ +#!/usr/bin/env python3 +"""Automatic Franka RGB-D camera calibration helpers. + +The default mode calibrates the fixed scene camera into the Franka base frame +(``panda_link0``) without markers: + +1. move the TCP through a small, conservative 3-D grid, +2. at each pose, toggle the Franka Hand while the arm is stationary, +3. segment the moving fingers in the scene RGB image and use aligned depth to + get a camera-frame point, +4. pair that point with the live robot TCP position and fit ``T_base_cam`` with + RANSAC Kabsch, +5. save the calibration record to + ``~/.cache/physical_agent/franka/camera_calibration/.json`` and ask + the running env server to reload it. + +This is the right first calibration for ``back_project`` because the scene +camera is fixed. The wrist camera is eye-in-hand; calibrating it correctly needs +hand-eye calibration (``T_tcp_cam``) using a fixed fiducial/ChArUco/AprilTag or a +scene-calibrated reference target observed from multiple wrist poses. This file +keeps the wrist record format ready, but does not invent an unsafe automatic +wrist calibration from one moving camera alone. + +Run the env server first, then run in the physicalagent env:: + + python deployment/franka/auto_calibrate_cameras.py --port 5599 --yes + +Offline math check:: + + python deployment/franka/auto_calibrate_cameras.py --self-test +""" +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path +from typing import Any + +import imageio.v2 as imageio +import numpy as np + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from deployment.franka import calibration as franka_calib # noqa: E402 +from deployment.lerobot import geometry as geom # noqa: E402 +from physical_agent.rpc_driver.socket import SocketRpcClient # noqa: E402 + +_DEFAULT_GRID_X = (0.46, 0.54, 0.60) +_DEFAULT_GRID_Y = (-0.08, 0.0, 0.08) +_DEFAULT_GRID_Z = (0.20, 0.27) +_DEFAULT_DOWN_EULER = [float(np.pi), 0.0, 0.0] + + +def _self_test() -> int: + """Validate Kabsch/RANSAC and motion-blob detection offline.""" + rng = np.random.default_rng(3) + q, _ = np.linalg.qr(rng.standard_normal((3, 3))) + if np.linalg.det(q) < 0: + q[:, 0] = -q[:, 0] + T_true = np.eye(4) + T_true[:3, :3] = q + T_true[:3, 3] = rng.normal(size=3) + + cam = rng.normal(size=(12, 3)) + base = geom.transform_points(T_true, cam) + rng.normal(scale=0.001, size=(12, 3)) + base[5] += [0.12, -0.08, 0.05] + T_est, rmse, inliers = geom.ransac_kabsch(cam, base, thresh_m=0.015) + fit_ok = np.allclose(T_est, T_true, atol=2e-2) and not bool(inliers[5]) + print( + "ransac_kabsch: " + f"rmse={rmse:.5f}m inliers={int(inliers.sum())}/12 " + f"outlier_excluded={not bool(inliers[5])} -> {'OK' if fit_ok else 'FAIL'}" + ) + + H, W = 480, 640 + rgb_open = np.zeros((H, W, 3), np.uint8) + rgb_closed = rgb_open.copy() + rgb_closed[210:235, 330:365] = 255 + depth = np.full((H, W), 0.55, np.float32) + K = np.array([[600.0, 0.0, 320.0], [0.0, 600.0, 240.0], [0.0, 0.0, 1.0]]) + det = _detect_tip_pixel_by_motion(rgb_open, rgb_closed, depth, K) + det_ok = det is not None and abs(det["pixel"][0] - 222) < 4 and abs(det["pixel"][1] - 347) < 4 + print(f"detect_tip: {det if det else 'None'} -> {'OK' if det_ok else 'FAIL'}") + return 0 if fit_ok and det_ok else 1 + + +def _detect_motion_blob_numpy( + rgb_open, + rgb_closed, + depth_m, + K, + *, + diff_thresh: int, + min_area: int, + max_area: int, +) -> dict | None: + """Locate the largest changed depth-valid blob without OpenCV.""" + a = np.asarray(rgb_open, dtype=np.float32).mean(axis=2) + b = np.asarray(rgb_closed, dtype=np.float32).mean(axis=2) + depth_m = np.asarray(depth_m, dtype=np.float64) + mask = (np.abs(a - b) >= int(diff_thresh)) & np.isfinite(depth_m) & (depth_m > 0) + if not np.any(mask): + return None + + try: + from scipy import ndimage + + labels, num = ndimage.label(mask) + best = None + best_area = 0 + for label in range(1, num + 1): + comp = labels == label + area = int(comp.sum()) + if min_area <= area <= max_area and area > best_area: + best = comp + best_area = area + if best is None: + return None + rows, cols = np.nonzero(best) + depths = depth_m[best] + except Exception: + rows, cols = np.nonzero(mask) + depths = depth_m[mask] + best_area = int(rows.size) + if not (min_area <= best_area <= max_area): + return None + + row = float(np.median(rows)) + col = float(np.median(cols)) + z = float(np.median(depths)) + return { + "pixel": [row, col], + "depth_m": z, + "area": best_area, + "xyz_cam": geom.backproject_pixel(K, col, row, z).tolist(), + } + + +def _save_debug_images( + *, + debug_dir: Path, + pose_idx: int, + camera: str, + rgb_open, + rgb_closed, + depth_m, +) -> dict[str, str]: + """Save open/closed/diff/depth images for diagnosing failed detections.""" + debug_dir.mkdir(parents=True, exist_ok=True) + prefix = debug_dir / f"{pose_idx:02d}_{camera}" + rgb_open = np.asarray(rgb_open, dtype=np.uint8) + rgb_closed = np.asarray(rgb_closed, dtype=np.uint8) + depth_m = np.asarray(depth_m, dtype=np.float32) + + diff = np.abs(rgb_open.astype(np.int16) - rgb_closed.astype(np.int16)).max(axis=2) + diff_img = np.clip(diff, 0, 255).astype(np.uint8) + + valid = np.isfinite(depth_m) & (depth_m > 0) + depth_img = np.zeros(depth_m.shape, dtype=np.uint8) + if np.any(valid): + lo, hi = np.percentile(depth_m[valid], [2, 98]) + if hi > lo: + depth_img[valid] = np.clip((depth_m[valid] - lo) / (hi - lo) * 255, 0, 255) + + paths = { + "open": str(prefix.with_name(prefix.name + "_open.png")), + "closed": str(prefix.with_name(prefix.name + "_closed.png")), + "diff": str(prefix.with_name(prefix.name + "_diff.png")), + "depth": str(prefix.with_name(prefix.name + "_depth.png")), + } + imageio.imwrite(paths["open"], rgb_open) + imageio.imwrite(paths["closed"], rgb_closed) + imageio.imwrite(paths["diff"], diff_img) + imageio.imwrite(paths["depth"], depth_img) + return paths + + +def _manual_pixel_from_terminal( + *, + image_path: str, + depth_m, + K, + patch_radius: int, +) -> dict | None: + """Prompt for a manual pixel and backproject it, or return None to skip.""" + print(f" manual fallback: inspect {image_path}") + print(" enter pixel as row,col (for example 240,320), or press Enter to skip") + text = input(" row,col> ").strip() + if not text: + return None + try: + row_s, col_s = text.replace(" ", "").split(",", 1) + row = int(round(float(row_s))) + col = int(round(float(col_s))) + except Exception: + print(" invalid pixel format; skipped") + return None + + z = geom.sample_depth_patch(depth_m, col, row, radius=patch_radius) + if not np.isfinite(z) or z <= 0: + print(f" no valid depth near ({row},{col}); skipped") + return None + p_cam = geom.backproject_pixel(K, col, row, z) + return { + "pixel": [float(row), float(col)], + "depth_m": float(z), + "area": int((2 * patch_radius + 1) ** 2), + "xyz_cam": p_cam.tolist(), + "manual": True, + } + + +def _detect_tip_pixel_by_motion( + rgb_open, + rgb_closed, + depth_m, + K, + *, + diff_thresh: int = 18, + min_area: int = 40, + max_area: int = 40000, +) -> dict | None: + """Detect gripper motion, preferring OpenCV but falling back gracefully.""" + try: + return geom.detect_tip_pixel_by_motion( + rgb_open, + rgb_closed, + depth_m, + K, + diff_thresh=diff_thresh, + min_area=min_area, + max_area=max_area, + ) + except ModuleNotFoundError as exc: + if exc.name != "cv2": + raise + return _detect_motion_blob_numpy( + rgb_open, + rgb_closed, + depth_m, + K, + diff_thresh=diff_thresh, + min_area=min_area, + max_area=max_area, + ) + + +def _parse_csv_floats(text: str, *, expected: int, name: str) -> tuple[float, ...]: + try: + values = tuple(float(part.strip()) for part in text.split(",") if part.strip()) + except ValueError as exc: + raise argparse.ArgumentTypeError(f"{name} must be comma-separated floats") from exc + if len(values) != expected: + raise argparse.ArgumentTypeError(f"{name} expects {expected} values, got {len(values)}") + return values + + +def _candidate_poses(args: argparse.Namespace) -> list[list[float]]: + xs = args.grid_x + ys = args.grid_y + zs = args.grid_z + poses = [[float(x), float(y), float(z)] for z in zs for y in ys for x in xs] + # Visit the center-ish pose first, then spread out. This makes early aborts + # less likely to leave the robot at a corner of the grid. + center = np.array([np.mean(xs), np.mean(ys), np.mean(zs)], dtype=np.float64) + poses.sort(key=lambda p: float(np.linalg.norm(np.asarray(p) - center))) + return poses + + +def _obs_camera(obs: dict, camera: str) -> tuple[np.ndarray, np.ndarray, dict]: + frames = obs.get("frames") or {} + depths = obs.get("depth") or {} + metas = obs.get("camera_meta") or {} + if camera not in frames: + raise RuntimeError(f"camera {camera!r} missing from observation frames") + if camera not in depths: + raise RuntimeError(f"camera {camera!r} missing from observation depth maps") + if camera not in metas: + raise RuntimeError(f"camera {camera!r} missing from observation metadata") + return ( + np.asarray(frames[camera], dtype=np.uint8), + np.asarray(depths[camera], dtype=np.float32), + dict(metas[camera]), + ) + + +def _tcp_point_from_pose(ee: dict, tcp_offset: tuple[float, float, float]) -> np.ndarray: + xyz = np.asarray(ee["xyz"], dtype=np.float64) + offset = np.asarray(tcp_offset, dtype=np.float64) + if np.allclose(offset, 0.0): + return xyz + from scipy.spatial.transform import Rotation as R + + quat = ee.get("quat_xyzw") + if quat is None: + raise RuntimeError("--tcp-offset requires get_ee_pose to return quat_xyzw") + return xyz + R.from_quat(np.asarray(quat, dtype=np.float64)).as_matrix() @ offset + + +def _detect_correspondence( + *, + client: SocketRpcClient, + camera: str, + tcp_offset: tuple[float, float, float], + diff_thresh: int, + min_area: int, + max_area: int, + settle_s: float, + detect_retries: int, + manual_on_fail: bool, + manual_always: bool, + manual_patch_radius: int, + debug_dir: Path, + pose_idx: int, +) -> dict[str, Any]: + """Toggle gripper once and return one cam/base correspondence.""" + det = None + rgb_closed = None + depth = None + meta = None + debug_paths: dict[str, str] = {} + attempts = max(1, int(detect_retries)) + for attempt in range(1, attempts + 1): + client.call("env.open_gripper", timeout_s=30.0) + time.sleep(settle_s) + obs_open = client.call("env.get_obs", timeout_s=30.0) + rgb_open, _, _ = _obs_camera(obs_open, camera) + + client.call("env.close_gripper", timeout_s=30.0) + time.sleep(settle_s) + obs_closed = client.call("env.get_obs", timeout_s=30.0) + rgb_closed, depth, meta = _obs_camera(obs_closed, camera) + + debug_paths = _save_debug_images( + debug_dir=debug_dir, + pose_idx=pose_idx * 10 + attempt, + camera=camera, + rgb_open=rgb_open, + rgb_closed=rgb_closed, + depth_m=depth, + ) + + if not manual_always: + det = _detect_tip_pixel_by_motion( + rgb_open, + rgb_closed, + depth, + np.asarray(meta["K"], dtype=np.float64), + diff_thresh=diff_thresh, + min_area=min_area, + max_area=max_area, + ) + if det is not None or manual_always: + break + + ee = client.call("env.get_ee_pose", timeout_s=15.0) + client.call("env.open_gripper", timeout_s=30.0) + + if (det is None or manual_always) and manual_on_fail: + assert rgb_closed is not None and depth is not None and meta is not None + det = _manual_pixel_from_terminal( + image_path=debug_paths.get("closed", ""), + depth_m=depth, + K=np.asarray(meta["K"], dtype=np.float64), + patch_radius=manual_patch_radius, + ) + if det is None: + raise RuntimeError( + "could not segment gripper motion in camera image; debug images: " + + json.dumps(debug_paths) + ) + + return { + "xyz_cam": np.asarray(det["xyz_cam"], dtype=np.float64), + "xyz_base": _tcp_point_from_pose(ee, tcp_offset), + "pixel": det["pixel"], + "depth_m": float(det["depth_m"]), + "area": int(det["area"]), + "manual": bool(det.get("manual", False)), + "debug_paths": debug_paths, + "ee": ee, + "camera_meta": meta, + } + + +def _calibrate_scene(args: argparse.Namespace) -> int: + if not args.yes and not args.no_save: + print( + "Refusing to move the robot without --yes. This calibration drives " + "the TCP through a small 3-D grid and toggles the gripper." + ) + return 2 + + client = SocketRpcClient(args.host, args.port) + meta_all = client.call("env.get_camera_meta", timeout_s=15.0) + if args.camera not in meta_all: + print(f"camera {args.camera!r} not available; found {sorted(meta_all)}") + return 2 + camera_meta = meta_all[args.camera] + serial = args.serial or camera_meta.get("serial") + if not serial: + print("could not determine camera serial; pass --serial") + return 2 + + poses = _candidate_poses(args) + print( + f"Calibrating fixed camera {args.camera!r} serial={serial} with up to " + f"{len(poses)} candidate poses; target valid points={args.n_points}." + ) + print("Clear the workspace. The gripper will move and open/close at each pose.") + + cam_pts: list[np.ndarray] = [] + base_pts: list[np.ndarray] = [] + records: list[dict[str, Any]] = [] + + try: + for idx, xyz in enumerate(poses, start=1): + if len(cam_pts) >= args.n_points: + break + print(f"\n[{idx}/{len(poses)}] move_to {np.round(xyz, 3).tolist()}") + move = client.call( + "env.move_to", + args=(xyz,), + kwargs={"euler_xyz": _DEFAULT_DOWN_EULER, "gripper": "open"}, + timeout_s=120.0, + ) + print(f" move: reached={move.get('reached')} err={move.get('pos_error_m')} final={move.get('final_xyz')}") + if move.get("error"): + print(f" skipped: {move['error']}") + continue + + try: + corr = _detect_correspondence( + client=client, + camera=args.camera, + tcp_offset=args.tcp_offset, + diff_thresh=args.diff_thresh, + min_area=args.min_area, + max_area=args.max_area, + settle_s=args.settle_s, + detect_retries=args.detect_retries, + manual_on_fail=args.manual_on_fail, + manual_always=args.manual_always, + manual_patch_radius=args.manual_patch_radius, + debug_dir=Path(args.debug_dir), + pose_idx=idx, + ) + except Exception as exc: + print(f" detection failed: {exc}") + continue + + cam_pts.append(corr["xyz_cam"]) + base_pts.append(corr["xyz_base"]) + records.append( + { + "target_xyz": xyz, + "pixel": corr["pixel"], + "depth_m": round(corr["depth_m"], 4), + "area": corr["area"], + "manual": corr["manual"], + "debug_paths": corr["debug_paths"], + "xyz_cam": np.round(corr["xyz_cam"], 5).tolist(), + "xyz_base": np.round(corr["xyz_base"], 5).tolist(), + "move": move, + } + ) + print( + " captured: " + f"pixel={np.round(corr['pixel'], 1).tolist()} " + f"depth={corr['depth_m']:.3f}m area={corr['area']} " + f"base={np.round(corr['xyz_base'], 3).tolist()} " + f"manual={corr['manual']}" + ) + finally: + try: + client.call("env.open_gripper", timeout_s=30.0) + except Exception: + pass + + if len(cam_pts) < 4: + print(f"Calibration failed: need >=4 correspondences, got {len(cam_pts)}") + print(json.dumps(records, indent=2, default=str)) + return 2 + + cam_arr = np.asarray(cam_pts, dtype=np.float64) + base_arr = np.asarray(base_pts, dtype=np.float64) + T_base_cam, rmse, inliers = geom.ransac_kabsch( + cam_arr, + base_arr, + thresh_m=args.ransac_thresh_m, + iters=args.ransac_iters, + min_inliers=min(4, len(cam_pts)), + seed=args.seed, + ) + n_inliers = int(inliers.sum()) + print( + f"\nFit complete: used={len(cam_pts)} inliers={n_inliers} " + f"RMSE={rmse * 1000:.1f} mm" + ) + if rmse > franka_calib.MAX_ACCEPTABLE_RMSE_M: + print( + "WARNING: RMSE exceeds the loader acceptance gate " + f"({franka_calib.MAX_ACCEPTABLE_RMSE_M * 1000:.0f} mm). " + "The server will reject this calibration unless you rerun with better detections." + ) + + result = { + "camera": args.camera, + "serial": serial, + "n_collected": len(cam_pts), + "n_inliers": n_inliers, + "rmse_m": rmse, + "T_base_cam": T_base_cam, + "records": records, + "inlier_mask": inliers.tolist(), + "tcp_offset": args.tcp_offset, + "method": "franka_markerless_gripper_motion", + } + + if args.no_save: + print("Not saved (--no-save). T_base_cam:") + print(json.dumps(result["T_base_cam"].tolist(), indent=2)) + return 0 + + path = franka_calib.save_scene_extrinsic( + serial, + T_base_cam, + K=np.asarray(camera_meta["K"], dtype=np.float64), + rmse_m=rmse, + num_points=len(cam_pts), + camera=args.camera, + n_inliers=n_inliers, + inlier_mask=inliers.tolist(), + method="franka_markerless_gripper_motion", + tcp_offset=args.tcp_offset, + correspondences=records, + ) + print(f"Saved T_base_cam -> {path}") + + reload_result = client.call( + "env.reload_camera_calibration", + kwargs={"camera": args.camera}, + timeout_s=15.0, + ) + print("Reload result:") + print(json.dumps(reload_result, indent=2, default=str)) + return 0 if rmse <= franka_calib.MAX_ACCEPTABLE_RMSE_M else 1 + + +def _explain_wrist() -> int: + print( + "Wrist camera calibration is hand-eye calibration (T_tcp_cam), not the " + "same fixed-camera problem as the scene camera. A reliable automated " + "script needs either a fixed fiducial/ChArUco/AprilTag board observed " + "from multiple wrist poses, or a scene-calibrated 3-D reference target. " + "This repository now supports loading/saving T_tcp_cam records, and " + "back_project(camera='wrist') will return panda_link0 xyz once such a " + "record exists." + ) + return 0 + + +def _build_argparser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, help="Franka env server transport port.") + parser.add_argument( + "--mode", + choices=["scene-auto", "wrist-info"], + default="scene-auto", + help="scene-auto calibrates fixed scene T_base_cam; wrist-info explains T_tcp_cam requirements.", + ) + parser.add_argument("--camera", default="scene", help="Camera name to calibrate (default: scene).") + parser.add_argument("--serial", default=None, help="Override saved RealSense serial.") + parser.add_argument("--n-points", type=int, default=8, help="Valid correspondences to collect.") + parser.add_argument("--yes", action="store_true", help="Confirm robot motion.") + parser.add_argument("--no-save", action="store_true", help="Fit but do not save/reload calibration.") + parser.add_argument("--self-test", action="store_true", help="Run offline math/detection self-test.") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--grid-x", type=lambda s: _parse_csv_floats(s, expected=3, name="grid-x"), default=_DEFAULT_GRID_X) + parser.add_argument("--grid-y", type=lambda s: _parse_csv_floats(s, expected=3, name="grid-y"), default=_DEFAULT_GRID_Y) + parser.add_argument("--grid-z", type=lambda s: _parse_csv_floats(s, expected=2, name="grid-z"), default=_DEFAULT_GRID_Z) + parser.add_argument( + "--tcp-offset", + type=lambda s: _parse_csv_floats(s, expected=3, name="tcp-offset"), + default=(0.0, 0.0, 0.0), + help="Optional calibration point offset in TCP frame, meters (default: 0,0,0).", + ) + parser.add_argument("--settle-s", type=float, default=0.4, help="Wait after gripper open/close captures.") + parser.add_argument("--diff-thresh", type=int, default=10) + parser.add_argument("--min-area", type=int, default=12) + parser.add_argument("--max-area", type=int, default=40000) + parser.add_argument("--detect-retries", type=int, default=2, + help="Open/close detection attempts per pose before fallback/skip.") + parser.add_argument("--manual-on-fail", action="store_true", + help="When auto detection fails, prompt for row,col on the saved closed image.") + parser.add_argument("--manual-always", action="store_true", + help="Always prompt for row,col instead of using auto detection.") + parser.add_argument("--manual-patch-radius", type=int, default=3, + help="Depth median patch radius for manually clicked pixels.") + parser.add_argument("--debug-dir", default="/tmp/franka_camera_calib_debug", + help="Directory for per-pose open/closed/diff/depth debug images.") + parser.add_argument("--ransac-thresh-m", type=float, default=0.02) + parser.add_argument("--ransac-iters", type=int, default=500) + return parser + + +def main() -> int: + args = _build_argparser().parse_args() + if args.manual_always: + args.manual_on_fail = True + if args.self_test: + return _self_test() + if args.mode == "wrist-info": + return _explain_wrist() + if args.port is None: + raise SystemExit("--port is required unless --self-test or --mode wrist-info") + if args.n_points < 4: + raise SystemExit("--n-points must be >= 4") + return _calibrate_scene(args) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/deployment/franka/calibrate_charuco_wrist.py b/deployment/franka/calibrate_charuco_wrist.py new file mode 100644 index 00000000..a3b22ea3 --- /dev/null +++ b/deployment/franka/calibrate_charuco_wrist.py @@ -0,0 +1,448 @@ +#!/usr/bin/env python3 +"""ChArUco-based wrist-camera hand-eye calibration for Franka. + +This estimates ``T_tcp_cam`` for the wrist camera. The ChArUco board must be +fixed in the scene while the wrist camera observes it from multiple robot poses. +The script detects the board in the wrist camera, reads the live TCP pose, then +uses OpenCV hand-eye calibration to solve camera-in-TCP. + +Typical workflow: + +1. Start the Franka env server. +2. Place the printed ChArUco board flat and rigid on the table. +3. Move the wrist camera so the board is visible in the wrist image. +4. Check detection: + + python deployment/franka/calibrate_charuco_wrist.py --port 5599 check + +5. Calibrate with a small automatic orbit around the current pose: + + python deployment/franka/calibrate_charuco_wrist.py --port 5599 calibrate --yes + +The scene camera is different: a ChArUco board gives ``T_scene_cam_board`` but +not ``T_base_scene_cam`` unless the board pose in ``panda_link0`` is known. Use +``auto_calibrate_cameras.py`` for markerless scene-to-base calibration, or add a +known board pose / touch-corner workflow for scene calibration. +""" +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path +from typing import Any + +import imageio.v2 as imageio +import numpy as np +from scipy.spatial.transform import Rotation as R + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from deployment.franka import calibration as franka_calib # noqa: E402 +from deployment.lerobot import geometry as geom # noqa: E402 +from physical_agent.rpc_driver.socket import SocketRpcClient # noqa: E402 + +_DEFAULT_BOARD_SPEC = ( + _REPO_ROOT / "resources" / "franka" / "calibration_boards" / "franka_charuco_7x5_25mm.json" +) +_DEFAULT_OFFSETS = ( + ([0.0, 0.0, 0.0], [0.0, 0.0, 0.0]), + ([0.025, 0.0, 0.0], [0.0, 0.0, 12.0]), + ([-0.025, 0.0, 0.0], [0.0, 0.0, -12.0]), + ([0.0, 0.025, 0.0], [0.0, 10.0, 0.0]), + ([0.0, -0.025, 0.0], [0.0, -10.0, 0.0]), + ([0.0, 0.0, 0.025], [8.0, 0.0, 0.0]), + ([0.0, 0.0, -0.015], [-8.0, 0.0, 0.0]), + ([0.02, 0.02, 0.015], [6.0, -8.0, 8.0]), + ([-0.02, -0.02, 0.015], [-6.0, 8.0, -8.0]), +) + + +def _require_cv2_aruco(): + try: + import cv2 + except ModuleNotFoundError as exc: + raise SystemExit( + "Missing dependency: cv2. Install with `uv pip install -e .[calibration]` " + "or run with an environment that has opencv-contrib-python-headless." + ) from exc + if not hasattr(cv2, "aruco"): + raise SystemExit("Installed cv2 lacks aruco; install opencv-contrib-python-headless.") + return cv2 + + +def _load_board_spec(path: Path) -> dict[str, Any]: + spec = json.loads(path.read_text()) + required = ["dictionary", "squares_x", "squares_y", "square_length_m", "marker_length_m"] + missing = [key for key in required if key not in spec] + if missing: + raise SystemExit(f"board spec missing keys: {missing}") + return spec + + +def _aruco_dictionary(cv2, dictionary_name: str): + aruco = cv2.aruco + key = dictionary_name.upper() + if not key.startswith("DICT_"): + key = f"DICT_{key}" + if not hasattr(aruco, key): + raise SystemExit(f"OpenCV does not know ArUco dictionary {dictionary_name!r}") + return aruco.getPredefinedDictionary(getattr(aruco, key)) + + +def _make_board(cv2, spec: dict[str, Any]): + aruco = cv2.aruco + dictionary = _aruco_dictionary(cv2, spec["dictionary"]) + squares = (int(spec["squares_x"]), int(spec["squares_y"])) + square = float(spec["square_length_m"]) + marker = float(spec["marker_length_m"]) + try: + return aruco.CharucoBoard(squares, square, marker, dictionary) + except Exception: + return aruco.CharucoBoard_create(squares[0], squares[1], square, marker, dictionary) + + +def _pose_to_matrix(xyz, quat_xyzw) -> np.ndarray: + T = np.eye(4, dtype=np.float64) + T[:3, :3] = R.from_quat(np.asarray(quat_xyzw, dtype=np.float64)).as_matrix() + T[:3, 3] = np.asarray(xyz, dtype=np.float64) + return T + + +def _pose_from_rvec_tvec(cv2, rvec, tvec) -> np.ndarray: + T = np.eye(4, dtype=np.float64) + R_cam_board, _ = cv2.Rodrigues(np.asarray(rvec, dtype=np.float64)) + T[:3, :3] = R_cam_board + T[:3, 3] = np.asarray(tvec, dtype=np.float64).reshape(3) + return T + + +def _detect_charuco_pose(cv2, image, K, dist_coeffs, board) -> dict | None: + aruco = cv2.aruco + gray = cv2.cvtColor(np.asarray(image, dtype=np.uint8), cv2.COLOR_RGB2GRAY) + marker_corners = [] + marker_ids = None + + if hasattr(aruco, "CharucoDetector"): + detector = aruco.CharucoDetector(board) + charuco_corners, charuco_ids, marker_corners, marker_ids = detector.detectBoard(gray) + count = 0 if charuco_ids is None else len(charuco_ids) + else: + params = aruco.DetectorParameters() + try: + detector = aruco.ArucoDetector(board.getDictionary(), params) + marker_corners, marker_ids, _ = detector.detectMarkers(gray) + except Exception: + marker_corners, marker_ids, _ = aruco.detectMarkers( + gray, board.getDictionary(), parameters=params + ) + if marker_ids is None or len(marker_ids) < 2: + return None + + try: + count, charuco_corners, charuco_ids = aruco.interpolateCornersCharuco( + marker_corners, + marker_ids, + gray, + board, + cameraMatrix=K, + distCoeffs=dist_coeffs, + ) + except TypeError: + count, charuco_corners, charuco_ids = aruco.interpolateCornersCharuco( + marker_corners, marker_ids, gray, board, K, dist_coeffs + ) + + if charuco_ids is None or int(count) < 6: + return None + + if hasattr(aruco, "estimatePoseCharucoBoard"): + rvec = np.zeros((3, 1), dtype=np.float64) + tvec = np.zeros((3, 1), dtype=np.float64) + ok, rvec, tvec = aruco.estimatePoseCharucoBoard( + charuco_corners, + charuco_ids, + board, + K, + dist_coeffs, + rvec, + tvec, + ) + else: + obj_points, img_points = board.matchImagePoints(charuco_corners, charuco_ids) + ok, rvec, tvec = cv2.solvePnP( + obj_points, + img_points, + K, + dist_coeffs, + flags=cv2.SOLVEPNP_ITERATIVE, + ) + if not ok: + return None + return { + "T_cam_board": _pose_from_rvec_tvec(cv2, rvec, tvec), + "n_markers": 0 if marker_ids is None else int(len(marker_ids)), + "n_corners": int(count), + "rvec": np.asarray(rvec, dtype=np.float64).reshape(3).tolist(), + "tvec": np.asarray(tvec, dtype=np.float64).reshape(3).tolist(), + } + + +def _draw_detection(cv2, image, K, dist_coeffs, board, out_path: Path) -> None: + aruco = cv2.aruco + canvas = np.asarray(image, dtype=np.uint8).copy() + gray = cv2.cvtColor(canvas, cv2.COLOR_RGB2GRAY) + charuco_corners = None + charuco_ids = None + if hasattr(aruco, "CharucoDetector"): + detector = aruco.CharucoDetector(board) + charuco_corners, charuco_ids, marker_corners, marker_ids = detector.detectBoard(gray) + else: + params = aruco.DetectorParameters() + try: + detector = aruco.ArucoDetector(board.getDictionary(), params) + marker_corners, marker_ids, _ = detector.detectMarkers(gray) + except Exception: + marker_corners, marker_ids, _ = aruco.detectMarkers( + gray, board.getDictionary(), parameters=params + ) + if marker_ids is not None and len(marker_ids) > 0: + aruco.drawDetectedMarkers(canvas, marker_corners, marker_ids) + if charuco_ids is not None and len(charuco_ids) > 0: + aruco.drawDetectedCornersCharuco(canvas, charuco_corners, charuco_ids) + out_path.parent.mkdir(parents=True, exist_ok=True) + imageio.imwrite(out_path, canvas) + + +def _capture(client: SocketRpcClient, camera: str) -> tuple[dict, np.ndarray, dict]: + obs = client.call("env.get_obs", timeout_s=30.0) + frames = obs.get("frames") or {} + meta = obs.get("camera_meta") or {} + if camera not in frames: + raise RuntimeError(f"camera {camera!r} missing from observation") + if camera not in meta: + raise RuntimeError(f"camera {camera!r} metadata missing") + ee = client.call("env.get_ee_pose", timeout_s=15.0) + return ee, np.asarray(frames[camera], dtype=np.uint8), dict(meta[camera]) + + +def _check(args: argparse.Namespace) -> int: + cv2 = _require_cv2_aruco() + board = _make_board(cv2, _load_board_spec(args.board_spec)) + client = SocketRpcClient(args.host, args.port) + ee, image, meta = _capture(client, args.camera) + K = np.asarray(meta["K"], dtype=np.float64) + dist = np.asarray(meta.get("dist_coeffs") or np.zeros(5), dtype=np.float64) + det = _detect_charuco_pose(cv2, image, K, dist, board) + out_path = Path(args.debug_dir) / f"{args.camera}_charuco_check.png" + _draw_detection(cv2, image, K, dist, board, out_path) + print(json.dumps({ + "camera": args.camera, + "serial": meta.get("serial"), + "debug_image": str(out_path), + "detected": det is not None, + "n_markers": None if det is None else det["n_markers"], + "n_corners": None if det is None else det["n_corners"], + "tcp_xyz": ee.get("xyz"), + }, indent=2)) + return 0 if det is not None else 2 + + +def _target_pose(start: dict, dxyz, drpy_deg) -> tuple[list[float], list[float]]: + xyz = np.asarray(start["xyz"], dtype=np.float64) + np.asarray(dxyz, dtype=np.float64) + euler = np.asarray(start["euler_xyz"], dtype=np.float64) + np.radians(np.asarray(drpy_deg, dtype=np.float64)) + return xyz.tolist(), euler.tolist() + + +def _collect_samples( + args: argparse.Namespace, +) -> tuple[ + list[np.ndarray], + list[np.ndarray], + list[np.ndarray], + list[np.ndarray], + list[dict], +]: + cv2 = _require_cv2_aruco() + board = _make_board(cv2, _load_board_spec(args.board_spec)) + client = SocketRpcClient(args.host, args.port) + start = client.call("env.get_ee_pose", timeout_s=15.0) + start_xyz = start["xyz"] + start_quat = start["quat_xyzw"] + + R_gripper2base: list[np.ndarray] = [] + t_gripper2base: list[np.ndarray] = [] + R_target2cam: list[np.ndarray] = [] + t_target2cam: list[np.ndarray] = [] + records: list[dict] = [] + + try: + for idx, (dxyz, drpy) in enumerate(_DEFAULT_OFFSETS, start=1): + if len(records) >= args.n_samples: + break + target_xyz, target_euler = _target_pose(start, dxyz, drpy) + print(f"\n[{idx}] target_xyz={np.round(target_xyz, 3).tolist()} drpy={drpy}") + move = client.call( + "env.move_to", + args=(target_xyz,), + kwargs={"euler_xyz": target_euler, "gripper": None}, + timeout_s=120.0, + ) + print(f" move: reached={move.get('reached')} err={move.get('pos_error_m')} final={move.get('final_xyz')}") + time.sleep(args.settle_s) + ee, image, meta = _capture(client, args.camera) + K = np.asarray(meta["K"], dtype=np.float64) + dist = np.asarray(meta.get("dist_coeffs") or np.zeros(5), dtype=np.float64) + det = _detect_charuco_pose(cv2, image, K, dist, board) + debug_path = Path(args.debug_dir) / f"{args.camera}_charuco_{idx:02d}.png" + _draw_detection(cv2, image, K, dist, board, debug_path) + if det is None: + print(f" detection failed (debug: {debug_path})") + continue + + T_base_tcp = _pose_to_matrix(ee["xyz"], ee["quat_xyzw"]) + T_cam_board = det["T_cam_board"] + R_gripper2base.append(T_base_tcp[:3, :3]) + t_gripper2base.append(T_base_tcp[:3, 3].reshape(3, 1)) + R_target2cam.append(T_cam_board[:3, :3]) + t_target2cam.append(T_cam_board[:3, 3].reshape(3, 1)) + records.append({ + "idx": idx, + "target_xyz": target_xyz, + "target_euler": target_euler, + "tcp_xyz": ee["xyz"], + "n_markers": det["n_markers"], + "n_corners": det["n_corners"], + "tvec": det["tvec"], + "debug_image": str(debug_path), + "move": move, + }) + print(f" captured: markers={det['n_markers']} corners={det['n_corners']} tvec={np.round(det['tvec'], 3).tolist()}") + finally: + try: + client.call( + "env.move_to", + args=(start_xyz,), + kwargs={"quat_xyzw": start_quat, "gripper": None}, + timeout_s=120.0, + ) + except Exception as exc: + print(f"warning: failed to return to start pose: {exc}") + + return ( + [np.asarray(r, dtype=np.float64) for r in R_gripper2base], + [np.asarray(t, dtype=np.float64) for t in t_gripper2base], + [np.asarray(r, dtype=np.float64) for r in R_target2cam], + [np.asarray(t, dtype=np.float64) for t in t_target2cam], + records, + ) + + +def _calibrate(args: argparse.Namespace) -> int: + if not args.yes: + print("Refusing to move the robot without --yes.") + return 2 + cv2 = _require_cv2_aruco() + board_spec = _load_board_spec(args.board_spec) + client = SocketRpcClient(args.host, args.port) + meta = client.call("env.get_camera_meta", timeout_s=15.0).get(args.camera) + if not meta: + print(f"camera {args.camera!r} not available") + return 2 + serial = args.serial or meta.get("serial") + if not serial: + print("could not determine wrist camera serial; pass --serial") + return 2 + + R_g2b, t_g2b, R_t2c, t_t2c, records = _collect_samples(args) + if len(records) < 5: + print(f"Need at least 5 valid board detections; got {len(records)}") + return 2 + + R_cam2tcp, t_cam2tcp = cv2.calibrateHandEye( + R_g2b, + t_g2b, + R_t2c, + t_t2c, + method=cv2.CALIB_HAND_EYE_TSAI, + ) + T_tcp_cam = np.eye(4, dtype=np.float64) + T_tcp_cam[:3, :3] = np.asarray(R_cam2tcp, dtype=np.float64) + T_tcp_cam[:3, 3] = np.asarray(t_cam2tcp, dtype=np.float64).reshape(3) + + residuals = [] + target_points_base = [] + for Rb, tb, Rc, tc in zip(R_g2b, t_g2b, R_t2c, t_t2c): + T_base_tcp = np.eye(4) + T_base_tcp[:3, :3] = Rb + T_base_tcp[:3, 3] = tb.reshape(3) + T_cam_board = np.eye(4) + T_cam_board[:3, :3] = Rc + T_cam_board[:3, 3] = tc.reshape(3) + T_base_board = T_base_tcp @ T_tcp_cam @ T_cam_board + target_points_base.append(T_base_board[:3, 3]) + target_points_base = np.asarray(target_points_base) + center = target_points_base.mean(axis=0) + residuals = np.linalg.norm(target_points_base - center, axis=1) + rmse = float(np.sqrt(np.mean(np.square(residuals)))) + + print(f"\nHand-eye fit: samples={len(records)} board-position RMSE={rmse * 1000:.1f} mm") + print("T_tcp_cam:") + print(json.dumps(T_tcp_cam.tolist(), indent=2)) + + if args.no_save: + return 0 + + path = franka_calib.save_wrist_extrinsic( + serial, + T_tcp_cam, + K=np.asarray(meta["K"], dtype=np.float64), + rmse_m=rmse, + num_points=len(records), + camera=args.camera, + method="charuco_hand_eye_tsai", + board_spec=board_spec, + records=records, + board_position_residuals_m=residuals.tolist(), + ) + print(f"Saved T_tcp_cam -> {path}") + reload_result = client.call("env.reload_camera_calibration", kwargs={"camera": args.camera}, timeout_s=15.0) + print("Reload result:") + print(json.dumps(reload_result, indent=2, default=str)) + return 0 if rmse <= franka_calib.MAX_ACCEPTABLE_RMSE_M else 1 + + +def _build_argparser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, required=True) + parser.add_argument("--camera", default="wrist") + parser.add_argument("--serial", default=None) + parser.add_argument("--board-spec", type=Path, default=_DEFAULT_BOARD_SPEC) + parser.add_argument("--debug-dir", type=Path, default=Path("/tmp/franka_charuco_wrist_debug")) + sub = parser.add_subparsers(dest="cmd", required=True) + + sub.add_parser("check", help="Capture one wrist frame and report board detection.") + + calib = sub.add_parser("calibrate", help="Move through a small orbit and solve T_tcp_cam.") + calib.add_argument("--yes", action="store_true") + calib.add_argument("--no-save", action="store_true") + calib.add_argument("--n-samples", type=int, default=8) + calib.add_argument("--settle-s", type=float, default=0.5) + return parser + + +def main() -> int: + args = _build_argparser().parse_args() + if args.cmd == "check": + return _check(args) + if args.cmd == "calibrate": + return _calibrate(args) + raise SystemExit(f"unknown command: {args.cmd}") + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/deployment/franka/calibration.py b/deployment/franka/calibration.py new file mode 100644 index 00000000..2f365c99 --- /dev/null +++ b/deployment/franka/calibration.py @@ -0,0 +1,120 @@ +"""Franka camera calibration loading helpers. + +Calibration records are stored per RealSense serial under +``~/.cache/physical_agent/franka/camera_calibration``. Supported records: + +- fixed scene camera: ``{"T_base_cam": [[...]], ...}`` +- wrist camera: ``{"T_tcp_cam": [[...]], ...}`` + +``T_base_cam`` maps camera-frame points into ``panda_link0``. ``T_tcp_cam`` maps +wrist-camera points into the live TCP frame; the env server composes it with the +current ``T_base_tcp`` for each observation. +""" +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +import numpy as np + +_CALIB_DIR = "~/.cache/physical_agent/franka/camera_calibration" +MAX_ACCEPTABLE_RMSE_M = 0.02 + + +def calib_path(serial: str) -> Path: + """Return the on-disk calibration path for a RealSense serial.""" + return Path(os.path.expanduser(_CALIB_DIR)) / f"{serial}.json" + + +def load_record(serial: str) -> dict | None: + """Load a camera calibration record, converting known transforms to arrays.""" + path = calib_path(serial) + if not path.is_file(): + return None + with open(path) as f: + data = json.load(f) + for key in ("T_base_cam", "T_tcp_cam"): + if key in data and data[key] is not None: + data[key] = np.asarray(data[key], dtype=np.float64) + data.setdefault("serial", str(serial)) + data["path"] = str(path) + return data + + +def _jsonable(value: Any) -> Any: + """Convert numpy values to JSON-compatible Python values.""" + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + if isinstance(value, dict): + return {key: _jsonable(val) for key, val in value.items()} + if isinstance(value, (list, tuple)): + return [_jsonable(val) for val in value] + return value + + +def save_record(serial: str, **fields: Any) -> Path: + """Save a calibration record for ``serial``. + + Known transform fields are ``T_base_cam`` for fixed cameras and + ``T_tcp_cam`` for wrist cameras. Extra diagnostics are preserved. + """ + path = calib_path(serial) + path.parent.mkdir(parents=True, exist_ok=True) + data = {"serial": str(serial), **fields} + with open(path, "w") as f: + json.dump(_jsonable(data), f, indent=2) + return path + + +def save_scene_extrinsic( + serial: str, + T_base_cam, + *, + K=None, + rmse_m: float | None = None, + num_points: int | None = None, + **extra: Any, +) -> Path: + """Persist a fixed-camera ``T_base_cam`` calibration.""" + fields: dict[str, Any] = {"T_base_cam": np.asarray(T_base_cam, dtype=np.float64)} + if K is not None: + fields["K"] = np.asarray(K, dtype=np.float64) + if rmse_m is not None: + fields["rmse_m"] = float(rmse_m) + if num_points is not None: + fields["num_points"] = int(num_points) + fields.update(extra) + return save_record(serial, **fields) + + +def save_wrist_extrinsic( + serial: str, + T_tcp_cam, + *, + K=None, + rmse_m: float | None = None, + num_points: int | None = None, + **extra: Any, +) -> Path: + """Persist a wrist-camera ``T_tcp_cam`` hand-eye calibration.""" + fields: dict[str, Any] = {"T_tcp_cam": np.asarray(T_tcp_cam, dtype=np.float64)} + if K is not None: + fields["K"] = np.asarray(K, dtype=np.float64) + if rmse_m is not None: + fields["rmse_m"] = float(rmse_m) + if num_points is not None: + fields["num_points"] = int(num_points) + fields.update(extra) + return save_record(serial, **fields) + + +def is_accepted(record: dict | None) -> bool: + """Return whether a record exists and passes the saved RMSE gate.""" + if record is None: + return False + rmse = record.get("rmse_m") + return rmse is None or float(rmse) <= MAX_ACCEPTABLE_RMSE_M diff --git a/deployment/franka/env_server.py b/deployment/franka/env_server.py new file mode 100644 index 00000000..5f93afdb --- /dev/null +++ b/deployment/franka/env_server.py @@ -0,0 +1,954 @@ +"""Standalone Franka (FR3) env host for the LLM-in-the-loop agent (env-only). + +Drives a Franka arm through the SERL cartesian-impedance ROS controller and the +``franka_gripper`` action topics, exposing agent-friendly *Cartesian* primitives +(``reset`` / ``get_obs`` / ``get_ee_pose`` / ``move_to`` / ``move_delta`` / +``open_gripper`` / ``close_gripper`` / ``get_spec``) over a pickle-framed TCP RPC +server (:class:`physical_agent.rpc_driver.socket.SocketRpcServer`) -- the same +wire protocol the LIBERO and LeRobot drivers use, so the agent side talks to all +three identically. + +Unlike the LIBERO driver, this server does **not** import any +``rlinf.envs.realworld`` module. Importing that package runs node-level ROS setup +side effects at import time (it kills any running ``roscore`` / ``rosmaster``) and +pulls in the Ray-based ``Worker`` stack. This driver instead talks to ROS directly +with ``rospy``, reusing only the *recipe* from +``rlinf.envs.realworld.franka.{franka_controller,franka_env}``: the impedance +controller channel names, the ``roslaunch`` bring-up, the ``franka_gripper`` +action messages, and the safety-box + pose-interpolation logic. + +Run it inside the RLinf ``.venv`` with the ``serl_franka_controllers`` catkin +workspace sourced (see ``deployment/franka/run_env_server.sh``):: + + source /home/franka/franka/RLinf/.venv/franka_catkin_ws/devel/setup.bash + /home/franka/franka/RLinf/.venv/bin/python deployment/franka/env_server.py \ + --output-dir /tmp/franka_run --robot-ip 172.16.0.2 + +Hardware defaults match the current bench: an FR3 at ``172.16.0.2`` with the +Franka Hand, plus two Intel RealSense D435I cameras. Every default is overridable +from the CLI. +""" +from __future__ import annotations + +import argparse +import json +import os +import sys +import threading +import time +from pathlib import Path +from typing import Any, Optional + +import numpy as np +from scipy.spatial.transform import Rotation as R +from scipy.spatial.transform import Slerp + +# Make ``physical_agent`` importable when this file is run from the RLinf .venv +# (which need not have physical_agent installed) -- the source tree is enough. +_PHYSICALAGENT_ROOT = Path(__file__).resolve().parents[2] +if str(_PHYSICALAGENT_ROOT) not in sys.path: + sys.path.insert(0, str(_PHYSICALAGENT_ROOT)) + +from physical_agent.rpc_driver.socket import SocketRpcServer # noqa: E402 +from physical_agent.utils.logging import get_logger, init_output_dir # noqa: E402 + +from deployment.franka import calibration as camera_calib # noqa: E402 + +logger = get_logger("franka_driver") + + +# --------------------------------------------------------------------------- +# Bench defaults (override from the CLI) +# --------------------------------------------------------------------------- + +_DEFAULT_ROBOT_IP = os.environ.get("FRANKA_ROBOT_IP", "172.16.0.2") +_DEFAULT_ROS_PKG = "serl_franka_controllers" + +# Two Intel RealSense D435I on the current bench. Order matters: the first camera +# is the primary/overview view (maps to the ToolResult ``_image_bytes`` slot), the +# second maps to ``_image_cam_bytes``. Names are what the agent sees. +_DEFAULT_CAMERAS: tuple[tuple[str, str], ...] = ( + ("scene", "142122070838"), + ("wrist", "141722078696"), +) +_CAM_WIDTH = 640 +_CAM_HEIGHT = 480 +_CAM_FPS = 30 + +# Conservative tabletop workspace box in the base (panda_link0) frame, meters. +# ``move_to`` clips targets to this so an LLM-supplied coordinate cannot drive the +# arm into the table / out of reach. Tuned around the collect-data reset pose +# ([0.5, 0, 0.1] with the gripper pointing down); widen/lower on your bench once +# you have confirmed the table height in the base frame. +_WORKSPACE_MIN = np.array([0.30, -0.50, 0.00], dtype=np.float64) +_WORKSPACE_MAX = np.array([1.10, 0.50, 0.50], dtype=np.float64) + +# Default "home" pose the agent's reset() drives to: above the table, gripper +# pointing straight down. Orientation is euler xyz (radians); rx = -pi points the +# Franka Hand down (matches the RLinf realworld collect configs). +_RESET_XYZ = np.array([0.50, 0.0, 0.25], dtype=np.float64) +_RESET_EULER = np.array([np.pi, 0.0, 0.0], dtype=np.float64) + +# The nominal orientation the arm holds; move_to keeps the current orientation +# unless the caller overrides it, and clips any requested orientation to a window +# around this so the wrist cannot flip into a self-collision. +_TARGET_EULER = _RESET_EULER.copy() +_EULER_WINDOW = np.array([0.6, 0.6, np.pi], dtype=np.float64) # +/- rad per axis + +# Motion smoothness / safety. The impedance controller tracks a streamed sequence +# of equilibrium poses; capping the per-step Cartesian and angular deltas and +# pacing at ``step_frequency`` keeps motions slow and gentle. +_STEP_FREQUENCY = 10.0 # Hz (equilibrium-pose publish rate) +_MAX_STEP_M = 0.01 # max Cartesian move per streamed setpoint (=> ~0.1 m/s) +_MAX_STEP_DEG = 5.0 # max orientation change per streamed setpoint +_MAX_MOVE_M = 0.60 # hard cap on a single move_to path length (safety) +_REACHED_TOL_M = 0.003 # move_to "reached" tolerance +_SETTLE_TIMEOUT_S = 2.0 # max time to hold the final setpoint while settling + +# franka_gripper widths (meters). +_GRIPPER_OPEN_WIDTH = 0.09 +_GRIPPER_GRASP_WIDTH = 0.01 +_GRIPPER_GRASP_FORCE = 130.0 +_GRIPPER_SPEED = 0.3 + + +# --------------------------------------------------------------------------- +# small pose helpers (scipy uses scalar-last quaternions: [x, y, z, w]) +# --------------------------------------------------------------------------- + + +def _euler_to_quat(euler_xyz) -> np.ndarray: + return R.from_euler("xyz", np.asarray(euler_xyz, dtype=np.float64)).as_quat() + + +def _quat_to_euler(quat_xyzw) -> np.ndarray: + return R.from_quat(np.asarray(quat_xyzw, dtype=np.float64)).as_euler("xyz") + + +def _pose_to_matrix(pose7: np.ndarray) -> np.ndarray: + """Convert ``[x, y, z, qx, qy, qz, qw]`` to ``T_base_tcp``.""" + pose7 = np.asarray(pose7, dtype=np.float64).reshape(-1)[:7] + T = np.eye(4, dtype=np.float64) + T[:3, :3] = R.from_quat(pose7[3:]).as_matrix() + T[:3, 3] = pose7[:3] + return T + + +def _clip_euler_window(euler_xyz: np.ndarray) -> np.ndarray: + """Clip an euler orientation to +/-``_EULER_WINDOW`` around ``_TARGET_EULER``. + + Wraps each axis into ``[-pi, pi]`` relative to the target first so the clip is + on the shortest angular distance, not the raw value. + """ + euler = np.asarray(euler_xyz, dtype=np.float64).copy() + delta = (euler - _TARGET_EULER + np.pi) % (2 * np.pi) - np.pi + delta = np.clip(delta, -_EULER_WINDOW, _EULER_WINDOW) + return _TARGET_EULER + delta + + +# --------------------------------------------------------------------------- +# RealSense color + depth camera +# --------------------------------------------------------------------------- + + +class RealSenseDepthCamera: + """Minimal RealSense RGB-D grabber with depth aligned to color.""" + + def __init__(self, name: str, serial: str, *, width: int, height: int, fps: int): + import pyrealsense2 as rs + + self.name = name + self.serial = str(serial) + self._rs = rs + self._pipeline = rs.pipeline() + cfg = rs.config() + cfg.enable_device(self.serial) + cfg.enable_stream(rs.stream.color, width, height, rs.format.rgb8, fps) + cfg.enable_stream(rs.stream.depth, width, height, rs.format.z16, fps) + self._profile = self._pipeline.start(cfg) + self._align = rs.align(rs.stream.color) + + depth_sensor = self._profile.get_device().first_depth_sensor() + self._depth_scale = float(depth_sensor.get_depth_scale()) + + color_stream = self._profile.get_stream( + rs.stream.color + ).as_video_stream_profile() + intr = color_stream.get_intrinsics() + self._K = np.array( + [[intr.fx, 0.0, intr.ppx], [0.0, intr.fy, intr.ppy], [0.0, 0.0, 1.0]], + dtype=np.float64, + ) + self._dist_coeffs = np.asarray(intr.coeffs, dtype=np.float64) + self._distortion_model = str(intr.model) + self._width = int(width) + self._height = int(height) + self._calib = camera_calib.load_record(self.serial) + self._calib_accepted = camera_calib.is_accepted(self._calib) + + # Drop the first few frames so auto-exposure settles. + for _ in range(5): + try: + self._pipeline.wait_for_frames(2000) + except Exception: + break + calib_kind = "uncalibrated" + if self._calib_accepted: + if self._calib and "T_base_cam" in self._calib: + calib_kind = "T_base_cam" + elif self._calib and "T_tcp_cam" in self._calib: + calib_kind = "T_tcp_cam" + elif self._calib is not None: + calib_kind = "rejected" + logger.info( + "camera '%s' (serial %s) RGB-D started; calibration=%s", + name, + self.serial, + calib_kind, + ) + + def read(self) -> tuple[np.ndarray, np.ndarray]: + """Return ``(rgb_uint8, depth_m_float32)`` with depth aligned to RGB.""" + frames = self._pipeline.wait_for_frames(2000) + frames = self._align.process(frames) + color = frames.get_color_frame() + depth = frames.get_depth_frame() + if not color or not depth: + raise RuntimeError(f"camera '{self.name}' incomplete frameset") + rgb = np.ascontiguousarray(np.asanyarray(color.get_data()), dtype=np.uint8) + depth_raw = np.asanyarray(depth.get_data()) + depth_m = np.ascontiguousarray( + depth_raw.astype(np.float32) * self._depth_scale + ) + return rgb, depth_m + + def meta(self, *, T_base_tcp: np.ndarray | None = None) -> dict: + """Return JSON-able camera metadata and any base-frame extrinsic.""" + T_base_cam = None + calibration_kind = None + if self._calib_accepted and self._calib is not None: + if "T_base_cam" in self._calib: + T_base_cam = self._calib["T_base_cam"] + calibration_kind = "T_base_cam" + elif "T_tcp_cam" in self._calib: + calibration_kind = "T_tcp_cam" + if T_base_tcp is not None: + T_base_cam = np.asarray(T_base_tcp, dtype=np.float64) @ self._calib[ + "T_tcp_cam" + ] + + return { + "name": self.name, + "serial": self.serial, + "frame": f"{self.name}_camera", + "width": self._width, + "height": self._height, + "K": self._K.tolist(), + "dist_coeffs": self._dist_coeffs.tolist(), + "distortion_model": self._distortion_model, + "depth_scale": self._depth_scale, + "calibrated": T_base_cam is not None, + "calibration_kind": calibration_kind, + "calibration_rmse_m": ( + None if self._calib is None else self._calib.get("rmse_m") + ), + "calibration_path": ( + None if self._calib is None else self._calib.get("path") + ), + "T_base_cam": None if T_base_cam is None else T_base_cam.tolist(), + } + + def reload_calibration(self) -> dict: + """Reload this camera's calibration record from disk.""" + self._calib = camera_calib.load_record(self.serial) + self._calib_accepted = camera_calib.is_accepted(self._calib) + meta = self.meta() + logger.info( + "camera '%s' calibration reloaded: calibrated=%s kind=%s", + self.name, + meta["calibrated"], + meta["calibration_kind"], + ) + return meta + + def close(self) -> None: + try: + self._pipeline.stop() + except Exception: + pass + + +# --------------------------------------------------------------------------- +# Franka ROS backend (plain rospy; no rlinf, no Ray) +# --------------------------------------------------------------------------- + + +class FrankaRobotBackend: + """Low-level Franka arm + gripper over ROS. + + Reuses the channel names, impedance ``roslaunch`` bring-up, and message + parsing from ``rlinf.envs.realworld.franka.franka_controller`` / + ``.common.gripper.franka_gripper`` -- reimplemented on plain ``rospy`` so this + driver stays free of the Ray ``Worker`` stack and the destructive + ``rlinf.envs.realworld`` import-time side effects. + """ + + _ARM_EQUILIBRIUM = "/cartesian_impedance_controller/equilibrium_pose" + _ARM_STATE = "/franka_state_controller/franka_states" + _ARM_RESET = "/franka_control/error_recovery/goal" + _GRIPPER_MOVE = "/franka_gripper/move/goal" + _GRIPPER_GRASP = "/franka_gripper/grasp/goal" + _GRIPPER_STATE = "/franka_gripper/joint_states" + + def __init__( + self, + *, + robot_ip: str, + ros_pkg: str = _DEFAULT_ROS_PKG, + load_gripper: bool = True, + node_name: str = "franka_agent_driver", + state_timeout_s: float = 4.0, + launch_timeout_s: float = 40.0, + ): + # Lazy ROS imports so the module can be imported (e.g. for --help) on a + # host without ROS on the path. + import geometry_msgs.msg as geom_msg + import rospy + from franka_gripper.msg import GraspActionGoal, MoveActionGoal + from franka_msgs.msg import ErrorRecoveryActionGoal, FrankaState + from sensor_msgs.msg import JointState + + self._rospy = rospy + self._geom_msg = geom_msg + self._FrankaState = FrankaState + self._ErrorRecoveryActionGoal = ErrorRecoveryActionGoal + self._MoveActionGoal = MoveActionGoal + self._GraspActionGoal = GraspActionGoal + self._JointState = JointState + + self._robot_ip = robot_ip + self._ros_pkg = ros_pkg + self._load_gripper = load_gripper + self._impedance_proc = None # only set if we launch it ourselves + + # Live state (updated by subscriber callbacks). + self._tcp_pose = np.zeros(7, dtype=np.float64) + self._tcp_pose[6] = 1.0 # unit quaternion + self._tcp_force = np.zeros(3, dtype=np.float64) + self._tcp_torque = np.zeros(3, dtype=np.float64) + self._state_seen = threading.Event() + self._gripper_width = _GRIPPER_OPEN_WIDTH + self._gripper_open = True + self._gripper_seen = threading.Event() + + self._ensure_master() + rospy.init_node(node_name, anonymous=True, disable_signals=True) + + # Publishers. + self._pub_equilibrium = rospy.Publisher( + self._ARM_EQUILIBRIUM, geom_msg.PoseStamped, queue_size=10 + ) + self._pub_reset = rospy.Publisher( + self._ARM_RESET, ErrorRecoveryActionGoal, queue_size=1 + ) + self._pub_gripper_move = rospy.Publisher( + self._GRIPPER_MOVE, MoveActionGoal, queue_size=1 + ) + self._pub_gripper_grasp = rospy.Publisher( + self._GRIPPER_GRASP, GraspActionGoal, queue_size=1 + ) + + # Subscribers. + rospy.Subscriber(self._ARM_STATE, FrankaState, self._on_state_msg) + rospy.Subscriber(self._GRIPPER_STATE, JointState, self._on_gripper_msg) + + self._ensure_impedance(state_timeout_s, launch_timeout_s) + + # -- lifecycle helpers ------------------------------------------------- + + def _ensure_master(self) -> None: + """Make sure a ROS master is reachable; launch ``roscore`` if not. + + Mirrors ``rlinf...common.ros.ros_controller.ROSController``: reuse a + running master, otherwise start one. ``rosmaster --core`` (what + ``roscore`` spawns) also counts as a running master. + """ + import psutil + + try: + import rosgraph + + if rosgraph.is_master_online(): + logger.info("ROS master already online at %s", os.environ.get( + "ROS_MASTER_URI", "http://localhost:11311")) + return + except Exception: + pass + + for proc in psutil.process_iter(["name"]): + if proc.info.get("name") in ("roscore", "rosmaster"): + logger.info("found running %s (pid %s)", proc.info["name"], proc.pid) + return + + logger.info("no ROS master found; starting `roscore`") + self._roscore_proc = psutil.Popen( + ["roscore"], stdout=sys.stdout, stderr=sys.stdout + ) + time.sleep(2.0) + + def _impedance_up(self, timeout_s: float) -> bool: + """Return True once a franka_states message arrives within ``timeout_s``.""" + return self._state_seen.wait(timeout=timeout_s) + + def _ensure_impedance(self, state_timeout_s: float, launch_timeout_s: float) -> None: + """Attach to a running cartesian-impedance controller, or launch one. + + We never launch a second controller: if ``franka_states`` is already + publishing we simply attach. Only when no state arrives do we + ``roslaunch serl_franka_controllers impedance.launch``. + """ + if self._impedance_up(state_timeout_s): + logger.info("attached to live cartesian-impedance controller") + return + + import psutil + + load_gripper = "true" if self._load_gripper else "false" + cmd = [ + "roslaunch", + self._ros_pkg, + "impedance.launch", + f"robot_ip:={self._robot_ip}", + f"load_gripper:={load_gripper}", + ] + logger.info("no live controller detected; starting `%s`", " ".join(cmd)) + self._impedance_proc = psutil.Popen(cmd, stdout=sys.stdout, stderr=sys.stdout) + + deadline = time.time() + launch_timeout_s + while time.time() < deadline: + if self._impedance_proc.poll() is not None: + raise RuntimeError( + "impedance roslaunch exited before becoming ready; run " + f"`{' '.join(cmd)}` manually to see the error (is " + f"{self._ros_pkg} on the ROS package path?)." + ) + if self._impedance_up(1.0): + logger.info("cartesian-impedance controller is up") + return + raise RuntimeError( + f"cartesian-impedance controller not ready after {launch_timeout_s}s" + ) + + # -- ROS callbacks ----------------------------------------------------- + + def _on_state_msg(self, msg) -> None: + tmatrix = np.array(list(msg.O_T_EE)).reshape(4, 4).T + quat = R.from_matrix(tmatrix[:3, :3].copy()).as_quat() + self._tcp_pose = np.concatenate([tmatrix[:3, 3], quat]) + self._tcp_force = np.array(list(msg.K_F_ext_hat_K)[:3]) + self._tcp_torque = np.array(list(msg.K_F_ext_hat_K)[3:]) + self._state_seen.set() + + def _on_gripper_msg(self, msg) -> None: + # joint_states reports both finger joints; their sum is the opening width. + self._gripper_width = float(np.sum(msg.position)) + self._gripper_open = self._gripper_width > 0.06 + self._gripper_seen.set() + + # -- arm --------------------------------------------------------------- + + def get_tcp_pose(self) -> np.ndarray: + """Current TCP pose ``[x, y, z, qx, qy, qz, qw]`` in the base frame.""" + return self._tcp_pose.copy() + + def get_state(self) -> dict: + return { + "tcp_pose": self._tcp_pose.copy(), + "tcp_force": self._tcp_force.copy(), + "tcp_torque": self._tcp_torque.copy(), + "gripper_width": self._gripper_width, + "gripper_open": self._gripper_open, + } + + def move_arm(self, pose7: np.ndarray) -> None: + """Publish one equilibrium pose ``[x, y, z, qx, qy, qz, qw]``.""" + pose7 = np.asarray(pose7, dtype=np.float64).reshape(-1) + assert pose7.shape[0] == 7, f"expected 7-D pose, got {pose7.shape}" + msg = self._geom_msg.PoseStamped() + msg.header.frame_id = "0" + msg.header.stamp = self._rospy.Time.now() + msg.pose.position = self._geom_msg.Point(pose7[0], pose7[1], pose7[2]) + msg.pose.orientation = self._geom_msg.Quaternion( + pose7[3], pose7[4], pose7[5], pose7[6] + ) + self._pub_equilibrium.publish(msg) + + def clear_errors(self) -> None: + self._pub_reset.publish(self._ErrorRecoveryActionGoal()) + + # -- gripper ----------------------------------------------------------- + + def open_gripper(self, speed: float = _GRIPPER_SPEED) -> None: + msg = self._MoveActionGoal() + msg.goal.width = _GRIPPER_OPEN_WIDTH + msg.goal.speed = speed + self._pub_gripper_move.publish(msg) + self._gripper_open = True + + def close_gripper( + self, speed: float = _GRIPPER_SPEED, force: float = _GRIPPER_GRASP_FORCE + ) -> None: + msg = self._GraspActionGoal() + msg.goal.width = _GRIPPER_GRASP_WIDTH + msg.goal.speed = speed + msg.goal.epsilon.inner = 1.0 + msg.goal.epsilon.outer = 1.0 + msg.goal.force = force + self._pub_gripper_grasp.publish(msg) + self._gripper_open = False + + def shutdown(self) -> None: + """Terminate only the controller we launched (never a pre-existing one).""" + if self._impedance_proc is not None and self._impedance_proc.poll() is None: + logger.info("terminating impedance controller we launched") + self._impedance_proc.terminate() + try: + self._impedance_proc.wait(timeout=10) + except Exception: + self._impedance_proc.kill() + + +# --------------------------------------------------------------------------- +# Agent-facing env facade +# --------------------------------------------------------------------------- + + +class FrankaAgentEnv: + """Cartesian-primitive facade the agent RPCs into. + + Observation:: + + {"state": {"tcp_xyz": (3,) float, # meters, base frame + "tcp_quat": (4,) float, # [x, y, z, w] + "tcp_euler": (3,) float, # radians xyz + "gripper_width": float, # meters + "gripper_open": bool}, + "frames": {: (H, W, 3) uint8, ...}} + + All values are plain numpy / python scalars so they pickle across the RPC + wire (the agent process does not import torch or ROS). + """ + + def __init__( + self, + backend: FrankaRobotBackend, + cameras: list[RealSenseDepthCamera], + *, + workspace_min: np.ndarray = _WORKSPACE_MIN, + workspace_max: np.ndarray = _WORKSPACE_MAX, + reset_xyz: np.ndarray = _RESET_XYZ, + reset_euler: np.ndarray = _RESET_EULER, + step_frequency: float = _STEP_FREQUENCY, + max_step_m: float = _MAX_STEP_M, + max_step_deg: float = _MAX_STEP_DEG, + max_move_m: float = _MAX_MOVE_M, + settle_timeout_s: float = _SETTLE_TIMEOUT_S, + ): + self._backend = backend + self._cameras = cameras + self._workspace_min = np.asarray(workspace_min, dtype=np.float64) + self._workspace_max = np.asarray(workspace_max, dtype=np.float64) + self._reset_xyz = np.asarray(reset_xyz, dtype=np.float64) + self._reset_quat = _euler_to_quat(reset_euler) + self._step_frequency = float(step_frequency) + self._max_step_m = float(max_step_m) + self._max_step_deg = float(max_step_deg) + self._max_move_m = float(max_move_m) + self._settle_timeout_s = float(settle_timeout_s) + + # -- observation ------------------------------------------------------- + + def _frames(self) -> dict[str, np.ndarray]: + out: dict[str, np.ndarray] = {} + for cam in self._cameras: + try: + rgb, _ = cam.read() + out[cam.name] = rgb + except Exception as e: + logger.warning("camera '%s' frame grab failed: %s", cam.name, e) + return out + + def _camera_observation( + self, T_base_tcp: np.ndarray + ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray], dict[str, dict]]: + frames: dict[str, np.ndarray] = {} + depths: dict[str, np.ndarray] = {} + meta: dict[str, dict] = {} + for cam in self._cameras: + try: + rgb, depth = cam.read() + frames[cam.name] = rgb + depths[cam.name] = depth + meta[cam.name] = cam.meta(T_base_tcp=T_base_tcp) + except Exception as e: + logger.warning("camera '%s' RGB-D grab failed: %s", cam.name, e) + return frames, depths, meta + + def _make_obs(self) -> dict: + pose = self._backend.get_tcp_pose() + st = self._backend.get_state() + frames, depths, camera_meta = self._camera_observation(_pose_to_matrix(pose)) + return { + "state": { + "tcp_xyz": pose[:3].astype(np.float64), + "tcp_quat": pose[3:].astype(np.float64), + "tcp_euler": _quat_to_euler(pose[3:]).astype(np.float64), + "gripper_width": float(st["gripper_width"]), + "gripper_open": bool(st["gripper_open"]), + }, + "frames": frames, + "depth": depths, + "camera_meta": camera_meta, + } + + def get_obs(self) -> dict: + """Return the current observation without moving the arm.""" + return self._make_obs() + + def get_ee_pose(self) -> dict: + """Return the current end-effector pose in the base frame.""" + pose = self._backend.get_tcp_pose() + return { + "xyz": [round(float(v), 4) for v in pose[:3]], + "quat_xyzw": [round(float(v), 4) for v in pose[3:]], + "euler_xyz": [round(float(v), 4) for v in _quat_to_euler(pose[3:])], + "frame": "panda_link0", + } + + def get_spec(self) -> dict: + """Static self-description for the agent side.""" + return { + "world_frame": "panda_link0", + "control": "cartesian_impedance", + "position_unit": "meters", + "orientation": "euler_xyz_radians (also quat xyzw)", + "workspace_min": [round(float(v), 3) for v in self._workspace_min], + "workspace_max": [round(float(v), 3) for v in self._workspace_max], + "camera_names": [c.name for c in self._cameras], + "depth_camera_names": [c.name for c in self._cameras], + "preferred_backproject_camera": "wrist", + "gripper": "binary (open_gripper / close_gripper); width in meters", + "reset_xyz": [round(float(v), 3) for v in self._reset_xyz], + } + + def get_camera_meta(self) -> dict: + """Return live per-camera intrinsics and base-frame extrinsics if calibrated.""" + T_base_tcp = _pose_to_matrix(self._backend.get_tcp_pose()) + return {cam.name: cam.meta(T_base_tcp=T_base_tcp) for cam in self._cameras} + + def reload_camera_calibration(self, camera: str | None = None) -> dict: + """Reload camera calibration records from disk.""" + requested = None if camera is None else str(camera) + reloaded: dict[str, dict] = {} + for cam in self._cameras: + if requested is not None and cam.name != requested: + continue + reloaded[cam.name] = cam.reload_calibration() + if requested is not None and requested not in reloaded: + return { + "error": f"unknown camera {requested!r}", + "available_cameras": [cam.name for cam in self._cameras], + } + return reloaded + + # -- motion ------------------------------------------------------------ + + def _clip_xyz(self, xyz: np.ndarray) -> tuple[np.ndarray, bool]: + clipped = np.clip(xyz, self._workspace_min, self._workspace_max) + return clipped, bool(np.any(clipped != xyz)) + + def _stream_to(self, target_xyz: np.ndarray, target_quat: np.ndarray) -> None: + """Stream interpolated equilibrium poses from the current pose to the + target, capping per-step Cartesian and angular deltas and pacing at + ``step_frequency`` so the impedance controller tracks a slow, smooth path. + """ + cur = self._backend.get_tcp_pose() + cur_xyz, cur_quat = cur[:3], cur[3:] + + dist = float(np.linalg.norm(target_xyz - cur_xyz)) + ang = float( + (R.from_quat(cur_quat).inv() * R.from_quat(target_quat)).magnitude() + ) + n_pos = int(np.ceil(dist / self._max_step_m)) if dist > 0 else 0 + n_ang = int(np.ceil(np.degrees(ang) / self._max_step_deg)) if ang > 0 else 0 + n = max(1, n_pos, n_ang) + + slerp = Slerp([0.0, 1.0], R.concatenate([R.from_quat(cur_quat), R.from_quat(target_quat)])) + period = 1.0 / self._step_frequency + for i in range(1, n + 1): + t = i / n + pos = cur_xyz + (target_xyz - cur_xyz) * t + quat = slerp([t])[0].as_quat() + self._backend.move_arm(np.concatenate([pos, quat])) + time.sleep(period) + + def _wait_until_reached( + self, target_xyz: np.ndarray, target_quat: np.ndarray + ) -> tuple[np.ndarray, float, bool]: + """Hold the final setpoint until the measured TCP pose is close enough.""" + target_pose = np.concatenate([target_xyz, target_quat]) + deadline = time.time() + self._settle_timeout_s + period = 1.0 / self._step_frequency + while True: + final = self._backend.get_tcp_pose() + pos_err = float(np.linalg.norm(final[:3] - target_xyz)) + if pos_err <= _REACHED_TOL_M or time.time() >= deadline: + return final, pos_err, pos_err <= _REACHED_TOL_M + self._backend.move_arm(target_pose) + time.sleep(period) + + def _apply_gripper(self, gripper: Optional[str]) -> None: + if gripper is None: + return + g = str(gripper).lower() + if g in ("open", "release"): + self._backend.open_gripper() + time.sleep(0.6) + elif g in ("close", "grasp"): + self._backend.close_gripper() + time.sleep(0.6) + else: + raise ValueError(f"gripper must be 'open' or 'close', got {gripper!r}") + + def move_to( + self, + xyz, + *, + euler_xyz=None, + quat_xyzw=None, + gripper: Optional[str] = None, + ) -> dict: + """Move the TCP to a base-frame ``xyz`` (meters), holding the current + orientation unless ``euler_xyz`` / ``quat_xyzw`` is given. + + The target is clipped to the workspace box and the orientation to a safe + window; the path is streamed as slow capped setpoints. Optionally set the + gripper ("open"/"close") first. Returns a log dict. + """ + self._backend.clear_errors() + cur = self._backend.get_tcp_pose() + + target_xyz = np.asarray(xyz, dtype=np.float64).reshape(-1)[:3] + target_xyz, clipped = self._clip_xyz(target_xyz) + + path_len = float(np.linalg.norm(target_xyz - cur[:3])) + if path_len > self._max_move_m: + return { + "reached": False, + "error": ( + f"requested move of {path_len:.3f} m exceeds the {self._max_move_m} m " + "single-move safety cap; issue smaller moves." + ), + "current_xyz": [round(float(v), 4) for v in cur[:3]], + } + + if quat_xyzw is not None: + target_quat = np.asarray(quat_xyzw, dtype=np.float64).reshape(-1)[:4] + target_quat = _euler_to_quat(_clip_euler_window(_quat_to_euler(target_quat))) + elif euler_xyz is not None: + target_quat = _euler_to_quat(_clip_euler_window(euler_xyz)) + else: + target_quat = cur[3:] + + self._apply_gripper(gripper) + self._stream_to(target_xyz, target_quat) + final, pos_err, reached = self._wait_until_reached(target_xyz, target_quat) + return { + "reached": reached, + "pos_error_m": round(pos_err, 4), + "clipped_to_workspace": clipped, + "target_xyz": [round(float(v), 4) for v in target_xyz], + "final_xyz": [round(float(v), 4) for v in final[:3]], + "final_euler": [round(float(v), 4) for v in _quat_to_euler(final[3:])], + "gripper_open": bool(self._backend.get_state()["gripper_open"]), + } + + def move_delta( + self, + *, + dxyz=None, + drpy_deg=None, + gripper: Optional[str] = None, + ) -> dict: + """Nudge the TCP by a relative ``dxyz`` (meters) and/or ``drpy_deg`` + (degrees, applied in the base frame), for fine alignment. + """ + cur = self._backend.get_tcp_pose() + target_xyz = cur[:3].copy() + if dxyz is not None: + target_xyz = target_xyz + np.asarray(dxyz, dtype=np.float64).reshape(-1)[:3] + + euler = _quat_to_euler(cur[3:]) + if drpy_deg is not None: + euler = euler + np.radians(np.asarray(drpy_deg, dtype=np.float64).reshape(-1)[:3]) + + return self.move_to( + target_xyz, euler_xyz=euler, gripper=gripper + ) + + def open_gripper(self) -> dict: + self._backend.open_gripper() + time.sleep(0.6) + return {"gripper_open": True, "gripper_width": self._backend.get_state()["gripper_width"]} + + def close_gripper(self) -> dict: + self._backend.close_gripper() + time.sleep(0.6) + return {"gripper_open": False, "gripper_width": self._backend.get_state()["gripper_width"]} + + # -- reset / teardown -------------------------------------------------- + + def reset(self) -> tuple[dict, dict]: + """Clear errors and drive the arm to its home pose; return ``(obs, {})``.""" + self._backend.clear_errors() + self._stream_to(self._reset_xyz, self._reset_quat) + time.sleep(0.5) + return self._make_obs(), {} + + def close(self) -> None: + for cam in self._cameras: + cam.close() + self._backend.shutdown() + + +# --------------------------------------------------------------------------- +# RPC dispatcher + parent watchdog (mirrors the LIBERO / LeRobot drivers) +# --------------------------------------------------------------------------- + + +_INITIAL_PPID = os.getppid() + + +def _start_parent_watchdog( + server: SocketRpcServer, shutdown_event: threading.Event, poll_s: float = 2.0 +) -> None: + """Shut the RPC server down if the agent (parent) process dies.""" + + def _watch() -> None: + while not shutdown_event.is_set(): + time.sleep(poll_s) + ppid = os.getppid() + if ppid != _INITIAL_PPID or ppid == 1: + logger.warning("parent died (ppid %s -> %s); stopping", _INITIAL_PPID, ppid) + shutdown_event.set() + threading.Thread(target=server.shutdown, daemon=True).start() + return + + threading.Thread(target=_watch, daemon=True).start() + + +def _build_dispatcher(env: FrankaAgentEnv, shutdown_event: threading.Event): + def dispatch(method: str, args: tuple, kwargs: dict): + if method.startswith("env."): + attr = method[len("env."):] + try: + return getattr(env, attr)(*args, **kwargs) + except Exception as e: + logger.warning("env method %s failed: %s", method, e) + raise + if method == "shutdown": + shutdown_event.set() + return {"ok": True} + raise ValueError(f"unknown RPC method: {method!r}") + + return dispatch + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def _parse_cameras(specs: list[str] | None) -> list[tuple[str, str]]: + if not specs: + return list(_DEFAULT_CAMERAS) + out: list[tuple[str, str]] = [] + for spec in specs: + name, _, serial = spec.partition(":") + if not name or not serial: + raise ValueError(f"--camera expects name:serial, got {spec!r}") + out.append((name, serial)) + return out + + +def _build_cameras(specs: list[tuple[str, str]]) -> list[RealSenseDepthCamera]: + cams: list[RealSenseDepthCamera] = [] + for name, serial in specs: + try: + cams.append( + RealSenseDepthCamera( + name, serial, width=_CAM_WIDTH, height=_CAM_HEIGHT, fps=_CAM_FPS + ) + ) + except Exception as e: + logger.warning("could not open camera '%s' (serial %s): %s", name, serial, e) + return cams + + +def main() -> int: + p = argparse.ArgumentParser(description="Standalone Franka env server") + p.add_argument("--output-dir", type=str, required=True) + p.add_argument("--robot-ip", type=str, default=_DEFAULT_ROBOT_IP) + p.add_argument("--ros-pkg", type=str, default=_DEFAULT_ROS_PKG) + p.add_argument("--no-gripper", action="store_true", help="Do not load the Franka Hand.") + p.add_argument( + "--camera", action="append", default=None, + help="Camera as name:serial (repeatable). Defaults to the two bench D435I.", + ) + p.add_argument("--transport-host", type=str, default="127.0.0.1") + p.add_argument("--transport-port", type=int, default=0) + args = p.parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + init_output_dir(args.output_dir) + logger.info( + "starting Franka env server: robot_ip=%s output_dir=%s", + args.robot_ip, args.output_dir, + ) + + backend = FrankaRobotBackend( + robot_ip=args.robot_ip, + ros_pkg=args.ros_pkg, + load_gripper=not args.no_gripper, + ) + cameras = _build_cameras(_parse_cameras(args.camera)) + env = FrankaAgentEnv(backend, cameras) + + shutdown_event = threading.Event() + dispatch = _build_dispatcher(env, shutdown_event) + server = SocketRpcServer((args.transport_host, args.transport_port), dispatch) + bound_host, bound_port = server.server_address + client_host = "127.0.0.1" if bound_host == "0.0.0.0" else bound_host + print( + json.dumps({ + "event": "transport_ready", "kind": "socket", + "host": client_host, "port": bound_port, + }), + flush=True, + ) + logger.info("RPC server listening on %s:%s", client_host, bound_port) + + _start_parent_watchdog(server, shutdown_event) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + shutdown_event.wait() + finally: + server.shutdown() + server.server_close() + env.close() + logger.info("driver exited cleanly") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/deployment/franka/generate_fiducial_board.py b/deployment/franka/generate_fiducial_board.py new file mode 100644 index 00000000..63f72ac8 --- /dev/null +++ b/deployment/franka/generate_fiducial_board.py @@ -0,0 +1,190 @@ +#!/usr/bin/env python3 +"""Generate a printable ChArUco fiducial board for Franka camera calibration. + +The output files are: + +- ``.png``: high-resolution board image with DPI metadata. +- ``.pdf``: printable single-page PDF at the same physical scale. +- ``.json``: board spec consumed by future calibration scripts. + +Print the PDF at 100% / actual size, not "fit to page". After printing, measure +one square and use the measured value if it differs from the requested size. + +Dependency: + uv pip install opencv-contrib-python-headless +""" +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image + +_REPO_ROOT = Path(__file__).resolve().parents[2] +_DEFAULT_OUT_DIR = _REPO_ROOT / "resources" / "franka" / "calibration_boards" + + +def _require_cv2_aruco(): + try: + import cv2 + except ModuleNotFoundError as exc: + raise SystemExit( + "Missing dependency: cv2. Install with:\n" + " uv pip install opencv-contrib-python-headless\n" + "or:\n" + " pip install opencv-contrib-python-headless" + ) from exc + if not hasattr(cv2, "aruco"): + raise SystemExit( + "Installed cv2 lacks the aruco module. Install opencv-contrib, not plain opencv:\n" + " uv pip install opencv-contrib-python-headless" + ) + return cv2 + + +def _aruco_dictionary(cv2, name: str): + aruco = cv2.aruco + key = f"DICT_{name.upper()}" if not name.upper().startswith("DICT_") else name.upper() + if not hasattr(aruco, key): + available = sorted(attr[5:] for attr in dir(aruco) if attr.startswith("DICT_")) + raise SystemExit(f"Unknown ArUco dictionary {name!r}. Available examples: {available[:12]}") + return aruco.getPredefinedDictionary(getattr(aruco, key)), key + + +def _make_charuco_board( + cv2, + *, + squares_x: int, + squares_y: int, + square_px: int, + marker_px: int, + dictionary, +): + aruco = cv2.aruco + try: + return aruco.CharucoBoard( + (int(squares_x), int(squares_y)), + float(square_px), + float(marker_px), + dictionary, + ) + except Exception: + return aruco.CharucoBoard_create( + int(squares_x), + int(squares_y), + float(square_px), + float(marker_px), + dictionary, + ) + + +def _draw_board(board, image_size: tuple[int, int], margin_px: int) -> np.ndarray: + """Draw a ChArUco board across OpenCV API versions.""" + width, height = image_size + if hasattr(board, "generateImage"): + img = board.generateImage((width, height), marginSize=int(margin_px), borderBits=1) + else: + img = board.draw((width, height), marginSize=int(margin_px), borderBits=1) + img = np.asarray(img, dtype=np.uint8) + if img.ndim == 3: + img = img[:, :, 0] + return img + + +def _save_outputs( + image: np.ndarray, + *, + out_dir: Path, + name: str, + dpi: int, + spec: dict[str, Any], +) -> dict[str, str]: + out_dir.mkdir(parents=True, exist_ok=True) + pil = Image.fromarray(image, mode="L") + png_path = out_dir / f"{name}.png" + pdf_path = out_dir / f"{name}.pdf" + json_path = out_dir / f"{name}.json" + pil.save(png_path, dpi=(dpi, dpi)) + pil.save(pdf_path, "PDF", resolution=float(dpi)) + json_path.write_text(json.dumps(spec, indent=2)) + return {"png": str(png_path), "pdf": str(pdf_path), "json": str(json_path)} + + +def _build_argparser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--out-dir", type=Path, default=_DEFAULT_OUT_DIR) + parser.add_argument("--name", default="franka_charuco_7x5_25mm") + parser.add_argument("--squares-x", type=int, default=7) + parser.add_argument("--squares-y", type=int, default=5) + parser.add_argument("--square-mm", type=float, default=25.0) + parser.add_argument("--marker-mm", type=float, default=18.0) + parser.add_argument("--margin-mm", type=float, default=12.0) + parser.add_argument("--dpi", type=int, default=300) + parser.add_argument("--dictionary", default="4X4_50", help="OpenCV ArUco dictionary suffix, e.g. 4X4_50 or APRILTAG_36h11.") + parser.add_argument("--check-deps", action="store_true", help="Only check for cv2.aruco and exit.") + return parser + + +def main() -> int: + args = _build_argparser().parse_args() + cv2 = _require_cv2_aruco() + if args.check_deps: + print(f"cv2 {cv2.__version__} with aruco: OK") + return 0 + + if args.squares_x < 2 or args.squares_y < 2: + raise SystemExit("--squares-x and --squares-y must be >= 2") + if not (0 < args.marker_mm < args.square_mm): + raise SystemExit("--marker-mm must be >0 and < --square-mm") + + px_per_mm = float(args.dpi) / 25.4 + square_px = int(round(args.square_mm * px_per_mm)) + marker_px = int(round(args.marker_mm * px_per_mm)) + margin_px = int(round(args.margin_mm * px_per_mm)) + board_width_px = args.squares_x * square_px + board_height_px = args.squares_y * square_px + image_size = (board_width_px + 2 * margin_px, board_height_px + 2 * margin_px) + + dictionary, dictionary_name = _aruco_dictionary(cv2, args.dictionary) + board = _make_charuco_board( + cv2, + squares_x=args.squares_x, + squares_y=args.squares_y, + square_px=square_px, + marker_px=marker_px, + dictionary=dictionary, + ) + image = _draw_board(board, image_size, margin_px) + + spec = { + "type": "charuco", + "dictionary": dictionary_name, + "squares_x": args.squares_x, + "squares_y": args.squares_y, + "square_length_m": args.square_mm / 1000.0, + "marker_length_m": args.marker_mm / 1000.0, + "margin_m": args.margin_mm / 1000.0, + "dpi": args.dpi, + "image_width_px": int(image_size[0]), + "image_height_px": int(image_size[1]), + "print_instructions": [ + "Print the PDF at 100% / actual size.", + "Disable fit-to-page or scaling.", + "Use matte paper and mount it flat to cardboard/foam board.", + "Measure one printed square and update square_length_m if needed.", + ], + } + paths = _save_outputs(image, out_dir=args.out_dir, name=args.name, dpi=args.dpi, spec=spec) + + physical_w_mm = image_size[0] / px_per_mm + physical_h_mm = image_size[1] / px_per_mm + print(json.dumps({"paths": paths, "physical_size_mm": [round(physical_w_mm, 1), round(physical_h_mm, 1)], "spec": spec}, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/deployment/franka/run_env_server.sh b/deployment/franka/run_env_server.sh new file mode 100755 index 00000000..498c7312 --- /dev/null +++ b/deployment/franka/run_env_server.sh @@ -0,0 +1,41 @@ +#!/usr/bin/env bash +# Launch the standalone Franka env server in the RLinf .venv with the +# serl_franka_controllers catkin workspace sourced. +# +# The agent (physical/.venv) connects to this over TCP. Two ways to use it: +# +# 1. Fixed port, then run the agent with --no-driver: +# bash deployment/franka/run_env_server.sh --transport-port 5599 +# # in the physical/.venv: +# python -m cli.main --env franka --no-driver --env-port 5599 ... +# +# 2. Let cli.main spawn it (it invokes this script); see start_franka_env_server. +# +# Override the machine-specific bits via env vars: +# FRANKA_CATKIN_SETUP catkin devel setup.bash (default: RLinf .venv workspace) +# RLINF_VENV_PYTHON python in the RLinf .venv +# FRANKA_ROBOT_IP robot IP (default 172.16.0.2) +set -euo pipefail + +FRANKA_CATKIN_SETUP="${FRANKA_CATKIN_SETUP:-/home/franka/franka/RLinf/.venv/franka_catkin_ws/devel/setup.bash}" +RLINF_VENV_PYTHON="${RLINF_VENV_PYTHON:-/home/franka/franka/RLinf/.venv/bin/python}" +FRANKA_ROBOT_IP="${FRANKA_ROBOT_IP:-172.16.0.2}" + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ENV_SERVER="${SCRIPT_DIR}/env_server.py" + +if [[ ! -f "${FRANKA_CATKIN_SETUP}" ]]; then + echo "!! catkin setup not found: ${FRANKA_CATKIN_SETUP}" >&2 + echo " set FRANKA_CATKIN_SETUP to your serl_franka_controllers workspace." >&2 + exit 1 +fi +if [[ ! -x "${RLINF_VENV_PYTHON}" ]]; then + echo "!! RLinf venv python not found: ${RLINF_VENV_PYTHON}" >&2 + echo " set RLINF_VENV_PYTHON to the RLinf .venv python." >&2 + exit 1 +fi + +# shellcheck disable=SC1090 +source "${FRANKA_CATKIN_SETUP}" + +exec "${RLINF_VENV_PYTHON}" "${ENV_SERVER}" --robot-ip "${FRANKA_ROBOT_IP}" "$@" diff --git a/deployment/lerobot/auto_calibrate_scene_cam.py b/deployment/lerobot/auto_calibrate_scene_cam.py new file mode 100644 index 00000000..2a2a2e75 --- /dev/null +++ b/deployment/lerobot/auto_calibrate_scene_cam.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python3 +"""Automatic, markerless scene-camera -> base calibration for the SO101. + +Triggers the env server's :meth:`auto_calibrate_scene_camera` routine, which: + +1. drives the gripper to a spread grid of base-frame positions (``move_to`` is + pure base-frame IK, so it needs no extrinsic), +2. at each pose toggles the gripper with the arm frozen and segments the motion + in the scene image to locate the tip (centroid + median depth -> camera + point); the achieved FK gives the base point, +3. fits ``T_base_cam`` with RANSAC Kabsch and saves it (hot-loaded by the + server, so back_project returns world coords immediately). + +No human input, no markers. Start the env server first, then run:: + + conda activate lerobot + python deployment/lerobot/auto_calibrate_scene_cam.py --port 53101 + +WARNING: this moves the arm through many poses. Clear the workspace first. + +Offline math check (no hardware):: + + python deployment/lerobot/auto_calibrate_scene_cam.py --self-test +""" +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import numpy as np + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from deployment.lerobot import geometry as geom # noqa: E402 +from rpent.rpc_driver.socket import SocketRpcClient # noqa: E402 + + +def _self_test() -> int: + """Validate RANSAC Kabsch + motion segmentation offline (no hardware).""" + rng = np.random.default_rng(0) + Q, _ = np.linalg.qr(rng.standard_normal((3, 3))) + if np.linalg.det(Q) < 0: + Q[:, 0] = -Q[:, 0] + T_true = np.eye(4) + T_true[:3, :3] = Q + T_true[:3, 3] = rng.standard_normal(3) + cam = rng.standard_normal((10, 3)) + base = geom.transform_points(T_true, cam) + rng.standard_normal((10, 3)) * 1e-3 + base[3] += [0.2, -0.15, 0.1] # inject an outlier + T_est, rmse, inliers = geom.ransac_kabsch(cam, base, thresh_m=0.02) + ok_fit = np.allclose(T_est, T_true, atol=2e-2) and not inliers[3] + print(f"ransac_kabsch: rmse={rmse:.5f}m inliers={int(inliers.sum())}/10 " + f"outlier_excluded={not inliers[3]} -> {'OK' if ok_fit else 'FAIL'}") + + # Synthetic two-frame motion: a blob that appears in the 'closed' frame. + H, W = 480, 640 + rgb_open = np.zeros((H, W, 3), np.uint8) + rgb_closed = rgb_open.copy() + rgb_closed[300:330, 400:430] = 255 # 'fingers' light up at (row~315,col~415) + depth = np.full((H, W), 0.4, np.float32) + K = np.array([[600, 0, 320], [0, 600, 240], [0, 0, 1]], float) + det = geom.detect_tip_pixel_by_motion(rgb_open, rgb_closed, depth, K) + ok_det = det is not None and abs(det["pixel"][0] - 314.5) < 3 and abs(det["pixel"][1] - 414.5) < 3 + print(f"detect_tip: {det if det else 'None'} -> {'OK' if ok_det else 'FAIL'}") + return 0 if (ok_fit and ok_det) else 1 + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--host", default="127.0.0.1") + ap.add_argument("--port", type=int, help="env server transport port.") + ap.add_argument("--n-points", type=int, default=10, + help="Target number of valid correspondences to collect.") + ap.add_argument("--no-save", action="store_true", + help="Compute T_base_cam but do not write it to disk.") + ap.add_argument("--self-test", action="store_true", + help="Run the offline math check and exit.") + args = ap.parse_args() + + if args.self_test: + return _self_test() + if args.port is None: + ap.error("--port is required (the env server's transport port).") + + print("This moves the arm through a grid of poses. Ensure the workspace is " + "clear. Starting...") + client = SocketRpcClient(args.host, args.port) + result = client.call( + "env.auto_calibrate_scene_camera", + kwargs={"n_points": args.n_points, "save": not args.no_save}, + timeout_s=600.0, + ) + + if "error" in result: + print(f"Calibration failed: {result['error']}") + if "poses" in result: + print(json.dumps(result["poses"], indent=2)) + return 2 + + print(f"\nUsed {result['n_used']} poses " + f"({result['n_inliers']} inliers), RMSE = {result['rmse_m'] * 1000:.1f} mm") + if result["rmse_m"] > 0.02: + print("WARNING: RMSE > 2 cm — check lighting / gripper visibility and rerun.") + if result.get("saved"): + print(f"Saved T_base_cam -> {result['path']}") + print("The server hot-loaded it; back_project now returns world coords.") + else: + print("Not saved (--no-save). T_base_cam:") + print(json.dumps(result["T_base_cam"], indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/deployment/lerobot/calibrate_scene_cam.py b/deployment/lerobot/calibrate_scene_cam.py new file mode 100644 index 00000000..6b66beb3 --- /dev/null +++ b/deployment/lerobot/calibrate_scene_cam.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +"""Touch / point-correspondence calibration of the scene camera -> arm base. + +Computes the fixed extrinsic ``T_base_cam`` that maps scene-camera points into +the SO101 ``base_link`` world frame, using only the arm's own FK + the scene +camera's aligned depth (no marker / no extra hardware). The result is saved via +:mod:`deployment.lerobot.calibration` and auto-loaded by the env server, after +which ``back_project`` returns world coordinates. + +Procedure (per correspondence): + +1. Free-drive the arm by hand so the gripper tip (``gripper_frame_link``, + roughly the point between the fingertips) rests at a distinct location that + is clearly visible to the scene camera. +2. The script reads the tip position in the base frame from FK + (``env.get_ee_pose``) and grabs the scene color + aligned depth + (``env.get_scene_frame``). +3. You click that same tip point in the color image; the script backprojects + the clicked pixel (median depth over a small patch) into the camera frame. + +After N>=4 non-coplanar points, a rigid Kabsch fit gives ``T_base_cam`` and the +fit RMSE (lower is better; aim for < ~1 cm). + +Run the env server first (note its --transport-port), then:: + + conda activate lerobot + python toolkits/lerobot/calibrate_scene_cam.py --port 53101 --num-points 6 + +Offline self-test of the math (no hardware):: + + python toolkits/lerobot/calibrate_scene_cam.py --self-test +""" +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import numpy as np + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from deployment.lerobot import calibration as scene_calib # noqa: E402 +from deployment.lerobot import geometry as geom # noqa: E402 +from rpent.rpc_driver.socket import SocketRpcClient # noqa: E402 + + +def _self_test() -> int: + """Validate the Kabsch pipeline on synthetic correspondences (no hardware).""" + rng = np.random.default_rng(0) + A = rng.standard_normal((3, 3)) + Q, _ = np.linalg.qr(A) + if np.linalg.det(Q) < 0: + Q[:, 0] = -Q[:, 0] + T_true = np.eye(4) + T_true[:3, :3] = Q + T_true[:3, 3] = rng.standard_normal(3) + + cam_pts = rng.standard_normal((8, 3)) + base_pts = geom.transform_points(T_true, cam_pts) + rng.standard_normal((8, 3)) * 1e-3 + T_est, rmse = geom.kabsch_umeyama(cam_pts, base_pts) + ok = np.allclose(T_est, T_true, atol=1e-2) + print(f"self-test: rmse={rmse:.5f}m recovered={'OK' if ok else 'FAIL'}") + return 0 if ok else 1 + + +def _click_pixel(color: np.ndarray, idx: int, total: int) -> tuple[float, float] | None: + """Show the color frame and return the clicked (col, row), or None if skipped.""" + import matplotlib.pyplot as plt + + fig = plt.figure(figsize=(8, 6)) + plt.imshow(color) + plt.title(f"[{idx}/{total}] Click the gripper TIP, then close is automatic. " + "Close window to skip.") + plt.tight_layout() + pts = plt.ginput(1, timeout=0) + plt.close(fig) + if not pts: + return None + return float(pts[0][0]), float(pts[0][1]) + + +def _collect(client: SocketRpcClient, num_points: int, patch_radius: int) -> tuple[np.ndarray, np.ndarray]: + """Collect (cam_point, base_point) correspondences interactively.""" + cam_pts: list[list[float]] = [] + base_pts: list[list[float]] = [] + + i = 0 + while len(cam_pts) < num_points: + i += 1 + input( + f"\n[{len(cam_pts) + 1}/{num_points}] Move the gripper tip to a distinct " + "scene point (vary x/y/z), hold it, then press Enter to capture..." + ) + ee = client.call("env.get_ee_pose", timeout_s=15) + if "error" in ee: + print(f" get_ee_pose failed: {ee['error']}") + continue + frame = client.call("env.get_scene_frame", timeout_s=15) + if "error" in frame: + print(f" get_scene_frame failed: {frame['error']}") + continue + + color = np.asarray(frame["color"], dtype=np.uint8) + depth = np.asarray(frame["depth"], dtype=np.float32) + K = np.asarray(frame["K"], dtype=np.float64) + + click = _click_pixel(color, len(cam_pts) + 1, num_points) + if click is None: + print(" skipped (no pixel clicked).") + continue + col, row = click + z = geom.sample_depth_patch(depth, int(round(col)), int(round(row)), radius=patch_radius) + if not np.isfinite(z) or z <= 0: + print(f" no valid depth at ({int(row)},{int(col)}); try another point/angle.") + continue + + p_cam = geom.backproject_pixel(K, col, row, z) + p_base = np.asarray(ee["xyz"], dtype=np.float64) + cam_pts.append(p_cam.tolist()) + base_pts.append(p_base.tolist()) + print(f" captured: cam={np.round(p_cam, 3).tolist()} " + f"base={np.round(p_base, 3).tolist()} depth={z:.3f}m") + + return np.asarray(cam_pts), np.asarray(base_pts) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--host", default="127.0.0.1", help="env server host.") + ap.add_argument("--port", type=int, help="env server transport port.") + ap.add_argument("--num-points", type=int, default=6, + help="Number of correspondences (>=4; more is better).") + ap.add_argument("--patch-radius", type=int, default=2, + help="Depth median patch radius (pixels) around each click.") + ap.add_argument("--serial", default=None, + help="Override scene serial for the saved file " + "(default: from env.get_scene_camera_meta).") + ap.add_argument("--self-test", action="store_true", + help="Run the offline Kabsch math check and exit.") + args = ap.parse_args() + + if args.self_test: + return _self_test() + if args.port is None: + ap.error("--port is required (the env server's transport port).") + if args.num_points < 4: + ap.error("need at least 4 correspondences for a stable fit.") + + client = SocketRpcClient(args.host, args.port) + meta = client.call("env.get_scene_camera_meta", timeout_s=15) + if "error" in meta: + print(f"scene camera not available: {meta['error']}") + return 2 + serial = args.serial or meta.get("serial") + if not serial: + print("could not determine scene camera serial; pass --serial.") + return 2 + print(f"Calibrating scene camera serial={serial}") + + # Free-drive so the operator can position the tip by hand. + client.call("env.set_torque", args=(False,), timeout_s=15) + print("Arm torque DISABLED — you can move it by hand. (Re-enabled at the end.)") + try: + cam_pts, base_pts = _collect(client, args.num_points, args.patch_radius) + finally: + client.call("env.set_torque", args=(True,), timeout_s=15) + print("Arm torque re-enabled.") + + T_base_cam, rmse = geom.kabsch_umeyama(cam_pts, base_pts) + print(f"\nFit complete: {len(cam_pts)} points, RMSE = {rmse * 1000:.1f} mm") + if rmse > 0.02: + print("WARNING: RMSE > 2 cm — consider recollecting with more spread / " + "better tip-pixel clicks.") + + path = scene_calib.save_extrinsic( + serial, T_base_cam, K=np.asarray(meta["K"]), rmse_m=rmse, num_points=len(cam_pts) + ) + print(f"Saved T_base_cam -> {path}") + print("Restart the env server (or it will pick this up next launch) so " + "back_project returns world coordinates.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/deployment/lerobot/calibration.py b/deployment/lerobot/calibration.py new file mode 100644 index 00000000..761580c3 --- /dev/null +++ b/deployment/lerobot/calibration.py @@ -0,0 +1,78 @@ +"""Load / save the scene-camera → base extrinsic ``T_base_cam``. + +``T_base_cam`` is the fixed rigid transform that maps a point in the scene +camera frame into the arm ``base_link`` world frame. It is produced once by the +touch/Kabsch calibration (``toolkits/lerobot/calibrate_scene_cam.py``) and +loaded by the env server so ``back_project`` can return world coordinates. + +Stored per camera serial under the LeRobot cache so it survives across runs. +""" +from __future__ import annotations + +import json +import os +from pathlib import Path + +import numpy as np + +_CALIB_DIR = "~/.cache/huggingface/lerobot/calibration/scene_cam" + +# A scene-cam -> base fit whose RMSE (meters) exceeds this is treated as a +# FAILED calibration: ``back_project`` would map pixels to badly wrong world +# coordinates (the arm then chases unreachable targets), so the driver refuses +# to save it or trust it on load. ~2 cm matches the auto-calibrator's existing +# "rerun" warning; a good touch/Kabsch fit is typically a few mm. +MAX_ACCEPTABLE_RMSE_M = 0.02 + + +def calib_path(serial: str) -> Path: + """Return the on-disk path of the extrinsic file for a camera ``serial``.""" + return Path(os.path.expanduser(_CALIB_DIR)) / f"{serial}.json" + + +def save_extrinsic( + serial: str, + T_base_cam, + *, + K=None, + rmse_m: float | None = None, + num_points: int | None = None, +) -> Path: + """Persist ``T_base_cam`` (and calibration diagnostics) for ``serial``.""" + path = calib_path(serial) + path.parent.mkdir(parents=True, exist_ok=True) + data: dict = { + "serial": serial, + "T_base_cam": np.asarray(T_base_cam, dtype=np.float64).tolist(), + } + if K is not None: + data["K"] = np.asarray(K, dtype=np.float64).tolist() + if rmse_m is not None: + data["rmse_m"] = float(rmse_m) + if num_points is not None: + data["num_points"] = int(num_points) + with open(path, "w") as f: + json.dump(data, f, indent=2) + return path + + +def load_extrinsic_record(serial: str) -> dict | None: + """Load the full saved calibration record for ``serial`` (or ``None``). + + Returns the parsed JSON with ``T_base_cam`` as a ``(4, 4)`` ndarray plus any + saved diagnostics (``rmse_m``, ``num_points``, ``K``). Use this when you need + to judge fit quality, not just apply the transform. + """ + path = calib_path(serial) + if not path.is_file(): + return None + with open(path) as f: + data = json.load(f) + data["T_base_cam"] = np.asarray(data["T_base_cam"], dtype=np.float64) + return data + + +def load_extrinsic(serial: str) -> np.ndarray | None: + """Load ``T_base_cam`` (4x4) for ``serial``, or ``None`` if not calibrated.""" + record = load_extrinsic_record(serial) + return None if record is None else record["T_base_cam"] diff --git a/deployment/lerobot/diagnose_motors.py b/deployment/lerobot/diagnose_motors.py new file mode 100644 index 00000000..49776d08 --- /dev/null +++ b/deployment/lerobot/diagnose_motors.py @@ -0,0 +1,173 @@ +"""Motor-health diagnostic for the SO101 follower arm. + +Reads each servo's temperature, voltage, current, load, hardware-error status, +and torque state WITHOUT moving the arm. Use it when the arm stops reaching +commanded positions (drifts / presses the table / grasps miss even though the +scene localization and IK are correct) to check whether a joint is overheated, +under-volted, faulted, or straining. + +Optionally (``--tracking-test``) it nudges each arm joint a few degrees around +its CURRENT pose and measures how far the achieved angle is from the command, +to pinpoint a joint that no longer tracks (slipped horn, weak motor, or +calibration drift). That part MOVES the arm, so put the arm in a safe, raised +pose first. + +Run with the env server STOPPED (it needs exclusive access to the motor bus):: + + conda activate lerobot + python deployment/lerobot/diagnose_motors.py # health only, no motion + python deployment/lerobot/diagnose_motors.py --tracking-test # also nudges each joint +""" +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +# Make ``rpent`` importable if ever needed; also keeps paths consistent +# with the driver when run from the repo root. +_ROOT = Path(__file__).resolve().parents[2] +if str(_ROOT) not in sys.path: + sys.path.insert(0, str(_ROOT)) + +_ARM_JOINTS = ("shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll") +_ALL_MOTORS = _ARM_JOINTS + ("gripper",) + + +def _read(bus, name: str, motor: str): + """Read one raw register value, returning a short error string on failure.""" + try: + return bus.read(name, motor, normalize=False) + except Exception as e: # noqa: BLE001 - diagnostic must never crash on one bad read + return f"err:{type(e).__name__}" + + +def _fmt(v, width: int) -> str: + return f"{v:>{width}}" if not isinstance(v, float) else f"{v:>{width}.1f}" + + +def _set_p_coefficient(robot, p: int) -> None: + """Set the position gain on the arm joints (to test tracking stiffness).""" + try: + with robot.bus.torque_disabled(): + for m in _ARM_JOINTS: + robot.bus.write("P_Coefficient", m, int(p)) + print(f"[set P_Coefficient={int(p)} on arm joints for this test]\n") + except Exception as e: # noqa: BLE001 + print(f"[could not set P_Coefficient: {e}]\n") + + +def _health(robot) -> None: + bus = robot.bus + try: + obs = robot.get_observation() + except Exception: + obs = {} + print("\n=== motor health (no motion) ===") + print( + f"{'motor':16}{'pos_deg':>9}{'temp_C':>8}{'volt_V':>8}" + f"{'load':>8}{'current':>9}{'status':>8}{'torque':>8}{'Pgain':>7}" + ) + for m in _ALL_MOTORS: + temp = _read(bus, "Present_Temperature", m) # deg C + volt = _read(bus, "Present_Voltage", m) # 0.1 V units + load = _read(bus, "Present_Load", m) + curr = _read(bus, "Present_Current", m) + stat = _read(bus, "Status", m) # 0 = OK; nonzero = HW error flags + torq = _read(bus, "Torque_Enable", m) + pgain = _read(bus, "P_Coefficient", m) + pos = obs.get(f"{m}.pos") + volt_v = round(volt / 10.0, 1) if isinstance(volt, (int, float)) else volt + pos_s = f"{pos:9.1f}" if isinstance(pos, (int, float)) else f"{str(pos):>9}" + print( + f"{m:16}{pos_s}{_fmt(temp, 8)}{_fmt(volt_v, 8)}" + f"{_fmt(load, 8)}{_fmt(curr, 9)}{_fmt(stat, 8)}{_fmt(torq, 8)}{_fmt(pgain, 7)}" + ) + print( + "\nInterpretation:\n" + " temp_C > ~55 -> overheating; Feetech servos derate torque and under-reach.\n" + " volt_V -> should match your supply and be steady; sag => weak torque.\n" + " status != 0 -> a hardware-error flag latched (overload/overheat/voltage).\n" + " load/current -> high while merely holding a light pose => a straining joint.\n" + ) + + +def _tracking_test(robot, nudge_deg: float) -> None: + base = robot.get_observation() + q0 = {m: float(base[f"{m}.pos"]) for m in _ARM_JOINTS} + g0 = float(base.get("gripper.pos", 50.0)) + + def send(qd: dict) -> None: + act = {f"{m}.pos": float(qd[m]) for m in _ARM_JOINTS} + act["gripper.pos"] = g0 + robot.send_action(act) + + print("=== per-joint tracking test (MOVES the arm ±%.0f deg around the current pose) ===" % nudge_deg) + worst = 0.0 + for j in _ARM_JOINTS: + for sgn in (+1.0, -1.0): + qd = dict(q0) + qd[j] = q0[j] + sgn * nudge_deg + send(qd) + time.sleep(0.9) + ach = float(robot.get_observation()[f"{j}.pos"]) + err = ach - qd[j] + worst = max(worst, abs(err)) + flag = " <-- POOR TRACKING" if abs(err) > 3.0 else "" + print(f" {j:16} cmd={qd[j]:7.1f} achieved={ach:7.1f} err={err:+6.1f} deg{flag}") + send(q0) + time.sleep(0.9) + print( + f"\nworst tracking error: {worst:.1f} deg\n" + " > ~3 deg on a free-space nudge = that joint is not tracking " + "(slipped horn / weak or faulted motor / calibration drift). Recalibrate " + "the follower or inspect that motor before running the agent again.\n" + ) + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--port", default="/dev/ttyACM1") + p.add_argument("--calibration-id", default="my_awesome_follower_arm") + p.add_argument( + "--tracking-test", action="store_true", + help="Also nudge each arm joint a few degrees and report the " + "command-vs-readback error. MOVES THE ARM — place it in a safe, " + "raised pose first.", + ) + p.add_argument("--nudge-deg", type=float, default=6.0) + p.add_argument( + "--p-coefficient", type=int, default=None, + help="Set the arm servos' P_Coefficient (position gain) before testing " + "(LeRobot uses 16; factory default 32). Sweep e.g. 32 then 48 to " + "see if tracking tightens.", + ) + args = p.parse_args() + + from lerobot.robots.so_follower import SO101Follower + from lerobot.robots.so_follower.config_so_follower import SO101FollowerConfig + + robot = SO101Follower( + SO101FollowerConfig( + port=args.port, + id=args.calibration_id, + use_degrees=True, + disable_torque_on_disconnect=False, + cameras={}, + ) + ) + robot.connect(calibrate=False) + try: + if args.p_coefficient is not None: + _set_p_coefficient(robot, args.p_coefficient) + _health(robot) + if args.tracking_test: + _tracking_test(robot, args.nudge_deg) + finally: + robot.disconnect() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/deployment/lerobot/env_server.py b/deployment/lerobot/env_server.py new file mode 100644 index 00000000..f22c64fc --- /dev/null +++ b/deployment/lerobot/env_server.py @@ -0,0 +1,1311 @@ +"""Standalone LeRobot SO101 env host for the LLM-in-the-loop agent (env-only). + +Drives a physical SO101 follower arm through LeRobot's synchronous Python +API (:class:`lerobot.robots.so_follower.SO101Follower`) and exposes a minimal +``reset`` / ``step`` gym-style surface over a pickle-framed TCP RPC server +(:class:`rpent.rpc_driver.socket.SocketRpcServer`) — the same wire +protocol the LIBERO driver uses, so the agent side talks to both identically. + +Unlike the LIBERO driver, this server does **not** wrap an RLinf env class: +importing ``rlinf.envs.realworld`` runs node-level ROS setup side effects at +import time (it kills ``roscore``/``rosmaster`` processes), which is +inappropriate for a standalone driver. Instead this file talks to LeRobot +directly, reusing the device / action / observation recipe from RLinf's +``rlinf.envs.realworld.so101.SO101Env``. + +Run it inside the ``lerobot`` conda env:: + + conda activate lerobot + python deployment/lerobot/env_server.py --output-dir /tmp/so101_run + +Hardware defaults match the current bench setup: follower on ``/dev/ttyACM1`` +(calibration id ``my_awesome_follower_arm``), an OpenCV hand/arm camera on +``/dev/video2``, and an Intel RealSense D405 scene camera (serial +``409122274720``). Every default is overridable from the CLI. + +Launched manually for now; wiring into ``cli/main.py`` (per-env client class + +driver script selection) is a separate step. +""" +from __future__ import annotations + +import argparse +import json +import os +import signal +import sys +import threading +import time +from pathlib import Path + +import numpy as np + +# Make ``rpent`` importable when this file is run as a script from +# the ``lerobot`` conda env (which need not have rpent installed). +_PHYSICALAGENT_ROOT = Path(__file__).resolve().parents[2] +if str(_PHYSICALAGENT_ROOT) not in sys.path: + sys.path.insert(0, str(_PHYSICALAGENT_ROOT)) + +from rpent.rpc_driver.socket import SocketRpcServer # noqa: E402 +from rpent.utils.logging import get_logger, init_output_dir # noqa: E402 + +from deployment.lerobot import calibration as scene_calib # noqa: E402 +from deployment.lerobot import geometry as geom # noqa: E402 +from deployment.lerobot.kinematics import SO101Kinematics # noqa: E402 +from deployment.lerobot.scene_camera import SceneCameraD405 # noqa: E402 + +logger = get_logger("lerobot_driver") + + +# --------------------------------------------------------------------------- +# SO101 joint / limit constants (mirrors rlinf SO101Env defaults) +# --------------------------------------------------------------------------- + +# Arm joints in bus-ID order; LeRobot keys obs/action by ``.pos``. +_ARM_JOINTS = ( + "shoulder_pan", + "shoulder_lift", + "elbow_flex", + "wrist_flex", + "wrist_roll", +) +_GRIPPER = "gripper" +_NUM_ARM_JOINTS = len(_ARM_JOINTS) + +# Arm joint limits in degrees (arm only, no gripper), aligned to the SO101 URDF +# (deg = rad * 180/pi). Matching the URDF matters for ``move_to``'s top-down +# mode: its IK solutions use wrist_flex up to ~95 deg and wrist_roll beyond +# +/-150 deg, so a tighter clip here would silently alter the solved pose (and +# re-tilt the gripper). placo already enforces these limits; this clip is a +# secondary safety net. +_JOINT_LIMIT_LOW_DEG = np.array([-110.0, -100.0, -96.8, -95.0, -157.2], dtype=np.float32) +_JOINT_LIMIT_HIGH_DEG = np.array([110.0, 100.0, 96.8, 95.0, 162.8], dtype=np.float32) +_GRIPPER_LIMIT_LOW = 0.0 +_GRIPPER_LIMIT_HIGH = 90.0 +# Home pose ``[q1..q5, gripper]``. The five arm joints use the all-zero +# calibrated pose (as requested). The gripper homes to a mid opening, NOT 0: +# driving the gripper to 0 stalls it closed against its mechanical stop and +# trips the motor's overload protection (observed: gripper motor id 6 dropped +# off the bus after a 0 command). Keep this strictly between the limits. +_RESET_QPOS = np.array([0.0, -100.0, 90.0, 65.0, 0.0, 50.0], dtype=np.float32) + +# Conservative tabletop workspace box in the base (world) frame, meters. +# ``move_to`` clips targets to this so an LLM-supplied coordinate can't drive +# the arm into the table / out of reach. Sized to the SO101's ~0.35 m reach and +# the observed pick region (the base origin sits above the table, so the plate/ +# table surface is around z ~ -0.06): x forward, y lateral, z up. move_to now +# closes the loop and actually reaches the commanded z, so the z floor must stop +# the fingertips just ABOVE the plate (grasps straddle the cube's upper body). +# Tune for your bench. +_WORKSPACE_MIN = np.array([0.08, -0.28, -0.055], dtype=np.float64) +_WORKSPACE_MAX = np.array([0.38, 0.28, 0.30], dtype=np.float64) + +# --- grasp control point (TCP) ------------------------------------------ +# IK/FK use the ``gripper_frame_link`` frame, which actually sits on the FIXED +# jaw's fingertip -- ~2.5 cm ABOVE the fingertip plane and offset toward the +# fixed side, NOT between the fingers. This vector (meters, in the +# gripper_frame_link LOCAL frame) shifts the controlled/reported point to the +# grasp point between the fingertips, so ``move_to`` targets and ``get_ee_pose`` +# refer to where an object is actually grasped. Derived from URDF + FK + a +# calibrated-depth fingertip measurement (~2 cm lateral toward the moving jaw, +# ~2.8 cm down to the fingertip plane). TUNE on hardware / via touch calibration +# if grasps land consistently off-centre; set to zeros for the raw frame. +_TCP_OFFSET_GRIPPER = np.array([-0.028, 0.043, 0.0282], dtype=np.float64) + +# --- motion speed / smoothness (safety) --------------------------------- +# The arm runs in position mode. Two knobs keep motions slow and gentle: +# * Servo acceleration: LeRobot's configure() maxes the Feetech "Acceleration" +# register at 254 (snappy). We override it with a gentler value (0-254; +# lower = softer ramps), applied to every motor after connect. +# * Software velocity cap: point-to-point motions (reset, move_to, +# move_joints_delta) are streamed as interpolated setpoints so no joint +# exceeds ``_MAX_JOINT_VEL_DEG_S`` (deg/s), instead of snapping to the target +# at full servo speed. Both are overridable from the CLI. +_MOTOR_ACCELERATION = 40 +_MAX_JOINT_VEL_DEG_S = 70.0 +_PACE_DT_S = 0.05 # setpoint-streaming period (20 Hz) +# Feetech position gain. LeRobot's configure() lowers P_Coefficient to 16 (from +# the factory default 32) "to avoid shakiness"; that soft gain lets gravity- +# loaded joints (esp. elbow_flex) under-reach their commanded angle, so the +# gripper lands short and low. We raise it so the arm actually holds commanded +# joints. Overridable from the CLI (raise for stiffer tracking; lower if the +# arm oscillates/buzzes). None = leave LeRobot's value untouched. +_POSITION_GAIN = 32 + +# move_to interpolates the full gripper POSE into fine steps (straight-line +# position + slerped orientation) and solves IK warm-started at each, so the tip +# tracks the Cartesian line and the wrist reorients smoothly -- avoiding wrist +# IK-branch flips that otherwise swing the long gripper through the table. +_CART_STEP_M = 0.01 # max Cartesian move per interpolated IK step +_REORIENT_STEP_DEG = 6.0 # max orientation change per interpolated IK step +# If an interpolated step's IK solution jumps more than this (deg) from the +# previous one, the path crossed an IK discontinuity (near-singular / infeasible +# top-down pose). move_to stops at the last safe pose rather than streaming the +# arm through the swing, which could drive the long gripper into the table. +_MAX_STEP_JOINT_JUMP_DEG = 20.0 +# Closed-loop correction: after the open-loop path, the real servos can settle +# short of the commanded pose (gravity sag under load), so move_to re-commands +# with Cartesian feed-forward until the achieved tip is within _CORRECTION_TOL_M +# or the budget runs out. Feed-forward is capped and workspace-clipped for +# safety (it can never drive below the workspace floor / into the plate). +_MAX_POSITION_CORRECTIONS = 3 +_CORRECTION_TOL_M = 0.008 +_MAX_CORRECTION_M = 0.06 + + +def _build_camera_configs(raw: dict[str, dict]) -> dict: + """Build LeRobot ``CameraConfig`` instances from user-facing dicts. + + Each value must contain a ``"type"`` key (``"opencv"`` or + ``"intelrealsense"``); the remaining keys are forwarded to the matching + ``CameraConfig`` subclass. Mirrors ``SO101Env._build_camera_configs``: + the config subclasses self-register on import, so the two camera modules + must be imported before ``get_choice_class`` can resolve the type name. + """ + import lerobot.cameras.opencv.configuration_opencv # noqa: F401 registers "opencv" + import lerobot.cameras.realsense.configuration_realsense # noqa: F401 registers "intelrealsense" + from lerobot.cameras.configs import CameraConfig + + result: dict = {} + for name, cfg in raw.items(): + cfg = dict(cfg) + cam_type = cfg.pop("type") + result[name] = CameraConfig.get_choice_class(cam_type)(**cfg) + return result + + +def _to_lerobot_action(arm_targets: np.ndarray, gripper_target: float) -> dict: + """Build the ``{motor}.pos`` dict LeRobot's ``send_action`` expects. + + LeRobot filters incoming keys via ``key.endswith(".pos")``; a missing + suffix silently drops that motor. + """ + action = {f"{name}.pos": float(arm_targets[i]) for i, name in enumerate(_ARM_JOINTS)} + action[f"{_GRIPPER}.pos"] = float(gripper_target) + return action + + +def _topdown_rotation(yaw: float) -> np.ndarray: + """Target gripper rotation for a top-down approach (in the base frame). + + Maps the gripper's local approach axis (local +z, the wrist->fingertips + direction) to world -z (straight down). ``yaw`` (radians) rotates the + gripper about the vertical: the gripper's local +x maps to the horizontal + direction ``(cos yaw, sin yaw, 0)``, which sets the jaw-line heading. + """ + c, s = np.cos(yaw), np.sin(yaw) + return np.array( + [[c, s, 0.0], [s, -c, 0.0], [0.0, 0.0, -1.0]], dtype=np.float64 + ) + + +def _grasp_point(T_base_gripper: np.ndarray) -> np.ndarray: + """World xyz of the grasp point (between the fingertips) for a + ``gripper_frame_link`` pose, applying ``_TCP_OFFSET_GRIPPER`` in the gripper + frame. ``gripper_frame_link`` itself is on the fixed jaw, so this is what + ``move_to`` should target and ``get_ee_pose`` should report for grasping. + """ + T = np.asarray(T_base_gripper, dtype=np.float64) + return T[:3, 3] + T[:3, :3] @ _TCP_OFFSET_GRIPPER + + +def _approach_tilt_deg(R: np.ndarray) -> float: + """Angle (deg) between the gripper approach axis (local +z) and world -z. + + 0 deg = pointing straight down. Used to verify a top-down ``move_to``. + """ + z_world = np.asarray(R, dtype=np.float64) @ np.array([0.0, 0.0, 1.0]) + return float(np.degrees(np.arccos(np.clip(-z_world[2], -1.0, 1.0)))) + + +def _rotation_angle_deg(R0: np.ndarray, R1: np.ndarray) -> float: + """Geodesic angle (deg) between two rotation matrices.""" + cos = ( + np.trace(np.asarray(R0, dtype=np.float64).T @ np.asarray(R1, dtype=np.float64)) + - 1.0 + ) / 2.0 + return float(np.degrees(np.arccos(np.clip(cos, -1.0, 1.0)))) + + +def _slerp_rotation(R0: np.ndarray, R1: np.ndarray, t: float) -> np.ndarray: + """Interpolate rotation ``R0`` -> ``R1`` by fraction ``t`` in [0, 1]. + + Rotates about the fixed axis of the relative rotation (Rodrigues) -- a + matrix slerp -- so the gripper reorients along one smooth shortest arc. + """ + R0 = np.asarray(R0, dtype=np.float64) + R1 = np.asarray(R1, dtype=np.float64) + R_rel = R0.T @ R1 + ang = np.arccos(np.clip((np.trace(R_rel) - 1.0) / 2.0, -1.0, 1.0)) + if ang < 1e-8: + return R0.copy() + axis = np.array( + [ + R_rel[2, 1] - R_rel[1, 2], + R_rel[0, 2] - R_rel[2, 0], + R_rel[1, 0] - R_rel[0, 1], + ], + dtype=np.float64, + ) / (2.0 * np.sin(ang)) + th = ang * float(t) + K = np.array( + [ + [0.0, -axis[2], axis[1]], + [axis[2], 0.0, -axis[0]], + [-axis[1], axis[0], 0.0], + ], + dtype=np.float64, + ) + return R0 @ (np.eye(3) + np.sin(th) * K + (1.0 - np.cos(th)) * (K @ K)) + + +class SO101LeRobotEnv: + """Minimal ``reset`` / ``step`` driver for a physical SO101 arm. + + Action: ``(6,)`` float array ``[q1..q5, gripper]`` of absolute joint + position targets in degrees. Arm targets are clipped to the configured + joint limits and the gripper to ``[0, 90]`` before being sent to the + motor bus. + + Observation:: + + {"state": {"joint_position": (5,) float32, # arm joints, degrees + "gripper_position": (1,) float32, # gripper opening + "ee_pose_base": (3,) float32, # gripper xyz in base (FK) + "ee_quat_base": (4,) float32}, # gripper quat wxyz (FK) + "frames": {"arm": (H, W, 3) uint8, "scene": (H, W, 3) uint8}, + "depth": {"scene": (H, W) float32}} # metric, aligned to color + + ``ee_pose_base`` / ``ee_quat_base`` are present only when FK is available; + ``scene`` frames/depth only when the scene camera is configured. All values + are plain numpy / floats so they pickle across the RPC wire (the agent + process does not import torch). + """ + + def __init__( + self, + *, + port: str, + calibration_id: str, + arm_camera_cfgs: dict[str, dict], + scene_serial: str | None = None, + scene_size: tuple[int, int] = (720, 1280), + scene_fps: int = 30, + urdf_path: str | None = None, + max_relative_target: float | None = None, + max_episode_steps: int = 200, + step_frequency: float = 30.0, + motor_acceleration: int | None = _MOTOR_ACCELERATION, + max_joint_vel_deg_s: float = _MAX_JOINT_VEL_DEG_S, + position_gain: int | None = _POSITION_GAIN, + auto_calibrate: bool = False, + ) -> None: + self._max_episode_steps = max_episode_steps + self._step_frequency = step_frequency + self._max_joint_vel_deg_s = float(max_joint_vel_deg_s) + self._pace_dt = _PACE_DT_S + self._num_steps = 0 + self._action_low = np.append(_JOINT_LIMIT_LOW_DEG, _GRIPPER_LIMIT_LOW).astype(np.float32) + self._action_high = np.append(_JOINT_LIMIT_HIGH_DEG, _GRIPPER_LIMIT_HIGH).astype(np.float32) + + from lerobot.robots.so_follower import SO101Follower + from lerobot.robots.so_follower.config_so_follower import SO101FollowerConfig + + robot_cfg = SO101FollowerConfig( + port=port, + id=calibration_id, + use_degrees=True, + max_relative_target=max_relative_target, + # Keep torque on at disconnect so the arm holds its parked pose. + disable_torque_on_disconnect=False, + cameras=_build_camera_configs(arm_camera_cfgs), + ) + self._robot = SO101Follower(robot_cfg) + # ``calibrate=False`` loads the on-disk calibration without ever + # prompting on stdin (a server must never block on input). + self._robot.connect(calibrate=auto_calibrate) + self._arm_camera_names = tuple(arm_camera_cfgs.keys()) + + # Soften motion: LeRobot's configure() sets the Feetech "Acceleration" + # register to its max (254). Override with a gentler value so the arm + # ramps smoothly rather than snapping (safety). Done with torque briefly + # off, mirroring LeRobot's own register writes. + self._motor_acceleration = motor_acceleration + if motor_acceleration is not None: + try: + with self._robot.bus.torque_disabled(): + for motor in self._robot.bus.motors: + # Keep the gripper snappy so grasps close promptly; only + # slow the (heavier, safety-relevant) arm joints. + if motor == _GRIPPER: + continue + self._robot.bus.write( + "Acceleration", motor, int(motor_acceleration) + ) + except Exception as e: + logger.warning( + "could not set motor Acceleration=%s (motions stay fast): %s", + motor_acceleration, e, + ) + + # Stiffen position holding: LeRobot lowers P_Coefficient to 16, which + # lets gravity-loaded arm joints under-reach their target. Raise it on + # the arm joints (leave the gripper as LeRobot set it). + self._position_gain = position_gain + if position_gain is not None: + try: + with self._robot.bus.torque_disabled(): + for motor in self._robot.bus.motors: + if motor == _GRIPPER: + continue + self._robot.bus.write( + "P_Coefficient", motor, int(position_gain) + ) + except Exception as e: + logger.warning( + "could not set motor P_Coefficient=%s (tracking may be soft): %s", + position_gain, e, + ) + + # Scene camera (fixed, depth) is managed directly via pyrealsense2 so we + # get depth aligned to color + intrinsics (LeRobot's wrapper gives + # neither). The arm camera stays under LeRobot above (color only). + self._scene_serial = scene_serial or None + self._scene_cam: SceneCameraD405 | None = None + if self._scene_serial: + scene_h, scene_w = scene_size + self._scene_cam = SceneCameraD405( + self._scene_serial, width=scene_w, height=scene_h, fps=scene_fps, + ) + + # Forward kinematics for end-effector pose in the base (world) frame. + self._kin: SO101Kinematics | None = None + try: + self._kin = SO101Kinematics(urdf_path=urdf_path) + except Exception as e: + logger.warning("FK unavailable (%s); ee_pose will be omitted", e) + + # Scene-cam -> base extrinsic, if it has been calibrated (touch/Kabsch). + # A fit whose saved RMSE is too large is REJECTED (treated as + # uncalibrated): a bad extrinsic makes back_project return badly wrong + # world coordinates, so the arm would chase unreachable targets. The + # operator must recalibrate rather than have the agent act on garbage. + self._T_base_cam = None + self._calib_rmse_m: float | None = None + if self._scene_serial: + record = scene_calib.load_extrinsic_record(self._scene_serial) + if record is not None: + self._calib_rmse_m = record.get("rmse_m") + rmse = self._calib_rmse_m + if rmse is not None and rmse > scene_calib.MAX_ACCEPTABLE_RMSE_M: + logger.warning( + "scene-cam extrinsic REJECTED: rmse=%.3fm > %.3fm limit; " + "back_project would be unreliable. Recalibrate with " + "deployment/lerobot/auto_calibrate_scene_cam.py " + "(or calibrate_scene_cam.py).", + rmse, scene_calib.MAX_ACCEPTABLE_RMSE_M, + ) + else: + self._T_base_cam = record["T_base_cam"] + + if self._T_base_cam is not None: + extr_status = ( + f"loaded (rmse={self._calib_rmse_m:.3f}m)" + if self._calib_rmse_m is not None else "loaded" + ) + elif self._calib_rmse_m is not None: + extr_status = f"rejected (rmse={self._calib_rmse_m:.3f}m)" + else: + extr_status = "uncalibrated" + + cam_list = list(self._arm_camera_names) + (["scene"] if self._scene_cam else []) + logger.info( + "SO101 connected on %s (calibration_id=%s); cameras=[%s]; " + "FK=%s; scene_extrinsic=%s", + port, calibration_id, ", ".join(cam_list) or "none", + "on" if self._kin else "off", + extr_status, + ) + logger.info( + "motion pacing: max_joint_vel=%.0f deg/s; motor_acceleration=%s", + self._max_joint_vel_deg_s, + self._motor_acceleration if self._motor_acceleration is not None + else "default (254)", + ) + logger.info( + "position gain: P_Coefficient=%s", + self._position_gain if self._position_gain is not None + else "default (LeRobot 16)", + ) + + # ------------------------------------------------------------------ + # gym-like surface + # ------------------------------------------------------------------ + + def _stream_joint_path(self, q_from, q_to, gripper_value) -> None: + """Glide the arm from ``q_from`` to ``q_to`` (arm-joint vectors, degrees). + + Streams interpolated joint setpoints so no joint moves faster than + ``self._max_joint_vel_deg_s`` (deg/s), instead of commanding the target + in one shot and letting the servos snap there at full speed. The gripper + is held at ``gripper_value`` throughout. + """ + q_from = np.asarray(q_from, dtype=np.float64).reshape(-1)[:_NUM_ARM_JOINTS] + q_to = np.asarray(q_to, dtype=np.float64).reshape(-1)[:_NUM_ARM_JOINTS] + max_delta = float(np.max(np.abs(q_to - q_from))) if q_to.size else 0.0 + per_step = max(1e-6, self._max_joint_vel_deg_s * self._pace_dt) + n = max(1, int(np.ceil(max_delta / per_step))) + for i in range(1, n + 1): + q = q_from + (q_to - q_from) * (i / n) + self._robot.send_action(_to_lerobot_action(q, gripper_value)) + time.sleep(self._pace_dt) + + def reset(self) -> tuple[dict, dict]: + """Send the arm to its rest pose (paced), reset the counter, return obs.""" + self._num_steps = 0 + obs = self._robot.get_observation() + cur = np.array( + [obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float64 + ) + self._stream_joint_path( + cur, _RESET_QPOS[:_NUM_ARM_JOINTS], float(_RESET_QPOS[_NUM_ARM_JOINTS]) + ) + time.sleep(0.3) + return self._get_observation(), {} + + def step(self, action) -> tuple[dict, float, bool, bool, dict]: + """Command one absolute joint target and return the gym 5-tuple. + + Returns ``(obs, reward, terminated, truncated, info)``. This minimal + driver computes no task reward (``reward == 0.0``) and never + auto-terminates; ``truncated`` flips once ``max_episode_steps`` is + reached. Task success / stopping is decided by the agent. + """ + t0 = time.time() + action = np.asarray(action, dtype=np.float32).reshape(-1) + expected = _NUM_ARM_JOINTS + 1 + if action.shape[0] != expected: + raise ValueError( + f"action must have {expected} entries [q1..q5, gripper]; " + f"got {action.shape[0]}" + ) + action = np.clip(action, self._action_low, self._action_high) + self._robot.send_action( + _to_lerobot_action(action[:_NUM_ARM_JOINTS], float(action[_NUM_ARM_JOINTS])) + ) + self._num_steps += 1 + + obs = self._get_observation() + truncated = self._num_steps >= self._max_episode_steps + + # Pace the control loop to the requested frequency. + dt = time.time() - t0 + time.sleep(max(0.0, (1.0 / self._step_frequency) - dt)) + return obs, 0.0, False, truncated, {} + + def get_spec(self) -> dict: + """Static self-description for the agent side (action bounds, cams).""" + cam_names = list(self._arm_camera_names) + (["scene"] if self._scene_cam else []) + return { + "action_dim": _NUM_ARM_JOINTS + 1, + "arm_joints": list(_ARM_JOINTS), + "action_low": self._action_low.tolist(), + "action_high": self._action_high.tolist(), + "camera_names": cam_names, + "scene_camera": "scene" if self._scene_cam else None, + "has_ee_pose": self._kin is not None, + "world_frame": "base_link", + "max_episode_steps": self._max_episode_steps, + } + + # ------------------------------------------------------------------ + # localization surface (base frame == world) + # ------------------------------------------------------------------ + + def get_ee_pose(self) -> dict: + """Live FK: gripper pose in the base (world) frame.""" + obs = self._robot.get_observation() + joints = np.array( + [obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float32 + ) + ee = self._compute_ee_pose(joints) + if ee is None: + return {"error": "FK unavailable (URDF/placo missing)"} + return { + "xyz": _grasp_point(ee["T"]).tolist(), + "quat_wxyz": ee["quat"].tolist(), + "joints_deg": joints.tolist(), + "T_base_gripper": ee["T"].tolist(), + "frame": "base_link", + } + + def get_scene_camera_meta(self) -> dict: + """Scene-camera intrinsics + depth scale + base extrinsic (if any).""" + if self._scene_cam is None: + return {"error": "scene camera not configured"} + meta = self._scene_cam.meta() + meta["frame"] = "scene_cam" + meta["calibrated"] = self._T_base_cam is not None + meta["rmse_m"] = self._calib_rmse_m + meta["T_base_cam"] = ( + self._T_base_cam.tolist() if self._T_base_cam is not None else None + ) + return meta + + def get_scene_frame(self) -> dict: + """Live scene color + metric depth (for calibration / ad-hoc queries).""" + if self._scene_cam is None: + return {"error": "scene camera not configured"} + rgb, depth = self._scene_cam.read() + return {"color": rgb, "depth": depth, "K": self._scene_cam.K.tolist()} + + def get_obs(self) -> dict: + """Return the current observation without moving the arm. + + Used to refresh the agent-side cache after a primitive (e.g. move_to) + that changes the world but does not itself return an observation. + """ + return self._get_observation() + + def set_torque(self, enabled: bool) -> dict: + """Enable/disable arm motor torque. + + Disabling lets the operator free-drive the arm by hand (used by the + touch/Kabsch scene-camera calibration). Joint encoders remain readable + with torque off, so FK still works. Re-enable to hold position. + """ + bus = self._robot.bus + try: + if enabled: + bus.enable_torque() + else: + bus.disable_torque() + except Exception as e: + return {"ok": False, "error": str(e)} + logger.info("arm torque %s", "enabled" if enabled else "disabled") + return {"ok": True, "torque_enabled": bool(enabled)} + + def _read_arm_joints(self) -> np.ndarray: + """Current arm joint angles (deg) from a fresh observation.""" + obs = self._robot.get_observation() + return np.array( + [obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float64 + ) + + def _servo_pose_path( + self, from_joints, R_start, p_start, R_end, p_end, grip_val, orient_w, + settle_s, + ) -> bool: + """Stream the gripper along an interpolated pose (straight-line position + + slerped orientation), warm-starting IK per fine step so solutions stay + on ONE branch and the tip tracks the line. Returns ``halted`` (an IK + discontinuity forced an early stop). Settles ``settle_s`` before + returning so the caller can measure the achieved pose. + """ + from_joints = np.asarray(from_joints, dtype=np.float64) + p_start = np.asarray(p_start, dtype=np.float64) + p_end = np.asarray(p_end, dtype=np.float64) + dist = float(np.linalg.norm(p_end - p_start)) + reorient_deg = _rotation_angle_deg(R_start, R_end) if orient_w > 0 else 0.0 + n_steps = max( + 1, + int(np.ceil(dist / _CART_STEP_M)), + int(np.ceil(reorient_deg / _REORIENT_STEP_DEG)), + ) + seed = from_joints.copy() + prev_q = from_joints.copy() + halted = False + for i in range(1, n_steps + 1): + frac = i / n_steps + T_des = np.eye(4) + T_des[:3, :3] = ( + _slerp_rotation(R_start, R_end, frac) if orient_w > 0 else R_end + ) + T_des[:3, 3] = p_start + (p_end - p_start) * frac + q = self._kin.ik( + seed, T_des, position_weight=1.0, orientation_weight=orient_w + ) + q_arm = np.clip( + q[:_NUM_ARM_JOINTS], _JOINT_LIMIT_LOW_DEG.astype(np.float64), + _JOINT_LIMIT_HIGH_DEG.astype(np.float64), + ) + # A large jump between consecutive fine-step solutions = the path + # crossed an IK discontinuity (near-singular / infeasible top-down + # pose); stop at the last safe pose rather than swinging through it. + if float(np.max(np.abs(q_arm - prev_q))) > _MAX_STEP_JOINT_JUMP_DEG: + halted = True + break + seed = q_arm + self._stream_joint_path(prev_q, q_arm, grip_val) + prev_q = q_arm + time.sleep(settle_s) + return halted + + def move_to( + self, + xyz, + *, + gripper: float | None = None, + approach: str = "free", + yaw_deg: float | None = None, + settle_s: float = 0.4, + pos_tol_m: float = 0.02, + tilt_tol_deg: float = 15.0, + max_corrections: int = _MAX_POSITION_CORRECTIONS, + ) -> dict: + """Move the gripper to a world-frame (base_link) XYZ via IK. + + ``approach`` selects the wrist-orientation policy: + + * ``"free"`` (default): position-only IK; the wrist settles at whatever + orientation placo converges to. Maximal reach, but the fingertips' + location relative to the returned EE point is unpredictable (the TCP + is ~0.1 m out along the gripper axis), so it is unreliable for + grasping. + * ``"down"``: top-down IK -- the gripper approach axis is driven to + vertical (pointing straight down) so the fingertips descend along + world -z, which makes grasping predictable. ``yaw_deg`` sets the + jaw-line heading about the vertical (0 = +x/forward); if ``None`` a + reachable yaw is searched automatically. + + The target is clipped to the workspace box and approached by + interpolating the full gripper pose (straight-line position + smoothly + slerped orientation) into fine, warm-started IK steps. It then CLOSES + THE LOOP: it measures the achieved tip and re-commands with feed-forward + (up to ``max_corrections`` times) to cancel the servos' steady-state sag + under load, so the tip lands on the commanded xyz -- callers should pass + the true target, NOT a hand-tuned over-shoot. Holds the current gripper + opening unless ``gripper`` is given. + + Returns a log dict with ``reached`` (position within ``pos_tol_m`` and, + for ``"down"``, approach tilt within ``tilt_tol_deg``), the + commanded/achieved xyz, the position error, and the achieved approach + tilt in degrees. + """ + if self._kin is None: + return {"error": "IK unavailable (URDF/placo missing)"} + approach = str(approach).lower() + if approach not in ("free", "down"): + return {"error": f"approach must be 'free' or 'down'; got {approach!r}"} + target = np.asarray(xyz, dtype=np.float64).reshape(-1) + if target.shape[0] != 3: + return {"error": f"xyz must be 3 numbers; got {target.shape[0]}"} + + clipped = np.clip(target, _WORKSPACE_MIN, _WORKSPACE_MAX) + was_clipped = not np.allclose(clipped, target, atol=1e-6) + + obs = self._robot.get_observation() + cur_joints = np.array( + [obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float64 + ) + cur_gripper = float(obs.get(f"{_GRIPPER}.pos", 0.0)) + grip_val = cur_gripper if gripper is None else float(gripper) + + T_cur = self._kin.fk(cur_joints) + p0 = T_cur[:3, 3] + R0 = T_cur[:3, :3] + + # Target orientation for the move: reorient to vertical for "down"; + # leave it to the IK (weight 0) for "free". + if approach == "down": + orient_w = 1.0 + yaw = ( + self._search_topdown_yaw(cur_joints, clipped) + if yaw_deg is None else float(np.radians(yaw_deg)) + ) + R_target = _topdown_rotation(yaw) + else: + orient_w = 0.0 + yaw = None + R_target = R0 # ignored (orientation_weight=0) + + # Held orientation: top-down for "down", the current (ignored) rotation + # for "free". Move the tip along the interpolated pose, then close the + # loop to cancel the servos' steady-state sag under load. + R_hold = R_target if approach == "down" else R0 + # move_to controls the GRASP POINT (between the fingertips). Convert that + # target into the gripper_frame_link goal the IK/servo actually drives. + # The offset is only well-defined for the fixed "down" orientation; + # "free" leaves the tip frame uncorrected (its final orientation, hence + # the offset direction, is unknown until IK converges). + tcp_off_world = ( + R_hold @ _TCP_OFFSET_GRIPPER if approach == "down" + else np.zeros(3, dtype=np.float64) + ) + halted = self._servo_pose_path( + cur_joints, R0, p0, R_hold, clipped - tcp_off_world, + grip_val, orient_w, settle_s, + ) + final_joints = self._read_arm_joints() + T_final = self._kin.fk(final_joints) + grasp_final = _grasp_point(T_final) if approach == "down" else T_final[:3, 3] + + # Closed-loop correction: the real arm can settle short of the commanded + # pose (gravity sag), even though the IK path is exact. Feed the residual + # forward and re-command (workspace-clipped + capped) until the achieved + # GRASP POINT is within tolerance or the correction budget runs out. + n_corr = 0 + ff = np.zeros(3, dtype=np.float64) + while ( + not halted + and n_corr < max_corrections + and float(np.linalg.norm(clipped - grasp_final)) > _CORRECTION_TOL_M + ): + ff = np.clip( + ff + (clipped - grasp_final), -_MAX_CORRECTION_M, _MAX_CORRECTION_M + ) + corr_goal = ( + np.clip(clipped + ff, _WORKSPACE_MIN, _WORKSPACE_MAX) - tcp_off_world + ) + halted = self._servo_pose_path( + final_joints, R_hold, T_final[:3, 3], R_hold, corr_goal, + grip_val, orient_w, settle_s, + ) + final_joints = self._read_arm_joints() + T_final = self._kin.fk(final_joints) + grasp_final = _grasp_point(T_final) if approach == "down" else T_final[:3, 3] + n_corr += 1 + + err = float(np.linalg.norm(grasp_final - clipped)) + tilt = _approach_tilt_deg(T_final[:3, :3]) + reached = ( + not halted + and err <= pos_tol_m + and (approach != "down" or tilt <= tilt_tol_deg) + ) + result = { + "reached": bool(reached), + "approach": approach, + "target_xyz": [round(float(v), 4) for v in target], + "commanded_xyz": [round(float(v), 4) for v in clipped], + "final_xyz": [round(float(v), 4) for v in grasp_final], + "pos_error_m": round(err, 4), + "approach_tilt_deg": round(tilt, 1), + "yaw_deg": (None if yaw is None else round(float(np.degrees(yaw)), 1)), + "clipped_to_workspace": was_clipped, + "pos_corrections": n_corr, + "halted_early": bool(halted), + "joints_deg": [round(float(v), 2) for v in final_joints], + "gripper": round(grip_val, 2), + } + if halted: + result["note"] = ( + "stopped partway: the straight-line top-down path crossed an " + "unreachable / near-singular pose. Try a nearer target, a " + "different yaw_deg, or move in smaller hops." + ) + elif not reached: + result["note"] = ( + f"settled {round(err * 1000)} mm short after {n_corr} " + "correction(s); the target may be past the arm's reach here. " + "Try a nearer / higher target." + ) + return result + + def _search_topdown_yaw(self, seed_joints, target, *, n_yaw: int = 12) -> float: + """Pick a top-down wrist yaw (rad) that reaches ``target`` best. + + Whether a vertical approach is reachable depends on the jaw-line yaw + (wrist_roll/shoulder_pan coupling), so we sweep candidate yaws, run IK + for each, and keep the one with the smallest FK position error (ties + broken by approach tilt). Pure CPU -- the arm does not move. + """ + best_yaw, best_key = 0.0, None + seed = np.asarray(seed_joints, dtype=np.float64) + for k in range(n_yaw): + yaw = 2.0 * np.pi * k / n_yaw + T_des = np.eye(4) + T_des[:3, :3] = _topdown_rotation(yaw) + T_des[:3, 3] = target + q = self._kin.ik(seed, T_des, position_weight=1.0, orientation_weight=1.0) + T_q = self._kin.fk(q[:_NUM_ARM_JOINTS]) + perr = float(np.linalg.norm(T_q[:3, 3] - target)) + tilt = _approach_tilt_deg(T_q[:3, :3]) + key = (round(perr, 4), round(tilt, 1)) + if best_key is None or key < best_key: + best_key, best_yaw = key, yaw + return best_yaw + + def move_joints_delta( + self, + delta_deg, + *, + gripper_delta: float | None = None, + max_step_deg: float = 15.0, + settle_s: float = 0.4, + ) -> dict: + """Nudge each arm joint by a relative amount (degrees). + + ``delta_deg`` is 5 values ``[d_pan, d_lift, d_elbow, d_wrist_flex, + d_wrist_roll]`` added to the current arm joints. Each entry is capped to + +/- ``max_step_deg`` and the result is clamped to the joint limits, so a + single call makes a small, safe adjustment. Optionally nudge the gripper + by ``gripper_delta`` (clamped to its limits). Use this for fine + alignment ``move_to`` cannot express -- e.g. tweak wrist_roll to line the + jaws up with an object, or descend a few millimetres -- reading the new + EE pose back from the result. + + Returns a log dict with the applied delta, achieved joints, gripper, and + (when FK is available) the new EE xyz and approach tilt. + """ + delta = np.asarray(delta_deg, dtype=np.float64).reshape(-1) + if delta.shape[0] != _NUM_ARM_JOINTS: + return { + "error": f"delta_deg must have {_NUM_ARM_JOINTS} entries " + f"[pan, lift, elbow, wrist_flex, wrist_roll]; got {delta.shape[0]}" + } + cap = abs(float(max_step_deg)) + delta = np.clip(delta, -cap, cap) + + obs = self._robot.get_observation() + cur_joints = np.array( + [obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float64 + ) + cur_gripper = float(obs.get(f"{_GRIPPER}.pos", 0.0)) + + target_joints = np.clip( + cur_joints + delta, + _JOINT_LIMIT_LOW_DEG.astype(np.float64), + _JOINT_LIMIT_HIGH_DEG.astype(np.float64), + ) + if gripper_delta is None: + grip_val = cur_gripper + else: + grip_val = float( + np.clip(cur_gripper + float(gripper_delta), + _GRIPPER_LIMIT_LOW, _GRIPPER_LIMIT_HIGH) + ) + + self._stream_joint_path(cur_joints, target_joints, grip_val) + time.sleep(settle_s) + + final_obs = self._robot.get_observation() + final_joints = np.array( + [final_obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float64 + ) + result: dict = { + "applied_delta_deg": [round(float(v), 2) for v in delta], + "joints_deg": [round(float(v), 2) for v in final_joints], + "gripper": round(grip_val, 2), + } + ee = self._compute_ee_pose(final_joints) + if ee is not None: + result["ee_xyz"] = [round(float(v), 4) for v in _grasp_point(ee["T"])] + result["approach_tilt_deg"] = round(_approach_tilt_deg(ee["T"][:3, :3]), 1) + return result + + # ------------------------------------------------------------------ + # automatic scene-camera calibration (markerless, gripper-motion) + # ------------------------------------------------------------------ + + def _set_gripper_hold(self, gripper_value: float) -> None: + """Set the gripper opening while freezing the arm at its current joints. + + Used during calibration so that between the two capture frames ONLY the + gripper fingers move (clean motion segmentation). + """ + obs = self._robot.get_observation() + q = np.array([obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float64) + self._robot.send_action(_to_lerobot_action(q, float(gripper_value))) + + @staticmethod + def _calibration_targets() -> list[list[float]]: + """A spread, non-coplanar grid of tip targets inside the workspace.""" + xs = [0.15, 0.22, 0.28] + ys = [-0.12, 0.0, 0.12] + zs = [0.12, 0.19] + return [[x, y, z] for z in zs for y in ys for x in xs] + + def auto_calibrate_scene_camera( + self, + *, + n_points: int = 10, + gripper_open: float = 90.0, + gripper_closed: float = 20.0, + settle_s: float = 0.8, + ransac_thresh_m: float = 0.015, + save: bool = True, + ) -> dict: + """Markerless automatic scene-cam -> base calibration. + + Drives the tip to a grid of base-frame positions (move_to needs no + extrinsic), and at each pose toggles the gripper with the arm frozen and + segments the motion in the scene image to locate the tip (centroid + + median depth -> camera point). The achieved FK gives the base point. + A RANSAC Kabsch fit then yields ``T_base_cam``, which is saved and + hot-loaded so back_project returns world coords immediately. + + Returns a summary dict (``n_used``, ``rmse_m``, per-pose diagnostics). + """ + if self._scene_cam is None: + return {"error": "scene camera not configured"} + if self._kin is None: + return {"error": "IK unavailable (URDF/placo missing)"} + + targets = self._calibration_targets() + cam_pts: list[list[float]] = [] + base_pts: list[list[float]] = [] + poses: list[dict] = [] + + for tgt in targets: + if len(cam_pts) >= n_points: + break + mv = self.move_to(tgt, gripper=gripper_open, settle_s=settle_s) + if "error" in mv or not mv.get("reached"): + poses.append({"target": tgt, "skipped": "unreachable"}) + continue + time.sleep(settle_s) + + base = np.asarray(mv["final_xyz"], dtype=np.float64) + self._scene_cam.read() # flush a frame + rgb_open, _ = self._scene_cam.read() + self._set_gripper_hold(gripper_closed) + time.sleep(settle_s) + rgb_closed, depth = self._scene_cam.read() + self._set_gripper_hold(gripper_open) # reopen for the next pose + + det = geom.detect_tip_pixel_by_motion( + rgb_open, rgb_closed, depth, self._scene_cam.K, + ) + if det is None: + poses.append({"target": tgt, "skipped": "no_tip_detected"}) + continue + cam_pts.append(det["xyz_cam"]) + base_pts.append(base.tolist()) + poses.append({"target": tgt, "base_xyz": base.round(4).tolist(), + "pixel": [round(v, 1) for v in det["pixel"]], + "depth_m": round(det["depth_m"], 4), "area": det["area"]}) + + if len(cam_pts) < 4: + return {"error": f"only {len(cam_pts)} usable points (need >= 4)", + "poses": poses} + + T, rmse, inliers = geom.ransac_kabsch( + cam_pts, base_pts, thresh_m=ransac_thresh_m, + ) + accepted = bool(rmse <= scene_calib.MAX_ACCEPTABLE_RMSE_M) + result = { + "n_targets": len(targets), + "n_used": len(cam_pts), + "n_inliers": int(np.asarray(inliers).sum()), + "rmse_m": round(float(rmse), 4), + "T_base_cam": T.tolist(), + "accepted": accepted, + "saved": False, + "poses": poses, + } + if not accepted: + # A high RMSE means the cam/base correspondences are inconsistent + # (poor tip detection, lighting, or occlusion). Saving it would + # silently corrupt every back_project, so refuse and ask for a rerun. + result["error"] = ( + f"calibration RMSE {rmse * 1000:.1f} mm exceeds the " + f"{scene_calib.MAX_ACCEPTABLE_RMSE_M * 1000:.0f} mm limit; not " + "saved. Clear the workspace, improve gripper visibility/lighting, " + "and rerun." + ) + logger.warning( + "scene-cam calibration REJECTED: rmse=%.4fm (> %.3fm); not saved", + rmse, scene_calib.MAX_ACCEPTABLE_RMSE_M, + ) + elif save: + path = scene_calib.save_extrinsic( + self._scene_serial, T, K=self._scene_cam.K, + rmse_m=rmse, num_points=int(np.asarray(inliers).sum()), + ) + self._T_base_cam = T # hot-load so back_project works immediately + self._calib_rmse_m = round(float(rmse), 4) + result["saved"] = True + result["path"] = str(path) + logger.info("scene-cam calibrated: rmse=%.4fm, saved %s", rmse, path) + + # Park the arm at rest after the sweep (paced, gentle). + try: + obs = self._robot.get_observation() + cur = np.array( + [obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float64 + ) + self._stream_joint_path( + cur, _RESET_QPOS[:_NUM_ARM_JOINTS], float(_RESET_QPOS[_NUM_ARM_JOINTS]) + ) + except Exception: + pass + return result + + def close(self) -> None: + """Park the arm at rest (paced + torque held) and disconnect cleanly.""" + try: + obs = self._robot.get_observation() + cur = np.array( + [obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float64 + ) + self._stream_joint_path( + cur, _RESET_QPOS[:_NUM_ARM_JOINTS], float(_RESET_QPOS[_NUM_ARM_JOINTS]) + ) + time.sleep(0.3) + except Exception as e: + logger.warning("failed to park arm on close: %s", e) + if self._scene_cam is not None: + try: + self._scene_cam.close() + except Exception as e: + logger.warning("error closing scene camera: %s", e) + try: + self._robot.disconnect() + logger.info("SO101 disconnected") + except Exception as e: + logger.warning("error disconnecting robot: %s", e) + + # ------------------------------------------------------------------ + # internals + # ------------------------------------------------------------------ + + def _get_observation(self) -> dict: + obs = self._robot.get_observation() + joint_position = np.array( + [obs.get(f"{n}.pos", 0.0) for n in _ARM_JOINTS], dtype=np.float32 + ) + gripper_position = np.array([obs.get(f"{_GRIPPER}.pos", 0.0)], dtype=np.float32) + + frames: dict = {} + for cam in self._arm_camera_names: + frame = obs.get(cam) + if frame is not None: + frames[cam] = np.ascontiguousarray(np.asarray(frame, dtype=np.uint8)) + + depth: dict = {} + if self._scene_cam is not None: + scene_rgb, scene_depth = self._scene_cam.read() + frames["scene"] = scene_rgb + depth["scene"] = scene_depth + + state: dict = { + "joint_position": joint_position, + "gripper_position": gripper_position, + } + ee = self._compute_ee_pose(joint_position) + if ee is not None: + state["ee_pose_base"] = _grasp_point(ee["T"]).astype(np.float32) + state["ee_quat_base"] = ee["quat"] + + out: dict = {"state": state, "frames": frames} + if depth: + out["depth"] = depth + return out + + def _compute_ee_pose(self, joints_deg) -> dict | None: + """FK -> gripper pose in the base (world) frame, or None if FK is off.""" + if self._kin is None: + return None + try: + T = self._kin.fk(joints_deg) + except Exception as e: + logger.warning("FK failed: %s", e) + return None + xyz = T[:3, 3].astype(np.float32) + quat = geom.rotation_to_quat(T[:3, :3]).astype(np.float32) + return {"xyz": xyz, "quat": quat, "T": T} + + +# --------------------------------------------------------------------------- +# RPC plumbing +# --------------------------------------------------------------------------- + +_INITIAL_PPID = os.getppid() + + +def _start_parent_watchdog( + server: SocketRpcServer, + shutdown_event: threading.Event, + poll_s: float = 2.0, +) -> None: + """Shut the RPC server down if the agent (parent) process dies.""" + + def _watch() -> None: + while not shutdown_event.is_set(): + time.sleep(poll_s) + ppid = os.getppid() + if ppid != _INITIAL_PPID or ppid == 1: + logger.warning( + "parent died (ppid %s -> %s); stopping RPC server", + _INITIAL_PPID, + ppid, + ) + shutdown_event.set() + threading.Thread(target=server.shutdown, daemon=True).start() + return + + threading.Thread(target=_watch, daemon=True).start() + + +def _build_dispatcher(env: SO101LeRobotEnv, shutdown_event: threading.Event): + """Route ``env.*`` / ``shutdown`` to the right callable.""" + + def dispatch(method: str, args: tuple, kwargs: dict): + if method.startswith("env."): + return getattr(env, method[len("env."):])(*args, **kwargs) + if method == "shutdown": + shutdown_event.set() + return {"ok": True} + raise ValueError(f"unknown RPC method: {method!r}") + + return dispatch + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def _build_argparser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description="Standalone LeRobot SO101 env server") + p.add_argument("--port", default="/dev/ttyACM1", + help="Serial port of the SO101 follower arm.") + p.add_argument("--calibration-id", default="my_awesome_follower_arm", + help="LeRobot calibration id (loads .json).") + p.add_argument("--arm-camera-path", default="/dev/video2", + help="OpenCV device path for the arm/hand camera " + "(empty string to disable).") + p.add_argument("--scene-camera-serial", default="409122274720", + help="Intel RealSense serial/name for the scene camera " + "(empty string to disable).") + p.add_argument("--camera-width", type=int, default=640) + p.add_argument("--camera-height", type=int, default=480) + p.add_argument("--camera-fps", type=int, default=30) + p.add_argument("--scene-camera-width", type=int, default=1280, + help="Scene (RealSense) color/depth width; default 1280 " + "(720p) for finer localization. <=0 falls back to " + "--camera-width.") + p.add_argument("--scene-camera-height", type=int, default=720, + help="Scene (RealSense) color/depth height; default 720. " + "<=0 falls back to --camera-height.") + p.add_argument("--no-cameras", action="store_true", + help="Disable all cameras (state-only observations).") + p.add_argument("--max-relative-target", type=float, default=None, + help="Per-step joint movement cap in degrees (safety). " + "Default: no cap (matches RLinf SO101Env).") + p.add_argument("--max-episode-steps", type=int, default=200) + p.add_argument("--step-frequency", type=float, default=30.0) + p.add_argument("--max-joint-vel", type=float, default=_MAX_JOINT_VEL_DEG_S, + help="Max joint speed (deg/s) for paced point-to-point moves " + "(reset/move_to/move_joints_delta). Lower = slower/safer.") + p.add_argument("--motor-acceleration", type=int, default=_MOTOR_ACCELERATION, + help="Feetech servo Acceleration register (0-254; LeRobot " + "default 254). Lower = gentler ramps. Set 254 for the " + "original snappy motion.") + p.add_argument("--position-gain", type=int, default=_POSITION_GAIN, + help="Feetech servo P_Coefficient / position gain (LeRobot " + "uses 16, factory default 32). Higher = stiffer, holds " + "commanded joints under load (fixes gripper landing " + "short/low); too high can buzz/oscillate. Applied to arm " + "joints only.") + p.add_argument("--auto-calibrate", action="store_true", + help="Allow LeRobot to run interactive calibration if the " + "arm is uncalibrated (may block on stdin). Off by default.") + p.add_argument("--urdf-path", default=None, + help="SO101 URDF for FK / EE pose. Default: " + "~/.cache/huggingface/lerobot/urdf/so101.urdf") + p.add_argument("--output-dir", required=True) + p.add_argument("--transport-port", type=int, default=0, + help="Socket transport port. 0 asks the OS for a free port.") + return p + + +def _build_arm_camera_cfgs(args: argparse.Namespace) -> dict[str, dict]: + """Assemble the LeRobot (arm/hand) camera spec from CLI args. + + The scene camera is NOT included here — it is managed directly via + pyrealsense2 (see :class:`SceneCameraD405`) so we get depth + intrinsics. + """ + if args.no_cameras or not args.arm_camera_path: + return {} + return { + "arm": { + "type": "opencv", + "index_or_path": args.arm_camera_path, + "width": args.camera_width, + "height": args.camera_height, + "fps": args.camera_fps, + } + } + + +def main() -> int: + args = _build_argparser().parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + init_output_dir(args.output_dir) + + arm_camera_cfgs = _build_arm_camera_cfgs(args) + scene_serial = None if args.no_cameras else (args.scene_camera_serial or None) + logger.info( + "starting SO101 env server: port=%s output_dir=%s arm_cams=%s scene=%s", + args.port, args.output_dir, list(arm_camera_cfgs), scene_serial, + ) + + env = SO101LeRobotEnv( + port=args.port, + calibration_id=args.calibration_id, + arm_camera_cfgs=arm_camera_cfgs, + scene_serial=scene_serial, + scene_size=( + args.scene_camera_height if args.scene_camera_height > 0 else args.camera_height, + args.scene_camera_width if args.scene_camera_width > 0 else args.camera_width, + ), + scene_fps=args.camera_fps, + urdf_path=args.urdf_path, + max_relative_target=args.max_relative_target, + max_episode_steps=args.max_episode_steps, + step_frequency=args.step_frequency, + motor_acceleration=args.motor_acceleration, + max_joint_vel_deg_s=args.max_joint_vel, + position_gain=args.position_gain, + auto_calibrate=args.auto_calibrate, + ) + + shutdown_event = threading.Event() + dispatch = _build_dispatcher(env, shutdown_event) + + server = SocketRpcServer(("127.0.0.1", args.transport_port), dispatch) + bound_host, bound_port = server.server_address + client_host = "127.0.0.1" if bound_host == "0.0.0.0" else bound_host + print( + json.dumps({ + "event": "transport_ready", + "kind": "socket", + "host": client_host, + "port": bound_port, + }), + flush=True, + ) + logger.info("RPC server listening on %s:%s", client_host, bound_port) + + # Park the arm on SIGTERM / SIGINT (launcher kill, Ctrl-C). The handler + # only flags shutdown; the actual parking runs in the ``finally`` below, + # never inside the signal context. (SIGKILL / kill -9 cannot be caught.) + def _handle_signal(signum, _frame): + logger.warning( + "received %s; parking arm and shutting down", + signal.Signals(signum).name, + ) + shutdown_event.set() + + for _sig in (signal.SIGTERM, signal.SIGINT): + signal.signal(_sig, _handle_signal) + + _start_parent_watchdog(server, shutdown_event) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + # Loop with a timeout so a signal landing during the wait is observed + # promptly even if it doesn't interrupt the blocking call. + while not shutdown_event.wait(timeout=1.0): + pass + finally: + env.close() + server.shutdown() + server.server_close() + logger.info("driver exited cleanly") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/deployment/lerobot/geometry.py b/deployment/lerobot/geometry.py new file mode 100644 index 00000000..a9d260a0 --- /dev/null +++ b/deployment/lerobot/geometry.py @@ -0,0 +1,256 @@ +"""Pure-numpy geometry helpers for SO101 scene-camera localization. + +No hardware / lerobot imports — safe to unit-test offline. All transforms use +the convention ``T_a_b`` = pose of frame ``b`` in frame ``a`` so that +``p_a = T_a_b @ [p_b; 1]``. The world frame is the arm ``base_link``. +""" +from __future__ import annotations + +import numpy as np + + +def backproject_pixel(K, col: float, row: float, depth_m: float) -> np.ndarray: + """Backproject a pixel + metric depth into the camera frame (meters). + + Args: + K: 3x3 pinhole intrinsics (of the stream the pixel was taken from; + for our driver, depth is aligned to color so the color ``K`` + applies to both). + col: pixel x (column, u). + row: pixel y (row, v). + depth_m: metric depth at ``(row, col)`` in meters. + + Returns: + ``(3,)`` point ``[x, y, z]`` in the camera frame. + """ + K = np.asarray(K, dtype=np.float64) + fx, fy = K[0, 0], K[1, 1] + cx, cy = K[0, 2], K[1, 2] + z = float(depth_m) + x = (float(col) - cx) * z / fx + y = (float(row) - cy) * z / fy + return np.array([x, y, z], dtype=np.float64) + + +def transform_points(T, pts) -> np.ndarray: + """Apply a 4x4 homogeneous transform to a point or array of points. + + Args: + T: 4x4 transform. + pts: ``(3,)`` or ``(N, 3)`` points. + + Returns: + Transformed points, same leading shape as ``pts``. + """ + T = np.asarray(T, dtype=np.float64) + pts = np.asarray(pts, dtype=np.float64) + single = pts.ndim == 1 + p = np.atleast_2d(pts) + ph = np.concatenate([p, np.ones((p.shape[0], 1))], axis=1) # (N, 4) + out = (ph @ T.T)[:, :3] + return out[0] if single else out + + +def invert_transform(T) -> np.ndarray: + """Invert a 4x4 rigid transform (R, t) -> (R^T, -R^T t).""" + T = np.asarray(T, dtype=np.float64) + R = T[:3, :3] + t = T[:3, 3] + out = np.eye(4) + out[:3, :3] = R.T + out[:3, 3] = -R.T @ t + return out + + +def kabsch_umeyama(src, dst) -> tuple[np.ndarray, float]: + """Best-fit rigid transform mapping ``src`` -> ``dst`` (no scaling). + + Solves for ``T`` minimizing ``sum_i || T @ src_i - dst_i ||^2`` using the + SVD (Kabsch/Umeyama) with a reflection guard. + + Args: + src: ``(N, 3)`` source points (e.g. camera-frame). + dst: ``(N, 3)`` destination points (e.g. base-frame). + + Returns: + ``(T_4x4, rmse_meters)``. + """ + src = np.asarray(src, dtype=np.float64) + dst = np.asarray(dst, dtype=np.float64) + if src.shape != dst.shape or src.ndim != 2 or src.shape[1] != 3: + raise ValueError("src and dst must both be (N, 3) with matching N") + if src.shape[0] < 3: + raise ValueError("need at least 3 correspondences (4+ recommended)") + + c_src = src.mean(axis=0) + c_dst = dst.mean(axis=0) + s = src - c_src + d = dst - c_dst + + H = s.T @ d + U, _, Vt = np.linalg.svd(H) + # Reflection guard: ensure a proper rotation (det = +1). + D = np.eye(3) + D[2, 2] = np.sign(np.linalg.det(Vt.T @ U.T)) + R = Vt.T @ D @ U.T + t = c_dst - R @ c_src + + T = np.eye(4) + T[:3, :3] = R + T[:3, 3] = t + + resid = transform_points(T, src) - dst + rmse = float(np.sqrt((resid ** 2).sum(axis=1).mean())) + return T, rmse + + +def ransac_kabsch( + src, + dst, + *, + thresh_m: float = 0.015, + iters: int = 300, + min_inliers: int = 4, + seed: int = 0, +) -> tuple[np.ndarray, float, np.ndarray]: + """Robust rigid fit ``src`` -> ``dst`` with RANSAC over :func:`kabsch_umeyama`. + + Drops outlier correspondences (e.g. a mis-detected tip). Samples 4 points, + fits, counts inliers within ``thresh_m``, keeps the best consensus, then + refits on all inliers. + + Returns ``(T_4x4, inlier_rmse_m, inlier_mask)``. Falls back to a plain fit + on all points if no good consensus is found. + """ + src = np.asarray(src, dtype=np.float64) + dst = np.asarray(dst, dtype=np.float64) + n = src.shape[0] + if n < 4: + T, rmse = kabsch_umeyama(src, dst) + return T, rmse, np.ones(n, dtype=bool) + + rng = np.random.default_rng(seed) + idx = np.arange(n) + best_mask = None + best_count = 0 + for _ in range(iters): + sample = rng.choice(idx, size=4, replace=False) + try: + T, _ = kabsch_umeyama(src[sample], dst[sample]) + except Exception: + continue + resid = np.linalg.norm(transform_points(T, src) - dst, axis=1) + mask = resid < thresh_m + count = int(mask.sum()) + if count > best_count: + best_count = count + best_mask = mask + + if best_mask is None or best_count < min_inliers: + T, rmse = kabsch_umeyama(src, dst) + return T, rmse, np.ones(n, dtype=bool) + + T, rmse = kabsch_umeyama(src[best_mask], dst[best_mask]) + return T, rmse, best_mask + +def rotation_to_quat(R) -> np.ndarray: + """Convert a 3x3 rotation matrix to a quaternion ``[w, x, y, z]``.""" + R = np.asarray(R, dtype=np.float64) + tr = np.trace(R) + if tr > 0: + s = np.sqrt(tr + 1.0) * 2 + w = 0.25 * s + x = (R[2, 1] - R[1, 2]) / s + y = (R[0, 2] - R[2, 0]) / s + z = (R[1, 0] - R[0, 1]) / s + elif R[0, 0] > R[1, 1] and R[0, 0] > R[2, 2]: + s = np.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2]) * 2 + w = (R[2, 1] - R[1, 2]) / s + x = 0.25 * s + y = (R[0, 1] + R[1, 0]) / s + z = (R[0, 2] + R[2, 0]) / s + elif R[1, 1] > R[2, 2]: + s = np.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2]) * 2 + w = (R[0, 2] - R[2, 0]) / s + x = (R[0, 1] + R[1, 0]) / s + y = 0.25 * s + z = (R[1, 2] + R[2, 1]) / s + else: + s = np.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1]) * 2 + w = (R[1, 0] - R[0, 1]) / s + x = (R[0, 2] + R[2, 0]) / s + y = (R[1, 2] + R[2, 1]) / s + z = 0.25 * s + return np.array([w, x, y, z], dtype=np.float64) + + +def sample_depth_patch(depth_m, col: int, row: int, radius: int = 2) -> float: + """Median of the valid (>0, finite) depths in a small patch around a pixel. + + Robustifies a single-pixel depth read (sensor noise / dropouts). Returns + ``nan`` if no valid depth is found in the patch. + """ + depth_m = np.asarray(depth_m, dtype=np.float64) + h, w = depth_m.shape[:2] + r0, r1 = max(0, row - radius), min(h, row + radius + 1) + c0, c1 = max(0, col - radius), min(w, col + radius + 1) + patch = depth_m[r0:r1, c0:c1].reshape(-1) + valid = patch[np.isfinite(patch) & (patch > 0)] + if valid.size == 0: + return float("nan") + return float(np.median(valid)) + + +def detect_tip_pixel_by_motion( + rgb_open, + rgb_closed, + depth_m, + K, + *, + diff_thresh: int = 18, + min_area: int = 40, + max_area: int = 40000, +) -> dict | None: + """Locate the gripper in the scene image via gripper-toggle motion. + + Given two scene frames that differ only by the gripper opening (arm held + still), the changed pixels are the gripper fingers. Returns the centroid of + the largest valid motion blob, the median depth over it, and the + backprojected camera-frame point — or ``None`` if no usable blob is found. + + ``cv2`` is imported lazily so the rest of this module stays import-light. + """ + import cv2 + + a = cv2.cvtColor(np.asarray(rgb_open, dtype=np.uint8), cv2.COLOR_RGB2GRAY) + b = cv2.cvtColor(np.asarray(rgb_closed, dtype=np.uint8), cv2.COLOR_RGB2GRAY) + diff = cv2.GaussianBlur(cv2.absdiff(a, b), (5, 5), 0) + _, mask = cv2.threshold(diff, int(diff_thresh), 255, cv2.THRESH_BINARY) + kernel = np.ones((3, 3), np.uint8) + mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel) + mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8)) + + num, labels, stats, centroids = cv2.connectedComponentsWithStats(mask, connectivity=8) + if num <= 1: + return None + + depth_m = np.asarray(depth_m, dtype=np.float64) + # Largest component first (skip background label 0). + for comp in np.argsort(stats[1:, cv2.CC_STAT_AREA])[::-1] + 1: + area = int(stats[comp, cv2.CC_STAT_AREA]) + if area < min_area or area > max_area: + continue + col, row = float(centroids[comp][0]), float(centroids[comp][1]) + dvals = depth_m[labels == comp] + dvals = dvals[np.isfinite(dvals) & (dvals > 0)] + if dvals.size == 0: + continue + z = float(np.median(dvals)) + p_cam = backproject_pixel(K, col, row, z) + return { + "pixel": [row, col], + "depth_m": z, + "area": area, + "xyz_cam": p_cam.tolist(), + } + return None diff --git a/deployment/lerobot/kinematics.py b/deployment/lerobot/kinematics.py new file mode 100644 index 00000000..588a46ca --- /dev/null +++ b/deployment/lerobot/kinematics.py @@ -0,0 +1,100 @@ +"""SO101 forward / inverse kinematics in the arm base frame. + +Thin wrapper over LeRobot's placo-based +:class:`lerobot.model.kinematics.RobotKinematics`, pinned to the SO101 arm +joints and the ``gripper_frame_link`` tip frame. Forward kinematics returns +``T_base_gripper`` (the gripper pose in the ``base_link`` world frame); inverse +kinematics maps a desired base-frame gripper pose back to joint targets. +""" +from __future__ import annotations + +import os + +import numpy as np + +# Arm joints in bus-ID order (no gripper); must match the URDF joint names. +_ARM_JOINTS = ( + "shoulder_pan", + "shoulder_lift", + "elbow_flex", + "wrist_flex", + "wrist_roll", +) +_DEFAULT_URDF = "~/.cache/huggingface/lerobot/urdf/so101.urdf" +_TARGET_FRAME = "gripper_frame_link" + + +class SO101Kinematics: + """FK/IK for the SO101 arm, expressed in the ``base_link`` world frame.""" + + def __init__(self, urdf_path: str | None = None, target_frame: str = _TARGET_FRAME): + from lerobot.model.kinematics import RobotKinematics + + urdf = os.path.expanduser(urdf_path or _DEFAULT_URDF) + if not os.path.isfile(urdf): + raise FileNotFoundError( + f"SO101 URDF not found at {urdf}. Set urdf_path or download the " + "SO101 URDF (see toolkits/lerobot/compute_ee_pose.py --urdf help)." + ) + self._kin = RobotKinematics( + urdf_path=urdf, + target_frame_name=target_frame, + joint_names=list(_ARM_JOINTS), + ) + + def fk(self, joints_deg) -> np.ndarray: + """Forward kinematics: 5 arm joint angles (deg) -> ``T_base_gripper`` (4x4).""" + q = np.asarray(joints_deg, dtype=np.float64).reshape(-1)[: len(_ARM_JOINTS)] + return np.asarray(self._kin.forward_kinematics(q), dtype=np.float64) + + def ik( + self, + current_joints_deg, + T_base_gripper_des, + *, + position_weight: float = 1.0, + orientation_weight: float = 0.01, + max_iters: int = 60, + pos_tol_m: float = 0.002, + orient_tol_deg: float = 2.0, + ) -> np.ndarray: + """Inverse kinematics: desired ``T_base_gripper`` (4x4) -> arm joints (deg). + + placo's solver advances one step per call, so a single solve undershoots + on large targets. We iterate (feeding each solution back as the seed) + until the FK error is within tolerance or ``max_iters`` is hit; the + result is the best effort (check the FK error if it matters). + + With ``orientation_weight <= 0`` only the position error (``pos_tol_m``) + gates convergence and the target rotation is ignored. With a positive + ``orientation_weight`` the solve also matches the rotation of + ``T_base_gripper_des``, and convergence additionally requires the + orientation error below ``orient_tol_deg`` (used by ``move_to``'s + top-down mode). + + ``current_joints_deg`` seeds the first iteration (its gripper entry, if + any, is preserved by LeRobot's solver). + """ + q = np.asarray(current_joints_deg, dtype=np.float64).reshape(-1) + T = np.asarray(T_base_gripper_des, dtype=np.float64) + target = T[:3, 3] + R_des = T[:3, :3] + n = len(_ARM_JOINTS) + for _ in range(max_iters): + q = np.asarray( + self._kin.inverse_kinematics( + q, T, + position_weight=position_weight, + orientation_weight=orientation_weight, + ), + dtype=np.float64, + ) + T_cur = self.fk(q[:n]) + if np.linalg.norm(T_cur[:3, 3] - target) >= pos_tol_m: + continue + if orientation_weight <= 0.0: + break + cos_ang = (np.trace(R_des.T @ T_cur[:3, :3]) - 1.0) / 2.0 + if np.degrees(np.arccos(np.clip(cos_ang, -1.0, 1.0))) < orient_tol_deg: + break + return q diff --git a/deployment/lerobot/scene_camera.py b/deployment/lerobot/scene_camera.py new file mode 100644 index 00000000..3f93ac74 --- /dev/null +++ b/deployment/lerobot/scene_camera.py @@ -0,0 +1,108 @@ +"""Direct pyrealsense2 scene camera for the SO101 env (depth aligned to color). + +LeRobot's RealSense wrapper does not align depth to color and does not expose +intrinsics, both of which we need for pixel -> 3D backprojection. So the scene +camera is managed here directly via ``pyrealsense2``: a single pipeline streams +color + depth, ``rs.align`` registers depth into the color frame, and the color +intrinsics + depth scale are read from the active profile. + +The arm (hand) camera stays under LeRobot (color only); only the fixed scene +camera needs depth. +""" +from __future__ import annotations + +import numpy as np + + +class SceneCameraD405: + """Color + depth (aligned to color) from an Intel RealSense (e.g. D405).""" + + def __init__( + self, + serial: str, + *, + width: int = 640, + height: int = 480, + fps: int = 30, + warmup_frames: int = 15, + ) -> None: + import pyrealsense2 as rs + + self._rs = rs + self._serial = str(serial) + self._width = int(width) + self._height = int(height) + + self._pipeline = rs.pipeline() + cfg = rs.config() + cfg.enable_device(self._serial) + cfg.enable_stream(rs.stream.color, self._width, self._height, rs.format.rgb8, int(fps)) + cfg.enable_stream(rs.stream.depth, self._width, self._height, rs.format.z16, int(fps)) + self._profile = self._pipeline.start(cfg) + + # Align depth into the color frame so depth[row, col] matches the + # color pixel (row, col). + self._align = rs.align(rs.stream.color) + + depth_sensor = self._profile.get_device().first_depth_sensor() + self._depth_scale = float(depth_sensor.get_depth_scale()) # meters / unit + + color_stream = self._profile.get_stream(rs.stream.color).as_video_stream_profile() + intr = color_stream.get_intrinsics() + self._K = np.array( + [[intr.fx, 0.0, intr.ppx], [0.0, intr.fy, intr.ppy], [0.0, 0.0, 1.0]], + dtype=np.float64, + ) + + for _ in range(max(0, warmup_frames)): + self._pipeline.wait_for_frames() + + # -- capture ----------------------------------------------------------- + + def read(self) -> tuple[np.ndarray, np.ndarray]: + """Return ``(color_rgb_uint8 [H,W,3], depth_m_float32 [H,W])``. + + Depth is metric (meters), aligned to the color frame; invalid/no-return + pixels are ``0.0``. + """ + frames = self._pipeline.wait_for_frames() + frames = self._align.process(frames) + color = frames.get_color_frame() + depth = frames.get_depth_frame() + if not color or not depth: + raise RuntimeError("scene camera: incomplete frameset (no color/depth)") + color_img = np.ascontiguousarray(np.asanyarray(color.get_data()), dtype=np.uint8) + depth_raw = np.asanyarray(depth.get_data()) # uint16, depth units + depth_m = (depth_raw.astype(np.float32) * self._depth_scale) + return color_img, np.ascontiguousarray(depth_m) + + # -- metadata ---------------------------------------------------------- + + @property + def K(self) -> np.ndarray: + return self._K + + @property + def depth_scale(self) -> float: + return self._depth_scale + + @property + def size(self) -> tuple[int, int]: + """``(height, width)``.""" + return (self._height, self._width) + + def meta(self) -> dict: + """JSON-able camera metadata (intrinsics, size, depth scale, serial).""" + return { + "serial": self._serial, + "K": self._K.tolist(), + "width": self._width, + "height": self._height, + "depth_scale": self._depth_scale, + } + + def close(self) -> None: + try: + self._pipeline.stop() + except Exception: + pass diff --git a/resources/franka/calibration_boards/franka_charuco_7x5_25mm.json b/resources/franka/calibration_boards/franka_charuco_7x5_25mm.json new file mode 100644 index 00000000..d21400ca --- /dev/null +++ b/resources/franka/calibration_boards/franka_charuco_7x5_25mm.json @@ -0,0 +1,18 @@ +{ + "type": "charuco", + "dictionary": "DICT_4X4_50", + "squares_x": 7, + "squares_y": 5, + "square_length_m": 0.025, + "marker_length_m": 0.018, + "margin_m": 0.012, + "dpi": 300, + "image_width_px": 2349, + "image_height_px": 1759, + "print_instructions": [ + "Print the PDF at 100% / actual size.", + "Disable fit-to-page or scaling.", + "Use matte paper and mount it flat to cardboard/foam board.", + "Measure one printed square and update square_length_m if needed." + ] +} \ No newline at end of file diff --git a/resources/franka/calibration_boards/franka_charuco_7x5_25mm.pdf b/resources/franka/calibration_boards/franka_charuco_7x5_25mm.pdf new file mode 100644 index 0000000000000000000000000000000000000000..d77409bbb381592f3f0fedf8b0a6668bbc2febb2 GIT binary patch literal 139890 zcmeEP30zF;`#+XMmMnE`X}VXbD`Y8a8m^nfCHvTFPzh1CR^<@MjfBv(72@Ki#+FDL z6?LsyD*8;z7>sjta#*J|NX1~%GMlwp z^R)((Khv`O^us}Qe>lft@p4=1bt(okoEEPW-`RX6+)>4VvCL-OQmw&53`eLKOt4*M zw-l~B%n1G%zih=i+cjDSzmJ>wyREgO4Se^vW$V^XwOupLapfvU2U{4yaB?Us9M{ZP zwa6MC=674?W!AQ4YZkew82n-DU)yYek5!-ks91Z*;x%4vRDSB(t$PohKDv6p_BAvbI&Aoek>h?h9d99XZ(*Eu=6xVpJ--r~L0XWO5?djj_E+aLJXfrH0GPMrMv)af(- zT)22C;&SAbt5I?BH*X~*{(Jk*<0ns_C8wl5Ps_~8e*GpVH}CDclG3vB58RKuib_!$ z)|wO>9s5|RcidO;Pay$?1QZfbNI)S0g#;85P)I-_0fht<5>QA$A%X8s0%wJe^PfbI z2uTWEod5c8YZGC@Yn3-Sz$0xXapa^@t zFHwZ`3?aHY#wS&F!%HV-i!i@OBCKs>X0EUuy+UWT2usKlVczM*WyDsRU=Lg?dmSEE z$a526<5UUr1$apmJ=0lWEW(cYh%oJNIxmA(d~Jg!{xDvI?P1Ua4*rUp|9E20w~og< z8S7^ij-S+Cbr#?9>FVwWOp^M!jv1pn;HcB0ij--lYE=Ii^enumt4l1$mf$=!N#uqR z7Vm8WoL@NbeOLS4O-a70SdWUG*gzeWf4Z3CyUwxevd7Fkmtgydoip0f_m*VdXt`s! za)|1K%k`+_>1a?df(gpdGtDU*dUf@4z=xH^(Y(o-4u>3QM7?8*)OX ztX|~lc;;<#3U`g${>W+C4`Ut=>y6EPRXVtJ{_J{|VwpnoDKuXLHQ&2iKDR_z^X2Y- z5#B-9ax5(;?urYxIqcx9rFZ*}c6&yqtfoAP)pNnP!(Z>pXLE~0*vU>Q|5BdW(@fVZ z7l84nqdHZDotqh4Ika!(2I?dTD_DdrV^!>97s7>?rfj4#1*Vy*my0m+!rEozh0g|8 z_U+46q23}fV|eK>iuINvED{&$?u;ht8APO@4SC^sCRZgwggu&MsYLyJOst+UK6LFu zOr&jP7wFrIu#nj?9P0cJY+*lk=>#VZ-$jHy-j@Hmo`<=jX2uo)Z0f|<4 zJwr$*TEchTGFg13c_J+DWVly7!rxfW2=k0#VsuntV${ir31vDyqP#6Q5(%6N?U6;o zYa;CJ?@p!lJdIT>e>nIs`BFsFSMXzTFg|Kn)+s4HyLa#8qe0s`QN-3x@Zqzft zb(-YF_`t-dq`|~IV@2js-ZdILG)KXnUM@t22rEB6vZ8rClekcL7=L~29gV9JAi~Dq zVp>qolRB*YLL(72I8t~WUiB!hXW_>Llh+MePUEV>bxlpJV(WSEe`%njSoUHab)qCC zxSkmXNqB6c>r@>lz#lv423e^t-;qEum8#=*3Buf)rjr=*ut>IN&kB$6=i z087(BgjG)K<*})rCpov_0ljqU2UWh3IjE7vtWz*cEKl6-@42_#a2KVMXI(jlrL7;W z8PnbR#%-+)M*bqK4QDeW%;(+bcROOau zdJosielhEg-tib*_%}OeHXLVcVTe_;dREU95q8mE5O9AS5qAO4I~|$|3O4EKhYZ^9 zrR<1^`8`!fnK0Adcka(iSbb$jcJEWC&w@2MML>f730G8?lhtu z7+P7UMc9_2SWp8-c{pzne91S2=)zVOVGmP|iLh4*G@`G+;u1|Mi_kS~GNU4WvtG{0 z@tUcsvHP5@Q*SM=;aWS_P6u15R2hy!{WQr@XzY|rt(Qza*{>`rbc1?;Q8kN2_|D6N z)sh%fVz;&Y?(Kd>g0Uy(S(^D3C#~0TxQgvKGopJ-T0*D@Q|OK#NT1H!<@zKplF?~? zIxleiLAM3ks}u1qL(UJ!u;*^77gD2|Saj0p z*7lJY-FiD?#?*7JrqbDfV zi6D*GpCv}@#NMo3pgrpTW1SLVFx9#TeT3dNFxAr-=it#BX$8vR=CDw19ka213`-8X zJ$y4v^;7{&^`iMqVUGx|5O$YGeP{MqP{Mi}NF^M`BK+p@saLWZ;M&leYIdrQ2QJCkwbeQMaFwf(|JlQa# z-%DVI5iB4hCF3NKlD%e3EzovFv<`%dtK+SgEuT2uyCc+c^+r{u_&tKa8FOn)Ggrpk;_%)7ELh&- zPSSeY_BXd{JW^_x*KN~Zw-YSRgjw+HXrwhiX@Ll17lU^Q4!~g#@Ce5CUrUHEa1Rr} z1=zY8e1et{Nqi9o1L+Y$bWe;0pP+LJy+}Y8EcgE!7l~PT-1q;-YA!PA;;KKnB-c0m z()H=J@*sZ2WYgjuGKr*nu$tolUo zrKe#}Tp9@x0W~w-O_)u0i($IxCeBz6yjWcU$5mz5MvZauq!mO&G#{ z9kdMmKy`nxA}28i!U2wO{=m->+lZ(jwd~HIG`_w_d7u`gJeF?ts&nLZA6}> zs^j~k$VE<>GGn2``v@QIhx2?Gj{u~$;aAsT@xdJokLb=$-jL*}d?d$|J2%+j?Gc~8 zTQpl6_b6x>GHn5V)6;PZFETUVv&D^%w6v%r*VA6!ZI!TjMcq52&(_lG1j3EbaqWQ$IJ##L9=&8TwncBK3KVO3PUHY(h_?zEOwL4+cxx?^2 zmm2mAmdaa=S7vg$2}AiUrt~qPTL8|43$Y)G+|p##VAJ+mZ%5NL!A;;!n;UtBob`ry?ERn z9vCyfeI%wp$g1Q{qi3S&2Qasi6UXW9y4Eas_OnZl7gJ*+?*$f4jmfyDwkq8G^`;x) zCNMos90SouMfJ#&vaTzNwEVtj7Y-f*bqdC4m2kin1uG(THrSHmBi5ErgT2ZhEW+-A zW#ACZpBJ;S=?)70qFUh*RO~=;bWp@8rIt|xf@TaWbB=sg4!u_J7zsiu1*GXov^P~?3cSJ@H6`lm8afZ}q6j=%ss6;nNKa_VqYfPr}2o;ZB?hsxK*n2ge z7wHu@WmFHf-NzJ8Bp8mh&De#6HQi)$)uSLG61G67#$Xfo*ky)* zt~x~|vyQ^lHl=mNuZ507A%7Hx#Fv2dP_barkT{>oR?9c(^wfs^u*l@LllJ(Z@L#p> zbZ*mamDw!XJ}0v9h4ca-?z(3-cmd+l7t#w@1iC~YmlWQB7XUb(*ZaT?5EtHn^-C7s zfc5K_5}C(r>(MjIQFYltf)(O?77tYWfrk90Vo z#)kAaB54wL0}6b^-iE)#+kgTe(%VP`ZzFtjyk?u?qRlbSgquzY{R`(C7jew{6kbjL zn3~=5(5WzSzNP76BA(!#t1&hU?0J7?v0C%)Jcu9nY{PB)K0yJ4J%fH7oMYZ>gp=-Z z|2xsyP?uCgJfzDP3C;momjup%q%H}Z1F|lObPmLIN#GoibxHHVIS>~nk!~8iryT}v znz%BlK@-FZA%4nt_uocnEE{<_Y&CcLTS3uaOW#A6@4Xv76?@g}MtI!!c{W2br?HDs zT*b3JfiAA%dCs+jDjsWphXSbL8FC~v5uAiftHVP8G z!I0>ML{0gMi0pFl7(MagNdPm*lTF%r# z!J0;-vG(rs%y-K>LiHj{JvsE`r-Gs1n+Io2QaeY&0Oikrl?#Y8J^q;rb_!>Br2d)e(=XE8ZpXD zVCb?LEWw_D3fPA83TefUUaA5rzyVMJq|KNeQKkgW)0ZWcq`(;Y&_C)Q0ZPLhxZx3} zbSf0*@RiH}tZvJq9rwL@;_XHIk@t57Iy^f2F|t#7!sE_rL(WZ|(rozca-(;_0Av8m z_2_=4+m5WQ`eovm8gW-0!WaFjqHcF^;5j)^c%1)KNG zsZ#f*CoQv&zw)LJeEa+(7Jp>qPFlFFIP$NAS!tI0{O|sjj zbBzmrYk&H|nlZY^gPy%5(IAZzV&ku{!qqX}CAhg!cjIYUe`-aakJ|S3t>#||<}G*p zzHztK*vL-9Mm(61{p1uNM|J{oq=i!wzXaB5uMncgs93O4+5>RpCSBm{&s$og4DQY_ z5`6-0PDvU)GbEZO!p3B*B?Mr2C)#f$(I+K7!j_lld=WOR3Vo75EARYJgcS)m67)$o z{NW_{I=c*)09Y2002Wn3gxv=#Kh(b@CEm|E6C$ z_xco?=?|E|lo(-%Y%U@@M;;_{=3A<5yE-rA30YDaBPj7w?3@Ofl%Q?v{F4vez#7BT zA&UfmgQ^ODq3)g-w)m*kPlQ^;A^@!7LT_t`&rDCb2bY(4n@M%6>wYlALzSQq@~NB5 zCnefM;0B3k@6iBwv3E3`hz0$h_uM`)?>IuZ6j|=3vR&2lXkbdSmWoa36Qzula`D;& z@K9kdvWC=xEKbEM`_8NIf*?N`dzZz6{CF^_DkXq2@X<&IdDKK5%~JiRypeE;bT^B^ zT0U^dr1qr?d-ivQSSkXBMgfgM=uI^Uc{HYw5TVenw*ZXEzCtm^Bxqu#B-o0^nTTzw z|B28y=ns@rEV~`Z4fG!ytf;U9MG+N3bNYULh&tn_2MJ8OE(WFr+-5Q#k&2rFoJJ;e zWFcdyvK=y~AU@&{|3$jo{Wgk&z_>F1G8W zH6eQ^F>B`3fmpM1n8sqL6xvb*R-aLRFmDie3gR+Eu=-NM@mC3sk|IPf`o>v-H?+As zoA(b{_h>y|T=xjsd0n#bQPc7O6h0odYjNu6*(w+tqRJSp)(~QoFN>>xw4o!13QVwf z+q>ry?Uq0VW~jmRXA%9EWI7m4{_TCG)0pi#Pw(q1-5##==kVLQnmK14nVPpOoA<(W zwPU|sNnLz)PDl;y-D1~R4KLLZVM$P$SxYK)t}U+atVgQni&#ZQqi!1*ywANB70@#$ zXk72hSc`sc!!gykP|3nXM86PvkaF?KXYI}I=U8NT_ury^19C4#`H$-3HLhp)Smkz5 zljxstxq;6_ma8^KCNjT2b;dQ&DtToxXJ+yT{XF&fvn~vazFT9vH~X`DY;TjC3wg(V z^BOzG&#q?V9##|{p-kJ_TCHicdhL6H##9s@f!+2+wSi4zxYboyKO$6n&?7n&0I=b!KbA~vpE%3OENC#3t#NWbr|`P1re~V zwj%7zH4*kJTX@-@w}e@Fh+PVHPx?~DM~#+-{ZkVS*RF)R2=fUy3@)TaR-~+4n^U1V zd+@evtFN-J9?re)3@cd%x{5Rc$v4putXvV+WkYv7E`b09;Gb+Ro@&gga4vHxdQPNe z_Z(fB0!iC=fQvrO5Mg^cL^{;{83-qHDkein;S#R+8-cOLuFLvE{4Dlg6Ff)B1gQbeR9Bj+k!)Dinn-kh}ZrO8n zm$ipq?;Jf6|18dcbgLQ(RZv{w6Kv!Ek1W8M?*bT;F~UHilSR2f0XT2Vh-bo=B5YJo zXSelb{=Z&4?9^-h!iFN9{=aPD&mq1Pa-dM@v!3uAZ|=vN8@@F*$l8`09TTMazXa!m6lXXiHHm$W@ z`Sb$p5o{mS4!SpPLWW`(xuiohfaw>~of=dVeWoefx~iF`&V6kFFC3&>v1ySScu_IjB_t0oWK9;;?tSiy6_^Cw=Q?*s)o^Ps+#^5Mm$V$>f~&33)-E0!3gxY2S`EHNs9 z{qyl-a^;dEYjHQ-KzsWx_lFefF4xo9sjQ1>?u&;e1d8sV-(TOl9FOTQ(bRi^#{+N&8r*TJsBX#DN(`|@CD9PDG)F6XX4j{G$J4Xm^!n(pl zmmbyqR!u=N=-F&&AqAzh=Vql4zRP)DhzI@%9V!BeCBgi4;88s*hdH#OlX5a9<)lv| zeD{;o|5Uz+K1YNtVOAW(3*qbbjE68BbjhX6C1eSnk2l&`RWDN}t37p8$!gD;c;(P} zB?u==;*QE8aj7iDgbJn}V?mmF8%Sw}r4WA0bMs`LN&le&(fw2@z*Gl2@Y6LskEYV)6H_KvafGC0)Im z`DAw1y@}#H~uD-SG3-u&k5s-3&=dKQ(`|*P@S~4@QclaA-d+$$L@2n_7u)g#(EPe(GW*PuMCYS;7iXVjW^ zF;jVyVN$F7K(2e&KqvcCFDADYhOdoC9@=*mQ@Le_h|*iIY-Ku9s4)=42$Tn91CBch z^*yYyn>bM4&!7D$=8-4;NX{hgnE>bKN6b`Sb-*-KfcoCtqvpb_UwR}xeGS4$k`z5Y zNzBR)_Kk#`UC#90ldpK3Pn@N(P>n9A2NvMv^+9czxf}l#Kg_xL!+O?CBRob`)AbduQxZ5cO7{CKa-$uvcRD>k zey~IDHHN<3Y}O2J^OiZIam(MTQwb|BZP*qMTf0Tao(uORPafUJrp)S(t9!;wSrG8Z z?2VP9>DM{gkTeez&FRooadPo&qq=!i9~E_8G{B3Hbr>l5de;)<+H!-nrDTcc{mt42 zXG#OMsf^7+4jPq0p|na-XSc+%K$>@7(I!GtVplcPMsd5ydQM--)I|~5P_Pi7nw_8= ztuHjBsx)^)92yobJb>pwLz+knVHk4^SR7=#Rup0!X`VEpi;6JhB^Dvd*ZX?@5vo0v zlo}jaVvX_*zK-9EtN{dL1o|;+kCgOd{@tk*rVr$T+#PZe#>?kcRcSX_RYe)wK;$o} zS{$LAl`r?hq>8u!<_|)JO7<2~IA1Sb(qR5tq+X_=9;@SS)cO+@2SX`=GIYaGg)-U-Oh4I{YA{y_SM~-hg!Q^{j@w77ZWQP{oqV_d;NB;9-9M@%pAXTS zzU|7WBU9G+M4oiP>(X9+7JD>zgAR|MX=xW16%*)qDqnwe;ojtuOS^8W4LLS-ineZd zldUo9z^;SfNZo#b@j~9;g0d+&^c~KIcGN>(YOZzA@AKkPig# zbqGX1E6LANkT&=V?tcfI8Fw}>dF08cf6rN^C$8@CFSaoJ*v8@7cFIqm-Jxj7Mp9+T z6D_OWM3kWfq$18y{Vvd|XXkjRmFRAc=;^4od|1@X%OefXoL*q&*w3%GS)0xS$5>I2VnbN!Q$4!9Cu_7?;cnKXWqHrCq1FI@kI??t6aW`{zC9ZXKA-4qLTu zn$;$t6rafT7E_8_=y>9i6TY7RE&$gLFg!u7fkZ5WS@~;dmI{)Gl8Q_ht#RdQqCt7m zA%}}P{-|=*h*;^TuIx>wMo>mHy5Ti~g5$ri{w14{Djcn_!N4k%#gNr(MUF6gzD+&@ z#U;hm@>ogI(uD&BNsUT#R5MaE$V2|vtW;}DAZn=eFfu40`5DR&g2O@)xPZt3ve8hL zC5MU>e1j-Iibe^|x`|?tMC#fE`&t)xSfWk_7*~Q_4P_>6f}z{MnzOY3Z+P>PBY|q5l0+F zGz}3rtI_uRi<8_^IU>Re$oydAXgvQ5zfrR9P@Wwup|50VMHq5A=cgH6oB!7RS)L9LthYXL;bhvib^-z{czzO@xFmCAI3*25h?Gwf)wY#>+ z+4A4lXX5+r-W`i^S&$aYYzfrwDhFvq+hagu359f9y2JvGy6;k)H^2u7kwS=$7Ro^W zeg-~%6-6WZ`18gZ1M}5yf5WF{zDJLo{kq`4Q*UGjRIOz&+oaC~pmdz}Kj zHQMd5$X+m7ZPxu^v!7nF-+JZohql+cKUXo_f6?#{hSsZnTG?6}>s6jMmVC2=TTaxK z{aIOa`vfrzwD$~J;;gH^$4xtE?KEU>M~Yp7++>KYgI5#y&a<4X6A4AZ@=`3 z#ef4VO3ih){P7?`M|%%LJK%x++Rw_8UVA$16N!=dEp_2$5!U*YV8VtJuKy(wW{V%M zNH@NIO%UynkdxrdmA6#FpAXWI~jB)VGb#%dGTCo>&#PFUBZz=)Ve~BR;*M{tf zp>9YQbcHY7Os4r6cg>v7Y|h)7#fhJ1IONvf3!>MUE$Vz>;Gf&Hy#LiIJIE_-1k~oM zXd!>kOnUowE_rMfJ{zUgxmxiwCd5Xx@x3*cX|E0nZ}d-7YR;Nojn1<*buW zH;li1^uo?z3zE7GUpRPh8=9ZTc~iC7uQd%zW`#x1<`O!|i@dbkXsqVk`Ny{vh7HA9 z?3fwcz?q%-!2q<${iyVJ2Vb7_Y<|UQ=exIWrl$7R8q^!y^%u#OUo)k)b~2FO_Q(u3 zy#Cw!1n1xUbk5%IqWD z4419-yO}+_*J!2n-f9&1sIL?PL?0Ara|$$&y!uAE>gy^kS^JT~9+RbfWI++lGFHVt zc3}~qBT_bsZ=qfe_^)gdR56vL&JV8a+n1}NhdL1(5F3^{50Zdz)qmymoKQ_F;Cjf; zxsv;Qy+lUUNM%w>iWP(%-yfx-r&vK%p2g>Kw2Bo}rUVC(LI+i?m%3UU&VkQ_&%ciq z_e zBJ7T%2wNFa^OM&my#?S`c;lJ*p7Z%Gg@jXDx*OA`^!#4uj3b4@d)B@3@8*g9F_}mI zHc(|`HRVs&Rz6RyyPHfq_Hh&ha_zClao(Ow030j+jQ&8ihsVVqUypXNB+}%R(%S(&&FrCX~_@!Yh zUV6J&-|e2Tc86~&V89QuzD%x)KmDUN?q-sV6u9G`O=>PIphJFLnBhGU$4I5)x_Zz8&32xi>P3>4@G7h3QR zf%)gjz9Yf{9rc%h|5$i=Ej#@r*_;+#N$^cAWpduB5Y+#nCJ-CPaD<`yE0X4DQzfbq z875eCC`1eDv_(ZiK6mm+jOv>IZBXJ?^22YNVS;0VVojpc9w0IKgQ>_}F0b2$O2t%( zk|L1*#mIc=W#uekLQ(PS9Y=xikOw1ji&8EOO0$PD0fmD7NoE3tbb@F9O2|Y~{WPH2 z6L^RS*bGuQRr6Do+WF)P%cs<+d5x*?=yLkk$mLDGDL1ASm6?PIM`4Fi+X1~rDHn!} z{G@XI{xXx)s}dcfA%%v>Hh+sG5h{@<5hP{3qr6!k>3n=z;Z+xw_lQLG3nVo>otI9Q zJXr$p78mO6j3J9N3q@FSC|WE|5qgE;<(n*LY}BS&D>653qR5}jeI-3#Qp`t3{4mL^N_)_vaN4G_3g@^V83kbOIA$*YbKsSpCvJB!T{c_ZjlHp~}q zL%9!F%18WS_grp~O9d*o3;dOmTofHxy+^gFR|+#b)x;cH}i_xr)|^eFm0R_#9P{BlR4lv_Q4+s z8es%Ill&ovou?Cyb9gxK8HtyMeDTr!up#kGD7O_e5z+~-dAUMp_XGfW2u>+3(S)*a z)>PuH4iV7#ICt{>a*ji(ASohh=un4*q5p_5j}J-93t-wQtyJlft96E=^x&Zqme(bY z=F*36-A{a66{{ZiLVsRV=uawlH(#A&-C?~u1E3~~d2hmb{|Yqp90j1z>lfN;joTq3#(dN#vJapZ}_>x6_d&{Fzst z8tn_>W=;C@rxC|hx~HCRF(_b1@Qj%d?mz4SxumiE*Amd-4XXRZOxpl-YZ;Nm7h!Pp z^nlz`A_;WsoI)=W&;`r=p_&Ag&%;T8h+-E5j_Dt|;v!8XlThn8j*{`POj@Tj`~6*` z&IZoD#?>xOnEHPFxLv)KLo`3Ut!1FsI%SYY&!;!hE!pcXEa`%ib+4%IMHOU|Ua8g{ z5Rtc9pV#;i3Ml;PY3hiZ-e2*AhcZ=z^bLc8p0O61hniD)1|19@o=P%;cvIPbera zw4&Ea)7O*1SA0|x&+y|(1JE3Xs$w)pEMHfG<_O)81YnLW2_7$U>G3<(g1ypI&=SRc znoibzase)+$}JK`xoMgS1sT^Ka1&luK>TFuLFg`SCc^qMh5z{Tma!`L;-x2no{_X> z6(9YtNMj51{!%9TvM=Kcg#m;P+qN|38QZRCSkWD86Nf?mnGf8^r2wEaMWgvrKqjR_ zjne zLELmQ{xn2{?K2>rDL(jNM54#$O~$iAp{mHgfSHb$4CWi8dL-#S5a!$`7H0URK4`W2 zuu3gxVlAw!;*1WI>& zdr$(D4$rzQ0ZL~|fzlsUsKOOU%nH?)^{H_svR=0b;DjdQ_rien1?P{CZeA-+Q|k^$ z1~l0VO5RZ~RR(4>{WJ+Pnz2)EIdwe5ll{s9XvTYxG~+D7cV1qt+t$n;9`_~57;C@t z=}!vySY5}N93To1PWdR5yaBR9iNS?`vvUAX9?2GlSOJGK)l;K*3uP}?^-e{LY)8Sq zLKh+f;?6-M%j%&sj^PyR0v~vLJhg8M1o`@}qZ|QHEY&x@u`nlk8^}2cV-bGy_|(Ly z^%T&3nj7hyaaEru5)(C#LHWZOdijc3<@Jy%FzLh-W|=;o&LJMLFzFR!uZbiar4+pf znnnlteFXBEau_vu5oErD;>Dm?8B|)bazG*lq*>e1p%gz1YqPb%8yJbsySs>2fRZ;oQE7GicN1e-PgmMOjKIt2bLs3 zqVoRR=3cpL|7q8J%)9vOZft>}3kg!_2_S`SIIr+~U|v&^SzqH@oG(s^fcw(iBbJTR0{0OCbN=3w)nZWd6N9#z+ls!W7+@|{?(sM;Kv)4fub zw`yABEHXJ<4L*J&7+OL>;pxIqrcf>7A>ce3=}E}3H&OIbYEqtQr0z$OPa|y)bS0FG zw5xuLDpr!qB1HNV6;L=0;r;{jiI?!o&s7QeEaEp&e8X~_f~YbaJ=QLhp>QNxrQn4e z5oVsvqEkICvhG5SlH%mk@R5}r-;D*1Dlm~*95KPd!}#l<z zir@~13$~F~Y=#$C%McWhQ9X8ekf;;A8MSlFURIY)mG`hR`p|-G9oy97ZLW@H9skjdUJi5Hbr_?tWO-cpM)>H zGg5z2xNh>uTe{DY^vy{Gsa$@RD3^)7S-W67mBE?HTLi<-2>L%_83^bnEi=q>iw>^W zd+mI)wAudNx>z$FV>M8r>5^{ROduO8wggk6i&=vy$sT=;i=-q58G7|(_Hy0&@}*v; zSN|z*dhaw+IemOKh7HF$44Sh}Dms*J(P%GTX4NvldH)<0+gEw+ANwLk@d9RX@YeDq z@2LB4e0<77C!H4f1&p=TaJU*YaMcYPo6hte6^o(VB2v(XWW9=Ka#bRrtZ$Mf9J>_! zy3&c3E_kuxk^4-2Z_;tj+;_OzCHqnQ2?Gjk%09*$FaCYS*q^)4h>gVv^$+y_W_=^( z*@{u}5f}pfUd`|6XYl)z6?QA97%uC-a(W;AHxpM)+_-vIhtb|1Z>4(o$2;AT)jgn? zD4Rq!)sZ5u&#GQ#sIYg_95U!?E=qn55FPSR`M!;EL}4>5X=X>)B4Q#9xVQ zSr9y;qg>VVHIdI&2kb2`o5Htn6*2gs6 zZ@W;CJ392~cXq%qfYdGf6hM*2ZsI0 za4#^wzJBllxoQ{SH}U4slJQ3IhnN6i2v@aFp@Rpf_x%T*QXGaEd@MW6Wf?Rf-pqa@ z?@;kwLOqGg6ikNLfe9x^gx!tYcX3G)H+=-B@G&Iiez@%fdQx6a{3_G2i99Awxzwbh zh)$v)4RQn`rRTJqM>3jUNP4=+#FBozC=4BTy<;sz?lJ0u> zj0Fj+rP{Ow$YytJo(@SR_8UN-{DpdZt~d8Z!r{Iw5?}qa@|zO~%K_zIiNgYQBf&pn z;2(O+1Pb=_YD>JNmmi;4t|GOd$dnId<8E~bgA_uJd2pogI=t#pTbmgDps&rPJaUy`9 z6sbJ~s*$>5VBM7e|4dj;swV*i+r)+=iVPVMJNge}b>z}QCwR-Ju!sMRY@PJ-A^@A#3$sEG3i&>}g(1=hJe+>K_R% zKm(0s3~lL8s(&kefke=wlu?XG0(VA|9i!OZP+3*PHg7ysw*d}MJ#?(HI0wYsf6oX3 ztKFH7ejBul{eAiyj%wS#@1*&bd`;IA16)RpxZSzgGymnO5c}C#;v;Oaj?OO`Rs~y2 zpq00W5jG}6l^_E+iDa>2qb^{d$8NuThP{dU zW1~sles>bo2bV!;Ki2E=J0T#pD=jZ&%mO!5e z0zaDrU|%;tn}h6#>~iQ2hQP?mDvY`tSGu7a1jEB;O5A%Rl(3txz-@XeXZ+x%Gv!)jCcVFBmE-1KlJS1SHIVinNQwp zNBWxXw*G5d>aoEcyR{mw-6<&H8vZqSZ?3j?%rZG&955&NMx|1O--82xuPn7?o4Fmj zH%;F`@6M^5-5cX}V)@Azfvu>4Yv#OQMhWX?(17AMFMV*^lVRg%O1x#e}?)%)4q;)7o91MFw~ za9-mm|7P82^tw*u<-;tMv?!pZB_|IQ`rIoe{@r<%6FK6=ePQgAUitn%WCl~R2T}8a zhAVv`Y*4zjn|4-DBa0cDK-ry&{g!BFQeSO8V%#2+yVN3IG~NE7XA%QvZCMcX6D zp-}Qg14>b={M(@GP_7*qiB+j=WZB~PQkz$grHxqU$U|M;^INJE>aGlVy)zQ(f@F+D z3Ah1v?6*iM!ZV_lA1n%~hxCIP#8uw3r_U>P57*f#H)1p+uape%Yge*t|84Cx%$sL- z{4#Z{(uLN>izA^qClD&Occ&5Uz~IdSPMR%6aHy9XK~La8U^n`P5M9{Hz^Ir4R^_V% z8qqg;9M>mhou?zwDXHMd>;!;+@+}T&@$;~wR7 zLAQ8luFH53OmkW}Y;MUjzl6Zk^O$3zY?eNrckjY!$DuR4X6`=RarJxsJ9)<=jk_8T zjTu?|;nvOb@h#^axp-T9hx)YN|1sAt#x!zb;7Xs9pFeL3*&%J56-ak;H=GBY0kx(e z^o{@ak6icQTm6@m#=UQQZLTo@8?|=#>^Qzdw^ggg(!5udi#I1qkA3~DruYn9#4|_; zQjaJ$@CkcwJSk+H@jTlL<{E(|`n_B4?KpOYmuj~#vb}YEPbJ5~`kY|oxI~PWI6#0p zcD|eXRTNa^OP9XP@1#5b%X~bu|9prYHV~hN(2Dj9v%e0fBQ&Gd4ZSQp=X+V7=sZ8W&64g zblw|U@rf8s{2^rd0wQUGHU7NmOjUXwG#I#@?oW&`6Z~r66I^b9zb=F;X@n4Jf$>R| zUD>71*%0fT29#=(GIND(ITfQ;Lqt{&(3m>KWrTN;;K&sSZSKVt_eROwNwa71Wxvc` zBPX@*kW{e$!o(#uWmbP&#kvG)mrtvheG4KajUskgOXwC`_;^@znT0}nVX!UdHBmW% z`#z!QO8khl>z?|6Bv;JSceXy?_v^Q^{Qk(axCitB$;sMV)m}__I<@(Zpeb0d<>Z3U z@>mHe^6TNi(#R$P60LBgeiOfmJk!GI^Lg`MaXRI9d6#zCkGuUXw_=a*ADfJb-*&3s z(7=YP!=g}Z=KolU^}Xd2r`mkkLB9xLkoTL6d3OQ;He-Ld&`Ik|tcOO&ktwTTg1;Ad ze=TRB#%h#JUDmQyYo6&S6TQZjYTcpHJ3z)ks2a(4Q@C4GGBsbLgUGeaNvWhzk!8x~ zZe$&R^qzGG^4(OR)g730z$iQc)e=P+xzg*@y*@gKUiVw#UsHrOYtDN0YiYa_IYfkE zio7AZM=a%&pb#O|1FG`^o%cd+V^r|4&ef{_Y-FN=LJFo|2{UmPSK~P(-Fm43d_=E< zM$Hd9^s-emGE1!8>6zK`vagcnR`ktL^b(S%ztVYudFD0xrXK4y=oLo?ykuHVH2mm4 zW0Cgap{;%!v*YKCzXA_uw=8*kudx~Z&w)Tls>^Zunp;BgGX zTlnY20?u>hHBPRbGk~KutbQ|K(9TmVhs}y6V1Q3st!Rt}38-)0EH2#oW{`F?uhZ$o zOuTKP(i^|0||p@kT@OW-2bBXl9IP z>^$oxQ(#}ByE&p~aND=Z^JZQiX~;U2etWf=UvINEgU+nc#?Gq@9@Ap9H>0sY>7|Mb zG#+jKKOjVVw!)-Qw61Ncu*5$(0mW_OqK00y{%8KX9(OHlR|BoI_9OLvJG!;l7Moyp znl+Ouv{~=_=&Bp4coF1D3+{qrVzk6=tGn49#CQE?I%1V!aDmGfaY~$mX_rd zao9rdgrVy{u*Luk9RM+)?{yF2LT_u(_tR7EQ7WZrrn1LKI6aHLlby-rxpQU{5BT^I z=gTh{!Edbf3a^ct&;Ki7R7xP!aVc65Hn#cVEMT02Hqd=cI-V!nKm%1;FtmZbf6c~* zp9(~2KWzkR@Q=<5hRHmgOY?}z{Mu!RA$}uP6U_zHh znmDEj&e~=;XItZIMSuo$fpj$khxcq3A0=ZmoM5o^n-j1nlX z1^Rf%fYQPVC`L#P`uS;jFwZo)=v>C#l|TZwbS3Bfo%f0RS0A?PZUiux_CS0w-Jh7r z=0kqpiU@$aDL!hrG}cKNAKUnEFAy_6CURL~#>cXug?_;JXgwl*=k@&5 zdTgqyT^9@!9*{&%cwd-sQzMVD2s>0(&mxICs5G^DkqEnJAh=KG0!L*;Xr2l*r;yA3 zMG)H}>oKVw6RA_=vXLe1I}LUc&!d9;@7*Q}*XZixsxQOZcEXNyTBs-DCvBFuAv z$yXhie2*Wjs{N?_)PObRLeJ#`r-%$XH^P&o7x5=R{aS$Ez*m7|_4rBFU$uWy%c%S) z!kz~bOM?09fEo8$xf1MbI^PVxZTtj$HGLZ4yFX?jLVKgwD;n#wA<8Rd2QGHXIQr?1o*P_Dp*BwxSD=Bh-2*)o~g5urB$ z|H~gN!tQ}Y9fJAuVm3D3-jQz~%2ENTG49bWxd`@-rW3KCf%Bf*C*~n$x9WK-W$BJ$ z779byuaCio=BoI^NpY8CMwNMH`9&_3r-XupwapmhsiO+=R43&bwgBPB)<;7oI}~_l zEkQc5Zs5c!S>kzrv$jcZT4P#SxbOg;vkZ8(EQDdqF%VP$`vN2taGy@Zod-#svCmOK z{H{uPU%CqsJUse2s=!b^Q4`M2g{PrTafO)E0{LL@<`=8V1_)ewZ!x($LynNZ^%lUA zwXaajk`*+uvY!5oRpPq}j3oe)H!d*{dWEygXE1J(kA^7m`m|P2Dxm2%8ma&k{V;g( zB*MIdH}t)wRZi1!2R-}49%(XW0g^WeltR6Pu`)`7TxLkFDu`;?(ZDzK5qjHzQkcd# zSM#@s1*y*wVM~}52k}Bsx%P~QaAnZwg|Ni5%PheEJd0QMomb&iWBY$vSK>DttPK#@ zTjZm#h&K4`pC7UdLAy-U7JI%w>UNir5rvxlI)^GnRsW|cM6|w;S*LpP9h+4LTEop( z6e1I4A)o@*av48W2hbOv8e%G0^%7->LBBr63nbzqnUzMN!F5)DjXx#*91yaVoC&fb zLNOA8Jkk9)YHs7YuhIOMh$d#AeuyyfWESZ+KGus<|STt{qK=SmYg_8U<&g$ zp5=7%-Ts3s2AbM@R7pM`b$(&ql~G5w4$O@_xn2j10@>~@|L@nZ{*3r}D@G>yUh*4| zd18)%e(C!3_0Y5W z_Ivd(%URo%O=Dl29pp3mx_&CeUJrZF8Vg1wO)#0&{U4YS6olf)29Ai>kc|HO#vvFa zuO~`uU+lk2PIomF^{3pFM^(!VAun>de#I~jvWac65*Qt_aTk;|1`qpnVspsDpMuc} zNfV2(v;N=PuqfgbiN%1B7CDt4$&eAnszaI@luHrW58#!G-f@`?-}D9#aHnQF0u#9l zkpVHH<0C7YOAdjBWS8G#?`T|=0Eo`rVp=>hMcMmb81)t315TNS``3ulKa|kEJaS*9 zvJmv=$e1T(sUJ#Ouxvi#2ubfQJzMKphu1fASCg2U(|BrM3t>-nPY9rjt2|}4I@0=u zbNm~IRe;+ffm5M9vPgJMguVUUsT9;GJQH#!$Tr52bqD24_*99HAp>}VR0*P8-~3&! z2DccqOBT%M+0Q-LW~%WTvqk#)XS%oSctvaH+Ua;)8{iC9H%bz81>nKWM`ROSD5aX2 zLNB%)R@H|W7}Ac-0LLtL9Y-+P?w$Kj=a0KF)6!;r@HrA}Q0QFrlfZiTi&m?5KY2Ra zlWY(TtxaqFga$2a;|cG{YeksS6THOmbN93=DxOXx#lCDG{UmytODiAGK1a(=UOuID zd3VU5gky2Hy=LxP6m($5TdUn)K=n4uC))yO@bYiO{2UnxG) ztl0%~m&(|KlLBN!j>y7CSBfgCKM2x%bMeQG2Y$w^?nyZ>v7nC zkOZFH^_{~OY#cxH@D4B4>{(z?*5Zf*Gx>8gb7$x9;F3pPJ3DJ?Jq~|=;=#WQukSQl z(R_-U#$n7Nb3P2(!!6rFiT6)_qSen=B8!tR8g_o?VxX}*b>f|Owu@W;ylSjQ&N;Yq zOJL4cIhaed{Tm1Zm%^R90Cjd>`(WOHLn7?Y5Tc_6P-mYB$6rM`5`E6Rv+9v@$=>P8 zGsDRf6I!j>uJbJP=E-wYnk{U#s!ihgvr)jk3RyR>Y~@L2R7)1vB-K!%%x5OemZv%R zWrybb&6s-W{H9B{JK01XSUmmZ3N!3{i{UM^Q&S${pQHYpCdZ;8Gr){18YR_%67JapKh_waw;zSB?tf|)ueSe9*}lHuC*=1?x7^XTs&%fZlt zdKB-#ksW_v*idub9gM`iVh81`r*u+xuWg&|vOBtd#JMR141 z1yk^n{SA;7X$Mn8Ae$;+jWG6O#_&Mp=&ZMd)GRL4-AOi?dUm0BWfu9Mk=XRA?FYw zUh(>^R#@=jTPDqZh-jC7p$KaZD5T;~E$L89Gh$1N5-}m~;;1S4ny=+Hm{A0jsLsIX zX+?tKD5JAA2Uzp9+z05_aUa7&;_*(#v&6i6ss$mnay( z*@l;c*YtWHTZHutA-X!=q*r#sOD973!0(X=Ya5wG6SkvQ=s1GcnGcj|C$kMY<(0i1 zdHEuDdv|WdZqNN{84;U1%xWekUjzdg-0A|_MSs9AjU!^G<9Vr}ZvmtAUqV6DU(E4f zmb^T0YW|$L`7xGZXU*IF%AL@GWSiT2`$mhfIe>gmFHt47(gb@VA!f4!$cS_!$pL(V`sr(wj7UJ58f(Xz|t?kON-vPU4x_g%=2p| z^PC(d!{j1C(xnWy$EIu>zaE0!7f#3hvJ-n76!EZO#GRK6Excpe z7w4RB6?P(zV;8g8CD0)4Z_jIc&tx9n9NT-hrMLL3{>ed0%qu>UyU%e!)e) zm$0jPSr^Xli?_e2xxSMog-tAM-eTf;j0|$>E z6W84RS3xLbAi;<=t}AVcmxB?z@(zasIB&HScpo=Xle~D!UJm_N2018xkvi2L# z^>jSLySspWG2mWA)YAcdyt%ut_88xHm0S0gPG{X^K_cG_#>Z9Nv|(s*(PnleJMx3- zd0}V)Ve#H3z?n0|<;9F04puP{kVI+F(~*7;(7cm=4;Ty*zXxdENxx?`_&s@~g8($| zq=Rrp>L5TGj&u+f;1d7f%SBt~=E?N!CI-ANGm0IqF0u=+2xQFjygBgHbKRatLKcST z3>-WmpV<$5tPhYEGk`2$S79{%_V;~xaGF!ZZ)5@&a?gLWizDczoOx>yLiT}k;Arl*t+52R-Sv3Cy&l#lm`8AbGvR6Z&xA}#UMiGIY4M4!$R zCd6-aA1KJbpLhK06thSS;!po=vIy;`TS`Kva$~xXTbq3IixW)WeJEa<0+2pB!m;_r zbtDoOLGHMtGzwmKlI2Cg)O}zHiHkhR52^PmcbJkjh6kiOI_rUhKhj^QyC;S%mGURI zsLneSwFJ^ziyywJASbnM3gT|cGqTp~$}MoIFq=3oEWK^@tU}Q8OA3sumgbjl!6x+~ zYi%UCB`)2*av%Q`na^rj%f9G4;?>r7FXyk{76Le_6TAOCw(seVj$>bW4Zj7}cw8H@ zL8Q7NUC?!$6ezVBLZCArf{kmbS_PGY|2;MXK95~T7fdE`PjtXN!JA6(Eiv#(agqN> zCi8kM@eauga^z0gXT2!b`W{~=%9Pzb!FkhRn-3q1jsc#hkyHRPKlIm~oV+rbGxJHI z!BIws^BKQgIK86Nu5O(UcOLTI-T;tM$`C_8a|{Z_qi94Qf8JPQs0#S)Z^&kNL4VbQ z^vsq(+N^SrMzn?RNU&A009ft26z2_qFNwibic2)1EK<*(yZEGE_Qi9TydEbT7i?{P zbLmy>!+Ni5_TW9v!-{DnKfr|#2JM%QydU>s`N2MiJVzci-Sh65&cds@*g4(q+SxCj z(G@G?2iq<7N8w0EflQ>QS1t2|d{b=a&(RPK?{o^NB(4=yS!yzGY6D&Xbk)u8w-$an%hQo6htesWMid ztR;ee2bqVFSy5HiH^SJjWq8%Ng?v0ERX0NU#FS%baP@SJQBjSy0q+Aa0D*|F>z+tQ zj{tH}4Tlm%BfVCwdmx}SwLY)$BYE&n(m{#wJkmcJAeH)lqF7=L+?|Rgh7y>*KksDN zI};7w>d@tBlM|m89WC{^>a1UwG5FWaP0C(ozr?I+FxIEBAQW*lwFZbvoJ?pC+DI`1 zNw(_eMnHkS!7N-;_3804q{Q*&upY`g2hbzd-CaNax^MgKt)B0D)WTt#K5f5QH(ibL zSBz7Nmc)uy%w%>3pn`I=poVJXVMM=NE>obzEYTyLm2G$~!VEMW_g%EBEab1wq5olL zoZI}0YQwN_-(UkY5RIo8fh?l-zdsD}pAZ^L54zl9t1hr7hDQ?5o)*4OyYuVEOFb9O zG;4eLmwz@a(|En3EM=3WkD@htqM{{yS9Ym0U?g@<`wZ%M-)qgGC6U`=GUSUgeY9J@6tv4{w&n_nq7G8?CTZ@^LlybT+5l4 zvHYjvd%-pv9mek-*KW(d|LUIfw4e&onNro9)(GT%`Yo~$6%{lg<8rKE5w?s~v5#F? z1nmJ*Hd5vdOg~HVSv$;!C=IUtg&F}wR1GxcK$sZiN2clzUAqtyXbZ~mmI2t5A;bIEogAY7~qlnc{7#3nRD?sIEGVspUbEBwUw} zBm^W{;h$!_?lNGc6D{GpZka5;(ma3+oDBD>A?-RMVQr)u9=wmL%&mJj>5wVtX^_K2 zan?``Lq8Xy>h^@KBOl!8G^1Bk_lNiEXQi$!15^npEmsB4>#Ew zvmFku@5a8;Pa=jwnB(UZybw4ZCiZ6SLY%8e`r5`?4032wqzlk!n^ztiFMKoX$*BU^ zlZ)mvg*_s;Vv15|)r2$+;7t^t2!(+nm&m2R<*GhUgw`?u#awuXUcO>hISfQDs|98C zHrlU~VG5EF2ZCP6x1t3QpYAJ^5T6E(EUTwUQw*nA7rLCa2d2#wXmZwn9SGcGQ%NK- zk$Vk#F#-13qL&*v7j^z2O=$h^#7LCKsYEgnn|gq)Fd4rGEF+FMe|&UvDvK8s3Z#`~ zTCT{mMG>?;jWyckIQ&ptbkn{UC!T?(9we9C?AMxYOJ+TbX8c6-NM7PKQ&mDuIeu)* z0j+x%Py7+OdTE70xa~m>>U2|;;1egF^xSI-NKRp2E2#_ieCkw_*=FZ zDBchf&htfQ4Gizz?pIW(dnwAlXJ*ilp5bG9_l0s{T0-a}a_}hsPmZNjeO{sLs8VGh z$JI}hkmDLV<(7*zhr`%cY<0`|w^lo5hAmmDGp6KruPtZ8PuO()?UmPv`{tizpRGeB zYh*FKIQe?W{IG;G$Cu9@IAqJ*0by1;{k>MT#fEI0GTP5p{CMOOq9Rd0f+>Jd5)`@x z6&sM_q?aCdsNffqPCQ|j>Er1f;t`9aX%+uNX%KX@zZHPVFu=$kAx_UV@)WSgo0Vxbg zrX#NlG;M|afI8#{j_;3Bfr@Ea4%1H~sWUdzcA$#Pqt;%1yu?4MU5hx2k+$p?o2}B> zXwZ3sN?af)Madn6Q03pqk3lTIdhqW*rU_B*?-lgZqiD-){k!egBo#WJe_J&8)WMR> zo4Yrz*`m1vANdh)4CF%prV+>2YH)GVHmA`oCM0jT|M#-MHP7e9a_2DbY#pYz{pzms zai_a$2kbOI3w(_YBMqcy;XPfqC(vyP&Qp`;TuUIAvG#W;aNzq6ITD)k_^QO2ZO}55 zuH8?1UCwh-|D}nF$Eae~_UN8K6|#(JALT#Qbr$icpX*o+n|8gMJ@e1C^8XAHYi}KH z!}E)ZH{Uq8BxEIScg)RY_JtD<*=f0>EGP6`#Z+$DA>svCt!N|qKiwf`@y@RE$Co%C z>9zH$9rxOtH(mO=u{1TBH@nl(eUKnO5ODoEK+)q}BsfA+^wfRm&+u@@j1jANq`mc> zJE`)rtyWnHO%^3Wda69qGu^8?I6{&p)QVMMVv%i@KUf zc|iwqh-I4RKrD07TzdJ;Sp|~m`cBOx(nev8WKQ#pG`%uhXa!3~NORX=I<3%PX(HE4*!>_+ZU3hz$FnIXrAGQ0ICTs}ms zLwt!Fk_Saj?DrgLzSh1=rsIxivyj-(Wg~~B^|Q-(bC-L1hI4XXSG$<ECu6BKd*# z+W>!E$gV=P1M4>lcZ7_P1lbJ*Yi&6dqZ~z8qW(GF>f_7lZH&jQX7RE(eEmPJT(Z|Y z_44BT14lFlq;4NrVZC%gso9@MU1076XJk0#rwg5XOlSB<>1%!hMrC?fS%6 zi)}&b#r~Te`x)J56a-FRt~&VC^EUx{);n74oj&r6AQXh#AP%_M{lqMTPgkDp_klC` z-Mg7v)%e=kYtJ_uKXIYW?&NErhhkRoP8LTdSoP-lLpkLXKOlS;1P@?F{C4PaoUddl zF=(XmS!1z$R@6z!#814t@du47FNcMA1|_tQv0*Rw=DuGT;=RT~yMxl|^8eTsBMrf- znA{I%4k8$@kR0TQa2BaENrfl3PM0XhI{&3m^b}b{iWDj;5Q);!=m=b$5lQ}RKtvc3 z$bf68lg&!u$p69ObqJbIbTkrWKUZyp|2~Wdl1HIRQDdgCC+Y<7B!4!<$Uxc?KeoLn zOAmpjqN-;@G?o2GZa`Ef4PaL&Gz((H6d?`qo~~=j%OB@Jd983s4+(4psouWpbPtNu z5}D7ccRR=#+d}$A4O2E0*$@<5203C1hX6e?#G#2kfl!AHTvF;XAo3K_)1~f86faNI zk%Jte9_rvke-lQNN(c0ap1Y6m|JJnRKf}e1;g`Q)FJ+^}>~ub=%}ke_+umUs?H%!X zlZ1Xsuei$OR(58U$v>|4|M2?*TqfWg2V5r5EV=*3D*K1lz-0owRnPzNqyf*9TLHQy z6nIys>FfDF3}^pm=vq}Zy|eUf+TGl*zWG7>8e^=EmzI12pIA2$d}5tmm&V!?HkprB zM^~(U`^S6DVHOV|i^uz?hKleoFyt_>OROYh!{UF!`+>KCUI*R=dMy8+*53W!Dz2CI z=KmA&f14WyJah8oZ}aKjcDGMYtGV;I(Bd4szhm4(x;B5pfxV7FbGVi#d($87>_X9_ z``fl;q<0=Zu2Lqbt;{^Zk$ZyUbJLx`fH^HM{wMO0ukoX4o6{2?@rGTxX0M!L5O-2a zFQI3ShJ~cc{9`-nNlpge?*Y@*@^sa?{~3S@Yj>pKKFRkf6Qx-Mr>e8FF#ia%1rEs# z^hDjMe|E2KP3wISE#0dhuzTa&tri!TY@N9%XA4BR7l>-;)N_{>5N8;>%hNFR?dtHH78)pNrHcmGG&mijlN4*ZX?0ez*7~pxG zJ@@~F*4Bg8pHOi{S-+T|Lv{41YrwZnLdsDbtCI$@!GX1d_a1cRTK(T; z`=0^t&Dy^I@8UQ{>@6tRcqi6s0_z|s*5NknRR<#WP}^D9wtZo3$iZEO<%nTWBa2QW zFtxyisW2bl7=#?)<^_&U9q__xEUUJ#wwj=AR;;m$wIPc!Pz~->;g3dmd_ox5CM~eF z-4Kq#a#kL=Sw)8l4akYn5Qh)SIV>E#H_(m2^U-(r4!JRHL~8rW53DH#ztjWAZ*SY* zx_13??d}tw{+;N!ar0Ta_00NxR_S|z1JNesp!!sLT|MxA%OBtkY5e~g7UTijsCB?` z-N#$%54!`;$Zp*A{4eOH$3EcbDX%Q(jBK{o^S?im|7RU<0PJ#p0oFN>fp{loF{HQI|S!aXj%Gx?jkeAA}3 zdv1LaUDI5+JkB&cR=yCnoYZi>&iOI@(Ya>eKEhist74Ze<@p%5^swobj8!4Z>)5U* z&)M#;KAPUM=}V- z0_?pt6eTdNI3{$kx?`~Pn9-9Y%&}N@WWcI|&oZcSv zM`VTX+pzh{nbqRdNl($Z6m@HfPq@N^sSJ-5WU9G$A6emZeJ`dr&xsZ(h@8;MzlKJ2qo9AB$kGjL9sz6x4Cgnjo`S1(33N@qU*UhaYVMQi)Bnj|TK8_rv7gH~^||LTutoID zVPJn4I+|)nZvmv$jfi)eIC66zZQtHykbEs$d8-B#6t-4utf5Q&Tb%o$@Py zIt_r()XZ4V+)Tm5$Vku745ZyTKd&S;ucTPP7_3V_q_Q9tXa~?7prwgLKq&=NsCiH~ zf)qmC=#!e4UXr0;Vh*$l>XAaI6$U^6RA>OT3eHW=RnUhx(ZZI8W-y~aAs91(5d>Kd1?6ypil&-f#QoBN3q#W<;bGfpu`U+{PF8l!HAU)#OCyjoWKChX>0pQZOsCWR( zzx=!#Jo7>l02G1;c)V)|IK6$a-H-nUhdZ;L&cb9o8>d_LSfx>u+VZN11hFtwE1B-r zPM`M%pM4bycwf{5aP{m%Ap8wf5AH2{^k|~E%Gg&mUXFuII-Y6D=|OQ!j&n*t$#W#x>?Xq`c1 z3_PQiD2@7DP57lFs1=K~v~)0v3yR+Cb-qW8o4C2_Ohj__Ol10K?StzGa`{XXnRn%= z5N8ZJ$1f*Z?tK(}_-N)5YL&uqGJe`I>&bZ(;hEo5vu9+&NTVWFh?994GU~$_kS4W`8Ii3Kd^Ji2dYOq$nx|6bM#9uXD&ER!HVfS$at$O1`a=R6%W< z9e12s4ID>J;$XqbtQv^y?zxo{q!HR=q@|=W$81ngkt-E*1x*^So}VoC{W_>?ww*E3 z!*$eX{nO7cjIaW(k!}4CG{~tjR1WzBro`(a<2oTCyOJnYtHKEjWMc>vm4!ZjeVS%B^Q@bm zhl|>W4O=FGL3DFk8FpTgOo+3+tQ{b02S$)X3YIM{rr?zlrdk(VEPuDyAQ4ilm&P3= z3k5~KXejJ)4y`y385QYzf@7_7e)X2AvhC4ore(9W1uQ{hpHL9l zd>b_wV_q{4CZZc;>nAd4U0_woZ&6Byv4mV=%|Z>8*eAZWVFtebid~@5>Ok`2CZ2$=c;Ob9m4R$fkMqw zKs>Vp4s~~h(Y~p~H%~anEC}YjhPUpo@u5(F_{n_5p2kv>jqb}lQsIYcYk3*6%(7|Q zsLsr3)Jn{cUF*ITc%NNVEE)+uFRX>1j=ZNc0wi7^vO0Qa4u%qeadSDuGE$YR`i{c*i*tYL3Vqf~FW* z`SWW@C1UkP5bG{MEuEDPBOEUrHIRPC?^RrQCNyL@vP5srMPw(dw;Z rB7=3y4+AiPTu$4*c>Z79yRVhu>D2*Kl&}4<&lD0VDG43V(k}lFu@lXW literal 0 HcmV?d00001 diff --git a/resources/lerobot/memory/MEMORY.md b/resources/lerobot/memory/MEMORY.md new file mode 100644 index 00000000..ed1108cc --- /dev/null +++ b/resources/lerobot/memory/MEMORY.md @@ -0,0 +1 @@ +- [For placing in the white plate at x≈0.25,y≈0.105, high vertical standoffs can be unreachable; move low/reachable over the plate and release around z≈-0.03.](plate_place_low_standoff.md) diff --git a/resources/lerobot/memory/plate_place_low_standoff.md b/resources/lerobot/memory/plate_place_low_standoff.md new file mode 100644 index 00000000..b32c986f --- /dev/null +++ b/resources/lerobot/memory/plate_place_low_standoff.md @@ -0,0 +1,7 @@ +--- +hook: "For placing in the white plate at x\u22480.25,y\u22480.105, high vertical standoffs can be unreachable; move low/reachable over the plate and release around z\u2248-0.03." +env: lerobot +updated: 2026-07-03 +--- + +In the green-cube-to-white-plate task, the plate interior back-projected near x=0.224–0.252, y=0.103–0.109, z≈-0.067. Carrying the grasped cube to a high down-oriented standoff at [0.245, 0.105, 0.080] with yaw 90 failed, settling ~45 mm short (final z≈0.037) with a reach note. A lower sequence was reachable and successful: [0.230,0.105,0.035] -> [0.240,0.106,0.015] -> [0.248,0.106,0.005], then lower to about [0.251,0.108,-0.030] and open. The cube remained in the plate after lifting the gripper. For this plate region, avoid insisting on a high vertical standoff; use reachable lower waypoints and verify with the scene/arm images. diff --git a/robots/lerobot/__init__.py b/robots/lerobot/__init__.py new file mode 100644 index 00000000..4d9f8ef9 --- /dev/null +++ b/robots/lerobot/__init__.py @@ -0,0 +1,73 @@ +"""LeRobot SO101 environment extension. + +Entry point for the env registry: :func:`get_env_spec` and +:func:`get_toolkit` are discovered by +:func:`rpent.envs.base._resolve_env` via +``importlib.import_module("robots.lerobot")`` — dropping this +package on disk is the entire registration step. +""" +from __future__ import annotations + +import sys +from typing import Any + +from rpent.deployment.launcher import EnvDriverContext +from rpent.envs.env_spec import EnvSpec +from rpent.envs.prompt_bundle import PromptBundle +from robots.lerobot.prompt import ( + system_prompt, + user_prompt, +) +from rpent.rpc_driver.base import RpcClient +from rpent.utils.config import get_repo_root + + +def get_env_spec() -> EnvSpec: + """Return the SO101 env identity + prompt bundle. + + Tool schemas, handlers, and the MCP allowlist live on the toolkit (see + :func:`get_toolkit`); driver launch config is attached here. + """ + return EnvSpec( + name="lerobot", + prompts=PromptBundle( + system=system_prompt, + user=user_prompt, + ), + build_command=_build_driver_command, + requires_suite_task=False, + ) + + +def _build_driver_command(context: EnvDriverContext) -> list[str]: + return [ + sys.executable, + str(get_repo_root() / "deployment" / "lerobot" / "env_server.py"), + "--output-dir", + str(context.output_dir), + "--max-episode-steps", + str(context.max_episode_steps), + ] + + +def get_toolkit( + *, + rpc_client: RpcClient, + driver_context: EnvDriverContext, + video_path: str | None = None, + dashboard: Any = None, +): + """Return the SO101 toolkit (common tools + SO101 primitives). + + ``driver_context`` is accepted for a uniform env extension API; SO101 does + not currently need any extra context beyond the RPC client. + """ + del driver_context + from robots.lerobot.env_client import LerobotEnvClient + from robots.lerobot.toolkit import LerobotToolkit + + return LerobotToolkit( + env=LerobotEnvClient(rpc_client), + video_path=video_path, + dashboard=dashboard, + ) diff --git a/robots/lerobot/env_client.py b/robots/lerobot/env_client.py new file mode 100644 index 00000000..43e9aa8a --- /dev/null +++ b/robots/lerobot/env_client.py @@ -0,0 +1,141 @@ +"""LeRobot SO101 env client that forwards calls over a driver RPC client. + +Mirrors the RPC surface exposed by ``deployment/lerobot/env_server.py`` +(:class:`SO101LeRobotEnv`): a minimal gym-style ``reset`` / ``step`` plus a +``get_spec`` self-description. Each method turns one agent-side call into one +RPC against the driver process via :class:`RpcClient`. + +The agent process does not import torch; the driver returns plain numpy / +floats, so values cross the wire unchanged. +""" +from __future__ import annotations + +from typing import Any + +import numpy as np + +from rpent.rpc_driver.base import RpcClient + + +# Per-method RPC timeouts (seconds). ``reset`` moves the arm to its rest pose +# (server sleeps ~1s); ``step`` advances one control tick. +_TIMEOUT_S = { + "default": 30.0, + "env.reset": 120.0, + "env.step": 60.0, + "env.get_spec": 30.0, + "env.get_ee_pose": 30.0, + "env.get_scene_camera_meta": 30.0, + "env.move_to": 120.0, + "env.move_joints_delta": 60.0, + "env.get_obs": 60.0, +} + + +class LerobotEnvClient: + """Remote stub for the SO101 LeRobot env protocol.""" + + def __init__(self, client: RpcClient): + self._client = client + self._spec: dict | None = None + + def reset(self) -> tuple[dict, Any]: + """Drive the arm to its rest pose and return ``(obs, info)``. + + ``obs`` is ``{"state": {"joint_position": (5,) float32, + "gripper_position": (1,) float32}, "frames": {: (H, W, 3) + uint8, ...}}``. + """ + return self._client.call("env.reset", timeout_s=_TIMEOUT_S["env.reset"]) + + def step(self, action) -> tuple[dict, float, bool, bool, dict]: + """Send one absolute joint-target command ``[q1..q5, gripper]``. + + Returns ``(obs, reward, terminated, truncated, info)``. Targets are + clipped to the arm's joint limits server-side. + """ + action = np.asarray(action, dtype=np.float32).reshape(-1) + return self._client.call( + "env.step", args=(action,), timeout_s=_TIMEOUT_S["env.step"] + ) + + def get_spec(self) -> dict: + """Return (and cache) the env's static self-description. + + Keys: ``action_dim``, ``arm_joints``, ``action_low``, + ``action_high``, ``camera_names``, ``max_episode_steps``. + """ + if self._spec is None: + self._spec = self._client.call( + "env.get_spec", timeout_s=_TIMEOUT_S["env.get_spec"] + ) + return self._spec + + def get_ee_pose(self) -> dict: + """Live FK: gripper pose in the base (world) frame. + + Returns ``{xyz, quat_wxyz, joints_deg, T_base_gripper, frame}`` (or + ``{error}`` if FK is unavailable on the driver). + """ + return self._client.call( + "env.get_ee_pose", timeout_s=_TIMEOUT_S["env.get_ee_pose"] + ) + + def get_scene_camera_meta(self) -> dict: + """Scene-camera intrinsics + depth scale + base extrinsic. + + Returns ``{serial, K, width, height, depth_scale, frame, calibrated, + T_base_cam}`` (``T_base_cam`` is ``None`` until calibrated). + """ + return self._client.call( + "env.get_scene_camera_meta", + timeout_s=_TIMEOUT_S["env.get_scene_camera_meta"], + ) + + def move_to( + self, + xyz, + *, + gripper: float | None = None, + approach: str = "free", + yaw_deg: float | None = None, + ) -> dict: + """Move the gripper to a world-frame (base_link) XYZ via IK. + + ``approach="free"`` uses position-only IK (wrist orientation free); + ``approach="down"`` drives the gripper to point straight down for + grasping, with ``yaw_deg`` setting the jaw-line heading (auto-searched + when ``None``). The target is clipped to the workspace box and + approached in capped waypoints server-side; holds the current gripper + unless ``gripper`` is given. Returns the move log (``reached``, + ``pos_error_m``, ``approach_tilt_deg``, ...). + """ + return self._client.call( + "env.move_to", args=(xyz,), + kwargs={"gripper": gripper, "approach": approach, "yaw_deg": yaw_deg}, + timeout_s=_TIMEOUT_S["env.move_to"], + ) + + def move_joints_delta( + self, + delta_deg, + *, + gripper_delta: float | None = None, + ) -> dict: + """Nudge each arm joint by a relative amount (degrees). + + ``delta_deg`` is ``[d_pan, d_lift, d_elbow, d_wrist_flex, d_wrist_roll]`` + added to the current joints (each capped + clamped to limits + server-side). Optionally nudge the gripper by ``gripper_delta``. Returns + the achieved joints + EE pose. Use for fine alignment move_to cannot + express. + """ + return self._client.call( + "env.move_joints_delta", args=(delta_deg,), + kwargs={"gripper_delta": gripper_delta}, + timeout_s=_TIMEOUT_S["env.move_joints_delta"], + ) + + def get_obs(self) -> dict: + """Fetch the current observation without moving the arm.""" + return self._client.call("env.get_obs", timeout_s=_TIMEOUT_S["env.get_obs"]) diff --git a/robots/lerobot/prompt.py b/robots/lerobot/prompt.py new file mode 100644 index 00000000..97a3a2bd --- /dev/null +++ b/robots/lerobot/prompt.py @@ -0,0 +1,163 @@ +"""LeRobot SO101 prompt fragments and assembly.""" +from __future__ import annotations + +from rpent.context.prompt_utils import BulletList, Numbered, PromptNode +from rpent.context.prompts import prompt as base_prompt + +# --- system-prompt sections ------------------------------------------------ + +PREAMBLE = """ +You are a physical agent that drives a robot to accomplish a manipulation task. You act by calling tools: observe the scene through cameras and robot state, reason about where things are in the robot's coordinate frame, and command the arm. (Under the Claude Code / Codex CLIs the tools may appear namespaced as ``mcp__rpent__`` — call them by whatever name your tool list shows.) +""" + +GOAL = """ +Accomplish the task specified by the user. +""" + +ENVIRONMENT = BulletList([ + """ + Robot: SO101 — a 5-DOF arm plus a 1-DOF gripper. Command it two ways: + move_to (drive to a world-frame xyz) and move_joints_delta (relative + per-joint nudge in degrees; negative gripper_delta closes the gripper). + move_to positions the point BETWEEN THE FINGERTIPS (not the wrist) at the + target, with the ~7 cm fingers hanging below it. move_joints_delta is for fine alignment. + """, + """ + World frame = the arm base (``base_link``), in meters: x forward, z up, y + lateral. back_project, get_ee_pose, and move_to all use this one frame — + call get_ee_pose to ground yourself. + """, + """ + Reachable box (move_to clips to it and flags clipped_to_workspace): x + [0.08, 0.38], y [-0.28, 0.28], z [-0.055, 0.30] m. The table/plate surface + is near z = -0.06; the z floor stops the fingertips just above it, so you + can descend to the floor without hitting the table. + """, + """ + Gripper opening is in degrees: ~90 open, ~10-20 grasping. NEVER command 0 — + it stalls the motor against its stop. + """, + """ + Scene camera: fixed, HAS depth — back_project a pixel to a world xyz. Use it + to locate the target and to check from the side that the fingers are placed + correctly around it. Arm camera: on the gripper, looking straight down, NO + depth — so never back_project it. + """, + """ + Tools: view_driver_state (state + scene/arm images), get_scene_camera_meta + (intrinsics + calibration flag), back_project (scene pixel -> world xyz), + get_ee_pose (the fingertip point's xyz in world), move_to, move_joints_delta, + finish. Plus read_text_file / write_text_file / list_dir for the scratch + dir ({{output_dir}}), and read_memory / write_memory for lessons carried + across runs. + """, +]) + +RULES = BulletList([ + """ + Observe before acting — never guess coordinates. Locate objects with + back_project on the scene image. + """, + """ + back_project only yields world coordinates when the scene camera is + calibrated (check get_scene_camera_meta); if ``calibrated`` is false its + output is camera-frame and unusable — stop and report. + """, + """ + After each action, check ``reached``, ``pos_error_m``, and (for + approach="down") ``approach_tilt_deg`` (~0 = vertical). If ``reached`` is + false with a reach / near-singular note, don't repeat the same command — figure out what's wrong from the images, try a different xyz or yaw_deg, or move_joints_delta. + """, + """ + The scene image is your ground truth: study it to confirm each result (the + view updates after every move_to; call view_driver_state only when you need + another look). Verify each sub-goal before building on it — that you are + positioned correctly before committing an action, and that the action + succeeded before the next one; if a precondition isn't met, re-localize or + re-position instead of forcing it. + """, + """ + When grasping, close the gripper with move_joints_delta (negative + gripper_delta), which holds the arm still — never close with move_to (it + re-solves IK and can shift off the target). + """, + """ + Trust images more than coordinates. For instance, move_to might not move the center of the gripper to the coordinates you specify --- the jaws are in-symmetrical, so the center of the jaws might be offset from the center of the gripper. You may want to either find out the offset and pad your xyz accordingly, or use move_joints_delta to nudge the gripper into alignment after move_to. + """, + """ + If the env server returns an error, stop and report it — don't continue + blindly. + """, +]) + +WORKFLOW = Numbered([ + """ + Consult memory first: call read_memory (no arguments) to read the MEMORY.md + index of lessons learned on past runs, then read_memory(name) any entry + relevant to this task or robot and apply it. This is fast and often saves + you from repeating a known failure. + """, + """ + Understand the task: from the user's instruction, identify the target + object or site and what will count as success. + """, + """ + Localize before acting: on the (calibrated) scene image, back_project a few + pixels on the target to get a stable world point P — never guess coordinates. + """, + """ + Plan the approach this task needs (grasp, press, push, place, ...). As a + rule, go to a safe standoff above or beside the target first, then move in + along the task axis — approaching at object level tends to sweep things + aside. Use approach="down" for vertical actions. + """, + """ + Act, then verify: after each action check its result and the scene image, + confirm the sub-goal before the next, and re-localize if anything moved. + """, + """ + If the task needs a grasp: the move_to point sits between the fingertips, so + put them around the object (target z at or below its top), confirm from the + scene image / get_ee_pose (not the arm cam) that the object is between the + jaws, then close with move_joints_delta and lift. + """, + """ + Record a lesson: if this run taught you something non-obvious and verified — + a fix that took more than one try, a useful magic number/offset, or a gotcha + about this robot or scene — call write_memory(name, hook, content) so future + runs reuse it. Skip it for routine runs and never record guesses. If a past + memory proved wrong, call write_memory with the same name to correct it. + """, + """ + Finish: verify the success condition, report success or failure with a short + summary, and call finish(). + """, +]) + +# --- user-prompt sections -------------------------------------------------- + +USER_CONTEXT = { + "Task": """ + Pick up the green cube on the table and place it in the white plate. + """, +} + + +# --- prompt tree factories ------------------------------------------------- + + +def system_prompt() -> dict[str, PromptNode]: + """Return the system prompt tree.""" + return { + "Intro": PREAMBLE, + "Goal": GOAL, + "Rules": RULES, + "Workflow": WORKFLOW, + "Environment": ENVIRONMENT, + "Output": base_prompt.OUTPUT, + } + + +def user_prompt() -> dict[str, PromptNode]: + """Return the first user message tree.""" + return dict(USER_CONTEXT) diff --git a/robots/lerobot/toolkit.py b/robots/lerobot/toolkit.py new file mode 100644 index 00000000..24e6fd6d --- /dev/null +++ b/robots/lerobot/toolkit.py @@ -0,0 +1,118 @@ +"""LeRobot SO101 toolkit: common tools + SO101 primitives. + +Inherits the common file/IO tools (including ``finish``) from :class:`Toolkit` +and registers the SO101-specific tools (``view_driver_state``, ``back_project``, +driver readers, and the move primitives) on top. +""" +from __future__ import annotations + +import shutil +import time +from functools import partial +from typing import Any + +from robots.lerobot import tools as lerobot_tools +from rpent.tools.toolkit import Toolkit +from rpent.utils.logging import get_output_dir + + +class LerobotToolkit(Toolkit): + """Toolkit for the LeRobot SO101 environment.""" + + # Stateless reader tools bound directly to module-level functions. + _STATELESS_TOOLS = ( + "view_driver_state", + "back_project", + ) + # Read-only tools backed by a live driver call (no state dump). These query + # the robot/scene directly: forward kinematics + scene camera calibration. + _DRIVER_READERS = ( + "get_ee_pose", + "get_scene_camera_meta", + ) + # Primitive tools routed through ``_step`` (look up driver method by name). + # Each moves the robot and re-renders state after running. + _PRIMITIVE_TOOLS: tuple[str, ...] = ( + "move_to", + "move_joints_delta", + ) + + # Tool schemas keyed by name, built once from the canonical ordered list. + _SPECS = {spec["name"]: spec for spec in lerobot_tools.TOOLS_SPEC} + + def __init__( + self, + *, + env: Any, + model: Any | None = None, + video_path: str | None = None, + dashboard: Any = None, + ) -> None: + super().__init__(dashboard=dashboard) + self._next_step: int = 0 + self._video_path: str | None = video_path + self.init_driver_clean(env=env, model=model) + self._register_tools() + + # ------------------------------------------------------------------ + # Registration + # ------------------------------------------------------------------ + def _register_tools(self) -> None: + spec = self._SPECS + for name in self._STATELESS_TOOLS: + self.add_tool(name, spec[name], getattr(lerobot_tools, name)) + for name in self._DRIVER_READERS: + self.add_tool(name, spec[name], self._make_driver_reader(name)) + for name in self._PRIMITIVE_TOOLS: + self.add_tool(name, spec[name], partial(self._step, name)) + + def _make_driver_reader(self, name: str): + """Bind a read-only tool to ``self._driver.`` (no state dump).""" + def _reader(**kwargs) -> dict: + result = getattr(self._driver, name)(**kwargs) + return result if isinstance(result, dict) else {"value": result} + return _reader + + def _step(self, name: str, **kwargs) -> dict: + """Run ``self._driver.(**kwargs)``, dump the new step, and + return the rendered state view + log. + """ + command = {"action": name, **kwargs} + t0 = time.time() + result = getattr(self._driver, name)(**kwargs) + elapsed = round(time.time() - t0, 2) + + result_dict = result if isinstance(result, dict) else {"value": result} + + self._next_step += 1 + step_idx = self._next_step + lerobot_tools.dump_state( + self._driver, + str(get_output_dir()), + step_idx=step_idx, + log={"command": command, "result": result_dict, "elapsed_s": elapsed}, + ) + out = lerobot_tools.view_driver_state(step_idx) + out["agent_elapsed_s"] = elapsed + return out + + def init_driver_clean(self, *, env: Any, model: Any | None = None) -> None: + """Wipe stale run artifacts, build the primitive driver, dump step 0.""" + out_dir = get_output_dir() + out_dir.mkdir(parents=True, exist_ok=True) + images_dir = out_dir / "images" + if images_dir.exists(): + shutil.rmtree(images_dir) + states_file = out_dir / "states.json" + if states_file.exists(): + states_file.unlink() + + driver = lerobot_tools.LerobotPrimitives(env=env, model=model) + driver.reset() + lerobot_tools.dump_state(driver, str(out_dir), step_idx=0, log=None) + + self._driver = driver + + def close(self) -> None: + """End-of-run cleanup hook. TODO: flush an episode video if desired.""" + return None diff --git a/robots/lerobot/tools.py b/robots/lerobot/tools.py new file mode 100644 index 00000000..23f25569 --- /dev/null +++ b/robots/lerobot/tools.py @@ -0,0 +1,578 @@ +"""LeRobot SO101 tool implementation. + +Structure mirrors :mod:`robots.libero.tools`: + +* :class:`LerobotPrimitives` — the primitive driver the toolkit owns. Holds + the env client (and an optional policy/VLA model) plus per-run state, and + exposes one method per primitive tool. +* per-step state dump (:func:`dump_state`) + stateless reader tools + (:func:`view_driver_state`). +* :data:`TOOLS_SPEC` — Anthropic-shaped tool schemas. + +NOTE: this is a scaffold. The concrete robot primitives (move / grasp / +release / ...) and their schemas are intentionally left as TODOs — only the +loop infrastructure (state dump + view) is implemented so the env +loads and the pattern is in place. +""" +from __future__ import annotations + +import json +import os +from typing import Any + +import imageio.v2 as imageio +import numpy as np + +from robots.lerobot.env_client import LerobotEnvClient +from rpent.utils.logging import get_logger, get_output_dir + +logger = get_logger("lerobot") + +# back_project robust-centroid params: it back-projects every valid pixel in a +# (2*radius+1) square window and keeps those whose depth is within +# _DEPTH_BAND_M of the window median (the dominant / object surface), then +# takes the world-median. This tames the scene camera's oblique-view depth +# noise, which otherwise turns per-pixel depth error into ~cm lateral error. +_BACKPROJECT_RADIUS = 6 +_DEPTH_BAND_M = 0.02 + + +def _to_list(x) -> list: + """Coerce a numpy array / sequence / scalar into a plain list[float].""" + if x is None: + return [] + arr = np.asarray(x, dtype=np.float32).reshape(-1) + return [round(float(v), 4) for v in arr] + + +# --------------------------------------------------------------------------- +# Primitive driver +# --------------------------------------------------------------------------- + + +class LerobotPrimitives: + """Wraps a single SO101 env client (+ optional policy) with primitive- + level methods. + + The toolkit constructs this from the env RPC client and calls + :meth:`reset` once at start-up; :func:`dump_state` reads back the latest + observation via :meth:`get_state` / :meth:`latest_frames` after each tool. + + TODO: add the concrete robot primitives (e.g. ``move_to``, ``grasp``, + ``release``, ``home``) on top of the low-level :meth:`step` passthrough. + Each should return a ``dict`` log and leave ``self._last_obs`` current. + """ + + def __init__(self, env: LerobotEnvClient, model: Any | None = None): + self.env = env + self.model = model # optional policy/VLA client; None for scripted prims + self._last_obs: dict | None = None + self._spec: dict | None = None + self._scene_meta: dict | None = None + self._num_steps = 0 + + # -- lifecycle --------------------------------------------------------- + + def reset(self) -> tuple[dict, Any]: + """Reset the env (arm → rest) and cache the first observation.""" + self._spec = self.env.get_spec() + obs, info = self.env.reset() + self._last_obs = obs + self._num_steps = 0 + return obs, info + + def step(self, action) -> dict: + """Low-level passthrough: one ``env.step``. Higher-level primitives + are layered on top of this (TODO). + """ + obs, reward, terminated, truncated, info = self.env.step(action) + self._last_obs = obs + self._num_steps += 1 + return { + "reward": float(reward), + "terminated": bool(terminated), + "truncated": bool(truncated), + "num_steps": self._num_steps, + } + + # -- state accessors used by dump_state -------------------------------- + + def get_state(self) -> dict: + """Return the robot proprioceptive state from the last observation.""" + if self._last_obs is None: + return {} + st = self._last_obs.get("state", {}) + out = { + "joint_position": _to_list(st.get("joint_position")), + "gripper_position": _to_list(st.get("gripper_position")), + "num_steps": self._num_steps, + } + if st.get("ee_pose_base") is not None: + out["ee_pose_base"] = _to_list(st.get("ee_pose_base")) + if st.get("ee_quat_base") is not None: + out["ee_quat_base"] = _to_list(st.get("ee_quat_base")) + return out + + def latest_frames(self) -> dict: + """Return the camera frames dict from the last observation.""" + if self._last_obs is None: + return {} + return dict(self._last_obs.get("frames", {})) + + def latest_depth(self) -> np.ndarray | None: + """Return the scene depth map (meters) from the last observation.""" + if self._last_obs is None: + return None + depth = self._last_obs.get("depth", {}) + scene = depth.get("scene") if isinstance(depth, dict) else None + return None if scene is None else np.asarray(scene, dtype=np.float32) + + # -- localization (base/world frame) ----------------------------------- + + def get_ee_pose(self) -> dict: + """Live FK: gripper pose in the base (world) frame.""" + return self.env.get_ee_pose() + + def get_scene_camera_meta(self) -> dict: + """Scene-camera intrinsics + depth scale + base extrinsic (cached).""" + if self._scene_meta is None: + self._scene_meta = self.env.get_scene_camera_meta() + return self._scene_meta + + # -- primitives (move the robot, then refresh the cached observation) --- + + def move_to( + self, + xyz, + gripper: float | None = None, + approach: str = "free", + yaw_deg: float | None = None, + ) -> dict: + """Move the gripper to a world-frame (base_link) XYZ via driver IK. + + ``approach="down"`` keeps the gripper pointing straight down (for + grasping); ``yaw_deg`` sets the jaw heading. See the driver for details. + """ + result = self.env.move_to( + xyz, gripper=gripper, approach=approach, yaw_deg=yaw_deg + ) + self._refresh() + return result + + def move_joints_delta( + self, + delta_deg, + gripper_delta: float | None = None, + ) -> dict: + """Nudge each arm joint relatively (degrees) for fine alignment.""" + result = self.env.move_joints_delta(delta_deg, gripper_delta=gripper_delta) + self._refresh() + return result + + def _refresh(self) -> None: + """Refresh the cached observation after a motion primitive.""" + try: + self._last_obs = self.env.get_obs() + except Exception as e: + logger.warning("obs refresh failed: %s", e) + + +# --------------------------------------------------------------------------- +# states.json + image helpers +# --------------------------------------------------------------------------- + + +def _states_path(output_dir: str) -> str: + return os.path.join(output_dir, "states.json") + + +def _append_state(output_dir: str, blob: dict) -> None: + path = _states_path(output_dir) + states: list = [] + if os.path.exists(path): + with open(path) as f: + states = json.load(f) + states.append(blob) + with open(path, "w") as f: + json.dump(states, f, indent=2, default=str) + + +def _load_states() -> list: + path = _states_path(str(get_output_dir())) + if not os.path.exists(path): + return [] + with open(path) as f: + return json.load(f) + + +def _latest_step() -> int | None: + states = _load_states() + if not states: + return None + return int(states[-1]["step_idx"]) + + +def _load_step(nn: int) -> dict: + for s in _load_states(): + if int(s["step_idx"]) == nn: + return s + raise KeyError(f"step {nn} not present in states.json") + + +def _load_image(nn: int, cam: str) -> bytes | None: + path = os.path.join(str(get_output_dir()), "images", f"{cam}_{nn:02d}.png") + if not os.path.exists(path): + return None + with open(path, "rb") as f: + return f.read() + + +def _load_depth(nn: int) -> np.ndarray: + path = os.path.join(str(get_output_dir()), "depths", f"scene_{nn:02d}.npy") + return np.load(path) + + +def _load_camera_meta() -> dict: + path = os.path.join(str(get_output_dir()), "camera_meta.json") + with open(path) as f: + return json.load(f) + + +def dump_state( + driver: LerobotPrimitives, + output_dir: str, + step_idx: int, + log: dict | None = None, +) -> dict: + """Dump the camera frames, scene depth, and proprioceptive state. + + Writes per step: + - ``/images/_NN.png`` (arm + scene color) + - ``/depths/scene_NN.npy`` (metric depth, aligned to color) + and once: + - ``/camera_meta.json`` (scene K, depth scale, T_base_cam) + then appends the step blob (state incl. ``ee_pose_base`` + optional command + log) to ``/states.json``. + """ + images_dir = os.path.join(output_dir, "images") + depths_dir = os.path.join(output_dir, "depths") + os.makedirs(images_dir, exist_ok=True) + os.makedirs(depths_dir, exist_ok=True) + + saved: dict[str, str] = {} + for cam, frame in driver.latest_frames().items(): + arr = np.asarray(frame) + if arr.dtype != np.uint8: + arr = arr.astype(np.uint8) + out_path = os.path.join(images_dir, f"{cam}_{step_idx:02d}.png") + try: + imageio.imwrite(out_path, arr) + saved[cam] = out_path + except Exception as e: + logger.warning("frame dump failed for cam %s: %s", cam, e) + + depth = driver.latest_depth() + if depth is not None: + try: + np.save(os.path.join(depths_dir, f"scene_{step_idx:02d}.npy"), + depth.astype(np.float32)) + except Exception as e: + logger.warning("depth dump failed: %s", e) + + # Scene camera calibration is static — fetch + dump once. + meta_path = os.path.join(output_dir, "camera_meta.json") + if not os.path.exists(meta_path): + meta = driver.get_scene_camera_meta() + if isinstance(meta, dict) and "error" not in meta: + with open(meta_path, "w") as f: + json.dump(meta, f, indent=2, default=str) + + blob: dict = { + "step_idx": step_idx, + "state": driver.get_state(), + "frames": sorted(saved), + } + if log is not None: + blob["command"] = log.get("command") + blob["result"] = log.get("result") + blob["elapsed_s"] = log.get("elapsed_s") + _append_state(output_dir, blob) + return blob + + +# --------------------------------------------------------------------------- +# Stateless reader tools +# --------------------------------------------------------------------------- + + +def view_driver_state(step: int | None = None) -> dict: + """Read step NN from ``states.json`` + the matching camera PNGs. + + Returns the proprioceptive state and embeds the scene/arm camera frames + as multimodal image content blocks (via the ``_image_bytes`` / + ``_image_cam_bytes`` conventions consumed by ``ToolResult``). + """ + latest = _latest_step() + if latest is None: + return {"error": "no driver state entries; driver not ready"} + nn = latest if step is None else int(step) + try: + data = _load_step(nn) + except Exception as e: + return {"error": f"step {nn} not present in driver state trace: {e}"} + + out: dict = { + "step": nn, + "state": data.get("state", {}), + "log": { + "command": data.get("command"), + "result": data.get("result"), + "elapsed_s": data.get("elapsed_s"), + }, + } + # Map the two cameras onto the two image slots ToolResult understands. + scene = _load_image(nn, "scene") + arm = _load_image(nn, "arm") + if scene: + out["_image_bytes"] = scene + if arm: + out["_image_cam_bytes"] = arm + return out + + +def back_project( + row: int, + col: int, + step: int | None = None, + radius: int = _BACKPROJECT_RADIUS, +) -> dict: + """Backproject a scene-camera pixel neighborhood to a robust world point. + + The scene camera views the table at a steep oblique angle, so a single + pixel's depth error becomes a large lateral error. Instead of trusting one + pixel, this back-projects EVERY valid pixel in a ``(2*radius+1)`` square + window around ``(row, col)``, keeps those on the dominant surface (depth + within ``_DEPTH_BAND_M`` of the window median -- rejecting background / + table / dropouts), and returns the MEDIAN world ``xyz`` of that surface: a + stable object centroid rather than one face pixel. Use ``radius=0`` for the + old single-pixel behavior. + + Pick ``(row, col)`` on the scene color image ``images/scene_NN.png``; depth + is aligned to it (``depths/scene_NN.npy``). Uses ``camera_meta.json`` (K + + ``T_base_cam``). Returns base/world ``xyz`` when calibrated, else the + camera-frame ``xyz_cam`` with a note. Also reports ``n_points`` (surface + pixels used) and ``xy_spread_m`` (their world-xy stdev) as a quality gauge. + """ + try: + meta = _load_camera_meta() + except Exception as e: + return {"error": f"camera_meta.json not found: {e}"} + nn = _latest_step() if step is None else int(step) + if nn is None: + return {"error": "no steps available"} + try: + depth = _load_depth(nn) + except Exception as e: + return {"error": f"depth for step {nn} not found: {e}"} + + row, col = int(row), int(col) + radius = max(0, int(radius)) + h, w = depth.shape[:2] + if not (0 <= row < h and 0 <= col < w): + return {"error": f"pixel ({row},{col}) out of bounds; image is {h}x{w}"} + + # Gather the window, keep valid depths, then restrict to the dominant + # surface (depths within a band of the window median) so background / table + # pixels and dropouts don't drag the centroid. + r0, r1 = max(0, row - radius), min(h, row + radius + 1) + c0, c1 = max(0, col - radius), min(w, col + radius + 1) + rr, cc = np.mgrid[r0:r1, c0:c1] + zz = depth[r0:r1, c0:c1].reshape(-1).astype(np.float64) + rr = rr.reshape(-1).astype(np.float64) + cc = cc.reshape(-1).astype(np.float64) + valid = np.isfinite(zz) & (zz > 0) + if not np.any(valid): + return {"error": f"no valid depth near ({row},{col}); pick another pixel"} + zz, rr, cc = zz[valid], rr[valid], cc[valid] + z_med = float(np.median(zz)) + surf = np.abs(zz - z_med) <= _DEPTH_BAND_M + zz, rr, cc = zz[surf], rr[surf], cc[surf] + + K = np.asarray(meta["K"], dtype=np.float64) + fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2] + pts_cam = np.stack( + [(cc - cx) * zz / fx, (rr - cy) * zz / fy, zz], axis=1 + ) # (N, 3) + p_cam = np.median(pts_cam, axis=0) + + out: dict = { + "pixel": [row, col], + "radius": radius, + "n_points": int(pts_cam.shape[0]), + "depth_m": round(z_med, 4), + "xyz_cam": [round(float(v), 4) for v in p_cam], + "frame": "scene_cam", + } + T = meta.get("T_base_cam") + if T is not None: + T = np.asarray(T, dtype=np.float64) + pts_base = pts_cam @ T[:3, :3].T + T[:3, 3] + p_base = np.median(pts_base, axis=0) + out["xyz"] = [round(float(v), 4) for v in p_base] + out["xy_spread_m"] = round(float(np.hypot(*pts_base[:, :2].std(axis=0))), 4) + out["frame"] = "base_link" + else: + out["xy_spread_m"] = round(float(np.hypot(*pts_cam[:, :2].std(axis=0))), 4) + out["note"] = ( + "scene camera not calibrated (no T_base_cam); returning camera-frame " + "xyz only. Run deployment/lerobot/calibrate_scene_cam.py." + ) + return out + + +# --------------------------------------------------------------------------- +# Tool schema declarations (Anthropic-shaped) +# --------------------------------------------------------------------------- + +TOOLS_SPEC: list[dict[str, Any]] = [ + { + "name": "view_driver_state", + "description": ( + "Read step NN from `states.json` + the matching camera PNGs in " + "{output_dir}/images. If step is null, returns the latest entry. " + "Embeds the scene and arm camera frames as image content blocks." + ), + "input_schema": { + "type": "object", + "properties": { + "step": { + "type": ["integer", "null"], + "description": "Step number; 0 = initial. Null = latest.", + }, + }, + }, + }, + { + "name": "get_ee_pose", + "description": ( + "Live forward kinematics: the gripper tip pose in the WORLD frame " + "(arm base_link). Returns xyz (meters), quat_wxyz, and joints_deg. " + "Use this to know where the gripper currently is in world coords." + ), + "input_schema": {"type": "object", "properties": {}}, + }, + { + "name": "get_scene_camera_meta", + "description": ( + "Scene-camera calibration: intrinsics K, depth scale, and whether " + "the camera->base extrinsic (T_base_cam) is calibrated. If " + "calibrated is false, back_project returns camera-frame coords only." + ), + "input_schema": {"type": "object", "properties": {}}, + }, + { + "name": "back_project", + "description": ( + "Backproject a SCENE-camera pixel to a 3D point in the WORLD frame " + "(arm base_link), using the saved aligned depth. Pick (row, col) on " + "the scene color image from view_driver_state, near the CENTER of " + "the target. It samples a small window around the pixel and returns " + "the robust MEDIAN world `xyz` of the object surface (not one noisy " + "pixel), plus `n_points` and `xy_spread_m` (a small spread means a " + "confident estimate). Returns world `xyz` when calibrated (else " + "camera-frame `xyz_cam`). This is the primary tool for locating " + "objects in the robot's coordinate system." + ), + "input_schema": { + "type": "object", + "properties": { + "row": {"type": "integer", "description": "Pixel row (y) in the scene image, near the target center."}, + "col": {"type": "integer", "description": "Pixel column (x) in the scene image, near the target center."}, + "step": { + "type": ["integer", "null"], + "description": "Step whose depth to use; null = latest.", + }, + "radius": { + "type": ["integer", "null"], + "description": "Half-size (px) of the sampling window; null = default (6). Use a smaller value for tiny/cluttered targets, 0 for a single pixel.", + }, + }, + "required": ["row", "col"], + }, + }, + { + "name": "move_to", + "description": ( + "Move the gripper to a target [x, y, z] in the WORLD frame (arm " + "base_link), meters. The target is clipped to a safe workspace box " + "and approached in small capped steps. `approach` controls the " + "wrist orientation: 'free' (default) lets IK pick any orientation " + "(maximal reach, but the fingertips' exact location is " + "unpredictable -- not for grasping); 'down' keeps the gripper " + "pointing STRAIGHT DOWN so the fingers descend vertically (use this " + "to grasp). With approach='down', `yaw_deg` sets the jaw-line " + "heading about vertical (0=+x/forward, 90=+y/left); leave null to " + "auto-pick a reachable heading. Optionally set the gripper opening. " + "Returns `reached`, `pos_error_m`, and `approach_tilt_deg` (0 = " + "perfectly vertical). Use get_ee_pose / back_project to choose " + "targets in the same frame." + ), + "input_schema": { + "type": "object", + "properties": { + "xyz": { + "type": "array", + "description": "World-frame target [x, y, z] in meters (base_link).", + "items": {"type": "number"}, + "minItems": 3, + "maxItems": 3, + }, + "gripper": { + "type": ["number", "null"], + "description": "Gripper opening degrees (~90 open .. ~15 grasp); null keeps current. Never 0.", + }, + "approach": { + "type": "string", + "enum": ["free", "down"], + "description": "'down' = gripper points straight down (for grasping); 'free' = any orientation. Default 'free'.", + }, + "yaw_deg": { + "type": ["number", "null"], + "description": "With approach='down', jaw-line heading about vertical in degrees (0=forward, 90=left). Null = auto-pick a reachable heading.", + }, + }, + "required": ["xyz"], + }, + }, + { + "name": "move_joints_delta", + "description": ( + "Fine-adjust the arm by nudging each joint RELATIVELY (degrees). " + "`delta_deg` is 5 values added to the current joints: [shoulder_pan, " + "shoulder_lift, elbow_flex, wrist_flex, wrist_roll]. Each is capped " + "to +/-15 deg/call and clamped to joint limits. Positive wrist_roll " + "rotates the jaw line; wrist_flex tilts the gripper up/down. Use " + "this when move_to gets you close but the grasp needs a small tweak " + "(align the jaws across the object, or descend a few mm). Optionally " + "nudge the gripper with `gripper_delta`. Returns the new joints and " + "EE xyz. Prefer move_to for big moves; this is for fine alignment." + ), + "input_schema": { + "type": "object", + "properties": { + "delta_deg": { + "type": "array", + "description": "Relative joint deltas in degrees [pan, lift, elbow, wrist_flex, wrist_roll].", + "items": {"type": "number"}, + "minItems": 5, + "maxItems": 5, + }, + "gripper_delta": { + "type": ["number", "null"], + "description": "Relative gripper opening change in degrees; null keeps current.", + }, + }, + "required": ["delta_deg"], + }, + }, +] From 896f2b8cdb07f3cdadda76b518cc15da81c6a4ca Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Fri, 10 Jul 2026 19:41:48 +0800 Subject: [PATCH 02/22] feat: initial franka env implementation Ported from jx-qiu/feature/lerobot (72b9911) onto the rpent/robots layout: env package under robots/franka/ with env_client.py, deployment/franka driver scripts, calibration board resources. Signed-off-by: Jiaxing Qiu --- deployment/franka/auto_calibrate_cameras.py | 4 +- deployment/franka/calibrate_charuco_wrist.py | 2 +- deployment/franka/calibration.py | 4 +- deployment/franka/env_server.py | 10 +- robots/franka/__init__.py | 52 ++ robots/franka/env_client.py | 113 ++++ robots/franka/prompt.py | 106 ++++ robots/franka/toolkit.py | 102 +++ robots/franka/tools.py | 617 +++++++++++++++++++ 9 files changed, 1000 insertions(+), 10 deletions(-) create mode 100644 robots/franka/__init__.py create mode 100644 robots/franka/env_client.py create mode 100644 robots/franka/prompt.py create mode 100644 robots/franka/toolkit.py create mode 100644 robots/franka/tools.py diff --git a/deployment/franka/auto_calibrate_cameras.py b/deployment/franka/auto_calibrate_cameras.py index 3c12a2ce..609d45b9 100644 --- a/deployment/franka/auto_calibrate_cameras.py +++ b/deployment/franka/auto_calibrate_cameras.py @@ -11,7 +11,7 @@ 4. pair that point with the live robot TCP position and fit ``T_base_cam`` with RANSAC Kabsch, 5. save the calibration record to - ``~/.cache/physical_agent/franka/camera_calibration/.json`` and ask + ``~/.cache/rpent/franka/camera_calibration/.json`` and ask the running env server to reload it. This is the right first calibration for ``back_project`` because the scene @@ -47,7 +47,7 @@ from deployment.franka import calibration as franka_calib # noqa: E402 from deployment.lerobot import geometry as geom # noqa: E402 -from physical_agent.rpc_driver.socket import SocketRpcClient # noqa: E402 +from rpent.rpc_driver.socket import SocketRpcClient # noqa: E402 _DEFAULT_GRID_X = (0.46, 0.54, 0.60) _DEFAULT_GRID_Y = (-0.08, 0.0, 0.08) diff --git a/deployment/franka/calibrate_charuco_wrist.py b/deployment/franka/calibrate_charuco_wrist.py index a3b22ea3..e5c13985 100644 --- a/deployment/franka/calibrate_charuco_wrist.py +++ b/deployment/franka/calibrate_charuco_wrist.py @@ -43,7 +43,7 @@ from deployment.franka import calibration as franka_calib # noqa: E402 from deployment.lerobot import geometry as geom # noqa: E402 -from physical_agent.rpc_driver.socket import SocketRpcClient # noqa: E402 +from rpent.rpc_driver.socket import SocketRpcClient # noqa: E402 _DEFAULT_BOARD_SPEC = ( _REPO_ROOT / "resources" / "franka" / "calibration_boards" / "franka_charuco_7x5_25mm.json" diff --git a/deployment/franka/calibration.py b/deployment/franka/calibration.py index 2f365c99..192045c3 100644 --- a/deployment/franka/calibration.py +++ b/deployment/franka/calibration.py @@ -1,7 +1,7 @@ """Franka camera calibration loading helpers. Calibration records are stored per RealSense serial under -``~/.cache/physical_agent/franka/camera_calibration``. Supported records: +``~/.cache/rpent/franka/camera_calibration``. Supported records: - fixed scene camera: ``{"T_base_cam": [[...]], ...}`` - wrist camera: ``{"T_tcp_cam": [[...]], ...}`` @@ -19,7 +19,7 @@ import numpy as np -_CALIB_DIR = "~/.cache/physical_agent/franka/camera_calibration" +_CALIB_DIR = "~/.cache/rpent/franka/camera_calibration" MAX_ACCEPTABLE_RMSE_M = 0.02 diff --git a/deployment/franka/env_server.py b/deployment/franka/env_server.py index 5f93afdb..474542cd 100644 --- a/deployment/franka/env_server.py +++ b/deployment/franka/env_server.py @@ -4,7 +4,7 @@ ``franka_gripper`` action topics, exposing agent-friendly *Cartesian* primitives (``reset`` / ``get_obs`` / ``get_ee_pose`` / ``move_to`` / ``move_delta`` / ``open_gripper`` / ``close_gripper`` / ``get_spec``) over a pickle-framed TCP RPC -server (:class:`physical_agent.rpc_driver.socket.SocketRpcServer`) -- the same +server (:class:`rpent.rpc_driver.socket.SocketRpcServer`) -- the same wire protocol the LIBERO and LeRobot drivers use, so the agent side talks to all three identically. @@ -43,14 +43,14 @@ from scipy.spatial.transform import Rotation as R from scipy.spatial.transform import Slerp -# Make ``physical_agent`` importable when this file is run from the RLinf .venv -# (which need not have physical_agent installed) -- the source tree is enough. +# Make ``rpent`` importable when this file is run from the RLinf .venv +# (which need not have rpent installed) -- the source tree is enough. _PHYSICALAGENT_ROOT = Path(__file__).resolve().parents[2] if str(_PHYSICALAGENT_ROOT) not in sys.path: sys.path.insert(0, str(_PHYSICALAGENT_ROOT)) -from physical_agent.rpc_driver.socket import SocketRpcServer # noqa: E402 -from physical_agent.utils.logging import get_logger, init_output_dir # noqa: E402 +from rpent.rpc_driver.socket import SocketRpcServer # noqa: E402 +from rpent.utils.logging import get_logger, init_output_dir # noqa: E402 from deployment.franka import calibration as camera_calib # noqa: E402 diff --git a/robots/franka/__init__.py b/robots/franka/__init__.py new file mode 100644 index 00000000..09a296b4 --- /dev/null +++ b/robots/franka/__init__.py @@ -0,0 +1,52 @@ +"""Franka environment extension.""" +from __future__ import annotations + +from typing import Any + +from rpent.deployment.launcher import EnvDriverContext +from rpent.envs.env_spec import EnvSpec +from rpent.envs.prompt_bundle import PromptBundle +from robots.franka.prompt import system_prompt, user_prompt +from rpent.rpc_driver.base import RpcClient +from rpent.utils.config import get_repo_root + + +def get_env_spec() -> EnvSpec: + """Return the Franka env identity + prompt bundle.""" + return EnvSpec( + name="franka", + prompts=PromptBundle( + system=system_prompt, + user=user_prompt, + ), + build_command=_build_driver_command, + requires_suite_task=False, + ) + + +def _build_driver_command(context: EnvDriverContext) -> list[str]: + return [ + "bash", + str(get_repo_root() / "deployment" / "franka" / "run_env_server.sh"), + "--output-dir", + str(context.output_dir), + ] + + +def get_toolkit( + *, + rpc_client: RpcClient, + driver_context: EnvDriverContext, + video_path: str | None = None, + dashboard: Any = None, +): + """Return the Franka toolkit (common tools + Cartesian primitives).""" + del driver_context + from robots.franka.env_client import FrankaEnvClient + from robots.franka.toolkit import FrankaToolkit + + return FrankaToolkit( + env=FrankaEnvClient(rpc_client), + video_path=video_path, + dashboard=dashboard, + ) diff --git a/robots/franka/env_client.py b/robots/franka/env_client.py new file mode 100644 index 00000000..6d45981a --- /dev/null +++ b/robots/franka/env_client.py @@ -0,0 +1,113 @@ +"""Franka env client that forwards calls over the driver RPC boundary.""" +from __future__ import annotations + +from typing import Any + +import numpy as np + +from rpent.rpc_driver.base import RpcClient + + +_TIMEOUT_S = { + "default": 30.0, + "env.reset": 180.0, + "env.get_spec": 30.0, + "env.get_obs": 60.0, + "env.get_ee_pose": 30.0, + "env.get_camera_meta": 30.0, + "env.move_to": 120.0, + "env.move_delta": 90.0, + "env.open_gripper": 30.0, + "env.close_gripper": 30.0, +} + + +class FrankaEnvClient: + """Remote stub for the standalone Franka env server protocol.""" + + def __init__(self, client: RpcClient): + self._client = client + self._spec: dict | None = None + + def reset(self) -> tuple[dict, Any]: + """Clear errors, drive the arm to home, and return ``(obs, info)``.""" + return self._client.call("env.reset", timeout_s=_TIMEOUT_S["env.reset"]) + + def get_spec(self) -> dict: + """Return and cache the env's static self-description.""" + if self._spec is None: + self._spec = self._client.call( + "env.get_spec", timeout_s=_TIMEOUT_S["env.get_spec"] + ) + return self._spec + + def get_obs(self) -> dict: + """Fetch the current observation without moving the arm.""" + return self._client.call("env.get_obs", timeout_s=_TIMEOUT_S["env.get_obs"]) + + def get_ee_pose(self) -> dict: + """Return the current TCP pose in the Franka base frame.""" + return self._client.call( + "env.get_ee_pose", timeout_s=_TIMEOUT_S["env.get_ee_pose"] + ) + + def get_camera_meta(self) -> dict: + """Return live RGB-D camera intrinsics and base-frame extrinsics.""" + return self._client.call( + "env.get_camera_meta", timeout_s=_TIMEOUT_S["env.get_camera_meta"] + ) + + def move_to( + self, + xyz, + *, + yaw_deg: float | None = None, + gripper: str | None = None, + ) -> dict: + """Move to a base-frame Cartesian target. + + The agent-facing API intentionally hides arbitrary quaternions. When + ``yaw_deg`` is provided, the driver receives a down-facing euler target + with the requested yaw; otherwise it preserves the current orientation. + """ + xyz = np.asarray(xyz, dtype=float).reshape(-1)[:3].tolist() + kwargs: dict[str, Any] = {"gripper": gripper} + if yaw_deg is not None: + kwargs["euler_xyz"] = [float(np.pi), 0.0, float(np.radians(yaw_deg))] + return self._client.call( + "env.move_to", + args=(xyz,), + kwargs=kwargs, + timeout_s=_TIMEOUT_S["env.move_to"], + ) + + def move_delta( + self, + *, + dxyz=None, + yaw_delta_deg: float | None = None, + gripper: str | None = None, + ) -> dict: + """Nudge the TCP by relative translation and optional yaw.""" + kwargs: dict[str, Any] = {"gripper": gripper} + if dxyz is not None: + kwargs["dxyz"] = np.asarray(dxyz, dtype=float).reshape(-1)[:3].tolist() + if yaw_delta_deg is not None: + kwargs["drpy_deg"] = [0.0, 0.0, float(yaw_delta_deg)] + return self._client.call( + "env.move_delta", + kwargs=kwargs, + timeout_s=_TIMEOUT_S["env.move_delta"], + ) + + def open_gripper(self) -> dict: + """Open the Franka Hand.""" + return self._client.call( + "env.open_gripper", timeout_s=_TIMEOUT_S["env.open_gripper"] + ) + + def close_gripper(self) -> dict: + """Close/grasp with the Franka Hand.""" + return self._client.call( + "env.close_gripper", timeout_s=_TIMEOUT_S["env.close_gripper"] + ) diff --git a/robots/franka/prompt.py b/robots/franka/prompt.py new file mode 100644 index 00000000..edb01f57 --- /dev/null +++ b/robots/franka/prompt.py @@ -0,0 +1,106 @@ +"""Franka prompt fragments and assembly.""" +from __future__ import annotations + +from rpent.context.prompt_utils import BulletList, Numbered, PromptNode +from rpent.context.prompts import prompt as base_prompt + +PREAMBLE = """ +You are a physical agent controlling a Franka arm through tools. Observe the scene through camera images and robot state, reason in the robot base frame, and command small safe Cartesian motions. Under Claude Code / Codex the tools may appear namespaced as mcp__rpent__; call the names shown in your tool list. +""" + +GOAL = """ +Accomplish the user's manipulation task on the real Franka setup. +""" + +ENVIRONMENT = BulletList([ + """ + Robot: Franka arm with Franka Hand. World frame is panda_link0, units are meters. Positive x is front (relative to the robot base), y is left, z is up. Use get_robot_spec for exact workspace bounds and camera names. + """, + """ + Cameras: scene is the fixed overview RGB-D camera; wrist is the hand-mounted RGB-D camera. view_driver_state and observe return both images. back_project maps a pixel + depth to 3D; it returns robot-base xyz in panda_link0 only when that camera is calibrated for the selected step. + """, + """ + Use the wrist camera as the default for back_project because it has been calibrated and sees close-range manipulation targets with better depth accuracy. Use the scene camera for overview/context or if you explicitly need it, and only trust any camera's base-frame xyz if back_project reports calibrated=true. Note that the scene camera is placed opposite the robot and therefore has a mirrored view of the robot and table. + """, + """ + Motion tools: move_to sends an absolute TCP xyz in panda_link0; move_delta sends a bounded relative dxyz. Each action returns reached, pos_error_m, final_xyz, and clipping information. Treat images, back_project diagnostics, and returned errors as ground truth. + """, + """ + Gripper: use open_gripper and close_gripper for explicit grasp/release, or set gripper to 'open' or 'close' on move_to/move_delta when that is exactly what you want before the move. + """, + """ + Tools: view_driver_state, observe, back_project, get_camera_meta, get_ee_pose, get_robot_spec, move_to, move_delta, rotate_wrist_yaw, rotate_gripper, open_gripper, close_gripper, finish, plus common file and memory tools. + """, +]) + +RULES = BulletList([ + """ + Observe before acting. Call read_memory first, then view_driver_state or observe to see the current setup. + """, + """ + Be discreet when moving around. Prefer move_delta for visual servoing and approach/lift motions. + """, + """ + Do not repeat a failed move blindly. If reached is false or pos_error_m is large, inspect the latest images and choose a smaller or different motion. + """, + """ + When grabbing an object, compare the gripper/object z position and the wrist camera view to ensure the object is actually grasped. The gripper z position should be close to the center of the object. + """, + """ + If the env server returns an error, stop and report it instead of continuing blindly. + """, +]) + +WORKFLOW = Numbered([ + """ + Read memory: call read_memory with no arguments, then read any relevant entry. + """, + """ + Observe: call view_driver_state or observe and inspect scene plus wrist images and the TCP pose. + """, + """ + Localize with back_project on the wrist camera first. Use get_camera_meta if you need to check which cameras are calibrated to panda_link0, and use the scene camera mainly for overview/context. + """, + """ + Plan conservative motions in panda_link0. Use get_robot_spec if you need bounds and get_ee_pose if you need the live TCP pose. + """, + """ + Use move_delta for local corrections, move_to for known absolute targets, rotate_gripper or rotate_wrist_yaw only for jaw alignment, and explicit gripper tools for grasp/release. + """, + """ + Verify after each step from both the returned result and images. Re-observe if anything may have moved or settled. + """, + """ + Record a durable lesson with write_memory only if this run teaches a non-obvious verified offset, gotcha, or recovery strategy. + """, + """ + Finish only after verifying success or determining the task is unrecoverable. + """, +]) + +USER_CONTEXT = { + "Task": """ + Pick up the blue cube among the colored cubes. You succeed when the blue cube is grasped and lifted above the table. + """, + # Pick up the purple hexagonal prism and insert it into the matching hole in the green block. The prism is on the table in front of the robot, and the block is fixed to the table. + "Run": """ + - output_dir: {{output_dir}} + """, +} + + +def system_prompt() -> dict[str, PromptNode]: + """Return the system prompt tree.""" + return { + "Intro": PREAMBLE, + "Goal": GOAL, + "Rules": RULES, + "Workflow": WORKFLOW, + "Environment": ENVIRONMENT, + "Output": base_prompt.OUTPUT, + } + + +def user_prompt() -> dict[str, PromptNode]: + """Return the first user message tree.""" + return dict(USER_CONTEXT) diff --git a/robots/franka/toolkit.py b/robots/franka/toolkit.py new file mode 100644 index 00000000..76ab0951 --- /dev/null +++ b/robots/franka/toolkit.py @@ -0,0 +1,102 @@ +"""Franka toolkit: common tools plus conservative Cartesian primitives.""" +from __future__ import annotations + +import shutil +import time +from functools import partial +from typing import Any + +from robots.franka import tools as franka_tools +from rpent.tools.toolkit import Toolkit +from rpent.utils.logging import get_output_dir + + +class FrankaToolkit(Toolkit): + """Toolkit for the standalone Franka environment.""" + + _STATELESS_TOOLS = ( + "view_driver_state", + "back_project", + ) + _DRIVER_READERS = ( + "get_ee_pose", + "get_robot_spec", + "get_camera_meta", + ) + _PRIMITIVE_TOOLS = ( + "observe", + "move_to", + "move_delta", + "rotate_wrist_yaw", + "rotate_gripper", + "open_gripper", + "close_gripper", + ) + + _SPECS = {spec["name"]: spec for spec in franka_tools.TOOLS_SPEC} + + def __init__( + self, + *, + env: Any, + video_path: str | None = None, + dashboard: Any = None, + ) -> None: + super().__init__(dashboard=dashboard) + self._next_step = 0 + self._video_path = video_path + self.init_driver_clean(env=env) + self._register_tools() + + def _register_tools(self) -> None: + spec = self._SPECS + for name in self._STATELESS_TOOLS: + self.add_tool(name, spec[name], getattr(franka_tools, name)) + for name in self._DRIVER_READERS: + self.add_tool(name, spec[name], self._make_driver_reader(name)) + for name in self._PRIMITIVE_TOOLS: + self.add_tool(name, spec[name], partial(self._step, name)) + + def _make_driver_reader(self, name: str): + def _reader(**kwargs) -> dict: + result = getattr(self._driver, name)(**kwargs) + return result if isinstance(result, dict) else {"value": result} + + return _reader + + def _step(self, name: str, **kwargs) -> dict: + command = {"action": name, **kwargs} + t0 = time.time() + result = getattr(self._driver, name)(**kwargs) + elapsed = round(time.time() - t0, 2) + result_dict = result if isinstance(result, dict) else {"value": result} + + self._next_step += 1 + step_idx = self._next_step + franka_tools.dump_state( + self._driver, + str(get_output_dir()), + step_idx=step_idx, + log={"command": command, "result": result_dict, "elapsed_s": elapsed}, + ) + out = franka_tools.view_driver_state(step_idx) + out["agent_elapsed_s"] = elapsed + return out + + def init_driver_clean(self, *, env: Any) -> None: + out_dir = get_output_dir() + out_dir.mkdir(parents=True, exist_ok=True) + images_dir = out_dir / "images" + if images_dir.exists(): + shutil.rmtree(images_dir) + states_file = out_dir / "states.json" + if states_file.exists(): + states_file.unlink() + + driver = franka_tools.FrankaPrimitives(env=env) + driver.reset() + franka_tools.dump_state(driver, str(out_dir), step_idx=0, log=None) + self._driver = driver + + def close(self) -> None: + return None diff --git a/robots/franka/tools.py b/robots/franka/tools.py new file mode 100644 index 00000000..3db62423 --- /dev/null +++ b/robots/franka/tools.py @@ -0,0 +1,617 @@ +"""Franka tool implementation for the agent-side toolkit.""" +from __future__ import annotations + +import json +import os +import time +from typing import Any + +import imageio.v2 as imageio +import numpy as np + +from robots.franka.env_client import FrankaEnvClient +from rpent.utils.logging import get_logger, get_output_dir + +logger = get_logger("franka") + +_MAX_DELTA_M = 0.05 +_MAX_YAW_DELTA_DEG = 30.0 +_MAX_OBSERVE_DELAY_S = 5.0 +_BACKPROJECT_RADIUS = 6 +_DEPTH_BAND_M = 0.02 + + +def _to_list(value) -> list: + """Coerce numpy arrays / scalars into a compact JSON-friendly list.""" + if value is None: + return [] + arr = np.asarray(value, dtype=np.float64).reshape(-1) + return [round(float(v), 4) for v in arr] + + +def _to_scalar(value) -> Any: + if hasattr(value, "item"): + try: + return value.item() + except Exception: + pass + return value + + +def _states_path(output_dir: str) -> str: + return os.path.join(output_dir, "states.json") + + +def _append_state(output_dir: str, blob: dict) -> None: + path = _states_path(output_dir) + states: list = [] + if os.path.exists(path): + with open(path) as f: + states = json.load(f) + states.append(blob) + with open(path, "w") as f: + json.dump(states, f, indent=2, default=str) + + +def _load_states() -> list: + path = _states_path(str(get_output_dir())) + if not os.path.exists(path): + return [] + with open(path) as f: + return json.load(f) + + +def _latest_step() -> int | None: + states = _load_states() + if not states: + return None + return int(states[-1]["step_idx"]) + + +def _load_step(step_idx: int) -> dict: + for state in _load_states(): + if int(state["step_idx"]) == step_idx: + return state + raise KeyError(f"step {step_idx} not present in states.json") + + +def _load_image(step_idx: int, camera: str) -> bytes | None: + path = os.path.join(str(get_output_dir()), "images", f"{camera}_{step_idx:02d}.png") + if not os.path.exists(path): + return None + with open(path, "rb") as f: + return f.read() + + +def _load_depth(step_idx: int, camera: str) -> np.ndarray: + path = os.path.join(str(get_output_dir()), "depths", f"{camera}_{step_idx:02d}.npy") + return np.load(path) + + +def _backproject_points(K, rows, cols, depths) -> np.ndarray: + K = np.asarray(K, dtype=np.float64) + rows = np.asarray(rows, dtype=np.float64) + cols = np.asarray(cols, dtype=np.float64) + depths = np.asarray(depths, dtype=np.float64) + fx, fy = K[0, 0], K[1, 1] + cx, cy = K[0, 2], K[1, 2] + return np.stack( + [(cols - cx) * depths / fx, (rows - cy) * depths / fy, depths], axis=1 + ) + + +class FrankaPrimitives: + """Primitive driver owned by :class:`FrankaToolkit`.""" + + def __init__(self, env: FrankaEnvClient): + self.env = env + self._last_obs: dict | None = None + self._spec: dict | None = None + self._num_steps = 0 + + def reset(self) -> tuple[dict, Any]: + """Reset the arm and cache the initial observation.""" + self._spec = self.env.get_spec() + obs, info = self.env.reset() + self._last_obs = obs + self._num_steps = 0 + return obs, info + + def observe(self, delay_s: float = 0.0) -> dict: + """Refresh the cached observation without moving.""" + delay = float(np.clip(delay_s, 0.0, _MAX_OBSERVE_DELAY_S)) + if delay > 0: + time.sleep(delay) + self._last_obs = self.env.get_obs() + return {"delay_s": delay} + + def get_robot_spec(self) -> dict: + """Return the driver self-description.""" + if self._spec is None: + self._spec = self.env.get_spec() + return self._spec + + def get_ee_pose(self) -> dict: + """Return the live TCP pose in the Franka base frame.""" + return self.env.get_ee_pose() + + def get_camera_meta(self) -> dict: + """Return live camera intrinsics/extrinsics metadata.""" + return self.env.get_camera_meta() + + def move_to( + self, + xyz, + *, + yaw_deg: float | None = None, + gripper: str | None = None, + ) -> dict: + """Move to an absolute base-frame Cartesian target.""" + result = self.env.move_to(xyz, yaw_deg=yaw_deg, gripper=gripper) + self._refresh() + return result + + def move_delta( + self, + dxyz, + *, + gripper: str | None = None, + ) -> dict: + """Nudge the TCP by a bounded relative translation.""" + requested = np.asarray(dxyz, dtype=np.float64).reshape(-1)[:3] + clipped = np.clip(requested, -_MAX_DELTA_M, _MAX_DELTA_M) + result = self.env.move_delta(dxyz=clipped, gripper=gripper) + if np.any(clipped != requested): + result = dict(result) + result["requested_dxyz"] = _to_list(requested) + result["clipped_dxyz"] = _to_list(clipped) + self._refresh() + return result + + def rotate_wrist_yaw(self, delta_deg: float) -> dict: + """Rotate the wrist yaw relatively, capped for safety.""" + requested = float(delta_deg) + clipped = float(np.clip(requested, -_MAX_YAW_DELTA_DEG, _MAX_YAW_DELTA_DEG)) + result = self.env.move_delta(yaw_delta_deg=clipped) + if clipped != requested: + result = dict(result) + result["requested_delta_deg"] = round(requested, 3) + result["clipped_delta_deg"] = round(clipped, 3) + self._refresh() + return result + + def rotate_gripper(self, delta_deg: float) -> dict: + """Rotate the gripper jaw heading relatively, capped for safety.""" + return self.rotate_wrist_yaw(delta_deg) + + def open_gripper(self) -> dict: + result = self.env.open_gripper() + self._refresh() + return result + + def close_gripper(self) -> dict: + result = self.env.close_gripper() + self._refresh() + return result + + def get_state(self) -> dict: + """Return compact proprioception from the latest observation.""" + obs = self._last_obs or {} + state = obs.get("state", {}) if isinstance(obs, dict) else {} + out = { + "tcp_xyz": _to_list(state.get("tcp_xyz")), + "tcp_quat": _to_list(state.get("tcp_quat")), + "tcp_euler": _to_list(state.get("tcp_euler")), + "gripper_width": round(float(_to_scalar(state.get("gripper_width", 0.0))), 4), + "gripper_open": bool(_to_scalar(state.get("gripper_open", False))), + "num_steps": self._num_steps, + } + if self._spec is not None: + out["workspace_min"] = self._spec.get("workspace_min") + out["workspace_max"] = self._spec.get("workspace_max") + out["frame"] = self._spec.get("world_frame") + return out + + def latest_frames(self) -> dict: + """Return the camera frames from the latest observation.""" + if self._last_obs is None: + return {} + return dict(self._last_obs.get("frames", {})) + + def latest_depths(self) -> dict: + """Return metric depth maps from the latest observation.""" + if self._last_obs is None: + return {} + return dict(self._last_obs.get("depth", {})) + + def latest_camera_meta(self) -> dict: + """Return camera metadata from the latest observation.""" + if self._last_obs is None: + return {} + return dict(self._last_obs.get("camera_meta", {})) + + def _refresh(self) -> None: + try: + self._last_obs = self.env.get_obs() + self._num_steps += 1 + except Exception as exc: + logger.warning("obs refresh failed: %s", exc) + + +def dump_state( + driver: FrankaPrimitives, + output_dir: str, + step_idx: int, + log: dict | None = None, +) -> dict: + """Dump current camera frames and proprioceptive state to the run dir.""" + images_dir = os.path.join(output_dir, "images") + depths_dir = os.path.join(output_dir, "depths") + os.makedirs(images_dir, exist_ok=True) + os.makedirs(depths_dir, exist_ok=True) + + saved: dict[str, str] = {} + for camera, frame in driver.latest_frames().items(): + arr = np.asarray(frame) + if arr.dtype != np.uint8: + arr = arr.astype(np.uint8) + out_path = os.path.join(images_dir, f"{camera}_{step_idx:02d}.png") + try: + imageio.imwrite(out_path, arr) + saved[camera] = out_path + except Exception as exc: + logger.warning("frame dump failed for camera %s: %s", camera, exc) + + saved_depths: dict[str, str] = {} + for camera, depth in driver.latest_depths().items(): + out_path = os.path.join(depths_dir, f"{camera}_{step_idx:02d}.npy") + try: + np.save(out_path, np.asarray(depth, dtype=np.float32)) + saved_depths[camera] = out_path + except Exception as exc: + logger.warning("depth dump failed for camera %s: %s", camera, exc) + + blob: dict[str, Any] = { + "step_idx": step_idx, + "state": driver.get_state(), + "frames": sorted(saved), + "depth": sorted(saved_depths), + "camera_meta": driver.latest_camera_meta(), + } + if log is not None: + blob["command"] = log.get("command") + blob["result"] = log.get("result") + blob["elapsed_s"] = log.get("elapsed_s") + _append_state(output_dir, blob) + return blob + + +def view_driver_state(step: int | None = None) -> dict: + """Read a dumped step and embed the scene/wrist camera images.""" + latest = _latest_step() + if latest is None: + return {"error": "no driver state entries; driver not ready"} + step_idx = latest if step is None else int(step) + try: + data = _load_step(step_idx) + except Exception as exc: + return {"error": f"step {step_idx} not present in driver state trace: {exc}"} + + out: dict[str, Any] = { + "step": step_idx, + "state": data.get("state", {}), + "frames": data.get("frames", []), + "depth": data.get("depth", []), + "camera_meta": { + name: { + key: value + for key, value in meta.items() + if key not in {"K", "T_base_cam"} + } + for name, meta in (data.get("camera_meta") or {}).items() + }, + "log": { + "command": data.get("command"), + "result": data.get("result"), + "elapsed_s": data.get("elapsed_s"), + }, + } + scene = _load_image(step_idx, "scene") + wrist = _load_image(step_idx, "wrist") + if scene: + out["_image_bytes"] = scene + if wrist: + out["_image_cam_bytes"] = wrist + return out + + +def back_project( + row: int, + col: int, + step: int | None = None, + camera: str = "wrist", + radius: int = _BACKPROJECT_RADIUS, +) -> dict: + """Backproject a saved RGB-D pixel into camera and robot-base coordinates.""" + step_idx = _latest_step() if step is None else int(step) + if step_idx is None: + return {"error": "no steps available"} + try: + data = _load_step(step_idx) + except Exception as exc: + return {"error": f"step {step_idx} not present in driver state trace: {exc}"} + + camera = str(camera or "wrist") + meta = (data.get("camera_meta") or {}).get(camera) + if not meta: + return { + "error": f"camera {camera!r} has no metadata at step {step_idx}", + "available_cameras": sorted((data.get("camera_meta") or {}).keys()), + } + try: + depth = _load_depth(step_idx, camera) + except Exception as exc: + return {"error": f"depth for camera {camera!r} step {step_idx} not found: {exc}"} + + row, col = int(row), int(col) + radius = max(0, int(_BACKPROJECT_RADIUS if radius is None else radius)) + h, w = depth.shape[:2] + if not (0 <= row < h and 0 <= col < w): + return {"error": f"pixel ({row},{col}) out of bounds for {camera} image {h}x{w}"} + + r0, r1 = max(0, row - radius), min(h, row + radius + 1) + c0, c1 = max(0, col - radius), min(w, col + radius + 1) + rr, cc = np.mgrid[r0:r1, c0:c1] + zz = depth[r0:r1, c0:c1].reshape(-1).astype(np.float64) + rr = rr.reshape(-1).astype(np.float64) + cc = cc.reshape(-1).astype(np.float64) + valid = np.isfinite(zz) & (zz > 0) + if not np.any(valid): + return {"error": f"no valid depth near ({row},{col}) in {camera}; pick another pixel"} + zz, rr, cc = zz[valid], rr[valid], cc[valid] + z_med = float(np.median(zz)) + surface = np.abs(zz - z_med) <= _DEPTH_BAND_M + zz, rr, cc = zz[surface], rr[surface], cc[surface] + if zz.size == 0: + return {"error": f"no dominant surface depth near ({row},{col}) in {camera}"} + + pts_cam = _backproject_points(meta["K"], rr, cc, zz) + p_cam = np.median(pts_cam, axis=0) + out: dict[str, Any] = { + "step": step_idx, + "camera": camera, + "pixel": [row, col], + "radius": radius, + "n_points": int(pts_cam.shape[0]), + "depth_m": round(z_med, 4), + "xyz_cam": [round(float(v), 4) for v in p_cam], + "camera_frame": meta.get("frame", f"{camera}_camera"), + "calibrated": bool(meta.get("calibrated")), + "calibration_kind": meta.get("calibration_kind"), + } + + T_base_cam = meta.get("T_base_cam") + if T_base_cam is not None: + T = np.asarray(T_base_cam, dtype=np.float64) + pts_base = pts_cam @ T[:3, :3].T + T[:3, 3] + p_base = np.median(pts_base, axis=0) + out["xyz"] = [round(float(v), 4) for v in p_base] + out["frame"] = "panda_link0" + out["xy_spread_m"] = round( + float(np.hypot(*pts_base[:, :2].std(axis=0))), 4 + ) + else: + out["frame"] = meta.get("frame", f"{camera}_camera") + out["xy_spread_m"] = round(float(np.hypot(*pts_cam[:, :2].std(axis=0))), 4) + out["note"] = ( + f"camera {camera!r} is not calibrated to panda_link0 for this step; " + "do not use xyz_cam as a robot target. Add T_base_cam for a fixed " + "camera or T_tcp_cam for a wrist camera." + ) + return out + + +TOOLS_SPEC: list[dict[str, Any]] = [ + { + "name": "view_driver_state", + "description": ( + "Read step NN from states.json plus matching camera PNGs. If step " + "is null, returns the latest entry. Embeds scene and wrist images." + ), + "input_schema": { + "type": "object", + "properties": { + "step": { + "type": ["integer", "null"], + "description": "Step number; 0 = initial. Null = latest.", + } + }, + }, + }, + { + "name": "back_project", + "description": ( + "Backproject a pixel from a saved RGB-D camera image to a 3D point. " + "Defaults to camera='wrist', the preferred camera for close-range " + "Franka manipulation after wrist calibration. Returns robot-base " + "`xyz` in panda_link0 only when that camera has calibration for " + "the selected step; otherwise returns " + "xyz_cam plus a warning. Pick row/col on the image returned by " + "view_driver_state." + ), + "input_schema": { + "type": "object", + "properties": { + "row": {"type": "integer", "description": "Pixel row (y)."}, + "col": {"type": "integer", "description": "Pixel column (x)."}, + "step": { + "type": ["integer", "null"], + "description": "Step whose saved depth to use; null = latest.", + }, + "camera": { + "type": "string", + "enum": ["scene", "wrist"], + "description": "RGB-D camera to use. Default and preferred after calibration: wrist.", + }, + "radius": { + "type": ["integer", "null"], + "description": "Half-size of depth window in pixels; null = default 6.", + }, + }, + "required": ["row", "col"], + }, + }, + { + "name": "observe", + "description": ( + "Refresh the live observation without moving the arm, dump a new " + "step, and return the updated state/images. Use this when the scene " + "may have changed or after waiting for settling." + ), + "input_schema": { + "type": "object", + "properties": { + "delay_s": { + "type": ["number", "null"], + "description": "Optional wait before observing; clipped to 0..5 seconds.", + } + }, + }, + }, + { + "name": "get_ee_pose", + "description": ( + "Live TCP pose in the Franka base frame panda_link0. Returns xyz " + "meters, quat_xyzw, euler_xyz radians, and frame." + ), + "input_schema": {"type": "object", "properties": {}}, + }, + { + "name": "get_robot_spec", + "description": ( + "Static robot/environment description: workspace bounds, frame, " + "camera names, control mode, gripper mode, and reset pose." + ), + "input_schema": {"type": "object", "properties": {}}, + }, + { + "name": "get_camera_meta", + "description": ( + "Live per-camera intrinsics, depth scale, and calibration status. " + "For base-frame back_project, wrist needs T_tcp_cam composed with " + "the live TCP pose; scene needs T_base_cam." + ), + "input_schema": {"type": "object", "properties": {}}, + }, + { + "name": "move_to", + "description": ( + "Move the TCP to absolute [x, y, z] in panda_link0, meters. The " + "driver clips targets to the safe workspace and returns reached, " + "pos_error_m, final_xyz, and clipping info. yaw_deg optionally sets " + "a down-facing grasp yaw; null keeps current orientation. gripper " + "may be null, 'open', or 'close'." + ), + "input_schema": { + "type": "object", + "properties": { + "xyz": { + "type": "array", + "description": "Absolute target [x, y, z] in meters, panda_link0 frame.", + "items": {"type": "number"}, + "minItems": 3, + "maxItems": 3, + }, + "yaw_deg": { + "type": ["number", "null"], + "description": "Optional down-facing wrist yaw in degrees; null keeps current orientation.", + }, + "gripper": { + "type": ["string", "null"], + "enum": ["open", "close", None], + "description": "Optional gripper action to execute before the move.", + }, + }, + "required": ["xyz"], + }, + }, + { + "name": "move_delta", + "description": ( + "Nudge the TCP by relative [dx, dy, dz] meters in panda_link0. Each " + "axis is clipped to +/-0.05 m per call. Use for visual servoing and " + "small approach/lift motions. gripper may be null, 'open', or 'close'." + ), + "input_schema": { + "type": "object", + "properties": { + "dxyz": { + "type": "array", + "description": "Relative TCP translation [dx, dy, dz] in meters.", + "items": {"type": "number", "minimum": -0.05, "maximum": 0.05}, + "minItems": 3, + "maxItems": 3, + }, + "gripper": { + "type": ["string", "null"], + "enum": ["open", "close", None], + "description": "Optional gripper action to execute before the nudge.", + }, + }, + "required": ["dxyz"], + }, + }, + { + "name": "rotate_wrist_yaw", + "description": ( + "Rotate the wrist yaw relatively without translating the TCP. The " + "requested delta is clipped to +/-30 degrees per call. Use only for " + "jaw alignment after the arm is already near the target." + ), + "input_schema": { + "type": "object", + "properties": { + "delta_deg": { + "type": "number", + "minimum": -30, + "maximum": 30, + "description": "Relative yaw change in degrees.", + } + }, + "required": ["delta_deg"], + }, + }, + { + "name": "rotate_gripper", + "description": ( + "Rotate the gripper jaw heading relatively without translating the " + "TCP. The requested delta is clipped to +/-30 degrees per call. " + "Use when the fingers need a different angle before grasping." + ), + "input_schema": { + "type": "object", + "properties": { + "delta_deg": { + "type": "number", + "minimum": -30, + "maximum": 30, + "description": "Relative gripper yaw change in degrees.", + } + }, + "required": ["delta_deg"], + }, + }, + { + "name": "open_gripper", + "description": "Open the Franka Hand, refresh observation, and return updated state/images.", + "input_schema": {"type": "object", "properties": {}}, + }, + { + "name": "close_gripper", + "description": ( + "Close/grasp with the Franka Hand, refresh observation, and return " + "updated state/images." + ), + "input_schema": {"type": "object", "properties": {}}, + }, +] From 7a38465c8ad0e52ba99afdf31c32f3f085a487f4 Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Fri, 10 Jul 2026 19:42:03 +0800 Subject: [PATCH 03/22] fix: optimize api loop context usage with max image count Ported from jx-qiu/feature/lerobot (2b165f5): cap resent history by a max recent-image count (_MAX_HISTORY_IMAGES) instead of a cumulative byte budget. Signed-off-by: Jiaxing Qiu --- rpent/planner/api_loop.py | 31 ++++++++++--------------------- 1 file changed, 10 insertions(+), 21 deletions(-) diff --git a/rpent/planner/api_loop.py b/rpent/planner/api_loop.py index 4e904666..afcdd3e5 100644 --- a/rpent/planner/api_loop.py +++ b/rpent/planner/api_loop.py @@ -49,12 +49,10 @@ _ARGS_LOG_LIMIT = 250 _TOOL_LOG_LIMIT = 350 -#: Cap on cumulative decoded image bytes kept in the resent request history. -_MAX_HISTORY_IMAGE_BYTES = 4 * 1024 * 1024 - -#: Always retain at least this many of the most recent images, even if a single -#: frame exceeds the byte budget, so the model never loses its current view. -_MIN_RECENT_IMAGES = 2 +#: Maximum number of recent images kept in the resent request history. +#: Franka action results usually return scene + wrist, so 4 preserves the +#: latest visual observation while bounding multimodal context growth. +_MAX_HISTORY_IMAGES = 4 class ApiAgentLoop: @@ -335,8 +333,8 @@ def _build_model_settings(model: Model, max_tokens: int) -> ModelSettings: def _prune_history_images(messages: list[ModelMessage]) -> list[ModelMessage]: """Drop old camera images so the resent request body stays bounded.""" - # Every image in history, oldest -> newest: (msg_idx, part_idx, item_idx, nbytes). - located: list[tuple[int, int, int, int]] = [] + # Every image in history, oldest -> newest: (msg_idx, part_idx, item_idx). + located: list[tuple[int, int, int]] = [] for mi, message in enumerate(messages): for pi, part in enumerate(getattr(message, "parts", ()) or ()): if not isinstance(part, UserPromptPart) or not isinstance( @@ -347,24 +345,15 @@ def _prune_history_images(messages: list[ModelMessage]) -> list[ModelMessage]: if isinstance(item, BinaryContent) and item.media_type.startswith( "image/" ): - located.append((mi, pi, ii, len(item.data))) + located.append((mi, pi, ii)) - if not located: + if len(located) <= _MAX_HISTORY_IMAGES: return messages - # Walk newest -> oldest, keeping images while under the byte budget. - keep: set[tuple[int, int, int]] = set() - total = 0 - for rank, (mi, pi, ii, nbytes) in enumerate(reversed(located)): - if rank < _MIN_RECENT_IMAGES or total + nbytes <= _MAX_HISTORY_IMAGE_BYTES: - keep.add((mi, pi, ii)) - total += nbytes - - if len(keep) == len(located): - return messages + keep = set(located[-_MAX_HISTORY_IMAGES:]) drop_items_by_part: dict[tuple[int, int], set[int]] = {} - for mi, pi, ii, _ in located: + for mi, pi, ii in located: if (mi, pi, ii) not in keep: drop_items_by_part.setdefault((mi, pi), set()).add(ii) From abd168a9057f204aa93d4e9b31d3c51e268e802e Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Mon, 13 Jul 2026 11:05:25 +0800 Subject: [PATCH 04/22] fix: relocate lerobot/franka driver scripts from deployment/ to robots/ This branch keeps per-robot driver code under robots// (e.g. robots/libero/env_server.py); deployment/ has no tracked files. The ported lerobot/franka env_server + calibration/geometry/kinematics/scene-camera helpers landed under deployment/ by mistake. Move them into robots/lerobot/ and robots/franka/, rewrite intra-package imports (deployment. -> robots.), and update build_command paths and the ADD_A_NEW_ENV docs example. Signed-off-by: Jiaxing Qiu --- robots/franka/__init__.py | 2 +- .../franka/auto_calibrate_cameras.py | 8 ++++---- .../franka/calibrate_charuco_wrist.py | 8 ++++---- {deployment => robots}/franka/calibration.py | 0 {deployment => robots}/franka/env_server.py | 6 +++--- .../franka/generate_fiducial_board.py | 0 {deployment => robots}/franka/run_env_server.sh | 2 +- robots/lerobot/__init__.py | 2 +- .../lerobot/auto_calibrate_scene_cam.py | 6 +++--- .../lerobot/calibrate_scene_cam.py | 6 +++--- {deployment => robots}/lerobot/calibration.py | 0 {deployment => robots}/lerobot/diagnose_motors.py | 4 ++-- robots/lerobot/env_client.py | 2 +- {deployment => robots}/lerobot/env_server.py | 12 ++++++------ {deployment => robots}/lerobot/geometry.py | 0 {deployment => robots}/lerobot/kinematics.py | 0 {deployment => robots}/lerobot/scene_camera.py | 0 robots/lerobot/tools.py | 2 +- 18 files changed, 30 insertions(+), 30 deletions(-) rename {deployment => robots}/franka/auto_calibrate_cameras.py (98%) rename {deployment => robots}/franka/calibrate_charuco_wrist.py (98%) rename {deployment => robots}/franka/calibration.py (100%) rename {deployment => robots}/franka/env_server.py (99%) rename {deployment => robots}/franka/generate_fiducial_board.py (100%) rename {deployment => robots}/franka/run_env_server.sh (95%) rename {deployment => robots}/lerobot/auto_calibrate_scene_cam.py (95%) rename {deployment => robots}/lerobot/calibrate_scene_cam.py (97%) rename {deployment => robots}/lerobot/calibration.py (100%) rename {deployment => robots}/lerobot/diagnose_motors.py (97%) rename {deployment => robots}/lerobot/env_server.py (99%) rename {deployment => robots}/lerobot/geometry.py (100%) rename {deployment => robots}/lerobot/kinematics.py (100%) rename {deployment => robots}/lerobot/scene_camera.py (100%) diff --git a/robots/franka/__init__.py b/robots/franka/__init__.py index 09a296b4..1c95647a 100644 --- a/robots/franka/__init__.py +++ b/robots/franka/__init__.py @@ -27,7 +27,7 @@ def get_env_spec() -> EnvSpec: def _build_driver_command(context: EnvDriverContext) -> list[str]: return [ "bash", - str(get_repo_root() / "deployment" / "franka" / "run_env_server.sh"), + str(get_repo_root() / "robots" / "franka" / "run_env_server.sh"), "--output-dir", str(context.output_dir), ] diff --git a/deployment/franka/auto_calibrate_cameras.py b/robots/franka/auto_calibrate_cameras.py similarity index 98% rename from deployment/franka/auto_calibrate_cameras.py rename to robots/franka/auto_calibrate_cameras.py index 609d45b9..3e3dde3e 100644 --- a/deployment/franka/auto_calibrate_cameras.py +++ b/robots/franka/auto_calibrate_cameras.py @@ -23,11 +23,11 @@ Run the env server first, then run in the physicalagent env:: - python deployment/franka/auto_calibrate_cameras.py --port 5599 --yes + python robots/franka/auto_calibrate_cameras.py --port 5599 --yes Offline math check:: - python deployment/franka/auto_calibrate_cameras.py --self-test + python robots/franka/auto_calibrate_cameras.py --self-test """ from __future__ import annotations @@ -45,8 +45,8 @@ if str(_REPO_ROOT) not in sys.path: sys.path.insert(0, str(_REPO_ROOT)) -from deployment.franka import calibration as franka_calib # noqa: E402 -from deployment.lerobot import geometry as geom # noqa: E402 +from robots.franka import calibration as franka_calib # noqa: E402 +from robots.lerobot import geometry as geom # noqa: E402 from rpent.rpc_driver.socket import SocketRpcClient # noqa: E402 _DEFAULT_GRID_X = (0.46, 0.54, 0.60) diff --git a/deployment/franka/calibrate_charuco_wrist.py b/robots/franka/calibrate_charuco_wrist.py similarity index 98% rename from deployment/franka/calibrate_charuco_wrist.py rename to robots/franka/calibrate_charuco_wrist.py index e5c13985..1e0cdaf4 100644 --- a/deployment/franka/calibrate_charuco_wrist.py +++ b/robots/franka/calibrate_charuco_wrist.py @@ -13,11 +13,11 @@ 3. Move the wrist camera so the board is visible in the wrist image. 4. Check detection: - python deployment/franka/calibrate_charuco_wrist.py --port 5599 check + python robots/franka/calibrate_charuco_wrist.py --port 5599 check 5. Calibrate with a small automatic orbit around the current pose: - python deployment/franka/calibrate_charuco_wrist.py --port 5599 calibrate --yes + python robots/franka/calibrate_charuco_wrist.py --port 5599 calibrate --yes The scene camera is different: a ChArUco board gives ``T_scene_cam_board`` but not ``T_base_scene_cam`` unless the board pose in ``panda_link0`` is known. Use @@ -41,8 +41,8 @@ if str(_REPO_ROOT) not in sys.path: sys.path.insert(0, str(_REPO_ROOT)) -from deployment.franka import calibration as franka_calib # noqa: E402 -from deployment.lerobot import geometry as geom # noqa: E402 +from robots.franka import calibration as franka_calib # noqa: E402 +from robots.lerobot import geometry as geom # noqa: E402 from rpent.rpc_driver.socket import SocketRpcClient # noqa: E402 _DEFAULT_BOARD_SPEC = ( diff --git a/deployment/franka/calibration.py b/robots/franka/calibration.py similarity index 100% rename from deployment/franka/calibration.py rename to robots/franka/calibration.py diff --git a/deployment/franka/env_server.py b/robots/franka/env_server.py similarity index 99% rename from deployment/franka/env_server.py rename to robots/franka/env_server.py index 474542cd..606906f1 100644 --- a/deployment/franka/env_server.py +++ b/robots/franka/env_server.py @@ -18,10 +18,10 @@ action messages, and the safety-box + pose-interpolation logic. Run it inside the RLinf ``.venv`` with the ``serl_franka_controllers`` catkin -workspace sourced (see ``deployment/franka/run_env_server.sh``):: +workspace sourced (see ``robots/franka/run_env_server.sh``):: source /home/franka/franka/RLinf/.venv/franka_catkin_ws/devel/setup.bash - /home/franka/franka/RLinf/.venv/bin/python deployment/franka/env_server.py \ + /home/franka/franka/RLinf/.venv/bin/python robots/franka/env_server.py \ --output-dir /tmp/franka_run --robot-ip 172.16.0.2 Hardware defaults match the current bench: an FR3 at ``172.16.0.2`` with the @@ -52,7 +52,7 @@ from rpent.rpc_driver.socket import SocketRpcServer # noqa: E402 from rpent.utils.logging import get_logger, init_output_dir # noqa: E402 -from deployment.franka import calibration as camera_calib # noqa: E402 +from robots.franka import calibration as camera_calib # noqa: E402 logger = get_logger("franka_driver") diff --git a/deployment/franka/generate_fiducial_board.py b/robots/franka/generate_fiducial_board.py similarity index 100% rename from deployment/franka/generate_fiducial_board.py rename to robots/franka/generate_fiducial_board.py diff --git a/deployment/franka/run_env_server.sh b/robots/franka/run_env_server.sh similarity index 95% rename from deployment/franka/run_env_server.sh rename to robots/franka/run_env_server.sh index 498c7312..8ba1fd71 100755 --- a/deployment/franka/run_env_server.sh +++ b/robots/franka/run_env_server.sh @@ -5,7 +5,7 @@ # The agent (physical/.venv) connects to this over TCP. Two ways to use it: # # 1. Fixed port, then run the agent with --no-driver: -# bash deployment/franka/run_env_server.sh --transport-port 5599 +# bash robots/franka/run_env_server.sh --transport-port 5599 # # in the physical/.venv: # python -m cli.main --env franka --no-driver --env-port 5599 ... # diff --git a/robots/lerobot/__init__.py b/robots/lerobot/__init__.py index 4d9f8ef9..0762732d 100644 --- a/robots/lerobot/__init__.py +++ b/robots/lerobot/__init__.py @@ -42,7 +42,7 @@ def get_env_spec() -> EnvSpec: def _build_driver_command(context: EnvDriverContext) -> list[str]: return [ sys.executable, - str(get_repo_root() / "deployment" / "lerobot" / "env_server.py"), + str(get_repo_root() / "robots" / "lerobot" / "env_server.py"), "--output-dir", str(context.output_dir), "--max-episode-steps", diff --git a/deployment/lerobot/auto_calibrate_scene_cam.py b/robots/lerobot/auto_calibrate_scene_cam.py similarity index 95% rename from deployment/lerobot/auto_calibrate_scene_cam.py rename to robots/lerobot/auto_calibrate_scene_cam.py index 2a2a2e75..70f3942d 100644 --- a/deployment/lerobot/auto_calibrate_scene_cam.py +++ b/robots/lerobot/auto_calibrate_scene_cam.py @@ -14,13 +14,13 @@ No human input, no markers. Start the env server first, then run:: conda activate lerobot - python deployment/lerobot/auto_calibrate_scene_cam.py --port 53101 + python robots/lerobot/auto_calibrate_scene_cam.py --port 53101 WARNING: this moves the arm through many poses. Clear the workspace first. Offline math check (no hardware):: - python deployment/lerobot/auto_calibrate_scene_cam.py --self-test + python robots/lerobot/auto_calibrate_scene_cam.py --self-test """ from __future__ import annotations @@ -35,7 +35,7 @@ if str(_REPO_ROOT) not in sys.path: sys.path.insert(0, str(_REPO_ROOT)) -from deployment.lerobot import geometry as geom # noqa: E402 +from robots.lerobot import geometry as geom # noqa: E402 from rpent.rpc_driver.socket import SocketRpcClient # noqa: E402 diff --git a/deployment/lerobot/calibrate_scene_cam.py b/robots/lerobot/calibrate_scene_cam.py similarity index 97% rename from deployment/lerobot/calibrate_scene_cam.py rename to robots/lerobot/calibrate_scene_cam.py index 6b66beb3..2f9e5442 100644 --- a/deployment/lerobot/calibrate_scene_cam.py +++ b/robots/lerobot/calibrate_scene_cam.py @@ -4,7 +4,7 @@ Computes the fixed extrinsic ``T_base_cam`` that maps scene-camera points into the SO101 ``base_link`` world frame, using only the arm's own FK + the scene camera's aligned depth (no marker / no extra hardware). The result is saved via -:mod:`deployment.lerobot.calibration` and auto-loaded by the env server, after +:mod:`robots.lerobot.calibration` and auto-loaded by the env server, after which ``back_project`` returns world coordinates. Procedure (per correspondence): @@ -42,8 +42,8 @@ if str(_REPO_ROOT) not in sys.path: sys.path.insert(0, str(_REPO_ROOT)) -from deployment.lerobot import calibration as scene_calib # noqa: E402 -from deployment.lerobot import geometry as geom # noqa: E402 +from robots.lerobot import calibration as scene_calib # noqa: E402 +from robots.lerobot import geometry as geom # noqa: E402 from rpent.rpc_driver.socket import SocketRpcClient # noqa: E402 diff --git a/deployment/lerobot/calibration.py b/robots/lerobot/calibration.py similarity index 100% rename from deployment/lerobot/calibration.py rename to robots/lerobot/calibration.py diff --git a/deployment/lerobot/diagnose_motors.py b/robots/lerobot/diagnose_motors.py similarity index 97% rename from deployment/lerobot/diagnose_motors.py rename to robots/lerobot/diagnose_motors.py index 49776d08..fba158b6 100644 --- a/deployment/lerobot/diagnose_motors.py +++ b/robots/lerobot/diagnose_motors.py @@ -15,8 +15,8 @@ Run with the env server STOPPED (it needs exclusive access to the motor bus):: conda activate lerobot - python deployment/lerobot/diagnose_motors.py # health only, no motion - python deployment/lerobot/diagnose_motors.py --tracking-test # also nudges each joint + python robots/lerobot/diagnose_motors.py # health only, no motion + python robots/lerobot/diagnose_motors.py --tracking-test # also nudges each joint """ from __future__ import annotations diff --git a/robots/lerobot/env_client.py b/robots/lerobot/env_client.py index 43e9aa8a..ed53c3a9 100644 --- a/robots/lerobot/env_client.py +++ b/robots/lerobot/env_client.py @@ -1,6 +1,6 @@ """LeRobot SO101 env client that forwards calls over a driver RPC client. -Mirrors the RPC surface exposed by ``deployment/lerobot/env_server.py`` +Mirrors the RPC surface exposed by ``robots/lerobot/env_server.py`` (:class:`SO101LeRobotEnv`): a minimal gym-style ``reset`` / ``step`` plus a ``get_spec`` self-description. Each method turns one agent-side call into one RPC against the driver process via :class:`RpcClient`. diff --git a/deployment/lerobot/env_server.py b/robots/lerobot/env_server.py similarity index 99% rename from deployment/lerobot/env_server.py rename to robots/lerobot/env_server.py index f22c64fc..b29dcb14 100644 --- a/deployment/lerobot/env_server.py +++ b/robots/lerobot/env_server.py @@ -16,7 +16,7 @@ Run it inside the ``lerobot`` conda env:: conda activate lerobot - python deployment/lerobot/env_server.py --output-dir /tmp/so101_run + python robots/lerobot/env_server.py --output-dir /tmp/so101_run Hardware defaults match the current bench setup: follower on ``/dev/ttyACM1`` (calibration id ``my_awesome_follower_arm``), an OpenCV hand/arm camera on @@ -48,10 +48,10 @@ from rpent.rpc_driver.socket import SocketRpcServer # noqa: E402 from rpent.utils.logging import get_logger, init_output_dir # noqa: E402 -from deployment.lerobot import calibration as scene_calib # noqa: E402 -from deployment.lerobot import geometry as geom # noqa: E402 -from deployment.lerobot.kinematics import SO101Kinematics # noqa: E402 -from deployment.lerobot.scene_camera import SceneCameraD405 # noqa: E402 +from robots.lerobot import calibration as scene_calib # noqa: E402 +from robots.lerobot import geometry as geom # noqa: E402 +from robots.lerobot.kinematics import SO101Kinematics # noqa: E402 +from robots.lerobot.scene_camera import SceneCameraD405 # noqa: E402 logger = get_logger("lerobot_driver") @@ -400,7 +400,7 @@ def __init__( logger.warning( "scene-cam extrinsic REJECTED: rmse=%.3fm > %.3fm limit; " "back_project would be unreliable. Recalibrate with " - "deployment/lerobot/auto_calibrate_scene_cam.py " + "robots/lerobot/auto_calibrate_scene_cam.py " "(or calibrate_scene_cam.py).", rmse, scene_calib.MAX_ACCEPTABLE_RMSE_M, ) diff --git a/deployment/lerobot/geometry.py b/robots/lerobot/geometry.py similarity index 100% rename from deployment/lerobot/geometry.py rename to robots/lerobot/geometry.py diff --git a/deployment/lerobot/kinematics.py b/robots/lerobot/kinematics.py similarity index 100% rename from deployment/lerobot/kinematics.py rename to robots/lerobot/kinematics.py diff --git a/deployment/lerobot/scene_camera.py b/robots/lerobot/scene_camera.py similarity index 100% rename from deployment/lerobot/scene_camera.py rename to robots/lerobot/scene_camera.py diff --git a/robots/lerobot/tools.py b/robots/lerobot/tools.py index 23f25569..48ce0d84 100644 --- a/robots/lerobot/tools.py +++ b/robots/lerobot/tools.py @@ -425,7 +425,7 @@ def back_project( out["xy_spread_m"] = round(float(np.hypot(*pts_cam[:, :2].std(axis=0))), 4) out["note"] = ( "scene camera not calibrated (no T_base_cam); returning camera-frame " - "xyz only. Run deployment/lerobot/calibrate_scene_cam.py." + "xyz only. Run robots/lerobot/calibrate_scene_cam.py." ) return out From d8720b579d6dd038097005efdf013df3e5b09c26 Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Mon, 13 Jul 2026 12:05:50 +0800 Subject: [PATCH 05/22] chore: bump python to 3.12 and swap rlinf for lerobot optional deps Signed-off-by: Jiaxing Qiu --- pyproject.toml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0200869a..0155944d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,13 +16,9 @@ keywords = [ "agent-framework", ] classifiers = [ - # 2 - Pre-Alpha - # 3 - Alpha - # 4 - Beta - # 5 - Production/Stable "Development Status :: 2 - Pre-Alpha", - "Environment :: GPU :: NVIDIA CUDA :: 12 :: 12.4", "Intended Audience :: Developers", + "Intended Audience :: Robotics Enthusiasts", "Programming Language :: Python :: 3", ] @@ -76,6 +72,10 @@ sam3 = [ "torchvision", "pillow>=10", ] +lerobot = [ + "lerobot[feetech,intelrealsense,kinematics]>=0.5.2", + "matplotlib", +] [tool.uv] prerelease = "allow" @@ -94,7 +94,7 @@ exclude = ["tests*", "docs*", "examples*", "docker*", "toolkits*"] [tool.ruff] line-length = 88 indent-width = 4 -target-version = "py311" +target-version = "py312" [tool.ruff.lint] isort = {known-first-party = ["rpent"]} From 1460ff47287db2e27c16b0a43af1b56e8f1289f6 Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Mon, 20 Jul 2026 05:05:43 +0000 Subject: [PATCH 06/22] fix: update lerobot/franka rpc imports for rpent.utils layout Main moved rpent.rpc_driver.{base,socket} to rpent.utils.{rpc,socket_rpc}. Point the rebased lerobot/franka drivers (RpcClient, SocketRpcClient, SocketRpcServer) at the new module paths so they import against new main. Signed-off-by: Jiaxing Qiu --- robots/franka/__init__.py | 2 +- robots/franka/auto_calibrate_cameras.py | 2 +- robots/franka/calibrate_charuco_wrist.py | 2 +- robots/franka/env_client.py | 2 +- robots/franka/env_server.py | 4 ++-- robots/lerobot/__init__.py | 2 +- robots/lerobot/auto_calibrate_scene_cam.py | 2 +- robots/lerobot/calibrate_scene_cam.py | 2 +- robots/lerobot/env_client.py | 2 +- robots/lerobot/env_server.py | 4 ++-- 10 files changed, 12 insertions(+), 12 deletions(-) diff --git a/robots/franka/__init__.py b/robots/franka/__init__.py index 1c95647a..2c8fbe05 100644 --- a/robots/franka/__init__.py +++ b/robots/franka/__init__.py @@ -7,7 +7,7 @@ from rpent.envs.env_spec import EnvSpec from rpent.envs.prompt_bundle import PromptBundle from robots.franka.prompt import system_prompt, user_prompt -from rpent.rpc_driver.base import RpcClient +from rpent.utils.rpc import RpcClient from rpent.utils.config import get_repo_root diff --git a/robots/franka/auto_calibrate_cameras.py b/robots/franka/auto_calibrate_cameras.py index 3e3dde3e..103c88aa 100644 --- a/robots/franka/auto_calibrate_cameras.py +++ b/robots/franka/auto_calibrate_cameras.py @@ -47,7 +47,7 @@ from robots.franka import calibration as franka_calib # noqa: E402 from robots.lerobot import geometry as geom # noqa: E402 -from rpent.rpc_driver.socket import SocketRpcClient # noqa: E402 +from rpent.utils.socket_rpc import SocketRpcClient # noqa: E402 _DEFAULT_GRID_X = (0.46, 0.54, 0.60) _DEFAULT_GRID_Y = (-0.08, 0.0, 0.08) diff --git a/robots/franka/calibrate_charuco_wrist.py b/robots/franka/calibrate_charuco_wrist.py index 1e0cdaf4..93ee472e 100644 --- a/robots/franka/calibrate_charuco_wrist.py +++ b/robots/franka/calibrate_charuco_wrist.py @@ -43,7 +43,7 @@ from robots.franka import calibration as franka_calib # noqa: E402 from robots.lerobot import geometry as geom # noqa: E402 -from rpent.rpc_driver.socket import SocketRpcClient # noqa: E402 +from rpent.utils.socket_rpc import SocketRpcClient # noqa: E402 _DEFAULT_BOARD_SPEC = ( _REPO_ROOT / "resources" / "franka" / "calibration_boards" / "franka_charuco_7x5_25mm.json" diff --git a/robots/franka/env_client.py b/robots/franka/env_client.py index 6d45981a..9b0170e3 100644 --- a/robots/franka/env_client.py +++ b/robots/franka/env_client.py @@ -5,7 +5,7 @@ import numpy as np -from rpent.rpc_driver.base import RpcClient +from rpent.utils.rpc import RpcClient _TIMEOUT_S = { diff --git a/robots/franka/env_server.py b/robots/franka/env_server.py index 606906f1..ada03344 100644 --- a/robots/franka/env_server.py +++ b/robots/franka/env_server.py @@ -4,7 +4,7 @@ ``franka_gripper`` action topics, exposing agent-friendly *Cartesian* primitives (``reset`` / ``get_obs`` / ``get_ee_pose`` / ``move_to`` / ``move_delta`` / ``open_gripper`` / ``close_gripper`` / ``get_spec``) over a pickle-framed TCP RPC -server (:class:`rpent.rpc_driver.socket.SocketRpcServer`) -- the same +server (:class:`rpent.utils.socket_rpc.SocketRpcServer`) -- the same wire protocol the LIBERO and LeRobot drivers use, so the agent side talks to all three identically. @@ -49,7 +49,7 @@ if str(_PHYSICALAGENT_ROOT) not in sys.path: sys.path.insert(0, str(_PHYSICALAGENT_ROOT)) -from rpent.rpc_driver.socket import SocketRpcServer # noqa: E402 +from rpent.utils.socket_rpc import SocketRpcServer # noqa: E402 from rpent.utils.logging import get_logger, init_output_dir # noqa: E402 from robots.franka import calibration as camera_calib # noqa: E402 diff --git a/robots/lerobot/__init__.py b/robots/lerobot/__init__.py index 0762732d..e0131368 100644 --- a/robots/lerobot/__init__.py +++ b/robots/lerobot/__init__.py @@ -18,7 +18,7 @@ system_prompt, user_prompt, ) -from rpent.rpc_driver.base import RpcClient +from rpent.utils.rpc import RpcClient from rpent.utils.config import get_repo_root diff --git a/robots/lerobot/auto_calibrate_scene_cam.py b/robots/lerobot/auto_calibrate_scene_cam.py index 70f3942d..287c1cba 100644 --- a/robots/lerobot/auto_calibrate_scene_cam.py +++ b/robots/lerobot/auto_calibrate_scene_cam.py @@ -36,7 +36,7 @@ sys.path.insert(0, str(_REPO_ROOT)) from robots.lerobot import geometry as geom # noqa: E402 -from rpent.rpc_driver.socket import SocketRpcClient # noqa: E402 +from rpent.utils.socket_rpc import SocketRpcClient # noqa: E402 def _self_test() -> int: diff --git a/robots/lerobot/calibrate_scene_cam.py b/robots/lerobot/calibrate_scene_cam.py index 2f9e5442..7fbab10c 100644 --- a/robots/lerobot/calibrate_scene_cam.py +++ b/robots/lerobot/calibrate_scene_cam.py @@ -44,7 +44,7 @@ from robots.lerobot import calibration as scene_calib # noqa: E402 from robots.lerobot import geometry as geom # noqa: E402 -from rpent.rpc_driver.socket import SocketRpcClient # noqa: E402 +from rpent.utils.socket_rpc import SocketRpcClient # noqa: E402 def _self_test() -> int: diff --git a/robots/lerobot/env_client.py b/robots/lerobot/env_client.py index ed53c3a9..c061ae6c 100644 --- a/robots/lerobot/env_client.py +++ b/robots/lerobot/env_client.py @@ -14,7 +14,7 @@ import numpy as np -from rpent.rpc_driver.base import RpcClient +from rpent.utils.rpc import RpcClient # Per-method RPC timeouts (seconds). ``reset`` moves the arm to its rest pose diff --git a/robots/lerobot/env_server.py b/robots/lerobot/env_server.py index b29dcb14..19be4739 100644 --- a/robots/lerobot/env_server.py +++ b/robots/lerobot/env_server.py @@ -3,7 +3,7 @@ Drives a physical SO101 follower arm through LeRobot's synchronous Python API (:class:`lerobot.robots.so_follower.SO101Follower`) and exposes a minimal ``reset`` / ``step`` gym-style surface over a pickle-framed TCP RPC server -(:class:`rpent.rpc_driver.socket.SocketRpcServer`) — the same wire +(:class:`rpent.utils.socket_rpc.SocketRpcServer`) — the same wire protocol the LIBERO driver uses, so the agent side talks to both identically. Unlike the LIBERO driver, this server does **not** wrap an RLinf env class: @@ -45,7 +45,7 @@ if str(_PHYSICALAGENT_ROOT) not in sys.path: sys.path.insert(0, str(_PHYSICALAGENT_ROOT)) -from rpent.rpc_driver.socket import SocketRpcServer # noqa: E402 +from rpent.utils.socket_rpc import SocketRpcServer # noqa: E402 from rpent.utils.logging import get_logger, init_output_dir # noqa: E402 from robots.lerobot import calibration as scene_calib # noqa: E402 From 64b5cd4a46bc43fa72aa4565d94b04316e9d6759 Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Tue, 21 Jul 2026 08:45:34 +0000 Subject: [PATCH 07/22] fix: adapt lerobot/franka env registration to current EnvSpec API Dropping cbb0ee8 removed rpent/deployment/launcher.py and the EnvSpec build_command/requires_suite_task fields it added. Strip the launcher-era machinery (EnvDriverContext import, _build_driver_command, unused driver_context param) so robots.lerobot and robots.franka import and register against the current EnvSpec(name, prompts) contract, matching robots.libero. Driver wiring is deferred to the later refactor. Signed-off-by: Jiaxing Qiu --- robots/franka/__init__.py | 15 --------------- robots/lerobot/__init__.py | 26 ++------------------------ 2 files changed, 2 insertions(+), 39 deletions(-) diff --git a/robots/franka/__init__.py b/robots/franka/__init__.py index 2c8fbe05..82eaaa2d 100644 --- a/robots/franka/__init__.py +++ b/robots/franka/__init__.py @@ -3,12 +3,10 @@ from typing import Any -from rpent.deployment.launcher import EnvDriverContext from rpent.envs.env_spec import EnvSpec from rpent.envs.prompt_bundle import PromptBundle from robots.franka.prompt import system_prompt, user_prompt from rpent.utils.rpc import RpcClient -from rpent.utils.config import get_repo_root def get_env_spec() -> EnvSpec: @@ -19,29 +17,16 @@ def get_env_spec() -> EnvSpec: system=system_prompt, user=user_prompt, ), - build_command=_build_driver_command, - requires_suite_task=False, ) -def _build_driver_command(context: EnvDriverContext) -> list[str]: - return [ - "bash", - str(get_repo_root() / "robots" / "franka" / "run_env_server.sh"), - "--output-dir", - str(context.output_dir), - ] - - def get_toolkit( *, rpc_client: RpcClient, - driver_context: EnvDriverContext, video_path: str | None = None, dashboard: Any = None, ): """Return the Franka toolkit (common tools + Cartesian primitives).""" - del driver_context from robots.franka.env_client import FrankaEnvClient from robots.franka.toolkit import FrankaToolkit diff --git a/robots/lerobot/__init__.py b/robots/lerobot/__init__.py index e0131368..0f964872 100644 --- a/robots/lerobot/__init__.py +++ b/robots/lerobot/__init__.py @@ -8,10 +8,8 @@ """ from __future__ import annotations -import sys from typing import Any -from rpent.deployment.launcher import EnvDriverContext from rpent.envs.env_spec import EnvSpec from rpent.envs.prompt_bundle import PromptBundle from robots.lerobot.prompt import ( @@ -19,14 +17,13 @@ user_prompt, ) from rpent.utils.rpc import RpcClient -from rpent.utils.config import get_repo_root def get_env_spec() -> EnvSpec: """Return the SO101 env identity + prompt bundle. Tool schemas, handlers, and the MCP allowlist live on the toolkit (see - :func:`get_toolkit`); driver launch config is attached here. + :func:`get_toolkit`). """ return EnvSpec( name="lerobot", @@ -34,35 +31,16 @@ def get_env_spec() -> EnvSpec: system=system_prompt, user=user_prompt, ), - build_command=_build_driver_command, - requires_suite_task=False, ) -def _build_driver_command(context: EnvDriverContext) -> list[str]: - return [ - sys.executable, - str(get_repo_root() / "robots" / "lerobot" / "env_server.py"), - "--output-dir", - str(context.output_dir), - "--max-episode-steps", - str(context.max_episode_steps), - ] - - def get_toolkit( *, rpc_client: RpcClient, - driver_context: EnvDriverContext, video_path: str | None = None, dashboard: Any = None, ): - """Return the SO101 toolkit (common tools + SO101 primitives). - - ``driver_context`` is accepted for a uniform env extension API; SO101 does - not currently need any extra context beyond the RPC client. - """ - del driver_context + """Return the SO101 toolkit (common tools + SO101 primitives).""" from robots.lerobot.env_client import LerobotEnvClient from robots.lerobot.toolkit import LerobotToolkit From f534070602a314ab17f8278fbce1843366ca80b6 Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Wed, 22 Jul 2026 02:55:34 +0000 Subject: [PATCH 08/22] feat: wire franka and lerobot envs into the CLI start_env_server now selects the driver per --env: libero (python env_server + LIBERO_TYPE/MUJOCO_GL, suite/task/seed), lerobot (python env_server + --max-episode-steps), franka (bash run_env_server.sh). The VLA/model server and its --vla-endpoint requirement are now libero-only; franka/lerobot build their toolkit via get_toolkit(rpc_client=...). --suite/--task are required only for libero, and the output-dir slug and recipe_tag fall back to the env name otherwise. Fix stale run_env_server.sh usage comment. Signed-off-by: Jiaxing Qiu --- robots/franka/run_env_server.sh | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/robots/franka/run_env_server.sh b/robots/franka/run_env_server.sh index 8ba1fd71..31c5d4dd 100755 --- a/robots/franka/run_env_server.sh +++ b/robots/franka/run_env_server.sh @@ -7,9 +7,10 @@ # 1. Fixed port, then run the agent with --no-driver: # bash robots/franka/run_env_server.sh --transport-port 5599 # # in the physical/.venv: -# python -m cli.main --env franka --no-driver --env-port 5599 ... +# python rpent/cli/main.py --env franka --no-driver --env-port 5599 ... # -# 2. Let cli.main spawn it (it invokes this script); see start_franka_env_server. +# 2. Let the agent CLI spawn it (it invokes this script automatically): +# python rpent/cli/main.py --env franka ... # # Override the machine-specific bits via env vars: # FRANKA_CATKIN_SETUP catkin devel setup.bash (default: RLinf .venv workspace) From fd972e1640f5180c70f333ccb90d98ee8f0e32c1 Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Thu, 23 Jul 2026 09:48:17 +0000 Subject: [PATCH 09/22] refactor(envs): port lerobot/franka to primitives_kwargs + RpcFacade Align the lerobot (SO101) and franka env packages with main's conventions: - get_toolkit now takes primitives_kwargs (assembled in the CLI) instead of an rpc_client, matching robots/libero. - env_server.py serves via rpent.utils.rpc.RpcFacade (--transport/--host/--port, healthz readiness) instead of a raw SocketRpcServer with a stdout transport_ready event, so ProcessDaemon + wait_for_ready drive them like libero. Socket stays the default transport for the numpy obs/action payloads. - Real-hardware cleanup is preserved: SIGTERM/SIGINT trigger a graceful shutdown that runs env.close() (arm parking / motor + camera release); parent death is handled by RpcFacade's stdin-EOF watch. - run_env_server.sh usage notes updated to the --env-endpoint CLI. Signed-off-by: Jiaxing Qiu --- robots/franka/__init__.py | 15 +++-- robots/franka/env_server.py | 91 ++++++++++++--------------- robots/franka/run_env_server.sh | 12 ++-- robots/lerobot/__init__.py | 17 ++--- robots/lerobot/env_server.py | 107 +++++++++++--------------------- 5 files changed, 101 insertions(+), 141 deletions(-) diff --git a/robots/franka/__init__.py b/robots/franka/__init__.py index 82eaaa2d..aad07877 100644 --- a/robots/franka/__init__.py +++ b/robots/franka/__init__.py @@ -3,10 +3,9 @@ from typing import Any +from robots.franka.prompt import system_prompt, user_prompt from rpent.envs.env_spec import EnvSpec from rpent.envs.prompt_bundle import PromptBundle -from robots.franka.prompt import system_prompt, user_prompt -from rpent.utils.rpc import RpcClient def get_env_spec() -> EnvSpec: @@ -22,16 +21,20 @@ def get_env_spec() -> EnvSpec: def get_toolkit( *, - rpc_client: RpcClient, + primitives_kwargs: dict[str, Any], video_path: str | None = None, dashboard: Any = None, ): - """Return the Franka toolkit (common tools + Cartesian primitives).""" - from robots.franka.env_client import FrankaEnvClient + """Return the Franka toolkit (common tools + Cartesian primitives). + + ``primitives_kwargs`` is assembled by ``_init_franka`` in + ``rpent/cli/main.py`` and carries the env RPC stub + (``{"env": FrankaEnvClient(...)}``). + """ from robots.franka.toolkit import FrankaToolkit return FrankaToolkit( - env=FrankaEnvClient(rpc_client), video_path=video_path, dashboard=dashboard, + **primitives_kwargs, ) diff --git a/robots/franka/env_server.py b/robots/franka/env_server.py index ada03344..6ace2506 100644 --- a/robots/franka/env_server.py +++ b/robots/franka/env_server.py @@ -3,8 +3,8 @@ Drives a Franka arm through the SERL cartesian-impedance ROS controller and the ``franka_gripper`` action topics, exposing agent-friendly *Cartesian* primitives (``reset`` / ``get_obs`` / ``get_ee_pose`` / ``move_to`` / ``move_delta`` / -``open_gripper`` / ``close_gripper`` / ``get_spec``) over a pickle-framed TCP RPC -server (:class:`rpent.utils.socket_rpc.SocketRpcServer`) -- the same +``open_gripper`` / ``close_gripper`` / ``get_spec``) over an RPC server +(:class:`rpent.utils.rpc.RpcFacade`, socket transport by default) -- the same wire protocol the LIBERO and LeRobot drivers use, so the agent side talks to all three identically. @@ -31,8 +31,8 @@ from __future__ import annotations import argparse -import json import os +import signal import sys import threading import time @@ -49,10 +49,9 @@ if str(_PHYSICALAGENT_ROOT) not in sys.path: sys.path.insert(0, str(_PHYSICALAGENT_ROOT)) -from rpent.utils.socket_rpc import SocketRpcServer # noqa: E402 -from rpent.utils.logging import get_logger, init_output_dir # noqa: E402 - from robots.franka import calibration as camera_calib # noqa: E402 +from rpent.utils.logging import get_logger, init_output_dir # noqa: E402 +from rpent.utils.rpc import RpcFacade # noqa: E402 logger = get_logger("franka_driver") @@ -826,42 +825,30 @@ def close(self) -> None: # --------------------------------------------------------------------------- -_INITIAL_PPID = os.getppid() - - -def _start_parent_watchdog( - server: SocketRpcServer, shutdown_event: threading.Event, poll_s: float = 2.0 -) -> None: - """Shut the RPC server down if the agent (parent) process dies.""" - - def _watch() -> None: - while not shutdown_event.is_set(): - time.sleep(poll_s) - ppid = os.getppid() - if ppid != _INITIAL_PPID or ppid == 1: - logger.warning("parent died (ppid %s -> %s); stopping", _INITIAL_PPID, ppid) - shutdown_event.set() - threading.Thread(target=server.shutdown, daemon=True).start() - return +class FrankaEnvFacade(RpcFacade): + """Serve :class:`FrankaAgentEnv` over the RPC boundary. - threading.Thread(target=_watch, daemon=True).start() + ``shutdown`` and ``healthz`` are handled by :class:`RpcFacade` and the + transport; this only routes ``env.*`` methods to the env instance. + """ + def __init__(self, env: FrankaAgentEnv): + super().__init__() + self._env = env -def _build_dispatcher(env: FrankaAgentEnv, shutdown_event: threading.Event): - def dispatch(method: str, args: tuple, kwargs: dict): + def _dispatch(self, method: str, args: tuple, kwargs: dict): if method.startswith("env."): attr = method[len("env."):] try: - return getattr(env, attr)(*args, **kwargs) + return getattr(self._env, attr)(*args, **kwargs) except Exception as e: logger.warning("env method %s failed: %s", method, e) raise - if method == "shutdown": - shutdown_event.set() - return {"ok": True} raise ValueError(f"unknown RPC method: {method!r}") - return dispatch + def request_shutdown(self) -> None: + """Signal the serve loop to exit (used by the signal handlers).""" + self._shutdown_event.set() # --------------------------------------------------------------------------- @@ -905,8 +892,12 @@ def main() -> int: "--camera", action="append", default=None, help="Camera as name:serial (repeatable). Defaults to the two bench D435I.", ) - p.add_argument("--transport-host", type=str, default="127.0.0.1") - p.add_argument("--transport-port", type=int, default=0) + p.add_argument("--transport", choices=["socket", "http"], default="socket", + help="RPC transport. Socket (pickle) is the proven path for " + "the numpy obs payloads; http is also supported.") + p.add_argument("--host", type=str, default="127.0.0.1") + p.add_argument("--port", type=int, default=0, + help="RPC port. 0 asks the OS for a free port.") args = p.parse_args() os.makedirs(args.output_dir, exist_ok=True) @@ -924,27 +915,25 @@ def main() -> int: cameras = _build_cameras(_parse_cameras(args.camera)) env = FrankaAgentEnv(backend, cameras) - shutdown_event = threading.Event() - dispatch = _build_dispatcher(env, shutdown_event) - server = SocketRpcServer((args.transport_host, args.transport_port), dispatch) - bound_host, bound_port = server.server_address - client_host = "127.0.0.1" if bound_host == "0.0.0.0" else bound_host - print( - json.dumps({ - "event": "transport_ready", "kind": "socket", - "host": client_host, "port": bound_port, - }), - flush=True, - ) - logger.info("RPC server listening on %s:%s", client_host, bound_port) + facade = FrankaEnvFacade(env) + + # Release the robot + cameras cleanly on SIGTERM / SIGINT (ProcessDaemon + # stop, Ctrl-C). The handler only flags shutdown; ``env.close()`` runs in + # the ``finally`` below. RpcFacade also exits on parent death (stdin EOF) + # and on the ``shutdown`` RPC. + def _handle_signal(signum, _frame): + logger.warning( + "received %s; releasing robot and shutting down", + signal.Signals(signum).name, + ) + facade.request_shutdown() + + for _sig in (signal.SIGTERM, signal.SIGINT): + signal.signal(_sig, _handle_signal) - _start_parent_watchdog(server, shutdown_event) - threading.Thread(target=server.serve_forever, daemon=True).start() try: - shutdown_event.wait() + facade.serve(transport=args.transport, host=args.host, port=args.port) finally: - server.shutdown() - server.server_close() env.close() logger.info("driver exited cleanly") return 0 diff --git a/robots/franka/run_env_server.sh b/robots/franka/run_env_server.sh index 31c5d4dd..247cc17f 100755 --- a/robots/franka/run_env_server.sh +++ b/robots/franka/run_env_server.sh @@ -2,12 +2,14 @@ # Launch the standalone Franka env server in the RLinf .venv with the # serl_franka_controllers catkin workspace sourced. # -# The agent (physical/.venv) connects to this over TCP. Two ways to use it: +# The agent (agent venv) connects to this over TCP. Two ways to use it: # -# 1. Fixed port, then run the agent with --no-driver: -# bash robots/franka/run_env_server.sh --transport-port 5599 -# # in the physical/.venv: -# python rpent/cli/main.py --env franka --no-driver --env-port 5599 ... +# 1. Fixed port, then attach the agent via --env-endpoint: +# bash robots/franka/run_env_server.sh --output-dir /tmp/franka_run \ +# --transport socket --host 127.0.0.1 --port 5599 +# # in the agent venv: +# python rpent/cli/main.py --env franka \ +# --env-endpoint socket://127.0.0.1:5599 ... # # 2. Let the agent CLI spawn it (it invokes this script automatically): # python rpent/cli/main.py --env franka ... diff --git a/robots/lerobot/__init__.py b/robots/lerobot/__init__.py index 0f964872..032ce69b 100644 --- a/robots/lerobot/__init__.py +++ b/robots/lerobot/__init__.py @@ -10,13 +10,12 @@ from typing import Any -from rpent.envs.env_spec import EnvSpec -from rpent.envs.prompt_bundle import PromptBundle from robots.lerobot.prompt import ( system_prompt, user_prompt, ) -from rpent.utils.rpc import RpcClient +from rpent.envs.env_spec import EnvSpec +from rpent.envs.prompt_bundle import PromptBundle def get_env_spec() -> EnvSpec: @@ -36,16 +35,20 @@ def get_env_spec() -> EnvSpec: def get_toolkit( *, - rpc_client: RpcClient, + primitives_kwargs: dict[str, Any], video_path: str | None = None, dashboard: Any = None, ): - """Return the SO101 toolkit (common tools + SO101 primitives).""" - from robots.lerobot.env_client import LerobotEnvClient + """Return the SO101 toolkit (common tools + SO101 primitives). + + ``primitives_kwargs`` is assembled by ``_init_lerobot`` in + ``rpent/cli/main.py`` and carries the env RPC stub + (``{"env": LerobotEnvClient(...)}``, optionally plus ``"model"``). + """ from robots.lerobot.toolkit import LerobotToolkit return LerobotToolkit( - env=LerobotEnvClient(rpc_client), video_path=video_path, dashboard=dashboard, + **primitives_kwargs, ) diff --git a/robots/lerobot/env_server.py b/robots/lerobot/env_server.py index 19be4739..88eae8cb 100644 --- a/robots/lerobot/env_server.py +++ b/robots/lerobot/env_server.py @@ -2,9 +2,10 @@ Drives a physical SO101 follower arm through LeRobot's synchronous Python API (:class:`lerobot.robots.so_follower.SO101Follower`) and exposes a minimal -``reset`` / ``step`` gym-style surface over a pickle-framed TCP RPC server -(:class:`rpent.utils.socket_rpc.SocketRpcServer`) — the same wire -protocol the LIBERO driver uses, so the agent side talks to both identically. +``reset`` / ``step`` gym-style surface over an RPC server +(:class:`rpent.utils.rpc.RpcFacade`, socket transport by default) — the same +wire protocol the LIBERO driver uses, so the agent side talks to both +identically. Unlike the LIBERO driver, this server does **not** wrap an RLinf env class: importing ``rlinf.envs.realworld`` runs node-level ROS setup side effects at @@ -23,17 +24,16 @@ ``/dev/video2``, and an Intel RealSense D405 scene camera (serial ``409122274720``). Every default is overridable from the CLI. -Launched manually for now; wiring into ``cli/main.py`` (per-env client class + -driver script selection) is a separate step. +Spawned by ``rpent/cli/main.py`` when ``--env lerobot`` is selected (requires +the ``lerobot`` extra in the agent venv), or run manually in the ``lerobot`` +conda env and attached to via ``--env-endpoint socket://host:port``. """ from __future__ import annotations import argparse -import json import os import signal import sys -import threading import time from pathlib import Path @@ -45,13 +45,12 @@ if str(_PHYSICALAGENT_ROOT) not in sys.path: sys.path.insert(0, str(_PHYSICALAGENT_ROOT)) -from rpent.utils.socket_rpc import SocketRpcServer # noqa: E402 -from rpent.utils.logging import get_logger, init_output_dir # noqa: E402 - from robots.lerobot import calibration as scene_calib # noqa: E402 from robots.lerobot import geometry as geom # noqa: E402 from robots.lerobot.kinematics import SO101Kinematics # noqa: E402 from robots.lerobot.scene_camera import SceneCameraD405 # noqa: E402 +from rpent.utils.logging import get_logger, init_output_dir # noqa: E402 +from rpent.utils.rpc import RpcFacade # noqa: E402 logger = get_logger("lerobot_driver") @@ -1110,45 +1109,25 @@ def _compute_ee_pose(self, joints_deg) -> dict | None: # RPC plumbing # --------------------------------------------------------------------------- -_INITIAL_PPID = os.getppid() - - -def _start_parent_watchdog( - server: SocketRpcServer, - shutdown_event: threading.Event, - poll_s: float = 2.0, -) -> None: - """Shut the RPC server down if the agent (parent) process dies.""" - - def _watch() -> None: - while not shutdown_event.is_set(): - time.sleep(poll_s) - ppid = os.getppid() - if ppid != _INITIAL_PPID or ppid == 1: - logger.warning( - "parent died (ppid %s -> %s); stopping RPC server", - _INITIAL_PPID, - ppid, - ) - shutdown_event.set() - threading.Thread(target=server.shutdown, daemon=True).start() - return - - threading.Thread(target=_watch, daemon=True).start() +class LerobotEnvFacade(RpcFacade): + """Serve :class:`SO101LeRobotEnv` over the RPC boundary. + ``shutdown`` and ``healthz`` are handled by :class:`RpcFacade` and the + transport; this only routes ``env.*`` methods to the env instance. + """ -def _build_dispatcher(env: SO101LeRobotEnv, shutdown_event: threading.Event): - """Route ``env.*`` / ``shutdown`` to the right callable.""" + def __init__(self, env: SO101LeRobotEnv): + super().__init__() + self._env = env - def dispatch(method: str, args: tuple, kwargs: dict): + def _dispatch(self, method: str, args: tuple, kwargs: dict): if method.startswith("env."): - return getattr(env, method[len("env."):])(*args, **kwargs) - if method == "shutdown": - shutdown_event.set() - return {"ok": True} + return getattr(self._env, method[len("env."):])(*args, **kwargs) raise ValueError(f"unknown RPC method: {method!r}") - return dispatch + def request_shutdown(self) -> None: + """Signal the serve loop to exit (used by the signal handlers).""" + self._shutdown_event.set() # --------------------------------------------------------------------------- @@ -1205,8 +1184,12 @@ def _build_argparser() -> argparse.ArgumentParser: help="SO101 URDF for FK / EE pose. Default: " "~/.cache/huggingface/lerobot/urdf/so101.urdf") p.add_argument("--output-dir", required=True) - p.add_argument("--transport-port", type=int, default=0, - help="Socket transport port. 0 asks the OS for a free port.") + p.add_argument("--transport", choices=["socket", "http"], default="socket", + help="RPC transport. Socket (pickle) is the proven path for " + "the numpy obs/action payloads; http is also supported.") + p.add_argument("--host", type=str, default="127.0.0.1") + p.add_argument("--port", type=int, default=0, + help="RPC port. 0 asks the OS for a free port.") return p @@ -1262,47 +1245,27 @@ def main() -> int: auto_calibrate=args.auto_calibrate, ) - shutdown_event = threading.Event() - dispatch = _build_dispatcher(env, shutdown_event) - - server = SocketRpcServer(("127.0.0.1", args.transport_port), dispatch) - bound_host, bound_port = server.server_address - client_host = "127.0.0.1" if bound_host == "0.0.0.0" else bound_host - print( - json.dumps({ - "event": "transport_ready", - "kind": "socket", - "host": client_host, - "port": bound_port, - }), - flush=True, - ) - logger.info("RPC server listening on %s:%s", client_host, bound_port) + facade = LerobotEnvFacade(env) # Park the arm on SIGTERM / SIGINT (launcher kill, Ctrl-C). The handler - # only flags shutdown; the actual parking runs in the ``finally`` below, - # never inside the signal context. (SIGKILL / kill -9 cannot be caught.) + # only flags shutdown; the actual parking runs in ``env.close()`` in the + # ``finally`` below, never inside the signal context. (SIGKILL / kill -9 + # cannot be caught.) RpcFacade also exits on parent death (stdin EOF) and + # on the ``shutdown`` RPC. def _handle_signal(signum, _frame): logger.warning( "received %s; parking arm and shutting down", signal.Signals(signum).name, ) - shutdown_event.set() + facade.request_shutdown() for _sig in (signal.SIGTERM, signal.SIGINT): signal.signal(_sig, _handle_signal) - _start_parent_watchdog(server, shutdown_event) - threading.Thread(target=server.serve_forever, daemon=True).start() try: - # Loop with a timeout so a signal landing during the wait is observed - # promptly even if it doesn't interrupt the blocking call. - while not shutdown_event.wait(timeout=1.0): - pass + facade.serve(transport=args.transport, host=args.host, port=args.port) finally: env.close() - server.shutdown() - server.server_close() logger.info("driver exited cleanly") return 0 From 72855d001db7a2a6c987e8ed1b3e8ec3652b33e8 Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Thu, 23 Jul 2026 09:51:16 +0000 Subject: [PATCH 10/22] docs(lerobot): update calibrate_scene_cam for the new env_server flags The env_server now uses --port (not --transport-port); also fix the stale toolkits/ path to robots/. Signed-off-by: Jiaxing Qiu --- robots/lerobot/calibrate_scene_cam.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/robots/lerobot/calibrate_scene_cam.py b/robots/lerobot/calibrate_scene_cam.py index 7fbab10c..95f68633 100644 --- a/robots/lerobot/calibrate_scene_cam.py +++ b/robots/lerobot/calibrate_scene_cam.py @@ -21,14 +21,14 @@ After N>=4 non-coplanar points, a rigid Kabsch fit gives ``T_base_cam`` and the fit RMSE (lower is better; aim for < ~1 cm). -Run the env server first (note its --transport-port), then:: +Run the env server first with a fixed ``--port``, then:: conda activate lerobot - python toolkits/lerobot/calibrate_scene_cam.py --port 53101 --num-points 6 + python robots/lerobot/calibrate_scene_cam.py --port 53101 --num-points 6 Offline self-test of the math (no hardware):: - python toolkits/lerobot/calibrate_scene_cam.py --self-test + python robots/lerobot/calibrate_scene_cam.py --self-test """ from __future__ import annotations From aeb4baae6c743933de061d081293bbcf8479b54f Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Thu, 23 Jul 2026 10:01:04 +0000 Subject: [PATCH 11/22] refactor(envs): default lerobot/franka RPC transport to http Match libero's default: spawn the lerobot/franka env_servers with --transport http and talk to them over HttpRpcClient (http_rpc base64-encodes the numpy obs/action payloads). Socket transport stays available via --transport socket or an explicit socket://host:port --env-endpoint. Signed-off-by: Jiaxing Qiu --- robots/franka/env_server.py | 8 ++++---- robots/lerobot/env_server.py | 10 +++++----- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/robots/franka/env_server.py b/robots/franka/env_server.py index 6ace2506..03e4c15c 100644 --- a/robots/franka/env_server.py +++ b/robots/franka/env_server.py @@ -4,7 +4,7 @@ ``franka_gripper`` action topics, exposing agent-friendly *Cartesian* primitives (``reset`` / ``get_obs`` / ``get_ee_pose`` / ``move_to`` / ``move_delta`` / ``open_gripper`` / ``close_gripper`` / ``get_spec``) over an RPC server -(:class:`rpent.utils.rpc.RpcFacade`, socket transport by default) -- the same +(:class:`rpent.utils.rpc.RpcFacade`, http transport by default) -- the same wire protocol the LIBERO and LeRobot drivers use, so the agent side talks to all three identically. @@ -892,9 +892,9 @@ def main() -> int: "--camera", action="append", default=None, help="Camera as name:serial (repeatable). Defaults to the two bench D435I.", ) - p.add_argument("--transport", choices=["socket", "http"], default="socket", - help="RPC transport. Socket (pickle) is the proven path for " - "the numpy obs payloads; http is also supported.") + p.add_argument("--transport", choices=["socket", "http"], default="http", + help="RPC transport (default http). Socket (pickle) is also " + "available for the numpy obs payloads.") p.add_argument("--host", type=str, default="127.0.0.1") p.add_argument("--port", type=int, default=0, help="RPC port. 0 asks the OS for a free port.") diff --git a/robots/lerobot/env_server.py b/robots/lerobot/env_server.py index 88eae8cb..e127c206 100644 --- a/robots/lerobot/env_server.py +++ b/robots/lerobot/env_server.py @@ -3,7 +3,7 @@ Drives a physical SO101 follower arm through LeRobot's synchronous Python API (:class:`lerobot.robots.so_follower.SO101Follower`) and exposes a minimal ``reset`` / ``step`` gym-style surface over an RPC server -(:class:`rpent.utils.rpc.RpcFacade`, socket transport by default) — the same +(:class:`rpent.utils.rpc.RpcFacade`, http transport by default) — the same wire protocol the LIBERO driver uses, so the agent side talks to both identically. @@ -26,7 +26,7 @@ Spawned by ``rpent/cli/main.py`` when ``--env lerobot`` is selected (requires the ``lerobot`` extra in the agent venv), or run manually in the ``lerobot`` -conda env and attached to via ``--env-endpoint socket://host:port``. +conda env and attached to via ``--env-endpoint http://host:port``. """ from __future__ import annotations @@ -1184,9 +1184,9 @@ def _build_argparser() -> argparse.ArgumentParser: help="SO101 URDF for FK / EE pose. Default: " "~/.cache/huggingface/lerobot/urdf/so101.urdf") p.add_argument("--output-dir", required=True) - p.add_argument("--transport", choices=["socket", "http"], default="socket", - help="RPC transport. Socket (pickle) is the proven path for " - "the numpy obs/action payloads; http is also supported.") + p.add_argument("--transport", choices=["socket", "http"], default="http", + help="RPC transport (default http). Socket (pickle) is also " + "available for the numpy obs/action payloads.") p.add_argument("--host", type=str, default="127.0.0.1") p.add_argument("--port", type=int, default=0, help="RPC port. 0 asks the OS for a free port.") From 6dc7a2ef77e2590c515c11afb63261964501c2ec Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Fri, 24 Jul 2026 03:14:36 +0000 Subject: [PATCH 12/22] feat(envs): implement EnvSpec runner hooks for lerobot/franka The env refactor (#41) made rpent/cli/main.py env-agnostic: env-specific CLI args and runtime init now live behind EnvSpec's add_cli_args / parse_config / init_runtime hooks. Port lerobot and franka to that contract, mirroring robots/libero: - get_env_spec() now wires _add_cli_args / _parse_config / _init_runtime (previously these envs returned a bare EnvSpec(name, prompts), which no longer satisfies the now-required hook fields). - _init_runtime moves the old CLI-side _init_lerobot / _init_franka spawn/attach logic into the env package (ProcessDaemon spawn with --transport http, or attach via --env-endpoint [socket|http]://host:port; lazy heavy imports). - _parse_config derives a real-robot run identity (recipe_tag/output_dir/ task_desc keyed by env name; no suite/task/seed) and returns dashboard_state=None (the dashboard is libero-shaped; a warning is logged if --dashboard is passed). - main.py: widen --env choices to libero | lerobot | franka (the only CLI change; all other wiring is env-agnostic). Signed-off-by: Jiaxing Qiu --- robots/franka/__init__.py | 146 ++++++++++++++++++++++++++++++++++-- robots/lerobot/__init__.py | 148 +++++++++++++++++++++++++++++++++++-- rpent/cli/main.py | 5 +- 3 files changed, 284 insertions(+), 15 deletions(-) diff --git a/robots/franka/__init__.py b/robots/franka/__init__.py index aad07877..c559df4c 100644 --- a/robots/franka/__init__.py +++ b/robots/franka/__init__.py @@ -1,21 +1,41 @@ """Franka environment extension.""" from __future__ import annotations -from typing import Any +import argparse +from datetime import datetime +from pathlib import Path +from typing import TYPE_CHECKING, Any from robots.franka.prompt import system_prompt, user_prompt -from rpent.envs.env_spec import EnvSpec +from rpent.envs.env_spec import EnvSpec, RunConfig from rpent.envs.prompt_bundle import PromptBundle +from rpent.utils.config import get_repo_root +from rpent.utils.logging import get_logger + +if TYPE_CHECKING: + from rpent.utils.daemon import ProcessDaemon + from rpent.utils.rpc import RpcClient + +logger = get_logger("franka") def get_env_spec() -> EnvSpec: - """Return the Franka env identity + prompt bundle.""" + """Return the Franka env identity, prompt bundle, and runner hooks. + + Tool schemas, handlers, and the MCP allowlist live on the toolkit (see + :func:`get_toolkit`). The three runner hooks (:func:`_add_cli_args` / + :func:`_parse_config` / :func:`_init_runtime`) keep ``rpent/cli/main.py`` + env-agnostic, mirroring :mod:`robots.libero`. + """ return EnvSpec( name="franka", prompts=PromptBundle( system=system_prompt, user=user_prompt, ), + add_cli_args=_add_cli_args, + parse_config=_parse_config, + init_runtime=_init_runtime, ) @@ -27,9 +47,8 @@ def get_toolkit( ): """Return the Franka toolkit (common tools + Cartesian primitives). - ``primitives_kwargs`` is assembled by ``_init_franka`` in - ``rpent/cli/main.py`` and carries the env RPC stub - (``{"env": FrankaEnvClient(...)}``). + ``primitives_kwargs`` is assembled by :func:`_init_runtime` and carries the + env RPC stub (``{"env": FrankaEnvClient(...)}``). """ from robots.franka.toolkit import FrankaToolkit @@ -38,3 +57,118 @@ def get_toolkit( dashboard=dashboard, **primitives_kwargs, ) + + +def _add_cli_args(parser: argparse.ArgumentParser, use_dashboard: bool) -> None: + """Register Franka CLI flags on the shared ``parser``. + + ``use_dashboard`` is unused: the Franka setup is a real robot with no + suite/task/seed, so there is nothing for the (libero-shaped) dashboard + launcher to fill in. + """ + del use_dashboard + parser.add_argument( + "--env-endpoint", default=None, + help="[protocol://]host:port of an existing franka env_server " + "(protocol=http|socket, defaults to http). If unset, it is spawned " + "via run_env_server.sh (RLinf .venv).", + ) + + +def _parse_config(args: argparse.Namespace) -> RunConfig: + """Derive per-run identifiers for a Franka run. + + Real robots have no suite/task/seed, so the run is identified by the env + name. The dashboard is currently libero-shaped, so it is not wired here. + """ + if getattr(args, "dashboard", False): + logger.warning( + "--dashboard is only supported for the libero env; " + "continuing without the live dashboard." + ) + + recipe_tag = "franka" + output_dir = args.output_dir + if output_dir is None: + timestamp = datetime.now().strftime("%Y%m%d-%H:%M:%S") + output_dir = get_repo_root() / "logs" / f"{timestamp}_franka" + output_dir = Path(output_dir) + + return RunConfig( + recipe_tag=recipe_tag, + output_dir=output_dir, + prompt_vars={"env_name": "franka", "recipe_tag": recipe_tag}, + dashboard_state=None, + task_desc={"env": "franka"}, + ) + + +def _parse_endpoint(endpoint: str) -> tuple[str, str, int]: + """Parse ``[protocol://]host:port`` into ``(protocol, host, port)``. + + Protocol defaults to ``http`` when the prefix is omitted. + """ + if "://" in endpoint: + protocol, _, rest = endpoint.partition("://") + else: + protocol, rest = "http", endpoint + host, _, port = rest.partition(":") + if not host or not port: + raise ValueError( + f"--env-endpoint must be [protocol://]host:port, got {endpoint!r}" + ) + return protocol, host, int(port) + + +def _init_runtime( + args: argparse.Namespace, + output_dir: Path, +) -> tuple[list[ProcessDaemon], dict[str, Any]]: + """Spawn (or attach to) the Franka env_server; build primitives_kwargs. + + The Franka driver runs in the RLinf ``.venv`` with the catkin workspace + sourced, so it is spawned via ``run_env_server.sh`` (not this interpreter). + Pass ``--env-endpoint`` to attach to an already-running server. No VLA + server — the Cartesian primitives are scripted. + + Heavy deps are imported lazily so a bare ``import robots.franka`` (for + ``get_env_spec`` / ``get_toolkit``) doesn't drag them in. + """ + from robots.franka.env_client import FrankaEnvClient + from rpent.utils.daemon import ProcessDaemon, pick_free_port + from rpent.utils.http_rpc import HttpRpcClient + from rpent.utils.rpc import wait_for_ready + from rpent.utils.socket_rpc import SocketRpcClient + + daemons: list[ProcessDaemon] = [] + if args.env_endpoint is None: + host, port = "127.0.0.1", pick_free_port() + env_daemon = ProcessDaemon( + name="env_server", + cmd=[ + "bash", + str(get_repo_root() / "robots" / "franka" / "run_env_server.sh"), + "--output-dir", str(output_dir), + "--transport", "http", + "--host", host, + "--port", str(port), + ], + log_path=str(Path(output_dir) / "env_server.log"), + ) + env_daemon.start() + daemons.append(env_daemon) + env_client: RpcClient = HttpRpcClient(f"http://{host}:{port}") + wait_for_ready(env_client) + else: + protocol, host, port = _parse_endpoint(args.env_endpoint) + if protocol == "socket": + env_client = SocketRpcClient(host, port) + elif protocol == "http": + env_client = HttpRpcClient(f"http://{host}:{port}") + else: + raise ValueError( + f"--env-endpoint protocol must be socket or http, got {protocol!r}" + ) + + primitives_kwargs = {"env": FrankaEnvClient(env_client)} + return daemons, primitives_kwargs diff --git a/robots/lerobot/__init__.py b/robots/lerobot/__init__.py index 032ce69b..15969208 100644 --- a/robots/lerobot/__init__.py +++ b/robots/lerobot/__init__.py @@ -8,21 +8,35 @@ """ from __future__ import annotations -from typing import Any +import argparse +import sys +from datetime import datetime +from pathlib import Path +from typing import TYPE_CHECKING, Any from robots.lerobot.prompt import ( system_prompt, user_prompt, ) -from rpent.envs.env_spec import EnvSpec +from rpent.envs.env_spec import EnvSpec, RunConfig from rpent.envs.prompt_bundle import PromptBundle +from rpent.utils.config import get_repo_root +from rpent.utils.logging import get_logger + +if TYPE_CHECKING: + from rpent.utils.daemon import ProcessDaemon + from rpent.utils.rpc import RpcClient + +logger = get_logger("lerobot") def get_env_spec() -> EnvSpec: - """Return the SO101 env identity + prompt bundle. + """Return the SO101 env identity, prompt bundle, and runner hooks. Tool schemas, handlers, and the MCP allowlist live on the toolkit (see - :func:`get_toolkit`). + :func:`get_toolkit`). The three runner hooks (:func:`_add_cli_args` / + :func:`_parse_config` / :func:`_init_runtime`) keep ``rpent/cli/main.py`` + env-agnostic, mirroring :mod:`robots.libero`. """ return EnvSpec( name="lerobot", @@ -30,6 +44,9 @@ def get_env_spec() -> EnvSpec: system=system_prompt, user=user_prompt, ), + add_cli_args=_add_cli_args, + parse_config=_parse_config, + init_runtime=_init_runtime, ) @@ -41,9 +58,9 @@ def get_toolkit( ): """Return the SO101 toolkit (common tools + SO101 primitives). - ``primitives_kwargs`` is assembled by ``_init_lerobot`` in - ``rpent/cli/main.py`` and carries the env RPC stub - (``{"env": LerobotEnvClient(...)}``, optionally plus ``"model"``). + ``primitives_kwargs`` is assembled by :func:`_init_runtime` and carries the + env RPC stub (``{"env": LerobotEnvClient(...)}``, optionally plus + ``"model"``). """ from robots.lerobot.toolkit import LerobotToolkit @@ -52,3 +69,120 @@ def get_toolkit( dashboard=dashboard, **primitives_kwargs, ) + + +def _add_cli_args(parser: argparse.ArgumentParser, use_dashboard: bool) -> None: + """Register SO101 CLI flags on the shared ``parser``. + + ``use_dashboard`` is unused: the SO101 is a real robot with no + suite/task/seed, so there is nothing for the (libero-shaped) dashboard + launcher to fill in. + """ + del use_dashboard + parser.add_argument("--max-episode-steps", type=int, default=200) + parser.add_argument( + "--env-endpoint", default=None, + help="[protocol://]host:port of an existing lerobot env_server " + "(protocol=http|socket, defaults to http). If unset, a local " + "env_server is spawned (requires the 'lerobot' extra in this venv).", + ) + + +def _parse_config(args: argparse.Namespace) -> RunConfig: + """Derive per-run identifiers for a SO101 run. + + Real robots have no suite/task/seed, so the run is identified by the env + name. The dashboard is currently libero-shaped, so it is not wired here. + """ + if getattr(args, "dashboard", False): + logger.warning( + "--dashboard is only supported for the libero env; " + "continuing without the live dashboard." + ) + + recipe_tag = "lerobot" + output_dir = args.output_dir + if output_dir is None: + timestamp = datetime.now().strftime("%Y%m%d-%H:%M:%S") + output_dir = get_repo_root() / "logs" / f"{timestamp}_lerobot" + output_dir = Path(output_dir) + + return RunConfig( + recipe_tag=recipe_tag, + output_dir=output_dir, + prompt_vars={"env_name": "lerobot", "recipe_tag": recipe_tag}, + dashboard_state=None, + task_desc={"env": "lerobot"}, + ) + + +def _parse_endpoint(endpoint: str) -> tuple[str, str, int]: + """Parse ``[protocol://]host:port`` into ``(protocol, host, port)``. + + Protocol defaults to ``http`` when the prefix is omitted. + """ + if "://" in endpoint: + protocol, _, rest = endpoint.partition("://") + else: + protocol, rest = "http", endpoint + host, _, port = rest.partition(":") + if not host or not port: + raise ValueError( + f"--env-endpoint must be [protocol://]host:port, got {endpoint!r}" + ) + return protocol, host, int(port) + + +def _init_runtime( + args: argparse.Namespace, + output_dir: Path, +) -> tuple[list[ProcessDaemon], dict[str, Any]]: + """Spawn (or attach to) the SO101 env_server; build primitives_kwargs. + + The SO101 driver needs the ``lerobot`` stack. It is spawned with this + interpreter (works when the venv has the ``lerobot`` extra); otherwise run + the driver in its own ``lerobot`` conda env and attach via + ``--env-endpoint``. No VLA server — the SO101 primitives are scripted. + + Heavy deps are imported lazily so a bare ``import robots.lerobot`` (for + ``get_env_spec`` / ``get_toolkit``) doesn't drag them in. + """ + from robots.lerobot.env_client import LerobotEnvClient + from rpent.utils.daemon import ProcessDaemon, pick_free_port + from rpent.utils.http_rpc import HttpRpcClient + from rpent.utils.rpc import wait_for_ready + from rpent.utils.socket_rpc import SocketRpcClient + + daemons: list[ProcessDaemon] = [] + if args.env_endpoint is None: + host, port = "127.0.0.1", pick_free_port() + env_daemon = ProcessDaemon( + name="env_server", + cmd=[ + sys.executable, + str(get_repo_root() / "robots" / "lerobot" / "env_server.py"), + "--output-dir", str(output_dir), + "--max-episode-steps", str(args.max_episode_steps), + "--transport", "http", + "--host", host, + "--port", str(port), + ], + log_path=str(Path(output_dir) / "env_server.log"), + ) + env_daemon.start() + daemons.append(env_daemon) + env_client: RpcClient = HttpRpcClient(f"http://{host}:{port}") + wait_for_ready(env_client) + else: + protocol, host, port = _parse_endpoint(args.env_endpoint) + if protocol == "socket": + env_client = SocketRpcClient(host, port) + elif protocol == "http": + env_client = HttpRpcClient(f"http://{host}:{port}") + else: + raise ValueError( + f"--env-endpoint protocol must be socket or http, got {protocol!r}" + ) + + primitives_kwargs = {"env": LerobotEnvClient(env_client)} + return daemons, primitives_kwargs diff --git a/rpent/cli/main.py b/rpent/cli/main.py index 2596793d..f8a3bffb 100644 --- a/rpent/cli/main.py +++ b/rpent/cli/main.py @@ -86,8 +86,9 @@ def _build_argparser() -> argparse.ArgumentParser: description="Standalone hybrid LLM-in-the-loop agent for LIBERO PRO", ) - ap.add_argument("--env", dest="env_name", required=True, choices=["libero"], - help="Environment backend: libero.") + ap.add_argument("--env", dest="env_name", required=True, + choices=["libero", "lerobot", "franka"], + help="Environment backend: libero | lerobot | franka.") # models ap.add_argument("--planner", default="api", From 93a53f101cbe4eeea596bb2f86c6277150830d14 Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Fri, 24 Jul 2026 03:24:49 +0000 Subject: [PATCH 13/22] align pyproject Signed-off-by: Jiaxing Qiu --- pyproject.toml | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0155944d..bd26af04 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,10 @@ keywords = [ "agent-framework", ] classifiers = [ + # 2 - Pre-Alpha + # 3 - Alpha + # 4 - Beta + # 5 - Production/Stable "Development Status :: 2 - Pre-Alpha", "Intended Audience :: Developers", "Intended Audience :: Robotics Enthusiasts", @@ -63,6 +67,10 @@ libero-plus = [ "rpent[libero]", "rlinf-liberoplus", ] +lerobot = [ + "lerobot[feetech,intelrealsense,kinematics]>=0.5.2", + "matplotlib", +] full = [ "rpent[rlinf,openpi,libero-pro,sam3]", ] @@ -72,10 +80,6 @@ sam3 = [ "torchvision", "pillow>=10", ] -lerobot = [ - "lerobot[feetech,intelrealsense,kinematics]>=0.5.2", - "matplotlib", -] [tool.uv] prerelease = "allow" From 1ad6a569b208d1ed42304785126fc398a0e099dd Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Mon, 27 Jul 2026 16:24:07 +0800 Subject: [PATCH 14/22] fix: lerobot uv sync conflict Signed-off-by: Jiaxing Qiu --- pyproject.toml | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index bd26af04..00b1d343 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -68,7 +68,7 @@ libero-plus = [ "rlinf-liberoplus", ] lerobot = [ - "lerobot[feetech,intelrealsense,kinematics]>=0.5.2", + "lerobot[feetech,intelrealsense,kinematics]>=0.4.4", "matplotlib", ] full = [ @@ -83,6 +83,14 @@ sam3 = [ [tool.uv] prerelease = "allow" +# lerobot is not co-installable with the RLinf stack (libero/openpi). +conflicts = [ + [{ extra = "lerobot" }, { extra = "openpi" }], + [{ extra = "lerobot" }, { extra = "libero" }], + [{ extra = "lerobot" }, { extra = "libero-pro" }], + [{ extra = "lerobot" }, { extra = "libero-plus" }], + [{ extra = "lerobot" }, { extra = "full" }], +] [tool.setuptools] include-package-data = true @@ -98,7 +106,7 @@ exclude = ["tests*", "docs*", "examples*", "docker*", "toolkits*"] [tool.ruff] line-length = 88 indent-width = 4 -target-version = "py312" +target-version = "py311" [tool.ruff.lint] isort = {known-first-party = ["rpent"]} From 7320580e03b5f6ec4a87937fb65b6f9d9ed08981 Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Mon, 27 Jul 2026 16:30:01 +0800 Subject: [PATCH 15/22] fix(lerobot): rename --port to --serial-port and simplify user prompt Signed-off-by: Jiaxing Qiu --- robots/lerobot/env_server.py | 10 +++++----- robots/lerobot/prompt.py | 8 ++------ 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/robots/lerobot/env_server.py b/robots/lerobot/env_server.py index e127c206..954eb795 100644 --- a/robots/lerobot/env_server.py +++ b/robots/lerobot/env_server.py @@ -1137,11 +1137,11 @@ def request_shutdown(self) -> None: def _build_argparser() -> argparse.ArgumentParser: p = argparse.ArgumentParser(description="Standalone LeRobot SO101 env server") - p.add_argument("--port", default="/dev/ttyACM1", + p.add_argument("--serial-port", default="/dev/ttyACM0", help="Serial port of the SO101 follower arm.") p.add_argument("--calibration-id", default="my_awesome_follower_arm", help="LeRobot calibration id (loads .json).") - p.add_argument("--arm-camera-path", default="/dev/video2", + p.add_argument("--arm-camera-path", default="/dev/v4l/by-id/usb-icSpring_icspring_camera-video-index0", help="OpenCV device path for the arm/hand camera " "(empty string to disable).") p.add_argument("--scene-camera-serial", default="409122274720", @@ -1221,12 +1221,12 @@ def main() -> int: arm_camera_cfgs = _build_arm_camera_cfgs(args) scene_serial = None if args.no_cameras else (args.scene_camera_serial or None) logger.info( - "starting SO101 env server: port=%s output_dir=%s arm_cams=%s scene=%s", - args.port, args.output_dir, list(arm_camera_cfgs), scene_serial, + "starting SO101 env server: serial_port=%s rpc_port=%s output_dir=%s arm_cams=%s scene=%s", + args.serial_port, args.port, args.output_dir, list(arm_camera_cfgs), scene_serial, ) env = SO101LeRobotEnv( - port=args.port, + port=args.serial_port, calibration_id=args.calibration_id, arm_camera_cfgs=arm_camera_cfgs, scene_serial=scene_serial, diff --git a/robots/lerobot/prompt.py b/robots/lerobot/prompt.py index 97a3a2bd..3e05c7bf 100644 --- a/robots/lerobot/prompt.py +++ b/robots/lerobot/prompt.py @@ -136,11 +136,7 @@ # --- user-prompt sections -------------------------------------------------- -USER_CONTEXT = { - "Task": """ - Pick up the green cube on the table and place it in the white plate. - """, -} +USER_CONTEXT = """Pick up the green cube on the table.""" # --- prompt tree factories ------------------------------------------------- @@ -160,4 +156,4 @@ def system_prompt() -> dict[str, PromptNode]: def user_prompt() -> dict[str, PromptNode]: """Return the first user message tree.""" - return dict(USER_CONTEXT) + return USER_CONTEXT From 7b65febb545e24d025c4de0c5e8a361d5dae9598 Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Tue, 28 Jul 2026 16:55:00 +0800 Subject: [PATCH 16/22] feat(lerobot): enhance auto-calibration with joint offset recovery and improved target grid Signed-off-by: Jiaxing Qiu --- robots/lerobot/auto_calibrate_scene_cam.py | 60 ++++++++++-- robots/lerobot/env_server.py | 109 +++++++++++++++------ robots/lerobot/geometry.py | 104 +++++++++++++++++++- 3 files changed, 231 insertions(+), 42 deletions(-) diff --git a/robots/lerobot/auto_calibrate_scene_cam.py b/robots/lerobot/auto_calibrate_scene_cam.py index 287c1cba..4453f258 100644 --- a/robots/lerobot/auto_calibrate_scene_cam.py +++ b/robots/lerobot/auto_calibrate_scene_cam.py @@ -3,13 +3,15 @@ Triggers the env server's :meth:`auto_calibrate_scene_camera` routine, which: -1. drives the gripper to a spread grid of base-frame positions (``move_to`` is - pure base-frame IK, so it needs no extrinsic), +1. drives the gripper to a wide, non-coplanar grid of base-frame positions at + varied wrist orientations (``move_to`` is pure base-frame IK, so it needs no + extrinsic), 2. at each pose toggles the gripper with the arm frozen and segments the motion - in the scene image to locate the tip (centroid + median depth -> camera - point); the achieved FK gives the base point, -3. fits ``T_base_cam`` with RANSAC Kabsch and saves it (hot-loaded by the - server, so back_project returns world coords immediately). + in the temporally median-filtered scene image to locate the moving jaw (blob + centroid + depth at that centroid -> camera point); FK gives the tip pose, +3. jointly fits ``T_base_cam`` and the constant centroid-vs-tip offset (RANSAC), + so the offset no longer inflates the residual, and saves it (hot-loaded by + the server, so back_project returns world coords immediately). No human input, no markers. Start the env server first, then run:: @@ -66,7 +68,45 @@ def _self_test() -> int: det = geom.detect_tip_pixel_by_motion(rgb_open, rgb_closed, depth, K) ok_det = det is not None and abs(det["pixel"][0] - 314.5) < 3 and abs(det["pixel"][1] - 414.5) < 3 print(f"detect_tip: {det if det else 'None'} -> {'OK' if ok_det else 'FAIL'}") - return 0 if (ok_fit and ok_det) else 1 + + # Joint extrinsic + tip-offset recovery (solve_extrinsic_with_offset): with + # varied per-pose rotations the solver should recover T_base_cam AND the + # constant local offset, and beat the offset-blind fit, despite an outlier. + rng3 = np.random.default_rng(2) + Qc, _ = np.linalg.qr(rng3.standard_normal((3, 3))) + if np.linalg.det(Qc) < 0: + Qc[:, 0] = -Qc[:, 0] + T_cam = np.eye(4) + T_cam[:3, :3] = Qc + T_cam[:3, 3] = rng3.standard_normal(3) + d_true = np.array([0.015, -0.010, 0.020]) # constant gripper-local offset + n = 24 + origins = rng3.uniform([-0.1, -0.2, 0.0], [0.35, 0.2, 0.25], size=(n, 3)) + rots = np.empty((n, 3, 3)) + for i in range(n): + Qr, _ = np.linalg.qr(rng3.standard_normal((3, 3))) + if np.linalg.det(Qr) < 0: + Qr[:, 0] = -Qr[:, 0] + rots[i] = Qr + feat_base = origins + np.einsum("nij,j->ni", rots, d_true) + cam_j = geom.transform_points(geom.invert_transform(T_cam), feat_base) + cam_j += rng3.standard_normal((n, 3)) * 1e-3 # 1 mm detection noise + cam_j[5] += np.array([0.18, -0.12, 0.15]) # inject an outlier + T_est_j, rmse_j, inl, d_est = geom.solve_extrinsic_with_offset( + cam_j, origins, rots, thresh_m=0.01, + ) + _, rmse0, _ = geom.ransac_kabsch(cam_j, origins, thresh_m=0.01) + ok_solver = ( + np.allclose(T_est_j[:3, :3], T_cam[:3, :3], atol=5e-3) + and np.allclose(T_est_j[:3, 3], T_cam[:3, 3], atol=5e-3) + and np.allclose(d_est, d_true, atol=5e-3) + and not bool(inl[5]) + ) + print(f"solve_offset: rmse={rmse_j * 1000:.2f}mm " + f"(offset-blind {rmse0 * 1000:.1f}mm) " + f"d_err={np.linalg.norm(d_est - d_true) * 1000:.2f}mm " + f"outlier_excluded={not bool(inl[5])} -> {'OK' if ok_solver else 'FAIL'}") + return 0 if (ok_fit and ok_det and ok_solver) else 1 def main() -> int: @@ -74,7 +114,7 @@ def main() -> int: formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--host", default="127.0.0.1") ap.add_argument("--port", type=int, help="env server transport port.") - ap.add_argument("--n-points", type=int, default=10, + ap.add_argument("--n-points", type=int, default=24, help="Target number of valid correspondences to collect.") ap.add_argument("--no-save", action="store_true", help="Compute T_base_cam but do not write it to disk.") @@ -104,6 +144,10 @@ def main() -> int: print(f"\nUsed {result['n_used']} poses " f"({result['n_inliers']} inliers), RMSE = {result['rmse_m'] * 1000:.1f} mm") + off = result.get("tip_offset_local_m") + if off: + print("Estimated tip-detector offset (gripper frame): " + f"[{off[0] * 1000:.1f}, {off[1] * 1000:.1f}, {off[2] * 1000:.1f}] mm") if result["rmse_m"] > 0.02: print("WARNING: RMSE > 2 cm — check lighting / gripper visibility and rerun.") if result.get("saved"): diff --git a/robots/lerobot/env_server.py b/robots/lerobot/env_server.py index 954eb795..f1dc303e 100644 --- a/robots/lerobot/env_server.py +++ b/robots/lerobot/env_server.py @@ -104,11 +104,9 @@ # fixed side, NOT between the fingers. This vector (meters, in the # gripper_frame_link LOCAL frame) shifts the controlled/reported point to the # grasp point between the fingertips, so ``move_to`` targets and ``get_ee_pose`` -# refer to where an object is actually grasped. Derived from URDF + FK + a -# calibrated-depth fingertip measurement (~2 cm lateral toward the moving jaw, -# ~2.8 cm down to the fingertip plane). TUNE on hardware / via touch calibration +# refer to where an object is actually grasped. TUNE on hardware / via touch calibration # if grasps land consistently off-centre; set to zeros for the raw frame. -_TCP_OFFSET_GRIPPER = np.array([-0.028, 0.043, 0.0282], dtype=np.float64) +_TCP_OFFSET_GRIPPER = np.array([0.025, -0.01, 0.02], dtype=np.float64) # --- motion speed / smoothness (safety) --------------------------------- # The arm runs in position mode. Two knobs keep motions slow and gentle: @@ -119,8 +117,8 @@ # move_joints_delta) are streamed as interpolated setpoints so no joint # exceeds ``_MAX_JOINT_VEL_DEG_S`` (deg/s), instead of snapping to the target # at full servo speed. Both are overridable from the CLI. -_MOTOR_ACCELERATION = 40 -_MAX_JOINT_VEL_DEG_S = 70.0 +_MOTOR_ACCELERATION = 30 +_MAX_JOINT_VEL_DEG_S = 60.0 _PACE_DT_S = 0.05 # setpoint-streaming period (20 Hz) # Feetech position gain. LeRobot's configure() lowers P_Coefficient to 16 (from # the factory default 32) "to avoid shakiness"; that soft gain lets gravity- @@ -908,16 +906,50 @@ def _set_gripper_hold(self, gripper_value: float) -> None: @staticmethod def _calibration_targets() -> list[list[float]]: - """A spread, non-coplanar grid of tip targets inside the workspace.""" - xs = [0.15, 0.22, 0.28] - ys = [-0.12, 0.0, 0.12] - zs = [0.12, 0.19] - return [[x, y, z] for z in zs for y in ys for x in xs] + """A wide, non-coplanar grid of tip targets inside the workspace. + + Free-orientation IK reaches these, giving a large spread in x/y/z AND a + variety of wrist orientations. That orientation variety is what makes + the constant tip-detector offset identifiable (see + ``geometry.solve_extrinsic_with_offset``), so the calibration keeps the + default free approach rather than a fixed top-down one. Ordered z-fastest + so an early stop (``n_points``) still spans all three heights. + """ + xs = [0.15, 0.21, 0.27, 0.33] + ys = [-0.16, -0.05, 0.05, 0.16] + zs = [0.08, 0.15, 0.22] + return [[x, y, z] for x in xs for y in ys for z in zs] + + def _capture_scene_median( + self, n_frames: int = 5 + ) -> tuple[np.ndarray, np.ndarray]: + """Grab several scene frames and return per-pixel temporal medians. + + Median-averaging across frames suppresses the camera's per-frame color + and depth noise, which otherwise becomes lateral error once a pixel is + back-projected through the oblique view. Depth invalids (<=0 / + non-finite) are ignored per pixel; pixels with no valid sample stay 0. + """ + import warnings + + rgbs: list[np.ndarray] = [] + depths: list[np.ndarray] = [] + for _ in range(max(1, int(n_frames))): + rgb, depth = self._scene_cam.read() + rgbs.append(np.asarray(rgb)) + d = np.asarray(depth, dtype=np.float32) + depths.append(np.where(np.isfinite(d) & (d > 0), d, np.nan)) + rgb_med = np.median(np.stack(rgbs, axis=0), axis=0).astype(np.uint8) + with warnings.catch_warnings(): # nanmedian warns on all-invalid pixels + warnings.simplefilter("ignore", category=RuntimeWarning) + depth_med = np.nanmedian(np.stack(depths, axis=0), axis=0) + depth_med = np.nan_to_num(depth_med, nan=0.0).astype(np.float32) + return rgb_med, depth_med def auto_calibrate_scene_camera( self, *, - n_points: int = 10, + n_points: int = 24, gripper_open: float = 90.0, gripper_closed: float = 20.0, settle_s: float = 0.8, @@ -926,14 +958,20 @@ def auto_calibrate_scene_camera( ) -> dict: """Markerless automatic scene-cam -> base calibration. - Drives the tip to a grid of base-frame positions (move_to needs no - extrinsic), and at each pose toggles the gripper with the arm frozen and - segments the motion in the scene image to locate the tip (centroid + - median depth -> camera point). The achieved FK gives the base point. - A RANSAC Kabsch fit then yields ``T_base_cam``, which is saved and - hot-loaded so back_project returns world coords immediately. - - Returns a summary dict (``n_used``, ``rmse_m``, per-pose diagnostics). + Drives the tip to a wide, non-coplanar grid of base-frame positions + (move_to needs no extrinsic) at varied wrist orientations. At each pose + it toggles the gripper with the arm frozen and segments the motion in + the (temporally median-filtered) scene image to locate the moving jaw + (blob centroid + depth AT that centroid -> camera point); FK gives the + tip pose (origin + rotation). A joint fit + (:func:`geometry.solve_extrinsic_with_offset`) then recovers both + ``T_base_cam`` and the constant offset between the detected centroid and + the tip frame -- so that offset no longer inflates the residual -- and + the extrinsic is saved and hot-loaded so back_project returns world + coords immediately. + + Returns a summary dict (``n_used``, ``rmse_m``, ``tip_offset_local_m``, + per-pose diagnostics). """ if self._scene_cam is None: return {"error": "scene camera not configured"} @@ -942,7 +980,8 @@ def auto_calibrate_scene_camera( targets = self._calibration_targets() cam_pts: list[list[float]] = [] - base_pts: list[list[float]] = [] + tip_origins: list[list[float]] = [] + tip_rots: list[list[list[float]]] = [] poses: list[dict] = [] for tgt in targets: @@ -954,12 +993,19 @@ def auto_calibrate_scene_camera( continue time.sleep(settle_s) - base = np.asarray(mv["final_xyz"], dtype=np.float64) - self._scene_cam.read() # flush a frame - rgb_open, _ = self._scene_cam.read() + # Free-orientation FK tip pose (origin + rotation) at this pose. The + # rotation is what lets the fit solve out the constant tip-detector + # offset (geometry.solve_extrinsic_with_offset), so orientation must + # vary across the grid -- hence the free (not top-down) approach. + T_tip = self._kin.fk(self._read_arm_joints()) + o_i = T_tip[:3, 3] + R_i = T_tip[:3, :3] + + self._scene_cam.read() # drop the in-flight frame from the move + rgb_open, _ = self._capture_scene_median() self._set_gripper_hold(gripper_closed) time.sleep(settle_s) - rgb_closed, depth = self._scene_cam.read() + rgb_closed, depth = self._capture_scene_median() self._set_gripper_hold(gripper_open) # reopen for the next pose det = geom.detect_tip_pixel_by_motion( @@ -969,8 +1015,9 @@ def auto_calibrate_scene_camera( poses.append({"target": tgt, "skipped": "no_tip_detected"}) continue cam_pts.append(det["xyz_cam"]) - base_pts.append(base.tolist()) - poses.append({"target": tgt, "base_xyz": base.round(4).tolist(), + tip_origins.append(o_i.tolist()) + tip_rots.append(R_i.tolist()) + poses.append({"target": tgt, "base_xyz": o_i.round(4).tolist(), "pixel": [round(v, 1) for v in det["pixel"]], "depth_m": round(det["depth_m"], 4), "area": det["area"]}) @@ -978,8 +1025,11 @@ def auto_calibrate_scene_camera( return {"error": f"only {len(cam_pts)} usable points (need >= 4)", "poses": poses} - T, rmse, inliers = geom.ransac_kabsch( - cam_pts, base_pts, thresh_m=ransac_thresh_m, + # Joint fit: recover T_base_cam AND the constant offset between the + # detected motion-blob centroid and the FK tip frame, so that offset no + # longer pollutes the residual (the old ~1 cm RMSE floor). + T, rmse, inliers, tip_offset = geom.solve_extrinsic_with_offset( + cam_pts, tip_origins, tip_rots, thresh_m=ransac_thresh_m, ) accepted = bool(rmse <= scene_calib.MAX_ACCEPTABLE_RMSE_M) result = { @@ -987,6 +1037,7 @@ def auto_calibrate_scene_camera( "n_used": len(cam_pts), "n_inliers": int(np.asarray(inliers).sum()), "rmse_m": round(float(rmse), 4), + "tip_offset_local_m": [round(float(v), 4) for v in tip_offset], "T_base_cam": T.tolist(), "accepted": accepted, "saved": False, diff --git a/robots/lerobot/geometry.py b/robots/lerobot/geometry.py index a9d260a0..198c8df3 100644 --- a/robots/lerobot/geometry.py +++ b/robots/lerobot/geometry.py @@ -153,6 +153,93 @@ def ransac_kabsch( T, rmse = kabsch_umeyama(src[best_mask], dst[best_mask]) return T, rmse, best_mask + +def solve_extrinsic_with_offset( + cam_pts, + tip_origins, + tip_rotations, + *, + thresh_m: float = 0.015, + offset_iters: int = 15, + outer_iters: int = 3, + seed: int = 0, +) -> tuple[np.ndarray, float, np.ndarray, np.ndarray]: + """Joint fit of ``T_base_cam`` AND a constant gripper-local tip offset. + + The markerless routine detects the moving jaw's motion-blob centroid in the + camera, but matches it against the FK tip-frame ORIGIN in the base frame. + Those are not the same physical point: the centroid is displaced from the + tip-frame origin by an (approximately) constant vector ``d`` expressed in + the gripper LOCAL frame. In a free-orientation grid that displacement points + a different way in the base frame at every pose, so a single rigid + ``T_base_cam`` cannot absorb it -- it lands in the residual and is the main + driver of the ~1 cm fit RMSE. + + This models the displacement explicitly. For pose ``i`` with tip-frame + origin ``o_i`` and rotation ``R_i`` (base frame, from FK) and detected + camera point ``c_i``:: + + T_base_cam @ c_i == o_i + R_i @ d + + It alternates two convex steps: (1) with ``d`` fixed, Kabsch-fit ``T`` to the + corrected targets ``b_i = o_i + R_i @ d``; (2) with ``T`` fixed, solve the + linear least squares ``R_i @ d = T @ c_i - o_i`` (closed form + ``d = mean_i R_i^T (T c_i - o_i)`` since each ``R_i`` is orthonormal). + Gross outliers (bad detections) are rejected with RANSAC. ``d`` is only + identifiable when the tip ROTATIONS vary across poses; with (near-)constant + orientation the displacement is indistinguishable from ``T``'s translation + and the solver returns ``d ~= 0`` (reducing to the plain rigid fit). + + Args: + cam_pts: ``(N, 3)`` detected points in the camera frame. + tip_origins: ``(N, 3)`` FK tip-frame origins in the base frame. + tip_rotations: ``(N, 3, 3)`` FK tip-frame rotations in the base frame. + thresh_m: RANSAC inlier threshold (m) on the offset-corrected fit. + offset_iters: max inner alternations per outer round. + outer_iters: rounds of (refine offset -> re-select inliers). + seed: RANSAC RNG seed. + + Returns: + ``(T_4x4, inlier_rmse_m, inlier_mask, d_local)``. + """ + cam = np.asarray(cam_pts, dtype=np.float64) + o = np.asarray(tip_origins, dtype=np.float64) + Rs = np.asarray(tip_rotations, dtype=np.float64) + n = cam.shape[0] + if n < 4 or Rs.shape != (n, 3, 3): + T, rmse = kabsch_umeyama(cam, o) + return T, rmse, np.ones(n, dtype=bool), np.zeros(3) + + # Stage 1: a loose RANSAC (tolerating the still-unknown offset) drops gross + # outliers -- mis-detected tips -- before we estimate the offset. + loose = max(thresh_m, 0.07) + T, rmse, mask = ransac_kabsch(cam, o, thresh_m=loose, seed=seed) + if int(mask.sum()) < 4: + mask = np.ones(n, dtype=bool) + + d = np.zeros(3) + for _ in range(outer_iters): + # Alternate: solve the offset from residuals, refit T to the corrected + # targets, until the offset stops moving. + for _ in range(offset_iters): + r = transform_points(T, cam[mask]) - o[mask] # (M, 3) + # mean_i R_i^T r_i (R_i^T r_i via 'mba,mb->ma' since R_i is (i,j)=(row,col)) + d_new = np.einsum("mba,mb->ma", Rs[mask], r).mean(axis=0) + b = o + np.einsum("nij,j->ni", Rs, d_new) # o_i + R_i d + T, _ = kabsch_umeyama(cam[mask], b[mask]) + if np.linalg.norm(d_new - d) < 1e-6: + d = d_new + break + d = d_new + # Re-select inliers with the tight threshold on the corrected targets. + b = o + np.einsum("nij,j->ni", Rs, d) + T, rmse, mask = ransac_kabsch(cam, b, thresh_m=thresh_m, seed=seed) + if int(mask.sum()) < 4: + mask = np.ones(n, dtype=bool) + T, rmse = kabsch_umeyama(cam, b) + return T, rmse, mask, d + + def rotation_to_quat(R) -> np.ndarray: """Convert a 3x3 rotation matrix to a quaternion ``[w, x, y, z]``.""" R = np.asarray(R, dtype=np.float64) @@ -241,11 +328,18 @@ def detect_tip_pixel_by_motion( if area < min_area or area > max_area: continue col, row = float(centroids[comp][0]), float(centroids[comp][1]) - dvals = depth_m[labels == comp] - dvals = dvals[np.isfinite(dvals) & (dvals > 0)] - if dvals.size == 0: - continue - z = float(np.median(dvals)) + # Depth AT the centroid (small patch), so the returned (col, row, z) all + # describe the SAME point. The blob spans a depth gradient under the + # oblique scene view, so a whole-blob median would pair the centroid + # pixel with some other pixel's depth and bias the back-projection. + z = sample_depth_patch(depth_m, int(round(col)), int(round(row)), radius=3) + if not np.isfinite(z): + # Centroid fell on a depth dropout -- fall back to the blob median. + dvals = depth_m[labels == comp] + dvals = dvals[np.isfinite(dvals) & (dvals > 0)] + if dvals.size == 0: + continue + z = float(np.median(dvals)) p_cam = backproject_pixel(K, col, row, z) return { "pixel": [row, col], From dace1076fc2ae9c9226baf873e90c3cb9cf14246 Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Mon, 3 Aug 2026 11:55:11 +0800 Subject: [PATCH 17/22] refactor(state): initial state refactors Signed-off-by: Jiaxing Qiu --- robots/franka/toolkit.py | 50 +++--- robots/franka/tools.py | 287 +++++++++---------------------- robots/lerobot/toolkit.py | 44 +++-- robots/lerobot/tools.py | 311 +++++++++------------------------- robots/libero/toolkit.py | 61 ++++--- robots/libero/tools.py | 248 ++++++++++++--------------- rpent/tools/common.py | 88 ++++++++++ rpent/tools/state.py | 348 ++++++++++++++++++++++++++++++++++++++ rpent/tools/toolkit.py | 8 +- 9 files changed, 802 insertions(+), 643 deletions(-) create mode 100644 rpent/tools/state.py diff --git a/robots/franka/toolkit.py b/robots/franka/toolkit.py index 76ab0951..e7212e5a 100644 --- a/robots/franka/toolkit.py +++ b/robots/franka/toolkit.py @@ -1,12 +1,12 @@ """Franka toolkit: common tools plus conservative Cartesian primitives.""" from __future__ import annotations -import shutil import time from functools import partial from typing import Any from robots.franka import tools as franka_tools +from rpent.tools.state import EnvState from rpent.tools.toolkit import Toolkit from rpent.utils.logging import get_output_dir @@ -14,10 +14,11 @@ class FrankaToolkit(Toolkit): """Toolkit for the standalone Franka environment.""" - _STATELESS_TOOLS = ( - "view_driver_state", - "back_project", - ) + # Streams wiped on init (LIBERO layout: scene=image/depth, wrist=image_wrist/depth_wrist). + _WIPE_STREAMS = ("image", "image_wrist", "depth", "depth_wrist") + # view_driver_state image slots: primary scene -> _image_bytes, wrist -> _image_cam_bytes. + _VIEW_IMAGE_SLOTS = {"_image_bytes": "image", "_image_cam_bytes": "image_wrist"} + _DRIVER_READERS = ( "get_ee_pose", "get_robot_spec", @@ -42,16 +43,28 @@ def __init__( video_path: str | None = None, dashboard: Any = None, ) -> None: - super().__init__(dashboard=dashboard) - self._next_step = 0 + # EnvState owns the trace + counter for this run (explicit output_dir, + # no process-global). The runner will own its lifecycle in a later cut; + # for now the toolkit constructs it from get_output_dir(). + state = EnvState(get_output_dir()) + super().__init__(dashboard=dashboard, state=state) self._video_path = video_path self.init_driver_clean(env=env) self._register_tools() def _register_tools(self) -> None: spec = self._SPECS - for name in self._STATELESS_TOOLS: - self.add_tool(name, spec[name], getattr(franka_tools, name)) + # view_driver_state / back_project read trace state -> bind to EnvState. + self.add_tool( + "view_driver_state", + spec["view_driver_state"], + partial(self._state.view, image_slots=self._VIEW_IMAGE_SLOTS), + ) + self.add_tool( + "back_project", + spec["back_project"], + partial(franka_tools.back_project, state=self._state), + ) for name in self._DRIVER_READERS: self.add_tool(name, spec[name], self._make_driver_reader(name)) for name in self._PRIMITIVE_TOOLS: @@ -71,31 +84,22 @@ def _step(self, name: str, **kwargs) -> dict: elapsed = round(time.time() - t0, 2) result_dict = result if isinstance(result, dict) else {"value": result} - self._next_step += 1 - step_idx = self._next_step + step_idx = self._state.next_step_idx franka_tools.dump_state( self._driver, - str(get_output_dir()), + self._state, step_idx=step_idx, log={"command": command, "result": result_dict, "elapsed_s": elapsed}, ) - out = franka_tools.view_driver_state(step_idx) + out = self._state.view(step_idx, image_slots=self._VIEW_IMAGE_SLOTS) out["agent_elapsed_s"] = elapsed return out def init_driver_clean(self, *, env: Any) -> None: - out_dir = get_output_dir() - out_dir.mkdir(parents=True, exist_ok=True) - images_dir = out_dir / "images" - if images_dir.exists(): - shutil.rmtree(images_dir) - states_file = out_dir / "states.json" - if states_file.exists(): - states_file.unlink() - + self._state.reset(wipe_streams=self._WIPE_STREAMS) driver = franka_tools.FrankaPrimitives(env=env) driver.reset() - franka_tools.dump_state(driver, str(out_dir), step_idx=0, log=None) + franka_tools.dump_state(driver, self._state, step_idx=0, log=None) self._driver = driver def close(self) -> None: diff --git a/robots/franka/tools.py b/robots/franka/tools.py index 3db62423..31327bf4 100644 --- a/robots/franka/tools.py +++ b/robots/franka/tools.py @@ -1,16 +1,15 @@ """Franka tool implementation for the agent-side toolkit.""" from __future__ import annotations -import json -import os import time from typing import Any -import imageio.v2 as imageio import numpy as np from robots.franka.env_client import FrankaEnvClient -from rpent.utils.logging import get_logger, get_output_dir +from rpent.tools.common import robust_surface_centroid +from rpent.tools.state import EnvState, StepRecord +from rpent.utils.logging import get_logger logger = get_logger("franka") @@ -18,7 +17,6 @@ _MAX_YAW_DELTA_DEG = 30.0 _MAX_OBSERVE_DELAY_S = 5.0 _BACKPROJECT_RADIUS = 6 -_DEPTH_BAND_M = 0.02 def _to_list(value) -> list: @@ -38,66 +36,9 @@ def _to_scalar(value) -> Any: return value -def _states_path(output_dir: str) -> str: - return os.path.join(output_dir, "states.json") - - -def _append_state(output_dir: str, blob: dict) -> None: - path = _states_path(output_dir) - states: list = [] - if os.path.exists(path): - with open(path) as f: - states = json.load(f) - states.append(blob) - with open(path, "w") as f: - json.dump(states, f, indent=2, default=str) - - -def _load_states() -> list: - path = _states_path(str(get_output_dir())) - if not os.path.exists(path): - return [] - with open(path) as f: - return json.load(f) - - -def _latest_step() -> int | None: - states = _load_states() - if not states: - return None - return int(states[-1]["step_idx"]) - - -def _load_step(step_idx: int) -> dict: - for state in _load_states(): - if int(state["step_idx"]) == step_idx: - return state - raise KeyError(f"step {step_idx} not present in states.json") - - -def _load_image(step_idx: int, camera: str) -> bytes | None: - path = os.path.join(str(get_output_dir()), "images", f"{camera}_{step_idx:02d}.png") - if not os.path.exists(path): - return None - with open(path, "rb") as f: - return f.read() - - -def _load_depth(step_idx: int, camera: str) -> np.ndarray: - path = os.path.join(str(get_output_dir()), "depths", f"{camera}_{step_idx:02d}.npy") - return np.load(path) - - -def _backproject_points(K, rows, cols, depths) -> np.ndarray: - K = np.asarray(K, dtype=np.float64) - rows = np.asarray(rows, dtype=np.float64) - cols = np.asarray(cols, dtype=np.float64) - depths = np.asarray(depths, dtype=np.float64) - fx, fy = K[0, 0], K[1, 1] - cx, cy = K[0, 2], K[1, 2] - return np.stack( - [(cols - cx) * depths / fx, (rows - cy) * depths / fy, depths], axis=1 - ) +# State-trace I/O, image/depth layout, and the robust back-projection math now +# live in :class:`rpent.tools.state.EnvState` (one owner per run). The thin +# wrappers below delegate to it. class FrankaPrimitives: @@ -240,89 +181,45 @@ def _refresh(self) -> None: def dump_state( driver: FrankaPrimitives, - output_dir: str, + state: EnvState, step_idx: int, log: dict | None = None, -) -> dict: - """Dump current camera frames and proprioceptive state to the run dir.""" - images_dir = os.path.join(output_dir, "images") - depths_dir = os.path.join(output_dir, "depths") - os.makedirs(images_dir, exist_ok=True) - os.makedirs(depths_dir, exist_ok=True) +) -> StepRecord: + """Dump camera frames + proprioceptive state via the EnvState owner. + + Maps the Franka cameras onto the LIBERO stream layout: scene -> ``image``/ + ``depth`` (the primary view) and wrist -> ``image_wrist``/``depth_wrist``. + """ + image_streams = {"scene": "image", "wrist": "image_wrist"} + depth_streams = {"scene": "depth", "wrist": "depth_wrist"} - saved: dict[str, str] = {} + saved_frames: list[str] = [] for camera, frame in driver.latest_frames().items(): - arr = np.asarray(frame) - if arr.dtype != np.uint8: - arr = arr.astype(np.uint8) - out_path = os.path.join(images_dir, f"{camera}_{step_idx:02d}.png") - try: - imageio.imwrite(out_path, arr) - saved[camera] = out_path - except Exception as exc: - logger.warning("frame dump failed for camera %s: %s", camera, exc) + stream = image_streams.get(camera, camera) + if state.save_image(step_idx, stream, frame): + saved_frames.append(stream) - saved_depths: dict[str, str] = {} + saved_depths: list[str] = [] for camera, depth in driver.latest_depths().items(): - out_path = os.path.join(depths_dir, f"{camera}_{step_idx:02d}.npy") - try: - np.save(out_path, np.asarray(depth, dtype=np.float32)) - saved_depths[camera] = out_path - except Exception as exc: - logger.warning("depth dump failed for camera %s: %s", camera, exc) - - blob: dict[str, Any] = { - "step_idx": step_idx, - "state": driver.get_state(), - "frames": sorted(saved), - "depth": sorted(saved_depths), - "camera_meta": driver.latest_camera_meta(), - } - if log is not None: - blob["command"] = log.get("command") - blob["result"] = log.get("result") - blob["elapsed_s"] = log.get("elapsed_s") - _append_state(output_dir, blob) - return blob - - -def view_driver_state(step: int | None = None) -> dict: - """Read a dumped step and embed the scene/wrist camera images.""" - latest = _latest_step() - if latest is None: - return {"error": "no driver state entries; driver not ready"} - step_idx = latest if step is None else int(step) - try: - data = _load_step(step_idx) - except Exception as exc: - return {"error": f"step {step_idx} not present in driver state trace: {exc}"} - - out: dict[str, Any] = { - "step": step_idx, - "state": data.get("state", {}), - "frames": data.get("frames", []), - "depth": data.get("depth", []), - "camera_meta": { - name: { - key: value - for key, value in meta.items() - if key not in {"K", "T_base_cam"} - } - for name, meta in (data.get("camera_meta") or {}).items() - }, - "log": { - "command": data.get("command"), - "result": data.get("result"), - "elapsed_s": data.get("elapsed_s"), - }, - } - scene = _load_image(step_idx, "scene") - wrist = _load_image(step_idx, "wrist") - if scene: - out["_image_bytes"] = scene - if wrist: - out["_image_cam_bytes"] = wrist - return out + stream = depth_streams.get(camera, camera) + if state.save_depth(step_idx, stream, depth): + saved_depths.append(stream) + + record = StepRecord( + step_idx=step_idx, + state=driver.get_state(), + frames=sorted(saved_frames), + depth=sorted(saved_depths), + camera_meta=driver.latest_camera_meta(), + command=log.get("command") if log else None, + result=log.get("result") if log else None, + elapsed_s=log.get("elapsed_s") if log else None, + ) + return state.append(record) + + +# view_driver_state now lives on EnvState.view (bound by the toolkit with the +# scene/wrist image slots); nothing module-level is needed here. def back_project( @@ -330,85 +227,65 @@ def back_project( col: int, step: int | None = None, camera: str = "wrist", - radius: int = _BACKPROJECT_RADIUS, + radius: int | None = _BACKPROJECT_RADIUS, + *, + state: EnvState, ) -> dict: - """Backproject a saved RGB-D pixel into camera and robot-base coordinates.""" - step_idx = _latest_step() if step is None else int(step) - if step_idx is None: + """Back-project a saved RGB-D pixel into camera and robot-base coordinates. + + Uses :func:`rpent.tools.common.robust_surface_centroid` (median of the + dominant surface in a window around the pixel) so oblique-view depth noise + does not become ~cm lateral error. Returns robot-base ``xyz`` in + ``panda_link0`` when the camera has ``T_base_cam`` for the step; otherwise + ``xyz_cam`` plus a warning. + """ + nn = state.latest_step if step is None else int(step) + if nn is None: return {"error": "no steps available"} try: - data = _load_step(step_idx) + rec = state.get(nn) except Exception as exc: - return {"error": f"step {step_idx} not present in driver state trace: {exc}"} + return {"error": f"step {nn} not present in driver state trace: {exc}"} camera = str(camera or "wrist") - meta = (data.get("camera_meta") or {}).get(camera) + meta = (rec.camera_meta or {}).get(camera) if not meta: return { - "error": f"camera {camera!r} has no metadata at step {step_idx}", - "available_cameras": sorted((data.get("camera_meta") or {}).keys()), + "error": f"camera {camera!r} has no metadata at step {nn}", + "available_cameras": sorted((rec.camera_meta or {}).keys()), } + depth_stream = "depth" if camera == "scene" else "depth_wrist" try: - depth = _load_depth(step_idx, camera) + depth = state.load_depth(nn, depth_stream) except Exception as exc: - return {"error": f"depth for camera {camera!r} step {step_idx} not found: {exc}"} - - row, col = int(row), int(col) - radius = max(0, int(_BACKPROJECT_RADIUS if radius is None else radius)) - h, w = depth.shape[:2] - if not (0 <= row < h and 0 <= col < w): - return {"error": f"pixel ({row},{col}) out of bounds for {camera} image {h}x{w}"} - - r0, r1 = max(0, row - radius), min(h, row + radius + 1) - c0, c1 = max(0, col - radius), min(w, col + radius + 1) - rr, cc = np.mgrid[r0:r1, c0:c1] - zz = depth[r0:r1, c0:c1].reshape(-1).astype(np.float64) - rr = rr.reshape(-1).astype(np.float64) - cc = cc.reshape(-1).astype(np.float64) - valid = np.isfinite(zz) & (zz > 0) - if not np.any(valid): - return {"error": f"no valid depth near ({row},{col}) in {camera}; pick another pixel"} - zz, rr, cc = zz[valid], rr[valid], cc[valid] - z_med = float(np.median(zz)) - surface = np.abs(zz - z_med) <= _DEPTH_BAND_M - zz, rr, cc = zz[surface], rr[surface], cc[surface] - if zz.size == 0: - return {"error": f"no dominant surface depth near ({row},{col}) in {camera}"} - - pts_cam = _backproject_points(meta["K"], rr, cc, zz) - p_cam = np.median(pts_cam, axis=0) - out: dict[str, Any] = { - "step": step_idx, - "camera": camera, - "pixel": [row, col], - "radius": radius, - "n_points": int(pts_cam.shape[0]), - "depth_m": round(z_med, 4), - "xyz_cam": [round(float(v), 4) for v in p_cam], - "camera_frame": meta.get("frame", f"{camera}_camera"), - "calibrated": bool(meta.get("calibrated")), - "calibration_kind": meta.get("calibration_kind"), - } - - T_base_cam = meta.get("T_base_cam") - if T_base_cam is not None: - T = np.asarray(T_base_cam, dtype=np.float64) - pts_base = pts_cam @ T[:3, :3].T + T[:3, 3] - p_base = np.median(pts_base, axis=0) - out["xyz"] = [round(float(v), 4) for v in p_base] - out["frame"] = "panda_link0" - out["xy_spread_m"] = round( - float(np.hypot(*pts_base[:, :2].std(axis=0))), 4 - ) + return {"error": f"depth for camera {camera!r} step {nn} not found: {exc}"} + + res = robust_surface_centroid( + depth, + meta["K"], + meta.get("T_base_cam"), + row, + col, + radius=_BACKPROJECT_RADIUS if radius is None else radius, + ) + if "error" in res: + return res + + res["step"] = nn + res["camera"] = camera + res["camera_frame"] = meta.get("frame", f"{camera}_camera") + res["calibrated"] = bool(meta.get("calibrated")) + res["calibration_kind"] = meta.get("calibration_kind") + if meta.get("T_base_cam") is not None: + res["frame"] = "panda_link0" else: - out["frame"] = meta.get("frame", f"{camera}_camera") - out["xy_spread_m"] = round(float(np.hypot(*pts_cam[:, :2].std(axis=0))), 4) - out["note"] = ( + res["frame"] = meta.get("frame", f"{camera}_camera") + res["note"] = ( f"camera {camera!r} is not calibrated to panda_link0 for this step; " "do not use xyz_cam as a robot target. Add T_base_cam for a fixed " "camera or T_tcp_cam for a wrist camera." ) - return out + return res TOOLS_SPEC: list[dict[str, Any]] = [ diff --git a/robots/lerobot/toolkit.py b/robots/lerobot/toolkit.py index 24e6fd6d..119da41d 100644 --- a/robots/lerobot/toolkit.py +++ b/robots/lerobot/toolkit.py @@ -6,12 +6,12 @@ """ from __future__ import annotations -import shutil import time from functools import partial from typing import Any from robots.lerobot import tools as lerobot_tools +from rpent.tools.state import EnvState from rpent.tools.toolkit import Toolkit from rpent.utils.logging import get_output_dir @@ -19,11 +19,8 @@ class LerobotToolkit(Toolkit): """Toolkit for the LeRobot SO101 environment.""" - # Stateless reader tools bound directly to module-level functions. - _STATELESS_TOOLS = ( - "view_driver_state", - "back_project", - ) + _WIPE_STREAMS = ("image", "image_arm", "depth") + _VIEW_IMAGE_SLOTS = {"_image_bytes": "image", "_image_cam_bytes": "image_arm"} # Read-only tools backed by a live driver call (no state dump). These query # the robot/scene directly: forward kinematics + scene camera calibration. _DRIVER_READERS = ( @@ -48,8 +45,8 @@ def __init__( video_path: str | None = None, dashboard: Any = None, ) -> None: - super().__init__(dashboard=dashboard) - self._next_step: int = 0 + state = EnvState(get_output_dir()) + super().__init__(dashboard=dashboard, state=state) self._video_path: str | None = video_path self.init_driver_clean(env=env, model=model) self._register_tools() @@ -59,8 +56,16 @@ def __init__( # ------------------------------------------------------------------ def _register_tools(self) -> None: spec = self._SPECS - for name in self._STATELESS_TOOLS: - self.add_tool(name, spec[name], getattr(lerobot_tools, name)) + self.add_tool( + "view_driver_state", + spec["view_driver_state"], + partial(self._state.view, image_slots=self._VIEW_IMAGE_SLOTS), + ) + self.add_tool( + "back_project", + spec["back_project"], + partial(lerobot_tools.back_project, state=self._state), + ) for name in self._DRIVER_READERS: self.add_tool(name, spec[name], self._make_driver_reader(name)) for name in self._PRIMITIVE_TOOLS: @@ -84,32 +89,23 @@ def _step(self, name: str, **kwargs) -> dict: result_dict = result if isinstance(result, dict) else {"value": result} - self._next_step += 1 - step_idx = self._next_step + step_idx = self._state.next_step_idx lerobot_tools.dump_state( self._driver, - str(get_output_dir()), + self._state, step_idx=step_idx, log={"command": command, "result": result_dict, "elapsed_s": elapsed}, ) - out = lerobot_tools.view_driver_state(step_idx) + out = self._state.view(step_idx, image_slots=self._VIEW_IMAGE_SLOTS) out["agent_elapsed_s"] = elapsed return out def init_driver_clean(self, *, env: Any, model: Any | None = None) -> None: """Wipe stale run artifacts, build the primitive driver, dump step 0.""" - out_dir = get_output_dir() - out_dir.mkdir(parents=True, exist_ok=True) - images_dir = out_dir / "images" - if images_dir.exists(): - shutil.rmtree(images_dir) - states_file = out_dir / "states.json" - if states_file.exists(): - states_file.unlink() - + self._state.reset(wipe_streams=self._WIPE_STREAMS) driver = lerobot_tools.LerobotPrimitives(env=env, model=model) driver.reset() - lerobot_tools.dump_state(driver, str(out_dir), step_idx=0, log=None) + lerobot_tools.dump_state(driver, self._state, step_idx=0, log=None) self._driver = driver diff --git a/robots/lerobot/tools.py b/robots/lerobot/tools.py index 48ce0d84..c551dd77 100644 --- a/robots/lerobot/tools.py +++ b/robots/lerobot/tools.py @@ -5,8 +5,8 @@ * :class:`LerobotPrimitives` — the primitive driver the toolkit owns. Holds the env client (and an optional policy/VLA model) plus per-run state, and exposes one method per primitive tool. -* per-step state dump (:func:`dump_state`) + stateless reader tools - (:func:`view_driver_state`). +* per-step state dump (:func:`dump_state`) plus reader tools bound to the + run's :class:`rpent.tools.state.EnvState`. * :data:`TOOLS_SPEC` — Anthropic-shaped tool schemas. NOTE: this is a scaffold. The concrete robot primitives (move / grasp / @@ -16,25 +16,18 @@ """ from __future__ import annotations -import json -import os from typing import Any -import imageio.v2 as imageio import numpy as np from robots.lerobot.env_client import LerobotEnvClient -from rpent.utils.logging import get_logger, get_output_dir +from rpent.tools.common import robust_surface_centroid +from rpent.tools.state import EnvState, StepRecord +from rpent.utils.logging import get_logger logger = get_logger("lerobot") -# back_project robust-centroid params: it back-projects every valid pixel in a -# (2*radius+1) square window and keeps those whose depth is within -# _DEPTH_BAND_M of the window median (the dominant / object surface), then -# takes the world-median. This tames the scene camera's oblique-view depth -# noise, which otherwise turns per-pixel depth error into ~cm lateral error. _BACKPROJECT_RADIUS = 6 -_DEPTH_BAND_M = 0.02 def _to_list(x) -> list: @@ -177,174 +170,60 @@ def _refresh(self) -> None: logger.warning("obs refresh failed: %s", e) -# --------------------------------------------------------------------------- -# states.json + image helpers -# --------------------------------------------------------------------------- - - -def _states_path(output_dir: str) -> str: - return os.path.join(output_dir, "states.json") - - -def _append_state(output_dir: str, blob: dict) -> None: - path = _states_path(output_dir) - states: list = [] - if os.path.exists(path): - with open(path) as f: - states = json.load(f) - states.append(blob) - with open(path, "w") as f: - json.dump(states, f, indent=2, default=str) - - -def _load_states() -> list: - path = _states_path(str(get_output_dir())) - if not os.path.exists(path): - return [] - with open(path) as f: - return json.load(f) - - -def _latest_step() -> int | None: - states = _load_states() - if not states: - return None - return int(states[-1]["step_idx"]) - - -def _load_step(nn: int) -> dict: - for s in _load_states(): - if int(s["step_idx"]) == nn: - return s - raise KeyError(f"step {nn} not present in states.json") - - -def _load_image(nn: int, cam: str) -> bytes | None: - path = os.path.join(str(get_output_dir()), "images", f"{cam}_{nn:02d}.png") - if not os.path.exists(path): - return None - with open(path, "rb") as f: - return f.read() - - -def _load_depth(nn: int) -> np.ndarray: - path = os.path.join(str(get_output_dir()), "depths", f"scene_{nn:02d}.npy") - return np.load(path) - - -def _load_camera_meta() -> dict: - path = os.path.join(str(get_output_dir()), "camera_meta.json") - with open(path) as f: - return json.load(f) - - def dump_state( driver: LerobotPrimitives, - output_dir: str, + state: EnvState, step_idx: int, log: dict | None = None, -) -> dict: +) -> StepRecord: """Dump the camera frames, scene depth, and proprioceptive state. Writes per step: - - ``/images/_NN.png`` (arm + scene color) - - ``/depths/scene_NN.npy`` (metric depth, aligned to color) - and once: - - ``/camera_meta.json`` (scene K, depth scale, T_base_cam) - then appends the step blob (state incl. ``ee_pose_base`` + optional command - log) to ``/states.json``. - """ - images_dir = os.path.join(output_dir, "images") - depths_dir = os.path.join(output_dir, "depths") - os.makedirs(images_dir, exist_ok=True) - os.makedirs(depths_dir, exist_ok=True) - - saved: dict[str, str] = {} - for cam, frame in driver.latest_frames().items(): - arr = np.asarray(frame) - if arr.dtype != np.uint8: - arr = arr.astype(np.uint8) - out_path = os.path.join(images_dir, f"{cam}_{step_idx:02d}.png") - try: - imageio.imwrite(out_path, arr) - saved[cam] = out_path - except Exception as e: - logger.warning("frame dump failed for cam %s: %s", cam, e) - - depth = driver.latest_depth() - if depth is not None: - try: - np.save(os.path.join(depths_dir, f"scene_{step_idx:02d}.npy"), - depth.astype(np.float32)) - except Exception as e: - logger.warning("depth dump failed: %s", e) - - # Scene camera calibration is static — fetch + dump once. - meta_path = os.path.join(output_dir, "camera_meta.json") - if not os.path.exists(meta_path): - meta = driver.get_scene_camera_meta() - if isinstance(meta, dict) and "error" not in meta: - with open(meta_path, "w") as f: - json.dump(meta, f, indent=2, default=str) - - blob: dict = { - "step_idx": step_idx, - "state": driver.get_state(), - "frames": sorted(saved), - } - if log is not None: - blob["command"] = log.get("command") - blob["result"] = log.get("result") - blob["elapsed_s"] = log.get("elapsed_s") - _append_state(output_dir, blob) - return blob - - -# --------------------------------------------------------------------------- -# Stateless reader tools -# --------------------------------------------------------------------------- + - ``images/image_NN.png`` (scene color) + - ``images_arm/image_arm_NN.png`` (arm color) + - ``depths/depth_NN.npy`` (scene metric depth, aligned to scene color) - -def view_driver_state(step: int | None = None) -> dict: - """Read step NN from ``states.json`` + the matching camera PNGs. - - Returns the proprioceptive state and embeds the scene/arm camera frames - as multimodal image content blocks (via the ``_image_bytes`` / - ``_image_cam_bytes`` conventions consumed by ``ToolResult``). + Scene calibration is stored on the step record. ``EnvState`` atomically + appends that record to ``states.json`` and advances the trace counter. """ - latest = _latest_step() - if latest is None: - return {"error": "no driver state entries; driver not ready"} - nn = latest if step is None else int(step) - try: - data = _load_step(nn) - except Exception as e: - return {"error": f"step {nn} not present in driver state trace: {e}"} - - out: dict = { - "step": nn, - "state": data.get("state", {}), - "log": { - "command": data.get("command"), - "result": data.get("result"), - "elapsed_s": data.get("elapsed_s"), - }, - } - # Map the two cameras onto the two image slots ToolResult understands. - scene = _load_image(nn, "scene") - arm = _load_image(nn, "arm") - if scene: - out["_image_bytes"] = scene - if arm: - out["_image_cam_bytes"] = arm - return out + image_streams = {"scene": "image", "arm": "image_arm"} + saved_frames: list[str] = [] + for camera, frame in driver.latest_frames().items(): + stream = image_streams.get(camera, f"image_{camera}") + if state.save_image(step_idx, stream, frame): + saved_frames.append(stream) + + saved_depths: list[str] = [] + depth = driver.latest_depth() + if depth is not None and state.save_depth(step_idx, "depth", depth): + saved_depths.append("depth") + + scene_meta = driver.get_scene_camera_meta() + camera_meta = ( + {"scene": scene_meta} + if isinstance(scene_meta, dict) and "error" not in scene_meta + else {} + ) + record = StepRecord( + step_idx=step_idx, + state=driver.get_state(), + frames=sorted(saved_frames), + depth=sorted(saved_depths), + camera_meta=camera_meta, + command=log.get("command") if log else None, + result=log.get("result") if log else None, + elapsed_s=log.get("elapsed_s") if log else None, + ) + return state.append(record) def back_project( row: int, col: int, step: int | None = None, - radius: int = _BACKPROJECT_RADIUS, + radius: int | None = _BACKPROJECT_RADIUS, + *, + state: EnvState, ) -> dict: """Backproject a scene-camera pixel neighborhood to a robust world point. @@ -352,77 +231,48 @@ def back_project( pixel's depth error becomes a large lateral error. Instead of trusting one pixel, this back-projects EVERY valid pixel in a ``(2*radius+1)`` square window around ``(row, col)``, keeps those on the dominant surface (depth - within ``_DEPTH_BAND_M`` of the window median -- rejecting background / - table / dropouts), and returns the MEDIAN world ``xyz`` of that surface: a + within a narrow band of the window median, rejecting background, table, + and dropouts), and returns the MEDIAN world ``xyz`` of that surface: a stable object centroid rather than one face pixel. Use ``radius=0`` for the old single-pixel behavior. - Pick ``(row, col)`` on the scene color image ``images/scene_NN.png``; depth - is aligned to it (``depths/scene_NN.npy``). Uses ``camera_meta.json`` (K + - ``T_base_cam``). Returns base/world ``xyz`` when calibrated, else the - camera-frame ``xyz_cam`` with a note. Also reports ``n_points`` (surface - pixels used) and ``xy_spread_m`` (their world-xy stdev) as a quality gauge. + Pick ``(row, col)`` on ``images/image_NN.png``; depth is aligned to it in + ``depths/depth_NN.npy``. Returns base/world ``xyz`` when calibrated, else + camera-frame ``xyz_cam`` with a note. """ - try: - meta = _load_camera_meta() - except Exception as e: - return {"error": f"camera_meta.json not found: {e}"} - nn = _latest_step() if step is None else int(step) + nn = state.latest_step if step is None else int(step) if nn is None: return {"error": "no steps available"} try: - depth = _load_depth(nn) - except Exception as e: - return {"error": f"depth for step {nn} not found: {e}"} - - row, col = int(row), int(col) - radius = max(0, int(radius)) - h, w = depth.shape[:2] - if not (0 <= row < h and 0 <= col < w): - return {"error": f"pixel ({row},{col}) out of bounds; image is {h}x{w}"} - - # Gather the window, keep valid depths, then restrict to the dominant - # surface (depths within a band of the window median) so background / table - # pixels and dropouts don't drag the centroid. - r0, r1 = max(0, row - radius), min(h, row + radius + 1) - c0, c1 = max(0, col - radius), min(w, col + radius + 1) - rr, cc = np.mgrid[r0:r1, c0:c1] - zz = depth[r0:r1, c0:c1].reshape(-1).astype(np.float64) - rr = rr.reshape(-1).astype(np.float64) - cc = cc.reshape(-1).astype(np.float64) - valid = np.isfinite(zz) & (zz > 0) - if not np.any(valid): - return {"error": f"no valid depth near ({row},{col}); pick another pixel"} - zz, rr, cc = zz[valid], rr[valid], cc[valid] - z_med = float(np.median(zz)) - surf = np.abs(zz - z_med) <= _DEPTH_BAND_M - zz, rr, cc = zz[surf], rr[surf], cc[surf] - - K = np.asarray(meta["K"], dtype=np.float64) - fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2] - pts_cam = np.stack( - [(cc - cx) * zz / fx, (rr - cy) * zz / fy, zz], axis=1 - ) # (N, 3) - p_cam = np.median(pts_cam, axis=0) - - out: dict = { - "pixel": [row, col], - "radius": radius, - "n_points": int(pts_cam.shape[0]), - "depth_m": round(z_med, 4), - "xyz_cam": [round(float(v), 4) for v in p_cam], - "frame": "scene_cam", - } - T = meta.get("T_base_cam") - if T is not None: - T = np.asarray(T, dtype=np.float64) - pts_base = pts_cam @ T[:3, :3].T + T[:3, 3] - p_base = np.median(pts_base, axis=0) - out["xyz"] = [round(float(v), 4) for v in p_base] - out["xy_spread_m"] = round(float(np.hypot(*pts_base[:, :2].std(axis=0))), 4) + record = state.get(nn) + except Exception as exc: + return {"error": f"step {nn} not present in driver state trace: {exc}"} + meta = (record.camera_meta or {}).get("scene") + if not meta: + return {"error": f"scene camera metadata not recorded for step {nn}"} + try: + depth = state.load_depth(nn, "depth") + except Exception as exc: + return {"error": f"depth for step {nn} not found: {exc}"} + + out = robust_surface_centroid( + depth, + meta["K"], + meta.get("T_base_cam"), + row, + col, + radius=_BACKPROJECT_RADIUS if radius is None else radius, + ) + if "error" in out: + return out + + out["step"] = nn + out["camera"] = "scene" + out["camera_frame"] = meta.get("frame", "scene_cam") + if meta.get("T_base_cam") is not None: out["frame"] = "base_link" else: - out["xy_spread_m"] = round(float(np.hypot(*pts_cam[:, :2].std(axis=0))), 4) + out["frame"] = meta.get("frame", "scene_cam") out["note"] = ( "scene camera not calibrated (no T_base_cam); returning camera-frame " "xyz only. Run robots/lerobot/calibrate_scene_cam.py." @@ -438,9 +288,10 @@ def back_project( { "name": "view_driver_state", "description": ( - "Read step NN from `states.json` + the matching camera PNGs in " - "{output_dir}/images. If step is null, returns the latest entry. " - "Embeds the scene and arm camera frames as image content blocks." + "Read step NN from `states.json` plus the matching scene image " + "`images/image_NN.png` and arm image " + "`images_arm/image_arm_NN.png`. If step is null, returns the " + "latest entry. Embeds both camera frames as image content blocks." ), "input_schema": { "type": "object", diff --git a/robots/libero/toolkit.py b/robots/libero/toolkit.py index 475c1e74..71345caa 100644 --- a/robots/libero/toolkit.py +++ b/robots/libero/toolkit.py @@ -5,13 +5,13 @@ """ from __future__ import annotations -import shutil import time from functools import partial from typing import Any from robots.libero import tools as libero_tools from rpent.dashboard.events import DashboardEventSink, ToolResultEvent +from rpent.tools.state import EnvState from rpent.tools.toolkit import ToolCancelled, Toolkit from rpent.utils.logging import get_logger, get_output_dir @@ -19,6 +19,22 @@ class LiberoToolkit(Toolkit): """Toolkit for the LIBERO environment.""" + _WIPE_STREAMS = ( + "image", + "image_cam", + "depth", + "world", + "image_wrist", + "depth_wrist", + "world_wrist", + "wrist_meta", + "image_cam_hi", + "world_hi", + "image_wrist_hi", + "world_wrist_hi", + "segments", + "action_videos", + ) # Tool schemas keyed by name (built once from the canonical ordered list # in libero_tools.TOOLS_SPEC) so each tool registers with its own spec. _SPECS = {spec["name"]: spec for spec in libero_tools.TOOLS_SPEC} @@ -30,8 +46,8 @@ def __init__( dashboard_events: DashboardEventSink, video_path: str | None = None, ) -> None: - super().__init__(dashboard_events=dashboard_events) - self._next_step: int = 0 + state = EnvState(get_output_dir()) + super().__init__(dashboard_events=dashboard_events, state=state) self._video_path: str | None = video_path self.init_primitives_clean(primitives_kwargs=primitives_kwargs) self._register_libero_tools() @@ -44,10 +60,14 @@ def _register_libero_tools(self) -> None: # Inspection tools do not advance environment state. Most are stateless # module functions; segment is bound to the primitives-owned SAM3 client. inspection_handlers = { - "view_driver_state": libero_tools.view_driver_state, - "view_camera_meta": libero_tools.view_camera_meta, - "back_project": libero_tools.back_project, - "segment": self._primitives.segment, + "view_driver_state": partial( + libero_tools.view_driver_state, state=self._state + ), + "view_camera_meta": partial( + libero_tools.view_camera_meta, state=self._state + ), + "back_project": partial(libero_tools.back_project, state=self._state), + "segment": partial(self._primitives.segment, state=self._state), } for name, handler in inspection_handlers.items(): self.add_tool(name, specs[name], handler) @@ -79,11 +99,11 @@ def _step(self, name: str, **kwargs) -> dict: else: result_dict = {"value": result} - self._next_step += 1 - step_idx = self._next_step - output_dir = get_output_dir() + step_idx = self._state.next_step_idx if self._dashboard_events.enabled: - video_dir = libero_tools.artifact_path(output_dir, "action_videos") + video_dir = libero_tools.artifact_path( + self._state.output_dir, "action_videos" + ) video_path = video_dir / f"step_{step_idx:02d}_{name}.mp4" try: self._primitives.save_frame_slice(start_frame, str(video_path), fps=20) @@ -93,11 +113,11 @@ def _step(self, name: str, **kwargs) -> dict: ) libero_tools.dump_state( self._primitives, - str(output_dir), + self._state, step_idx=step_idx, log={"command": command, "result": result_dict, "elapsed_s": elapsed}, ) - out = libero_tools.view_driver_state(step_idx) + out = libero_tools.view_driver_state(step_idx, state=self._state) out["agent_elapsed_s"] = elapsed if result_dict.get("interrupted"): out.update(result_dict) @@ -109,14 +129,9 @@ def init_primitives_clean( primitives_kwargs: dict[str, Any], ) -> None: """Wipe stale run artifacts, build the LiberoPrimitives, dump step 0.""" - out_dir = get_output_dir() - out_dir.mkdir(parents=True, exist_ok=True) - for sub in libero_tools.ARTIFACT_DIRECTORIES: - target = out_dir / sub - if target.exists(): - shutil.rmtree(target) + self._state.reset(wipe_streams=self._WIPE_STREAMS) + out_dir = self._state.output_dir for target in ( - libero_tools.artifact_path(out_dir, "states"), libero_tools.artifact_path(out_dir, "metadata", camera="agentview", resolution="low"), libero_tools.artifact_path(out_dir, "episode_video"), ): @@ -129,11 +144,11 @@ def init_primitives_clean( ) primitives.reset() primitives.start_recording() - libero_tools.dump_state(primitives, str(out_dir), step_idx=0, log=None) + libero_tools.dump_state(primitives, self._state, step_idx=0, log=None) self._dashboard_events.emit( ToolResultEvent( name="view_driver_state", - result=libero_tools.view_driver_state(0), + result=libero_tools.view_driver_state(0, state=self._state), ) ) @@ -155,4 +170,4 @@ def close(self) -> None: def write_recipe(self, recipe_tag: str) -> str: """Write the LIBERO recipe JSONL from the dumped state trace.""" - return libero_tools.write_recipe_from_states(str(get_output_dir()), recipe_tag) + return libero_tools.write_recipe_from_states(self._state, recipe_tag) diff --git a/robots/libero/tools.py b/robots/libero/tools.py index 44145227..6b95f6cf 100644 --- a/robots/libero/tools.py +++ b/robots/libero/tools.py @@ -11,7 +11,8 @@ import numpy as np from robots.libero.env_client import LiberoEnvClient -from rpent.utils.logging import get_logger, get_output_dir +from rpent.tools.state import EnvState, StepRecord +from rpent.utils.logging import get_logger from rpent.utils.sam3_client import Sam3Client from rpent.utils.vla_client import VLAClient @@ -662,6 +663,8 @@ def segment( step: int | None = None, point: list[int] | None = None, min_score: float = 0.2, + *, + state: EnvState, ) -> dict: """Call SAM3 on an existing image artifact without advancing the env. @@ -669,7 +672,7 @@ def segment( artifacts. Errors are structured so the agent can continue with image inspection and ``back_project``. """ - nn = _latest_step() if step is None else int(step) + nn = state.latest_step if step is None else int(step) if nn is None: return {"error": "no state entries; cannot select segment image"} @@ -681,7 +684,7 @@ def segment( return {"error": "segment needs exactly one of prompt or point"} try: image_path, world_path, artifact_pairs = _select_segment_artifacts( - nn, camera + state, nn, camera ) except ValueError as e: return {"error": str(e)} @@ -720,7 +723,7 @@ def segment( "fallback": "Use manual visual localization and back_project.", } - out_dir = get_output_dir() + out_dir = state.output_dir segment_path, overlay_candidate_path, segment_index = ( _next_segment_artifact_paths(out_dir, nn) ) @@ -785,53 +788,25 @@ def segment( return result -# --------------------------------------------------------------------------- -# State artifacts -# --------------------------------------------------------------------------- - - -def _append_state(output_dir: str, blob: dict) -> None: - """Append *blob* to ``/states.json`` atomically.""" - path = artifact_path(output_dir, "states") - tmp_path = path.parent / f"{path.name}.tmp" - if path.exists(): - with open(path) as f: - states = json.load(f) - else: - states = [] - states.append(blob) - with open(tmp_path, "w") as f: - json.dump(states, f, indent=2) - os.replace(tmp_path, path) - - -def write_recipe_from_states(output_dir: str, recipe_tag: str) -> str: +def write_recipe_from_states(state: EnvState, recipe_tag: str) -> str: """Find a command sequence that gets ``libero_terminated=True``. Export non-error LIBERO primitive commands from ``states.json`` and successful segment calls from ``segments/segment_*.json``. """ - states_path = artifact_path(output_dir, "states") - if states_path.exists(): - with open(states_path) as f: - states = json.load(f) - else: - states = [] - command_events = [] - for step_idx, entry in enumerate(states): - if not entry: - continue - command = entry.get("command") + for record in state.records(): + command = record.command if command is None: continue if command.get("action") not in PRIMITIVE_TOOL_NAMES: continue - result = entry.get("result") + result = record.result if isinstance(result, dict) and result.get("error"): continue - command_events.append(((step_idx, -1), command)) + command_events.append(((record.step_idx, -1), command)) + output_dir = state.output_dir for artifact in artifact_path(output_dir, "segments").glob("segment_*.json"): with artifact.open() as f: segment = json.load(f) @@ -853,7 +828,7 @@ def write_recipe_from_states(output_dir: str, recipe_tag: str) -> str: event_order = (source_step, int(segment["segment_index"])) command_events.append((event_order, command)) - recipe_path = os.path.join(output_dir, f"recipe_{recipe_tag}.jsonl") + recipe_path = os.path.join(state.output_dir, f"recipe_{recipe_tag}.jsonl") tmp_path = recipe_path + ".tmp" command_events.sort(key=lambda event: event[0]) with open(tmp_path, "w") as f: @@ -889,8 +864,8 @@ def _world_from_depth(depth_metric: np.ndarray, camera_meta: dict) -> np.ndarray return (camera_points @ extrinsic.T)[..., :3] -def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, - log: dict | None = None) -> dict: +def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, + log: dict | None = None) -> StepRecord: """Dump state snapshot, images, and depth for step *step_idx*. Writes: @@ -911,6 +886,7 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, ``command``, ``result``, and ``elapsed_s`` fields are merged into the step blob so a single entry captures everything. """ + output_dir = env_state.output_dir for directory in ARTIFACT_DIRECTORIES: (Path(output_dir) / directory).mkdir(parents=True, exist_ok=True) @@ -918,6 +894,7 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, wrist_world_map = None agent_world_map_hi = None wrist_world_map_hi = None + wrist_meta = None # Reuse one raw observation snapshot for state and per-step artifacts. raw = primitives.env.raw_obs() # Expose robot proprioception and object names, but never privileged @@ -1056,6 +1033,7 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, "w", ) as f: json.dump(wmeta_out, f, indent=2) + wrist_meta = wmeta except Exception as e: logger.warning("wrist depth/world dump failed: %s", e) @@ -1129,25 +1107,39 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, except FileNotFoundError: pass - blob = { - "step_idx": step_idx, - "libero_terminated": primitives.env.episode_terminated, - "episode_truncated": primitives.env.episode_truncated, - "task_language": primitives.env.get_task_language(), - "state": state, - "world_map": agent_world_map, - "wrist_world_map": wrist_world_map, - "world_map_hi": agent_world_map_hi, - "wrist_world_map_hi": wrist_world_map_hi, - } - # Merge the execution log (command + result + elapsed_s) into the - # state blob so a single entry captures everything for the step. - if log is not None: - blob["command"] = log.get("command") - blob["result"] = log.get("result") - blob["elapsed_s"] = log.get("elapsed_s") - _append_state(output_dir, blob) - return blob + frame_streams = [ + stream + for stream in ("image", "image_cam", "image_wrist") + if env_state.artifact_path(step_idx, stream).exists() + ] + depth_streams = [ + stream + for stream in ("depth", "depth_wrist") + if env_state.artifact_path(step_idx, stream).exists() + ] + camera_meta = {"agentview": agentview_meta} + if wrist_meta is not None: + camera_meta["wrist"] = wrist_meta + record = StepRecord( + step_idx=step_idx, + state=state, + frames=frame_streams, + depth=depth_streams, + camera_meta=camera_meta, + command=log.get("command") if log else None, + result=log.get("result") if log else None, + elapsed_s=log.get("elapsed_s") if log else None, + extras={ + "libero_terminated": primitives.env.episode_terminated, + "episode_truncated": primitives.env.episode_truncated, + "task_language": primitives.env.get_task_language(), + "world_map": agent_world_map, + "wrist_world_map": wrist_world_map, + "world_map_hi": agent_world_map_hi, + "wrist_world_map_hi": wrist_world_map_hi, + }, + ) + return env_state.append(record) # --------------------------------------------------------------------------- @@ -1517,38 +1509,9 @@ def dump_state(primitives: LiberoPrimitives, output_dir: str, step_idx: int, ] -# --------------------------------------------------------------------------- -# State trace readers -# --------------------------------------------------------------------------- - - -def _load_states() -> list: - """Return the parsed state trace from the local output dir.""" - path = artifact_path(get_output_dir(), "states") - if not path.exists(): - return [] - with open(path) as f: - return json.load(f) - - -def _latest_step() -> int | None: - states = _load_states() - if not states: - return None - return states[-1]["step_idx"] - - -def _load_step(nn: int) -> dict: - """Look up the state blob for step ``nn`` from states.json.""" - for entry in _load_states(): - if entry.get("step_idx") == nn: - return entry - raise FileNotFoundError(f"step {nn} not present in states.json") - - -def _load_image_path(nn: int, kind: str) -> str | None: +def _load_image_path(state: EnvState, nn: int, kind: str) -> str | None: """Return the path to a dumped state image. None if not present.""" - out_dir = get_output_dir() + out_dir = state.output_dir if kind == "agent": path = artifact_path(out_dir, "policy_image", step=nn, camera="agentview", resolution="low") elif kind == "camera": @@ -1562,8 +1525,12 @@ def _load_image_path(nn: int, kind: str) -> str | None: return str(path) -def _load_camera_meta(camera: str = "agentview", nn: int | None = None) -> dict: - out_dir = get_output_dir() +def _load_camera_meta( + state: EnvState, + camera: str = "agentview", + nn: int | None = None, +) -> dict: + out_dir = state.output_dir if camera == "agentview": path = artifact_path(out_dir, "metadata", camera="agentview", resolution="low") elif camera == "wrist" and nn is not None: @@ -1571,71 +1538,74 @@ def _load_camera_meta(camera: str = "agentview", nn: int | None = None) -> dict: else: raise ValueError("camera must be 'agentview' or 'wrist' with nn") if not path.exists(): - raise FileNotFoundError(f"{path.name} not found in {out_dir}") + raise FileNotFoundError(f"{path.name} not found in {state.output_dir}") with open(path) as f: return json.load(f) -def _load_depth(camera: str, nn: int) -> np.ndarray: - out_dir = get_output_dir() - if camera not in ("agentview", "wrist"): +def _load_depth(state: EnvState, camera: str, nn: int) -> np.ndarray: + if camera == "agentview": + stream = "depth" + elif camera == "wrist": + stream = "depth_wrist" + else: raise ValueError("camera must be 'agentview' or 'wrist'") - path = artifact_path(out_dir, "depth", step=nn, camera=camera, resolution="low") - if not path.exists(): - raise FileNotFoundError(f"{path.name} not found in {out_dir}") - depth = np.load(path) + depth = state.load_depth(nn, stream) if depth.ndim == 3: depth = depth[..., 0] return depth -def view_driver_state(step: int | None = None) -> dict: - latest = _latest_step() +def view_driver_state(step: int | None = None, *, state: EnvState) -> dict: + latest = state.latest_step if latest is None: return {"error": "no state entries; env not ready"} nn = latest if step is None else int(step) try: - data = _load_step(nn) + record = state.get(nn) except Exception as e: return {"error": f"step {nn} not present in state trace: {e}"} - out: dict = {"step": nn} - out["task_language"] = data.get("task_language") - # Default to {} (not the entire blob) when "state" is missing — otherwise - # command/result/vla_desync would bleed into the "state" field, confusing - # the LLM about what robot state actually contains. - out["state"] = data.get("state", {}) - out["libero_terminated"] = data.get("libero_terminated") - out["episode_truncated"] = data.get("episode_truncated") - out["world_map"] = data.get("world_map") - out["wrist_world_map"] = data.get("wrist_world_map") - out["world_map_hi"] = data.get("world_map_hi") - out["wrist_world_map_hi"] = data.get("wrist_world_map_hi") + extras = record.extras + out: dict = {"step": nn, "state": record.state} + out["task_language"] = extras.get("task_language") + out["libero_terminated"] = extras.get("libero_terminated") + out["episode_truncated"] = extras.get("episode_truncated") + out["world_map"] = extras.get("world_map") + out["wrist_world_map"] = extras.get("wrist_world_map") + out["world_map_hi"] = extras.get("world_map_hi") + out["wrist_world_map_hi"] = extras.get("wrist_world_map_hi") out["log"] = { - "command": data.get("command"), - "result": data.get("result"), - "elapsed_s": data.get("elapsed_s"), + "command": record.command, + "result": record.result, + "elapsed_s": record.elapsed_s, } for field, kind in ( ("image_path", "agent"), ("image_cam_path", "camera"), ("image_wrist_path", "wrist"), ): - image_path = _load_image_path(nn, kind) + image_path = _load_image_path(state, nn, kind) if image_path: out[field] = image_path for field, camera in ( ("image_cam_hi_path", "agentview"), ("image_wrist_hi_path", "wrist"), ): - image_path = artifact_path(get_output_dir(), "image", step=nn, camera=camera, resolution="high") + image_path = artifact_path( + state.output_dir, + "image", + step=nn, + camera=camera, + resolution="high", + ) if image_path.exists(): out[field] = str(image_path) return out -def _select_segment_artifacts(nn: int, camera: str): - out_dir = get_output_dir() +def _select_segment_artifacts(state: EnvState, nn: int, camera: str): + out_dir = state.output_dir if camera not in ("agentview", "wrist"): raise ValueError(f"unknown segment camera: {camera}") pairs = [ @@ -1645,7 +1615,6 @@ def _select_segment_artifacts(nn: int, camera: str): ) for resolution in ("high", "low") ] - for image_path, world_path in pairs: if image_path.exists() and world_path.exists(): return image_path, world_path, pairs @@ -1737,19 +1706,24 @@ def _write_segment_overlay(image_path: Path, mask: np.ndarray, return False -def view_camera_meta(camera: str = "agentview", step: int | None = None) -> dict: +def view_camera_meta( + camera: str = "agentview", + step: int | None = None, + *, + state: EnvState, +) -> dict: """Read camera calibration metadata for localization.""" if camera not in ("agentview", "wrist"): return {"error": f"bad camera '{camera}' (use 'agentview' or 'wrist')"} nn = None if camera == "wrist": - nn = _latest_step() if step is None else int(step) + nn = state.latest_step if step is None else int(step) if nn is None: return {"error": "no wrist metadata available"} try: - meta = _load_camera_meta(camera, nn) + meta = _load_camera_meta(state, camera, nn) except Exception as e: return {"error": f"{camera} camera metadata not found: {e}"} @@ -1768,6 +1742,8 @@ def back_project( col_range: list | None = None, z_min: float | None = None, z_max: float | None = None, + *, + state: EnvState, ) -> dict: """Look up a pixel's world XYZ in the precomputed world map.""" if camera not in ("agentview", "wrist"): @@ -1784,22 +1760,21 @@ def back_project( ) } - latest = _latest_step() - nn = latest if step is None else int(step) + nn = state.latest_step if step is None else int(step) if nn is None: return {"error": "no depth/world-map files available"} try: - data = _load_step(nn) + record = state.get(nn) except Exception as e: return {"error": f"step {nn} not present in state trace: {e}"} if camera == "agentview": - hi_artifact = data.get("world_map_hi") - low_artifact = data.get("world_map") + hi_artifact = record.extras.get("world_map_hi") + low_artifact = record.extras.get("world_map") else: - hi_artifact = data.get("wrist_world_map_hi") - low_artifact = data.get("wrist_world_map") + hi_artifact = record.extras.get("wrist_world_map_hi") + low_artifact = record.extras.get("wrist_world_map") source_artifact = hi_artifact if resolution == "high" else low_artifact if not source_artifact: return { @@ -1810,8 +1785,7 @@ def back_project( } try: - world_path = artifact_path(get_output_dir(), "world", step=nn, camera=camera, resolution=resolution) - world_map = np.load(world_path) + world_map = np.load(state.output_dir / source_artifact) except Exception as e: return { "error": ( @@ -1897,7 +1871,7 @@ def back_project( depth_m = None if source_artifact == low_artifact: try: - depth = _load_depth(camera, nn) + depth = _load_depth(state, camera, nn) except Exception as e: return {"error": f"{camera} depth not found for step {nn}: {e}"} depth_m = float(depth[row, col]) diff --git a/rpent/tools/common.py b/rpent/tools/common.py index fd959566..54e02ce9 100644 --- a/rpent/tools/common.py +++ b/rpent/tools/common.py @@ -3,10 +3,16 @@ import os from pathlib import Path +from typing import Any + +import numpy as np from rpent.utils.config import get_repo_root from rpent.utils.logging import get_output_dir +_BACKPROJECT_RADIUS = 6 +_DEPTH_BAND_M = 0.02 + TOOLS_SPEC: list[dict] = [ { "name": "read_text_file", @@ -134,6 +140,88 @@ def finish(status: str, summary: str) -> dict: return {"_finish": True, "status": status, "summary": summary} +def backproject_points(K, rows, cols, depths) -> np.ndarray: + """Back-project pixel coords + depths to camera-frame XYZ, shape (N, 3).""" + K = np.asarray(K, dtype=np.float64) + rows = np.asarray(rows, dtype=np.float64) + cols = np.asarray(cols, dtype=np.float64) + depths = np.asarray(depths, dtype=np.float64) + fx, fy = K[0, 0], K[1, 1] + cx, cy = K[0, 2], K[1, 2] + return np.stack( + [(cols - cx) * depths / fx, (rows - cy) * depths / fy, depths], axis=1 + ) + + +def robust_surface_centroid( + depth: np.ndarray, + K, + T_base_cam, + row: int, + col: int, + *, + radius: int = _BACKPROJECT_RADIUS, + band: float = _DEPTH_BAND_M, +) -> dict: + """Back-project a pixel neighbourhood to a robust 3D point. + + Back-projects every valid pixel in a ``(2*radius+1)`` window, keeps those on + the dominant surface (depth within ``band`` of the window median, rejecting + background / table / dropouts), and returns the median point + diagnostics. + Returns world ``xyz`` when ``T_base_cam`` is given, otherwise camera-frame + ``xyz_cam``. + """ + row, col = int(row), int(col) + radius = max(0, int(radius)) + height, width = depth.shape[:2] + if not (0 <= row < height and 0 <= col < width): + return { + "error": ( + f"pixel ({row},{col}) out of bounds; image is {height}x{width}" + ) + } + row_start, row_end = max(0, row - radius), min(height, row + radius + 1) + col_start, col_end = max(0, col - radius), min(width, col + radius + 1) + rows, cols = np.mgrid[row_start:row_end, col_start:col_end] + depths = depth[row_start:row_end, col_start:col_end].reshape(-1).astype( + np.float64 + ) + rows = rows.reshape(-1).astype(np.float64) + cols = cols.reshape(-1).astype(np.float64) + valid = np.isfinite(depths) & (depths > 0) + if not np.any(valid): + return {"error": f"no valid depth near ({row},{col}); pick another pixel"} + depths, rows, cols = depths[valid], rows[valid], cols[valid] + median_depth = float(np.median(depths)) + surface = np.abs(depths - median_depth) <= band + depths, rows, cols = depths[surface], rows[surface], cols[surface] + if depths.size == 0: + return {"error": f"no dominant surface depth near ({row},{col})"} + + camera_points = backproject_points(K, rows, cols, depths) + camera_point = np.median(camera_points, axis=0) + out: dict[str, Any] = { + "pixel": [row, col], + "radius": radius, + "n_points": int(camera_points.shape[0]), + "depth_m": round(median_depth, 4), + "xyz_cam": [round(float(value), 4) for value in camera_point], + } + if T_base_cam is not None: + transform = np.asarray(T_base_cam, dtype=np.float64) + base_points = camera_points @ transform[:3, :3].T + transform[:3, 3] + base_point = np.median(base_points, axis=0) + out["xyz"] = [round(float(value), 4) for value in base_point] + out["xy_spread_m"] = round( + float(np.hypot(*base_points[:, :2].std(axis=0))), 4 + ) + else: + out["xy_spread_m"] = round( + float(np.hypot(*camera_points[:, :2].std(axis=0))), 4 + ) + return out + + TOOL_HANDLERS: dict = { "read_text_file": read_text_file, "write_text_file": write_text_file, diff --git a/rpent/tools/state.py b/rpent/tools/state.py new file mode 100644 index 00000000..573fde6c --- /dev/null +++ b/rpent/tools/state.py @@ -0,0 +1,348 @@ +"""EnvState: the per-run state-trace owner. + +A **step** is one motion-primitive tool call that produced a dumped state +snapshot, indexed from ``0`` where ``0`` is the post-``reset()`` baseline +(no command). Non-motion tool calls (``get_ee_pose``, ``back_project``, +``view_driver_state``, file/memory tools) are NOT steps and leave the trace +untouched. + +``EnvState`` replaces the per-env ``_append_state`` / ``_load_states`` / +``_latest_step`` / ``_load_step`` / ``_load_image`` / ``_load_depth`` free +functions (and the toolkit's ``_next_step`` counter) with one explicit owner +constructed with an ``output_dir`` (no process-global). The toolkit and the +reader tools (``view_driver_state``, ``back_project``) hold a non-owning +reference to one ``EnvState`` per run; the composition root owns its lifecycle. + +Artefact naming follows the LIBERO layout. Each saved artefact is a named +*stream* turned into a path by a fixed rule:: + + stream "image" -> images/image_NN.png + stream "image_wrist" -> images_wrist/image_wrist_NN.png + stream "depth" -> depths/depth_NN.npy + stream "depth_wrist" -> depths_wrist/depth_wrist_NN.npy + stream "world" -> world/world_NN.npy (libero) + stream "wrist_meta" -> wrist_meta/wrist_meta_NN.json + +i.e. the stream name is the file prefix; the directory pluralises the +artefact type (``image``->``images``, ``depth``->``depths``; ``world`` and +``wrist_meta`` stay singular, matching libero) and appends the camera suffix. +Franka maps scene->``image``/``depth`` and wrist->``image_wrist``/ +``depth_wrist``; lerobot maps scene->``image``/``depth`` and +arm->``image_arm``. No per-env layout class is needed: the env's thin +``dump_state`` wrapper just picks the stream names. +""" +from __future__ import annotations + +import json +import os +import shutil +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Iterable + +import imageio.v2 as imageio +import numpy as np + +from rpent.utils.logging import get_logger + +logger = get_logger("env_state") + +# Artefact type (the leading segment of a stream name) -> file extension. +_ARTIFACT_EXT: dict[str, str] = { + "image": ".png", + "depth": ".npy", + "world": ".npy", + "wrist_meta": ".json", +} + +# Artefact type -> directory name (plural where libero pluralises). +_PLURAL: dict[str, str] = { + "image": "images", + "depth": "depths", +} + +def _split_stream(stream: str) -> tuple[str, str]: + for artifact_type in sorted(_ARTIFACT_EXT, key=len, reverse=True): + if stream == artifact_type: + return artifact_type, "" + prefix = artifact_type + "_" + if stream.startswith(prefix): + return artifact_type, stream[len(prefix):] + art, sep, suffix = stream.partition("_") + return art, (suffix if sep else "") + + +def stream_dir(output_dir: Path, stream: str) -> Path: + """Directory for a stream's artefacts (e.g. ``images_wrist``).""" + art, suffix = _split_stream(stream) + base = _PLURAL.get(art, art) + return output_dir / (base if not suffix else f"{base}_{suffix}") + + +def stream_path(output_dir: Path, stream: str, step_idx: int, ext: str | None = None) -> Path: + """Full path for one stream artefact at ``step_idx``.""" + art, _ = _split_stream(stream) + if ext is None: + ext = _ARTIFACT_EXT.get(art, ".bin") + return stream_dir(output_dir, stream) / f"{stream}_{step_idx:02d}{ext}" + + +# --------------------------------------------------------------------------- +# StepRecord +# --------------------------------------------------------------------------- + + +@dataclass +class StepRecord: + """One dumped step: the unit the LLM reads back via ``view_driver_state``. + + ``step_idx`` is the motion-primitive index (0 = post-reset baseline, which + has ``command``/``result``/``elapsed_s`` = None). ``frames``/``depth`` are + the stream names saved for this step (e.g. ``["image", "image_wrist"]``). + ``camera_meta`` is the per-camera metadata dict at capture time (intrinsics + + extrinsics + calibration status). ``extras`` absorbs env-specific fields + (libero's world maps, ``libero_terminated``, task_language, ...). + """ + + step_idx: int + state: dict + command: dict | None = None + result: dict | None = None + elapsed_s: float | None = None + frames: list[str] = field(default_factory=list) + depth: list[str] = field(default_factory=list) + camera_meta: dict | None = None + extras: dict[str, Any] = field(default_factory=dict) + + def to_blob(self) -> dict[str, Any]: + blob: dict[str, Any] = { + "step_idx": self.step_idx, + "state": self.state, + "frames": self.frames, + "depth": self.depth, + } + if self.camera_meta is not None: + blob["camera_meta"] = self.camera_meta + if self.command is not None: + blob["command"] = self.command + if self.result is not None: + blob["result"] = self.result + if self.elapsed_s is not None: + blob["elapsed_s"] = self.elapsed_s + if self.extras: + blob["extras"] = self.extras + return blob + + @classmethod + def from_blob(cls, blob: dict) -> "StepRecord": + return cls( + step_idx=int(blob.get("step_idx", -1)), + state=blob.get("state", {}), + command=blob.get("command"), + result=blob.get("result"), + elapsed_s=blob.get("elapsed_s"), + frames=list(blob.get("frames", [])), + depth=list(blob.get("depth", [])), + camera_meta=blob.get("camera_meta"), + extras=dict(blob.get("extras") or {}), + ) + + +# --------------------------------------------------------------------------- +# EnvState +# --------------------------------------------------------------------------- + + +class EnvState: + """Owns the on-disk state trace for one run. + + Constructed with an explicit ``output_dir`` (the composition root passes + it; ``EnvState`` never reaches into a process-global). Owns the step + counter, the ``states.json`` trace (atomic append), and the image/depth + artefact layout. The toolkit and the reader tools (``view_driver_state``, + ``back_project``) hold a non-owning reference to one ``EnvState`` per run. + """ + + def __init__(self, output_dir: Path | str): + self.output_dir = Path(output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + self._steps: list[dict] = self._read_states() + self._next_step = self._derive_next_step() + + # -- internal: states.json I/O ---------------------------------------- + + def _states_path(self) -> Path: + return self.output_dir / "states.json" + + def _read_states(self) -> list[dict]: + path = self._states_path() + if not path.exists(): + return [] + try: + with open(path) as f: + arr = json.load(f) + return [s for s in arr if isinstance(s, dict)] if isinstance(arr, list) else [] + except Exception as e: + logger.warning("could not parse %s: %s; starting fresh", path, e) + return [] + + def _write_states_atomically(self, steps: list[dict]) -> None: + path = self._states_path() + tmp = path.with_name(path.name + ".tmp") + with open(tmp, "w") as f: + json.dump(steps, f, indent=2, default=str) + os.replace(tmp, path) + + def _derive_next_step(self) -> int: + if not self._steps: + return 0 + return max(int(s["step_idx"]) for s in self._steps if "step_idx" in s) + 1 + + # -- lifecycle -------------------------------------------------------- + + def reset(self, *, wipe_streams: Iterable[str] = ()) -> None: + """Wipe stale artefact dirs + ``states.json`` for a fresh run. + + Resets the in-memory trace and counter. Step 0 (baseline) is dumped by + the caller right after via :meth:`append`. + """ + self.output_dir.mkdir(parents=True, exist_ok=True) + for stream in wipe_streams: + d = stream_dir(self.output_dir, stream) + if d.exists(): + shutil.rmtree(d) + states = self._states_path() + if states.exists(): + states.unlink() + self._steps = [] + self._next_step = 0 + + # -- counter ---------------------------------------------------------- + + @property + def next_step_idx(self) -> int: + """The step_idx the next dumped step will get (does NOT advance).""" + return self._next_step + + @property + def latest_step(self) -> int | None: + """The highest step_idx successfully written (None if trace is empty).""" + if not self._steps: + return None + return int(self._steps[-1]["step_idx"]) + + # -- writing ---------------------------------------------------------- + + def save_image(self, step_idx: int, stream: str, frame) -> str | None: + """Save one RGB frame under ``_.png``; return ``stream`` on success.""" + path = stream_path(self.output_dir, stream, step_idx) + path.parent.mkdir(parents=True, exist_ok=True) + arr = np.asarray(frame) + if arr.dtype != np.uint8: + arr = arr.astype(np.uint8) + try: + imageio.imwrite(path, arr) + return stream + except Exception as e: + logger.warning("frame dump failed for stream %s: %s", stream, e) + return None + + def save_depth(self, step_idx: int, stream: str, depth) -> str | None: + """Save one depth map under ``_.npy``; return ``stream`` on success.""" + path = stream_path(self.output_dir, stream, step_idx) + path.parent.mkdir(parents=True, exist_ok=True) + try: + np.save(path, np.asarray(depth, dtype=np.float32)) + return stream + except Exception as e: + logger.warning("depth dump failed for stream %s: %s", stream, e) + return None + + def artifact_path( + self, + step_idx: int, + stream: str, + *, + ext: str | None = None, + ) -> Path: + """Return the canonical path for one step artifact.""" + return stream_path(self.output_dir, stream, step_idx, ext=ext) + + def append(self, record: StepRecord) -> StepRecord: + """Atomically append ``record`` to ``states.json`` and advance the counter. + + The counter and ``states.json`` are kept in sync: this write is the + single source of truth for "latest step". If the write raises, the + counter is NOT advanced (so the next caller reuses the same + ``step_idx``) -- no desync between the in-memory counter and disk. + """ + blob = record.to_blob() + self._write_states_atomically(self._steps + [blob]) + self._steps.append(blob) + self._next_step = max(self._next_step, record.step_idx + 1) + return record + + # -- reading ---------------------------------------------------------- + + def get(self, step_idx: int) -> StepRecord: + """Look up the step record for ``step_idx``.""" + for blob in self._steps: + if int(blob.get("step_idx", -1)) == step_idx: + return StepRecord.from_blob(blob) + raise KeyError(f"step {step_idx} not present in states.json") + + def records(self) -> list[StepRecord]: + """Return the persisted step records in trace order.""" + return [StepRecord.from_blob(blob) for blob in self._steps] + + def load_image_bytes(self, step_idx: int, stream: str) -> bytes | None: + path = stream_path(self.output_dir, stream, step_idx) + if not path.exists(): + return None + return path.read_bytes() + + def load_depth(self, step_idx: int, stream: str) -> np.ndarray: + path = stream_path(self.output_dir, stream, step_idx) + return np.load(path) + + # -- LLM-facing view (ex-view_driver_state) --------------------------- + + def view(self, step: int | None = None, *, image_slots: dict[str, str] | None = None) -> dict: + """Build the ``view_driver_state`` dict for one step. + + ``image_slots`` maps a ToolResult image slot (``_image_bytes``, + ``_image_cam_bytes``, ``_image_wrist_bytes``) to a stream name whose + saved PNG should be embedded as bytes. The env decides which cameras + map to which slots. + """ + latest = self.latest_step + if latest is None: + return {"error": "no driver state entries; driver not ready"} + nn = latest if step is None else int(step) + try: + rec = self.get(nn) + except Exception as e: + return {"error": f"step {nn} not present in driver state trace: {e}"} + out: dict[str, Any] = { + "step": nn, + "state": rec.state, + "frames": rec.frames, + "depth": rec.depth, + "camera_meta": { + name: {k: v for k, v in meta.items() if k not in {"K", "T_base_cam"}} + for name, meta in (rec.camera_meta or {}).items() + }, + "log": { + "command": rec.command, + "result": rec.result, + "elapsed_s": rec.elapsed_s, + }, + } + if rec.extras: + out["extras"] = rec.extras + if image_slots: + for slot, stream in image_slots.items(): + b = self.load_image_bytes(nn, stream) + if b: + out[slot] = b + return out diff --git a/rpent/tools/toolkit.py b/rpent/tools/toolkit.py index 31a7cbca..a747020f 100644 --- a/rpent/tools/toolkit.py +++ b/rpent/tools/toolkit.py @@ -106,10 +106,16 @@ class Toolkit: :meth:`close` to release env-side primitives / servers at the end of the run. """ - def __init__(self, *, dashboard_events: DashboardEventSink) -> None: + def __init__( + self, + *, + dashboard_events: DashboardEventSink, + state: Any = None, + ) -> None: # name -> (spec, handler) self._tools: dict[str, tuple[dict[str, Any], Callable[..., dict[str, Any]]]] = {} self._dashboard_events = dashboard_events + self._state = state self._operation_lock = threading.Lock() self._active_operation: _ToolOperation | None = None self._register_common_tools() From cc1a4ca993df5608a1a2a983531a4ed3947e6d54 Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Tue, 4 Aug 2026 08:59:09 +0000 Subject: [PATCH 18/22] refactor(state): centralize environment artifact lifecycle in EnvState Signed-off-by: Jiaxing Qiu --- .../rst_source/development/add_primitive.rst | 6 +- .../rst_source/development/add_robot.rst | 24 +- docs/source-en/rst_source/quickstart.rst | 15 +- docs/source-en/rst_source/usage/libero.rst | 6 +- .../rst_source/development/add_primitive.rst | 4 +- .../rst_source/development/add_robot.rst | 25 +- docs/source-zh/rst_source/quickstart.rst | 6 +- docs/source-zh/rst_source/usage/libero.rst | 6 +- robots/franka/__init__.py | 2 - robots/franka/toolkit.py | 19 +- robots/franka/tools.py | 101 +-- robots/lerobot/__init__.py | 2 - robots/lerobot/toolkit.py | 18 +- robots/lerobot/tools.py | 110 ++- robots/libero/__init__.py | 2 - robots/libero/guides/env_calibration.md | 44 +- robots/libero/guides/pro_hybrid_guide.md | 619 ++++---------- robots/libero/guides/strict_hybrid_guide.md | 783 +++++------------- robots/libero/prompts/system.py | 117 +-- robots/libero/prompts/user.py | 10 +- robots/libero/toolkit.py | 80 +- robots/libero/tools.py | 710 ++++++---------- rpent/cli/main.py | 1 - rpent/dashboard/server.py | 15 +- rpent/dashboard/state.py | 32 +- rpent/planner/api_loop.py | 6 +- rpent/tools/state.py | 635 ++++++++------ rpent/utils/sam3_client.py | 10 +- 28 files changed, 1293 insertions(+), 2115 deletions(-) diff --git a/docs/source-en/rst_source/development/add_primitive.rst b/docs/source-en/rst_source/development/add_primitive.rst index dcefdd5b..1148f94c 100644 --- a/docs/source-en/rst_source/development/add_primitive.rst +++ b/docs/source-en/rst_source/development/add_primitive.rst @@ -167,8 +167,10 @@ Design principles for a new primitive the state dump reflecting the post-action world. Don't let the primitive return before the render finishes. - **Return small dicts.** Tool return values are fed back to the LLM - as text. Store larger content, such as images, depth data, and - ``states.json``, in the state dump instead. + as text. Save larger observations through ``EnvState.save``; ``EnvState`` + automatically records each logical base name in its owned + ``StepRecord.artifacts`` set. Expose images through ``view_driver_state`` and + geometry through environment tools rather than returning raw paths. - **Guardrails belong in env_server**, not in the toolkit. The LLM can and will call any tool with any arguments; workspace bounds and safety clamps must be enforced on the server side. diff --git a/docs/source-en/rst_source/development/add_robot.rst b/docs/source-en/rst_source/development/add_robot.rst index 2111a9f7..0a2b5497 100644 --- a/docs/source-en/rst_source/development/add_robot.rst +++ b/docs/source-en/rst_source/development/add_robot.rst @@ -222,21 +222,24 @@ Anthropic-style tool definitions (``name``, ``description``, ``input_schema``), plus any module-level functions referenced by the toolkit (e.g. ``view_driver_state``, ``back_project``, ``finish``). -**Per-step state dump** — ``dump_state(primitives, output_dir, step_idx, log)`` -serializes whatever state the agent will read back via the ``view_*`` tools -(images, depths, JSON state, camera meta) into ``output_dir``. +**Per-step state dump** — ``dump_state(driver, env_state, log)`` opens +``env_state.record_step(...)`` and receives the allocated step index. Save +large observations through ``env_state.save(..., step=step_idx)``. ``EnvState`` +owns and commits the ``StepRecord`` and adds every successfully saved base name +to its flat ``artifacts`` set automatically. Readers use the canonical artifact +filenames rather than maintaining a parallel observation index. **Toolkit class** — subclass ``rpent.tools.toolkit.Toolkit``: -- build the primitives in ``__init__`` through a custom initialization - helper (named ``init_primitives_clean`` in LIBERO; it wipes stale - ``images/`` etc., constructs the primitives, and dumps step 0), +- build the primitive driver in ``__init__`` through a custom initialization + helper (named ``init_primitives_clean`` in LIBERO; it calls + ``EnvState.reset()``, constructs the primitives, and dumps step 0), - register each tool with ``self.add_tool(name, spec, handler)`` — stateless readers (``view_driver_state``, ``finish``, …) bind directly to module-level functions; primitive tools route through ``_step(name, **kwargs)`` which calls ``getattr(self._primitives, name)(**kwargs)`` and re-renders state, -- override ``close()`` to write any remaining agent-side artifacts (e.g. the - LIBERO toolkit saves the agentview MP4 there). +- override ``close()`` to save remaining agent-side artifacts through + ``EnvState`` (for example ``state.save("episode.mp4", frames, step=None)``). ``primitives_kwargs`` (forwarded from ``__init__.py:get_toolkit``) is the dict the toolkit passes verbatim to your primitives' ``__init__`` — typically @@ -246,8 +249,9 @@ Conventions worth keeping ------------------------- - ``output_dir`` is the working directory that the runner creates for each - run. Images, depths, ``states.json``, transcripts, ``episode.mp4``, and other - artifacts go there. + run. Environment observations are owned by ``EnvState``; callers use logical + base names and never construct storage paths. Transcripts and other + run-management outputs share the same run directory. - Tool definitions use the Anthropic format (``name`` / ``description`` / ``input_schema``). Every tool registered with ``self.add_tool(...)`` is exposed to all planners. diff --git a/docs/source-en/rst_source/quickstart.rst b/docs/source-en/rst_source/quickstart.rst index b9ba30b5..8fb27a7d 100644 --- a/docs/source-en/rst_source/quickstart.rst +++ b/docs/source-en/rst_source/quickstart.rst @@ -103,14 +103,11 @@ A successful run: by the elapsed time, token usage, and path to the run record. 3. With the Dashboard enabled, also streams agent output, camera views, the action timeline, and clip replays to the Dashboard. -4. By default, artifacts are saved under - ``logs/__t_s/``. They include - ``transcript_*.json`` (run record), ``states.json`` (one record per - environment step), ``recipe_*.jsonl`` (action sequence), and - ``episode.mp4`` (episode video). - -After the run, inspect the final record in ``states.json``: -``libero_terminated`` set to ``true`` means LIBERO judged the task complete. -You can also open ``episode.mp4`` to review the run. +4. By default, artifacts are saved under ``logs/__t_s/``. They include ``transcript_*.json`` (run record), ``states.json`` (the versioned ``EnvState`` manifest), ``recipe_*.jsonl`` (action sequence), and ``episode.mp4`` (episode video). Step artifact files use flat, zero-padded step prefixes with a minimum width of two digits and are managed internally by ``EnvState``. + +Inspect the final state through the Dashboard or +``view_driver_state(step=-1)``. Its top-level ``libero_terminated`` value is the +benchmark outcome. ``states.json`` is internal ``EnvState`` storage and should +not be parsed by callers. You can also open ``episode.mp4`` to review the run. If something goes wrong, inspect the four log files described at the bottom of :doc:`installation`. diff --git a/docs/source-en/rst_source/usage/libero.rst b/docs/source-en/rst_source/usage/libero.rst index 4fd25442..b2f7682b 100644 --- a/docs/source-en/rst_source/usage/libero.rst +++ b/docs/source-en/rst_source/usage/libero.rst @@ -147,8 +147,10 @@ Physical action tools advance the environment and record new state and images. coordinates. - ``segment(prompt=... / point=..., ...)`` — use SAM3 to segment an existing image with a text or point prompt. -- ``view_driver_state(step=None)`` — read an existing state and image record. -- ``view_camera_meta(camera=..., step=None)`` — read existing camera metadata. +- ``view_driver_state(step=-1)`` — read a recorded state and its embedded + observation images. Step ``0`` is initial; ``-1`` is latest. +- ``view_camera_meta(camera=..., step=-1)`` — read camera metadata for a + recorded step. Step ``-1`` is latest. - ``finish(status, summary)`` — end the current run. These tools do not advance the environment. diff --git a/docs/source-zh/rst_source/development/add_primitive.rst b/docs/source-zh/rst_source/development/add_primitive.rst index 909a7e82..2ba6ce24 100644 --- a/docs/source-zh/rst_source/development/add_primitive.rst +++ b/docs/source-zh/rst_source/development/add_primitive.rst @@ -152,7 +152,9 @@ primitives 方法，以及调用完成后的状态快照。区别仅在于方法 - **每个工具执行结束后都要保存新的状态快照。** 下一轮需要读取动作执行后的 环境状态，因此原语不能在渲染完成前返回。 - **工具只返回简短的字典。** 返回值会以文本形式提供给 LLM；图像、深度数据和 - ``states.json`` 等较大的内容则通过状态快照提供。 + 其他大型观测应通过 ``EnvState.save`` 保存；``EnvState`` 会把每个逻辑基础 + 文件名自动加入其持有的 ``StepRecord.artifacts`` 集合。图像通过 + ``view_driver_state`` 提供，几何数据通过环境工具访问，不返回原始路径。 - **安全限制由 ``env_server`` 强制执行。** LLM 可能使用任意参数调用工具， 因此工作空间边界和安全限制不能只依赖 toolkit。 diff --git a/docs/source-zh/rst_source/development/add_robot.rst b/docs/source-zh/rst_source/development/add_robot.rst index 995a13e7..0f76750c 100644 --- a/docs/source-zh/rst_source/development/add_robot.rst +++ b/docs/source-zh/rst_source/development/add_robot.rst @@ -209,21 +209,23 @@ Anthropic API 的工具定义格式，包含 ``name``、``description`` 和 ``input_schema``），以及 toolkit 引用的模块级函数，例如 ``view_driver_state``、``back_project`` 和 ``finish``。 -**每步状态 dump** —— ``dump_state(primitives, output_dir, step_idx, log)`` 把 agent -之后会通过 ``view_*`` 工具读回的所有状态 (图像、深度、JSON 状态、camera meta) -序列化到 ``output_dir``。 +**每步状态 dump** —— ``dump_state(driver, env_state, log)`` 通过 +``env_state.record_step(...)`` 创建由 ``EnvState`` 持有的步骤，并取得分配的 +step index。大型观测通过 ``env_state.save(..., step=step_idx)`` 保存。每次保存 +成功后，``EnvState`` 会自动把基础文件名加入该 ``StepRecord`` 的扁平 +``artifacts`` 集合并最终提交记录；读取方直接使用规范化的工件文件名。 **Toolkit 类** 继承 ``rpent.tools.toolkit.Toolkit``： -- 在 ``__init__`` 中通过自定义的初始化辅助方法构建 primitives（LIBERO - 中的方法名为 ``init_primitives_clean``；它会清理过期的 ``images/`` 等目录、 - 构造原语并 dump 第 0 步）, +- 在 ``__init__`` 中通过自定义的初始化辅助方法构建 primitive driver（LIBERO + 中的方法名为 ``init_primitives_clean``；它会调用 ``EnvState.reset()``、构造 + 原语并 dump 第 0 步）, - 用 ``self.add_tool(name, spec, handler)`` 注册每个工具。无状态的读取工具 （如 ``view_driver_state``、``finish``）直接绑定模块级函数；原语工具通过 ``_step(name, **kwargs)`` 调用。``_step`` 使用 - ``getattr(self._primitives, name)(**kwargs)`` 调用 primitives 方法并重新渲染状态； -- 重写 ``close()``，将 agent 侧生成的文件写入磁盘（例如 LIBERO toolkit - 在这里保存 agentview MP4）。 + ``getattr(self._primitives, name)(**kwargs)`` 调用 driver 方法并重新渲染状态； +- 重写 ``close()``，通过 ``EnvState`` 保存 agent 侧剩余工件（例如 + ``state.save("episode.mp4", frames, step=None)``）。 ``primitives_kwargs`` 由 ``__init__.py:get_toolkit`` 转发给 toolkit，再原样传入 primitives 的 ``__init__``。其中通常包含 @@ -232,8 +234,9 @@ primitives 的 ``__init__``。其中通常包含 建议遵循的约定 -------------- -- ``output_dir`` 是 runner 为单次运行创建的临时目录。图像、深度数据、 - ``states.json``、transcript 和 ``episode.mp4`` 等工件都写入该目录。 +- ``output_dir`` 是 runner 为单次运行创建的工作目录。环境观测由 + ``EnvState`` 管理；调用方只使用逻辑基础文件名，不自行拼接存储路径。 + transcript 等运行管理输出与环境工件共享该目录。 - 工具定义使用 Anthropic API 格式（``name`` / ``description`` / ``input_schema``）。 每个用 ``self.add_tool(...)`` 注册的工具都会暴露给所有 planner。 diff --git a/docs/source-zh/rst_source/quickstart.rst b/docs/source-zh/rst_source/quickstart.rst index b9cf71c7..9d870b55 100644 --- a/docs/source-zh/rst_source/quickstart.rst +++ b/docs/source-zh/rst_source/quickstart.rst @@ -95,7 +95,9 @@ LIBERO-PRO 仿真资源。下面以 LIBERO-PRO 和 ``claude_code`` planner 1. 终端会先显示 ``env_server``、``vla_server`` 和 ``sam3_server`` 的启动信息。 2. 智能体的逐轮输出和工具调用会显示在终端中；运行结束时还会显示耗时、token 用量和运行记录的路径。 3. 启用 Dashboard 后，智能体的输出、相机视图、动作时间线和片段回放也会实时显示在 Dashboard 中。 -4. 默认输出目录为 ``logs/__t_s/``，其中包含 ``transcript_*.json``\ （运行记录）、``states.json``\ （每个环境步的记录）、``recipe_*.jsonl``\ （动作序列）和 ``episode.mp4``\ （回合录像）。 +4. 默认输出目录为 ``logs/__t_s/``，其中包含 ``transcript_*.json``\ （运行记录）、``states.json``\ （带版本号的 ``EnvState`` 清单）、``recipe_*.jsonl``\ （动作序列）和 ``episode.mp4``\ （回合录像）。每步工件文件采用至少两位、零填充的步骤前缀扁平命名，并由 ``EnvState`` 在内部管理。 -运行结束后，查看 ``states.json`` 的最后一条记录：``libero_terminated`` 为 ``true`` 表示 LIBERO 已判定任务完成；也可以打开 ``episode.mp4`` 复核运行过程。 +通过 Dashboard 或 ``view_driver_state(step=-1)`` 查看最终状态；其顶层 +``libero_terminated`` 即为基准任务结果。``states.json`` 是 ``EnvState`` 的内部 +存储，调用方不应直接解析。也可以打开 ``episode.mp4`` 复核运行过程。 出问题时，参考 :doc:`installation` 页底部提到的四份日志文件。 diff --git a/docs/source-zh/rst_source/usage/libero.rst b/docs/source-zh/rst_source/usage/libero.rst index 2675458b..b66e06ac 100644 --- a/docs/source-zh/rst_source/usage/libero.rst +++ b/docs/source-zh/rst_source/usage/libero.rst @@ -141,8 +141,10 @@ LIBERO 工具分为物理动作工具和只读工具。 - ``back_project(row, col, ...)`` —— 将图像像素反投影到世界坐标。 - ``segment(prompt=... / point=..., ...)`` —— 通过 SAM3 对已有图像进行文本或 点提示分割。 -- ``view_driver_state(step=None)`` —— 读取已有的状态和图像记录。 -- ``view_camera_meta(camera=..., step=None)`` —— 读取已有的相机元数据。 +- ``view_driver_state(step=-1)`` —— 读取已记录的状态和内嵌观测图像；第 0 步为 + 初始状态，``-1`` 表示最新状态。 +- ``view_camera_meta(camera=..., step=-1)`` —— 读取指定步骤的相机元数据； + ``-1`` 表示最新状态。 - ``finish(status, summary)`` —— 结束当前运行。 这些工具不会推进环境。 diff --git a/robots/franka/__init__.py b/robots/franka/__init__.py index c559df4c..ec67a0b7 100644 --- a/robots/franka/__init__.py +++ b/robots/franka/__init__.py @@ -42,7 +42,6 @@ def get_env_spec() -> EnvSpec: def get_toolkit( *, primitives_kwargs: dict[str, Any], - video_path: str | None = None, dashboard: Any = None, ): """Return the Franka toolkit (common tools + Cartesian primitives). @@ -53,7 +52,6 @@ def get_toolkit( from robots.franka.toolkit import FrankaToolkit return FrankaToolkit( - video_path=video_path, dashboard=dashboard, **primitives_kwargs, ) diff --git a/robots/franka/toolkit.py b/robots/franka/toolkit.py index e7212e5a..de96df81 100644 --- a/robots/franka/toolkit.py +++ b/robots/franka/toolkit.py @@ -14,10 +14,11 @@ class FrankaToolkit(Toolkit): """Toolkit for the standalone Franka environment.""" - # Streams wiped on init (LIBERO layout: scene=image/depth, wrist=image_wrist/depth_wrist). - _WIPE_STREAMS = ("image", "image_wrist", "depth", "depth_wrist") # view_driver_state image slots: primary scene -> _image_bytes, wrist -> _image_cam_bytes. - _VIEW_IMAGE_SLOTS = {"_image_bytes": "image", "_image_cam_bytes": "image_wrist"} + _VIEW_IMAGE_SLOTS = { + "_image_bytes": "scene.png", + "_image_cam_bytes": "wrist.png", + } _DRIVER_READERS = ( "get_ee_pose", @@ -40,7 +41,6 @@ def __init__( self, *, env: Any, - video_path: str | None = None, dashboard: Any = None, ) -> None: # EnvState owns the trace + counter for this run (explicit output_dir, @@ -48,7 +48,6 @@ def __init__( # for now the toolkit constructs it from get_output_dir(). state = EnvState(get_output_dir()) super().__init__(dashboard=dashboard, state=state) - self._video_path = video_path self.init_driver_clean(env=env) self._register_tools() @@ -84,22 +83,20 @@ def _step(self, name: str, **kwargs) -> dict: elapsed = round(time.time() - t0, 2) result_dict = result if isinstance(result, dict) else {"value": result} - step_idx = self._state.next_step_idx - franka_tools.dump_state( + record = franka_tools.dump_state( self._driver, self._state, - step_idx=step_idx, log={"command": command, "result": result_dict, "elapsed_s": elapsed}, ) - out = self._state.view(step_idx, image_slots=self._VIEW_IMAGE_SLOTS) + out = self._state.view(record.step_idx, image_slots=self._VIEW_IMAGE_SLOTS) out["agent_elapsed_s"] = elapsed return out def init_driver_clean(self, *, env: Any) -> None: - self._state.reset(wipe_streams=self._WIPE_STREAMS) + self._state.reset() driver = franka_tools.FrankaPrimitives(env=env) driver.reset() - franka_tools.dump_state(driver, self._state, step_idx=0, log=None) + franka_tools.dump_state(driver, self._state, log=None) self._driver = driver def close(self) -> None: diff --git a/robots/franka/tools.py b/robots/franka/tools.py index 31327bf4..c4d2b08c 100644 --- a/robots/franka/tools.py +++ b/robots/franka/tools.py @@ -182,40 +182,38 @@ def _refresh(self) -> None: def dump_state( driver: FrankaPrimitives, state: EnvState, - step_idx: int, log: dict | None = None, ) -> StepRecord: - """Dump camera frames + proprioceptive state via the EnvState owner. - - Maps the Franka cameras onto the LIBERO stream layout: scene -> ``image``/ - ``depth`` (the primary view) and wrist -> ``image_wrist``/``depth_wrist``. - """ - image_streams = {"scene": "image", "wrist": "image_wrist"} - depth_streams = {"scene": "depth", "wrist": "depth_wrist"} - - saved_frames: list[str] = [] - for camera, frame in driver.latest_frames().items(): - stream = image_streams.get(camera, camera) - if state.save_image(step_idx, stream, frame): - saved_frames.append(stream) - - saved_depths: list[str] = [] - for camera, depth in driver.latest_depths().items(): - stream = depth_streams.get(camera, camera) - if state.save_depth(step_idx, stream, depth): - saved_depths.append(stream) - - record = StepRecord( - step_idx=step_idx, + """Dump camera artifacts and proprioceptive state through ``EnvState``.""" + log = log or {} + with state.record_step( state=driver.get_state(), - frames=sorted(saved_frames), - depth=sorted(saved_depths), - camera_meta=driver.latest_camera_meta(), - command=log.get("command") if log else None, - result=log.get("result") if log else None, - elapsed_s=log.get("elapsed_s") if log else None, - ) - return state.append(record) + command=log.get("command"), + result=log.get("result"), + elapsed_s=log.get("elapsed_s"), + ) as step_idx: + for camera, frame in driver.latest_frames().items(): + state.save( + f"{camera}.png", + frame, + step=step_idx, + ) + + for camera, depth in driver.latest_depths().items(): + state.save( + f"{camera}_depth.npy", + depth, + step=step_idx, + ) + + for camera, camera_meta in driver.latest_camera_meta().items(): + state.save( + f"{camera}_metadata.json", + camera_meta, + step=step_idx, + ) + + return state.get(step_idx) # view_driver_state now lives on EnvState.view (bound by the toolkit with the @@ -225,7 +223,7 @@ def dump_state( def back_project( row: int, col: int, - step: int | None = None, + step: int = -1, camera: str = "wrist", radius: int | None = _BACKPROJECT_RADIUS, *, @@ -239,24 +237,29 @@ def back_project( ``panda_link0`` when the camera has ``T_base_cam`` for the step; otherwise ``xyz_cam`` plus a warning. """ - nn = state.latest_step if step is None else int(step) - if nn is None: - return {"error": "no steps available"} try: - rec = state.get(nn) + record = state.get(step) except Exception as exc: - return {"error": f"step {nn} not present in driver state trace: {exc}"} + return {"error": f"state step not available: {exc}"} + nn = record.step_idx camera = str(camera or "wrist") - meta = (rec.camera_meta or {}).get(camera) - if not meta: + metadata_name = f"{camera}_metadata.json" + depth_name = f"{camera}_depth.npy" + if metadata_name not in record.artifacts: return { "error": f"camera {camera!r} has no metadata at step {nn}", - "available_cameras": sorted((rec.camera_meta or {}).keys()), + "available_cameras": sorted( + name.removesuffix("_metadata.json") + for name in record.artifacts + if name.endswith("_metadata.json") + ), } - depth_stream = "depth" if camera == "scene" else "depth_wrist" try: - depth = state.load_depth(nn, depth_stream) + meta = state.load(metadata_name, step=nn) + if depth_name not in record.artifacts: + raise FileNotFoundError(depth_name) + depth = state.load(depth_name, step=nn) except Exception as exc: return {"error": f"depth for camera {camera!r} step {nn} not found: {exc}"} @@ -292,15 +295,16 @@ def back_project( { "name": "view_driver_state", "description": ( - "Read step NN from states.json plus matching camera PNGs. If step " - "is null, returns the latest entry. Embeds scene and wrist images." + "Read one recorded state and its observation artifacts. Step -1 " + "selects the latest entry. Embeds scene and wrist images." ), "input_schema": { "type": "object", "properties": { "step": { - "type": ["integer", "null"], - "description": "Step number; 0 = initial. Null = latest.", + "type": "integer", + "default": -1, + "description": "Step number; 0 = initial, -1 = latest.", } }, }, @@ -322,8 +326,9 @@ def back_project( "row": {"type": "integer", "description": "Pixel row (y)."}, "col": {"type": "integer", "description": "Pixel column (x)."}, "step": { - "type": ["integer", "null"], - "description": "Step whose saved depth to use; null = latest.", + "type": "integer", + "default": -1, + "description": "Step whose saved depth to use; -1 = latest.", }, "camera": { "type": "string", diff --git a/robots/lerobot/__init__.py b/robots/lerobot/__init__.py index 15969208..62c3b574 100644 --- a/robots/lerobot/__init__.py +++ b/robots/lerobot/__init__.py @@ -53,7 +53,6 @@ def get_env_spec() -> EnvSpec: def get_toolkit( *, primitives_kwargs: dict[str, Any], - video_path: str | None = None, dashboard: Any = None, ): """Return the SO101 toolkit (common tools + SO101 primitives). @@ -65,7 +64,6 @@ def get_toolkit( from robots.lerobot.toolkit import LerobotToolkit return LerobotToolkit( - video_path=video_path, dashboard=dashboard, **primitives_kwargs, ) diff --git a/robots/lerobot/toolkit.py b/robots/lerobot/toolkit.py index 119da41d..88b7069c 100644 --- a/robots/lerobot/toolkit.py +++ b/robots/lerobot/toolkit.py @@ -19,8 +19,10 @@ class LerobotToolkit(Toolkit): """Toolkit for the LeRobot SO101 environment.""" - _WIPE_STREAMS = ("image", "image_arm", "depth") - _VIEW_IMAGE_SLOTS = {"_image_bytes": "image", "_image_cam_bytes": "image_arm"} + _VIEW_IMAGE_SLOTS = { + "_image_bytes": "scene.png", + "_image_cam_bytes": "arm.png", + } # Read-only tools backed by a live driver call (no state dump). These query # the robot/scene directly: forward kinematics + scene camera calibration. _DRIVER_READERS = ( @@ -42,12 +44,10 @@ def __init__( *, env: Any, model: Any | None = None, - video_path: str | None = None, dashboard: Any = None, ) -> None: state = EnvState(get_output_dir()) super().__init__(dashboard=dashboard, state=state) - self._video_path: str | None = video_path self.init_driver_clean(env=env, model=model) self._register_tools() @@ -89,23 +89,21 @@ def _step(self, name: str, **kwargs) -> dict: result_dict = result if isinstance(result, dict) else {"value": result} - step_idx = self._state.next_step_idx - lerobot_tools.dump_state( + record = lerobot_tools.dump_state( self._driver, self._state, - step_idx=step_idx, log={"command": command, "result": result_dict, "elapsed_s": elapsed}, ) - out = self._state.view(step_idx, image_slots=self._VIEW_IMAGE_SLOTS) + out = self._state.view(record.step_idx, image_slots=self._VIEW_IMAGE_SLOTS) out["agent_elapsed_s"] = elapsed return out def init_driver_clean(self, *, env: Any, model: Any | None = None) -> None: """Wipe stale run artifacts, build the primitive driver, dump step 0.""" - self._state.reset(wipe_streams=self._WIPE_STREAMS) + self._state.reset() driver = lerobot_tools.LerobotPrimitives(env=env, model=model) driver.reset() - lerobot_tools.dump_state(driver, self._state, step_idx=0, log=None) + lerobot_tools.dump_state(driver, self._state, log=None) self._driver = driver diff --git a/robots/lerobot/tools.py b/robots/lerobot/tools.py index c551dd77..89bfd7ab 100644 --- a/robots/lerobot/tools.py +++ b/robots/lerobot/tools.py @@ -173,54 +173,46 @@ def _refresh(self) -> None: def dump_state( driver: LerobotPrimitives, state: EnvState, - step_idx: int, log: dict | None = None, ) -> StepRecord: - """Dump the camera frames, scene depth, and proprioceptive state. - - Writes per step: - - ``images/image_NN.png`` (scene color) - - ``images_arm/image_arm_NN.png`` (arm color) - - ``depths/depth_NN.npy`` (scene metric depth, aligned to scene color) - - Scene calibration is stored on the step record. ``EnvState`` atomically - appends that record to ``states.json`` and advances the trace counter. - """ - image_streams = {"scene": "image", "arm": "image_arm"} - saved_frames: list[str] = [] - for camera, frame in driver.latest_frames().items(): - stream = image_streams.get(camera, f"image_{camera}") - if state.save_image(step_idx, stream, frame): - saved_frames.append(stream) - - saved_depths: list[str] = [] - depth = driver.latest_depth() - if depth is not None and state.save_depth(step_idx, "depth", depth): - saved_depths.append("depth") - - scene_meta = driver.get_scene_camera_meta() - camera_meta = ( - {"scene": scene_meta} - if isinstance(scene_meta, dict) and "error" not in scene_meta - else {} - ) - record = StepRecord( - step_idx=step_idx, + """Dump camera artifacts and proprioceptive state through ``EnvState``.""" + log = log or {} + with state.record_step( state=driver.get_state(), - frames=sorted(saved_frames), - depth=sorted(saved_depths), - camera_meta=camera_meta, - command=log.get("command") if log else None, - result=log.get("result") if log else None, - elapsed_s=log.get("elapsed_s") if log else None, - ) - return state.append(record) + command=log.get("command"), + result=log.get("result"), + elapsed_s=log.get("elapsed_s"), + ) as step_idx: + for camera, frame in driver.latest_frames().items(): + state.save( + f"{camera}.png", + frame, + step=step_idx, + ) + + depth = driver.latest_depth() + if depth is not None: + state.save( + "scene_depth.npy", + depth, + step=step_idx, + ) + + scene_meta = driver.get_scene_camera_meta() + if isinstance(scene_meta, dict) and "error" not in scene_meta: + state.save( + "scene_metadata.json", + scene_meta, + step=step_idx, + ) + + return state.get(step_idx) def back_project( row: int, col: int, - step: int | None = None, + step: int = -1, radius: int | None = _BACKPROJECT_RADIUS, *, state: EnvState, @@ -236,22 +228,24 @@ def back_project( stable object centroid rather than one face pixel. Use ``radius=0`` for the old single-pixel behavior. - Pick ``(row, col)`` on ``images/image_NN.png``; depth is aligned to it in - ``depths/depth_NN.npy``. Returns base/world ``xyz`` when calibrated, else + Pick ``(row, col)`` on the saved ``scene.png`` observation; depth is + aligned in ``scene_depth.npy``. Returns base/world ``xyz`` when calibrated, else camera-frame ``xyz_cam`` with a note. """ - nn = state.latest_step if step is None else int(step) - if nn is None: - return {"error": "no steps available"} try: - record = state.get(nn) + record = state.get(step) except Exception as exc: - return {"error": f"step {nn} not present in driver state trace: {exc}"} - meta = (record.camera_meta or {}).get("scene") - if not meta: + return {"error": f"state step not available: {exc}"} + nn = record.step_idx + metadata_name = "scene_metadata.json" + depth_name = "scene_depth.npy" + if metadata_name not in record.artifacts: return {"error": f"scene camera metadata not recorded for step {nn}"} try: - depth = state.load_depth(nn, "depth") + meta = state.load(metadata_name, step=nn) + if depth_name not in record.artifacts: + raise FileNotFoundError(depth_name) + depth = state.load(depth_name, step=nn) except Exception as exc: return {"error": f"depth for step {nn} not found: {exc}"} @@ -288,17 +282,16 @@ def back_project( { "name": "view_driver_state", "description": ( - "Read step NN from `states.json` plus the matching scene image " - "`images/image_NN.png` and arm image " - "`images_arm/image_arm_NN.png`. If step is null, returns the " - "latest entry. Embeds both camera frames as image content blocks." + "Read one recorded state and its observation artifacts. Step -1 " + "selects the latest entry. Embeds scene and arm camera frames." ), "input_schema": { "type": "object", "properties": { "step": { - "type": ["integer", "null"], - "description": "Step number; 0 = initial. Null = latest.", + "type": "integer", + "default": -1, + "description": "Step number; 0 = initial, -1 = latest.", }, }, }, @@ -340,8 +333,9 @@ def back_project( "row": {"type": "integer", "description": "Pixel row (y) in the scene image, near the target center."}, "col": {"type": "integer", "description": "Pixel column (x) in the scene image, near the target center."}, "step": { - "type": ["integer", "null"], - "description": "Step whose depth to use; null = latest.", + "type": "integer", + "default": -1, + "description": "Step whose depth to use; -1 = latest.", }, "radius": { "type": ["integer", "null"], diff --git a/robots/libero/__init__.py b/robots/libero/__init__.py index 7bf9323e..a5fd060e 100644 --- a/robots/libero/__init__.py +++ b/robots/libero/__init__.py @@ -46,7 +46,6 @@ def get_toolkit( *, primitives_kwargs: dict[str, Any], dashboard_events: DashboardEventSink, - video_path: str | None = None, ): """Return the LIBERO toolkit (common tools + LIBERO primitives).""" from robots.libero.toolkit import LiberoToolkit @@ -54,7 +53,6 @@ def get_toolkit( return LiberoToolkit( primitives_kwargs=primitives_kwargs, dashboard_events=dashboard_events, - video_path=video_path, ) diff --git a/robots/libero/guides/env_calibration.md b/robots/libero/guides/env_calibration.md index 1dbab00d..886fd711 100644 --- a/robots/libero/guides/env_calibration.md +++ b/robots/libero/guides/env_calibration.md @@ -2,12 +2,15 @@ ## Current LIBERO MCP Runtime Contract -Use this file as a calibration reference only. For current MCP-based runs, use -structured MCP tools, do not issue file-based protocol commands, and do not manually manage -`env_server.py`. Do not read BDDL files or hidden task definition files to infer -coordinates. Do not expect object world coordinates in `states.json`; localize -objects through images_cam + depth/back_project, segment, and wrist/high-res -artifacts when available. +Use this file as a calibration reference for structured-tool runs. The runner +owns the environment server and the `EnvState` lifecycle. Do not issue +file-based driver commands, inspect observation storage directly, or read BDDL +files for coordinates. + +Start with `view_driver_state({"step": 0})`. It returns the initial robot state, +top-level task language, logical observation references, and embedded camera +images. Use `back_project` or `segment` for geometry and `view_camera_meta` for +calibration. Step `-1` selects the latest record. Measured 2026-05-20 on `libero_10_with_mug` t0 (LIVING_ROOM frame) and t8 (KITCHEN frame). All probes use `move_to` with `gripper=-1` and tight @@ -17,9 +20,9 @@ Measured 2026-05-20 on `libero_10_with_mug` t0 (LIVING_ROOM frame) and t8 Each task scene uses one of the table fixtures below, which sets the entire world-frame z origin. The OSC workspace and all pick/place altitudes shift -accordingly. **Check `states.json[0].state.robot0_eef_pos[2]` in the initial -state and branch on it.** Do not read BDDL files for this; use runtime state and -visual evidence. +accordingly. **Check `state.robot0_eef_pos[2]` in the result of +`view_driver_state({"step": 0})` and branch on it.** Do not read BDDL files for +this; use runtime state and visual evidence. | Fixture | eef home z | Table top z | Used by tasks | |---|---|---|---| @@ -119,7 +122,7 @@ limit 1.15). My libero_10 t0 used z=0.95 for travel — safe and consistent. ## Practical rules going forward -1. **Always read `states.json[0].state.robot0_eef_pos[2]` before computing any z target.** +1. **Always inspect the initial returned `state.robot0_eef_pos[2]` before computing any z target.** ≈ 0.68 → LIVING_ROOM; ≈ 1.17 → KITCHEN; ≈ 0.26 → OBJECT. Use the matching frame table above. 2. **Never command an eef z below the per-frame floor.** Going to z=0.42 @@ -143,24 +146,25 @@ limit 1.15). My libero_10 t0 used z=0.95 for travel — safe and consistent. `move_to` + `set_gripper` (last resort; unreliable for objects <6 cm — see `resources/libero/memory/feedback_scripted_pick_limits.md`). -## Calibration log files +## Calibration Records + +Each calibration motion returns its state, command result, elapsed time, and +embedded observation. Use that tool result immediately. To revisit a recorded +step, call `view_driver_state({"step": N})`; use `-1` for the latest step. -Raw probe logs are preserved in: -- `{output_dir}/states.json` (one step entry per command — the per-command - audit; each entry has `command`, `result`, `state`, `elapsed_s`). -- Only kept for the most recent agent session; reproduce by re-running - the calibration with the snippet in the next section. +The internal `states.json` file is a versioned manifest owned by `EnvState`, not +a list for manual indexing. Calibration analysis should use structured tool +results rather than parsing storage files. ## Reproducer -Legacy calibration notes below describe the old file-protocol flow. In the -current MCP runtime, use the runner-managed environment and call structured MCP -tools instead. +Use the runner-managed environment and call structured tools: ```bash # For each z in 0.65 .. 0.42, call the MCP tool: move_to({"xyz": [-0.20, 0.10, 0.65], "gripper": -1, "tol": 0.008, "step_clip": 0.010, "max_steps": 80}) -# Then read states.json entry NN for final_eef_pos & final_dist_m. +# Inspect the returned result for final_eef_pos and final_dist_m. +# Call view_driver_state({"step": -1}) only when the latest state is needed again. ``` diff --git a/robots/libero/guides/pro_hybrid_guide.md b/robots/libero/guides/pro_hybrid_guide.md index d7999405..7da40264 100644 --- a/robots/libero/guides/pro_hybrid_guide.md +++ b/robots/libero/guides/pro_hybrid_guide.md @@ -1,502 +1,197 @@ -# LIBERO-Pro Hybrid (Pi0.5 + LLM-in-the-loop) — Perception-Isolated Guide - -You are picking up the **LIBERO-Pro** evaluation track in **perception-isolated** -mode — the only mode in this repository (the legacy oracle-state mode is not -included here). - -> **Pi0.5 only does the grasp (`pi0_pick`). The LLM (you) handles every motion -> (`move_to`), every release, sequencing, retries — and you do not get GT object -> coordinates. You localize objects yourself from the depth + camera calibration -> the runtime dumps each step.** - -This document layers on the base playbook. Read it first: - -- [`strict_hybrid_guide.md`](./strict_hybrid_guide.md) — the perception protocol: - back-projection localization, the perception artifacts - (`images_cam/image_cam_NN.png` / `depths/depth_NN.npy` / `camera_meta.json`), - the primitive vocabulary, the Rules (0/1/2/4/5), and the `strict_perception` - audit format. **This is the source of truth for *how you localize*.** -- **This file** — LIBERO-Pro–specific setup, the four perturbation axes, the Pi0 - fullshot baseline, the frame split, and how perception isolation changes the - P2-swap story. - -> Your task comes from `states.json[0].task_language`; the BDDL is FORBIDDEN. -> Read the authoritative instruction (the BDDL `:language` tag, coord-free) that -> the runtime injects and obey it verbatim. Do **not** scrape the BDDL or import -> the benchmark: that is error-prone (wrong task-map index → wrong task) and a -> perception-isolation breach (the `:init` block holds the GT coordinates this -> mode withholds). For swap fixtures, localize them **visually** (§3.5), never -> from `:init`. You never hand-roll projection math — `back_project` applies -> `K⁻¹` + the cam→world extrinsic and returns the surface `world_xyz` under a -> pixel; just use it. - -Whenever this guide says "see the perception protocol" or "the localization -snippet", it means `strict_hybrid_guide.md`. **Every rule there applies here**; -this guide only *adds* PRO constraints and tooling. Runner contract: call the -structured MCP tools (the runner owns the env server — do not start/stop it, and -issue no file-based commands); it is a single-episode run — no `reset`/`exit`, -recover in place or write an honest failure audit and `finish`. - -## 0. What's different from oracle PRO mode (read this first) - -| | oracle mode (not in this repo) | **perception (this guide)** | -|---|---|---| -| how launched | oracle-state run | `rpent/cli/main.py --libero-type pro` (or `LIBERO_TYPE=pro`); perception artifacts always dumped, coords withheld | -| `states.json` objects | full `objects:{name:[x,y,z]}` | **`object_names:[…]` only — NO coords** | -| how you learn the task | env prompt / scrape BDDL | **`states.json[0].task_language`** (authoritative `:language`, coord-free) — never read the BDDL | -| extra obs artifacts | agentview RGB only | **+ `images_cam/`, `depths/`, `world/` (agentview); `images_wrist/`, `depths_wrist/`, `world_wrist/` (wrist); hi-res pairs; `camera_meta.json`** — all via `back_project` | -| cameras | agentview only | **agentview (fixed, ~1m → ±8-13cm) + eye-in-hand wrist (moves with gripper, ±1-2cm when <20cm to target)** — see §3.6 | -| how P2 swap is solved | read swapped coords from `states.json[0]` | **localize the swapped objects by `back_project`** | -| how a swapped *fixture* site is found | read the swap BDDL `:init` block | **localize the fixture visually** (see §3.5) | -| run budget | short | **larger** — perception localization + manipulation is slower (raise `--max-turns` / `--planner-timeout-s`) | -| audit `regime` | `strict` | `strict_perception` | -| which image you pick pixels in | agentview RGB | `images_cam/image_cam_NN.png` (or hi-res) for **pixel-picking**; `images/image_NN.png` only for a sanity glance | - -**The single most important conceptual shift.** In oracle PRO mode the headline -is "the hybrid beats Pi0 on P2 because it reads the *swapped* coordinates straight -out of `states.json[0]`, while Pi0 is prompt-/memory-blind." In **perception** -mode there are no coordinates to read — so the hybrid's P2 win now comes from -**seeing where the object is and back-projecting it**. This is a *stronger* claim -(no oracle state at all), and it is the whole point of running PRO in perception -mode: P1 (Task) is still won by reading the *language*; P2 (Position) is now won by -*perception*, not by an oracle. - -## 1. Why LIBERO-Pro - -LIBERO-Pro -([paper](https://arxiv.org/pdf/2510.03827), [repo](https://github.com/Zxy-MLlab/LIBERO-PRO)) -perturbs each base task along five axes; all end-to-end VLAs (OpenVLA / Pi0 / -Pi0.5 / UniVLA) collapse on the two strongest: - -| Axis | Suffix | Paper column | What changes | Headline result | -|---|---|---|---|---| -| **Task** | `_task` | **P1** | Instruction + goal predicate inverted | All VLAs ≈ 0.0 | -| **Position** | `_swap` | **P2** | Object **and fixture** initial positions swapped | All VLAs 0.0–0.4 | -| Semantic | `_lan` | — | Instruction paraphrased; goal unchanged | VLAs handle (memorize visual) | -| Object | `_object` | — | Object appearance / colour / scale | VLA visual policy stressed | -| Environment | `_environment` | — | Table / scene swapped | Visual policy stressed | - -The agentic hybrid wins on **P1 and P2** by routing the language channel and the -*perceived* spatial-state channel through the LLM. Object and Environment -perturbations enter through the Pi0 vision channel and the hybrid inherits the -VLA's weakness there — declare that upfront, don't oversell. **In perception -mode, the Object/Environment axes also stress your own localization** (a -recolored or rescaled object is harder to pixel-pick), so lean on the -multi-pixel-median tip from the perception protocol. - -## 2. Setup (do these once on a fresh checkout) - -Run the idempotent installer from the repo root: - -```bash -bash scripts/install_libero_pro_plus.sh -``` +# LIBERO-Pro Hybrid Perception Guide + +This guide extends [strict_hybrid_guide.md](./strict_hybrid_guide.md) for the +LIBERO-Pro evaluation tracks. Read the strict guide first; its runtime contract, +localization discipline, and single-attempt rules apply unchanged. -It does all four steps below: the liberopro editable install, applies the -benchmark-registration patch, syncs the authoritative HF dataset snapshot, and -verifies with the `get_benchmark(...).get_task(0).language` check. The perception -observables (depth + both cameras + hi-res) are unconditional once the runner -launches with `--libero-type pro`; there is nothing perception-specific to -install beyond the PRO setup itself. +LIBERO-Pro perturbs object placement, fixture placement, spatial relations, and +task language. The agent must solve the observed scene rather than replaying a +seed-specific command sequence. -### 2.1. LIBERO-PRO repo +## Start Of Run -Cloned at `${LIBERO_PRO_PATH:-/path/to/LIBERO-PRO}/` from -`https://github.com/RLinf/LIBERO-PRO.git` and installed editable into the openpi -venv: +Call: -```bash -python -m pip show liberopro -# Name: liberopro Version: 0.1.0 Location: ${LIBERO_PRO_PATH:-/path/to/LIBERO-PRO} +```json +{"step": 0} ``` -### 2.2. Apply the benchmark-registration patch +through `view_driver_state`. From the returned tool result: -The upstream `__init__.py` does **not** expose the 16 perturbation suites through -`get_benchmark()`. Our patch -[`scripts/liberopro_register_perturbations.patch`](../../../../scripts/liberopro_register_perturbations.patch) -adds them and overrides `Task.language` to read each BDDL's actual `:language` -tag (so the perturbed instruction reaches Pi0 / hybrid). +- read top-level `task_language` verbatim; +- inspect `state.robot0_eef_pos` to identify the scene frame; +- inspect object names only as a scene inventory, never as coordinates; +- inspect the embedded agentview image for semantic identity and relations; +- inspect the wrist image only when it contains useful close geometry. -```bash -cd ${LIBERO_PRO_PATH:-/path/to/LIBERO-PRO} -git apply /scripts/liberopro_register_perturbations.patch -``` +Step `-1` means the latest record. Do not inspect or index `states.json` +directly; it is an internal manifest. -If already applied (likely), `git status -s` shows clean. If you reinstall -liberopro, re-apply. +## Perturbation Axes -### 2.3. Huggingface dataset (authoritative) +### Object perturbation -The LIBERO-PRO git repo ships **incomplete / broken** init files for several -perturbation suites (e.g. `libero_spatial_swap` has 0 BDDLs; some -`libero_spatial_task` `.pruned_init` files are 0 bytes). Treat the git repo as -unreliable for perturbation data. The full, correct set lives on Huggingface -([`zhouxueyang/LIBERO-Pro`](https://huggingface.co/datasets/zhouxueyang/LIBERO-Pro)), -persisted locally at: +Objects move relative to their seed positions. Re-localize every target and +destination from the current images. Seed recipes remain useful for primitive +ordering, prompt wording, safe heights, and known failure modes, but never for +coordinates. -``` -${LIBEROPRO_DATASET_PATH:-/path/to/liberopro_hf}/ -├── bddl_files/ 16 perturbation suites, 10 BDDLs each -└── init_files/ 16 perturbation suites, 10 init files each -``` +### Spatial perturbation -Covers `{libero_spatial, libero_object, libero_goal, libero_10} × {swap, task, -lan, object}`. The installer syncs this into the liberopro install (overwriting -the broken upstream files); if the persistent copy is gone, re-download with: +Relations such as left/right, on-top-of, or next-to may select a different +instance than in the reference scene. Identify the candidate satisfying the +current visual relation, then localize it with `back_project` or `segment`. -```bash -python -c " -from huggingface_hub import snapshot_download -snapshot_download(repo_id='zhouxueyang/LIBERO-Pro', repo_type='dataset', - local_dir='${LIBEROPRO_DATASET_PATH:-/path/to/liberopro_hf}', - allow_patterns=['bddl_files/**','init_files/**'])" -``` +### Goal perturbation -### 2.4. Verify +The destination or requested interaction changes. Re-read `task_language` and +classify the destination semantically. Do not infer the goal from the suite, +task index, object list, or a sibling result file. -```bash -LIBERO_TYPE=pro python -c " -import liberopro.liberopro.benchmark as bench -for n in ['libero_spatial_task','libero_spatial_swap','libero_spatial_lan']: - b = bench.get_benchmark(n)(); t = b.get_task(0) - print(f'{n} t0: {t.language!r} trials={len(b.get_task_init_states(0))}')" -``` +### Task-language perturbation -Expected: -``` -libero_spatial_task t0: 'Pick the akita black bowl not between the plate and the ramekin and place it on the plate' trials=50 -libero_spatial_swap t0: 'Pick the akita black bowl between the plate and the ramekin and place it on the plate' trials=50 -libero_spatial_lan t0: 'lift the black bowl between the plate and ramekin and set it on the plate' trials=50 -``` +The authoritative instruction is the current top-level `task_language`. Do not +read BDDL files. They combine language with hidden initialization data and +therefore violate perception isolation. -## 3. PRO-specific environment gotchas +## Scene Frame Selection -Everything in `strict_hybrid_guide.md` applies. The following are **additional** -PRO constraints (most live in detail at [`env_calibration.md`](./env_calibration.md)). +The initial end-effector height distinguishes the principal scene frames: -### 3.1. Three scene frames, picked per-task +| Initial EEF z | Frame | Typical fixtures | +|---|---|---| +| approximately 0.26 m | object/low-table | grocery and basket tasks | +| approximately 0.68 m | living-room table | plates, baskets, pudding | +| approximately 1.17 m | kitchen table | stove, cabinet, drawer, microwave | -PRO scenes use one of three table fixtures; the eef home z differs by up to -~0.9 m. +Use `view_driver_state({"step": 0})["state"]["robot0_eef_pos"][2]` as the +measurement. Then use the matching safe-height guidance from +[env_calibration.md](./env_calibration.md). -| Fixture | eef home z | Table top z | xy reachable | Where | -|---|---|---|---|---| -| `living_room_table` | ≈ 0.68 | ≈ 0.43 | `(x∈±0.30, y∈±0.30)` | basket / plate / pudding | -| `kitchen_table` | ≈ 1.17 | ≈ 0.90 | `(x∈±0.30, y∈±0.30)` | stove / cabinet / drawer / microwave | -| `object` (low table) | ≈ 0.26 | ≈ 0.0 | `(x∈±0.30, y∈±0.30)` | `libero_object` grocery-into-basket | +## Mandatory Perception Table -**Mandatory check at session start: read `states.json[0].state.robot0_eef_pos[2]`** -(via `view_driver_state({"step": 0})`). ≈ 0.68 → LIVING_ROOM; ≈ 1.17 → KITCHEN; -≈ 0.26 → OBJECT. Pick `pre_pos_z` / `carry_z` / `release_z` accordingly (per-item -OBJECT-frame altitudes are in -`resources/libero/memory/project_libero_object_pro_done.md`). Sending a -wrong-frame z (e.g. KITCHEN coordinates while the env is in LIVING_ROOM frame) -crashes the env worker (EOFError, silent state loss). +Before manipulating, create one row for every task-relevant entity: -> **Perception note.** This proprioceptive z is *not* an object coordinate — it's -> the robot's own pose, which `states.json` still gives you in perception mode. -> Reading it to pick the frame is fine. It also doubles as a free depth sanity -> check: your back-projected table z should sit ≈0.25 m below eef home z (≈0.43 in -> LIVING_ROOM, ≈0.90 in KITCHEN). If a back-projection lands far from that, you -> picked a wrong pixel. +| field | meaning | +|---|---| +| role | target, destination, support, fixture, or relation landmark | +| semantic evidence | visual properties and relation establishing identity | +| agentview pixels | several interior pixels or a verified segment mask | +| agentview xyz | robust projection from the global camera | +| wrist refinement | accepted, rejected, basket-confirmed, or unnecessary | +| final xyz | coordinate used for planning | +| uncertainty | duplicates, occlusion, rim bias, label ambiguity, etc. | -For `libero_spatial_*` you are always in **KITCHEN frame**. Standard heights: +Do not begin manipulation while a required row lacks a defensible identity or +coordinate. -``` -pre_pos_z = 1.05 # ~7 cm above objects at z≈0.97 -carry_z = 1.10 # safe traversal, well under upper limit 1.15 -release_z = 1.01 # ~7 cm above plate top at z≈0.90 -``` +## Swapped Objects And Fixtures -These are *altitude* knobs (robot-frame), not object positions — you still derive -the object's **xy** by `back_project` every time. - -### 3.2. xy single-step ±0.30 cap - -OSC flips IK branches if you command `|x|>0.30` or `|y|>0.30` in a single -`move_to` (the eef lands in the wrong half-space and corrupts the run). **Never -command beyond ±0.30 in a single move** — split into carry-z waypoints. (Detail -in `env_calibration.md`.) - -### 3.3. Slow long-distance carry for swap variants - -P2 swap can move a large object 15 cm across the table. `step_clip=0.025` lets it -slip in the gripper and the object ends up centimetres off target. Mitigation -(proven on `_swap` t0, carries over directly): - -- `carry_z = 1.15` (higher than usual 1.10) -- `step_clip = 0.020` (slower) -- **Re-localize mid-travel** instead of trusting a cached `object_xyz - eef_xyz` - offset: `Read images_cam/image_cam_NN.png` (or the hi-res `images_cam_hi/`) - mid-carry and call `back_project` on the carried object's pixel; if the - perceived offset drifted >5 mm from post-pick, re-pre-position and re-`pi0_pick`. - -### 3.4. Task language is from BDDL, not filename - -After the patch, `get_task(i).language` returns the perturbed `:language` tag, and -the env passes it to Pi0 as the prompt (surfaced to you as -`states.json[0].task_language`). For `_task` and `_lan` this is the perturbed -instruction. **Don't override it** — falsifying the VLA's prompt-blindness is the -point. You read the same language to decide *which* object to localize and place. - -### 3.5. ⚠ Swap moves FIXTURES too — and you must localize them visually - -This is the biggest perception-mode trap, and it is **specific to `_swap`**. The -P2 perturbation does not only swap loose objects; for `libero_goal_swap` it swaps -entire **fixtures** (stove ↔ cabinet ↔ wine_rack), so a goal predicate like -`On(bowl, flat_stove_1_cook_region)` now points at wherever the *stove* was -relocated to, and there are no coordinates in `states.json` to read. See -`resources/libero/memory/feedback_swap_perturbs_fixtures.md`. - -In **oracle** mode the documented fix is to read the swap BDDL `:init` block and -recompute the fixture site's world coordinates. **That is forbidden here** — the -`:init` block is ground-truth geometry. In **perception** mode you instead: - -1. Identify the target fixture by name from the (perturbed) task language and - `state.object_names`. -2. **Localize the fixture's predicate site visually** in - `images_cam_hi/image_cam_hi_NN.png`: pick pixels on the stove's cook region / - cabinet top surface / rack top shelf, then `back_project` (sample 3–5 pixels, - median the xy — fixtures are large and the surface you want is the *placement* - surface, not the nearest edge). -3. Carry target xy = perceived site xy; descend to the perceived site z + a small - clearance, `release`, then **retreat with gripper open** — the predicate often - fires *during* the settle, not at the release step itself. - -So the swap-fixture problem becomes "find the fixture in the image" rather than -"read where the BDDL put it." Loose-object swaps (e.g. `libero_spatial_swap`, -where only bowls/plates move) are simpler: just localize each object by -`back_project` as usual; their new positions fall straight out of the depth. - -### 3.6. Two cameras — agentview = IDENTITY, wrist = GEOMETRY - -Do **not** restate the two-camera protocol here — the strict guide's **First-step -perception protocol** is the source of truth: agentview / agentview-hi is the -semantic identity authority (decides *which* object/surface satisfies the task -language + relation); wrist / wrist-hi refines geometry for the *same* candidate -(accept only within ~3–5 cm of the agentview anchor, never average, basket/cavity -excepted). Both maps share one world frame, read via -`back_project({"camera":"wrist"})`; the optional `segment` honours both cameras -(`{"camera":"agentview"|"wrist"}`, `min_score` default 0.2, or `point:[row,col]`). - -**PRO implication:** `_swap` (and `_task`) can *invert* target or destination -semantics, so **never reuse base-task or recipe coordinates** — the object -satisfying the relation may now sit where the base task never put it. Let -agentview decide *what*; the near-vertical wrist is BAD at identity (locks onto -look-alikes) and only refines *where*. - -### 3.6c. Mandatory pre-task perception pass - -Also owned by the strict guide's First-step perception protocol: before ANY -pick/place, build the localization table (one row per task-relevant entity) and -pass the FINAL READY CHECK. **PRO implication:** in single-attempt mode a -wrong-target first grab is unrecoverable, and under `_swap` the "right" target is -exactly the one you'd get wrong by habit — so localizing all entities up front is -cheap insurance, not optional. - -### 3.6b. Hi-res perception channel (1024×1024) - -Also owned by the strict guide's Hi-res perception channel section (1024×1024 -pairs each step; prefer the hi image for identification; `back_project` defaults -`resolution="high"`; never mix pixel grids). **PRO implication:** hi-res fixes -*which* object you point at — decisive for the recolored/rescaled Object axis and -for telling swapped same-shape groceries apart — but does NOT change metric -accuracy or replace the wrist coarse→fine refinement. **Identity caveat:** SAM3 -gives two same-shaped, different-brand objects the SAME category (it labelled the -tomato-sauce and alphabet-soup cans identically) — use its mask only as a category -candidate, then READ the label yourself in the hi-res crop to assign identity; -never let SAM3's category settle a brand choice. - -## 4. The four-cell experiment per (base task, seed) - -For each base task you claim coverage on, generate four runs — every run is a -perception cell: - -| Suite | Variant | Perception-mode expectation | -|---|---|---| -| `libero_spatial` | base sanity | Pi0 and hybrid both pass | -| `libero_spatial_task` | **P1 Task** | Pi0 ✗ (picks base target); hybrid ✓ (LLM flips target from instruction, then localizes it) | -| `libero_spatial_swap` | **P2 Position** | Pi0 mixed; hybrid ✓ — **by localizing the swapped object/fixture from depth, not from state** | -| `libero_spatial_lan` | Semantic | Both pass (paraphrase invariant) | +For swap-style perturbations: -Replace `spatial` with `object`, `goal`, or `10` for the other base suites. +1. Identify both swapped entities in the embedded agentview image. +2. Classify each by appearance and current relation, not expected seed layout. +3. Localize each independently. +4. Verify the chosen target still satisfies `task_language`. +5. Re-localize after any contact that could move either entity. -### 4.1. Hybrid run — the MCP runner +The runtime withholds privileged coordinates, so there is no coordinate field +to fall back on. The current images and geometry tools are the source of truth. -Launch a cell with the CLI; the runner owns `env_server.py`, exposes the -structured tools, and runs single-attempt: +## Destination Classification -```bash -python rpent/cli/main.py --env libero --suite --task --seed \ - --libero-type pro --planner claude_code --model claude-opus-4-8 +Pro scenes frequently contain look-alike surfaces. Before placement, explicitly +classify all plausible destinations: -# e.g. --suite libero_spatial_task 0 (P1) · libero_spatial_swap 0 (P2) · -# libero_goal_swap 2 (P2 fixture swap) · libero_10_task 5 (long horizon) -``` +- plate versus stove burner; +- cabinet top versus drawer opening; +- basket interior versus rim; +- microwave cavity versus door or surrounding counter; +- movable lid versus fixed fixture surface. -`--libero-type pro` may be given as `LIBERO_TYPE=pro` instead. +Only after semantic classification should you call `back_project` or `segment` +for coordinates. -Audit + recipe land in the run's `output_dir` (`output_dir` and `recipe_tag` -arrive in your first message): +## Mid-Carry Re-Localization -``` -{output_dir}/{recipe_tag}.json <- you write this (write_text_file) -{output_dir}/recipe_{recipe_tag}.jsonl <- exported automatically by the runner -``` +Long-horizon Pro tasks often move or occlude objects during earlier steps. +Before each new pick or placement: -Do NOT write into `resources/libero/results_*_pert/` — that tree is a **read-only -seed-0 reference corpus**, not a write target. - -### 4.2. Environment server is runner-owned - -There is no manual server to launch and no REPL to drive: the MCP runner starts, -manages, and tears down `env_server.py` for you and blocks each tool call until -the next `states.json` entry is dumped. Do not start, stop, or background it, and -do not poll for readiness. - -### 4.3. Pi0 fullshot baseline - -The baseline is Pi0.5 driving the task end-to-end with the runtime's own -(perturbed) `task_language`. There is **no standalone baseline CLI in this repo**; -the numbers to compare against are the recorded full-shot results (the team's -`SUCCESS_RATES` table). **Do not invent a `pi0_baseline.py` path.** - -Pi0 never sees object coords in either mode, so there is no perception variant of -the baseline. Expected behavior: - -- **P1 (task):** Pi0 "succeeds at the wrong task" — it picks the *base*-task - target object, places it on the plate, and `libero_terminated=False` because the - goal predicate names a different object. This is exactly the gap the hybrid - closes. -- **P2 (position):** Pi0 picks / places at the *base* (un-swapped) location. - -### 4.4. Audit JSON — PRO + perception fields - -Start from the `strict_perception` audit schema in the perception protocol, and -add the PRO fields: - -```jsonc -{ - "suite": "libero_spatial_swap", - "task_id": 0, - "seed": 0, - "regime": "strict_perception", - "perturbation_type": "swap (P2 Position perturbation)", - "perturbed_task_language": "", - "perturbation_semantics": "", - "expected_baseline_behavior": "", - "strategy_notes": "HOW you localized — which pixel(s) in images_cam, back-projected world xyz; for swap, how you found the relocated object/fixture", - "pick_result": { /* the pi0_pick step's result */ }, - "final_state": { /* latest states.json entry's `state` field */ }, - "libero_terminated": true -} -``` +1. call `view_driver_state({"step": -1})` if the previous primitive result is + no longer in context; +2. inspect the newest embedded images; +3. re-localize any entity that may have moved; +4. update the working perception table; +5. verify the remaining primitive order still matches `task_language`. -`strategy_notes` **must** describe the localization (pixel → `back_project` → -world xyz). For swap cells, explicitly note that the relocated object/fixture was -found by perception, not by reading coords. If unrecoverable after honest -exploration, write `libero_terminated: false` with what you tried and which step -failed — never warp (teleport primitives are deleted; see Rule 4 in the -perception protocol). Write the audit with `write_text_file` to -`{output_dir}/{recipe_tag}.json`, then call `finish`. - -## 5. Perception-protocol rules — PRO clarifications - -- **Rule 0 (use images for reasoning).** Even more critical under PRO: P2 swap can - move a large object/fixture clear across the table, and you have *no* - coordinates to fall back on. `images_cam/image_cam_NN.png` + depth are your only - spatial truth — open `images_cam/` and describe the scene before deciding - targets. -- **Rule 1 (no `pi0_end_to_end`).** Pi0 does the grasp via `pi0_pick`; the LLM - scripts every motion + release. Under PRO this is doubly important — handing - back to Pi0 means handing back to the prompt-blind / memorized-place habit you - are trying to falsify. -- **Rule 2 (single-episode current run).** This is a one-shot eval: do not call - `reset` or `exit`. Recover *within* the episode when safe (re-localize, - re-pre-position, re-`pi0_pick`, walk the prompt ladder, - `rotate_pitch`/`move_pose`); otherwise write an honest stuck/failure audit and - `finish`. `_swap` typically needs an in-episode retry — document it. -- **Rule 4 (no teleport).** `set_object_pose`, `articulate_to`, `js_move_to`, - `carry_object` are deleted. A goal past OSC reach with no physical approach → - honest `libero_terminated:false`. -- **Rule 5 (assume solvable).** A localization that moves the gripper into thin - air means you picked a wrong pixel (a reflection, a rim, a decoy object under - the Object perturbation), not that the cell is unreachable. Re-look, re-pick, - re-`back_project` before concluding failure. - -## 6. Existing corpus +Never assume the initial coordinate remains valid after contact, release, or a +fixture interaction. -``` -robots/libero/guides/ -├── pro_hybrid_guide.md <- this file -├── strict_hybrid_guide.md <- perception protocol + Rules (source of truth) -└── env_calibration.md <- OSC frame bounds + safe altitudes -scripts/ -└── liberopro_register_perturbations.patch -resources/libero/memory/ <- MEMORY.md index + feedback_*/project_* notes -resources/libero/results_spatial_pert/ <- read-only seed-0 reference corpus -resources/libero/results_{object,goal,10}_pert/ <- same, other suites (seed-0) -``` +## Contact Tasks -Before you start, **read the auto-memory**: `resources/libero/memory/MEMORY.md` -(one-line hooks, auto-injected via CLAUDE.md). For perception PRO cells always -open `feedback_no_teleport_rule.md` and — for any `_swap` cell — -`feedback_swap_perturbs_fixtures.md` (what swaps, and why you re-find the -relocated fixture visually). For bowl→plate spatial tasks also read -`feedback_bowl_eef_y_offset.md`; for cluttered picks, `feedback_pi0_pick_full_prompt.md`; -after two failed retries, `feedback_failure_forensics.md`. The -`resources/libero/results_*_pert/` recipes are **inputs** (technique priors) — -consult them for prompt ladders, staging, and target zones, but never reuse their -coordinates (re-derive every xyz from THIS scene) and never write there. - -## 7. What to do next (priority order) - -1. **Extend spatial to all 10 tasks at seed 0**, four perception cells each - (base / `_task` / `_swap` / `_lan`). For hybrid runs, use the seed-0 reference - recipes in `resources/libero/results_spatial_pert/` as *technique* starting - points — the pick step is usually identical; the place target changes for - `_swap`, the target object changes for `_task`. Never reuse their coordinates. -2. **Scale to seeds beyond 0** (50 trials per task). Recipes must re-localize per - scene; a perception recipe is *data-flow* (perceive → plan → act), never - hard-coded xyz. -3. **Replicate on `libero_object`, `libero_goal`, `libero_10`.** Frame split - applies (read `states.json[0].state.robot0_eef_pos[2]`). For - `libero_goal_swap`, apply §3.5 — localize the **swapped fixture** visually. -4. **Aggregate into a main table** `(suite × perturbation × {Pi0, hybrid})`. The - headline number is the conditional: of the seeds Pi0 fails on, what fraction - does the *perception* hybrid solve? Because there is no oracle state anywhere in - the hybrid's reasoning here, this is the strongest single statistic the agentic - decomposition can claim. - -## 8. Quick reference for a brand-new session - -```bash -# 1. Sanity-check the liberopro patch (perturbed language must show) -LIBERO_TYPE=pro python -c \ - "import liberopro.liberopro.benchmark as b; print(b.get_benchmark('libero_spatial_task')().get_task(0).language)" -# -> must read 'Pick the akita black bowl not between ...' (the perturbed text) - -# 2. Read the auto-memory: resources/libero/memory/MEMORY.md - -# 3. Launch a perception cell (runner owns env_server; single-attempt) -python rpent/cli/main.py --env libero --suite libero_spatial_swap --task --seed 0 \ - --libero-type pro --planner claude_code --model claude-opus-4-8 -``` +Use `pi0_doubled` for learned contact behavior such as turning a knob or +opening/closing a drawer. Its success flag mirrors the benchmark termination +predicate, so an intermediate contact can be useful even when success is false. +Inspect state and image evidence after every contact attempt. + +Use short scripted alignment motions around contact skills. Avoid long blind +pushes, which can destabilize MuJoCo or move the end effector into an invalid IK +branch. + +## Planning Across Multiple Objects + +For multi-object tasks: + +- process objects in an order that minimizes collision and occlusion; +- preserve already-correct placements; +- route later carries around placed objects rather than over them; +- re-confirm each source and destination immediately before use; +- verify intermediate relations visually rather than assuming they held. + +When order is specified by the task, obey it even if another order appears +easier. + +## Reference Results And Memory + +Reference results and memory entries provide reusable strategy: + +- prompt ladders; +- manipulation ordering; +- safe approach and carry heights; +- fixture-specific contact patterns; +- known failure and recovery modes. + +They do not provide valid coordinates for the current run. Re-derive all +positions through perception. + +## Outcome And Audit + +The current benchmark outcome is top-level `libero_terminated` in the latest +environment tool result. It is not stored inside `state`. + +The audit should record: + +- exact perturbed `task_language`; +- perturbation type when known; +- semantic identification evidence; +- agentview and accepted wrist localization results; +- primitive sequence and recovery decisions; +- memory files consulted; +- final state and `libero_terminated`. + +Recipe export is runtime-managed from recorded primitives and successful +segmentation events. Artifact identifiers returned by tools are for audit and +traceability, not manual file access. + +## Quick Checklist -Then, inside the run: - -4. `view_driver_state({"step": 0})` → read `state.robot0_eef_pos[2]` to pick the - frame (§3.1). `Read images_cam/image_cam_00.png` (or hi-res) and - `view_camera_meta`. Localize the target (and, for `_swap`, the relocated - object/fixture) with `back_project` — run the mandatory pre-task perception pass - (§3.6c). Plan, then execute one structured tool at a time. -5. `write_text_file` the audit to `{output_dir}/{recipe_tag}.json` - (`regime: strict_perception`) before `finish`; the recipe - `{output_dir}/recipe_{recipe_tag}.jsonl` is exported automatically by the runner. - -When in doubt about *how to localize* or a primitive, the source of truth is -[`strict_hybrid_guide.md`](./strict_hybrid_guide.md); about *PRO setup / -perturbation semantics*, see -[`scripts/install_libero_pro_plus.sh`](../../../../scripts/install_libero_pro_plus.sh) -and §2. +- Read strict guide and relevant memory. +- Call `view_driver_state({"step": 0})`. +- Select the scene frame from initial EEF z. +- Read top-level `task_language`. +- Build the complete perception table. +- Re-derive all coordinates for this scene. +- Use agentview for identity and wrist for consistent refinement. +- Re-localize after contacts and placements. +- Check top-level `libero_terminated` after each relevant action. +- Write an honest audit and call `finish` without resetting. diff --git a/robots/libero/guides/strict_hybrid_guide.md b/robots/libero/guides/strict_hybrid_guide.md index 683feabe..3518a0a7 100644 --- a/robots/libero/guides/strict_hybrid_guide.md +++ b/robots/libero/guides/strict_hybrid_guide.md @@ -1,611 +1,206 @@ -# Strict Hybrid LLM + Pi0.5 — Perception-Isolated Guide - -You are taking over a hybrid LIBERO experiment in **perception-isolated** mode -— the only mode in this repository (the legacy oracle-state mode, where the -state JSON carried GT object coordinates, is not included here). - -> **Pi0.5 only does the grasp (`pi0_pick`). The LLM (you) handles every motion -> (`move_to`), every release, sequencing, retries — and you do not get GT -> object coordinates. You localize objects yourself from the depth + camera -> calibration the toolkit dumps each step.** - -## What's different from legacy oracle mode (read this first) - -| | legacy oracle mode (not in this repo) | **perception (this guide)** | -|---|---|---| -| how object coords are withheld | (none — full GT coords in `state`) | **the runner withholds them: `states.json` carries `object_names` only, no coordinates** | -| `states.json` objects | full `objects:{name:[x,y,z]}` | **`object_names:[…]` only — NO coords** | -| how you learn the task | env prompt / BDDL | **`states.json[NN].task_language`** (authoritative `:language`, coord-free) — never scrape the BDDL | -| extra obs artifacts | `images/image_NN.png` only | **+ `images_cam/`, `depths/`, `world/` (agentview); `images_wrist/`, `depths_wrist/`, `world_wrist/`, `wrist_meta/` (wrist); top-level `camera_meta.json`; hi-res `images_cam_hi/`, `world_hi/`, `images_wrist_hi/`, `world_wrist_hi/`** | -| cameras | agentview only | **agentview (fixed, ~1m → ±8–13 cm) + eye-in-hand wrist (moves with gripper, ±1–2 cm when <20 cm to target)** — coarse→fine, see below | -| how you get an object's xyz | read `state["objects"][name]` | **pick the object pixel, then `back_project({"row":ROW,"col":COL,"step":NN})` (K⁻¹+extrinsic already done); refine with the wrist map up close** | -| how you confirm the grasp | GT object-lift oracle | **grasp-only, no oracle — you judge the grasp from gripper width + wrist cam (Rule 1 / 1b)** | -| how you pick which of two identical objects | read their distinct coords | **by SPATIAL RELATION from `task_language` (elevation / left-right), never by `_1`/`_2` name** | -| cell budget | 600 (short suites) | **1200** — perceptual localization + manipulation is slower | -| audit `regime` | `strict` | `strict_perception` | - -How localization works (the core of this mode): - -The toolkit already back-projects EVERY pixel for you into world coordinates and -saves them as `world/world_NN.npy` (agentview) and `world_wrist/world_wrist_NN.npy` -(wrist) — both in the SAME world frame. **You do NOT write back-projection math; -you pick a pixel and call `back_project`, which indexes the map for you.** - -> **Path convention:** below, a bare file name refers to the file inside its own -> `output_dir` subdirectory — e.g. `image_cam_hi_NN.png` lives at -> `images_cam_hi/image_cam_hi_NN.png`, `world_NN.npy` at `world/world_NN.npy`, -> `image_wrist_hi_NN.png` at `images_wrist_hi/image_wrist_hi_NN.png` (see the obs -> artifacts row above for the full directory list). `NN` is the zero-padded step -> index. In practice you rarely hand-build these paths — each primitive tool -> already returns the resolved `image_cam_hi_path` etc. for the new step. - -1. Read `image_cam_hi_NN.png` — agentview RGB **in the calibration frame** - (vertical-flip of the raw buffer). This is the image you pick pixels in. - `image_NN.png` is the Pi0-rotation frame; **do not pick pixels there**. -2. Find the target object visually → pixel `(row, col)`. -3. Read its world xyz directly: `back_project({"row":ROW,"col":COL,"step":NN})`. - Sample 3–5 pixels on the object's top surface and median the returned xy - (robust to one mis-picked rim/edge pixel). That's the object's surface point - in world frame. - -### Hi-res perception channel (ON BY DEFAULT, 1024×1024) - -The toolkit dumps, every step, a 1024×1024 pair per camera IN ADDITION to the -256 files: `images_cam_hi/image_cam_hi_NN.png` + `world_hi/world_hi_NN.npy` -(agentview) and `images_wrist_hi/image_wrist_hi_NN.png` + -`world_wrist_hi/world_wrist_hi_NN.npy` (wrist). -**PREFER the hi pair for looking and identification** — a far object spans 4× -more pixels, and package text/label art is actually legible (e.g. "Cream -Cheese" vs "BUTTER" boxes are readable at 1024 and indistinguishable smears -at 256). If the hi files are absent, everything below works unchanged at 256. - -- A hi-res pixel `(row,col)` indexes ONLY the hi-res world map (`back_project` - default `resolution:"high"`; float16, same world frame). Never index a 1024 - pixel into the 256 map or vice versa — pass `resolution:"low"` only for a - pixel taken from a 256 image (convert by dividing/multiplying by 4 if needed). -- The `segment` command automatically uses the hi-res frame when present - (its `centroid_pixel`/`box` are then in 1024 coords). -- Metric accuracy of mask-median localization is the SAME at both resolutions - (the residual ~2 cm error is the surface-vs-center offset, not pixel size) — - the hi channel is for **identification**, not for replacing the wrist-cam - fine-localization protocol. -- Disk note: hi files keep only the LAST ~5 steps (rolling window); the 256 - history is complete. If the hi files are absent, everything below works - unchanged at 256. - -### Two cameras — agentview = IDENTITY, wrist = GEOMETRY - -**Roles are NOT symmetric (this is the core discipline, from the 80-task -localization sweep):** -- **agentview** (`image_cam_hi_NN.png` + `world_hi/world_hi_NN.npy`, ~1 m away): - the **semantic IDENTITY authority** — decides WHICH object/surface satisfies - the task language + spatial relation. At 1024 a far object spans enough pixels - to read labels/shape. Its metric precision is only ±8–13 cm, but identity is - its job, not millimetres. -- **wrist** (`image_wrist_hi_NN.png` + `world_wrist_hi/world_wrist_hi_NN.npy`, - <20 cm → ±1–2 cm): a **GEOMETRY refinement** camera for the SAME candidate - agentview already chose. It is near-vertical and **bad at identity** — it - cannot read side labels or tell duplicate/similar items apart, and in failed - probes it locked onto a look-alike hundreds of pixels away. **NEVER let the - wrist freely re-identify a non-basket target.** - -Protocol (non-basket objects/surfaces): -1. **Identity (agentview)** — in `image_cam_hi_NN.png` choose the target by RGB / - label / shape / global spatial relation; sample 3–8 pixels on it and median - `back_project` → the **identity anchor** (rough xy ok). -2. **Approach** — `move_to` ~15–20 cm directly above that anchor xy (toolkit - re-renders both cams after every primitive). -3. **Geometry refine (wrist)** — pick the SAME candidate's pixel in - `image_wrist_hi_NN.png`, `back_project` it with `camera:"wrist"`. **Accept - this corrected xy ONLY if it is within ~3–5 cm of the agentview anchor**; if - it jumps >5 cm, REJECT it (it hit a look-alike/background) and keep the - agentview xyz. The wrist may sharpen coordinates but may NOT override the - agentview semantic choice. Never average the two. -- **Basket / cavity special case:** for `basket`/cavity the wrist MAY also - confirm/refine the true interior centre (basket failures are rim/edge bias, - not semantic confusion). - -### Mandatory pre-task perception pass — localize EVERYTHING, THEN act - -Before ANY pick/place, build a localization table in your reasoning — one row -per task-relevant entity (every movable target, every destination/support/ -fixture, every relation landmark named by `task_language`), each with: -`name_or_role · agentview_evidence (why this candidate) · agentview_pixels · -agentview_xyz (median back_project on hi-res) · wrist_refine -accepted|rejected|basket_confirmed · final_xyz · uncertainty`. Do the -**FINAL READY CHECK** (every entity has a final xyz; non-basket wrist -refinements are spatially consistent with agentview; basket points are -interior-centred) — only then start manipulating. Rationale: in single-attempt -mode a **wrong-target first grab is unrecoverable** (it tips/displaces both the -grabbed object and the target zone), so the cheap insurance is to identify all -entities up front instead of recovering later. - -> The underlying math (already done for you): `cam = [(col−cx)·z/fx, -> (row−cy)·z/fy, z]` with `z=depth[row,col]`, then `P_world = -> extrinsic_cam2world @ [cam,1]`. You must invert `K` BEFORE the extrinsic — the -> old `E @ [col·z, row·z, z, 1]` recipe (no `K⁻¹`) is **wrong** (metres off). -> The `world/` and `world_wrist/` maps bake in the correct `K⁻¹`+extrinsic, and -> `back_project` reads them for you — that is precisely why `back_project` -> exists; you never write this math yourself. Forward world→pixel was verified -> 5/5 vs GT at design time (plate Δ=6 mm). The wrist map may be all-table/null -> early (the wrist only sees gripper + table until you move it over a target). - -## Before you start: READ THE AUTO-MEMORY - -Operating wisdom lives in the in-repo memory: +# Strict Hybrid LLM + Pi0.5 Perception Guide +This guide describes the perception-isolated LIBERO runtime. The agent receives +robot proprioception, object names, task language, and camera observations. It +does not receive privileged object coordinates. + +Pi0.5 performs grasping through `pi0_pick`. The LLM owns semantic perception, +localization, scripted motion, release, verification, and recovery within the +current episode. + +## Runtime Contract + +Every motion primitive appends a `StepRecord`. The record contains: + +- `step_idx`: sequential motion-step number; +- `state`: robot proprioception and coordinate-free scene information; +- `artifacts`: sorted logical base names for all files captured at the step; +- `command`, `result`, and `elapsed_s`; +- `extras`: LIBERO outcome fields such as `task_language`, + `libero_terminated`, and `episode_truncated`. + +Artifact storage is private to `EnvState`. Do not construct filenames or read +observation files manually. Use the structured tools: + +- `view_driver_state`: state, artifact names, log, and embedded images; +- `view_camera_meta`: calibration data for one camera and step; +- `back_project`: matching image pixel or region to world coordinates; +- `segment`: text- or point-prompted mask plus projected `world_xyz`. + +Step `0` is the initial observation. Step `-1` means the latest observation. +When a tool omits `step`, its schema default is `-1`. + +`states.json` is an internal versioned manifest. It is not an agent-facing state +API and should not be indexed directly. + +## Camera Roles + +`view_driver_state` embeds the best available images for the selected step: + +- **policy image**: the Pi0-oriented agentview input; +- **agentview image**: fixed global view, high resolution when available; +- **wrist image**: eye-in-hand view, high resolution when available. + +Use agentview for semantic identity and global relationships. Use wrist for +close-range geometry after the target identity has already been anchored in +agentview. The wrist view must not freely replace the semantic decision when +similar objects are nearby. + +Pixels passed to `back_project` must come from the same camera and resolution +specified in the call. The tool selects the matching world map internally. + +## Initial Perception Pass + +Before the first manipulation primitive: + +1. Call `view_driver_state({"step": 0})`. +2. Read the returned top-level `task_language` verbatim. +3. Identify every movable target, destination, support, and relation landmark + named by the task. +4. Inspect the embedded agentview image and classify candidates by color, + shape, label, container type, and spatial relation. +5. Localize each chosen candidate with several interior pixels and + `back_project`, or use `segment` when a stable text/point prompt exists. +6. Record a working table with semantic evidence, sampled pixels, projected + coordinates, uncertainty, and the selected final coordinate. + +Do not start manipulation until every required target and destination has a +defensible identity and coordinate estimate. + +## Semantic Identity Before Geometry + +Depth and world coordinates cannot distinguish visually similar surfaces. A +plate, stove burner, pot lid, and cabinet top can all be flat circular or planar +regions at similar heights. + +Classify the destination in RGB before localizing it: + +- **plate**: ceramic disc with a clean rim, often white or ringed; +- **stove region**: darker fixture surface, coil, grate, or burner pattern; +- **basket**: open container whose interior center differs from the rim; +- **cabinet or drawer**: fixture geometry associated with the task noun. + +When duplicate objects exist, choose by the relation in `task_language`, not by +internal object suffixes. + +## Coarse-to-Fine Localization + +For each non-basket object: + +1. Select the semantic candidate in agentview. +2. Project three to eight interior pixels and take a robust median. +3. Move approximately 15-20 cm above the agentview anchor. +4. Inspect the wrist image and refine the same physical candidate. +5. Accept the wrist estimate only when it remains within roughly 3-5 cm of the + agentview anchor. Otherwise reject it and retain the global estimate. + +For baskets and cavities, use agentview for global identity and wrist for the +interior center. Avoid rim pixels. `back_project` region mode can summarize a +bounded pixel window with optional `z_min` and `z_max` filtering. + +## Segmentation + +`segment` does not move the robot. Supply exactly one of: + +```json +{"prompt": "the black bowl on the stove", "camera": "agentview", "step": -1} ``` -resources/libero/memory/MEMORY.md + +```json +{"point": [420, 615], "camera": "agentview", "step": -1} ``` -Scan the ~30 one-line hooks. For perception cells **always** open: -- `feedback_no_teleport_rule.md` — the deleted primitives. -- `feedback_redo_cell_timeout_1200.md` — why the cell budget is long here. -(The grasp is oracle-free too: there is no GT-lift oracle — you judge the grasp -from gripper width + the wrist cam, see Rule 1 / 1b.) +The result includes the mask score, projected `world_xyz`, logical artifact +identifiers for audit, and an embedded overlay image when available. Inspect +the overlay before trusting the projection. -For bowl→plate spatial tasks always also read `feedback_bowl_eef_y_offset.md` -(bowl-eef y-offset 4.5 cm: place at `eef_y = plate_y + 0.045`, not `plate_y`). +Use plain visual phrases. Remove benchmark or brand names that the segmenter +cannot ground. If segmentation fails, choose pixels manually and call +`back_project`. -## Rule 1 — `pi0_pick` is grasp-only (no oracle) +## Grasping -`pi0_end_to_end` is FORBIDDEN. `pi0_pick` is for the grasp only; you script every -`move_to` and every `release`. Use: +Pre-position above the localized object before calling Pi0.5: -``` -pi0_pick({ - "prompt": "", +```json +{ + "prompt": "pick up the short red-label can", "max_chunks": 20, "lift_thresh": 0.05, "gripper_closed_thresh": 0.06 -}) +} ``` -`pi0_pick` takes **no object-tracking / oracle argument** — it is grasp-only -and reads NO GT object pose. Passing a name would do nothing even if you tried. -You judge "did I grab the target?" yourself (Rule 1b). NEVER let Pi0 finish the -place — YOU do every `move_to` and the `release`. - -## Rule 1b — JUDGE THE GRASP from perception, NOT from a name - -After a pick, decide "did I grab the target?" from two coord-free signals: - -- **Gripper** (`state.robot0_gripper_qpos` from the latest `states.json` - entry): fingers closed but NOT fully shut (~0.01–0.05 gap) ⇒ holding an - object; fully closed (~0.0) ⇒ grasped air. -- **Wrist cam**: after the lift, read `image_wrist_hi_NN.png` — the target - should be raised into the gripper and the spot it came from now empty (its - surface z jumps up by your lift distance). If needed, `back_project` wrist - pixels to confirm the surface z jumped. - -`pi0_pick`'s returned `success` (an eef-lift + gripper-closure heuristic, pure -proprioception) is a HINT, not proof — confirm with the wrist cam before -carrying. The only authoritative TASK-success signal is -`state.libero_terminated` (the benchmark predicate), which is name-independent. - -## Rule 4 — NO TELEPORT primitives (physics-only) - -`set_object_pose`, `articulate_to`, `js_move_to`, `carry_object` are **deleted -from the codebase**. They are not callable and are not in the tool list. If a -goal is past OSC reach and no physical approach works, write an honest -`libero_terminated:false` audit — never warp. - -## Rule 0 — Use images for reasoning, not just JSON state - -After every primitive tool call, inspect the returned state and image paths -(the tool returns them — no separate read needed). Read the new -`image_cam_hi_NN.png` path (calibration frame — the one you pick pixels in) -and, when close to a target, the `image_wrist_hi_NN.png` path. Even more -important here than in oracle mode: this is *the* signal you use to find -objects. - -`image_cam_hi_NN.png` (and its 256 counterpart `image_cam_NN.png`) is the -calibration-frame RGB — same scene as `image_NN.png` but vertically flipped so -that pixel coordinates align with the camera matrices in `camera_meta.json`. -Pick object pixels from these; `states.json` alone gives only proprioception + -object names. - -## Rule 2 — SINGLE EPISODE, NO RESET - -This is a **one-shot** evaluation: you get **exactly ONE episode**. Do NOT reset -and do NOT restart the episode. You MAY recover *within* this one episode -(re-localize — objects may have moved; re-pre-position; re-`pi0_pick` a missed -grasp; walk the next rung of the Pi0 prompt ladder in Rule 3; re-firm the grip; -`rotate_pitch`/`move_pose`) — that is all one continuous attempt. But the -instant you would want to start over, **STOP instead and write the audit** -(success or an honest `libero_terminated:false`), then call `finish`. Never -warp; never reset. - -## Rule 5 — Assume every task is physically solvable - -Same as oracle mode. A localization that "looks right" but moves the gripper -into thin air usually means your pixel was on a wrong surface (e.g. picked the -bowl's reflection on the table). Re-look at `image_cam_hi_NN.png`, pick a -different pixel firmly on the target's top, re-`back_project`. Don't conclude -"unreachable" until you've validated localization. - -## Rule 6 — Your task is `task_language`; the BDDL is FORBIDDEN - -Each `states.json[NN]` entry carries a **`task_language`** field — the -authoritative task instruction (the BDDL's `:language` tag, which contains -**no** coordinates). **Read it and obey it verbatim.** Do not infer the task -from object names, from sibling recipes, or by guessing a `task_map` index — -that produced *wrong-task* runs (an agent solving a task it was never assigned). - -> ⚠️ **Never read the BDDL files, import the benchmark, or query env object -> poses.** The BDDL is the one place the task language *and* the `:init` -> ground-truth coordinates live together — reading it to get the language also -> leaks the coordinates this mode exists to withhold (a perception-isolation -> breach; observed on several swap cells in earlier runs). You already have the -> task from `task_language`; you get object **positions** ONLY by depth -> back-projection. The runtime strips coords from `state` but does **not** -> sandbox the BDDL — that discipline is on you. - -## Rule 7 — Ground the target by its spatial RELATION, not its name - -When the task names a relation ("the bowl **ON THE COOKIES BOX**", "the mug -**LEFT OF** the plate"), the target is whichever object SATISFIES that relation -in the scene — find it by perception. Identical objects (`akita_black_bowl_1` vs -`_2`) carry NO perceptual difference in their names, so the name can't choose -for you; the RELATION disambiguates: - -- "on the cookies box" ⇒ the bowl that is **elevated** (~0.03–0.06 m above the - table, on the box) — distinguish it by its higher world-z from `back_project` - vs the table-level bowl. -- "left/right/front/back of X" ⇒ compare back-projected world xy to X's xy. - -Pick the target purely from where things ARE. (You never need to know which -`_N` name the target is — no primitive in this mode asks for one.) - -## Mental model - -1. **The runner (`rpent/cli/main.py`) owns a long-lived env server** — Pi0.5 + a - single-env LIBERO sim. It launches and manages the server; you do NOT start, - stop, or restart it. -2. **You call one structured MCP tool per step.** The tool BLOCKS until the - toolkit runs that one primitive and dumps the new step, then RETURNS the new - state + log + image paths. There is no file bus and no polling — the tool's - return value IS your signal. Each dumped step appends an entry to - `states.json` (entry `[NN]`) and writes `images/image_NN.png` + - `images_cam/image_cam_NN.png` + `depths/depth_NN.npy` + `world/world_NN.npy` - (+ the wrist and hi-res dirs) and, once, top-level `camera_meta.json`. -3. **You read the returned state, localize via `back_project`, decide the next - move, and call the next tool.** - -## Launch a session - -The runner (`rpent/cli/main.py`) launches and owns the env server (Pi0.5 + single-env -sim) — do not start/stop it. You call MCP tools; begin by reading step 0 via -`view_driver_state({"step":0})`. - -## The perception artifacts you read each step - -| artifact | what's in it | -|---|---| -| `states.json` (entry `[NN]`) | `step_idx`, `libero_terminated`, **`task_language` (your authoritative task instruction — the BDDL `:language` tag, coord-free; obey it verbatim)**, `state.{robot0_eef_pos, robot0_eef_quat, robot0_gripper_qpos, object_names}`, and the merged `command`/`result`/`elapsed_s` for that step. **No object coordinates.** Read via `view_driver_state({"step": NN})` (omit `step` = latest). | -| `images/image_NN.png` | RGB in Pi0 frame (180° rotated). *Do not pick pixels here for back-projection.* | -| `images_cam/image_cam_NN.png` | agentview RGB in **calibration frame** (vertical flip). Pick object pixels HERE (256 grid → `back_project` with `resolution:"low"`). | -| `images_cam_hi/image_cam_hi_NN.png` | **HI-RES 1024×1024** agentview RGB, calibration frame. **PREFER this for looking / identification**; `back_project` its pixels at the default `resolution:"high"`. | -| `depths/depth_NN.npy` | `(256, 256) float32` agentview metric depth (m), calibration frame. Same row/col as `image_cam_NN.png`. | -| `world/world_NN.npy` | `(256, 256, 3) float32` — **precomputed agentview world xyz per pixel** (K⁻¹+extrinsic done). Prefer `back_project`; read this manually only for debugging. Fixed cam (~1 m). | -| `world_hi/world_hi_NN.npy` | `(1024, 1024, 3) float16` — precomputed agentview world xyz per hi-res pixel. | -| `images_wrist/image_wrist_NN.png` (+ `images_wrist_hi/`) | **wrist (eye-in-hand) RGB**, calibration frame. Moves with the gripper. | -| `depths_wrist/depth_wrist_NN.npy` | wrist metric depth (m). | -| `world_wrist/world_wrist_NN.npy` (+ `world_wrist_hi/`) | `(256, 256, 3)` / `(1024, 1024, 3)` **precomputed wrist world xyz per pixel**, SAME world frame as agentview. ±1–2 cm when <20 cm to target. May be all-table/null until you move over the target. | -| `wrist_meta/wrist_meta_NN.json` | wrist intrinsics + extrinsic **for THAT step only** (the wrist cam moves). Read via `view_camera_meta({"camera":"wrist","step":NN})`. | -| `camera_meta.json` (top-level) | agentview `intrinsic_K` (3×3), `extrinsic_cam2world` (4×4), `depth_near/far`, projection recipe. Read via `view_camera_meta({"camera":"agentview"})`. | -| `segments/segment_NN_XX.json` | (after a completed `segment` response) `{found, mode, prompt|point, camera, source_step, score, box, mask_shape, centroid_pixel, world_xyz, n_pixels}` — SAM3's top mask back-projected via the matching world map. `world_xyz` is a robust MEDIAN over the whole mask. A valid no-detection response carries `{found:false, error}`. `XX` is a per-step index. | -| `segments/segment_overlay_NN_XX.png` | (only after a successful `segment` call) the segmented mask tinted red on the source image — read it to confirm SAM3 grabbed the right object. | - -The command + its result + `elapsed_s` are merged INTO the `states.json[NN]` -entry (and echoed in each primitive tool's return value) — there is no separate -per-step log file to read. - -## Localization with `back_project` (coarse agentview → fine wrist) - -You index the precomputed map through `back_project` — no K⁻¹ math. Coarse -(agentview) then, once the gripper is parked over the target, fine (wrist): +Treat `pi0_pick.success` as a hint. Verify the grasp from: -``` -# COARSE: agentview hi-res (default resolution:"high") — choose which object / rough xy -back_project({"row": ROW, "col": COL, "step": NN}) +- end-effector lift; +- nonzero gripper opening consistent with holding an object; +- wrist or agentview evidence that the correct object moved with the gripper; +- an empty or changed source location. -# FINE: wrist — after move_to ~15-20cm above the target, pick its pixel in -# image_wrist_hi_NN.png and back_project it (±1-2cm; refines the SAME candidate). -back_project({"row": ROW, "col": COL, "step": NN, "camera": "wrist"}) -``` +If the grasp failed, recover in the same episode by re-localizing, +re-positioning, and improving the visual prompt. Do not reset in single-attempt +evaluation mode. -(Pass `"resolution":"low"` only when `ROW,COL` came from a 256 image.) - -Tips: - -- Sample 3–5 pixels on the object (centre + a couple of edge pixels) and - median the back-projected xy — robust to a single mis-picked pixel. Avoid - pixels on the thin rim/edge or the gap to the table (those index a - background/edge depth and give a world point metres away). -- The returned point is the **visible surface** under your chosen pixel. - For a flat object (plate, basket) that surface ≈ the place target. For a - bowl/bottle, the surface is the top of the object; the rim's xy is what - you want for `release`, not the grasp's eef_y (apply - `feedback_bowl_eef_y_offset` — bowl `eef_y_target = perceived_plate_y + 0.045`). -- For the **table z** (when you need a known floor to compare against), - `back_project` a pixel on bare table near the object. - -## The command vocabulary - -Call one structured tool per step. The full primitive set (only the *control -signals* are listed here) — every one blocks until the step is dumped and -returns the new state + log + image paths: - -```jsonc -// === physics-only primitives (the entire allowed set) ===================== - -// Scripted EEF servo. action_scale=0.05 is the env's units; step_clip caps -// per-step Δxyz (m) BEFORE division by action_scale — smaller = slower. -move_to({"xyz": [x, y, z], "gripper": -1, - "tol": 0.012, "step_clip": 0.025, "max_steps": 80, - "action_scale": 0.05, "target_yaw": null}) // gripper: -1 open, +1 close - -// Pi0.5 closed-loop pick. Grasp-only — no oracle/tracking arg (Rule 1). -// You judge the grasp from gripper width + wrist cam (Rule 1b). -pi0_pick({"prompt": "pick up the X", "max_chunks": 20, - "lift_thresh": 0.05, "gripper_closed_thresh": 0.06}) - -// Pi0.5 for a contact skill (knob turn, drawer/door open-close — rare here). -// success mirrors libero_terminated only; inspect image/state for intermediates. -pi0_doubled({"prompt": "turn off the stove", "max_chunks": 20}) - -// Open gripper to place. Triggers libero termination if the On/In predicate met. -release({"max_steps": 20}) - -// Hold pose + drive gripper. Use to firm a grip mid-carry. -set_gripper({"gripper": 1, "steps": 5}) // -1 open, +1 close - -// Wrist yaw (world-z). Provide target_yaw (absolute) OR delta_yaw (relative). -rotate_wrist({"target_yaw": 0.0, "gripper": 1, - "max_steps": 40, "tol": 0.02, "step_clip": 0.10}) - -// Tilt eef pitch (axis-angle X). Cavity entry / micro-aiming. target_pitch OR -// delta_pitch. Use before threading a narrow opening whose face normal is ±y. -rotate_pitch({"target_pitch": 0.9, "gripper": 1, - "max_steps": 40, "tol": 0.02, "step_clip": 0.10}) - -// Co-vary xyz + pitch + yaw — threads cabinet-front IK singularity that move_to -// walls at. gripper defaults to -1 (OPEN) — pass gripper:1 while holding. -move_pose({"xyz": [x, y, z], "target_pitch": 0.0, "target_yaw": 0.0, - "gripper": 1, "step_clip": 0.02, "pitch_step": 0.08, "yaw_step": 0.08, - "tol": 0.012, "ori_tol": 0.05, "max_steps": 150}) - -// SAM3-grounded localization — does NOT move the robot and does NOT -// replace manual back-projection. Segments the most-recent dumped image for the -// prompt, back-projects the mask through the matching world map, and writes -// segments/segment_NN_XX.json {score, box, centroid_pixel, world_xyz (robust MEDIAN over -// the whole mask), n_pixels} + segments/segment_overlay_NN_XX.png. Read world_xyz -// directly instead of eyeballing a pixel. camera "wrist" uses the wrist world -// map for fine refinement (move the eef over the target first). Pass -// "point":[row,col] for a point prompt instead of text; prompt and point are -// mutually exclusive. min_score default 0.2. -// RPent manages the SAM3 service. If one call fails or finds nothing, the result is -// an {"error":...,"fallback":...} dict — fall back to picking a pixel in -// image_cam_hi_NN.png and calling back_project. -segment({"prompt": "the black bowl on the stove", "camera": "agentview", - "point": null, "min_score": 0.2}) - -// === FORBIDDEN — DELETED FROM THE CODE — DO NOT EMIT ============ -// set_object_pose, articulate_to, js_move_to, carry_object -``` +## Scripted Motion + +Use short, staged `move_to` commands. Do not issue a single horizontal move +larger than approximately 0.30 m. Long commands can switch IK branches and move +the end effector into the wrong half-space. + +After every motion: + +1. Inspect the returned `result` and final distance. +2. Inspect the new state and embedded images. +3. Re-localize anything that may have moved. +4. Continue only when the observed state matches the plan. + +Use `rotate_wrist`, `rotate_pitch`, or `move_pose` when orientation matters. +Use `pi0_doubled` for short learned contact interactions such as knobs, +buttons, doors, or drawers. Alternate it with short, capped scripted alignment +motions; never use one long blind push. + +## Placement + +Carry objects at a collision-safe height. Reconfirm the destination before +release, especially when plate/burner or rim/interior confusion is possible. + +For open containers: + +- localize the interior rather than the nearest rim; +- place above the interior center; +- lower enough to avoid a high-energy drop; +- inspect the post-release observation and outcome flag. + +If the task predicate does not fire, reclassify the destination before assuming +the grasp failed. Wrong-surface placement is a common cause. + +## Completion And Audit + +The current outcome is the top-level `libero_terminated` value returned by the +latest environment tool. Do not look for it inside `state`. + +On completion, write the requested audit with: + +- suite, task, seed, and evaluation regime; +- exact memory files consulted; +- semantic identification and localization strategy; +- grasp and placement evidence; +- final state and `libero_terminated`; +- honest failure details when the task remains incomplete. + +Then call `finish`. Recipe generation is handled by the runtime from recorded +primitive and successful segmentation events. + +## Final Checklist -### How to use `segment` (SAM3) — practical tips - -`segment` is the fastest way to localize: one call gives you a `world_xyz` -that is a robust MEDIAN over the whole object mask (hundreds of pixels), which -beats eyeballing 3–5 pixels. Workflow: - -1. Read `image_cam_hi_NN.png`, decide the target by the task's spatial RELATION. -2. Call `segment` with a **plain visual phrase + the relation**, then read the - returned `world_xyz` and the `segments/segment_overlay_NN_XX.png` to confirm the mask - is on the right object before you move. -3. **Two camera views — YOUR choice via the `"camera"` field** (`segment` works - on either; default `"agentview"`): - - `"camera":"agentview"` — the fixed ~1 m cam (±8–13 cm). Use it to choose - WHICH object (global layout / spatial relation) and get a rough xy. - - `"camera":"wrist"` — the eye-in-hand cam (±1–2 cm), reads the wrist world - map in the SAME world frame. Use it to PRECISELY localize once the gripper - is parked over the target. ⚠ It is null / all-table until you `move_to` - ~15–20 cm over the target (the wrist only sees the gripper+table before - that), so segment agentview first, approach, THEN segment wrist. - Typical coarse→fine: agentview select → `move_to` above → wrist refine → grasp. - You decide when the wrist refinement is worth it (small / closely-spaced / - identical objects benefit most; a big isolated plate may not need it). - -**Prompt phrasing (important — SAM3 is sensitive):** -- ✅ Use plain colour + shape + spatial relation: `"the black bowl on the stove"`, - `"the white plate"`, `"the bowl on the cookies box"`. -- ❌ Do NOT use the object's internal/brand NAME from `object_names`/BDDL: - `"the akita black bowl"` scores ~0.03 (no detection) because SAM3 can't ground - "akita"; the same object as `"the black bowl on the stove"` scores ~0.76. Strip - proper nouns (`akita`, `glazed_rim_porcelain_…`) — say what it *looks like*. -- For two identical objects, the relation in the prompt (`…on the stove`, - `…left of the plate`) usually steers SAM3 to the right instance; verify via the - overlay and the world-z (an object *on* a fixture has a higher z than a - table-level twin). The two-camera relation protocol still applies when SAM3 - can't disambiguate from text alone. -- If `segment` returns `{"error":...}` (low score / service down), walk the - prompt (drop the brand word, simplify, add the relation) or fall back to - manual pixel → `back_project`. - -## The strict-hybrid recipe (perception variant) - -A typical bowl→plate cell looks like: - -1. `view_driver_state({"step":0})`; inspect `task_language`, - `image_cam_hi_00.png`, and `camera_meta.json`. -2. Identify the target by the task's spatial RELATION, not its name (Rule 7). -3. Localize it — pick its pixel in `image_cam_hi_00.png`, - `back_project({"row":ROW,"col":COL,"step":0})` (coarse). Refine with the - wrist cam after pre-positioning (fine). -4. **Pre-position** ~15–20 cm above it: `move_to([obj_x, obj_y, carry_z])`, - gripper open. Re-`back_project` a wrist pixel here and correct the xy before - grasping. -5. `pi0_pick` with the right prompt (start "pick up the {object}", escalate per - the Pi0 ladder — see below). Confirm the grasp via gripper + wrist (Rule 1b). -6. `set_gripper({"gripper":1,"steps":5})` to firm the grip. -7. Localize the placement region (basket / plate / drawer slot) the same way. -8. `move_to([place_x, place_y, carry_z])` to traverse at constant height. -9. Optionally descend `move_to([place_x, place_y, place_z])`. -10. `release` — predicate (`On`/`In`) checks → `libero_terminated=True` if hit. -11. Light retreat (`move_to` upward) so the next step's image is clean. - -> **Predicate fire timing.** Most LIBERO `On(X, Y)` predicates fire on -> `release` if `X` is above `Y`'s region. `In(X, container)` needs `X` to -> have actually entered the container's volume before release. If a release -> fires `term=False` but the object is on top of the right region, descend -> 1–2 cm more and re-`release`. - -## Pi0 prompt ladder (Rule 3 — Pi0 IS the delivery service) - -Try in order; each rung uses a slightly more specific prompt or a re-pre-pos. - -1. `"pick up the {object}"` — generic sub-instruction. -2. The full `task_language` verbatim (e.g. `"Pick the akita black bowl on the - cookies box and place it on the plate"`). -3. Add a spatial qualifier (`"…on the wooden cabinet"`, `"…next to the - basket"`). -4. Re-pre-position 5 cm lower or shifted, then retry rung 2 or 3. - -Empirically Pi0 sometimes needs the **full prompt with the spatial -qualifier** for elevated picks (stove, cabinet-top, drawer). See -`feedback_pi0_pick_full_prompt.md` and the prompt-ladder note in MEMORY. - -## Key hyperparameters - -- Single-step `xyz` within ±0.30 m of current eef or OSC flips IK. Split - long traversals into 2–3 carry-z waypoints. -- `lift_thresh`: 0.05 (flat/stable) / 0.08 (slippery tall bottles). -- `step_clip`: 0.025 (empty / box) / 0.015 (cans) / 0.012 (tall bottles). -- Frame z (from `state.robot0_eef_pos[2]` at step 00): - ≈ 0.68 → LIVING_ROOM, ≈ 1.17 → KITCHEN, ≈ 0.26 → OBJECT. -- BOWL: `eef_y_target = perceived_plate_y + 0.045` (bowl-eef y-offset). -- TALL BOTTLES: carry at `z=0.30`, release without descending. -- Approach high-then-vertical; recover by re-`pi0_pick`, not by hovering. - -## Reading state - -After every primitive tool call, the return value already carries the new -`state`, the merged `command`/`result`/`elapsed_s` log, and the image paths — -you do **not** need a separate read. When you need an older step, call -`view_driver_state({"step":NN})` (omit `step` for the latest): - -1. Check `result.success`, `final_dist_m`, `peak_lift_m`, etc. in the returned - log. -2. Read the returned `image_cam_hi_NN.png` path → visual confirmation; pick - pixels for any new localization. -3. Check `state.robot0_eef_pos` / `robot0_gripper_qpos` / `libero_terminated` - in the returned `state`. - -You **do not** open the depth maps yourself — feed a pixel to `back_project`. -Don't call `view_driver_state` immediately after a primitive that already -returned the new state. - -## Common failure modes - -- **Pick missed (gripper closed empty).** `result.peak_lift_m` < `lift_thresh`, - `min_gripper_opening` ≈ 0. Re-pre-position 1–2 cm lower or shifted; retry - Pi0 with next prompt-ladder rung. -- **Object slipped mid-carry.** `release` returns `term=False` and the object - is no longer where you `release`d it. `release`, re-pre-pos above it, - `pi0_pick` again, traverse again. -- **OSC stuck.** `move_to` returns `final_dist_m > 0.05` at `max_steps` and - same xy twice → try `rotate_pitch`, split into more waypoints, or - `move_pose` (co-varying) for the cabinet-front singularity. -- **Placement off because localization was wrong.** The release puts the - object on bare table instead of on the plate. Re-read `image_cam_hi_NN.png`, - pick a different pixel on the target region (sample 3 pixels on the plate's - flat top, median the back-projected xy), redo. Tip: if depth at your chosen - pixel is much closer than the table z, you picked the camera's near edge / - an object rim — pick again. - -## Verifying strict compliance - -Before saving the audit, confirm your command history is physics-only. The -teleport primitives are not even in the tool list, but audit it anyway: read -the `command` field of every `states.json` entry (via -`view_driver_state({"step":NN})` per step, or `read_text_file` on -`{output_dir}/states.json`) and confirm every `command.action` is one of the -allowed physics primitives (`move_to`, `pi0_pick`, `pi0_doubled`, `release`, -`set_gripper`, `rotate_wrist`, `rotate_pitch`, `move_pose`) — no -`set_object_pose` / `articulate_to` / `js_move_to` / `carry_object` appears. - -## Persisting successful runs as audit JSONs - -When `state.libero_terminated == true`: - -a. The working command recipe (`{output_dir}/recipe_{recipe_tag}.jsonl`) is - **auto-exported by the runner** from the non-error primitive commands in - `states.json` plus successful segment calls recorded in - `segments/segment_*.json`, merged in execution order — you do NOT - hand-write it. -b. Write a minimal audit JSON with `write_text_file` to - `{output_dir}/{recipe_tag}.json` with at least: `suite`, `task_id`, `seed`, - `regime: "strict_perception"`, `strategy_notes` (mention HOW you localized — - which pixel, depth, back-projected world xyz), `pick_result` (the `result` - from your `pi0_pick`), `final_state` (the latest `states.json` entry's - `state` field), `libero_terminated: true`. -c. Call `finish({"status":"success","summary":"…"})`. - -If unrecoverable after honest exploration in this one episode, write -`{output_dir}/{recipe_tag}.json` with `libero_terminated: false` + -`strategy_notes` describing what you tried, the back-projected xyz you used, and -which step failed. Then call `finish` (NO reset, NO second attempt). - -> The `regime: "strict_perception"` tag distinguishes these audits from the -> oracle-state `strict` regime in mixed-mode datasets. - -## Iteration heuristics - -- After 2 failed retries on the same step, **stop tuning numerics** and - inspect images: read `image_cam_hi_NN.png` at pick / pre-release / - post-release. The visual disagreement is usually the bug (you picked the - wrong pixel, or Pi0 grabbed the decoy). This is the lesson of - `feedback_failure_forensics.md` — applies even more strongly here. -- If a release reports `term=False` but the object is visibly on the target - region, descend 1–2 cm more and `release` again; predicate often needs - contact, not just hover. -- Once you use `back_project`, trust it: if a localization "feels off", it's - because you picked the wrong pixel — not because the depth / calibration is - wrong. (`back_project` inverts `K` before the extrinsic; the raw - `E @ [col·z, row·z, z, 1]` form skips it and is wrong — which is why you use - `back_project` and never hand-roll the math.) - -## What "strict_perception" means concretely - -- **No GT object coordinates anywhere in your reasoning.** The `state` you - read has none; the only legitimate sources of object xyz are the camera - images + depth + the precomputed `world/` and `world_wrist/` maps (via - `back_project`). -- **No teleport primitives.** The four are deleted. -- **Pi0 only does the grasp.** You script every motion + release. -- **Fully oracle-free, including the grasp.** There is no GT-lift oracle — - `pi0_pick` reads NO GT object pose and takes no tracking argument. You judge - the grasp from gripper width + the wrist cam, and TASK success from - `state.libero_terminated` (the benchmark predicate). -- **Single attempt.** One episode, no reset (Rule 2). -- The expected audit `regime` is `strict_perception`. - -## Reference cases - -The seed-0 sweep results live under `resources/libero/results_*_pert/` -(`results_10_pert`, `results_object_pert`, `results_spatial_pert`, -`results_goal_pert` — PRO swap+task, t0–t9). Each solved cell has an audit JSON -+ a `recipe_{tag}.jsonl` command sequence. The consistent winning pattern: -localize → pre-pos → `pi0_pick` → `set_gripper` → move → `release` in 6–12 -commands. - -When you write a new audit, browse a sibling cell's `recipe_{tag}.jsonl` as a -*technique* template — but never paste its xyz; re-derive every position via -`back_project` from THIS scene's depth. - -Begin by reading `resources/libero/memory/MEMORY.md`, then call -`view_driver_state({"step":0})` and inspect the returned `image_cam_hi_00.png` -(+ `camera_meta.json` via `view_camera_meta`); localize the target object via -`back_project`, then plan and execute. +- Initial state read with `step: 0`. +- Task language copied from the returned tool result. +- All targets and destinations semantically identified in agentview. +- Coordinates obtained through `back_project` or `segment`. +- Wrist refinements checked against the agentview anchor. +- Grasp confirmed visually and proprioceptively. +- Long motion split into safe waypoints. +- Destination reclassified immediately before release. +- Latest result checked for top-level `libero_terminated`. +- Audit written before `finish`. diff --git a/robots/libero/prompts/system.py b/robots/libero/prompts/system.py index 16bd1e93..70591e30 100644 --- a/robots/libero/prompts/system.py +++ b/robots/libero/prompts/system.py @@ -10,7 +10,7 @@ > persistence / up to N attempts" instruction anywhere below).** This is a > ONE-SHOT evaluation: you get **exactly ONE episode**. You MUST NOT call > `reset`, and you must not restart the episode. Plan carefully, then execute -> your single best manipulation sequence toward `state.libero_terminated == true`. +> your single best manipulation sequence toward `libero_terminated == true`. > You MAY recover *within* this one episode (re-pre-position, re-`pi0_pick` a > missed grasp, walk the Pi0 prompt ladder, `rotate_pitch`/`move_pose`) — that is > all one continuous attempt — but the instant you would want to reset/start over, @@ -70,7 +70,7 @@ weak at reading side labels or distinguishing similar grocery items (ketchup/BBQ/tomato sauce, soup cans, cream cheese/butter). Do NOT let the wrist freely re-identify a non-basket target; it often locks onto a look-alike. - Instead: choose the target from `image_cam_hi_NN.png`, compute its agentview + Instead: choose the target from the high-resolution agentview image, compute its xyz, move over that candidate, project/track that SAME candidate in wrist, and refine only its surface/center coordinates. SAM3 scores ~0.02-0.06 on brand nouns ("alphabet soup", "tomato sauce") — prompt by colour+shape ("the short @@ -146,58 +146,28 @@ - Under some runtimes these same tools may appear namespaced; call the actual tool name shown in your tool list, preserving the same arguments and semantics. -The toolkit writes artifacts in `{{output_dir}}/`: - -- `{{output_dir}}/states.json` — top-level JSON array; each entry has - `step_idx`, `task_language`, `libero_terminated`, `state` (robot - proprioception + object_names; NO object coordinates), `command`, `result`, - `elapsed_s`, and world-map path fields when available. -- `{{output_dir}}/images/image_NN.png` — agentview RGB, 180°-rotated (Pi0 frame; - do NOT use for back-projection). -- `{{output_dir}}/images_cam/image_cam_NN.png` — agentview RGB in the CALIBRATION - frame; use for low-resolution pixel checks. -- `{{output_dir}}/depths/depth_NN.npy` — agentview metric depth (meters), - calibration frame. -- `{{output_dir}}/world/world_NN.npy` — HxWx3 precomputed world xyz per 256px - agentview pixel. Prefer `back_project`; read this manually only for debugging - or if the tool is unavailable. -- `{{output_dir}}/images_wrist/image_wrist_NN.png` — wrist RGB, calibration frame. -- `{{output_dir}}/depths_wrist/depth_wrist_NN.npy` — wrist metric depth (meters). -- `{{output_dir}}/world_wrist/world_wrist_NN.npy` — wrist world xyz map in the - SAME world frame as agentview. -- `{{output_dir}}/wrist_meta/wrist_meta_NN.json` — wrist intrinsics + extrinsic - FOR THAT STEP ONLY (the wrist cam moves, so it changes every step). -- `{{output_dir}}/images_cam_hi/image_cam_hi_NN.png` — HI-RES (1024x1024) - agentview RGB in calibration frame. USE THIS to inspect the scene and identify - objects — a far object spans 4x more pixels than at 256. -- `{{output_dir}}/world_hi/world_hi_NN.npy` — 1024x1024x3 float16 precomputed - world xyz per hi-res agentview pixel. Prefer `back_project`; if you manually - inspect it, never index a low-res pixel into this grid or vice versa. -- `{{output_dir}}/images_wrist_hi/image_wrist_hi_NN.png` / - `{{output_dir}}/world_wrist_hi/world_wrist_hi_NN.npy` — same hi-res pair for - the WRIST cam. - ⚠ Hi-res pixel (row,col) indexes ONLY the hi-res world map (and 256 pixel -> - 256 map). Don't mix grids; if you must convert, divide hi coords by 4. - ⚠ Hi-res files keep only the LAST 5 STEPS (disk); for older before/after - comparisons use the 256 files or `states.json` history. -- `{{output_dir}}/camera_meta.json` — agentview intrinsics K, cam->world - extrinsic, projection recipe. -- `{{output_dir}}/action_videos/step_NN_.mp4` — per-action clips generated - when the Dashboard is enabled. - -NN is zero-padded sequential (`00`, `01`, `02`, ...). Initial state step `00` is -dumped before you begin. Use `view_driver_state({"step": 0})` to read it.""" - -GOAL = """YOUR GOAL: produce `state.libero_terminated == true` in ONE episode. ⛔ NO +The driver records a state and an `observation` dictionary for every motion +step. Observation entries name the available policy, agentview, wrist, depth, +world-map, and metadata artifacts, but storage paths are internal to the +runtime. Do not construct or read artifact paths manually. + +Use `view_driver_state` to retrieve a state. It embeds the policy image and the +best available agentview and wrist images, preferring high resolution. Use +`view_camera_meta`, `back_project`, and `segment` to consume metadata, depth, +and world maps. These tools guarantee that the selected camera, resolution, +and step use matching artifacts. + +Step `0` is the initial state. Step `-1` selects the latest state.""" + +GOAL = """YOUR GOAL: produce top-level `libero_terminated == true` in ONE episode. ⛔ NO `reset`, NO retry (SINGLE-ATTEMPT MODE — see the override at the very top; it supersedes any reset/retry wording in the Rules below).""" RULES = """Rule 0 — USE IMAGES. After every primitive tool call, inspect the returned state - and image paths. If you need a state again, call `view_driver_state`. Read the - new `image_cam_hi_NN.png` path (calibration frame — the one you pick pixels in) - and, when close to a target, the `image_wrist_hi_NN.png` path. The image is - your spatial-reasoning input; `states.json` only gives proprioception + object - names. + and embedded images. If you need a state again, call `view_driver_state`. + Use agentview for global layout and wrist for close-range geometry. The image + is your spatial-reasoning input; the JSON state only gives proprioception + + object names. Rule 1 — Pi0 is ONLY for the grasp. Use: pi0_pick({ @@ -215,19 +185,19 @@ Rule 1b — JUDGE THE GRASP from perception, NOT from a name. After a pick, decide "did I grab the target?" from two coord-free signals: • GRIPPER (proprioception): `state.robot0_gripper_qpos` from the latest - `states.json` entry — fingers closed but NOT fully shut (~0.01–0.05 gap) + state record — fingers closed but NOT fully shut (~0.01–0.05 gap) ⇒ holding an object; fully closed (~0.0) ⇒ grasped air. - • WRIST CAM: Read `image_wrist_hi_NN.png` after lifting. The target should + • WRIST CAM: inspect the returned wrist image after lifting. The target should now be raised into the gripper, and the spot it came from should be EMPTY. Compare before/after wrist or agentview evidence; if needed, use `back_project` on wrist pixels to confirm the target surface z jumped up. `pi0_pick.success` (eef-lift + gripper-closure heuristic) is a HINT, not proof — always confirm with the wrist cam before carrying. -Rule 2 — Inspect THEN act. Call `view_driver_state({"step": 0})`, read the - returned high-resolution image path(s), and inspect the relevant memory/guides - BEFORE your first primitive. **Your task is `states.json[0]["task_language"]` - — read it and obey it verbatim.** This is the authoritative instruction (the BDDL +Rule 2 — Inspect THEN act. Call `view_driver_state({"step": 0})`, inspect the + returned high-resolution images, and inspect the relevant memory/guides + BEFORE your first primitive. **Your task is the returned `task_language`; + read it and obey it verbatim.** This is the authoritative instruction (the BDDL `:language` tag). Do NOT infer the task from object names, from sibling recipes, or by guessing a task_map index — those caused wrong-task runs in the past. @@ -257,7 +227,7 @@ plate, a stove burner/cook-region, a wooden-cabinet top, and a pot lid all read as "flat disc at table height" in back-projected coordinates. They are only separable in the RGB. So before you carry-and-release onto a surface, look at - `image_cam_hi_NN.png` (and the wrist `image_wrist_hi_NN.png` once close) and + returned agentview image (and the wrist image once close) and NAME each candidate surface: • PLATE ⇒ ceramic disc, usually white, with a clean raised rim (often colored concentric rings). This is the place target for "place it on the plate". @@ -292,8 +262,7 @@ LOCALIZATION = """This is the core of perception-isolated mode. To find where an object is: -1. Look at `image_cam_hi_NN.png` (1024x1024 — PREFER THIS; fall back to the - 256 `image_cam_NN.png` only if the hi file is absent) and find the target +1. Look at the returned agentview image (high resolution when available) and find the target object's pixel (row, col). (row = vertical/y from top, col = horizontal/x from left.) 2. Call `back_project` on that pixel: @@ -316,7 +285,7 @@ metres away. Pick pixels firmly on the object's top surface.) ALWAYS apply the manipulation offsets from memory to the PERCEIVED position -(e.g. BOWL: eef_y = plate_y + 0.045). Verify visually in image_cam after moving.""" +(e.g. BOWL: eef_y = plate_y + 0.045). Verify visually in agentview after moving.""" PERCEPTION_ALGORITHM = """This is the default perception algorithm for EVERY cell (from the 80-task localization sweep: `agentview_identity_wrist_geometry_except_basket`). @@ -335,11 +304,11 @@ ALGORITHM (run this BEFORE manipulating): -1. From `states.json[0]["task_language"]` + `image_cam_hi_00.png` + +1. From the initial `task_language` + returned agentview image + object_names, infer the task-relevant TARGETS and DESTINATIONS (language only; never BDDL/poses). -2. GLOBAL SEMANTIC PASS (agentview hi-res): in `image_cam_hi_NN.png` choose each +2. GLOBAL SEMANTIC PASS (agentview hi-res): in the returned agentview image choose each target/destination candidate by RGB, label/shape, and global spatial relation. For duplicates (two bowls/plates/mugs) pick by RELATION (on stove, on cookie box, left/right/front/back), not `_1/_2`. For sauce/can/box groceries use the @@ -348,7 +317,7 @@ burner vs cabinet/drawer vs basket) semantically in RGB here. 3. COARSE XYZ (agentview): pick 3-8 pixels firmly on the chosen candidate in - `image_cam_hi_NN.png`, call `back_project` on the SAME pixels, take the median. + agentview image, call `back_project` on the SAME pixels, take the median. Avoid edges/holes/shadows/table-gaps. This median is the IDENTITY ANCHOR for that entity. @@ -431,8 +400,8 @@ command lists. Re-derive every coordinate from THIS scene. """, """INSPECT INITIAL STATE: call `view_driver_state({"step": 0})`; inspect -`task_language`, object_names, eef pose, `image_cam_hi_00.png`, -`image_wrist_hi_00.png` if useful, and `camera_meta.json`. Identify ALL target + `task_language`, object_names, eef pose, the returned agentview and wrist images, + and call `view_camera_meta` if needed. Identify ALL target objects, destination surfaces, and relation landmarks named by task_language. """, """RUN THE MANDATORY PRE-TASK PERCEPTION PASS (FIRST-STEP ALGORITHM above) — @@ -450,10 +419,9 @@ pi0_pick({"prompt": "...", "max_chunks": 20, ...}) release({}) -Each primitive tool blocks until the next `states.json` entry is dumped and -returns the new state view + log + image paths. Then inspect the returned state -+ high-resolution image paths (+ `back_project` as needed), decide, repeat -with NN=02, 03, ... +Each primitive tool blocks until the next state record is dumped and returns the +new state view, log, and embedded images. Inspect them, use `back_project` as +needed, decide, and repeat. """, """ALLOWED PRIMITIVES (physics-only; full schemas in the tool list/guides): `move_to`, `pi0_pick`, `pi0_doubled`, `release`, `set_gripper`, @@ -475,16 +443,17 @@ SAM3 localization aid — `segment` (no robot motion): instead of eyeballing a pixel, call `segment({"prompt":"the black bowl on the cookies box", "camera":"agentview"})`. It runs SAM3 on the current image, back-projects the -mask via the matching world map, and writes `segments/segment_NN_XX.json` with a robust -median `world_xyz` (+ a `segments/segment_overlay_NN_XX.png` to confirm the right -object). Use `camera":"wrist"` (after parking the eef ~15–20 cm over the +mask via the matching world map, and returns a robust median `world_xyz` plus +an embedded overlay image for visual confirmation. The logical +`segment_artifact` and `overlay_artifact` names are audit references, not paths +to open manually. Use `camera":"wrist"` (after parking the eef ~15–20 cm over the target) for ±1–2 cm refinement, or `"point":[row,col]` for a point prompt. Text `prompt` and `point` are mutually exclusive; provide exactly one. ⚠ PROMPT PHRASING (SAM3 is sensitive): use a plain colour+shape+RELATION phrase, NEVER the internal/brand name from `object_names`/BDDL. `"the akita black bowl"` scores ~0.03 (SAM3 can't ground "akita") whereas `"the black bowl on the stove"` scores ~0.76. Strip proper nouns (akita, glazed_rim_porcelain_…) — say what it -LOOKS LIKE + where it is. Always inspect the returned overlay path to confirm +LOOKS LIKE + where it is. Always inspect the returned overlay image to confirm the mask landed on the right object before moving. This is a CONVENIENCE alternative to manual back-projection — if it returns `{"error":..., "fallback":...}` (server down / low score / no detection), walk @@ -499,7 +468,7 @@ unrecoverable within this one episode, do NOT reset — write an honest stuck-audit (`libero_terminated:false`) and call `finish`. Never warp. """, - """WHEN state.libero_terminated == True: + """WHEN `libero_terminated == true` in the latest tool result: a. Write audit `{{output_dir}}/{{recipe_tag}}.json` with: suite, task_id, seed, regime:"strict_perception", strategy_notes (incl. how you localized), pick_result, final_state (latest state's `state`), diff --git a/robots/libero/prompts/user.py b/robots/libero/prompts/user.py index 668a4db4..63d171ac 100644 --- a/robots/libero/prompts/user.py +++ b/robots/libero/prompts/user.py @@ -10,9 +10,11 @@ - recipe: {{output_dir}}/recipe_{{recipe_tag}}.jsonl""" -MODE = """Use the high-resolution image paths returned by view_driver_state and -back_project to localize objects before motion.""" +MODE = """Inspect the embedded high-resolution images returned by +view_driver_state, then use back_project or segment to localize objects before +motion.""" -BEGIN = """read MEMORY.md, the guides, then `view_driver_state({"step":0})` and the -returned high-resolution images. Localize the target, then plan and execute.""" +BEGIN = """Read MEMORY.md and the guides, then call +`view_driver_state({"step": 0})` and inspect its embedded images. Localize every +task-relevant entity before planning and execution.""" diff --git a/robots/libero/toolkit.py b/robots/libero/toolkit.py index 71345caa..b6ae17d5 100644 --- a/robots/libero/toolkit.py +++ b/robots/libero/toolkit.py @@ -19,22 +19,6 @@ class LiberoToolkit(Toolkit): """Toolkit for the LIBERO environment.""" - _WIPE_STREAMS = ( - "image", - "image_cam", - "depth", - "world", - "image_wrist", - "depth_wrist", - "world_wrist", - "wrist_meta", - "image_cam_hi", - "world_hi", - "image_wrist_hi", - "world_wrist_hi", - "segments", - "action_videos", - ) # Tool schemas keyed by name (built once from the canonical ordered list # in libero_tools.TOOLS_SPEC) so each tool registers with its own spec. _SPECS = {spec["name"]: spec for spec in libero_tools.TOOLS_SPEC} @@ -44,11 +28,9 @@ def __init__( *, primitives_kwargs: dict[str, Any], dashboard_events: DashboardEventSink, - video_path: str | None = None, ) -> None: state = EnvState(get_output_dir()) super().__init__(dashboard_events=dashboard_events, state=state) - self._video_path: str | None = video_path self.init_primitives_clean(primitives_kwargs=primitives_kwargs) self._register_libero_tools() @@ -99,28 +81,36 @@ def _step(self, name: str, **kwargs) -> dict: else: result_dict = {"value": result} - step_idx = self._state.next_step_idx + record = libero_tools.dump_state( + self._primitives, + self._state, + log={"command": command, "result": result_dict, "elapsed_s": elapsed}, + ) + action_video_name = None if self._dashboard_events.enabled: - video_dir = libero_tools.artifact_path( - self._state.output_dir, "action_videos" - ) - video_path = video_dir / f"step_{step_idx:02d}_{name}.mp4" try: - self._primitives.save_frame_slice(start_frame, str(video_path), fps=20) + frames = self._primitives.frame_slice(start_frame) + if frames: + candidate = f"action_{name}.mp4" + if self._state.save( + candidate, + frames, + step=record.step_idx, + fps=20, + ): + action_video_name = candidate except Exception as e: get_logger("libero_toolkit").warning( - f"failed to save action clip to {video_path}: {e}" + "failed to save action clip for step %s: %s", + record.step_idx, + e, ) - libero_tools.dump_state( - self._primitives, - self._state, - step_idx=step_idx, - log={"command": command, "result": result_dict, "elapsed_s": elapsed}, - ) - out = libero_tools.view_driver_state(step_idx, state=self._state) + out = libero_tools.view_driver_state(record.step_idx, state=self._state) out["agent_elapsed_s"] = elapsed if result_dict.get("interrupted"): out.update(result_dict) + if action_video_name is not None: + out["action_video_artifact"] = action_video_name return out def init_primitives_clean( @@ -129,14 +119,7 @@ def init_primitives_clean( primitives_kwargs: dict[str, Any], ) -> None: """Wipe stale run artifacts, build the LiberoPrimitives, dump step 0.""" - self._state.reset(wipe_streams=self._WIPE_STREAMS) - out_dir = self._state.output_dir - for target in ( - libero_tools.artifact_path(out_dir, "metadata", camera="agentview", resolution="low"), - libero_tools.artifact_path(out_dir, "episode_video"), - ): - if target.exists(): - target.unlink() + self._state.reset() primitives = libero_tools.LiberoPrimitives( check_cancelled=self.raise_if_cancelled, @@ -144,28 +127,25 @@ def init_primitives_clean( ) primitives.reset() primitives.start_recording() - libero_tools.dump_state(primitives, self._state, step_idx=0, log=None) + libero_tools.dump_state(primitives, self._state, log=None) self._dashboard_events.emit( ToolResultEvent( name="view_driver_state", - result=libero_tools.view_driver_state(0, state=self._state), + result=libero_tools.view_driver_state(0, state=self._state), ) ) self._primitives = primitives def close(self) -> None: - """Flush the agent-side video buffer to disk (end-of-run). - """ - if self._video_path is None: - return + """Flush the agent-side video buffer through ``EnvState``.""" try: - self._primitives.stop_recording_and_save(self._video_path) + frames = self._primitives.stop_recording() + if frames: + self._state.save("episode.mp4", frames, step=None, fps=20) except Exception as e: - # The runner is in the cleanup path; never let a video save - # abort it. get_logger("libero_toolkit").warning( - f"failed to save video to {self._video_path}: {e}" + f"failed to save episode video: {e}" ) def write_recipe(self, recipe_tag: str) -> str: diff --git a/robots/libero/tools.py b/robots/libero/tools.py index 6b95f6cf..7f511e8e 100644 --- a/robots/libero/tools.py +++ b/robots/libero/tools.py @@ -1,13 +1,9 @@ """LIBERO + OpenPI tool implementation.""" from __future__ import annotations -import json -import os from collections.abc import Callable -from pathlib import Path from typing import Any -import imageio.v2 as imageio import numpy as np from robots.libero.env_client import LiberoEnvClient @@ -18,70 +14,6 @@ logger = get_logger("libero") -ARTIFACT_LAYOUT: dict[tuple[str | None, str | None, str], str] = { - ("agentview", "low", "policy_image"): "images/image_{step:02d}.png", - ("agentview", "low", "image"): "images_cam/image_cam_{step:02d}.png", - ("agentview", "low", "depth"): "depths/depth_{step:02d}.npy", - ("agentview", "low", "world"): "world/world_{step:02d}.npy", - ("agentview", "low", "metadata"): "camera_meta.json", - ("agentview", "high", "image"): "images_cam_hi/image_cam_hi_{step:02d}.png", - ("agentview", "high", "world"): "world_hi/world_hi_{step:02d}.npy", - ("wrist", "low", "image"): "images_wrist/image_wrist_{step:02d}.png", - ("wrist", "low", "depth"): "depths_wrist/depth_wrist_{step:02d}.npy", - ("wrist", "low", "world"): "world_wrist/world_wrist_{step:02d}.npy", - ("wrist", "low", "metadata"): "wrist_meta/wrist_meta_{step:02d}.json", - ("wrist", "high", "image"): "images_wrist_hi/image_wrist_hi_{step:02d}.png", - ("wrist", "high", "world"): "world_wrist_hi/world_wrist_hi_{step:02d}.npy", - (None, None, "states"): "states.json", - (None, None, "episode_video"): "episode.mp4", - (None, None, "segments"): "segments", - (None, None, "action_videos"): "action_videos", -} -ARTIFACT_DIRECTORIES: tuple[str, ...] = ( - "images", - "images_cam", - "depths", - "world", - "images_cam_hi", - "world_hi", - "images_wrist", - "depths_wrist", - "world_wrist", - "wrist_meta", - "images_wrist_hi", - "world_wrist_hi", - "segments", - "action_videos", -) - - -def artifact_path( - output_dir: str | os.PathLike[str], - kind: str, - *, - step: int | None = None, - camera: str | None = None, - resolution: str | None = None, -) -> Path: - """Resolve one artifact path from the shared layout. - - ``kind`` identifies the artifact; the remaining fields are optional - qualifiers and must be passed by keyword to avoid mixing them up. - """ - return Path(output_dir) / _artifact_relative_path(step, camera, resolution, kind) - - -def _artifact_relative_path( - step: int | None, - camera: str | None, - resolution: str | None, - kind: str, -) -> str: - template = ARTIFACT_LAYOUT[(camera, resolution, kind)] - if step is None and "{step" in template: - raise ValueError(f"{kind} artifact requires a step") - return template.format(step=step) - def _normalize_xyz(xyz): """Coerce an LLM-supplied xyz into a length-3 list[float].""" @@ -117,7 +49,7 @@ def __init__( self._last_obs_eef_z = None self._last_obs_gripper = None # Per-env-step frame buffer for diagnostic video rendering. - # Toggled via start_recording() / stop_recording_and_save(). + # Toggled via start_recording() / stop_recording(). self._recording = False self._frames = [] @@ -132,22 +64,14 @@ def record_frame(self, obs): def recorded_frame_count(self) -> int: return len(self._frames) - def stop_recording_and_save(self, path: str, fps: int = 20): - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - n = len(self._frames) - if n > 0: - imageio.mimwrite(path, self._frames, fps=fps) + def stop_recording(self) -> list[np.ndarray]: + frames = list(self._frames) self._recording = False self._frames = [] - return {"path": path, "n_frames": n} + return frames - def save_frame_slice(self, start: int, path: str, fps: int = 20): - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - frames = list(self._frames[int(start):]) - n = len(frames) - if n > 0: - imageio.mimwrite(path, frames, fps=fps) - return {"path": path, "n_frames": n, "fps": fps} + def frame_slice(self, start: int) -> list[np.ndarray]: + return list(self._frames[int(start):]) def set_obs(self, obs): self._last_obs = obs @@ -660,7 +584,7 @@ def segment( self, prompt: str = "", camera: str = "agentview", - step: int | None = None, + step: int = -1, point: list[int] | None = None, min_score: float = 0.2, *, @@ -672,9 +596,11 @@ def segment( artifacts. Errors are structured so the agent can continue with image inspection and ``back_project``. """ - nn = state.latest_step if step is None else int(step) - if nn is None: - return {"error": "no state entries; cannot select segment image"} + try: + record = state.get(step) + except Exception as exc: + return {"error": f"state step not available: {exc}"} + nn = record.step_idx camera = camera or "agentview" prompt = prompt.strip() @@ -683,26 +609,27 @@ def segment( if has_prompt == has_point: return {"error": "segment needs exactly one of prompt or point"} try: - image_path, world_path, artifact_pairs = _select_segment_artifacts( - state, nn, camera + image_name, world_name, artifact_pairs = _select_segment_artifacts( + state, record, camera ) except ValueError as e: return {"error": str(e)} - if image_path is None: + if image_name is None or world_name is None: return { "error": "complete segment artifacts not found", "step": nn, "camera": camera, - "checked_paths": [ - str(path) + "checked_artifacts": [ + name for image, world in artifact_pairs - for path in (image, world) + for name in (image, world) + if name ], } try: data = self._sam3_client.segment( - image_path, + state.load_bytes(image_name, step=nn), text_prompt=prompt if has_prompt else None, point=point, min_score=min_score, @@ -712,36 +639,41 @@ def segment( "error": str(e), "step": nn, "camera": camera, - "image_path": str(image_path), + "image_artifact": image_name, } except Exception as e: return { "error": f"segmentation service call failed: {e}", "step": nn, "camera": camera, - "image_path": str(image_path), + "image_artifact": image_name, "fallback": "Use manual visual localization and back_project.", } - out_dir = state.output_dir - segment_path, overlay_candidate_path, segment_index = ( - _next_segment_artifact_paths(out_dir, nn) - ) - overlay_path = None + segment_index = _next_segment_index(record) + segment_name = f"segment_{segment_index:02d}.json" + overlay_name = f"segment_overlay_{segment_index:02d}.png" + saved_overlay = None mask = data.mask if data.found and isinstance(mask, np.ndarray): - if world_path is None or not world_path.exists(): + try: + world_map = state.load(world_name, step=nn) + except Exception as exc: world_result = { "world_xyz": None, - "world_error": "world map artifact not found for selected image", - "expected_world_path": str(world_path) if world_path else None, + "world_error": f"world map artifact not available: {exc}", + "expected_world_artifact": world_name, } else: - world_result = _mask_to_world(mask, np.load(world_path)) - world_result["world_path"] = str(world_path) - overlay_path = overlay_candidate_path - if not _write_segment_overlay(image_path, mask, overlay_path): - overlay_path = None + world_result = _mask_to_world(mask, world_map) + world_result["world_artifact"] = world_name + overlay = _make_segment_overlay(state.load(image_name, step=nn), mask) + if overlay is not None and state.save( + overlay_name, + overlay, + step=nn, + ): + saved_overlay = overlay_name else: world_result = { "world_xyz": None, @@ -754,7 +686,7 @@ def segment( "camera": camera, "source_step": nn, "segment_index": segment_index, - "image_path": str(image_path), + "image_artifact": image_name, "min_score": min_score, "score": round(float(data.score), 3) if data.score is not None else None, "box": data.box, @@ -767,14 +699,18 @@ def segment( if not data.found: segment_blob["error"] = data.reason or "SAM3 found no mask" segment_blob.update(world_result) - segment_path.write_text(json.dumps(segment_blob, indent=2, default=str)) + state.save( + segment_name, + segment_blob, + step=nn, + ) result = { "found": data.found, "step": nn, "camera": camera, - "image_path": str(image_path), - "segment_path": str(segment_path), + "image_artifact": image_name, + "segment_artifact": segment_name, "score": segment_blob["score"], "box": segment_blob["box"], "world_xyz": segment_blob["world_xyz"], @@ -783,59 +719,57 @@ def segment( if "error" in segment_blob: result["error"] = segment_blob["error"] result["fallback"] = "Use manual visual localization and back_project." - if overlay_path is not None and overlay_path.exists(): - result["overlay_path"] = str(overlay_path) + if saved_overlay is not None: + result["overlay_artifact"] = saved_overlay + result["_image_bytes"] = state.load_bytes(saved_overlay, step=nn) return result def write_recipe_from_states(state: EnvState, recipe_tag: str) -> str: """Find a command sequence that gets ``libero_terminated=True``. - Export non-error LIBERO primitive commands from ``states.json`` and - successful segment calls from ``segments/segment_*.json``. + Export non-error LIBERO primitive commands and successful segment calls. """ command_events = [] for record in state.records(): command = record.command - if command is None: - continue - if command.get("action") not in PRIMITIVE_TOOL_NAMES: - continue result = record.result - if isinstance(result, dict) and result.get("error"): - continue - command_events.append(((record.step_idx, -1), command)) - - output_dir = state.output_dir - for artifact in artifact_path(output_dir, "segments").glob("segment_*.json"): - with artifact.open() as f: - segment = json.load(f) - if segment.get("error"): - continue - if segment["mode"] == "text": - command = { - "action": "segment", - "prompt": segment["prompt"], - "camera": segment["camera"], - } - else: - command = { - "action": "segment", - "point": segment["point"], - "camera": segment["camera"], - } - source_step = int(segment["source_step"]) - event_order = (source_step, int(segment["segment_index"])) - command_events.append((event_order, command)) + if ( + command is not None + and command.get("action") in PRIMITIVE_TOOL_NAMES + and not (isinstance(result, dict) and result.get("error")) + ): + command_events.append(((record.step_idx, -1), command)) + + for name in sorted(record.artifacts): + if not (name.startswith("segment_") and name.endswith(".json")): + continue + segment = state.load(name, step=record.step_idx) + if segment.get("error"): + continue + if segment["mode"] == "text": + segment_command = { + "action": "segment", + "prompt": segment["prompt"], + "camera": segment["camera"], + } + else: + segment_command = { + "action": "segment", + "point": segment["point"], + "camera": segment["camera"], + } + event_order = (record.step_idx, int(segment["segment_index"])) + command_events.append((event_order, segment_command)) - recipe_path = os.path.join(state.output_dir, f"recipe_{recipe_tag}.jsonl") - tmp_path = recipe_path + ".tmp" command_events.sort(key=lambda event: event[0]) - with open(tmp_path, "w") as f: - for _, command in command_events: - f.write(json.dumps(command, separators=(",", ":")) + "\n") - os.replace(tmp_path, recipe_path) - return recipe_path + recipe_name = f"recipe_{recipe_tag}.jsonl" + state.save( + recipe_name, + [command for _, command in command_events], + step=None, + ) + return recipe_name def _metric_depth(depth: Any, camera_meta: dict) -> np.ndarray: @@ -864,41 +798,13 @@ def _world_from_depth(depth_metric: np.ndarray, camera_meta: dict) -> np.ndarray return (camera_points @ extrinsic.T)[..., :3] -def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, - log: dict | None = None) -> StepRecord: - """Dump state snapshot, images, and depth for step *step_idx*. - - Writes: - - ``/images/image_NN.png`` (Pi0-frame agentview) - - ``/images_cam/image_cam_NN.png`` (calibration-frame agentview) - - ``/depths/depth_NN.npy`` (metric depth, meters) - - ``/world/world_NN.npy`` (agentview world xyz map) - - ``/images_wrist/image_wrist_NN.png`` - - ``/depths_wrist/depth_wrist_NN.npy`` - - ``/world_wrist/world_wrist_NN.npy`` - - ``/wrist_meta/wrist_meta_NN.json`` - - high-res ``images_cam_hi`` / ``world_hi`` artifacts - - high-res ``images_wrist_hi`` / ``world_wrist_hi`` artifacts - - ``/camera_meta.json`` (static, once) - - appends the step blob to ``/states.json`` - - If *log* is provided (the return value of :func:`execute`), its - ``command``, ``result``, and ``elapsed_s`` fields are merged into the - step blob so a single entry captures everything. - """ - output_dir = env_state.output_dir - for directory in ARTIFACT_DIRECTORIES: - (Path(output_dir) / directory).mkdir(parents=True, exist_ok=True) - - agent_world_map = None - wrist_world_map = None - agent_world_map_hi = None - wrist_world_map_hi = None - wrist_meta = None - # Reuse one raw observation snapshot for state and per-step artifacts. +def dump_state( + primitives: LiberoPrimitives, + env_state: EnvState, + log: dict | None = None, +) -> StepRecord: + """Save one Libero observation through its owned state record.""" raw = primitives.env.raw_obs() - # Expose robot proprioception and object names, but never privileged - # object coordinates; the agent must localize through visual artifacts. state = { "robot0_eef_pos": [float(x) for x in raw["robot0_eef_pos"]], "robot0_eef_quat": [float(x) for x in raw["robot0_eef_quat"]], @@ -909,9 +815,32 @@ def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, if k.endswith("_pos") and "robot0" not in k and "to_robot" not in k ), } - imageio.imwrite( - artifact_path(output_dir, "policy_image", step=step_idx, camera="agentview", resolution="low"), + log = log or {} + with env_state.record_step( + state=state, + command=log.get("command"), + result=log.get("result"), + elapsed_s=log.get("elapsed_s"), + extras={ + "libero_terminated": primitives.env.episode_terminated, + "episode_truncated": primitives.env.episode_truncated, + "task_language": primitives.env.get_task_language(), + }, + ) as step_idx: + _save_observation_artifacts(primitives, env_state, step_idx, raw) + return env_state.get(step_idx) + + +def _save_observation_artifacts( + primitives: LiberoPrimitives, + env_state: EnvState, + step_idx: int, + raw: dict[str, Any], +) -> None: + env_state.save( + "agentview_policy.png", primitives._last_obs["main_images"], + step=step_idx, ) # --- camera calibration (static for agentview): fetch metadata as needed --- @@ -920,8 +849,7 @@ def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, height=256, width=256, ) or {} - camera_meta_path = artifact_path(output_dir, "metadata", camera="agentview", resolution="low") - if agentview_meta and not camera_meta_path.exists(): + if agentview_meta: cam_meta_out = dict(agentview_meta) cam_meta_out["projection"] = ( "Prefer the back_project(row, col, step=NN) MCP tool; it " @@ -930,24 +858,30 @@ def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, "256x256 calibration-frame image." ) cam_meta_out["note"] = ( - "depth_NN.npy is in this camera frame (vertical-flipped raw " - "buffer). image_NN.png is rotated 180deg (Pi0 convention) and " - "is NOT in the same frame as depth/K." + "The agentview_depth.npy observation is aligned with agentview.png. " + "agentview_policy.png uses the Pi0 orientation and must not supply " + "pixels for back-projection." + ) + env_state.save( + "agentview_metadata.json", + cam_meta_out, + step=step_idx, ) - with open(camera_meta_path, "w") as f: - json.dump(cam_meta_out, f, indent=2) # --- per-step RGB in the depth/K frame (vertical-flip of the raw buffer) --- - # The agent picks object pixels HERE (same frame as depth_NN.npy + K), so - # pixel -> depth -> back-project is direct. (image_NN.png is the 180°-rotated - # Pi0-convention frame and must NOT be used for back-projection.) + # Agentview pixels align with the matching depth and calibration. The policy + # image uses Pi0 orientation and must not supply back-projection pixels. try: ci = raw.get("agentview_image") if ci is not None: ci = np.asarray(ci) if ci.dtype != np.uint8: ci = ci.astype(np.uint8) - imageio.imwrite(artifact_path(output_dir, "image", step=step_idx, camera="agentview", resolution="low"), ci[::-1]) + env_state.save( + "agentview.png", + ci[::-1], + step=step_idx, + ) except Exception as e: logger.warning("image_cam dump failed: %s", e) @@ -960,21 +894,19 @@ def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, # depth map in this frame. VERIFIED 5/5: projecting each GT object # world pos via M lands on a pixel whose depth_flip[row,col] matches # the object's surface depth (plate Δ6mm, cookies Δ14mm). So - # pixel(row,col) in depth_NN.npy back-projects correctly with - # camera_meta.json (NOT the same frame as the 180°-rotated - # image_NN.png — see camera_meta note). + # Pixels in agentview.png align with agentview_depth.npy and the + # per-step agentview metadata saved with the record artifacts. d = _metric_depth(d, agentview_meta)[::-1] - np.save( - artifact_path(output_dir, "depth", step=step_idx, camera="agentview", resolution="low"), + env_state.save( + "agentview_depth.npy", d.astype(np.float32), + step=step_idx, ) world = _world_from_depth(d, agentview_meta).astype(np.float32) - np.save( - artifact_path(output_dir, "world", step=step_idx, camera="agentview", resolution="low"), + env_state.save( + "agentview_world.npy", world, - ) - agent_world_map = _artifact_relative_path( - step_idx, "agentview", "low", "world" + step=step_idx, ) except Exception as e: logger.warning("depth dump failed: %s", e) @@ -988,7 +920,11 @@ def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, wimg = np.asarray(wimg) if wimg.dtype != np.uint8: wimg = wimg.astype(np.uint8) - imageio.imwrite(artifact_path(output_dir, "image", step=step_idx, camera="wrist", resolution="low"), wimg[::-1]) + env_state.save( + "wrist.png", + wimg[::-1], + step=step_idx, + ) except Exception as e: logger.warning("wrist image dump failed: %s", e) @@ -1008,32 +944,30 @@ def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, logger.warning("wrist camera meta missing; skipping wrist depth/world") else: wdpt_metric = _metric_depth(wdpt_arr, wmeta)[::-1] - np.save( - artifact_path(output_dir, "depth", step=step_idx, camera="wrist", resolution="low"), + env_state.save( + "wrist_depth.npy", wdpt_metric.astype(np.float32), + step=step_idx, ) world_w = _world_from_depth(wdpt_metric, wmeta).astype(np.float32) - np.save( - artifact_path(output_dir, "world", step=step_idx, camera="wrist", resolution="low"), + env_state.save( + "wrist_world.npy", world_w, - ) - wrist_world_map = _artifact_relative_path( - step_idx, "wrist", "low", "world" + step=step_idx, ) wmeta_out = dict(wmeta) wmeta_out["note"] = ( "MOVING camera: extrinsic_cam2world is for THIS step " - "only. world_wrist_NN.npy[row,col] gives world " - "(x,y,z) for that pixel, in the SAME world frame as " - "agentview world_NN.npy." + "only. The matching wrist world-map observation gives world " + "(x,y,z) for that pixel in the same world frame as the " + "agentview world-map artifact." + ) + env_state.save( + "wrist_metadata.json", + wmeta_out, + step=step_idx, ) - with open( - artifact_path(output_dir, "metadata", step=step_idx, camera="wrist", resolution="low"), - "w", - ) as f: - json.dump(wmeta_out, f, indent=2) - wrist_meta = wmeta except Exception as e: logger.warning("wrist depth/world dump failed: %s", e) @@ -1047,20 +981,19 @@ def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, meta_hi = primitives.env.get_camera_meta("agentview", 1024, 1024) if meta_hi is None: raise RuntimeError("agentview camera metadata missing") - imageio.imwrite( - artifact_path(output_dir, "image", step=step_idx, camera="agentview", resolution="high"), + env_state.save( + "agentview_high.png", np.asarray(rgb_hi)[::-1], + step=step_idx, ) world_hi = _world_from_depth( _metric_depth(depth_hi, meta_hi)[::-1], meta_hi, ).astype(np.float16) - np.save( - artifact_path(output_dir, "world", step=step_idx, camera="agentview", resolution="high"), + env_state.save( + "agentview_world_high.npy", world_hi, - ) - agent_world_map_hi = _artifact_relative_path( - step_idx, "agentview", "high", "world" + step=step_idx, ) except Exception as e: logger.warning("agentview high-res dump failed: %s", e) @@ -1077,70 +1010,23 @@ def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, ) if meta_wrist_hi is None: raise RuntimeError("robot0_eye_in_hand camera metadata missing") - imageio.imwrite( - artifact_path(output_dir, "image", step=step_idx, camera="wrist", resolution="high"), + env_state.save( + "wrist_high.png", np.asarray(rgb_wrist_hi)[::-1], + step=step_idx, ) world_wrist_hi = _world_from_depth( _metric_depth(depth_wrist_hi, meta_wrist_hi)[::-1], meta_wrist_hi, ).astype(np.float16) - np.save( - artifact_path(output_dir, "world", step=step_idx, camera="wrist", resolution="high"), + env_state.save( + "wrist_world_high.npy", world_wrist_hi, - ) - wrist_world_map_hi = _artifact_relative_path( - step_idx, "wrist", "high", "world" + step=step_idx, ) except Exception as e: logger.warning("wrist high-res dump failed: %s", e) - for old_step in range(max(0, int(step_idx) - 4)): - for path in ( - artifact_path(output_dir, "image", step=old_step, camera="agentview", resolution="high"), - artifact_path(output_dir, "world", step=old_step, camera="agentview", resolution="high"), - artifact_path(output_dir, "image", step=old_step, camera="wrist", resolution="high"), - artifact_path(output_dir, "world", step=old_step, camera="wrist", resolution="high"), - ): - try: - os.unlink(path) - except FileNotFoundError: - pass - - frame_streams = [ - stream - for stream in ("image", "image_cam", "image_wrist") - if env_state.artifact_path(step_idx, stream).exists() - ] - depth_streams = [ - stream - for stream in ("depth", "depth_wrist") - if env_state.artifact_path(step_idx, stream).exists() - ] - camera_meta = {"agentview": agentview_meta} - if wrist_meta is not None: - camera_meta["wrist"] = wrist_meta - record = StepRecord( - step_idx=step_idx, - state=state, - frames=frame_streams, - depth=depth_streams, - camera_meta=camera_meta, - command=log.get("command") if log else None, - result=log.get("result") if log else None, - elapsed_s=log.get("elapsed_s") if log else None, - extras={ - "libero_terminated": primitives.env.episode_terminated, - "episode_truncated": primitives.env.episode_truncated, - "task_language": primitives.env.get_task_language(), - "world_map": agent_world_map, - "wrist_world_map": wrist_world_map, - "world_map_hi": agent_world_map_hi, - "wrist_world_map_hi": wrist_world_map_hi, - }, - ) - return env_state.append(record) - # --------------------------------------------------------------------------- # Tool schema declarations (Anthropic-shaped canonical schema) @@ -1161,16 +1047,9 @@ def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, { "name": "view_driver_state", "description": ( - "Read step NN from `states.json` + the matching " - "state images in {{output_dir}}. If step is " - "null, returns the latest entry. Each entry contains the robot " - "state, libero_terminated flag, command log, and result. Returns " - "available PNG paths in this stable " - "order: 1) `images/image_NN.png` (Pi0-frame agentview), " - "2) `images_cam/image_cam_NN.png` (calibration-frame agentview), " - "3) `images_wrist/image_wrist_NN.png` (calibration-frame wrist). " - "High-resolution calibration-frame images are returned as file " - "paths, not embedded as image bytes. " + "Read one recorded state and its observation artifacts. Step -1 " + "selects the latest entry. Embeds policy, agentview, and wrist " + "images when available. " "Use the calibration-frame images for pixel back-projection; JSON " "state alone is not enough. Use agentview for global tabletop " "layout and object locations; use wrist for close-range details " @@ -1180,8 +1059,9 @@ def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, "type": "object", "properties": { "step": { - "type": ["integer", "null"], - "description": "Step number; 0 = initial. Null = latest.", + "type": "integer", + "default": -1, + "description": "Step number; 0 = initial, -1 = latest.", }, }, }, @@ -1372,9 +1252,7 @@ def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, { "name": "view_camera_meta", "description": ( - "Read camera calibration metadata from the output dir. " - "camera='agentview' reads static camera_meta.json. " - "camera='wrist' reads the per-step wrist metadata." + "Read per-step camera calibration metadata from recorded artifacts." ), "input_schema": { "type": "object", @@ -1385,8 +1263,9 @@ def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, "description": "Camera metadata to read (default agentview).", }, "step": { - "type": ["integer", "null"], - "description": "Wrist metadata step to use (default latest).", + "type": "integer", + "default": -1, + "description": "Metadata step to use; -1 = latest.", }, }, }, @@ -1412,8 +1291,9 @@ def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, "description": "Artifact camera to use (default agentview).", }, "step": { - "type": ["integer", "null"], - "description": "Step NN to segment; null = latest.", + "type": "integer", + "default": -1, + "description": "Step to segment; -1 = latest.", }, "point": { "type": ["array", "null"], @@ -1439,7 +1319,7 @@ def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, "selected camera's precomputed world map. Row 0 = top of image, " "col 0 = left. Returns world_xyz in meters.\n\n" "USE THIS to find where an object is in the world — look at " - "the high-resolution paths returned by view_driver_state " + "the embedded high-resolution image returned by view_driver_state " "to pick a pixel on the target object, then call back_project. " "The default resolution is high (1024x1024). Pass " "resolution='low' only for pixels from the embedded/standard " @@ -1468,8 +1348,9 @@ def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, "description": "Pixel column (0=left) in the selected resolution image.", }, "step": { - "type": ["integer", "null"], - "description": "Depth/world-map step to use (default latest). 0 for initial.", + "type": "integer", + "default": -1, + "description": "Depth/world-map step; 0 = initial, -1 = latest.", }, "camera": { "type": "string", @@ -1509,128 +1390,68 @@ def dump_state(primitives: LiberoPrimitives, env_state: EnvState, step_idx: int, ] -def _load_image_path(state: EnvState, nn: int, kind: str) -> str | None: - """Return the path to a dumped state image. None if not present.""" - out_dir = state.output_dir - if kind == "agent": - path = artifact_path(out_dir, "policy_image", step=nn, camera="agentview", resolution="low") - elif kind == "camera": - path = artifact_path(out_dir, "image", step=nn, camera="agentview", resolution="low") - elif kind == "wrist": - path = artifact_path(out_dir, "image", step=nn, camera="wrist", resolution="low") - else: - raise ValueError(f"unknown image kind: {kind}") - if not path.exists(): - return None - return str(path) - - -def _load_camera_meta( - state: EnvState, - camera: str = "agentview", - nn: int | None = None, -) -> dict: - out_dir = state.output_dir - if camera == "agentview": - path = artifact_path(out_dir, "metadata", camera="agentview", resolution="low") - elif camera == "wrist" and nn is not None: - path = artifact_path(out_dir, "metadata", step=nn, camera="wrist", resolution="low") - else: - raise ValueError("camera must be 'agentview' or 'wrist' with nn") - if not path.exists(): - raise FileNotFoundError(f"{path.name} not found in {state.output_dir}") - with open(path) as f: - return json.load(f) - - -def _load_depth(state: EnvState, camera: str, nn: int) -> np.ndarray: - if camera == "agentview": - stream = "depth" - elif camera == "wrist": - stream = "depth_wrist" - else: - raise ValueError("camera must be 'agentview' or 'wrist'") - depth = state.load_depth(nn, stream) - if depth.ndim == 3: - depth = depth[..., 0] - return depth - - -def view_driver_state(step: int | None = None, *, state: EnvState) -> dict: - latest = state.latest_step - if latest is None: - return {"error": "no state entries; env not ready"} - nn = latest if step is None else int(step) +def view_driver_state(step: int = -1, *, state: EnvState) -> dict: try: - record = state.get(nn) - except Exception as e: - return {"error": f"step {nn} not present in state trace: {e}"} + record = state.get(step) + except Exception as exc: + return {"error": f"state step not available: {exc}"} + nn = record.step_idx extras = record.extras - out: dict = {"step": nn, "state": record.state} + out: dict = { + "step": nn, + "state": record.state, + "artifacts": sorted(record.artifacts), + } out["task_language"] = extras.get("task_language") out["libero_terminated"] = extras.get("libero_terminated") out["episode_truncated"] = extras.get("episode_truncated") - out["world_map"] = extras.get("world_map") - out["wrist_world_map"] = extras.get("wrist_world_map") - out["world_map_hi"] = extras.get("world_map_hi") - out["wrist_world_map_hi"] = extras.get("wrist_world_map_hi") out["log"] = { "command": record.command, "result": record.result, "elapsed_s": record.elapsed_s, } - for field, kind in ( - ("image_path", "agent"), - ("image_cam_path", "camera"), - ("image_wrist_path", "wrist"), - ): - image_path = _load_image_path(state, nn, kind) - if image_path: - out[field] = image_path - for field, camera in ( - ("image_cam_hi_path", "agentview"), - ("image_wrist_hi_path", "wrist"), + for slot, names in ( + ("_image_bytes", ("agentview_policy.png",)), + ("_image_cam_bytes", ("agentview_high.png", "agentview.png")), + ("_image_wrist_bytes", ("wrist_high.png", "wrist.png")), ): - image_path = artifact_path( - state.output_dir, - "image", - step=nn, - camera=camera, - resolution="high", - ) - if image_path.exists(): - out[field] = str(image_path) + name = next((name for name in names if name in record.artifacts), None) + if name: + try: + out[slot] = state.load_bytes(name, step=nn) + except FileNotFoundError: + pass return out -def _select_segment_artifacts(state: EnvState, nn: int, camera: str): - out_dir = state.output_dir +def _select_segment_artifacts( + state: EnvState, + record: StepRecord, + camera: str, +) -> tuple[str | None, str | None, list[tuple[str | None, str | None]]]: if camera not in ("agentview", "wrist"): raise ValueError(f"unknown segment camera: {camera}") pairs = [ - ( - artifact_path(out_dir, "image", step=nn, camera=camera, resolution=resolution), - artifact_path(out_dir, "world", step=nn, camera=camera, resolution=resolution), - ) - for resolution in ("high", "low") + (f"{camera}_high.png", f"{camera}_world_high.npy"), + (f"{camera}.png", f"{camera}_world.npy"), ] - for image_path, world_path in pairs: - if image_path.exists() and world_path.exists(): - return image_path, world_path, pairs + for image_name, world_name in pairs: + if ( + image_name in record.artifacts + and world_name in record.artifacts + and state.exists(image_name, step=record.step_idx) + and state.exists(world_name, step=record.step_idx) + ): + return image_name, world_name, pairs return None, None, pairs -def _next_segment_artifact_paths(out_dir: Path, nn: int): - segments_dir = artifact_path(out_dir, "segments") - segments_dir.mkdir(parents=True, exist_ok=True) +def _next_segment_index(record: StepRecord) -> int: idx = 0 - while True: - segment_path = segments_dir / f"segment_{nn:02d}_{idx:02d}.json" - overlay_path = segments_dir / f"segment_overlay_{nn:02d}_{idx:02d}.png" - if not segment_path.exists() and not overlay_path.exists(): - return segment_path, overlay_path, idx + while f"segment_{idx:02d}.json" in record.artifacts: idx += 1 + return idx def _mask_to_world(mask: np.ndarray, world_map: np.ndarray, @@ -1687,28 +1508,25 @@ def _mask_to_world(mask: np.ndarray, world_map: np.ndarray, return result -def _write_segment_overlay(image_path: Path, mask: np.ndarray, - overlay_path: Path) -> bool: - try: - image = imageio.imread(image_path) - if image.ndim != 3 or image.shape[:2] != mask.shape: - return False - overlay = image.copy() - red = np.zeros_like(overlay) - red[..., 0] = 255 - overlay[mask] = ( - 0.55 * overlay[mask].astype(np.float32) - + 0.45 * red[mask].astype(np.float32) - ).astype(np.uint8) - imageio.imwrite(overlay_path, overlay) - return overlay_path.exists() - except Exception: - return False +def _make_segment_overlay( + image: np.ndarray, + mask: np.ndarray, +) -> np.ndarray | None: + if image.ndim != 3 or image.shape[:2] != mask.shape: + return None + overlay = image.copy() + red = np.zeros_like(overlay) + red[..., 0] = 255 + overlay[mask] = ( + 0.55 * overlay[mask].astype(np.float32) + + 0.45 * red[mask].astype(np.float32) + ).astype(np.uint8) + return overlay def view_camera_meta( camera: str = "agentview", - step: int | None = None, + step: int = -1, *, state: EnvState, ) -> dict: @@ -1716,26 +1534,24 @@ def view_camera_meta( if camera not in ("agentview", "wrist"): return {"error": f"bad camera '{camera}' (use 'agentview' or 'wrist')"} - nn = None - if camera == "wrist": - nn = state.latest_step if step is None else int(step) - if nn is None: - return {"error": "no wrist metadata available"} - try: - meta = _load_camera_meta(state, camera, nn) + record = state.get(step) + metadata_name = f"{camera}_metadata.json" + if metadata_name not in record.artifacts: + raise FileNotFoundError(metadata_name) + meta = state.load(metadata_name, step=record.step_idx) except Exception as e: return {"error": f"{camera} camera metadata not found: {e}"} if camera == "agentview": return {"camera": "agentview", "camera_meta": meta} - return {"camera": "wrist", "step": nn, "camera_meta": meta} + return {"camera": "wrist", "step": record.step_idx, "camera_meta": meta} def back_project( row: int | None = None, col: int | None = None, - step: int | None = None, + step: int = -1, camera: str = "agentview", resolution: str = "high", row_range: list | None = None, @@ -1760,23 +1576,16 @@ def back_project( ) } - nn = state.latest_step if step is None else int(step) - if nn is None: - return {"error": "no depth/world-map files available"} - try: - record = state.get(nn) + record = state.get(step) except Exception as e: - return {"error": f"step {nn} not present in state trace: {e}"} + return {"error": f"state step not available: {e}"} + nn = record.step_idx - if camera == "agentview": - hi_artifact = record.extras.get("world_map_hi") - low_artifact = record.extras.get("world_map") - else: - hi_artifact = record.extras.get("wrist_world_map_hi") - low_artifact = record.extras.get("wrist_world_map") + hi_artifact = f"{camera}_world_high.npy" + low_artifact = f"{camera}_world.npy" source_artifact = hi_artifact if resolution == "high" else low_artifact - if not source_artifact: + if source_artifact not in record.artifacts: return { "error": ( f"{camera} {resolution}-resolution world map not recorded " @@ -1785,7 +1594,7 @@ def back_project( } try: - world_map = np.load(state.output_dir / source_artifact) + world_map = state.load(str(source_artifact), step=nn) except Exception as e: return { "error": ( @@ -1871,7 +1680,12 @@ def back_project( depth_m = None if source_artifact == low_artifact: try: - depth = _load_depth(state, camera, nn) + depth_artifact = f"{camera}_depth.npy" + if depth_artifact not in record.artifacts: + raise FileNotFoundError(depth_artifact) + depth = state.load(depth_artifact, step=nn) + if depth.ndim == 3: + depth = depth[..., 0] except Exception as e: return {"error": f"{camera} depth not found for step {nn}: {e}"} depth_m = float(depth[row, col]) diff --git a/rpent/cli/main.py b/rpent/cli/main.py index f8a3bffb..bc62c165 100644 --- a/rpent/cli/main.py +++ b/rpent/cli/main.py @@ -220,7 +220,6 @@ def main() -> int: toolkit = get_toolkit( env_name, primitives_kwargs=primitives_kwargs, - video_path=str(Path(output_dir) / "episode.mp4"), dashboard_events=dashboard_events, ) diff --git a/rpent/dashboard/server.py b/rpent/dashboard/server.py index e485d950..f76c5f06 100644 --- a/rpent/dashboard/server.py +++ b/rpent/dashboard/server.py @@ -259,10 +259,11 @@ def api_frame( @app.get("/api/run/video") def api_video(run: str) -> Response: live = self._resolve(run) - if live is None or not live.has_video(): + video = live.video() if live else None + if video is None: return Response(status_code=404) - return FileResponse( - live.video_path, + return Response( + video, media_type="video/mp4", headers={"Cache-Control": "no-store, max-age=0"}, ) @@ -270,11 +271,11 @@ def api_video(run: str) -> Response: @app.get("/api/run/action-video") def api_action_video(run: str, step: int) -> Response: live = self._resolve(run) - path = live.action_video_path(step) if live else None - if path is None: + video = live.action_video(step) if live else None + if video is None: return Response(status_code=404) - return FileResponse( - path, + return Response( + video, media_type="video/mp4", headers={"Cache-Control": "no-store, max-age=0"}, ) diff --git a/rpent/dashboard/state.py b/rpent/dashboard/state.py index 7617cdea..ab08b087 100644 --- a/rpent/dashboard/state.py +++ b/rpent/dashboard/state.py @@ -502,7 +502,14 @@ def _apply_tool_result(self, event: ToolResultEvent) -> None: result = event.result if not isinstance(result, dict): return - self._apply_frame_paths(result) + frames = { + "camera": result.get("_image_cam_bytes") or result.get("_image_bytes"), + "wrist": result.get("_image_wrist_bytes"), + } + self._update_frames( + step=result.get("step"), + frames={kind: data for kind, data in frames.items() if data}, + ) log = result.get("log") if not isinstance(log, dict): return @@ -521,11 +528,8 @@ def _apply_tool_result(self, event: ToolResultEvent) -> None: "result": log.get("result"), "elapsed_s": log.get("elapsed_s"), "terminated": terminated, - "has_action_video": ( - self.output_dir - / "action_videos" - / f"step_{step:02d}_{command.get('action', name)}.mp4" - ).exists(), + "action_video_artifact": result.get("action_video_artifact"), + "has_action_video": bool(result.get("action_video_artifact")), } with self._lock: self._timeline.append(item) @@ -675,19 +679,21 @@ def frame(self, kind: str) -> bytes | None: with self._lock: return self._frames.get(kind) - def action_video_path(self, step: int) -> Path | None: + def action_video(self, step: int) -> bytes | None: with self._lock: for item in self._timeline: if int(item.get("step", -1)) != int(step): continue - video_path = ( - self.output_dir - / "action_videos" - / f"step_{int(step):02d}_{item.get('action', '')}.mp4" - ) - return video_path if video_path.exists() else None + artifact = item.get("action_video_artifact") + if not artifact: + return None + video_path = self.output_dir / f"{int(step):02d}_{artifact}" + return video_path.read_bytes() if video_path.exists() else None return None + def video(self) -> bytes | None: + return self.video_path.read_bytes() if self.video_path.exists() else None + def has_video(self) -> bool: with self._lock: return ( diff --git a/rpent/planner/api_loop.py b/rpent/planner/api_loop.py index afcdd3e5..1e20d11b 100644 --- a/rpent/planner/api_loop.py +++ b/rpent/planner/api_loop.py @@ -408,7 +408,11 @@ def _build_tools(toolkit: Toolkit, *, no_images: bool = False) -> list[Tool]: def read_image(path: str) -> ToolReturn: - """Read a local image path returned by an RPent tool as visual input.""" + """Read an explicitly provided local image file as visual input. + + Environment observations are already embedded by ``view_driver_state``; + this helper is only for other user-selected local files. + """ return ToolReturn( return_value=path, content=[BinaryContent.from_path(path)], diff --git a/rpent/tools/state.py b/rpent/tools/state.py index 573fde6c..a24d15d9 100644 --- a/rpent/tools/state.py +++ b/rpent/tools/state.py @@ -1,44 +1,15 @@ -"""EnvState: the per-run state-trace owner. - -A **step** is one motion-primitive tool call that produced a dumped state -snapshot, indexed from ``0`` where ``0`` is the post-``reset()`` baseline -(no command). Non-motion tool calls (``get_ee_pose``, ``back_project``, -``view_driver_state``, file/memory tools) are NOT steps and leave the trace -untouched. - -``EnvState`` replaces the per-env ``_append_state`` / ``_load_states`` / -``_latest_step`` / ``_load_step`` / ``_load_image`` / ``_load_depth`` free -functions (and the toolkit's ``_next_step`` counter) with one explicit owner -constructed with an ``output_dir`` (no process-global). The toolkit and the -reader tools (``view_driver_state``, ``back_project``) hold a non-owning -reference to one ``EnvState`` per run; the composition root owns its lifecycle. - -Artefact naming follows the LIBERO layout. Each saved artefact is a named -*stream* turned into a path by a fixed rule:: - - stream "image" -> images/image_NN.png - stream "image_wrist" -> images_wrist/image_wrist_NN.png - stream "depth" -> depths/depth_NN.npy - stream "depth_wrist" -> depths_wrist/depth_wrist_NN.npy - stream "world" -> world/world_NN.npy (libero) - stream "wrist_meta" -> wrist_meta/wrist_meta_NN.json - -i.e. the stream name is the file prefix; the directory pluralises the -artefact type (``image``->``images``, ``depth``->``depths``; ``world`` and -``wrist_meta`` stay singular, matching libero) and appends the camera suffix. -Franka maps scene->``image``/``depth`` and wrist->``image_wrist``/ -``depth_wrist``; lerobot maps scene->``image``/``depth`` and -arm->``image_arm``. No per-env layout class is needed: the env's thin -``dump_state`` wrapper just picks the stream names. -""" +"""Per-run environment state and artifact storage.""" from __future__ import annotations +import copy +import fnmatch import json import os -import shutil +from collections.abc import Iterator +from contextlib import contextmanager from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Iterable +from typing import Any import imageio.v2 as imageio import numpy as np @@ -47,82 +18,47 @@ logger = get_logger("env_state") -# Artefact type (the leading segment of a stream name) -> file extension. -_ARTIFACT_EXT: dict[str, str] = { - "image": ".png", - "depth": ".npy", - "world": ".npy", - "wrist_meta": ".json", -} - -# Artefact type -> directory name (plural where libero pluralises). -_PLURAL: dict[str, str] = { - "image": "images", - "depth": "depths", -} - -def _split_stream(stream: str) -> tuple[str, str]: - for artifact_type in sorted(_ARTIFACT_EXT, key=len, reverse=True): - if stream == artifact_type: - return artifact_type, "" - prefix = artifact_type + "_" - if stream.startswith(prefix): - return artifact_type, stream[len(prefix):] - art, sep, suffix = stream.partition("_") - return art, (suffix if sep else "") - - -def stream_dir(output_dir: Path, stream: str) -> Path: - """Directory for a stream's artefacts (e.g. ``images_wrist``).""" - art, suffix = _split_stream(stream) - base = _PLURAL.get(art, art) - return output_dir / (base if not suffix else f"{base}_{suffix}") - - -def stream_path(output_dir: Path, stream: str, step_idx: int, ext: str | None = None) -> Path: - """Full path for one stream artefact at ``step_idx``.""" - art, _ = _split_stream(stream) - if ext is None: - ext = _ARTIFACT_EXT.get(art, ".bin") - return stream_dir(output_dir, stream) / f"{stream}_{step_idx:02d}{ext}" - - -# --------------------------------------------------------------------------- -# StepRecord -# --------------------------------------------------------------------------- +_MANIFEST_NAME = "states.json" +_MANIFEST_VERSION = 2 +_IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg"} +_TEXT_SUFFIXES = {".txt", ".md"} +_SUPPORTED_SUFFIXES = _IMAGE_SUFFIXES | { + ".npy", + ".json", + ".jsonl", + ".mp4", + ".bin", +} | _TEXT_SUFFIXES + + +def _json_default(value: Any) -> Any: + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + if isinstance(value, Path): + return str(value) + raise TypeError(f"{type(value).__name__} is not JSON serializable") @dataclass class StepRecord: - """One dumped step: the unit the LLM reads back via ``view_driver_state``. - - ``step_idx`` is the motion-primitive index (0 = post-reset baseline, which - has ``command``/``result``/``elapsed_s`` = None). ``frames``/``depth`` are - the stream names saved for this step (e.g. ``["image", "image_wrist"]``). - ``camera_meta`` is the per-camera metadata dict at capture time (intrinsics - + extrinsics + calibration status). ``extras`` absorbs env-specific fields - (libero's world maps, ``libero_terminated``, task_language, ...). - """ + """One motion step and the artifact base names captured for it.""" step_idx: int - state: dict + state: dict[str, Any] + artifacts: set[str] = field(default_factory=set) command: dict | None = None result: dict | None = None elapsed_s: float | None = None - frames: list[str] = field(default_factory=list) - depth: list[str] = field(default_factory=list) - camera_meta: dict | None = None extras: dict[str, Any] = field(default_factory=dict) def to_blob(self) -> dict[str, Any]: blob: dict[str, Any] = { "step_idx": self.step_idx, "state": self.state, - "frames": self.frames, - "depth": self.depth, + "artifacts": sorted(self.artifacts), } - if self.camera_meta is not None: - blob["camera_meta"] = self.camera_meta if self.command is not None: blob["command"] = self.command if self.result is not None: @@ -134,215 +70,382 @@ def to_blob(self) -> dict[str, Any]: return blob @classmethod - def from_blob(cls, blob: dict) -> "StepRecord": + def from_blob(cls, blob: dict[str, Any]) -> "StepRecord": return cls( - step_idx=int(blob.get("step_idx", -1)), - state=blob.get("state", {}), + step_idx=int(blob["step_idx"]), + state=dict(blob.get("state") or {}), + artifacts={str(name) for name in blob.get("artifacts") or []}, command=blob.get("command"), result=blob.get("result"), elapsed_s=blob.get("elapsed_s"), - frames=list(blob.get("frames", [])), - depth=list(blob.get("depth", [])), - camera_meta=blob.get("camera_meta"), extras=dict(blob.get("extras") or {}), ) -# --------------------------------------------------------------------------- -# EnvState -# --------------------------------------------------------------------------- - - class EnvState: - """Owns the on-disk state trace for one run. - - Constructed with an explicit ``output_dir`` (the composition root passes - it; ``EnvState`` never reaches into a process-global). Owns the step - counter, the ``states.json`` trace (atomic append), and the image/depth - artefact layout. The toolkit and the reader tools (``view_driver_state``, - ``back_project``) hold a non-owning reference to one ``EnvState`` per run. - """ + """Own a run's step trace and all state-related files in its output root.""" def __init__(self, output_dir: Path | str): - self.output_dir = Path(output_dir) - self.output_dir.mkdir(parents=True, exist_ok=True) - self._steps: list[dict] = self._read_states() - self._next_step = self._derive_next_step() - - # -- internal: states.json I/O ---------------------------------------- - - def _states_path(self) -> Path: - return self.output_dir / "states.json" + self._output_dir = Path(output_dir) + self._output_dir.mkdir(parents=True, exist_ok=True) + self._steps: list[StepRecord] = [] + self._pending_step: StepRecord | None = None + self._run_artifacts: set[str] = set() + self._next_step = 0 + self.reset() + + # -- private file resolution ----------------------------------------- + + @staticmethod + def _validate_name(name: str) -> str: + if not isinstance(name, str) or not name: + raise ValueError("artifact name must be a non-empty string") + path = Path(name) + if path.name != name or path.is_absolute() or name in {".", ".."}: + raise ValueError(f"artifact name must be a base filename: {name!r}") + if name == _MANIFEST_NAME: + raise ValueError(f"{_MANIFEST_NAME!r} is reserved") + if path.suffix.lower() not in _SUPPORTED_SUFFIXES: + raise ValueError(f"unsupported artifact suffix: {path.suffix or ''}") + return name + + def _artifact_file(self, name: str, step: int | None) -> Path: + name = self._validate_name(name) + if step is None: + return self._output_dir / name + if step < 0: + raise ValueError("artifact writes require a nonnegative step") + return self._output_dir / f"{step:02d}_{name}" + + def _manifest_file(self) -> Path: + return self._output_dir / _MANIFEST_NAME + + def _temporary_file(self, destination: Path) -> Path: + return destination.with_name( + f".{destination.stem}.tmp{destination.suffix}" + ) - def _read_states(self) -> list[dict]: - path = self._states_path() - if not path.exists(): - return [] + def _record_for_write(self, step: int) -> tuple[StepRecord, bool]: + if self._pending_step is not None and self._pending_step.step_idx == step: + return self._pending_step, False + for record in self._steps: + if record.step_idx == step: + return record, True + raise KeyError(f"step {step} not present in state trace") + + # -- manifest -------------------------------------------------------- + + def _write_manifest(self) -> None: + destination = self._manifest_file() + temporary = self._temporary_file(destination) + manifest = { + "version": _MANIFEST_VERSION, + "run_artifacts": sorted(self._run_artifacts), + "steps": [record.to_blob() for record in self._steps], + } try: - with open(path) as f: - arr = json.load(f) - return [s for s in arr if isinstance(s, dict)] if isinstance(arr, list) else [] - except Exception as e: - logger.warning("could not parse %s: %s; starting fresh", path, e) - return [] - - def _write_states_atomically(self, steps: list[dict]) -> None: - path = self._states_path() - tmp = path.with_name(path.name + ".tmp") - with open(tmp, "w") as f: - json.dump(steps, f, indent=2, default=str) - os.replace(tmp, path) - - def _derive_next_step(self) -> int: - if not self._steps: - return 0 - return max(int(s["step_idx"]) for s in self._steps if "step_idx" in s) + 1 - - # -- lifecycle -------------------------------------------------------- - - def reset(self, *, wipe_streams: Iterable[str] = ()) -> None: - """Wipe stale artefact dirs + ``states.json`` for a fresh run. - - Resets the in-memory trace and counter. Step 0 (baseline) is dumped by - the caller right after via :meth:`append`. - """ - self.output_dir.mkdir(parents=True, exist_ok=True) - for stream in wipe_streams: - d = stream_dir(self.output_dir, stream) - if d.exists(): - shutil.rmtree(d) - states = self._states_path() - if states.exists(): - states.unlink() + with temporary.open("w") as file: + json.dump(manifest, file, indent=2, default=_json_default) + os.replace(temporary, destination) + finally: + temporary.unlink(missing_ok=True) + + # -- lifecycle and counters ----------------------------------------- + + def reset(self) -> None: + """Remove state-owned artifacts and start a fresh trace.""" + self._output_dir.mkdir(parents=True, exist_ok=True) + records = list(self._steps) + if self._pending_step is not None: + records.append(self._pending_step) + for record in records: + for name in record.artifacts: + self._artifact_file(name, record.step_idx).unlink(missing_ok=True) + for name in self._run_artifacts: + self._artifact_file(name, None).unlink(missing_ok=True) + self._manifest_file().unlink(missing_ok=True) self._steps = [] + self._pending_step = None + self._run_artifacts = set() self._next_step = 0 - # -- counter ---------------------------------------------------------- - @property def next_step_idx(self) -> int: - """The step_idx the next dumped step will get (does NOT advance).""" return self._next_step @property def latest_step(self) -> int | None: - """The highest step_idx successfully written (None if trace is empty).""" if not self._steps: return None - return int(self._steps[-1]["step_idx"]) + return self._steps[-1].step_idx - # -- writing ---------------------------------------------------------- - - def save_image(self, step_idx: int, stream: str, frame) -> str | None: - """Save one RGB frame under ``_.png``; return ``stream`` on success.""" - path = stream_path(self.output_dir, stream, step_idx) - path.parent.mkdir(parents=True, exist_ok=True) - arr = np.asarray(frame) - if arr.dtype != np.uint8: - arr = arr.astype(np.uint8) - try: - imageio.imwrite(path, arr) - return stream - except Exception as e: - logger.warning("frame dump failed for stream %s: %s", stream, e) + def _resolve_read_step(self, step: int | None) -> int | None: + if step is None: return None - - def save_depth(self, step_idx: int, stream: str, depth) -> str | None: - """Save one depth map under ``_.npy``; return ``stream`` on success.""" - path = stream_path(self.output_dir, stream, step_idx) - path.parent.mkdir(parents=True, exist_ok=True) + if step == -1: + latest = self.latest_step + if latest is None: + raise LookupError("no steps available") + return latest + if step < 0: + raise ValueError("step must be -1, None, or nonnegative") + return step + + # -- generic artifact operations ------------------------------------ + + def save( + self, + name: str, + value: Any, + *, + step: int | None, + **options: Any, + ) -> str | None: + """Serialize ``value`` according to ``name`` and return its base name.""" + destination = self._artifact_file(name, step) + temporary = self._temporary_file(destination) + suffix = destination.suffix.lower() + record: StepRecord | None = None + committed = False + if step is not None: + record, committed = self._record_for_write(step) try: - np.save(path, np.asarray(depth, dtype=np.float32)) - return stream - except Exception as e: - logger.warning("depth dump failed for stream %s: %s", stream, e) + if suffix in _IMAGE_SUFFIXES: + array = np.asarray(value) + if array.dtype != np.uint8: + array = array.astype(np.uint8) + imageio.imwrite(temporary, array) + elif suffix == ".npy": + np.save(temporary, np.asarray(value)) + elif suffix == ".json": + with temporary.open("w") as file: + json.dump(value, file, indent=2, default=_json_default) + elif suffix == ".jsonl": + with temporary.open("w") as file: + if isinstance(value, str): + file.write(value) + else: + for item in value: + file.write( + json.dumps(item, default=_json_default) + "\n" + ) + elif suffix == ".mp4": + if isinstance(value, (bytes, bytearray, memoryview)): + temporary.write_bytes(bytes(value)) + else: + imageio.mimwrite( + temporary, + list(value), + fps=int(options.get("fps", 20)), + ) + elif suffix in _TEXT_SUFFIXES: + temporary.write_text(str(value)) + elif suffix == ".bin": + temporary.write_bytes(bytes(value)) + else: + raise ValueError(f"unsupported artifact suffix: {suffix}") + os.replace(temporary, destination) + if record is not None: + record.artifacts.add(name) + if committed: + self._write_manifest() + else: + self._run_artifacts.add(name) + self._write_manifest() + return name + except Exception as exc: + logger.warning("failed to save artifact %s: %s", name, exc) return None + finally: + temporary.unlink(missing_ok=True) + + def load(self, name: str, *, step: int | None = -1) -> Any: + """Load an artifact; ``step=-1`` selects the latest recorded step.""" + resolved_step = self._resolve_read_step(step) + source = self._artifact_file(name, resolved_step) + suffix = source.suffix.lower() + if suffix in _IMAGE_SUFFIXES: + return imageio.imread(source) + if suffix == ".npy": + return np.load(source) + if suffix == ".json": + with source.open() as file: + return json.load(file) + if suffix == ".jsonl": + with source.open() as file: + return [json.loads(line) for line in file if line.strip()] + if suffix in _TEXT_SUFFIXES: + return source.read_text() + return source.read_bytes() + + def load_bytes(self, name: str, *, step: int | None = -1) -> bytes: + resolved_step = self._resolve_read_step(step) + return self._artifact_file(name, resolved_step).read_bytes() + + def exists(self, name: str, *, step: int | None = -1) -> bool: + try: + resolved_step = self._resolve_read_step(step) + except LookupError: + return False + return self._artifact_file(name, resolved_step).exists() + + def remove(self, name: str, *, step: int | None) -> bool: + destination = self._artifact_file(name, step) + if not destination.exists(): + return False + record: StepRecord | None = None + committed = False + if step is not None: + record, committed = self._record_for_write(step) + destination.unlink() + if step is None and name in self._run_artifacts: + self._run_artifacts.remove(name) + self._write_manifest() + elif record is not None: + was_recorded = name in record.artifacts + record.artifacts.discard(name) + if was_recorded and committed: + self._write_manifest() + return True + + def list( + self, + pattern: str = "*", + *, + step: int | None = -1, + ) -> list[str]: + resolved_step = self._resolve_read_step(step) + if resolved_step is None: + names = [ + name + for name in self._run_artifacts + if self._artifact_file(name, None).exists() + ] + else: + record = self.get(resolved_step) + names = [ + name + for name in record.artifacts + if self._artifact_file(name, resolved_step).exists() + ] + return sorted(name for name in names if fnmatch.fnmatch(name, pattern)) + + def list_all(self, pattern: str = "*") -> list[tuple[int | None, str]]: + artifacts: list[tuple[int | None, str]] = [ + (None, name) for name in self.list(pattern, step=None) + ] + for record in self._steps: + artifacts.extend( + (record.step_idx, name) + for name in record.artifacts + if fnmatch.fnmatch(name, pattern) + and self._artifact_file(name, record.step_idx).exists() + ) + return sorted( + artifacts, + key=lambda item: (-1 if item[0] is None else item[0], item[1]), + ) - def artifact_path( + # -- step records ---------------------------------------------------- + + @contextmanager + def record_step( self, - step_idx: int, - stream: str, *, - ext: str | None = None, - ) -> Path: - """Return the canonical path for one step artifact.""" - return stream_path(self.output_dir, stream, step_idx, ext=ext) - - def append(self, record: StepRecord) -> StepRecord: - """Atomically append ``record`` to ``states.json`` and advance the counter. - - The counter and ``states.json`` are kept in sync: this write is the - single source of truth for "latest step". If the write raises, the - counter is NOT advanced (so the next caller reuses the same - ``step_idx``) -- no desync between the in-memory counter and disk. - """ - blob = record.to_blob() - self._write_states_atomically(self._steps + [blob]) - self._steps.append(blob) - self._next_step = max(self._next_step, record.step_idx + 1) - return record - - # -- reading ---------------------------------------------------------- - - def get(self, step_idx: int) -> StepRecord: - """Look up the step record for ``step_idx``.""" - for blob in self._steps: - if int(blob.get("step_idx", -1)) == step_idx: - return StepRecord.from_blob(blob) - raise KeyError(f"step {step_idx} not present in states.json") + state: dict[str, Any], + command: dict | None = None, + result: dict | None = None, + elapsed_s: float | None = None, + extras: dict[str, Any] | None = None, + ) -> Iterator[int]: + if self._pending_step is not None: + raise RuntimeError("a step record is already open") + record = StepRecord( + step_idx=self._next_step, + state=copy.deepcopy(state), + command=copy.deepcopy(command), + result=copy.deepcopy(result), + elapsed_s=elapsed_s, + extras=copy.deepcopy(extras or {}), + ) + self._pending_step = record + try: + yield record.step_idx + except BaseException: + raise + else: + self._steps.append(record) + try: + self._write_manifest() + except Exception: + self._steps.pop() + raise + self._next_step = record.step_idx + 1 + finally: + self._pending_step = None + + def get(self, step: int = -1) -> StepRecord: + resolved_step = self._resolve_read_step(step) + if resolved_step is None: + raise ValueError("step records are not run-level artifacts") + for record in self._steps: + if record.step_idx == resolved_step: + return copy.deepcopy(record) + raise KeyError(f"step {resolved_step} not present in state trace") def records(self) -> list[StepRecord]: - """Return the persisted step records in trace order.""" - return [StepRecord.from_blob(blob) for blob in self._steps] + return copy.deepcopy(self._steps) - def load_image_bytes(self, step_idx: int, stream: str) -> bytes | None: - path = stream_path(self.output_dir, stream, step_idx) - if not path.exists(): - return None - return path.read_bytes() - - def load_depth(self, step_idx: int, stream: str) -> np.ndarray: - path = stream_path(self.output_dir, stream, step_idx) - return np.load(path) - - # -- LLM-facing view (ex-view_driver_state) --------------------------- - - def view(self, step: int | None = None, *, image_slots: dict[str, str] | None = None) -> dict: - """Build the ``view_driver_state`` dict for one step. - - ``image_slots`` maps a ToolResult image slot (``_image_bytes``, - ``_image_cam_bytes``, ``_image_wrist_bytes``) to a stream name whose - saved PNG should be embedded as bytes. The env decides which cameras - map to which slots. - """ - latest = self.latest_step - if latest is None: - return {"error": "no driver state entries; driver not ready"} - nn = latest if step is None else int(step) + # -- LLM-facing view ------------------------------------------------- + + def view( + self, + step: int = -1, + *, + image_slots: dict[str, str] | None = None, + ) -> dict[str, Any]: try: - rec = self.get(nn) - except Exception as e: - return {"error": f"step {nn} not present in driver state trace: {e}"} + record = self.get(step) + except Exception as exc: + return {"error": f"state step not available: {exc}"} + + metadata: dict[str, Any] = {} + metadata_suffix = "_metadata.json" + for name in sorted(record.artifacts): + if not name.endswith(metadata_suffix): + continue + key = name.removesuffix(metadata_suffix) + try: + loaded = self.load(name, step=record.step_idx) + if isinstance(loaded, dict): + loaded = { + field: value + for field, value in loaded.items() + if field not in {"K", "T_base_cam"} + } + metadata[key] = loaded + except Exception as exc: + metadata[key] = {"error": str(exc)} + out: dict[str, Any] = { - "step": nn, - "state": rec.state, - "frames": rec.frames, - "depth": rec.depth, - "camera_meta": { - name: {k: v for k, v in meta.items() if k not in {"K", "T_base_cam"}} - for name, meta in (rec.camera_meta or {}).items() - }, + "step": record.step_idx, + "state": record.state, + "artifacts": sorted(record.artifacts), + "camera_meta": metadata, "log": { - "command": rec.command, - "result": rec.result, - "elapsed_s": rec.elapsed_s, + "command": record.command, + "result": record.result, + "elapsed_s": record.elapsed_s, }, } - if rec.extras: - out["extras"] = rec.extras + if record.extras: + out["extras"] = record.extras if image_slots: - for slot, stream in image_slots.items(): - b = self.load_image_bytes(nn, stream) - if b: - out[slot] = b + for slot, name in image_slots.items(): + if name not in record.artifacts: + continue + try: + out[slot] = self.load_bytes(name, step=record.step_idx) + except FileNotFoundError: + continue return out diff --git a/rpent/utils/sam3_client.py b/rpent/utils/sam3_client.py index f8a80037..22c23977 100644 --- a/rpent/utils/sam3_client.py +++ b/rpent/utils/sam3_client.py @@ -5,7 +5,6 @@ import base64 import io from dataclasses import dataclass -from pathlib import Path from typing import Any import imageio.v2 as imageio @@ -35,7 +34,7 @@ def __init__(self, client: RpcClient, *, timeout_s: float = 120.0) -> None: def segment( self, - image_path: str | Path, + image: bytes | bytearray | memoryview | np.ndarray, *, text_prompt: str | None = None, point: list[int] | None = None, @@ -55,7 +54,12 @@ def segment( if not 0.0 <= float(min_score) <= 1.0: raise ValueError("min_score must be between 0 and 1") - image_bytes = Path(image_path).read_bytes() + if isinstance(image, np.ndarray): + buffer = io.BytesIO() + imageio.imwrite(buffer, image, format="png") + image_bytes = buffer.getvalue() + else: + image_bytes = bytes(image) body: dict[str, Any] = { "image_base64": base64.b64encode(image_bytes).decode("ascii"), "min_score": float(min_score), From 118c94ff81046283ac633a4c6b82066a02bc97c8 Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Wed, 5 Aug 2026 03:30:40 +0000 Subject: [PATCH 19/22] refactor(toolkit): centralize state capture in base Toolkit.execute_tool Signed-off-by: Jiaxing Qiu --- robots/franka/toolkit.py | 36 ++++++++++------ robots/franka/tools.py | 23 ++++------ robots/lerobot/toolkit.py | 43 ++++++++++-------- robots/lerobot/tools.py | 15 +++---- robots/libero/toolkit.py | 64 +++++++++++++-------------- rpent/tools/toolkit.py | 91 +++++++++++++++++++++++++++++++++------ 6 files changed, 170 insertions(+), 102 deletions(-) diff --git a/robots/franka/toolkit.py b/robots/franka/toolkit.py index de96df81..14c57493 100644 --- a/robots/franka/toolkit.py +++ b/robots/franka/toolkit.py @@ -1,7 +1,6 @@ """Franka toolkit: common tools plus conservative Cartesian primitives.""" from __future__ import annotations -import time from functools import partial from typing import Any @@ -58,16 +57,28 @@ def _register_tools(self) -> None: "view_driver_state", spec["view_driver_state"], partial(self._state.view, image_slots=self._VIEW_IMAGE_SLOTS), + captures_state=False, ) self.add_tool( "back_project", spec["back_project"], partial(franka_tools.back_project, state=self._state), + captures_state=False, ) for name in self._DRIVER_READERS: - self.add_tool(name, spec[name], self._make_driver_reader(name)) + self.add_tool( + name, + spec[name], + self._make_driver_reader(name), + captures_state=False, + ) for name in self._PRIMITIVE_TOOLS: - self.add_tool(name, spec[name], partial(self._step, name)) + self.add_tool( + name, + spec[name], + getattr(self._driver, name), + captures_state=True, + ) def _make_driver_reader(self, name: str): def _reader(**kwargs) -> dict: @@ -76,20 +87,21 @@ def _reader(**kwargs) -> dict: return _reader - def _step(self, name: str, **kwargs) -> dict: - command = {"action": name, **kwargs} - t0 = time.time() - result = getattr(self._driver, name)(**kwargs) - elapsed = round(time.time() - t0, 2) - result_dict = result if isinstance(result, dict) else {"value": result} - + def get_state( + self, + *, + command: dict[str, Any], + result: dict[str, Any], + elapsed_s: float, + ) -> dict[str, Any]: + self._driver._refresh(increment_step=command["action"] != "observe") record = franka_tools.dump_state( self._driver, self._state, - log={"command": command, "result": result_dict, "elapsed_s": elapsed}, + log={"command": command, "result": result, "elapsed_s": elapsed_s}, ) out = self._state.view(record.step_idx, image_slots=self._VIEW_IMAGE_SLOTS) - out["agent_elapsed_s"] = elapsed + out["agent_elapsed_s"] = elapsed_s return out def init_driver_clean(self, *, env: Any) -> None: diff --git a/robots/franka/tools.py b/robots/franka/tools.py index c4d2b08c..56738c69 100644 --- a/robots/franka/tools.py +++ b/robots/franka/tools.py @@ -59,11 +59,10 @@ def reset(self) -> tuple[dict, Any]: return obs, info def observe(self, delay_s: float = 0.0) -> dict: - """Refresh the cached observation without moving.""" + """Wait before the toolkit captures a fresh observation.""" delay = float(np.clip(delay_s, 0.0, _MAX_OBSERVE_DELAY_S)) if delay > 0: time.sleep(delay) - self._last_obs = self.env.get_obs() return {"delay_s": delay} def get_robot_spec(self) -> dict: @@ -88,9 +87,7 @@ def move_to( gripper: str | None = None, ) -> dict: """Move to an absolute base-frame Cartesian target.""" - result = self.env.move_to(xyz, yaw_deg=yaw_deg, gripper=gripper) - self._refresh() - return result + return self.env.move_to(xyz, yaw_deg=yaw_deg, gripper=gripper) def move_delta( self, @@ -106,7 +103,6 @@ def move_delta( result = dict(result) result["requested_dxyz"] = _to_list(requested) result["clipped_dxyz"] = _to_list(clipped) - self._refresh() return result def rotate_wrist_yaw(self, delta_deg: float) -> dict: @@ -118,7 +114,6 @@ def rotate_wrist_yaw(self, delta_deg: float) -> dict: result = dict(result) result["requested_delta_deg"] = round(requested, 3) result["clipped_delta_deg"] = round(clipped, 3) - self._refresh() return result def rotate_gripper(self, delta_deg: float) -> dict: @@ -126,14 +121,10 @@ def rotate_gripper(self, delta_deg: float) -> dict: return self.rotate_wrist_yaw(delta_deg) def open_gripper(self) -> dict: - result = self.env.open_gripper() - self._refresh() - return result + return self.env.open_gripper() def close_gripper(self) -> dict: - result = self.env.close_gripper() - self._refresh() - return result + return self.env.close_gripper() def get_state(self) -> dict: """Return compact proprioception from the latest observation.""" @@ -171,10 +162,12 @@ def latest_camera_meta(self) -> dict: return {} return dict(self._last_obs.get("camera_meta", {})) - def _refresh(self) -> None: + def _refresh(self, *, increment_step: bool = True) -> None: + """Refresh the cached observation for toolkit state capture.""" try: self._last_obs = self.env.get_obs() - self._num_steps += 1 + if increment_step: + self._num_steps += 1 except Exception as exc: logger.warning("obs refresh failed: %s", exc) diff --git a/robots/lerobot/toolkit.py b/robots/lerobot/toolkit.py index 88b7069c..32c061d4 100644 --- a/robots/lerobot/toolkit.py +++ b/robots/lerobot/toolkit.py @@ -6,7 +6,6 @@ """ from __future__ import annotations -import time from functools import partial from typing import Any @@ -29,8 +28,7 @@ class LerobotToolkit(Toolkit): "get_ee_pose", "get_scene_camera_meta", ) - # Primitive tools routed through ``_step`` (look up driver method by name). - # Each moves the robot and re-renders state after running. + # Primitive tools move the robot; get_state refreshes and records state. _PRIMITIVE_TOOLS: tuple[str, ...] = ( "move_to", "move_joints_delta", @@ -60,16 +58,28 @@ def _register_tools(self) -> None: "view_driver_state", spec["view_driver_state"], partial(self._state.view, image_slots=self._VIEW_IMAGE_SLOTS), + captures_state=False, ) self.add_tool( "back_project", spec["back_project"], partial(lerobot_tools.back_project, state=self._state), + captures_state=False, ) for name in self._DRIVER_READERS: - self.add_tool(name, spec[name], self._make_driver_reader(name)) + self.add_tool( + name, + spec[name], + self._make_driver_reader(name), + captures_state=False, + ) for name in self._PRIMITIVE_TOOLS: - self.add_tool(name, spec[name], partial(self._step, name)) + self.add_tool( + name, + spec[name], + getattr(self._driver, name), + captures_state=True, + ) def _make_driver_reader(self, name: str): """Bind a read-only tool to ``self._driver.`` (no state dump).""" @@ -78,24 +88,21 @@ def _reader(**kwargs) -> dict: return result if isinstance(result, dict) else {"value": result} return _reader - def _step(self, name: str, **kwargs) -> dict: - """Run ``self._driver.(**kwargs)``, dump the new step, and - return the rendered state view + log. - """ - command = {"action": name, **kwargs} - t0 = time.time() - result = getattr(self._driver, name)(**kwargs) - elapsed = round(time.time() - t0, 2) - - result_dict = result if isinstance(result, dict) else {"value": result} - + def get_state( + self, + *, + command: dict[str, Any], + result: dict[str, Any], + elapsed_s: float, + ) -> dict[str, Any]: + self._driver._refresh() record = lerobot_tools.dump_state( self._driver, self._state, - log={"command": command, "result": result_dict, "elapsed_s": elapsed}, + log={"command": command, "result": result, "elapsed_s": elapsed_s}, ) out = self._state.view(record.step_idx, image_slots=self._VIEW_IMAGE_SLOTS) - out["agent_elapsed_s"] = elapsed + out["agent_elapsed_s"] = elapsed_s return out def init_driver_clean(self, *, env: Any, model: Any | None = None) -> None: diff --git a/robots/lerobot/tools.py b/robots/lerobot/tools.py index 89bfd7ab..346d6663 100644 --- a/robots/lerobot/tools.py +++ b/robots/lerobot/tools.py @@ -132,7 +132,7 @@ def get_scene_camera_meta(self) -> dict: self._scene_meta = self.env.get_scene_camera_meta() return self._scene_meta - # -- primitives (move the robot, then refresh the cached observation) --- + # -- primitives (move the robot; toolkit capture refreshes observation) -- def move_to( self, @@ -146,11 +146,9 @@ def move_to( ``approach="down"`` keeps the gripper pointing straight down (for grasping); ``yaw_deg`` sets the jaw heading. See the driver for details. """ - result = self.env.move_to( + return self.env.move_to( xyz, gripper=gripper, approach=approach, yaw_deg=yaw_deg ) - self._refresh() - return result def move_joints_delta( self, @@ -158,12 +156,13 @@ def move_joints_delta( gripper_delta: float | None = None, ) -> dict: """Nudge each arm joint relatively (degrees) for fine alignment.""" - result = self.env.move_joints_delta(delta_deg, gripper_delta=gripper_delta) - self._refresh() - return result + return self.env.move_joints_delta( + delta_deg, + gripper_delta=gripper_delta, + ) def _refresh(self) -> None: - """Refresh the cached observation after a motion primitive.""" + """Refresh the cached observation for toolkit state capture.""" try: self._last_obs = self.env.get_obs() except Exception as e: diff --git a/robots/libero/toolkit.py b/robots/libero/toolkit.py index b6ae17d5..89a5c1da 100644 --- a/robots/libero/toolkit.py +++ b/robots/libero/toolkit.py @@ -5,14 +5,13 @@ """ from __future__ import annotations -import time from functools import partial from typing import Any from robots.libero import tools as libero_tools from rpent.dashboard.events import DashboardEventSink, ToolResultEvent from rpent.tools.state import EnvState -from rpent.tools.toolkit import ToolCancelled, Toolkit +from rpent.tools.toolkit import Toolkit from rpent.utils.logging import get_logger, get_output_dir @@ -52,46 +51,40 @@ def _register_libero_tools(self) -> None: "segment": partial(self._primitives.segment, state=self._state), } for name, handler in inspection_handlers.items(): - self.add_tool(name, specs[name], handler) - # Primitive tools: each goes through _step, which looks up the - # matching primitive method via getattr at call time. + self.add_tool( + name, + specs[name], + handler, + captures_state=False, + ) for name in libero_tools.PRIMITIVE_TOOL_NAMES: - self.add_tool(name, specs[name], partial(self._step, name)) - - def _step(self, name: str, **kwargs) -> dict: - """Run ``self._primitives.(**kwargs)``, dump the new step, and - return the rendered state view + log. - """ - command = {"action": name, **kwargs} - t0 = time.time() - start_frame = self._primitives.recorded_frame_count() - try: - result = getattr(self._primitives, name)(**kwargs) - self.raise_if_cancelled() - except ToolCancelled as exc: - result = { - "error": str(exc), - "code": "tool_cancelled", - "interrupted": True, - } - elapsed = round(time.time() - t0, 2) - - if isinstance(result, dict): - result_dict = result - else: - result_dict = {"value": result} + self.add_tool( + name, + specs[name], + getattr(self._primitives, name), + captures_state=True, + ) + def get_state( + self, + *, + command: dict[str, Any], + result: dict[str, Any], + elapsed_s: float, + ) -> dict[str, Any]: + frame_start = self._action_frame_cursor + self._action_frame_cursor = self._primitives.recorded_frame_count() record = libero_tools.dump_state( self._primitives, self._state, - log={"command": command, "result": result_dict, "elapsed_s": elapsed}, + log={"command": command, "result": result, "elapsed_s": elapsed_s}, ) action_video_name = None if self._dashboard_events.enabled: try: - frames = self._primitives.frame_slice(start_frame) + frames = self._primitives.frame_slice(frame_start) if frames: - candidate = f"action_{name}.mp4" + candidate = f"action_{command['action']}.mp4" if self._state.save( candidate, frames, @@ -106,9 +99,9 @@ def _step(self, name: str, **kwargs) -> dict: e, ) out = libero_tools.view_driver_state(record.step_idx, state=self._state) - out["agent_elapsed_s"] = elapsed - if result_dict.get("interrupted"): - out.update(result_dict) + out["agent_elapsed_s"] = elapsed_s + if result.get("interrupted"): + out.update(result) if action_video_name is not None: out["action_video_artifact"] = action_video_name return out @@ -127,6 +120,7 @@ def init_primitives_clean( ) primitives.reset() primitives.start_recording() + self._action_frame_cursor = primitives.recorded_frame_count() libero_tools.dump_state(primitives, self._state, log=None) self._dashboard_events.emit( ToolResultEvent( diff --git a/rpent/tools/toolkit.py b/rpent/tools/toolkit.py index a747020f..7502b1f3 100644 --- a/rpent/tools/toolkit.py +++ b/rpent/tools/toolkit.py @@ -7,8 +7,10 @@ from __future__ import annotations import base64 +import inspect import json import threading +import time import traceback from collections.abc import Callable from dataclasses import dataclass, field @@ -112,8 +114,10 @@ def __init__( dashboard_events: DashboardEventSink, state: Any = None, ) -> None: - # name -> (spec, handler) - self._tools: dict[str, tuple[dict[str, Any], Callable[..., dict[str, Any]]]] = {} + self._tools: dict[ + str, + tuple[dict[str, Any], Callable[..., Any], bool], + ] = {} self._dashboard_events = dashboard_events self._state = state self._operation_lock = threading.Lock() @@ -128,7 +132,9 @@ def add_tool( self, name: str, spec: dict[str, Any], - handler: Callable[..., dict[str, Any]], + handler: Callable[..., Any], + *, + captures_state: bool, ) -> None: """Register one tool under ``name`` with its schema and handler. @@ -138,8 +144,10 @@ def add_tool( ``description``, ``input_schema``). handler: Callable invoked with the tool's input kwargs; returns a result dict. + captures_state: Whether the tool updates environment state + that must be captured after the handler returns or raises. """ - self._tools[name] = (spec, handler) + self._tools[name] = (spec, handler, captures_state) def _register_common_tools(self) -> None: """Register the file/IO tools shared by every run.""" @@ -147,7 +155,12 @@ def _register_common_tools(self) -> None: for spec in common.TOOLS_SPEC: name = spec["name"] - self.add_tool(name, spec, common.TOOL_HANDLERS[name]) + self.add_tool( + name, + spec, + common.TOOL_HANDLERS[name], + captures_state=False, + ) # ------------------------------------------------------------------ # Planner-facing API @@ -156,39 +169,89 @@ def _register_common_tools(self) -> None: def get_tools_spec(self) -> list[dict[str, Any]]: """Return the tool schemas the LLM sees.""" return substitute( - [spec for spec, _ in self._tools.values()] + [spec for spec, _, _ in self._tools.values()] ) def execute_tool(self, name: str, input_dict: dict[str, Any]) -> ToolResult: """Dispatch a tool call to its registered handler.""" entry = self._tools.get(name) if entry is None: - return ToolResult(name=name, result={"error": f"unknown tool: {name}"}) - handler = entry[1] + return self._finish_tool(name, {"error": f"unknown tool: {name}"}) + _, handler, captures_state = entry with self._operation_lock: if self._active_operation is not None: - return ToolResult( - name=name, - result={"error": "another tool operation is still active"}, + return self._finish_tool( + name, + {"error": "another tool operation is still active"}, ) operation = _ToolOperation() self._active_operation = operation try: + started = time.perf_counter() + failed = False try: result = handler(**input_dict) except TypeError as e: - result = {"error": f"bad arguments for {name}: {e}", "got": input_dict} + return self._finish_tool( + name, + {"error": f"bad arguments for {name}: {e}", "got": input_dict}, + ) + except ToolCancelled as e: + result = { + "error": str(e), + "code": "tool_cancelled", + "interrupted": True, + } + failed = True except Exception as e: result = {"error": str(e), "traceback": traceback.format_exc()} - self._dashboard_events.emit(ToolResultEvent(name=name, result=result)) - return ToolResult(name=name, result=result) + failed = True + + if captures_state: + elapsed_s = round(time.perf_counter() - started, 2) + result_dict = result if isinstance(result, dict) else {"value": result} + command = {"action": name, **input_dict} + try: + captured = self.get_state( + command=command, + result=result_dict, + elapsed_s=elapsed_s, + ) + except Exception as e: + captured = result_dict + captured["state_capture_error"] = str(e) + captured.setdefault( + "error", f"failed to capture state after {name}: {e}" + ) + captured.setdefault("traceback", traceback.format_exc()) + result = captured + if failed: + result.setdefault("error", result_dict["error"]) + if "traceback" in result_dict: + result.setdefault("traceback", result_dict["traceback"]) + + return self._finish_tool(name, result) finally: with self._operation_lock: self._active_operation = None operation.done_event.set() + def _finish_tool(self, name: str, result: Any) -> ToolResult: + self._dashboard_events.emit(ToolResultEvent(name=name, result=result)) + return ToolResult(name=name, result=result) + + def get_state( + self, + *, + command: dict[str, Any], + result: dict[str, Any], + elapsed_s: float, + ) -> dict[str, Any]: + """Capture and return the observation produced by a stateful tool.""" + raise NotImplementedError + # ------------------------------------------------------------------ # Server lifecycle hooks (overridden by env toolkits) # ------------------------------------------------------------------ From d8c0916be3923269adb97e76d1730b7f85e20521 Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Wed, 5 Aug 2026 09:08:19 +0000 Subject: [PATCH 20/22] refactor(toolkit): mark state-advancing primitives with @updatestate Signed-off-by: Jiaxing Qiu --- .../rst_source/development/add_primitive.rst | 24 +++-- .../rst_source/development/add_robot.rst | 11 ++- .../rst_source/development/add_primitive.rst | 23 +++-- .../rst_source/development/add_robot.rst | 8 +- robots/franka/toolkit.py | 70 ++++----------- robots/franka/tools.py | 8 ++ robots/lerobot/toolkit.py | 66 ++++---------- robots/lerobot/tools.py | 3 + robots/libero/toolkit.py | 37 ++++---- robots/libero/tools.py | 35 +++++--- rpent/tools/state.py | 87 ++++++++++--------- rpent/tools/toolkit.py | 48 +++++++--- 12 files changed, 207 insertions(+), 213 deletions(-) diff --git a/docs/source-en/rst_source/development/add_primitive.rst b/docs/source-en/rst_source/development/add_primitive.rst index 1148f94c..297851e5 100644 --- a/docs/source-en/rst_source/development/add_primitive.rst +++ b/docs/source-en/rst_source/development/add_primitive.rst @@ -35,7 +35,7 @@ call. They differ only in how the method is implemented. Add a scripted primitive ------------------------ -Adding a scripted primitive usually involves three steps: +Adding a scripted primitive usually involves two steps: 1. **Add a method to the primitives.** Add the method to the current environment's primitives class, such as @@ -43,16 +43,25 @@ Adding a scripted primitive usually involves three steps: the tool-call arguments, performs the work, usually through one or more ``self._env.step(...)`` calls, and returns a small log ``dict``. + Mark the method with :func:`~rpent.tools.toolkit.updatestate` so the + toolkit re-renders state (``get_state``) automatically after it runs: + .. code-block:: python + from rpent.tools.toolkit import updatestate + + @updatestate def open_drawer(self, dx: float = 0.15) -> dict: # Move end-effector back by dx while gripper is closed. for _ in range(N): self._env.step(build_open_drawer_chunk(dx)) return {"ok": True, "dx": dx} + Read-only tools (``view_driver_state``, ``back_project``, ``segment``, + ...) are simply left unmarked -- the toolkit skips state capture for them. + 2. **Add the tool schema.** Add an entry to ``TOOLS_SPEC`` in - ``toolkit.py``: + ``robots//tools.py``: .. code-block:: python @@ -67,13 +76,10 @@ Adding a scripted primitive usually involves three steps: }, } -3. **Register the tool in the toolkit.** Route it through the toolkit's - ``_step`` helper so that state is re-rendered after execution: - - .. code-block:: python - - self.add_tool("open_drawer", OPEN_DRAWER_SPEC, - lambda **kw: self._step("open_drawer", **kw)) +Once both exist, the toolkit registers the tool automatically: it iterates +``TOOLS_SPEC`` and binds each spec to the matching primitive-driver method +(e.g. ``getattr(self._primitives, name)``); ``@updatestate`` decides whether +state is captured -- no explicit ``add_tool`` call is needed. After these steps, the ``api``, ``claude_code``, and ``codex`` planners can all call the primitive without any other code changes. diff --git a/docs/source-en/rst_source/development/add_robot.rst b/docs/source-en/rst_source/development/add_robot.rst index 0a2b5497..7f61d588 100644 --- a/docs/source-en/rst_source/development/add_robot.rst +++ b/docs/source-en/rst_source/development/add_robot.rst @@ -223,10 +223,13 @@ plus any module-level functions referenced by the toolkit (e.g. ``view_driver_state``, ``back_project``, ``finish``). **Per-step state dump** — ``dump_state(driver, env_state, log)`` opens -``env_state.record_step(...)`` and receives the allocated step index. Save -large observations through ``env_state.save(..., step=step_idx)``. ``EnvState`` -owns and commits the ``StepRecord`` and adds every successfully saved base name -to its flat ``artifacts`` set automatically. Readers use the canonical artifact +``env_state.record_step(...)`` and receives the allocated step index; the +``StepRecord`` is appended and committed immediately. Save large observations +through ``env_state.save(...)`` — inside a ``record_step`` block the ``step`` +argument may be omitted (it defaults to the new step), pass an explicit +``step=`` to target a different step, and ``step=None`` for run-level +artifacts. ``EnvState`` adds every successfully saved base name to the step's +flat ``artifacts`` set automatically. Readers use the canonical artifact filenames rather than maintaining a parallel observation index. **Toolkit class** — subclass ``rpent.tools.toolkit.Toolkit``: diff --git a/docs/source-zh/rst_source/development/add_primitive.rst b/docs/source-zh/rst_source/development/add_primitive.rst index 2ba6ce24..5d475ea3 100644 --- a/docs/source-zh/rst_source/development/add_primitive.rst +++ b/docs/source-zh/rst_source/development/add_primitive.rst @@ -33,22 +33,31 @@ primitives 方法，以及调用完成后的状态快照。区别仅在于方法 添加一个脚本化原语 ------------------ -添加脚本化原语通常需要以下三个步骤： +添加脚本化原语通常需要以下两个步骤： 1. **在 primitives 中添加方法。** 在当前环境的 primitives 类（如 ``LiberoPrimitives``、``MyRobotPrimitives``）中添加 一个方法。该方法接收工具调用的参数，执行一次或多次 ``self._env.step(...)``，并返回一个简短的日志字典。 + 为该方法加上 :func:`~rpent.tools.toolkit.updatestate` 装饰器， + toolkit 会在其执行后自动重新渲染状态（``get_state``）： + .. code-block:: python + from rpent.tools.toolkit import updatestate + + @updatestate def open_drawer(self, dx: float = 0.15) -> dict: # 保持夹爪闭合，沿 -x 方向后拉 dx 米。 for _ in range(N): self._env.step(build_open_drawer_chunk(dx)) return {"ok": True, "dx": dx} -2. **添加工具定义。** 在 ``toolkit.py`` 的 ``TOOLS_SPEC`` 中新增一项： + 只读工具（``view_driver_state``、``back_project``、``segment`` 等） + 无需装饰——toolkit 会跳过它们的状态捕获。 + +2. **添加工具定义。** 在 ``robots//tools.py`` 的 ``TOOLS_SPEC`` 中新增一项： .. code-block:: python @@ -63,13 +72,9 @@ primitives 方法，以及调用完成后的状态快照。区别仅在于方法 }, } -3. **在 toolkit 中注册工具。** 通过 toolkit 的 ``_step`` 辅助函数运行 - 该工具，使其在执行结束后自动重新渲染状态： - - .. code-block:: python - - self.add_tool("open_drawer", OPEN_DRAWER_SPEC, - lambda **kw: self._step("open_drawer", **kw)) +两者就位后，toolkit 会自动注册该工具：它遍历 ``TOOLS_SPEC``，把每个定义 +绑定到对应的 primitive driver 方法（如 ``getattr(self._primitives, name)``）， +由 ``@updatestate`` 决定是否捕获状态——无需显式调用 ``add_tool``。 完成以上步骤后，``api``、``claude_code`` 和 ``codex`` 三种 planner 都可以调用该工具，无需修改其他代码。 diff --git a/docs/source-zh/rst_source/development/add_robot.rst b/docs/source-zh/rst_source/development/add_robot.rst index 0f76750c..394d9e90 100644 --- a/docs/source-zh/rst_source/development/add_robot.rst +++ b/docs/source-zh/rst_source/development/add_robot.rst @@ -211,9 +211,11 @@ Anthropic API 的工具定义格式，包含 ``name``、``description`` 和 **每步状态 dump** —— ``dump_state(driver, env_state, log)`` 通过 ``env_state.record_step(...)`` 创建由 ``EnvState`` 持有的步骤，并取得分配的 -step index。大型观测通过 ``env_state.save(..., step=step_idx)`` 保存。每次保存 -成功后，``EnvState`` 会自动把基础文件名加入该 ``StepRecord`` 的扁平 -``artifacts`` 集合并最终提交记录；读取方直接使用规范化的工件文件名。 +step index；该 ``StepRecord`` 会被立即追加并提交。大型观测通过 +``env_state.save(...)`` 保存——在 ``record_step`` 块内可省略 ``step`` 参数 +（默认指向刚创建的步骤），传显式 ``step=`` 可指定其它步骤，``step=None`` +用于运行级工件。每次保存成功后，``EnvState`` 会自动把基础文件名加入该 +``StepRecord`` 的扁平 ``artifacts`` 集合；读取方直接使用规范化的工件文件名。 **Toolkit 类** 继承 ``rpent.tools.toolkit.Toolkit``： diff --git a/robots/franka/toolkit.py b/robots/franka/toolkit.py index 14c57493..45741645 100644 --- a/robots/franka/toolkit.py +++ b/robots/franka/toolkit.py @@ -19,23 +19,6 @@ class FrankaToolkit(Toolkit): "_image_cam_bytes": "wrist.png", } - _DRIVER_READERS = ( - "get_ee_pose", - "get_robot_spec", - "get_camera_meta", - ) - _PRIMITIVE_TOOLS = ( - "observe", - "move_to", - "move_delta", - "rotate_wrist_yaw", - "rotate_gripper", - "open_gripper", - "close_gripper", - ) - - _SPECS = {spec["name"]: spec for spec in franka_tools.TOOLS_SPEC} - def __init__( self, *, @@ -51,41 +34,24 @@ def __init__( self._register_tools() def _register_tools(self) -> None: - spec = self._SPECS - # view_driver_state / back_project read trace state -> bind to EnvState. - self.add_tool( - "view_driver_state", - spec["view_driver_state"], - partial(self._state.view, image_slots=self._VIEW_IMAGE_SLOTS), - captures_state=False, - ) - self.add_tool( - "back_project", - spec["back_project"], - partial(franka_tools.back_project, state=self._state), - captures_state=False, - ) - for name in self._DRIVER_READERS: - self.add_tool( - name, - spec[name], - self._make_driver_reader(name), - captures_state=False, - ) - for name in self._PRIMITIVE_TOOLS: - self.add_tool( - name, - spec[name], - getattr(self._driver, name), - captures_state=True, - ) - - def _make_driver_reader(self, name: str): - def _reader(**kwargs) -> dict: - result = getattr(self._driver, name)(**kwargs) - return result if isinstance(result, dict) else {"value": result} - - return _reader + # Read-only tools whose handlers aren't driver methods (they need the + # run's EnvState bound in). Every other spec binds to its primitive- + # driver method; @updatestate on the method decides state capture. + state_handlers = { + "view_driver_state": partial( + self._state.view, image_slots=self._VIEW_IMAGE_SLOTS + ), + "back_project": partial(franka_tools.back_project, state=self._state), + } + for spec in franka_tools.TOOLS_SPEC: + name = spec["name"] + if name in state_handlers: + handler = state_handlers[name] + else: + handler = getattr(self._driver, name, None) + if handler is None: + continue # spec without a backing driver method + self.add_tool(name, spec, handler) def get_state( self, diff --git a/robots/franka/tools.py b/robots/franka/tools.py index 56738c69..d379f2a7 100644 --- a/robots/franka/tools.py +++ b/robots/franka/tools.py @@ -9,6 +9,7 @@ from robots.franka.env_client import FrankaEnvClient from rpent.tools.common import robust_surface_centroid from rpent.tools.state import EnvState, StepRecord +from rpent.tools.toolkit import updatestate from rpent.utils.logging import get_logger logger = get_logger("franka") @@ -58,6 +59,7 @@ def reset(self) -> tuple[dict, Any]: self._num_steps = 0 return obs, info + @updatestate def observe(self, delay_s: float = 0.0) -> dict: """Wait before the toolkit captures a fresh observation.""" delay = float(np.clip(delay_s, 0.0, _MAX_OBSERVE_DELAY_S)) @@ -79,6 +81,7 @@ def get_camera_meta(self) -> dict: """Return live camera intrinsics/extrinsics metadata.""" return self.env.get_camera_meta() + @updatestate def move_to( self, xyz, @@ -89,6 +92,7 @@ def move_to( """Move to an absolute base-frame Cartesian target.""" return self.env.move_to(xyz, yaw_deg=yaw_deg, gripper=gripper) + @updatestate def move_delta( self, dxyz, @@ -105,6 +109,7 @@ def move_delta( result["clipped_dxyz"] = _to_list(clipped) return result + @updatestate def rotate_wrist_yaw(self, delta_deg: float) -> dict: """Rotate the wrist yaw relatively, capped for safety.""" requested = float(delta_deg) @@ -116,13 +121,16 @@ def rotate_wrist_yaw(self, delta_deg: float) -> dict: result["clipped_delta_deg"] = round(clipped, 3) return result + @updatestate def rotate_gripper(self, delta_deg: float) -> dict: """Rotate the gripper jaw heading relatively, capped for safety.""" return self.rotate_wrist_yaw(delta_deg) + @updatestate def open_gripper(self) -> dict: return self.env.open_gripper() + @updatestate def close_gripper(self) -> dict: return self.env.close_gripper() diff --git a/robots/lerobot/toolkit.py b/robots/lerobot/toolkit.py index 32c061d4..75d184b4 100644 --- a/robots/lerobot/toolkit.py +++ b/robots/lerobot/toolkit.py @@ -22,20 +22,6 @@ class LerobotToolkit(Toolkit): "_image_bytes": "scene.png", "_image_cam_bytes": "arm.png", } - # Read-only tools backed by a live driver call (no state dump). These query - # the robot/scene directly: forward kinematics + scene camera calibration. - _DRIVER_READERS = ( - "get_ee_pose", - "get_scene_camera_meta", - ) - # Primitive tools move the robot; get_state refreshes and records state. - _PRIMITIVE_TOOLS: tuple[str, ...] = ( - "move_to", - "move_joints_delta", - ) - - # Tool schemas keyed by name, built once from the canonical ordered list. - _SPECS = {spec["name"]: spec for spec in lerobot_tools.TOOLS_SPEC} def __init__( self, @@ -53,40 +39,24 @@ def __init__( # Registration # ------------------------------------------------------------------ def _register_tools(self) -> None: - spec = self._SPECS - self.add_tool( - "view_driver_state", - spec["view_driver_state"], - partial(self._state.view, image_slots=self._VIEW_IMAGE_SLOTS), - captures_state=False, - ) - self.add_tool( - "back_project", - spec["back_project"], - partial(lerobot_tools.back_project, state=self._state), - captures_state=False, - ) - for name in self._DRIVER_READERS: - self.add_tool( - name, - spec[name], - self._make_driver_reader(name), - captures_state=False, - ) - for name in self._PRIMITIVE_TOOLS: - self.add_tool( - name, - spec[name], - getattr(self._driver, name), - captures_state=True, - ) - - def _make_driver_reader(self, name: str): - """Bind a read-only tool to ``self._driver.`` (no state dump).""" - def _reader(**kwargs) -> dict: - result = getattr(self._driver, name)(**kwargs) - return result if isinstance(result, dict) else {"value": result} - return _reader + # Read-only tools whose handlers aren't driver methods (they need the + # run's EnvState bound in). Every other spec binds to its primitive- + # driver method; @updatestate on the method decides state capture. + state_handlers = { + "view_driver_state": partial( + self._state.view, image_slots=self._VIEW_IMAGE_SLOTS + ), + "back_project": partial(lerobot_tools.back_project, state=self._state), + } + for spec in lerobot_tools.TOOLS_SPEC: + name = spec["name"] + if name in state_handlers: + handler = state_handlers[name] + else: + handler = getattr(self._driver, name, None) + if handler is None: + continue # spec without a backing driver method + self.add_tool(name, spec, handler) def get_state( self, diff --git a/robots/lerobot/tools.py b/robots/lerobot/tools.py index 346d6663..81a8b227 100644 --- a/robots/lerobot/tools.py +++ b/robots/lerobot/tools.py @@ -23,6 +23,7 @@ from robots.lerobot.env_client import LerobotEnvClient from rpent.tools.common import robust_surface_centroid from rpent.tools.state import EnvState, StepRecord +from rpent.tools.toolkit import updatestate from rpent.utils.logging import get_logger logger = get_logger("lerobot") @@ -134,6 +135,7 @@ def get_scene_camera_meta(self) -> dict: # -- primitives (move the robot; toolkit capture refreshes observation) -- + @updatestate def move_to( self, xyz, @@ -150,6 +152,7 @@ def move_to( xyz, gripper=gripper, approach=approach, yaw_deg=yaw_deg ) + @updatestate def move_joints_delta( self, delta_deg, diff --git a/robots/libero/toolkit.py b/robots/libero/toolkit.py index 89a5c1da..95b9253f 100644 --- a/robots/libero/toolkit.py +++ b/robots/libero/toolkit.py @@ -18,10 +18,6 @@ class LiberoToolkit(Toolkit): """Toolkit for the LIBERO environment.""" - # Tool schemas keyed by name (built once from the canonical ordered list - # in libero_tools.TOOLS_SPEC) so each tool registers with its own spec. - _SPECS = {spec["name"]: spec for spec in libero_tools.TOOLS_SPEC} - def __init__( self, *, @@ -37,10 +33,12 @@ def __init__( # Registration # ------------------------------------------------------------------ def _register_libero_tools(self) -> None: - specs = self._SPECS - # Inspection tools do not advance environment state. Most are stateless - # module functions; segment is bound to the primitives-owned SAM3 client. - inspection_handlers = { + # Read-only tools whose handlers aren't primitive methods (they need + # the run's EnvState bound in, or -- like segment -- must stay + # read-only despite being a primitives method). Every other spec binds + # to its primitive-driver method; @updatestate on the method decides + # whether state is captured. + state_handlers = { "view_driver_state": partial( libero_tools.view_driver_state, state=self._state ), @@ -50,20 +48,15 @@ def _register_libero_tools(self) -> None: "back_project": partial(libero_tools.back_project, state=self._state), "segment": partial(self._primitives.segment, state=self._state), } - for name, handler in inspection_handlers.items(): - self.add_tool( - name, - specs[name], - handler, - captures_state=False, - ) - for name in libero_tools.PRIMITIVE_TOOL_NAMES: - self.add_tool( - name, - specs[name], - getattr(self._primitives, name), - captures_state=True, - ) + for spec in libero_tools.TOOLS_SPEC: + name = spec["name"] + if name in state_handlers: + handler = state_handlers[name] + else: + handler = getattr(self._primitives, name, None) + if handler is None: + continue # spec without a backing primitive method + self.add_tool(name, spec, handler) def get_state( self, diff --git a/robots/libero/tools.py b/robots/libero/tools.py index 7f511e8e..ff8cb901 100644 --- a/robots/libero/tools.py +++ b/robots/libero/tools.py @@ -8,6 +8,7 @@ from robots.libero.env_client import LiberoEnvClient from rpent.tools.state import EnvState, StepRecord +from rpent.tools.toolkit import updatestate from rpent.utils.logging import get_logger from rpent.utils.sam3_client import Sam3Client from rpent.utils.vla_client import VLAClient @@ -124,6 +125,7 @@ def _vlm_chunk(self, instruction: str): if original_task is not None: self._last_obs["task_descriptions"] = original_task + @updatestate def pi0_pick( self, prompt: str, @@ -199,6 +201,7 @@ def pi0_pick( }, } + @updatestate def pi0_doubled( self, prompt: str, @@ -240,6 +243,7 @@ def pi0_doubled( }, } + @updatestate def move_to( self, xyz, @@ -304,6 +308,7 @@ def move_to( "libero_terminated": self.env.episode_terminated, } + @updatestate def rotate_wrist( self, *, @@ -377,6 +382,7 @@ def _yaw_of(quat_xyzw): "libero_terminated": self.env.episode_terminated, } + @updatestate def rotate_pitch( self, *, @@ -458,6 +464,7 @@ def _pitch_of(quat_xyzw): "libero_terminated": self.env.episode_terminated, } + @updatestate def move_pose( self, xyz, @@ -528,6 +535,7 @@ def _yaw_of(q): "libero_terminated": self.env.episode_terminated, } + @updatestate def release( self, *, @@ -556,6 +564,7 @@ def release( "libero_terminated": self.env.episode_terminated, } + @updatestate def set_gripper( self, *, @@ -725,6 +734,19 @@ def segment( return result +def _is_primitive_action(name: object) -> bool: + """Whether ``name`` is a state-advancing LIBERO primitive. + + A primitive is any ``@updatestate``-marked method on + :class:`LiberoPrimitives`; read-only tools (``view_driver_state``, + ``back_project``, ``segment``, ...) and non-strings read as ``False``. + """ + if not isinstance(name, str): + return False + method = getattr(LiberoPrimitives, name, None) + return method is not None and bool(getattr(method, "_captures_state", False)) + + def write_recipe_from_states(state: EnvState, recipe_tag: str) -> str: """Find a command sequence that gets ``libero_terminated=True``. @@ -736,7 +758,7 @@ def write_recipe_from_states(state: EnvState, recipe_tag: str) -> str: result = record.result if ( command is not None - and command.get("action") in PRIMITIVE_TOOL_NAMES + and _is_primitive_action(command.get("action")) and not (isinstance(result, dict) and result.get("error")) ): command_events.append(((record.step_idx, -1), command)) @@ -1032,17 +1054,6 @@ def _save_observation_artifacts( # Tool schema declarations (Anthropic-shaped canonical schema) # --------------------------------------------------------------------------- -PRIMITIVE_TOOL_NAMES: tuple[str, ...] = ( - "move_to", - "pi0_pick", - "pi0_doubled", - "release", - "set_gripper", - "rotate_wrist", - "rotate_pitch", - "move_pose", -) - TOOLS_SPEC = [ { "name": "view_driver_state", diff --git a/rpent/tools/state.py b/rpent/tools/state.py index a24d15d9..3d43677b 100644 --- a/rpent/tools/state.py +++ b/rpent/tools/state.py @@ -89,8 +89,8 @@ def __init__(self, output_dir: Path | str): self._output_dir = Path(output_dir) self._output_dir.mkdir(parents=True, exist_ok=True) self._steps: list[StepRecord] = [] - self._pending_step: StepRecord | None = None self._run_artifacts: set[str] = set() + self._open_count = 0 self._next_step = 0 self.reset() @@ -125,12 +125,11 @@ def _temporary_file(self, destination: Path) -> Path: f".{destination.stem}.tmp{destination.suffix}" ) - def _record_for_write(self, step: int) -> tuple[StepRecord, bool]: - if self._pending_step is not None and self._pending_step.step_idx == step: - return self._pending_step, False + def _record_for(self, step: int) -> StepRecord: + """Return the live step record at ``step`` for in-place updates.""" for record in self._steps: if record.step_idx == step: - return record, True + return record raise KeyError(f"step {step} not present in state trace") # -- manifest -------------------------------------------------------- @@ -155,18 +154,15 @@ def _write_manifest(self) -> None: def reset(self) -> None: """Remove state-owned artifacts and start a fresh trace.""" self._output_dir.mkdir(parents=True, exist_ok=True) - records = list(self._steps) - if self._pending_step is not None: - records.append(self._pending_step) - for record in records: + for record in self._steps: for name in record.artifacts: self._artifact_file(name, record.step_idx).unlink(missing_ok=True) for name in self._run_artifacts: self._artifact_file(name, None).unlink(missing_ok=True) self._manifest_file().unlink(missing_ok=True) self._steps = [] - self._pending_step = None self._run_artifacts = set() + self._open_count = 0 self._next_step = 0 @property @@ -198,17 +194,23 @@ def save( name: str, value: Any, *, - step: int | None, + step: int | None = -1, **options: Any, ) -> str | None: - """Serialize ``value`` according to ``name`` and return its base name.""" + """Serialize ``value`` according to ``name`` and return its base name. + + ``step`` defaults to ``-1`` (the most recently recorded step), so calls + made inside (or right after) a :meth:`record_step` block attach to that + step without an explicit index. Pass an ``int`` to target a specific + step, or ``None`` for a run-level artifact such as an episode video. + """ + step = self._resolve_read_step(step) destination = self._artifact_file(name, step) temporary = self._temporary_file(destination) suffix = destination.suffix.lower() - record: StepRecord | None = None - committed = False - if step is not None: - record, committed = self._record_for_write(step) + record: StepRecord | None = ( + self._record_for(step) if step is not None else None + ) try: if suffix in _IMAGE_SUFFIXES: array = np.asarray(value) @@ -247,11 +249,9 @@ def save( os.replace(temporary, destination) if record is not None: record.artifacts.add(name) - if committed: - self._write_manifest() else: self._run_artifacts.add(name) - self._write_manifest() + self._write_manifest() return name except Exception as exc: logger.warning("failed to save artifact %s: %s", name, exc) @@ -293,19 +293,12 @@ def remove(self, name: str, *, step: int | None) -> bool: destination = self._artifact_file(name, step) if not destination.exists(): return False - record: StepRecord | None = None - committed = False - if step is not None: - record, committed = self._record_for_write(step) + if step is None: + self._run_artifacts.discard(name) + else: + self._record_for(step).artifacts.discard(name) destination.unlink() - if step is None and name in self._run_artifacts: - self._run_artifacts.remove(name) - self._write_manifest() - elif record is not None: - was_recorded = name in record.artifacts - record.artifacts.discard(name) - if was_recorded and committed: - self._write_manifest() + self._write_manifest() return True def list( @@ -358,7 +351,13 @@ def record_step( elapsed_s: float | None = None, extras: dict[str, Any] | None = None, ) -> Iterator[int]: - if self._pending_step is not None: + """Append a new step record and yield its index. + + The step is committed to the trace immediately, so any subsequent + :meth:`save` without an explicit ``step`` attaches to it. If the block + raises, the step and any artifacts written to it are rolled back. + """ + if self._open_count: raise RuntimeError("a step record is already open") record = StepRecord( step_idx=self._next_step, @@ -368,21 +367,27 @@ def record_step( elapsed_s=elapsed_s, extras=copy.deepcopy(extras or {}), ) - self._pending_step = record + self._steps.append(record) + self._next_step = record.step_idx + 1 + self._open_count += 1 + self._write_manifest() try: yield record.step_idx except BaseException: - raise - else: - self._steps.append(record) + for name in list(record.artifacts): + self._artifact_file(name, record.step_idx).unlink(missing_ok=True) + if self._steps and self._steps[-1] is record: + self._steps.pop() + self._next_step = record.step_idx try: self._write_manifest() - except Exception: - self._steps.pop() - raise - self._next_step = record.step_idx + 1 + except Exception as exc: + logger.warning( + "failed to rewrite manifest after step rollback: %s", exc + ) + raise finally: - self._pending_step = None + self._open_count -= 1 def get(self, step: int = -1) -> StepRecord: resolved_step = self._resolve_read_step(step) diff --git a/rpent/tools/toolkit.py b/rpent/tools/toolkit.py index 7502b1f3..22d30871 100644 --- a/rpent/tools/toolkit.py +++ b/rpent/tools/toolkit.py @@ -7,7 +7,6 @@ from __future__ import annotations import base64 -import inspect import json import threading import time @@ -30,6 +29,37 @@ class ToolCancelled(Exception): """Raised when an environment reaches a safe cancellation boundary.""" +def updatestate(func): + """Mark a tool handler as one that advances environment state. + + The toolkit captures a fresh observation (:meth:`Toolkit.get_state`) + after a handler carrying this marker runs (or raises). Apply it to + primitive-driver methods that move the robot; read-only tools (file IO, + ``view_driver_state``, ``back_project``, ``segment``, ...) are left + unmarked and skip state capture. + + The marker is read off the underlying function, so registering a bound + method (``getattr(self._driver, name)``) inherits it automatically. + ``functools.partial`` does not delegate attribute access, so a partial + wrapping a marked method is treated as read-only (intentional -- see the + LIBERO ``segment`` tool). + """ + func._captures_state = True + return func + + +def _captures_state(handler: Callable[..., Any]) -> bool: + """Whether ``handler`` was marked with :func:`updatestate`. + + Resolves through ``__func__`` so bound methods (the common case for + primitive-driver tools) report the marker set on their underlying + function. Plain functions, closures, and ``functools.partial`` objects + do not delegate, so undecorated read-only handlers read as ``False``. + """ + target = getattr(handler, "__func__", handler) + return bool(getattr(target, "_captures_state", False)) + + @dataclass class ToolResult: """Result of executing one tool call. @@ -133,8 +163,6 @@ def add_tool( name: str, spec: dict[str, Any], handler: Callable[..., Any], - *, - captures_state: bool, ) -> None: """Register one tool under ``name`` with its schema and handler. @@ -143,11 +171,10 @@ def add_tool( spec: Anthropic-shaped tool schema dict (``name``, ``description``, ``input_schema``). handler: Callable invoked with the tool's input kwargs; returns - a result dict. - captures_state: Whether the tool updates environment state - that must be captured after the handler returns or raises. + a result dict. Decorate state-advancing primitives with + :func:`updatestate`. """ - self._tools[name] = (spec, handler, captures_state) + self._tools[name] = (spec, handler, _captures_state(handler)) def _register_common_tools(self) -> None: """Register the file/IO tools shared by every run.""" @@ -155,12 +182,7 @@ def _register_common_tools(self) -> None: for spec in common.TOOLS_SPEC: name = spec["name"] - self.add_tool( - name, - spec, - common.TOOL_HANDLERS[name], - captures_state=False, - ) + self.add_tool(name, spec, common.TOOL_HANDLERS[name]) # ------------------------------------------------------------------ # Planner-facing API From 099e55419bd6bf226d3da81c1a7f01bf49f8d345 Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Wed, 5 Aug 2026 09:22:14 +0000 Subject: [PATCH 21/22] chore: several renames Signed-off-by: Jiaxing Qiu --- .../rst_source/development/add_primitive.rst | 6 +++--- .../rst_source/development/add_robot.rst | 6 +++--- .../rst_source/development/interfaces.rst | 2 +- docs/source-en/rst_source/quickstart.rst | 2 +- docs/source-en/rst_source/usage/libero.rst | 4 ++-- .../rst_source/development/add_primitive.rst | 6 +++--- .../rst_source/development/add_robot.rst | 6 +++--- .../rst_source/development/interfaces.rst | 2 +- docs/source-zh/rst_source/quickstart.rst | 2 +- docs/source-zh/rst_source/usage/libero.rst | 4 ++-- robots/franka/prompt.py | 8 ++++---- robots/franka/toolkit.py | 6 +++--- robots/franka/tools.py | 6 +++--- robots/lerobot/prompt.py | 4 ++-- robots/lerobot/toolkit.py | 6 +++--- robots/lerobot/tools.py | 4 ++-- robots/libero/guides/env_calibration.md | 8 ++++---- robots/libero/guides/pro_hybrid_guide.md | 8 ++++---- robots/libero/guides/strict_hybrid_guide.md | 6 +++--- robots/libero/prompts/system.py | 12 ++++++------ robots/libero/prompts/user.py | 4 ++-- robots/libero/toolkit.py | 12 ++++++------ robots/libero/tools.py | 10 +++++----- rpent/context/prompts/prompt.py | 2 +- rpent/planner/api_loop.py | 4 ++-- rpent/tools/toolkit.py | 16 ++++++++-------- scripts/codex_proxy/litellm_callbacks.py | 2 +- 27 files changed, 79 insertions(+), 79 deletions(-) diff --git a/docs/source-en/rst_source/development/add_primitive.rst b/docs/source-en/rst_source/development/add_primitive.rst index 297851e5..947a38bf 100644 --- a/docs/source-en/rst_source/development/add_primitive.rst +++ b/docs/source-en/rst_source/development/add_primitive.rst @@ -44,7 +44,7 @@ Adding a scripted primitive usually involves two steps: more ``self._env.step(...)`` calls, and returns a small log ``dict``. Mark the method with :func:`~rpent.tools.toolkit.updatestate` so the - toolkit re-renders state (``get_state``) automatically after it runs: + toolkit re-renders state (``get_env_state``) automatically after it runs: .. code-block:: python @@ -57,7 +57,7 @@ Adding a scripted primitive usually involves two steps: self._env.step(build_open_drawer_chunk(dx)) return {"ok": True, "dx": dx} - Read-only tools (``view_driver_state``, ``back_project``, ``segment``, + Read-only tools (``view_env_state``, ``back_project``, ``segment``, ...) are simply left unmarked -- the toolkit skips state capture for them. 2. **Add the tool schema.** Add an entry to ``TOOLS_SPEC`` in @@ -175,7 +175,7 @@ Design principles for a new primitive - **Return small dicts.** Tool return values are fed back to the LLM as text. Save larger observations through ``EnvState.save``; ``EnvState`` automatically records each logical base name in its owned - ``StepRecord.artifacts`` set. Expose images through ``view_driver_state`` and + ``StepRecord.artifacts`` set. Expose images through ``view_env_state`` and geometry through environment tools rather than returning raw paths. - **Guardrails belong in env_server**, not in the toolkit. The LLM can and will call any tool with any arguments; workspace bounds diff --git a/docs/source-en/rst_source/development/add_robot.rst b/docs/source-en/rst_source/development/add_robot.rst index 7f61d588..b3578486 100644 --- a/docs/source-en/rst_source/development/add_robot.rst +++ b/docs/source-en/rst_source/development/add_robot.rst @@ -220,7 +220,7 @@ state needed for the current run. It exposes one method per primitive tool **Tool definitions and handlers** — a module-level ``TOOLS_SPEC`` list of Anthropic-style tool definitions (``name``, ``description``, ``input_schema``), plus any module-level functions referenced by the toolkit (e.g. -``view_driver_state``, ``back_project``, ``finish``). +``view_env_state``, ``back_project``, ``finish``). **Per-step state dump** — ``dump_state(driver, env_state, log)`` opens ``env_state.record_step(...)`` and receives the allocated step index; the @@ -238,7 +238,7 @@ filenames rather than maintaining a parallel observation index. helper (named ``init_primitives_clean`` in LIBERO; it calls ``EnvState.reset()``, constructs the primitives, and dumps step 0), - register each tool with ``self.add_tool(name, spec, handler)`` — stateless - readers (``view_driver_state``, ``finish``, …) bind directly to module-level + readers (``view_env_state``, ``finish``, …) bind directly to module-level functions; primitive tools route through ``_step(name, **kwargs)`` which calls ``getattr(self._primitives, name)(**kwargs)`` and re-renders state, - override ``close()`` to save remaining agent-side artifacts through @@ -260,7 +260,7 @@ Conventions worth keeping exposed to all planners. - Server-side return values must be picklable and torch-free. - Each primitive tool dumps a fresh state snapshot after running so the next - ``view_driver_state`` call reflects the post-action world. + ``view_env_state`` call reflects the post-action world. - Treat ``dump_state`` as the source of truth for what the agent sees — any new modality (e.g. tactile, force) goes through it. diff --git a/docs/source-en/rst_source/development/interfaces.rst b/docs/source-en/rst_source/development/interfaces.rst index bf6a6a59..02d865a0 100644 --- a/docs/source-en/rst_source/development/interfaces.rst +++ b/docs/source-en/rst_source/development/interfaces.rst @@ -90,7 +90,7 @@ Subclass ``Toolkit`` in ``robots//toolkit.py`` and register env tools with ends; optional ``_image_bytes`` (etc.) to return camera images. The base class already registers common file tools; call ``super().__init__()`` then -``add_tool`` for env tools. Per-step state and ``view_driver_state`` are in +``add_tool`` for env tools. Per-step state and ``view_env_state`` are in :doc:`add_primitive`. Inter-process communication diff --git a/docs/source-en/rst_source/quickstart.rst b/docs/source-en/rst_source/quickstart.rst index 8fb27a7d..ac49ab68 100644 --- a/docs/source-en/rst_source/quickstart.rst +++ b/docs/source-en/rst_source/quickstart.rst @@ -106,7 +106,7 @@ A successful run: 4. By default, artifacts are saved under ``logs/__t_s/``. They include ``transcript_*.json`` (run record), ``states.json`` (the versioned ``EnvState`` manifest), ``recipe_*.jsonl`` (action sequence), and ``episode.mp4`` (episode video). Step artifact files use flat, zero-padded step prefixes with a minimum width of two digits and are managed internally by ``EnvState``. Inspect the final state through the Dashboard or -``view_driver_state(step=-1)``. Its top-level ``libero_terminated`` value is the +``view_env_state(step=-1)``. Its top-level ``libero_terminated`` value is the benchmark outcome. ``states.json`` is internal ``EnvState`` storage and should not be parsed by callers. You can also open ``episode.mp4`` to review the run. If something goes wrong, inspect the four log files described at the diff --git a/docs/source-en/rst_source/usage/libero.rst b/docs/source-en/rst_source/usage/libero.rst index b2f7682b..7e90c120 100644 --- a/docs/source-en/rst_source/usage/libero.rst +++ b/docs/source-en/rst_source/usage/libero.rst @@ -116,7 +116,7 @@ What runs where transports (HTTP or socket). It returns only the top compressed PNG mask. - **toolkit** (``robots/libero/toolkit.py``) — defines the tools the LLM can call: ``pi0_pick`` (fed to Pi0.5), ``move_to``, - ``rotate_wrist``, ``back_project``, ``view_driver_state``, + ``rotate_wrist``, ``back_project``, ``view_env_state``, ``finish``, … Tools the planner can call @@ -147,7 +147,7 @@ Physical action tools advance the environment and record new state and images. coordinates. - ``segment(prompt=... / point=..., ...)`` — use SAM3 to segment an existing image with a text or point prompt. -- ``view_driver_state(step=-1)`` — read a recorded state and its embedded +- ``view_env_state(step=-1)`` — read a recorded state and its embedded observation images. Step ``0`` is initial; ``-1`` is latest. - ``view_camera_meta(camera=..., step=-1)`` — read camera metadata for a recorded step. Step ``-1`` is latest. diff --git a/docs/source-zh/rst_source/development/add_primitive.rst b/docs/source-zh/rst_source/development/add_primitive.rst index 5d475ea3..59022970 100644 --- a/docs/source-zh/rst_source/development/add_primitive.rst +++ b/docs/source-zh/rst_source/development/add_primitive.rst @@ -41,7 +41,7 @@ primitives 方法，以及调用完成后的状态快照。区别仅在于方法 ``self._env.step(...)``，并返回一个简短的日志字典。 为该方法加上 :func:`~rpent.tools.toolkit.updatestate` 装饰器， - toolkit 会在其执行后自动重新渲染状态（``get_state``）： + toolkit 会在其执行后自动重新渲染状态（``get_env_state``）： .. code-block:: python @@ -54,7 +54,7 @@ primitives 方法，以及调用完成后的状态快照。区别仅在于方法 self._env.step(build_open_drawer_chunk(dx)) return {"ok": True, "dx": dx} - 只读工具（``view_driver_state``、``back_project``、``segment`` 等） + 只读工具（``view_env_state``、``back_project``、``segment`` 等） 无需装饰——toolkit 会跳过它们的状态捕获。 2. **添加工具定义。** 在 ``robots//tools.py`` 的 ``TOOLS_SPEC`` 中新增一项： @@ -159,7 +159,7 @@ primitives 方法，以及调用完成后的状态快照。区别仅在于方法 - **工具只返回简短的字典。** 返回值会以文本形式提供给 LLM；图像、深度数据和 其他大型观测应通过 ``EnvState.save`` 保存；``EnvState`` 会把每个逻辑基础 文件名自动加入其持有的 ``StepRecord.artifacts`` 集合。图像通过 - ``view_driver_state`` 提供，几何数据通过环境工具访问，不返回原始路径。 + ``view_env_state`` 提供，几何数据通过环境工具访问，不返回原始路径。 - **安全限制由 ``env_server`` 强制执行。** LLM 可能使用任意参数调用工具， 因此工作空间边界和安全限制不能只依赖 toolkit。 diff --git a/docs/source-zh/rst_source/development/add_robot.rst b/docs/source-zh/rst_source/development/add_robot.rst index 394d9e90..e646a37d 100644 --- a/docs/source-zh/rst_source/development/add_robot.rst +++ b/docs/source-zh/rst_source/development/add_robot.rst @@ -207,7 +207,7 @@ toolkit 模块通常包含四部分： **工具定义和处理函数** 包括模块级的 ``TOOLS_SPEC`` 列表（列表元素采用 Anthropic API 的工具定义格式，包含 ``name``、``description`` 和 ``input_schema``），以及 toolkit 引用的模块级函数，例如 -``view_driver_state``、``back_project`` 和 ``finish``。 +``view_env_state``、``back_project`` 和 ``finish``。 **每步状态 dump** —— ``dump_state(driver, env_state, log)`` 通过 ``env_state.record_step(...)`` 创建由 ``EnvState`` 持有的步骤，并取得分配的 @@ -223,7 +223,7 @@ step index；该 ``StepRecord`` 会被立即追加并提交。大型观测通过 中的方法名为 ``init_primitives_clean``；它会调用 ``EnvState.reset()``、构造 原语并 dump 第 0 步）, - 用 ``self.add_tool(name, spec, handler)`` 注册每个工具。无状态的读取工具 - （如 ``view_driver_state``、``finish``）直接绑定模块级函数；原语工具通过 + （如 ``view_env_state``、``finish``）直接绑定模块级函数；原语工具通过 ``_step(name, **kwargs)`` 调用。``_step`` 使用 ``getattr(self._primitives, name)(**kwargs)`` 调用 driver 方法并重新渲染状态； - 重写 ``close()``，通过 ``EnvState`` 保存 agent 侧剩余工件（例如 @@ -244,7 +244,7 @@ primitives 的 ``__init__``。其中通常包含 每个用 ``self.add_tool(...)`` 注册的工具都会暴露给所有 planner。 - 环境侧的返回值必须可 pickle，且不包含 torch 对象。 - 每个原语工具执行后要 dump 一次新的状态快照, 这样下一次 - ``view_driver_state`` 看到的是动作后的世界。 + ``view_env_state`` 看到的是动作后的世界。 - ``dump_state`` 是 Agent 获取环境状态的唯一数据来源；任何新的模态 （例如触觉、力）都通过它提供。 diff --git a/docs/source-zh/rst_source/development/interfaces.rst b/docs/source-zh/rst_source/development/interfaces.rst index d28c3d34..822b02fa 100644 --- a/docs/source-zh/rst_source/development/interfaces.rst +++ b/docs/source-zh/rst_source/development/interfaces.rst @@ -87,7 +87,7 @@ Planner 需要回传相机图时可设 ``_image_bytes`` 等字段。 基类已注册公共文件工具；子类 ``super().__init__()`` 后追加本环境工具即可。逐步状态与 -``view_driver_state`` 见 :doc:`add_primitive`。 +``view_env_state`` 见 :doc:`add_primitive`。 进程间通信 ---------- diff --git a/docs/source-zh/rst_source/quickstart.rst b/docs/source-zh/rst_source/quickstart.rst index 9d870b55..2b153eb0 100644 --- a/docs/source-zh/rst_source/quickstart.rst +++ b/docs/source-zh/rst_source/quickstart.rst @@ -97,7 +97,7 @@ LIBERO-PRO 仿真资源。下面以 LIBERO-PRO 和 ``claude_code`` planner 3. 启用 Dashboard 后，智能体的输出、相机视图、动作时间线和片段回放也会实时显示在 Dashboard 中。 4. 默认输出目录为 ``logs/__t_s/``，其中包含 ``transcript_*.json``\ （运行记录）、``states.json``\ （带版本号的 ``EnvState`` 清单）、``recipe_*.jsonl``\ （动作序列）和 ``episode.mp4``\ （回合录像）。每步工件文件采用至少两位、零填充的步骤前缀扁平命名，并由 ``EnvState`` 在内部管理。 -通过 Dashboard 或 ``view_driver_state(step=-1)`` 查看最终状态；其顶层 +通过 Dashboard 或 ``view_env_state(step=-1)`` 查看最终状态；其顶层 ``libero_terminated`` 即为基准任务结果。``states.json`` 是 ``EnvState`` 的内部 存储，调用方不应直接解析。也可以打开 ``episode.mp4`` 复核运行过程。 出问题时，参考 :doc:`installation` 页底部提到的四份日志文件。 diff --git a/docs/source-zh/rst_source/usage/libero.rst b/docs/source-zh/rst_source/usage/libero.rst index b66e06ac..940de6da 100644 --- a/docs/source-zh/rst_source/usage/libero.rst +++ b/docs/source-zh/rst_source/usage/libero.rst @@ -112,7 +112,7 @@ LIBERO-PRO 核心套件一览 排名第一的压缩 PNG mask。 - **toolkit（工具集）** （``robots/libero/toolkit.py``）—— 定义 LLM 能调用的工具：``pi0_pick``（交给 Pi0.5）、``move_to``、``rotate_wrist``、 - ``back_project``、``view_driver_state``、``finish``… + ``back_project``、``view_env_state``、``finish``… Planner 能调用的工具 -------------------- @@ -141,7 +141,7 @@ LIBERO 工具分为物理动作工具和只读工具。 - ``back_project(row, col, ...)`` —— 将图像像素反投影到世界坐标。 - ``segment(prompt=... / point=..., ...)`` —— 通过 SAM3 对已有图像进行文本或 点提示分割。 -- ``view_driver_state(step=-1)`` —— 读取已记录的状态和内嵌观测图像；第 0 步为 +- ``view_env_state(step=-1)`` —— 读取已记录的状态和内嵌观测图像；第 0 步为 初始状态，``-1`` 表示最新状态。 - ``view_camera_meta(camera=..., step=-1)`` —— 读取指定步骤的相机元数据； ``-1`` 表示最新状态。 diff --git a/robots/franka/prompt.py b/robots/franka/prompt.py index edb01f57..3f2ef9c1 100644 --- a/robots/franka/prompt.py +++ b/robots/franka/prompt.py @@ -17,7 +17,7 @@ Robot: Franka arm with Franka Hand. World frame is panda_link0, units are meters. Positive x is front (relative to the robot base), y is left, z is up. Use get_robot_spec for exact workspace bounds and camera names. """, """ - Cameras: scene is the fixed overview RGB-D camera; wrist is the hand-mounted RGB-D camera. view_driver_state and observe return both images. back_project maps a pixel + depth to 3D; it returns robot-base xyz in panda_link0 only when that camera is calibrated for the selected step. + Cameras: scene is the fixed overview RGB-D camera; wrist is the hand-mounted RGB-D camera. view_env_state and observe return both images. back_project maps a pixel + depth to 3D; it returns robot-base xyz in panda_link0 only when that camera is calibrated for the selected step. """, """ Use the wrist camera as the default for back_project because it has been calibrated and sees close-range manipulation targets with better depth accuracy. Use the scene camera for overview/context or if you explicitly need it, and only trust any camera's base-frame xyz if back_project reports calibrated=true. Note that the scene camera is placed opposite the robot and therefore has a mirrored view of the robot and table. @@ -29,13 +29,13 @@ Gripper: use open_gripper and close_gripper for explicit grasp/release, or set gripper to 'open' or 'close' on move_to/move_delta when that is exactly what you want before the move. """, """ - Tools: view_driver_state, observe, back_project, get_camera_meta, get_ee_pose, get_robot_spec, move_to, move_delta, rotate_wrist_yaw, rotate_gripper, open_gripper, close_gripper, finish, plus common file and memory tools. + Tools: view_env_state, observe, back_project, get_camera_meta, get_ee_pose, get_robot_spec, move_to, move_delta, rotate_wrist_yaw, rotate_gripper, open_gripper, close_gripper, finish, plus common file and memory tools. """, ]) RULES = BulletList([ """ - Observe before acting. Call read_memory first, then view_driver_state or observe to see the current setup. + Observe before acting. Call read_memory first, then view_env_state or observe to see the current setup. """, """ Be discreet when moving around. Prefer move_delta for visual servoing and approach/lift motions. @@ -56,7 +56,7 @@ Read memory: call read_memory with no arguments, then read any relevant entry. """, """ - Observe: call view_driver_state or observe and inspect scene plus wrist images and the TCP pose. + Observe: call view_env_state or observe and inspect scene plus wrist images and the TCP pose. """, """ Localize with back_project on the wrist camera first. Use get_camera_meta if you need to check which cameras are calibrated to panda_link0, and use the scene camera mainly for overview/context. diff --git a/robots/franka/toolkit.py b/robots/franka/toolkit.py index 45741645..9d16ddb0 100644 --- a/robots/franka/toolkit.py +++ b/robots/franka/toolkit.py @@ -13,7 +13,7 @@ class FrankaToolkit(Toolkit): """Toolkit for the standalone Franka environment.""" - # view_driver_state image slots: primary scene -> _image_bytes, wrist -> _image_cam_bytes. + # view_env_state image slots: primary scene -> _image_bytes, wrist -> _image_cam_bytes. _VIEW_IMAGE_SLOTS = { "_image_bytes": "scene.png", "_image_cam_bytes": "wrist.png", @@ -38,7 +38,7 @@ def _register_tools(self) -> None: # run's EnvState bound in). Every other spec binds to its primitive- # driver method; @updatestate on the method decides state capture. state_handlers = { - "view_driver_state": partial( + "view_env_state": partial( self._state.view, image_slots=self._VIEW_IMAGE_SLOTS ), "back_project": partial(franka_tools.back_project, state=self._state), @@ -53,7 +53,7 @@ def _register_tools(self) -> None: continue # spec without a backing driver method self.add_tool(name, spec, handler) - def get_state( + def get_env_state( self, *, command: dict[str, Any], diff --git a/robots/franka/tools.py b/robots/franka/tools.py index d379f2a7..c906ae69 100644 --- a/robots/franka/tools.py +++ b/robots/franka/tools.py @@ -217,7 +217,7 @@ def dump_state( return state.get(step_idx) -# view_driver_state now lives on EnvState.view (bound by the toolkit with the +# view_env_state now lives on EnvState.view (bound by the toolkit with the # scene/wrist image slots); nothing module-level is needed here. @@ -294,7 +294,7 @@ def back_project( TOOLS_SPEC: list[dict[str, Any]] = [ { - "name": "view_driver_state", + "name": "view_env_state", "description": ( "Read one recorded state and its observation artifacts. Step -1 " "selects the latest entry. Embeds scene and wrist images." @@ -319,7 +319,7 @@ def back_project( "`xyz` in panda_link0 only when that camera has calibration for " "the selected step; otherwise returns " "xyz_cam plus a warning. Pick row/col on the image returned by " - "view_driver_state." + "view_env_state." ), "input_schema": { "type": "object", diff --git a/robots/lerobot/prompt.py b/robots/lerobot/prompt.py index 3e05c7bf..320ae853 100644 --- a/robots/lerobot/prompt.py +++ b/robots/lerobot/prompt.py @@ -44,7 +44,7 @@ depth — so never back_project it. """, """ - Tools: view_driver_state (state + scene/arm images), get_scene_camera_meta + Tools: view_env_state (state + scene/arm images), get_scene_camera_meta (intrinsics + calibration flag), back_project (scene pixel -> world xyz), get_ee_pose (the fingertip point's xyz in world), move_to, move_joints_delta, finish. Plus read_text_file / write_text_file / list_dir for the scratch @@ -70,7 +70,7 @@ """, """ The scene image is your ground truth: study it to confirm each result (the - view updates after every move_to; call view_driver_state only when you need + view updates after every move_to; call view_env_state only when you need another look). Verify each sub-goal before building on it — that you are positioned correctly before committing an action, and that the action succeeded before the next one; if a precondition isn't met, re-localize or diff --git a/robots/lerobot/toolkit.py b/robots/lerobot/toolkit.py index 75d184b4..4a243c97 100644 --- a/robots/lerobot/toolkit.py +++ b/robots/lerobot/toolkit.py @@ -1,7 +1,7 @@ """LeRobot SO101 toolkit: common tools + SO101 primitives. Inherits the common file/IO tools (including ``finish``) from :class:`Toolkit` -and registers the SO101-specific tools (``view_driver_state``, ``back_project``, +and registers the SO101-specific tools (``view_env_state``, ``back_project``, driver readers, and the move primitives) on top. """ from __future__ import annotations @@ -43,7 +43,7 @@ def _register_tools(self) -> None: # run's EnvState bound in). Every other spec binds to its primitive- # driver method; @updatestate on the method decides state capture. state_handlers = { - "view_driver_state": partial( + "view_env_state": partial( self._state.view, image_slots=self._VIEW_IMAGE_SLOTS ), "back_project": partial(lerobot_tools.back_project, state=self._state), @@ -58,7 +58,7 @@ def _register_tools(self) -> None: continue # spec without a backing driver method self.add_tool(name, spec, handler) - def get_state( + def get_env_state( self, *, command: dict[str, Any], diff --git a/robots/lerobot/tools.py b/robots/lerobot/tools.py index 81a8b227..e1b8294e 100644 --- a/robots/lerobot/tools.py +++ b/robots/lerobot/tools.py @@ -282,7 +282,7 @@ def back_project( TOOLS_SPEC: list[dict[str, Any]] = [ { - "name": "view_driver_state", + "name": "view_env_state", "description": ( "Read one recorded state and its observation artifacts. Step -1 " "selects the latest entry. Embeds scene and arm camera frames." @@ -321,7 +321,7 @@ def back_project( "description": ( "Backproject a SCENE-camera pixel to a 3D point in the WORLD frame " "(arm base_link), using the saved aligned depth. Pick (row, col) on " - "the scene color image from view_driver_state, near the CENTER of " + "the scene color image from view_env_state, near the CENTER of " "the target. It samples a small window around the pixel and returns " "the robust MEDIAN world `xyz` of the object surface (not one noisy " "pixel), plus `n_points` and `xy_spread_m` (a small spread means a " diff --git a/robots/libero/guides/env_calibration.md b/robots/libero/guides/env_calibration.md index 886fd711..c3eb4203 100644 --- a/robots/libero/guides/env_calibration.md +++ b/robots/libero/guides/env_calibration.md @@ -7,7 +7,7 @@ owns the environment server and the `EnvState` lifecycle. Do not issue file-based driver commands, inspect observation storage directly, or read BDDL files for coordinates. -Start with `view_driver_state({"step": 0})`. It returns the initial robot state, +Start with `view_env_state({"step": 0})`. It returns the initial robot state, top-level task language, logical observation references, and embedded camera images. Use `back_project` or `segment` for geometry and `view_camera_meta` for calibration. Step `-1` selects the latest record. @@ -21,7 +21,7 @@ Measured 2026-05-20 on `libero_10_with_mug` t0 (LIVING_ROOM frame) and t8 Each task scene uses one of the table fixtures below, which sets the entire world-frame z origin. The OSC workspace and all pick/place altitudes shift accordingly. **Check `state.robot0_eef_pos[2]` in the result of -`view_driver_state({"step": 0})` and branch on it.** Do not read BDDL files for +`view_env_state({"step": 0})` and branch on it.** Do not read BDDL files for this; use runtime state and visual evidence. | Fixture | eef home z | Table top z | Used by tasks | @@ -150,7 +150,7 @@ limit 1.15). My libero_10 t0 used z=0.95 for travel — safe and consistent. Each calibration motion returns its state, command result, elapsed time, and embedded observation. Use that tool result immediately. To revisit a recorded -step, call `view_driver_state({"step": N})`; use `-1` for the latest step. +step, call `view_env_state({"step": N})`; use `-1` for the latest step. The internal `states.json` file is a versioned manifest owned by `EnvState`, not a list for manual indexing. Calibration analysis should use structured tool @@ -166,5 +166,5 @@ move_to({"xyz": [-0.20, 0.10, 0.65], "gripper": -1, "tol": 0.008, "step_clip": 0.010, "max_steps": 80}) # Inspect the returned result for final_eef_pos and final_dist_m. -# Call view_driver_state({"step": -1}) only when the latest state is needed again. +# Call view_env_state({"step": -1}) only when the latest state is needed again. ``` diff --git a/robots/libero/guides/pro_hybrid_guide.md b/robots/libero/guides/pro_hybrid_guide.md index 7da40264..7b3054cb 100644 --- a/robots/libero/guides/pro_hybrid_guide.md +++ b/robots/libero/guides/pro_hybrid_guide.md @@ -16,7 +16,7 @@ Call: {"step": 0} ``` -through `view_driver_state`. From the returned tool result: +through `view_env_state`. From the returned tool result: - read top-level `task_language` verbatim; - inspect `state.robot0_eef_pos` to identify the scene frame; @@ -64,7 +64,7 @@ The initial end-effector height distinguishes the principal scene frames: | approximately 0.68 m | living-room table | plates, baskets, pudding | | approximately 1.17 m | kitchen table | stove, cabinet, drawer, microwave | -Use `view_driver_state({"step": 0})["state"]["robot0_eef_pos"][2]` as the +Use `view_env_state({"step": 0})["state"]["robot0_eef_pos"][2]` as the measurement. Then use the matching safe-height guidance from [env_calibration.md](./env_calibration.md). @@ -117,7 +117,7 @@ for coordinates. Long-horizon Pro tasks often move or occlude objects during earlier steps. Before each new pick or placement: -1. call `view_driver_state({"step": -1})` if the previous primitive result is +1. call `view_env_state({"step": -1})` if the previous primitive result is no longer in context; 2. inspect the newest embedded images; 3. re-localize any entity that may have moved; @@ -186,7 +186,7 @@ traceability, not manual file access. ## Quick Checklist - Read strict guide and relevant memory. -- Call `view_driver_state({"step": 0})`. +- Call `view_env_state({"step": 0})`. - Select the scene frame from initial EEF z. - Read top-level `task_language`. - Build the complete perception table. diff --git a/robots/libero/guides/strict_hybrid_guide.md b/robots/libero/guides/strict_hybrid_guide.md index 3518a0a7..417b380e 100644 --- a/robots/libero/guides/strict_hybrid_guide.md +++ b/robots/libero/guides/strict_hybrid_guide.md @@ -22,7 +22,7 @@ Every motion primitive appends a `StepRecord`. The record contains: Artifact storage is private to `EnvState`. Do not construct filenames or read observation files manually. Use the structured tools: -- `view_driver_state`: state, artifact names, log, and embedded images; +- `view_env_state`: state, artifact names, log, and embedded images; - `view_camera_meta`: calibration data for one camera and step; - `back_project`: matching image pixel or region to world coordinates; - `segment`: text- or point-prompted mask plus projected `world_xyz`. @@ -35,7 +35,7 @@ API and should not be indexed directly. ## Camera Roles -`view_driver_state` embeds the best available images for the selected step: +`view_env_state` embeds the best available images for the selected step: - **policy image**: the Pi0-oriented agentview input; - **agentview image**: fixed global view, high resolution when available; @@ -53,7 +53,7 @@ specified in the call. The tool selects the matching world map internally. Before the first manipulation primitive: -1. Call `view_driver_state({"step": 0})`. +1. Call `view_env_state({"step": 0})`. 2. Read the returned top-level `task_language` verbatim. 3. Identify every movable target, destination, support, and relation landmark named by the task. diff --git a/robots/libero/prompts/system.py b/robots/libero/prompts/system.py index 70591e30..d93b8df8 100644 --- a/robots/libero/prompts/system.py +++ b/robots/libero/prompts/system.py @@ -141,7 +141,7 @@ - Call the real structured tools exposed by the runtime. - Use bare tool names in this prompt: `move_to`, `pi0_pick`, `release`, `set_gripper`, `rotate_wrist`, `rotate_pitch`, `move_pose`, `pi0_doubled`, - `view_driver_state`, `view_camera_meta`, `back_project`, `segment`, + `view_env_state`, `view_camera_meta`, `back_project`, `segment`, `read_text_file`, `write_text_file`, `list_dir`, `finish`. - Under some runtimes these same tools may appear namespaced; call the actual tool name shown in your tool list, preserving the same arguments and semantics. @@ -151,7 +151,7 @@ world-map, and metadata artifacts, but storage paths are internal to the runtime. Do not construct or read artifact paths manually. -Use `view_driver_state` to retrieve a state. It embeds the policy image and the +Use `view_env_state` to retrieve a state. It embeds the policy image and the best available agentview and wrist images, preferring high resolution. Use `view_camera_meta`, `back_project`, and `segment` to consume metadata, depth, and world maps. These tools guarantee that the selected camera, resolution, @@ -164,7 +164,7 @@ supersedes any reset/retry wording in the Rules below).""" RULES = """Rule 0 — USE IMAGES. After every primitive tool call, inspect the returned state - and embedded images. If you need a state again, call `view_driver_state`. + and embedded images. If you need a state again, call `view_env_state`. Use agentview for global layout and wrist for close-range geometry. The image is your spatial-reasoning input; the JSON state only gives proprioception + object names. @@ -194,7 +194,7 @@ `pi0_pick.success` (eef-lift + gripper-closure heuristic) is a HINT, not proof — always confirm with the wrist cam before carrying. -Rule 2 — Inspect THEN act. Call `view_driver_state({"step": 0})`, inspect the +Rule 2 — Inspect THEN act. Call `view_env_state({"step": 0})`, inspect the returned high-resolution images, and inspect the relevant memory/guides BEFORE your first primitive. **Your task is the returned `task_language`; read it and obey it verbatim.** This is the authoritative instruction (the BDDL @@ -399,7 +399,7 @@ with older/oracle assumptions; do NOT copy coordinates and do NOT replay stale command lists. Re-derive every coordinate from THIS scene. """, - """INSPECT INITIAL STATE: call `view_driver_state({"step": 0})`; inspect + """INSPECT INITIAL STATE: call `view_env_state({"step": 0})`; inspect `task_language`, object_names, eef pose, the returned agentview and wrist images, and call `view_camera_meta` if needed. Identify ALL target objects, destination surfaces, and relation landmarks named by task_language. @@ -488,7 +488,7 @@ OUTPUT_DISCIPLINE = """- Brief reasoning before each tool call (1-2 sentences): observation → decision. - Don't re-read files already in this session. -- Don't call `view_driver_state` immediately after a primitive tool already +- Don't call `view_env_state` immediately after a primitive tool already returned the new state. - Save the audit BEFORE calling `finish`. - Stop immediately after writing the audit and calling `finish`. Do not chat further.""" diff --git a/robots/libero/prompts/user.py b/robots/libero/prompts/user.py index 63d171ac..a71b2444 100644 --- a/robots/libero/prompts/user.py +++ b/robots/libero/prompts/user.py @@ -11,10 +11,10 @@ MODE = """Inspect the embedded high-resolution images returned by -view_driver_state, then use back_project or segment to localize objects before +view_env_state, then use back_project or segment to localize objects before motion.""" BEGIN = """Read MEMORY.md and the guides, then call -`view_driver_state({"step": 0})` and inspect its embedded images. Localize every +`view_env_state({"step": 0})` and inspect its embedded images. Localize every task-relevant entity before planning and execution.""" diff --git a/robots/libero/toolkit.py b/robots/libero/toolkit.py index 95b9253f..6b94cc25 100644 --- a/robots/libero/toolkit.py +++ b/robots/libero/toolkit.py @@ -39,8 +39,8 @@ def _register_libero_tools(self) -> None: # to its primitive-driver method; @updatestate on the method decides # whether state is captured. state_handlers = { - "view_driver_state": partial( - libero_tools.view_driver_state, state=self._state + "view_env_state": partial( + libero_tools.view_env_state, state=self._state ), "view_camera_meta": partial( libero_tools.view_camera_meta, state=self._state @@ -58,7 +58,7 @@ def _register_libero_tools(self) -> None: continue # spec without a backing primitive method self.add_tool(name, spec, handler) - def get_state( + def get_env_state( self, *, command: dict[str, Any], @@ -91,7 +91,7 @@ def get_state( record.step_idx, e, ) - out = libero_tools.view_driver_state(record.step_idx, state=self._state) + out = libero_tools.view_env_state(record.step_idx, state=self._state) out["agent_elapsed_s"] = elapsed_s if result.get("interrupted"): out.update(result) @@ -117,8 +117,8 @@ def init_primitives_clean( libero_tools.dump_state(primitives, self._state, log=None) self._dashboard_events.emit( ToolResultEvent( - name="view_driver_state", - result=libero_tools.view_driver_state(0, state=self._state), + name="view_env_state", + result=libero_tools.view_env_state(0, state=self._state), ) ) diff --git a/robots/libero/tools.py b/robots/libero/tools.py index ff8cb901..a143a898 100644 --- a/robots/libero/tools.py +++ b/robots/libero/tools.py @@ -738,13 +738,13 @@ def _is_primitive_action(name: object) -> bool: """Whether ``name`` is a state-advancing LIBERO primitive. A primitive is any ``@updatestate``-marked method on - :class:`LiberoPrimitives`; read-only tools (``view_driver_state``, + :class:`LiberoPrimitives`; read-only tools (``view_env_state``, ``back_project``, ``segment``, ...) and non-strings read as ``False``. """ if not isinstance(name, str): return False method = getattr(LiberoPrimitives, name, None) - return method is not None and bool(getattr(method, "_captures_state", False)) + return method is not None and bool(getattr(method, "_updates_state", False)) def write_recipe_from_states(state: EnvState, recipe_tag: str) -> str: @@ -1056,7 +1056,7 @@ def _save_observation_artifacts( TOOLS_SPEC = [ { - "name": "view_driver_state", + "name": "view_env_state", "description": ( "Read one recorded state and its observation artifacts. Step -1 " "selects the latest entry. Embeds policy, agentview, and wrist " @@ -1330,7 +1330,7 @@ def _save_observation_artifacts( "selected camera's precomputed world map. Row 0 = top of image, " "col 0 = left. Returns world_xyz in meters.\n\n" "USE THIS to find where an object is in the world — look at " - "the embedded high-resolution image returned by view_driver_state " + "the embedded high-resolution image returned by view_env_state " "to pick a pixel on the target object, then call back_project. " "The default resolution is high (1024x1024). Pass " "resolution='low' only for pixels from the embedded/standard " @@ -1401,7 +1401,7 @@ def _save_observation_artifacts( ] -def view_driver_state(step: int = -1, *, state: EnvState) -> dict: +def view_env_state(step: int = -1, *, state: EnvState) -> dict: try: record = state.get(step) except Exception as exc: diff --git a/rpent/context/prompts/prompt.py b/rpent/context/prompts/prompt.py index b3c9e47b..43e6dad0 100644 --- a/rpent/context/prompts/prompt.py +++ b/rpent/context/prompts/prompt.py @@ -5,7 +5,7 @@ OUTPUT = BulletList([ "Brief reasoning before each tool call (1-2 sentences): observation -> decision.", - "Don't re-read files already in this session. Don't view_driver_state right after a primitive tool already returned the state.", + "Don't re-read files already in this session. Don't view_env_state right after a primitive tool already returned the state.", "Numerical coords in 3 decimals are enough.", "Save artifacts BEFORE calling finish. Stop immediately after writing the audit; do not chat further.", ]) diff --git a/rpent/planner/api_loop.py b/rpent/planner/api_loop.py index 1e20d11b..6470dadf 100644 --- a/rpent/planner/api_loop.py +++ b/rpent/planner/api_loop.py @@ -410,7 +410,7 @@ def _build_tools(toolkit: Toolkit, *, no_images: bool = False) -> list[Tool]: def read_image(path: str) -> ToolReturn: """Read an explicitly provided local image file as visual input. - Environment observations are already embedded by ``view_driver_state``; + Environment observations are already embedded by ``view_env_state``; this helper is only for other user-selected local files. """ return ToolReturn( @@ -423,7 +423,7 @@ def read_image_text_only(path: str) -> str: """``read_image`` stub for ``--no-images``: acknowledge, send no bytes.""" return ( f"{path} exists, but image input is disabled (--no-images, text-only " - "model). Reason from textual state instead: view_driver_state, " + "model). Reason from textual state instead: view_env_state, " "back_project, and the numeric fields in tool results." ) diff --git a/rpent/tools/toolkit.py b/rpent/tools/toolkit.py index 22d30871..9c424020 100644 --- a/rpent/tools/toolkit.py +++ b/rpent/tools/toolkit.py @@ -32,10 +32,10 @@ class ToolCancelled(Exception): def updatestate(func): """Mark a tool handler as one that advances environment state. - The toolkit captures a fresh observation (:meth:`Toolkit.get_state`) + The toolkit captures a fresh observation (:meth:`Toolkit.get_env_state`) after a handler carrying this marker runs (or raises). Apply it to primitive-driver methods that move the robot; read-only tools (file IO, - ``view_driver_state``, ``back_project``, ``segment``, ...) are left + ``view_env_state``, ``back_project``, ``segment``, ...) are left unmarked and skip state capture. The marker is read off the underlying function, so registering a bound @@ -44,11 +44,11 @@ def updatestate(func): wrapping a marked method is treated as read-only (intentional -- see the LIBERO ``segment`` tool). """ - func._captures_state = True + func._updates_state = True return func -def _captures_state(handler: Callable[..., Any]) -> bool: +def _updates_state(handler: Callable[..., Any]) -> bool: """Whether ``handler`` was marked with :func:`updatestate`. Resolves through ``__func__`` so bound methods (the common case for @@ -57,7 +57,7 @@ def _captures_state(handler: Callable[..., Any]) -> bool: do not delegate, so undecorated read-only handlers read as ``False``. """ target = getattr(handler, "__func__", handler) - return bool(getattr(target, "_captures_state", False)) + return bool(getattr(target, "_updates_state", False)) @dataclass @@ -174,7 +174,7 @@ def add_tool( a result dict. Decorate state-advancing primitives with :func:`updatestate`. """ - self._tools[name] = (spec, handler, _captures_state(handler)) + self._tools[name] = (spec, handler, _updates_state(handler)) def _register_common_tools(self) -> None: """Register the file/IO tools shared by every run.""" @@ -236,7 +236,7 @@ def execute_tool(self, name: str, input_dict: dict[str, Any]) -> ToolResult: result_dict = result if isinstance(result, dict) else {"value": result} command = {"action": name, **input_dict} try: - captured = self.get_state( + captured = self.get_env_state( command=command, result=result_dict, elapsed_s=elapsed_s, @@ -264,7 +264,7 @@ def _finish_tool(self, name: str, result: Any) -> ToolResult: self._dashboard_events.emit(ToolResultEvent(name=name, result=result)) return ToolResult(name=name, result=result) - def get_state( + def get_env_state( self, *, command: dict[str, Any], diff --git a/scripts/codex_proxy/litellm_callbacks.py b/scripts/codex_proxy/litellm_callbacks.py index 23aca5f2..5d54e05b 100644 --- a/scripts/codex_proxy/litellm_callbacks.py +++ b/scripts/codex_proxy/litellm_callbacks.py @@ -92,7 +92,7 @@ def _extract_mcp_namespace_map(data: dict) -> dict[str, str]: Example:: {"back_project": "mcp__rpent", - "view_driver_state": "mcp__rpent"} + "view_env_state": "mcp__rpent"} """ tools = data.get("tools") if not isinstance(tools, list): From 5790fc38f77b485ab761c061192fba2f3c86fa07 Mon Sep 17 00:00:00 2001 From: Jiaxing Qiu Date: Thu, 6 Aug 2026 03:12:30 +0000 Subject: [PATCH 22/22] refactor(toolkit): publish dashboard steps as StepRecords, introduce _FRAME_ARTIFACTS Signed-off-by: Jiaxing Qiu --- robots/franka/__init__.py | 5 +-- robots/franka/toolkit.py | 13 ++++++-- robots/lerobot/__init__.py | 5 +-- robots/lerobot/toolkit.py | 14 +++++--- robots/libero/toolkit.py | 25 ++++++-------- robots/libero/tools.py | 4 +-- rpent/cli/dashboard.py | 1 - rpent/dashboard/events.py | 10 ++++++ rpent/dashboard/state.py | 67 +++++++++++++++++++++++++++++++++++--- rpent/tools/state.py | 4 +++ rpent/tools/toolkit.py | 23 +++++++++++-- 11 files changed, 134 insertions(+), 37 deletions(-) diff --git a/robots/franka/__init__.py b/robots/franka/__init__.py index ec67a0b7..8d3a403b 100644 --- a/robots/franka/__init__.py +++ b/robots/franka/__init__.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any from robots.franka.prompt import system_prompt, user_prompt +from rpent.dashboard.events import DashboardEventSink from rpent.envs.env_spec import EnvSpec, RunConfig from rpent.envs.prompt_bundle import PromptBundle from rpent.utils.config import get_repo_root @@ -42,7 +43,7 @@ def get_env_spec() -> EnvSpec: def get_toolkit( *, primitives_kwargs: dict[str, Any], - dashboard: Any = None, + dashboard_events: DashboardEventSink, ): """Return the Franka toolkit (common tools + Cartesian primitives). @@ -52,7 +53,7 @@ def get_toolkit( from robots.franka.toolkit import FrankaToolkit return FrankaToolkit( - dashboard=dashboard, + dashboard_events=dashboard_events, **primitives_kwargs, ) diff --git a/robots/franka/toolkit.py b/robots/franka/toolkit.py index 9d16ddb0..2080f768 100644 --- a/robots/franka/toolkit.py +++ b/robots/franka/toolkit.py @@ -5,6 +5,7 @@ from typing import Any from robots.franka import tools as franka_tools +from rpent.dashboard.events import DashboardEventSink from rpent.tools.state import EnvState from rpent.tools.toolkit import Toolkit from rpent.utils.logging import get_output_dir @@ -18,18 +19,23 @@ class FrankaToolkit(Toolkit): "_image_bytes": "scene.png", "_image_cam_bytes": "wrist.png", } + # Per-env artifact names for the dashboard's live frame images. + _FRAME_ARTIFACTS = { + "camera": "scene.png", + "wrist": "wrist.png", + } def __init__( self, *, env: Any, - dashboard: Any = None, + dashboard_events: DashboardEventSink, ) -> None: # EnvState owns the trace + counter for this run (explicit output_dir, # no process-global). The runner will own its lifecycle in a later cut; # for now the toolkit constructs it from get_output_dir(). state = EnvState(get_output_dir()) - super().__init__(dashboard=dashboard, state=state) + super().__init__(dashboard_events=dashboard_events, state=state) self.init_driver_clean(env=env) self._register_tools() @@ -74,8 +80,9 @@ def init_driver_clean(self, *, env: Any) -> None: self._state.reset() driver = franka_tools.FrankaPrimitives(env=env) driver.reset() - franka_tools.dump_state(driver, self._state, log=None) + record = franka_tools.dump_state(driver, self._state, log=None) self._driver = driver + self._publish_step(record) def close(self) -> None: return None diff --git a/robots/lerobot/__init__.py b/robots/lerobot/__init__.py index 62c3b574..d10c95ea 100644 --- a/robots/lerobot/__init__.py +++ b/robots/lerobot/__init__.py @@ -18,6 +18,7 @@ system_prompt, user_prompt, ) +from rpent.dashboard.events import DashboardEventSink from rpent.envs.env_spec import EnvSpec, RunConfig from rpent.envs.prompt_bundle import PromptBundle from rpent.utils.config import get_repo_root @@ -53,7 +54,7 @@ def get_env_spec() -> EnvSpec: def get_toolkit( *, primitives_kwargs: dict[str, Any], - dashboard: Any = None, + dashboard_events: DashboardEventSink, ): """Return the SO101 toolkit (common tools + SO101 primitives). @@ -64,7 +65,7 @@ def get_toolkit( from robots.lerobot.toolkit import LerobotToolkit return LerobotToolkit( - dashboard=dashboard, + dashboard_events=dashboard_events, **primitives_kwargs, ) diff --git a/robots/lerobot/toolkit.py b/robots/lerobot/toolkit.py index 4a243c97..30024f66 100644 --- a/robots/lerobot/toolkit.py +++ b/robots/lerobot/toolkit.py @@ -10,6 +10,7 @@ from typing import Any from robots.lerobot import tools as lerobot_tools +from rpent.dashboard.events import DashboardEventSink from rpent.tools.state import EnvState from rpent.tools.toolkit import Toolkit from rpent.utils.logging import get_output_dir @@ -22,16 +23,21 @@ class LerobotToolkit(Toolkit): "_image_bytes": "scene.png", "_image_cam_bytes": "arm.png", } + # Per-env artifact names for the dashboard's live frame images. + _FRAME_ARTIFACTS = { + "camera": "scene.png", + "wrist": "arm.png", + } def __init__( self, *, env: Any, model: Any | None = None, - dashboard: Any = None, + dashboard_events: DashboardEventSink, ) -> None: state = EnvState(get_output_dir()) - super().__init__(dashboard=dashboard, state=state) + super().__init__(dashboard_events=dashboard_events, state=state) self.init_driver_clean(env=env, model=model) self._register_tools() @@ -80,9 +86,9 @@ def init_driver_clean(self, *, env: Any, model: Any | None = None) -> None: self._state.reset() driver = lerobot_tools.LerobotPrimitives(env=env, model=model) driver.reset() - lerobot_tools.dump_state(driver, self._state, log=None) - + record = lerobot_tools.dump_state(driver, self._state, log=None) self._driver = driver + self._publish_step(record) def close(self) -> None: """End-of-run cleanup hook. TODO: flush an episode video if desired.""" diff --git a/robots/libero/toolkit.py b/robots/libero/toolkit.py index 6b94cc25..23a315cb 100644 --- a/robots/libero/toolkit.py +++ b/robots/libero/toolkit.py @@ -9,7 +9,7 @@ from typing import Any from robots.libero import tools as libero_tools -from rpent.dashboard.events import DashboardEventSink, ToolResultEvent +from rpent.dashboard.events import DashboardEventSink from rpent.tools.state import EnvState from rpent.tools.toolkit import Toolkit from rpent.utils.logging import get_logger, get_output_dir @@ -18,6 +18,11 @@ class LiberoToolkit(Toolkit): """Toolkit for the LIBERO environment.""" + _FRAME_ARTIFACTS = { + "camera": "agentview.png", + "wrist": "wrist.png", + } + def __init__( self, *, @@ -72,19 +77,17 @@ def get_env_state( self._state, log={"command": command, "result": result, "elapsed_s": elapsed_s}, ) - action_video_name = None if self._dashboard_events.enabled: try: frames = self._primitives.frame_slice(frame_start) if frames: candidate = f"action_{command['action']}.mp4" - if self._state.save( + self._state.save( candidate, frames, step=record.step_idx, fps=20, - ): - action_video_name = candidate + ) except Exception as e: get_logger("libero_toolkit").warning( "failed to save action clip for step %s: %s", @@ -95,8 +98,6 @@ def get_env_state( out["agent_elapsed_s"] = elapsed_s if result.get("interrupted"): out.update(result) - if action_video_name is not None: - out["action_video_artifact"] = action_video_name return out def init_primitives_clean( @@ -114,15 +115,9 @@ def init_primitives_clean( primitives.reset() primitives.start_recording() self._action_frame_cursor = primitives.recorded_frame_count() - libero_tools.dump_state(primitives, self._state, log=None) - self._dashboard_events.emit( - ToolResultEvent( - name="view_env_state", - result=libero_tools.view_env_state(0, state=self._state), - ) - ) - + record = libero_tools.dump_state(primitives, self._state, log=None) self._primitives = primitives + self._publish_step(record) def close(self) -> None: """Flush the agent-side video buffer through ``EnvState``.""" diff --git a/robots/libero/tools.py b/robots/libero/tools.py index a143a898..2f66eb76 100644 --- a/robots/libero/tools.py +++ b/robots/libero/tools.py @@ -844,7 +844,7 @@ def dump_state( result=log.get("result"), elapsed_s=log.get("elapsed_s"), extras={ - "libero_terminated": primitives.env.episode_terminated, + "terminated": primitives.env.episode_terminated, "episode_truncated": primitives.env.episode_truncated, "task_language": primitives.env.get_task_language(), }, @@ -1415,7 +1415,7 @@ def view_env_state(step: int = -1, *, state: EnvState) -> dict: "artifacts": sorted(record.artifacts), } out["task_language"] = extras.get("task_language") - out["libero_terminated"] = extras.get("libero_terminated") + out["libero_terminated"] = extras.get("terminated") out["episode_truncated"] = extras.get("episode_truncated") out["log"] = { "command": record.command, diff --git a/rpent/cli/dashboard.py b/rpent/cli/dashboard.py index a367820b..3ae39131 100644 --- a/rpent/cli/dashboard.py +++ b/rpent/cli/dashboard.py @@ -150,7 +150,6 @@ def _run_dashboard_task( toolkit = get_toolkit( args.env_name, primitives_kwargs=primitives_kwargs, - video_path=str(output_dir / "episode.mp4"), dashboard_events=state, ) planner = build_planner( diff --git a/rpent/dashboard/events.py b/rpent/dashboard/events.py index 5784624d..0c6cbfae 100644 --- a/rpent/dashboard/events.py +++ b/rpent/dashboard/events.py @@ -39,6 +39,15 @@ class ToolResultEvent: result: Any +@dataclass(frozen=True, slots=True) +class StepRecordEvent: + """Publish one recorded environment step and its artifact context.""" + + record: Any + env_state: Any + frame_artifacts: dict[str, str] + + @dataclass(frozen=True, slots=True) class RunStartedEvent: """Mark startup complete and the agent run active.""" @@ -49,6 +58,7 @@ class RunStartedEvent: | UsageEvent | RuntimeStatusEvent | ToolResultEvent + | StepRecordEvent | RunStartedEvent ) diff --git a/rpent/dashboard/state.py b/rpent/dashboard/state.py index ab08b087..0a1179b8 100644 --- a/rpent/dashboard/state.py +++ b/rpent/dashboard/state.py @@ -11,6 +11,7 @@ DashboardEvent, RunStartedEvent, RuntimeStatusEvent, + StepRecordEvent, ToolResultEvent, TranscriptEvent, UsageEvent, @@ -26,6 +27,7 @@ if TYPE_CHECKING: from rpent.dashboard.commands import TaskCommand + from rpent.tools.state import EnvState, StepRecord RUNTIME_COMPONENTS = ("env", "vla", "sam3") RUNTIME_STATUSES = {"pending", "starting", "ready", "failed"} @@ -78,6 +80,8 @@ def __init__( self._timeline: list[dict[str, Any]] = [] self._frames: dict[str, bytes] = {} self._frame_idx = -1 + self.env_state: EnvState | None = None + self.frame_artifacts: dict[str, str] = {} self._accepting_input = False self._planner_activity: PlannerActivity = "starting" self._interrupt_requested = False @@ -268,6 +272,8 @@ def _begin_task_locked( self._timeline = [] self._frames = {} self._frame_idx = -1 + self.env_state = None + self.frame_artifacts = {} self._accepting_input = False self._planner_activity = "starting" self._interrupt_requested = False @@ -477,6 +483,11 @@ def emit(self, event: DashboardEvent) -> None: if isinstance(event, ToolResultEvent): self._apply_tool_result(event) return + if isinstance(event, StepRecordEvent): + self.env_state = event.env_state + self.frame_artifacts = dict(event.frame_artifacts) + self.on_step(event.record) + return if isinstance(event, RunStartedEvent): self._start() return @@ -535,6 +546,46 @@ def _apply_tool_result(self, event: ToolResultEvent) -> None: self._timeline.append(item) self._terminated = self._terminated or terminated + def on_step(self, record: StepRecord) -> None: + """Project one recorded environment step into frames and timeline.""" + self._update_step_frames(record) + command = record.command + if not isinstance(command, dict) or not command.get("action"): + return + terminated = bool(record.extras.get("terminated")) + action_video = next( + (name for name in sorted(record.artifacts) if name.endswith(".mp4")), + None, + ) + item = { + "step": record.step_idx, + "action": str(command.get("action")), + "args": {key: value for key, value in command.items() if key != "action"}, + "result": record.result, + "elapsed_s": record.elapsed_s, + "terminated": terminated, + "action_video_artifact": action_video, + "has_action_video": action_video is not None, + } + with self._lock: + self._timeline.append(item) + self._terminated = self._terminated or terminated + + def _update_step_frames(self, record: StepRecord) -> None: + """Load dashboard frame bytes from the step's canonical artifacts.""" + env_state = self.env_state + if env_state is None: + return + frames: dict[str, bytes] = {} + for kind, artifact in self.frame_artifacts.items(): + if kind not in FRAME_KINDS or artifact not in record.artifacts: + continue + try: + frames[kind] = env_state.load_bytes(artifact, step=record.step_idx) + except FileNotFoundError: + continue + self._update_frames(step=record.step_idx, frames=frames) + def _apply_frame_paths(self, result: dict[str, Any]) -> None: path_keys = { "camera": "image_cam_path", @@ -680,16 +731,22 @@ def frame(self, kind: str) -> bytes | None: return self._frames.get(kind) def action_video(self, step: int) -> bytes | None: + env_state = self.env_state + if env_state is None: + return None with self._lock: + artifact = None for item in self._timeline: if int(item.get("step", -1)) != int(step): continue artifact = item.get("action_video_artifact") - if not artifact: - return None - video_path = self.output_dir / f"{int(step):02d}_{artifact}" - return video_path.read_bytes() if video_path.exists() else None - return None + break + if not artifact: + return None + try: + return env_state.load_bytes(artifact, step=int(step)) + except FileNotFoundError: + return None def video(self) -> bytes | None: return self.video_path.read_bytes() if self.video_path.exists() else None diff --git a/rpent/tools/state.py b/rpent/tools/state.py index 3d43677b..a059cb77 100644 --- a/rpent/tools/state.py +++ b/rpent/tools/state.py @@ -175,6 +175,10 @@ def latest_step(self) -> int | None: return None return self._steps[-1].step_idx + def latest_record(self) -> StepRecord | None: + """Return the most recently recorded step (live reference, no copy).""" + return self._steps[-1] if self._steps else None + def _resolve_read_step(self, step: int | None) -> int | None: if step is None: return None diff --git a/rpent/tools/toolkit.py b/rpent/tools/toolkit.py index 9c424020..76882f81 100644 --- a/rpent/tools/toolkit.py +++ b/rpent/tools/toolkit.py @@ -13,11 +13,14 @@ import traceback from collections.abc import Callable from dataclasses import dataclass, field -from typing import Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar -from rpent.dashboard.events import DashboardEventSink, ToolResultEvent +from rpent.dashboard.events import DashboardEventSink, StepRecordEvent from rpent.utils.templates import substitute +if TYPE_CHECKING: + from rpent.tools.state import StepRecord + @dataclass(slots=True) class _ToolOperation: @@ -235,6 +238,7 @@ def execute_tool(self, name: str, input_dict: dict[str, Any]) -> ToolResult: elapsed_s = round(time.perf_counter() - started, 2) result_dict = result if isinstance(result, dict) else {"value": result} command = {"action": name, **input_dict} + record: StepRecord | None = None try: captured = self.get_env_state( command=command, @@ -248,11 +252,15 @@ def execute_tool(self, name: str, input_dict: dict[str, Any]) -> ToolResult: "error", f"failed to capture state after {name}: {e}" ) captured.setdefault("traceback", traceback.format_exc()) + else: + record = self._state.latest_record() result = captured if failed: result.setdefault("error", result_dict["error"]) if "traceback" in result_dict: result.setdefault("traceback", result_dict["traceback"]) + if record is not None: + self._publish_step(record) return self._finish_tool(name, result) finally: @@ -260,8 +268,17 @@ def execute_tool(self, name: str, input_dict: dict[str, Any]) -> ToolResult: self._active_operation = None operation.done_event.set() + def _publish_step(self, record: StepRecord) -> None: + """Publish one recorded environment step to the dashboard sink.""" + self._dashboard_events.emit( + StepRecordEvent( + record=record, + env_state=self._state, + frame_artifacts=dict(getattr(type(self), "_FRAME_ARTIFACTS", {})), + ) + ) + def _finish_tool(self, name: str, result: Any) -> ToolResult: - self._dashboard_events.emit(ToolResultEvent(name=name, result=result)) return ToolResult(name=name, result=result) def get_env_state(