| |
| """Drive the Newton/MJWarp bridge from a trained pi0.5 checkpoint instead of |
| a human teleop device - the ROS-topic equivalent of |
| task1_mujoco/.../teleop/run_policy.py's LeRobot-checkpoint branch, adapted |
| to this task's `/isaac/*` (state, camera) / `/bridge/*` (commands) topics. |
| |
| Requires the bridge launched with --enable_cameras (see |
| isaaclab_fr3duo_newton_bridge.py) so the three /isaac/*_camera/image_raw |
| topics this node subscribes to actually exist. |
| |
| Known gap, stated plainly rather than silently ignored: the trained |
| policy's 23-dim action also includes 6-dof base cmd_vel and a spine-height |
| target (see baseline/preprocess.py's dataset action feature names), but |
| there is no ROS command topic in this bridge for either yet (base motion |
| only accepts discrete pedal tokens, spine height only accepts local |
| keyboard control) - this adapter computes those values but only actuates |
| the two 7-dof arms and two grippers. |
| |
| python scripts/adapters/policy_adapter.py \\ |
| --policy-path baseline/ckpt/pi05_lora_20260917_1702/checkpoints/240000/pretrained_model \\ |
| --dataset-root ../data/pi05 \\ |
| --video-dir baseline/results/eval/rollout_<ts> \\ |
| --duration-sec 60 |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import rclpy |
| import torch |
| from rclpy.node import Node |
| from sensor_msgs.msg import Image, JointState |
|
|
| BASELINE_DIR = Path(__file__).resolve().parents[2] / "baseline" |
| sys.path.insert(0, str(BASELINE_DIR)) |
| from preprocess import COMPACT_STATE_NAMES, fit_resolution, pad_to_size |
|
|
| LEFT_ARM_JOINTS = [f"left_fr3v2_joint{i}" for i in range(1, 8)] |
| RIGHT_ARM_JOINTS = [f"right_fr3v2_joint{i}" for i in range(1, 8)] |
| LEFT_GRIPPER_JOINT_NAME = "left_robotiq_opening" |
| RIGHT_GRIPPER_JOINT_NAME = "right_robotiq_opening" |
| IMAGE_TOPICS = { |
| "head": "/isaac/head_camera/image_raw", |
| "wrist_left": "/isaac/left_wrist_camera/image_raw", |
| "wrist_right": "/isaac/right_wrist_camera/image_raw", |
| } |
| assert len(COMPACT_STATE_NAMES) == 16 |
|
|
|
|
| def _resize_with_pad_uint8(frame: np.ndarray, target: tuple[int, int]) -> np.ndarray: |
| """(H,W,3) uint8 -> resize preserving aspect ratio to fit inside `target` |
| (W,H), then center-pad with black - identical algorithm to |
| preprocess.py's fit_resolution+_resize_to_uint8+pad_to_size sequence |
| used to derive the training images, just operating on a live uint8 |
| array instead of a decoded video tensor.""" |
| h, w = frame.shape[:2] |
| inner_w, inner_h = fit_resolution((w, h), target) |
| if (inner_w, inner_h) != (w, h): |
| tensor = torch.from_numpy(frame).permute(2, 0, 1).float() / 255.0 |
| resized = torch.nn.functional.interpolate( |
| tensor.unsqueeze(0), size=(inner_h, inner_w), mode="bilinear", align_corners=False, antialias=True |
| )[0] |
| frame = (resized.clamp(0.0, 1.0) * 255.0).round().to(torch.uint8).permute(1, 2, 0).numpy() |
| frame = np.ascontiguousarray(frame) |
| return pad_to_size(frame, target) |
|
|
|
|
| class PolicyAdapterNode(Node): |
| def __init__(self): |
| super().__init__("policy_adapter") |
| self._left_arm_pos: list[float] | None = None |
| self._right_arm_pos: list[float] | None = None |
| self._left_gripper_pos: float | None = None |
| self._right_gripper_pos: float | None = None |
| self._images: dict[str, np.ndarray] = {} |
|
|
| self.create_subscription(JointState, "/isaac/left_joint_states", self._on_left_arm, 10) |
| self.create_subscription(JointState, "/isaac/right_joint_states", self._on_right_arm, 10) |
| self.create_subscription(JointState, "/isaac/left_robotiq_joint_states", self._on_left_gripper, 10) |
| self.create_subscription(JointState, "/isaac/right_robotiq_joint_states", self._on_right_gripper, 10) |
| for key, topic in IMAGE_TOPICS.items(): |
| self.create_subscription(Image, topic, lambda msg, key=key: self._on_image(key, msg), 10) |
|
|
| self._left_arm_pub = self.create_publisher(JointState, "/bridge/left_joint_commands", 10) |
| self._right_arm_pub = self.create_publisher(JointState, "/bridge/right_joint_commands", 10) |
| self._left_gripper_pub = self.create_publisher(JointState, "/bridge/left_robotiq_joint_commands", 10) |
| self._right_gripper_pub = self.create_publisher(JointState, "/bridge/right_robotiq_joint_commands", 10) |
| self.get_logger().info("policy_adapter listening on /isaac state+camera topics") |
|
|
| def _on_left_arm(self, msg: JointState) -> None: |
| self._left_arm_pos = list(msg.position[:7]) |
|
|
| def _on_right_arm(self, msg: JointState) -> None: |
| self._right_arm_pos = list(msg.position[:7]) |
|
|
| def _on_left_gripper(self, msg: JointState) -> None: |
| if msg.position: |
| self._left_gripper_pos = float(msg.position[0]) |
|
|
| def _on_right_gripper(self, msg: JointState) -> None: |
| if msg.position: |
| self._right_gripper_pos = float(msg.position[0]) |
|
|
| def _on_image(self, key: str, msg: Image) -> None: |
| arr = np.frombuffer(bytes(msg.data), dtype=np.uint8).reshape(msg.height, msg.width, 3) |
| self._images[key] = arr.copy() |
|
|
| def ready(self) -> bool: |
| return ( |
| self._left_arm_pos is not None |
| and self._right_arm_pos is not None |
| and self._left_gripper_pos is not None |
| and self._right_gripper_pos is not None |
| and all(key in self._images for key in IMAGE_TOPICS) |
| ) |
|
|
| def build_state(self) -> np.ndarray: |
| |
| |
| return np.array( |
| self._left_arm_pos + [self._left_gripper_pos] + self._right_arm_pos + [self._right_gripper_pos], |
| dtype=np.float32, |
| ) |
|
|
| def latest_images(self) -> dict[str, np.ndarray]: |
| return dict(self._images) |
|
|
| def publish_action(self, action: np.ndarray) -> None: |
| stamp = self.get_clock().now().to_msg() |
|
|
| def arm_msg(names: list[str], positions: list[float]) -> JointState: |
| msg = JointState() |
| msg.header.stamp = stamp |
| msg.name = names |
| msg.position = positions |
| return msg |
|
|
| def gripper_msg(name: str, width_percent: float) -> JointState: |
| msg = JointState() |
| msg.header.stamp = stamp |
| msg.name = [name] |
| |
| |
| |
| |
| msg.position = [max(0.0, min(1.0, float(width_percent)))] |
| return msg |
|
|
| self._left_arm_pub.publish(arm_msg(LEFT_ARM_JOINTS, [float(v) for v in action[0:7]])) |
| self._left_gripper_pub.publish(gripper_msg(LEFT_GRIPPER_JOINT_NAME, action[7])) |
| self._right_arm_pub.publish(arm_msg(RIGHT_ARM_JOINTS, [float(v) for v in action[8:15]])) |
| self._right_gripper_pub.publish(gripper_msg(RIGHT_GRIPPER_JOINT_NAME, action[15])) |
|
|
|
|
| def build_parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument("--policy-path", required=True, help="checkpoint dir, e.g. .../checkpoints/240000/pretrained_model") |
| parser.add_argument("--dataset-repo-id", default="ebim-benchmark/ebim-task1-realrobotdata-lerobot-part1") |
| parser.add_argument("--dataset-root", required=True, help="local dataset root used to read fps/image size") |
| parser.add_argument("--task", default="cable routing", help="language task string the model was trained with") |
| parser.add_argument("--control-hz", type=float, default=0.0, help="0 = use the dataset's own fps") |
| parser.add_argument("--image-size", default="224x224") |
| parser.add_argument("--duration-sec", type=float, default=60.0) |
| parser.add_argument("--video-dir", required=True, help="directory to write one mp4 per camera into") |
| parser.add_argument("--video-fps", type=float, default=0.0, help="0 = same as --control-hz") |
| parser.add_argument( |
| "--bf16-vision", |
| action=argparse.BooleanOptionalAction, |
| default=True, |
| help="build the policy on CPU and cast the vision encoder (kept float32 by lerobot's pi0.5) to bfloat16 " |
| "before moving to the GPU - ~0.8 GB less VRAM (pi0.5 peaks ~8.9 GB instead of ~9.7 GB), which is what " |
| "lets it share a 16 GB GPU with the bridge + cable world", |
| ) |
| return parser |
|
|
|
|
| def main() -> int: |
| args = build_parser().parse_args() |
|
|
| import imageio |
| from lerobot.common.control_utils import predict_action |
| from lerobot.configs.policies import PreTrainedConfig |
| from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata |
| from lerobot.policies.factory import make_policy, make_pre_post_processors |
| from lerobot.utils.device_utils import get_safe_torch_device |
|
|
| w, h = (int(v) for v in args.image_size.lower().split("x")) |
| image_size = (w, h) |
|
|
| print(f"[policy_adapter] loading checkpoint from {args.policy_path}") |
| ds_meta = LeRobotDatasetMetadata(repo_id=args.dataset_repo_id, root=args.dataset_root) |
| policy_cfg = PreTrainedConfig.from_pretrained(args.policy_path) |
| policy_cfg.pretrained_path = args.policy_path |
| if args.bf16_vision: |
| |
| |
| |
| target_device = policy_cfg.device |
| policy_cfg.device = "cpu" |
| policy = make_policy(policy_cfg, ds_meta=ds_meta) |
| for name, module in policy.named_modules(): |
| if name.endswith("vision_tower") or name.endswith("multi_modal_projector"): |
| module.to(dtype=torch.bfloat16) |
| policy_cfg.device = target_device |
| policy.to(target_device) |
| else: |
| policy = make_policy(policy_cfg, ds_meta=ds_meta) |
| preprocessor, postprocessor = make_pre_post_processors(policy_cfg, pretrained_path=args.policy_path) |
| device = get_safe_torch_device(policy_cfg.device) |
| policy.reset() |
| control_hz = args.control_hz or float(ds_meta.fps) |
| video_fps = args.video_fps or control_hz |
| print(f"[policy_adapter] {policy_cfg.type} loaded on {device}, control_hz={control_hz:g}, task={args.task!r}") |
|
|
| video_dir = Path(args.video_dir) |
| video_dir.mkdir(parents=True, exist_ok=True) |
| writers = { |
| |
| |
| |
| |
| key: imageio.get_writer(str(video_dir / f"{key}.mp4"), fps=video_fps, codec="libx264", is_batch=False) |
| for key in IMAGE_TOPICS |
| } |
|
|
| rclpy.init() |
| node = PolicyAdapterNode() |
|
|
| print("[policy_adapter] waiting for first state+image messages...") |
| wait_start = time.monotonic() |
| while not node.ready(): |
| rclpy.spin_once(node, timeout_sec=0.1) |
| if time.monotonic() - wait_start > 30.0: |
| node.destroy_node() |
| rclpy.shutdown() |
| raise SystemExit("Timed out waiting for /isaac state+camera topics - is the bridge running with --enable_cameras?") |
| print("[policy_adapter] first observation received, starting rollout") |
|
|
| control_dt = 1.0 / control_hz |
| steps = int(round(args.duration_sec * control_hz)) |
| action = None |
| try: |
| for step in range(steps): |
| tick_start = time.monotonic() |
| rclpy.spin_once(node, timeout_sec=0.0) |
|
|
| state = node.build_state() |
| raw_images = node.latest_images() |
| resized = {key: _resize_with_pad_uint8(img, image_size) for key, img in raw_images.items()} |
| for key, frame in resized.items(): |
| writers[key].append_data(frame) |
|
|
| obs = {"observation.state": state} |
| obs.update({f"observation.images.{key}": frame for key, frame in resized.items()}) |
| action = predict_action(obs, policy, device, preprocessor, postprocessor, use_amp=False, task=args.task) |
| action = action.squeeze(0).numpy() |
| node.publish_action(action) |
|
|
| if step % int(control_hz) == 0: |
| print(f"[policy_adapter] step {step}/{steps} gripper_l={action[7]:.3f} gripper_r={action[15]:.3f}") |
|
|
| sleep = control_dt - (time.monotonic() - tick_start) |
| if sleep > 0: |
| time.sleep(sleep) |
| finally: |
| for writer in writers.values(): |
| writer.close() |
| node.destroy_node() |
| rclpy.shutdown() |
| print(f"[policy_adapter] rollout finished, videos written to {video_dir}") |
|
|
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
