#!/usr/bin/env python3 """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_ \\ --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 # noqa: E402 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: # order matches baseline/preprocess.py's COMPACT_STATE_NAMES exactly: # left arm x7, left gripper, right arm x7, right gripper. 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] # same clamp-to-[0,1] convention as scripts/adapters/gello_to_bridge.py's # make_gripper_command - the action feature is literally named # "*_gripper_target_width_percent_value", the same quantity GELLO's # target_gripper_width_percent topic carries. 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 # see run_policy.py's note: required or weights stay random-init if args.bf16_vision: # Build on CPU so the float32 vision tower never lands on the GPU, then # cast only it: casting the whole policy breaks pi0.5's action # projection (float32 noise into a bfloat16 Linear). 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 = { # is_batch=False: this imageio/av version's append_data() otherwise # treats a single (H,W,3) frame's leading axis as a batch of frames - # verified empirically, produces a "could not broadcast" shape error # on the very first frame without this. 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())