Add ROS2 node for diff-drive robot deployment

vla_micro_ros2_node.py

@@ -0,0 +1,193 @@
+#!/usr/bin/env python3
+"""
+OpenVLA-Micro ROS2 node for diff-drive robot (Nova Carter).
+
+Subscribes:
+  /rgb/image_raw       — Kinect RGB
+  /depth/image_raw     — Kinect depth (obstacle avoidance)
+  /vla_instruction     — text command
+
+Publishes:
+  /cmd_vel             — Twist
+
+Run:
+  python3 vla_micro_ros2_node.py
+
+First-time setup will load & cache model components (~30s load, once).
+Subsequent starts load in ~3-5s.
+"""
+import os, time, threading
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+import numpy as np
+from PIL import Image
+import rclpy
+from rclpy.node import Node
+from sensor_msgs.msg import Image as RosImage
+from geometry_msgs.msg import Twist
+from std_msgs.msg import String
+from cv_bridge import CvBridge
+import torch
+
+CACHE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "model_cache")
+
+# ── Diff-drive params ───────────────────────────────────────────
+VX_GAIN = 0.15
+WZ_GAIN = 0.8
+VX_MAX = 0.3
+WZ_MAX = 0.6
+CTRL_RATE = 3.0
+STOP_TIMEOUT = 1.0
+
+# ── Obstacle avoidance ──────────────────────────────────────────
+OBSTACLE_MIN_DIST = 0.3
+OBSTACLE_BRAKE_DIST = 0.8
+DEPTH_CENTER_STRIP = 0.3
+
+
+def _load_model(device="cpu"):
+    """Load model — uses cache if available, else downloads & caches."""
+    from fast_load import load_openvla_micro
+
+    cache = os.path.expanduser(CACHE_DIR)
+    if not os.path.isdir(cache):
+        # First run: download from HF and cache components
+        print("First-time setup: downloading & caching model...", flush=True)
+        from modeling_openvla_micro import OpenVLAMicro
+        model = OpenVLAMicro.from_pretrained("theguy21/openvla-micro", device=device)
+        import shutil
+        shutil.copytree("model_cache", cache, dirs_exist_ok=True)
+    return load_openvla_micro(cache, device=device)
+
+
+class OpenVLAMicroNode(Node):
+    def __init__(self):
+        super().__init__("vla_micro_node")
+        self.cmd_pub = self.create_publisher(Twist, "/cmd_vel", 10)
+
+        self.instruction = None
+        self.latest_rgb = None
+        self.latest_depth = None
+        self.lock = threading.Lock()
+        self.br = CvBridge()
+
+        self.create_subscription(String, "/vla_instruction", self.instruction_cb, 10)
+        self.create_subscription(RosImage, "/rgb/image_raw", self.image_cb, 10)
+        self.create_subscription(RosImage, "/depth/image_raw", self.depth_cb, 10)
+
+        self.get_logger().info("Loading OpenVLA-Micro...")
+        t0 = time.time()
+        self.model = _load_model(device="cpu")
+        n_params = sum(p.numel() for p in self.model.parameters()) / 1e6
+        self.get_logger().info(
+            f"Loaded in {time.time()-t0:.0f}s ({n_params:.0f}M params)"
+        )
+
+        self.last_cmd_time = time.time()
+        self.timer = self.create_timer(1.0 / CTRL_RATE, self.tick)
+
+    def instruction_cb(self, msg):
+        with self.lock:
+            self.instruction = msg.data.strip()
+        self.get_logger().info(f"Instr: {self.instruction}")
+
+    def image_cb(self, msg):
+        with self.lock:
+            self.latest_rgb = self.br.imgmsg_to_cv2(msg, desired_encoding="rgb8")
+
+    def depth_cb(self, msg):
+        with self.lock:
+            self.latest_depth = self.br.imgmsg_to_cv2(msg, desired_encoding="16UC1")
+
+    def check_obstacles(self):
+        depth = self.latest_depth
+        if depth is None:
+            return 99.0, 99.0, 99.0
+        h, w = depth.shape
+        top = int(h * (0.5 - DEPTH_CENTER_STRIP / 2))
+        bot = int(h * (0.5 + DEPTH_CENTER_STRIP / 2))
+        center = depth[top:bot, :]
+        valid = center[center > 0]
+        if len(valid) == 0:
+            return 99.0, 99.0, 99.0
+        valid_m = valid.astype(np.float32) / 1000.0
+        closest = float(np.min(valid_m))
+        left = center[:, : w // 3]
+        right = center[:, 2 * w // 3 :]
+        left_m = left[left > 0].astype(np.float32) / 1000.0
+        right_m = right[right > 0].astype(np.float32) / 1000.0
+        return (
+            closest,
+            float(np.min(left_m)) if len(left_m) > 0 else 99.0,
+            float(np.min(right_m)) if len(right_m) > 0 else 99.0,
+        )
+
+    def stop(self):
+        self.cmd_pub.publish(Twist())
+        self.last_cmd_time = time.time()
+
+    def tick(self):
+        with self.lock:
+            instr = self.instruction
+            rgb = self.latest_rgb
+
+        if instr is None or rgb is None:
+            if time.time() - self.last_cmd_time > STOP_TIMEOUT:
+                self.stop()
+            return
+
+        pil_img = Image.fromarray(rgb, mode="RGB")
+
+        t0 = time.perf_counter()
+        with torch.no_grad():
+            action = self.model.predict_action(pil_img, instr)
+        dt = time.perf_counter() - t0
+
+        # 7-DoF: [x, y, z, roll, pitch, yaw, gripper] → diff-drive
+        vx = float(action[0]) * VX_GAIN
+        wz = float(action[3]) * WZ_GAIN
+
+        vx = max(-VX_MAX, min(VX_MAX, vx))
+        wz = max(-WZ_MAX, min(WZ_MAX, wz))
+
+        # Obstacle avoidance
+        closest, left_min, right_min = self.check_obstacles()
+        if closest < OBSTACLE_MIN_DIST:
+            self.get_logger().warn(f"Obstacle {closest:.2f}m — stop")
+            self.stop()
+            return
+        if closest < OBSTACLE_BRAKE_DIST:
+            brake = (closest - OBSTACLE_MIN_DIST) / (OBSTACLE_BRAKE_DIST - OBSTACLE_MIN_DIST)
+            brake = max(0.0, min(1.0, brake))
+            vx *= brake
+            if left_min < right_min and abs(wz) < 0.3:
+                wz = 0.5
+            elif right_min < left_min and abs(wz) < 0.3:
+                wz = -0.5
+            self.get_logger().info(f"Brake {brake:.2f} obs={closest:.2f}m")
+
+        cmd = Twist()
+        cmd.linear.x = vx
+        cmd.angular.z = wz
+        self.cmd_pub.publish(cmd)
+        self.last_cmd_time = time.time()
+
+        self.get_logger().info(
+            f"vx={vx:.3f} wz={wz:.3f} ({dt*1000:.0f}ms) [{instr[:40]}]"
+        )
+
+
+def main():
+    rclpy.init()
+    node = OpenVLAMicroNode()
+    try:
+        rclpy.spin(node)
+    except KeyboardInterrupt:
+        node.stop()
+    node.destroy_node()
+    rclpy.shutdown()
+
+
+if __name__ == "__main__":
+    main()