foodforthought-cli 0.2.7__py3-none-any.whl → 0.3.0__py3-none-any.whl

ate/robot/llm_system_id.py

@@ -0,0 +1,654 @@
+#!/usr/bin/env python3
+"""
+LLM-Based Robot System Identification
+
+No markers. No calibration. Just:
+1. Wiggle each motor
+2. Show before/after frames to multimodal LLM
+3. LLM describes what moved and how
+4. Generate URDF with semantic labels
+
+This bridges real robots to Artifex Desktop.
+
+Usage:
+    python -m ate.robot.llm_system_id --port /dev/cu.usbserial-10 --motors 11,12,13
+"""
+
+import os
+import time
+import json
+import base64
+import io
+from dataclasses import dataclass, field
+from typing import List, Dict, Optional, Tuple, Any
+from pathlib import Path
+from enum import Enum
+
+try:
+    import cv2
+    import numpy as np
+    HAS_CV = True
+except ImportError:
+    HAS_CV = False
+
+try:
+    from anthropic import Anthropic
+    HAS_ANTHROPIC = True
+except ImportError:
+    HAS_ANTHROPIC = False
+
+# Import unified JointType from shared types module
+from ate.robot.types import JointType
+
+
+@dataclass
+class JointDiscovery:
+    """What we learned about a joint from the LLM."""
+    motor_id: int
+    joint_type: JointType
+    body_part: str  # "arm", "leg", "gripper", etc.
+    joint_name: str  # "shoulder", "elbow", "knee", etc.
+    axis_description: str  # "horizontal", "vertical", etc.
+    range_degrees: float  # Approximate range of motion
+    parent_link: str  # What this is attached to
+    observations: List[str]  # Raw LLM observations
+    confidence: float
+
+    # Images for reference
+    frame_min: Optional[np.ndarray] = None
+    frame_max: Optional[np.ndarray] = None
+
+    def to_dict(self) -> dict:
+        return {
+            "motor_id": self.motor_id,
+            "joint_type": self.joint_type.value,
+            "body_part": self.body_part,
+            "joint_name": self.joint_name,
+            "axis_description": self.axis_description,
+            "range_degrees": self.range_degrees,
+            "parent_link": self.parent_link,
+            "observations": self.observations,
+            "confidence": self.confidence,
+        }
+
+
+class LLMSystemIdentifier:
+    """
+    Use multimodal LLM to identify robot structure.
+
+    The LLM acts as a "robot coach" that watches the robot move
+    and describes what it sees, enabling us to build a URDF.
+    """
+
+    def __init__(
+        self,
+        camera_index: int = 0,
+        model: str = "claude-sonnet-4-20250514",
+    ):
+        if not HAS_CV:
+            raise ImportError("OpenCV required: pip install opencv-python")
+        if not HAS_ANTHROPIC:
+            raise ImportError("Anthropic SDK required: pip install anthropic")
+
+        self.camera_index = camera_index
+        self.model = model
+        self.client = Anthropic()
+
+        # Camera
+        self.cap = None
+
+        # Results
+        self.discoveries: Dict[int, JointDiscovery] = {}
+        self.robot_overview: Optional[str] = None
+
+    def connect_camera(self) -> bool:
+        """Open camera connection."""
+        self.cap = cv2.VideoCapture(self.camera_index)
+        if not self.cap.isOpened():
+            print(f"Failed to open camera {self.camera_index}")
+            return False
+
+        # Set resolution
+        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
+        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
+
+        # Warm up
+        for _ in range(5):
+            self.cap.read()
+
+        return True
+
+    def disconnect(self):
+        """Close camera."""
+        if self.cap:
+            self.cap.release()
+            self.cap = None
+
+    def capture_frame(self) -> np.ndarray:
+        """Capture a single frame."""
+        if not self.cap:
+            return np.array([])
+
+        # Discard stale frames
+        for _ in range(3):
+            self.cap.read()
+
+        ret, frame = self.cap.read()
+        return frame if ret else np.array([])
+
+    def frame_to_base64(self, frame: np.ndarray) -> str:
+        """Convert frame to base64 for LLM."""
+        _, buffer = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 85])
+        return base64.standard_b64encode(buffer).decode('utf-8')
+
+    def analyze_robot_overview(self, frame: np.ndarray) -> str:
+        """Get initial overview of the robot."""
+        print("\n  Analyzing robot structure...")
+
+        image_data = self.frame_to_base64(frame)
+
+        response = self.client.messages.create(
+            model=self.model,
+            max_tokens=1024,
+            messages=[{
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image",
+                        "source": {
+                            "type": "base64",
+                            "media_type": "image/jpeg",
+                            "data": image_data,
+                        },
+                    },
+                    {
+                        "type": "text",
+                        "text": """Analyze this robot. Describe:
+
+1. What type of robot is this? (quadruped, arm, humanoid, etc.)
+2. How many limbs/appendages do you see?
+3. What are the main body parts? (body, legs, arm, gripper, head, etc.)
+4. Estimate the approximate size.
+
+Be concise but specific. This will help us understand the kinematic structure."""
+                    }
+                ],
+            }],
+        )
+
+        return response.content[0].text
+
+    def analyze_motor_motion(
+        self,
+        motor_id: int,
+        frame_before: np.ndarray,
+        frame_after: np.ndarray,
+        pwm_before: int,
+        pwm_after: int,
+    ) -> JointDiscovery:
+        """
+        Analyze what happened when a motor moved.
+
+        Shows before/after frames to LLM and asks it to describe the motion.
+        """
+        print(f"\n  Analyzing motor {motor_id} motion...")
+
+        before_b64 = self.frame_to_base64(frame_before)
+        after_b64 = self.frame_to_base64(frame_after)
+
+        response = self.client.messages.create(
+            model=self.model,
+            max_tokens=1024,
+            messages=[{
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": f"I'm testing motor {motor_id} on a robot. Here are two images showing BEFORE (PWM={pwm_before}) and AFTER (PWM={pwm_after}) positions."
+                    },
+                    {
+                        "type": "text",
+                        "text": "BEFORE position:"
+                    },
+                    {
+                        "type": "image",
+                        "source": {
+                            "type": "base64",
+                            "media_type": "image/jpeg",
+                            "data": before_b64,
+                        },
+                    },
+                    {
+                        "type": "text",
+                        "text": "AFTER position:"
+                    },
+                    {
+                        "type": "image",
+                        "source": {
+                            "type": "base64",
+                            "media_type": "image/jpeg",
+                            "data": after_b64,
+                        },
+                    },
+                    {
+                        "type": "text",
+                        "text": """Analyze the motion. Answer these questions:
+
+1. WHAT MOVED? (e.g., "the gripper", "front-left leg", "arm shoulder")
+2. HOW DID IT MOVE? (e.g., "rotated downward", "opened wider", "extended forward")
+3. JOINT TYPE: Is this a rotation (revolute) or linear slide (prismatic)?
+4. AXIS: What axis does it rotate/move around? (horizontal/vertical/forward-back)
+5. RANGE: Approximately how many degrees did it rotate? (or mm if linear)
+6. PARENT: What body part is this attached to?
+7. CONFIDENCE: How certain are you? (high/medium/low)
+
+If nothing visibly moved, say "NO VISIBLE MOTION" and explain why (maybe the camera angle, or the motor isn't connected).
+
+Format your response as:
+BODY_PART: [part name]
+JOINT_NAME: [specific joint name]
+MOTION: [description]
+JOINT_TYPE: [revolute/prismatic/fixed]
+AXIS: [description]
+RANGE_DEGREES: [number]
+PARENT: [parent link name]
+CONFIDENCE: [high/medium/low]
+NOTES: [any other observations]"""
+                    }
+                ],
+            }],
+        )
+
+        # Parse response
+        text = response.content[0].text
+        return self._parse_motion_analysis(motor_id, text, frame_before, frame_after)
+
+    def _parse_motion_analysis(
+        self,
+        motor_id: int,
+        text: str,
+        frame_min: np.ndarray,
+        frame_max: np.ndarray,
+    ) -> JointDiscovery:
+        """Parse LLM response into structured JointDiscovery."""
+
+        lines = text.strip().split('\n')
+        parsed = {}
+
+        for line in lines:
+            if ':' in line:
+                key, value = line.split(':', 1)
+                parsed[key.strip().upper()] = value.strip()
+
+        # Extract values with defaults
+        body_part = parsed.get('BODY_PART', 'unknown').lower()
+        joint_name = parsed.get('JOINT_NAME', f'joint_{motor_id}').lower()
+
+        # Parse joint type
+        joint_type_str = parsed.get('JOINT_TYPE', 'unknown').lower()
+        if 'revolute' in joint_type_str or 'rotat' in joint_type_str:
+            joint_type = JointType.REVOLUTE
+        elif 'prismatic' in joint_type_str or 'linear' in joint_type_str:
+            joint_type = JointType.PRISMATIC
+        elif 'fixed' in joint_type_str or 'no' in joint_type_str:
+            joint_type = JointType.FIXED
+        else:
+            joint_type = JointType.UNKNOWN
+
+        # Parse range
+        range_str = parsed.get('RANGE_DEGREES', '0')
+        try:
+            range_degrees = float(''.join(c for c in range_str if c.isdigit() or c == '.'))
+        except ValueError:
+            range_degrees = 0.0
+
+        # Parse confidence
+        conf_str = parsed.get('CONFIDENCE', 'medium').lower()
+        if 'high' in conf_str:
+            confidence = 0.9
+        elif 'low' in conf_str:
+            confidence = 0.4
+        else:
+            confidence = 0.7
+
+        return JointDiscovery(
+            motor_id=motor_id,
+            joint_type=joint_type,
+            body_part=body_part,
+            joint_name=joint_name,
+            axis_description=parsed.get('AXIS', 'unknown'),
+            range_degrees=range_degrees,
+            parent_link=parsed.get('PARENT', 'base_link'),
+            observations=[
+                parsed.get('MOTION', ''),
+                parsed.get('NOTES', ''),
+            ],
+            confidence=confidence,
+            frame_min=frame_min,
+            frame_max=frame_max,
+        )
+
+    def run_wiggle_test(
+        self,
+        motor_id: int,
+        pwm_min: int = 500,
+        pwm_max: int = 2500,
+        cmd_func=None,
+    ) -> JointDiscovery:
+        """
+        Run wiggle test on a single motor.
+
+        1. Move to min position, capture frame
+        2. Move to max position, capture frame
+        3. Ask LLM to analyze the motion
+        """
+        print(f"\n--- Wiggle Test: Motor {motor_id} ---")
+
+        # Move to min, capture
+        print(f"  Moving to PWM {pwm_min}...")
+        if cmd_func:
+            cmd_func(f"dog.set_servo({motor_id}, {pwm_min}, 800)")
+        time.sleep(1.0)
+        frame_min = self.capture_frame()
+
+        # Show frame
+        if frame_min.size > 0:
+            cv2.imshow("System ID", frame_min)
+            cv2.waitKey(100)
+
+        # Move to max, capture
+        print(f"  Moving to PWM {pwm_max}...")
+        if cmd_func:
+            cmd_func(f"dog.set_servo({motor_id}, {pwm_max}, 800)")
+        time.sleep(1.0)
+        frame_max = self.capture_frame()
+
+        # Show frame
+        if frame_max.size > 0:
+            cv2.imshow("System ID", frame_max)
+            cv2.waitKey(100)
+
+        # Save frames for debugging
+        cv2.imwrite(f"/tmp/motor_{motor_id}_min.jpg", frame_min)
+        cv2.imwrite(f"/tmp/motor_{motor_id}_max.jpg", frame_max)
+
+        # Return to center
+        print(f"  Returning to center...")
+        if cmd_func:
+            cmd_func(f"dog.set_servo({motor_id}, 1500, 500)")
+        time.sleep(0.5)
+
+        # Analyze with LLM
+        discovery = self.analyze_motor_motion(
+            motor_id,
+            frame_min,
+            frame_max,
+            pwm_min,
+            pwm_max,
+        )
+
+        print(f"\n  Result:")
+        print(f"    Body part: {discovery.body_part}")
+        print(f"    Joint: {discovery.joint_name}")
+        print(f"    Type: {discovery.joint_type.value}")
+        print(f"    Range: ~{discovery.range_degrees:.0f}°")
+        print(f"    Confidence: {discovery.confidence:.0%}")
+
+        return discovery
+
+    def identify_robot(
+        self,
+        motor_ids: List[int],
+        cmd_func=None,
+    ) -> Dict[int, JointDiscovery]:
+        """
+        Full robot identification.
+
+        1. Get overview of robot
+        2. Wiggle each motor
+        3. Build kinematic understanding
+        """
+        print("\n" + "=" * 60)
+        print("LLM-BASED ROBOT SYSTEM IDENTIFICATION")
+        print("=" * 60)
+
+        # Initial frame and overview
+        frame = self.capture_frame()
+        if frame.size == 0:
+            print("ERROR: Cannot capture from camera!")
+            return {}
+
+        cv2.imshow("System ID", frame)
+        cv2.waitKey(100)
+        cv2.imwrite("/tmp/robot_overview.jpg", frame)
+
+        self.robot_overview = self.analyze_robot_overview(frame)
+        print(f"\n  Robot Overview:\n{self.robot_overview}")
+
+        # Test each motor
+        for motor_id in motor_ids:
+            discovery = self.run_wiggle_test(motor_id, cmd_func=cmd_func)
+            self.discoveries[motor_id] = discovery
+
+        return self.discoveries
+
+    def generate_urdf(self, robot_name: str = "discovered_robot") -> str:
+        """
+        Generate URDF from discoveries.
+
+        This creates a valid URDF that can be loaded into Artifex Desktop
+        for visualization and simulation.
+        """
+        lines = [
+            '<?xml version="1.0"?>',
+            f'<robot name="{robot_name}">',
+            '  <!-- Generated by LLM System Identification -->',
+            f'  <!-- Robot Overview: {self.robot_overview[:100] if self.robot_overview else "Unknown"}... -->',
+            '',
+            '  <!-- ========== LINKS ========== -->',
+            '',
+            '  <!-- Base Link -->',
+            '  <link name="base_link">',
+            '    <visual>',
+            '      <geometry>',
+            '        <box size="0.15 0.10 0.05"/>',
+            '      </geometry>',
+            '      <material name="gray">',
+            '        <color rgba="0.5 0.5 0.5 1"/>',
+            '      </material>',
+            '    </visual>',
+            '    <collision>',
+            '      <geometry>',
+            '        <box size="0.15 0.10 0.05"/>',
+            '      </geometry>',
+            '    </collision>',
+            '    <inertial>',
+            '      <mass value="0.5"/>',
+            '      <inertia ixx="0.001" ixy="0" ixz="0" iyy="0.001" iyz="0" izz="0.001"/>',
+            '    </inertial>',
+            '  </link>',
+            '',
+        ]
+
+        # Sort discoveries to build kinematic chain
+        # For now, assume simple chain: base → joint1 → joint2 → ...
+        parent_link = "base_link"
+
+        for motor_id, discovery in sorted(self.discoveries.items()):
+            link_name = f"{discovery.body_part}_{discovery.joint_name}".replace(" ", "_")
+            joint_name = f"motor_{motor_id}_{discovery.joint_name}".replace(" ", "_")
+
+            # Determine axis from description
+            axis = "0 0 1"  # Default Z-axis
+            axis_desc = discovery.axis_description.lower()
+            if 'horizontal' in axis_desc or 'side' in axis_desc:
+                axis = "1 0 0"  # X-axis
+            elif 'forward' in axis_desc or 'pitch' in axis_desc:
+                axis = "0 1 0"  # Y-axis
+            elif 'vertical' in axis_desc or 'yaw' in axis_desc:
+                axis = "0 0 1"  # Z-axis
+
+            # Joint limits
+            range_rad = discovery.range_degrees * 3.14159 / 180.0
+            lower = -range_rad / 2
+            upper = range_rad / 2
+
+            # Add joint
+            joint_type = discovery.joint_type.value
+            if joint_type == "unknown":
+                joint_type = "revolute"
+
+            lines.extend([
+                f'  <!-- Motor {motor_id}: {discovery.body_part} {discovery.joint_name} -->',
+                f'  <!-- Observations: {"; ".join(discovery.observations)} -->',
+                f'  <joint name="{joint_name}" type="{joint_type}">',
+                f'    <parent link="{parent_link}"/>',
+                f'    <child link="{link_name}"/>',
+                f'    <origin xyz="0.05 0 0" rpy="0 0 0"/>',
+                f'    <axis xyz="{axis}"/>',
+            ])
+
+            if joint_type == "revolute":
+                lines.append(f'    <limit lower="{lower:.3f}" upper="{upper:.3f}" effort="10" velocity="3"/>')
+
+            lines.extend([
+                '  </joint>',
+                '',
+                f'  <link name="{link_name}">',
+                '    <visual>',
+                '      <origin xyz="0.025 0 0"/>',
+                '      <geometry>',
+                '        <cylinder radius="0.01" length="0.05"/>',
+                '      </geometry>',
+                '      <material name="blue">',
+                '        <color rgba="0.2 0.2 0.8 1"/>',
+                '      </material>',
+                '    </visual>',
+                '    <collision>',
+                '      <origin xyz="0.025 0 0"/>',
+                '      <geometry>',
+                '        <cylinder radius="0.01" length="0.05"/>',
+                '      </geometry>',
+                '    </collision>',
+                '    <inertial>',
+                '      <mass value="0.1"/>',
+                '      <inertia ixx="0.0001" ixy="0" ixz="0" iyy="0.0001" iyz="0" izz="0.0001"/>',
+                '    </inertial>',
+                '  </link>',
+                '',
+            ])
+
+            parent_link = link_name
+
+        lines.append('</robot>')
+
+        return '\n'.join(lines)
+
+    def save_results(self, output_dir: Path):
+        """Save all results."""
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        # Save JSON
+        data = {
+            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+            "robot_overview": self.robot_overview,
+            "discoveries": {
+                str(k): v.to_dict() for k, v in self.discoveries.items()
+            },
+        }
+
+        json_path = output_dir / "system_id_results.json"
+        with open(json_path, 'w') as f:
+            json.dump(data, f, indent=2)
+        print(f"\nSaved results to: {json_path}")
+
+        # Save URDF
+        urdf = self.generate_urdf()
+        urdf_path = output_dir / "discovered_robot.urdf"
+        with open(urdf_path, 'w') as f:
+            f.write(urdf)
+        print(f"Saved URDF to: {urdf_path}")
+
+        # Save before/after images
+        for motor_id, discovery in self.discoveries.items():
+            if discovery.frame_min is not None:
+                cv2.imwrite(str(output_dir / f"motor_{motor_id}_min.jpg"), discovery.frame_min)
+            if discovery.frame_max is not None:
+                cv2.imwrite(str(output_dir / f"motor_{motor_id}_max.jpg"), discovery.frame_max)
+
+        print(f"Saved motion images to: {output_dir}")
+
+        return urdf_path
+
+
+def main():
+    """Main entry point."""
+    import argparse
+
+    parser = argparse.ArgumentParser(description="LLM-Based Robot System Identification")
+    parser.add_argument("--port", default="/dev/cu.usbserial-10", help="Robot serial port")
+    parser.add_argument("--motors", default="11,12,13", help="Motor IDs to test")
+    parser.add_argument("--camera", type=int, default=0, help="Camera index")
+    parser.add_argument("--output", default="/tmp/llm_system_id", help="Output directory")
+    parser.add_argument("--demo", action="store_true", help="Demo mode (no robot)")
+
+    args = parser.parse_args()
+
+    motor_ids = [int(x) for x in args.motors.split(",")]
+
+    # Create identifier
+    identifier = LLMSystemIdentifier(camera_index=args.camera)
+
+    if not identifier.connect_camera():
+        print("Failed to connect camera!")
+        return
+
+    cmd_func = None
+    ser = None
+
+    if not args.demo:
+        # Connect to robot
+        import serial
+        ser = serial.Serial(args.port, 115200, timeout=2)
+        time.sleep(0.5)
+        ser.read(ser.in_waiting)
+
+        def cmd(command: str):
+            ser.write(f"{command}\r\n".encode())
+            time.sleep(0.3)
+            return ser.read(ser.in_waiting).decode('utf-8', errors='ignore')
+
+        cmd("from HW_MechDog import MechDog")
+        time.sleep(1.5)
+        cmd("dog = MechDog()")
+        time.sleep(1.5)
+
+        cmd_func = cmd
+
+    try:
+        # Run identification
+        discoveries = identifier.identify_robot(motor_ids, cmd_func=cmd_func)
+
+        if discoveries:
+            # Save results
+            identifier.save_results(Path(args.output))
+
+            print("\n" + "=" * 60)
+            print("IDENTIFICATION COMPLETE")
+            print("=" * 60)
+            print(f"\nDiscovered {len(discoveries)} joints:")
+            for motor_id, disc in discoveries.items():
+                print(f"  Motor {motor_id}: {disc.body_part} {disc.joint_name} ({disc.joint_type.value})")
+
+            print(f"\nURDF saved to: {args.output}/discovered_robot.urdf")
+            print("You can import this into Artifex Desktop!")
+
+    finally:
+        identifier.disconnect()
+        if ser:
+            ser.close()
+        cv2.destroyAllWindows()
+
+
+if __name__ == "__main__":
+    main()