diff --git a/.gitignore b/.gitignore index b47e22cbfad..b9211a4df39 100644 --- a/.gitignore +++ b/.gitignore @@ -15,6 +15,7 @@ ### Environments & Dependencies ### .env .venv +.venv-lerobot-test env/ venv/ env.bak/ diff --git a/docs/source/Brewie.mdx b/docs/source/Brewie.mdx new file mode 100644 index 00000000000..befbe6b6560 --- /dev/null +++ b/docs/source/Brewie.mdx @@ -0,0 +1,255 @@ +# Brewie + +The Brewie robot is a humanoid robot based on the Hiwonder-manufactured AiNex model, redesigned to use ROS 1 - noetic (Robot Operating System) for communication and control. + +## Hardware Overview + +The Brewie robot is a humanoid robot based on the Hiwonder-manufactured AiNex model. This modification features: + +- **Humanoid Design**: Bipedal walking robot with human-like proportions +- **Dual Manipulators**: Equipped with two articulated arms for manipulation tasks +- **Locomotion**: Capable of walking and navigating in its environment +- **Vision System**: Features a movable camera for visual perception +- **Control Unit**: Powered by a Raspberry Pi 5 as the main computing platform +- **Network Connectivity**: Can be controlled via local network connection, eliminating the need for direct physical connections + +## Prerequisites + +- ROS environment running on the robot +- Network connectivity between control computer and robot + +## Install LeRobot + +Follow the [installation instructions](./installation) to install LeRobot. + +Install LeRobot with Brewie dependencies: + +```bash +pip install -e ".[brewie]" +``` + +## ROS Configuration + +The Brewie robot uses ROS for communication instead of direct motor control. Ensure your ROS environment is properly configured: + +### ROS Master Setup + +The robot runs a ROS master that needs to be accessible from your control computer: + +```bash +# On the robot (Raspberry Pi) +roscore + +# Or with custom IP/port +export ROS_MASTER_URI=http://192.168.1.100:11311 +roscore +``` + +### Network Configuration + +Ensure both devices are on the same network and can communicate: + +```bash +# Test connectivity from control computer +ping +``` + +## Robot Configuration + +### Connection Parameters + +Configure the robot connection in your code: + +```python +from lerobot.robots.brewie.config_Brewie import BrewieConfig + +config = BrewieConfig( + master_ip="192.168.1.100", # Robot's IP address + master_port=9090, # ROS bridge port + servo_duration=0.1, # Movement duration + max_relative_target=50.0 # Safety limit +) +``` + +### Servo Mapping + +The robot uses the following servo ID mapping: + +| ID | Joint Name | +|----|------------| +| 13 | left_shoulder_pan | +| 14 | right_shoulder_pan | +| 15 | left_shoulder_lift | +| 16 | right_shoulder_lift | +| 17 | left_forearm_roll | +| 18 | right_forearm_roll | +| 19 | left_forearm_pitch | +| 20 | right_forearm_pitch | +| 21 | left_gripper | +| 22 | right_gripper | + +## ROS Interface + +### Services + +#### Get Servo Positions +- **Name:** `/ros_robot_controller/bus_servo/get_position` +- **Type:** `ros_robot_controller/GetBusServosPosition` +- **Request:** `{'id': [13, 14, 15, 16, 17, 18, 19, 20, 21, 22]}` +- **Response:** `{'success': True, 'position': [{'id': 13, 'position': 872}, ...]}` + +### Topics + +#### Set Servo Positions +- **Name:** `/ros_robot_controller/bus_servo/set_position` +- **Type:** `ros_robot_controller/SetBusServosPosition` +- **Message:** `{'duration': 0.1, 'position': [{'id': 13, 'position': 500}, ...]}` + +#### Camera Image +- **Name:** `/camera/image_raw/compressed` +- **Type:** `sensor_msgs/CompressedImage` + +## Basic Usage + +### Initialize Robot + +```python +from lerobot.robots.brewie.Brewie_base import BrewieBase +from lerobot.robots.brewie.config_Brewie import BrewieConfig + +config = BrewieConfig( + master_ip="192.168.1.100", + master_port=9090 +) + +robot = BrewieBase(config) +robot.connect() + +# Get current observation +observation = robot.get_observation() + +# Send action command +action = { + 'left_shoulder_pan': 500, + 'right_shoulder_pan': 500, + # ... other joints +} +robot.send_action(action) + +robot.disconnect() +``` + +## Dataset Recording + +### Using Configuration-Based Recording + +The recommended way to record datasets with Brewie is using the configuration-based approach: + +#### 1. Setup HuggingFace Token + +```bash +export HUGGINGFACE_TOKEN=your_token_here +``` + +#### 2. Configure Recording Settings + +Edit `examples/brewie/record_config.py`: + +```python +config = RecordingConfig( + hf_username="your_username", # Replace with your username + dataset_name="brewie_manipulation", + ros_master_ip="192.168.1.100", # Robot's IP address + ros_master_port=9090, + num_episodes=5, + episode_time_sec=30, + task_description="Brewie robot manipulation task", + task_category="manipulation", + difficulty_level="beginner" +) +``` + +#### 3. Run Recording + +```bash +python examples/brewie/record_with_config.py +``` + +### Configuration Options + +The `record_config.py` file provides several predefined configurations: + +- **Quick Demo**: `RecordingConfig.quick_demo()` - 2 short episodes for testing +- **Full Dataset**: `RecordingConfig.full_dataset()` - 20 episodes for training +- **Pick & Place**: `RecordingConfig.pick_place_task()` - Specialized for grasping tasks +- **Assembly**: `RecordingConfig.assembly_task()` - Complex manipulation tasks + +### Advanced Recording Features + +- **Resume Recording**: Continue adding episodes to existing datasets +- **Automatic Hub Upload**: Push datasets to HuggingFace Hub automatically +- **Multi-threaded Image Recording**: Optimized for video capture +- **Interactive Mode Selection**: Choose between creating new or continuing existing datasets + +### Recording Controls + +During recording: +- **ENTER**: Start/continue episode recording +- **ESC**: Stop recording +- **R**: Rewrite current episode + +## Compatibility + +The Brewie implementation is fully compatible with LeRobot standards: + +- Maintains `Robot` class interface +- Supports dataset collection and training +- Compatible with all LeRobot policies +- Preserves observation and action formats + +## Troubleshooting + +### Connection Issues + +1. **Check ROS Master**: Ensure ROS master is running on the robot +2. **Network Connectivity**: Verify both devices are on the same network +3. **Firewall**: Check if firewall is blocking ROS bridge port (9090) +4. **ROS Environment**: Ensure `ROS_MASTER_URI` is set correctly + +### Recording Issues + +1. **Token Authentication**: Verify HuggingFace token is valid and has write permissions +2. **Dataset Permissions**: Check if you have access to create datasets on Hub +3. **Memory Issues**: Reduce FPS or episode duration if running out of memory +4. **Video Sync**: Lower FPS if experiencing video synchronization issues + +### Performance Optimization + +For better performance: +- Reduce FPS to 15-20 for stability +- Use fewer image writer threads (2-4) +- Ensure sufficient disk space for video recording +- Close unnecessary applications during recording + +## Examples + +See files in `examples/brewie/`: +- `record_config.py` - Configuration file for recording +- `record_with_config.py` - Main recording script with configuration support +- `README.md` - Additional usage examples and documentation + +## Testing + +Tests are available in `tests/robots/test_brewie_ros.py` to verify: +- Robot initialization and connection +- ROS service and topic communication +- Observation retrieval and action sending +- Proper disconnection and cleanup + +## Requirements + +- Python 3.8+ +- ROS (Robot Operating System) +- `roslibpy` library for Python-ROS communication +- OpenCV for image processing +- HuggingFace Hub for dataset management \ No newline at end of file diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 1055975d7bc..9825a4ee0c2 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -105,6 +105,8 @@ title: OMX - local: openarm title: OpenArm + - local: Brewie + title: Brewie title: "Robots" - sections: - local: phone_teleop diff --git a/examples/brewie/record_config.py b/examples/brewie/record_config.py new file mode 100644 index 00000000000..0425ff26945 --- /dev/null +++ b/examples/brewie/record_config.py @@ -0,0 +1,342 @@ +#!/usr/bin/env python + +""" +Configuration file for Brewie dataset recording. + +This file contains all settings for dataset recording. +Modify values in this file before running record.py + +IMPORTANT: hf_token is now obtained from HUGGINGFACE_TOKEN environment variable +or from command line arguments. Do not store tokens in code! + +USAGE EXAMPLES: + +1. With environment variable (recommended): + export HUGGINGFACE_TOKEN=your_token_here + python examples/brewie/record_with_config.py + +2. With command line argument: + python examples/brewie/record_with_config.py --hf-token your_token_here + +3. Interactive input (token will be requested at startup): + python examples/brewie/record_with_config.py + +4. Alternative environment variable: + export HF_TOKEN=your_token_here + python examples/brewie/record_with_config.py + +SECURITY: +- Never commit tokens to code +- Use environment variables for production +- Tokens are entered hidden (not displayed in terminal) +""" + +import os +import sys +import getpass +from dataclasses import dataclass +from typing import Optional + +def get_hf_token() -> str: + """ + Get HuggingFace token from environment variable or prompt user. + + Token retrieval order: + 1. HUGGINGFACE_TOKEN environment variable + 2. HF_TOKEN environment variable + 3. Command line argument --hf-token + 4. Interactive input (hidden) + + Returns: + str: HuggingFace token + + Raises: + ValueError: If token not found and user cancelled input + """ + def validate_token(token: str) -> str: + """Validate HuggingFace token format.""" + if not token or not token.strip(): + raise ValueError("Token cannot be empty") + + token = token.strip() + + # Basic validation of HuggingFace token format + if not token.startswith("hf_"): + print("Warning: Token does not start with 'hf_'. Make sure this is a valid HuggingFace token.") + + if len(token) < 10: + raise ValueError("Token is too short. Check token correctness.") + + return token + + # 1. Check environment variables + token = os.getenv("HUGGINGFACE_TOKEN") or os.getenv("HF_TOKEN") + if token: + return validate_token(token) + + # 2. Check command line arguments + if "--hf-token" in sys.argv: + try: + token_index = sys.argv.index("--hf-token") + if token_index + 1 < len(sys.argv): + token = sys.argv[token_index + 1] + return validate_token(token) + except (ValueError, IndexError): + pass + + # 3. Prompt user + print("HuggingFace token not found in environment variables.") + print("Set HUGGINGFACE_TOKEN environment variable or enter token:") + print(" export HUGGINGFACE_TOKEN=your_token_here") + print() + + try: + token = getpass.getpass("Enter your HuggingFace token (hidden input): ") + return validate_token(token) + except KeyboardInterrupt: + print("\nInput cancelled by user") + raise ValueError("Token not provided") + +@dataclass +class RecordingConfig: + """Configuration for Brewie dataset recording.""" + + # ============================================================================= + # HUGGINGFACE SETTINGS + # ============================================================================= + + # Your HuggingFace credentials + hf_username: str = "your_username" # Replace with your username + # hf_token is now obtained dynamically from environment variables or input + + # Dataset name (will be created as username/dataset_name) + dataset_name: str = "hit_detection" + + # ============================================================================= + # RECORDING SETTINGS + # ============================================================================= + + # Number of episodes to record + num_episodes: int = 5 + + # Recording frequency (frames per second) + # WARNING: After adding new sensors, it's recommended to reduce FPS + # to prevent video synchronization issues + fps: int = 20 # Reduced from 30 to 20 for stability + + # Duration of each episode in seconds + episode_time_sec: int = 30 + + # Reset time between episodes in seconds + reset_time_sec: int = 5 + + # ============================================================================= + # TASK DESCRIPTION + # ============================================================================= + + # Description of the task the robot will perform + task_description: str = "Brewie robot manipulation demonstration" + + # Additional metadata + task_category: str = "manipulation" # manipulation, pick_place, assembly, etc. + difficulty_level: str = "beginner" # beginner, intermediate, advanced + + # ============================================================================= + # ROBOT SETTINGS + # ============================================================================= + + # ROS connection parameters + ros_master_ip: str = "localhost" + ros_master_port: int = 9090 + + # Safety settings + max_relative_target: float = 50.0 # Maximum relative movement per step + servo_duration: float = 0.1 # Duration for servo movements + + # ============================================================================= + # DATASET SETTINGS + # ============================================================================= + + # Use video in dataset + use_videos: bool = True + + # Number of threads for image recording + image_writer_threads: int = 4 + + # ============================================================================= + # ADDITIONAL SETTINGS + # ============================================================================= + + # Display data during recording + display_data: bool = True + + # Session name for visualization + session_name: str = "brewie_record" + + # Automatically push to Hub after recording + auto_push_to_hub: bool = True + + # Continue recording in existing dataset (add new episodes) + resume_existing_dataset: bool = False + + def get_hf_token(self) -> str: + """ + Get HuggingFace token for this configuration. + + Returns: + str: HuggingFace token + """ + return get_hf_token() + + # ============================================================================= + # PREDEFINED CONFIGURATIONS + # ============================================================================= + + @classmethod + def quick_demo(cls) -> "RecordingConfig": + """Quick demo - 2 short episodes.""" + return cls( + hf_username="forroot", # REQUIRED: replace with your username + # hf_token is obtained dynamically from environment variables + ros_master_ip="192.168.20.21", + ros_master_port=9090, + num_episodes=2, + fps=20, + episode_time_sec=15, + reset_time_sec=3, + task_description="Fast demo of robot movements", + task_category="demo", + difficulty_level="beginner", + resume_existing_dataset=True + ) + + @classmethod + def detection_aim(cls) -> "RecordingConfig": + """racking and aiming at an enemy robot for fire. FAST MODE""" + return cls( + hf_username="forroot", # REQUIRED: replace with your username + dataset_name ="detection_aim", + ros_master_ip="192.168.20.21", + ros_master_port=9090, + num_episodes=2, + fps=20, + episode_time_sec=15, + reset_time_sec=3, + task_description="Tracking and aiming at an enemy robot for fire", + task_category="aim", + difficulty_level="beginner", + resume_existing_dataset=True + ) + @classmethod + def hit_detection(cls) -> "RecordingConfig": + return cls( + hf_username="forroot", # REQUIRED: replace with your username + dataset_name ="TERST2", + ros_master_ip="192.168.20.23", + ros_master_port=9090, + num_episodes=2, + fps=20, + episode_time_sec=2, + reset_time_sec=3, + task_description="Testing observation of hit data FIRE button True = human hit verification", + task_category="hit", + difficulty_level="beginner", + resume_existing_dataset=False + ) + + @classmethod + def resume_demo(cls) -> "RecordingConfig": + """Demonstration of continuing recording in existing dataset.""" + return cls( + hf_username="forroot", # REQUIRED: replace with your username + # hf_token is obtained dynamically from environment variables + ros_master_ip="192.168.20.21", + ros_master_port=9090, + num_episodes=3, + episode_time_sec=20, + reset_time_sec=5, + task_description="Additional episodes for existing dataset", + task_category="demo", + difficulty_level="beginner", + resume_existing_dataset=True # Enable resume recording mode + ) + + @classmethod + def full_dataset(cls) -> "RecordingConfig": + """Full dataset - many episodes for training.""" + return cls( + num_episodes=20, + episode_time_sec=60, + reset_time_sec=10, + task_description="Full set of demonstrations for training", + task_category="manipulation", + difficulty_level="intermediate" + ) + + @classmethod + def pick_place_task(cls) -> "RecordingConfig": + """Configuration for pick and place task.""" + return cls( + num_episodes=10, + episode_time_sec=45, + reset_time_sec=8, + task_description="Grasping and placing objects", + task_category="pick_place", + difficulty_level="intermediate" + ) + + @classmethod + def assembly_task(cls) -> "RecordingConfig": + """Configuration for assembly task.""" + return cls( + num_episodes=15, + episode_time_sec=90, + reset_time_sec=15, + task_description="Assembly of parts by robot", + task_category="assembly", + difficulty_level="advanced" + ) + + @classmethod + def optimized_with_sensors(cls) -> "RecordingConfig": + """Optimized configuration for working with new sensors.""" + return cls( + hf_username="forroot", # REQUIRED: replace with your username + ros_master_ip="192.168.20.21", + ros_master_port=9090, + num_episodes=5, + episode_time_sec=30, + reset_time_sec=5, + fps=15, # Reduced frequency for stability + task_description="Optimized recording with new sensors", + task_category="demo", + difficulty_level="beginner", + resume_existing_dataset=True, + image_writer_threads=2, # Fewer threads for stability + use_videos=True + ) + +# ============================================================================= +# CONFIGURATION SELECTION +# ============================================================================= + +# Choose one of the predefined configurations or create your own +#config = RecordingConfig.optimized_with_sensors() # Recommended for new sensors +config = RecordingConfig.hit_detection() +# config = RecordingConfig.resume_demo() # For continuing recording in existing dataset +# config = RecordingConfig.full_dataset() +# config = RecordingConfig.pick_place_task() +# config = RecordingConfig.assembly_task() + +# Or create your own configuration +''' +config = RecordingConfig( + hf_username="your_username", # REQUIRED: replace with your username + # hf_token is obtained automatically from environment variables + dataset_name="brewie_my_task", + num_episodes=5, + episode_time_sec=30, + task_description="My task for Brewie robot", + resume_existing_dataset=False # True to continue recording in existing dataset +) +''' \ No newline at end of file diff --git a/examples/brewie/record_with_config.py b/examples/brewie/record_with_config.py new file mode 100644 index 00000000000..d81946974c8 --- /dev/null +++ b/examples/brewie/record_with_config.py @@ -0,0 +1,426 @@ +#!/usr/bin/env python + +""" +Example of recording a dataset for Brewie robot using configuration file. + +This script uses record_config.py to configure all recording parameters. +Modify settings in record_config.py before running this script. + +FEATURES: +- Automatic detection of existing datasets +- Ability to continue recording in existing dataset (adding new episodes) +- Interactive recording mode selection +- Support for all standard LeRobot functions +- Secure HuggingFace token retrieval from environment variables + +Usage: + # With environment variable + export HUGGINGFACE_TOKEN=your_token_here + python examples/brewie/record_with_config.py + + # With command line argument + python examples/brewie/record_with_config.py --hf-token your_token_here + + # Interactive input (token will be requested at startup) + python examples/brewie/record_with_config.py + +Operating modes: +1. Creating new dataset (default) +2. Continuing recording in existing dataset (automatically offered when detected) +3. Forced continuation of recording (via resume_existing_dataset=True in config) +""" + +import os +import shutil +import sys +from pathlib import Path + +# Add path to lerobot modules +sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src")) + +from lerobot.datasets.lerobot_dataset import LeRobotDataset +from lerobot.datasets.utils import hw_to_dataset_features +from lerobot.processor import make_default_processors +from lerobot.scripts.lerobot_record import record_loop +from lerobot.robots.brewie.config_Brewie import BrewieConfig +from lerobot.robots.brewie.Brewie_base import BrewieBase +from lerobot.teleoperators.keyboard import KeyboardTeleop, KeyboardTeleopConfig +from lerobot.utils.constants import ACTION, OBS_STR +from lerobot.utils.control_utils import init_keyboard_listener +from lerobot.utils.utils import log_say as _log_say +from lerobot.utils.visualization_utils import init_rerun + +# Import configuration +from record_config import config + +def log_say(text: str, play_sounds: bool = True, blocking: bool = False): + """ + Wrapper for log_say that duplicates information to console. + + Args: + text: Text to output + play_sounds: Whether to play sound (passed to original log_say) + blocking: Blocking mode (passed to original log_say) + """ + # Output to console + print(f"[LOG] {text}") + + # Call original log_say function + _log_say(text, play_sounds, blocking) + +def check_dataset_exists(dataset_repo_id: str) -> bool: + """Check if dataset exists locally or on Hub.""" + try: + # Attempt to load existing dataset + existing_dataset = LeRobotDataset(dataset_repo_id) + return True + except Exception: + return False + +def get_existing_dataset_info(dataset_repo_id: str) -> dict: + """Get information about existing dataset.""" + try: + existing_dataset = LeRobotDataset(dataset_repo_id) + return { + "exists": True, + "num_episodes": existing_dataset.num_episodes, + "fps": existing_dataset.fps, + "robot_type": existing_dataset.meta.robot_type, + "features": list(existing_dataset.features.keys()) + } + except Exception as e: + return { + "exists": False, + "error": str(e) + } + +def validate_config(): + """Check configuration correctness.""" + errors = [] + + if config.hf_username == "your_username": + errors.append("Need to specify your HuggingFace username in config.hf_username") + + # Check token through new method + try: + hf_token = config.get_hf_token() + if not hf_token or hf_token.strip() == "": + errors.append("Failed to get HuggingFace token") + except ValueError as e: + errors.append(f"Error getting HuggingFace token: {e}") + + if not config.dataset_name: + errors.append("Need to specify dataset name in config.dataset_name") + + if config.num_episodes <= 0: + errors.append("Number of episodes must be greater than 0") + + if config.episode_time_sec <= 0: + errors.append("Episode duration must be greater than 0") + + if errors: + print("Configuration errors:") + for error in errors: + print(f" - {error}") + print("\nFix errors in record_config.py file and run script again.") + print("\nTo set HuggingFace token use:") + print(" export HUGGINGFACE_TOKEN=your_token_here") + print("or pass token via command line argument:") + print(" python record_with_config.py --hf-token your_token_here") + return False + + return True + +def print_config_summary(dataset_repo_id: str, existing_dataset_info: dict = None): + """Print configuration summary.""" + print("=" * 60) + print("BREWIE DATASET RECORDING CONFIGURATION") + print("=" * 60) + print(f"Dataset: {dataset_repo_id}") + + if existing_dataset_info and existing_dataset_info.get("exists"): + print(f"MODE: Continue recording in existing dataset") + print(f"Existing episodes: {existing_dataset_info['num_episodes']}") + print(f"Episodes to add: {config.num_episodes}") + print(f"Total episodes: {existing_dataset_info['num_episodes'] + config.num_episodes}") + else: + print(f"MODE: Create new dataset") + print(f"Episodes: {config.num_episodes}") + + print(f"Task: {config.task_description}") + print(f"Category: {config.task_category}") + print(f"Difficulty: {config.difficulty_level}") + print(f"Episode duration: {config.episode_time_sec}s") + print(f"Reset time: {config.reset_time_sec}s") + print(f"Recording frequency: {config.fps} FPS") + print(f"ROS Master: {config.ros_master_ip}:{config.ros_master_port}") + print("=" * 60) + +def main(): + """Main dataset recording function.""" + + # Check configuration + if not validate_config(): + return + + # Create dataset ID + dataset_repo_id = f"{config.hf_username}/{config.dataset_name}" + + # Check dataset existence + existing_dataset_info = get_existing_dataset_info(dataset_repo_id) + + # Automatic resume mode detection if not explicitly set + should_resume = config.resume_existing_dataset + if existing_dataset_info.get("exists") and not config.resume_existing_dataset: + print(f"\nFound existing dataset: {dataset_repo_id}") + print(f"Existing episodes: {existing_dataset_info['num_episodes']}") + print(f"FPS: {existing_dataset_info['fps']}") + print(f"Robot type: {existing_dataset_info['robot_type']}") + + response = input("\nContinue recording in existing dataset? (y/N): ").strip().lower() + if response in ['y', 'yes']: + should_resume = True + print("Mode: Continue recording in existing dataset") + else: + print("Mode: Create new dataset (existing will be overwritten)") + + # Print configuration summary + if should_resume and existing_dataset_info.get("exists"): + print_config_summary(dataset_repo_id, existing_dataset_info) + else: + print_config_summary(dataset_repo_id) + + # Confirmation to start + response = input("Continue recording with these settings? (y/N): ").strip().lower() + if response not in ['y', 'yes']: + print("Recording cancelled.") + return + + # ============================================================================= + # CREATE CONFIGURATIONS + # ============================================================================= + + # Robot configuration + robot_config = BrewieConfig( + master_ip=config.ros_master_ip, + master_port=config.ros_master_port, + servo_duration=config.servo_duration, + max_relative_target=config.max_relative_target, + ) + + # Teleoperator configuration + keyboard_config = KeyboardTeleopConfig() + + # ============================================================================= + # DEVICE INITIALIZATION + # ============================================================================= + + robot = BrewieBase(robot_config) + keyboard = KeyboardTeleop(keyboard_config) + + # Processors for record_loop + teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors() + + # ============================================================================= + # DATASET SETUP + # ============================================================================= + + # Dataset features configuration + action_features = hw_to_dataset_features(robot.action_features, ACTION) + obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR) + dataset_features = {**action_features, **obs_features} + + # Create or load dataset + if should_resume and existing_dataset_info.get("exists"): + log_say("Loading existing dataset to continue recording...") + dataset = LeRobotDataset( + repo_id=dataset_repo_id, + batch_encoding_size=1, # Use default value + ) + + # Start image writer for existing dataset + if hasattr(robot, "cameras") and len(robot.cameras) > 0: + dataset.start_image_writer( + num_processes=0, # Use default value + num_threads=config.image_writer_threads, + ) + + log_say(f"Dataset loaded. Existing episodes: {dataset.num_episodes}") + else: + log_say("Creating new dataset...") + try: + dataset = LeRobotDataset.create( + repo_id=dataset_repo_id, + fps=config.fps, + features=dataset_features, + robot_type=robot.name, + use_videos=config.use_videos, + image_writer_threads=config.image_writer_threads, + ) + except FileExistsError: + log_say("Dataset folder already exists. Removing existing folder and creating new dataset...") + # Remove existing dataset folder + dataset_root = Path.home() / ".cache" / "huggingface" / "lerobot" / dataset_repo_id + if dataset_root.exists(): + shutil.rmtree(dataset_root) + # Create new dataset + dataset = LeRobotDataset.create( + repo_id=dataset_repo_id, + fps=config.fps, + features=dataset_features, + robot_type=robot.name, + use_videos=config.use_videos, + image_writer_threads=config.image_writer_threads, + ) + + # ============================================================================= + # DEVICE CONNECTION + # ============================================================================= + + log_say("Connecting to Brewie robot...") + try: + robot.connect() + except Exception as e: + log_say(f"Error connecting to robot: {e}") + return + + log_say("Connecting to teleoperator...") + try: + keyboard.connect() + except Exception as e: + log_say(f"Error connecting to teleoperator: {e}") + robot.disconnect() + return + + # Initialize visualization + init_rerun(session_name=config.session_name) + + # Initialize keyboard listener + listener, events = init_keyboard_listener() + + # Check connections + if not robot.is_connected: + log_say("ERROR: Brewie robot not connected!") + keyboard.disconnect() + listener.stop() + return + + if not keyboard.is_connected: + log_say("ERROR: Teleoperator not connected!") + robot.disconnect() + listener.stop() + return + + log_say("All devices connected successfully!") + log_say("Controls:") + log_say(" - ENTER: Start/continue episode recording") + log_say(" - ESC: Stop recording") + log_say(" - R: Rewrite current episode") + + # ============================================================================= + # EPISODE RECORDING LOOP + # ============================================================================= + + # Determine starting episode number + start_episode = dataset.num_episodes if should_resume else 0 + total_episodes_to_record = config.num_episodes + recorded_episodes = 0 + + log_say(f"Starting recording. Will record {total_episodes_to_record} episodes") + if should_resume: + log_say(f"Continuing from episode {start_episode}") + + try: + while recorded_episodes < total_episodes_to_record and not events["stop_recording"]: + current_episode_num = start_episode + recorded_episodes + 1 + log_say(f"Recording episode {current_episode_num} ({recorded_episodes + 1}/{total_episodes_to_record})") + log_say("Press ENTER to start episode recording...") + input() + + # Start recording loop + record_loop( + robot=robot, + events=events, + fps=config.fps, + dataset=dataset, + teleop=keyboard, + control_time_s=config.episode_time_sec, + single_task=config.task_description, + display_data=config.display_data, + teleop_action_processor=teleop_action_processor, + robot_action_processor=robot_action_processor, + robot_observation_processor=robot_observation_processor, + ) + + # Environment reset logic + if not events["stop_recording"] and ( + (recorded_episodes < total_episodes_to_record - 1) or events["rerecord_episode"] + ): + log_say("Resetting environment...") + record_loop( + robot=robot, + events=events, + fps=config.fps, + teleop=keyboard, + control_time_s=config.reset_time_sec, + single_task="Robot position reset", + display_data=config.display_data, + teleop_action_processor=teleop_action_processor, + robot_action_processor=robot_action_processor, + robot_observation_processor=robot_observation_processor, + ) + + # Handle episode rewrite + if events["rerecord_episode"]: + log_say("Rewriting episode...") + events["rerecord_episode"] = False + events["exit_early"] = False + dataset.clear_episode_buffer() + continue + + # Save episode + dataset.save_episode() + recorded_episodes += 1 + current_episode_num = start_episode + recorded_episodes + log_say(f"Episode {current_episode_num} saved") + + except KeyboardInterrupt: + log_say("Recording interrupted by user") + except Exception as e: + log_say(f"Error during recording: {e}") + + # ============================================================================= + # COMPLETION AND HUB UPLOAD + # ============================================================================= + + total_episodes_in_dataset = dataset.num_episodes + if should_resume: + log_say(f"Recording completed! Added {recorded_episodes} new episodes") + log_say(f"Total episodes in dataset: {total_episodes_in_dataset}") + else: + log_say(f"Recording completed! Saved {recorded_episodes} episodes") + + if recorded_episodes > 0 and config.auto_push_to_hub: + log_say("Uploading dataset to HuggingFace Hub...") + try: + dataset.push_to_hub() + log_say(f"Dataset successfully uploaded to: https://huggingface.co/datasets/{dataset_repo_id}") + except Exception as e: + log_say(f"Error uploading to Hub: {e}") + log_say("Dataset saved locally") + elif recorded_episodes > 0: + log_say("Dataset saved locally (auto_push_to_hub = False)") + + # Disconnect devices + log_say("Disconnecting devices...") + try: + robot.disconnect() + keyboard.disconnect() + listener.stop() + except Exception as e: + log_say(f"Error during disconnection: {e}") + + log_say("Dataset recording completed!") + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index f4fb7d24967..cbfe967babc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -120,6 +120,7 @@ unitree_g1 = [ "casadi>=3.6.0,<4.0.0", ] reachy2 = ["reachy2_sdk>=1.0.15,<1.1.0"] +brewie = ["roslibpy>=1.4.0,<2.0.0"] kinematics = ["lerobot[placo-dep]"] intelrealsense = [ "pyrealsense2>=2.55.1.6486,<2.57.0 ; sys_platform != 'darwin'", diff --git a/scripts/setup_test_env.ps1 b/scripts/setup_test_env.ps1 new file mode 100644 index 00000000000..dfe39438a0f --- /dev/null +++ b/scripts/setup_test_env.ps1 @@ -0,0 +1,40 @@ +# Setup test environment for LeRobot +# Run from repo root: .\scripts\setup_test_env.ps1 + +$ErrorActionPreference = "Continue" # Don't stop on warnings + +Write-Host "=== LeRobot Test Environment Setup ===" -ForegroundColor Cyan + +# 1. Create venv +$venvPath = ".venv-lerobot-test" +if (Test-Path $venvPath) { + Write-Host "Removing existing $venvPath..." -ForegroundColor Yellow + Remove-Item -Recurse -Force $venvPath +} + +Write-Host "Creating virtual environment..." -ForegroundColor Green +python -m venv $venvPath + +# 2. Activate and upgrade pip +& "$venvPath\Scripts\Activate.ps1" +python -m pip install --upgrade pip + +# 3. Install lerobot with test extras (ignore warnings) +Write-Host "Installing lerobot with test dependencies (this may take several minutes)..." -ForegroundColor Green +$env:PYTHONWARNINGS = "ignore" +pip install -e ".[test]" 2>&1 | Out-Host + +if ($LASTEXITCODE -ne 0) { + Write-Host "Install failed. Trying minimal install..." -ForegroundColor Yellow + pip install -e . 2>&1 | Out-Host + pip install pytest pytest-timeout pytest-cov 2>&1 | Out-Host +} + +# 4. Verify +Write-Host "`nVerifying installation..." -ForegroundColor Green +python -c "import lerobot; print('lerobot:', lerobot.__version__)" +python -c "from lerobot.robots.brewie import BrewieConfig, BrewieBase; print('Brewie: OK')" + +Write-Host "`n=== Setup complete ===" -ForegroundColor Cyan +Write-Host "Activate: .\.venv-lerobot-test\Scripts\Activate.ps1" +Write-Host "Run tests: pytest tests -v --maxfail=5" diff --git a/src/lerobot/robots/brewie/Brewie_base.py b/src/lerobot/robots/brewie/Brewie_base.py new file mode 100644 index 00000000000..54de24d8927 --- /dev/null +++ b/src/lerobot/robots/brewie/Brewie_base.py @@ -0,0 +1,895 @@ +#!/usr/bin/env python + +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import time +import threading +import base64 +from functools import cached_property +from typing import Any +import numpy as np +import cv2 +from io import BytesIO + +import roslibpy +from roslibpy import Ros, Service, ServiceRequest, Topic, Message + +from lerobot.cameras.utils import make_cameras_from_configs +from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError + +from ..robot import Robot +from ..utils import ensure_safe_goal_position +from .config_Brewie import BrewieConfig + +logger = logging.getLogger(__name__) + + +class JoystickSubscriber: + """ + Class for subscribing to ROS topic /joy and receiving joystick data. + """ + + def __init__(self, ros_client, joy_topic): + """ + Initialize joystick subscriber. + + Args: + ros_client: ROS client for connection + joy_topic: ROS topic with joystick data + """ + self.last_joy_data = None + self.client = ros_client + self.joy_topic = joy_topic + self.joy_lock = threading.Lock() + self.last_message = None + + def on_joy_received(self, message): + """ + Callback that is called when a new joystick message is received. + + Args: + message: ROS message with joystick data + """ + try: + with self.joy_lock: + self.last_message = message + # Extract joystick data + joy_data = self._extract_joy_data(message) + if joy_data is not None: + self.last_joy_data = joy_data + else: + logger.warning("[JoystickSubscriber] Failed to extract joystick data from message") + except Exception as e: + logger.error(f"[JoystickSubscriber] Error in on_joy_received: {e}") + + def _extract_joy_data(self, message): + """ + Extract joystick data from ROS message. + + Args: + message: ROS message with joystick data + + Returns: + dict or None: Joystick data or None on error + """ + try: + if message is None: + logger.warning("[JoystickSubscriber] Message is None") + return None + + # Extract axes and buttons + axes = message.get('axes', []) + buttons = message.get('buttons', []) + + if not axes and not buttons: + logger.warning("[JoystickSubscriber] No axes or buttons data in message") + return None + + # Create structured data + joy_data = { + 'axes': axes, + 'buttons': buttons, + 'timestamp': time.time() + } + + return joy_data + + except Exception as e: + logger.error(f"[JoystickSubscriber] Error extracting joystick data: {e}") + return None + + def get_last_joy_data(self): + """ + Method that returns the latest joystick data. + + Returns: + dict or None: Latest joystick data or None if no data available + """ + with self.joy_lock: + if self.last_joy_data is None: + logger.debug("[JoystickSubscriber] No joystick data received") + return None + return self.last_joy_data.copy() + + def get_last_message(self): + """ + Returns the last received message. + + Returns: + dict or None: Last ROS message or None + """ + with self.joy_lock: + return self.last_message + + def subscribe(self): + """Subscribe to joystick data topic.""" + try: + self.joy_topic.subscribe(self.on_joy_received) + logger.info("[JoystickSubscriber] Successfully subscribed to joystick topic") + except Exception as e: + logger.error(f"[JoystickSubscriber] Failed to subscribe to joystick topic: {e}") + raise + + +class IMUSubscriber: + """ + Class for subscribing to ROS topic /imu and receiving IMU data. + """ + + def __init__(self, ros_client, imu_topic): + """ + Initialize IMU subscriber. + + Args: + ros_client: ROS client for connection + imu_topic: ROS topic with IMU data + """ + self.last_imu_data = None + self.client = ros_client + self.imu_topic = imu_topic + self.imu_lock = threading.Lock() + self.last_message = None + + def on_imu_received(self, message): + """ + Callback that is called when a new IMU message is received. + + Args: + message: ROS message with IMU data + """ + try: + with self.imu_lock: + self.last_message = message + # Extract IMU data + imu_data = self._extract_imu_data(message) + if imu_data is not None: + self.last_imu_data = imu_data + else: + logger.warning("[IMUSubscriber] Failed to extract IMU data from message") + except Exception as e: + logger.error(f"[IMUSubscriber] Error in on_imu_received: {e}") + + def _extract_imu_data(self, message): + """ + Extract IMU data from ROS message. + + Args: + message: ROS message with IMU data + + Returns: + dict or None: IMU data or None on error + """ + try: + if message is None: + logger.warning("[IMUSubscriber] Message is None") + return None + + # Extract orientation + orientation = message.get('orientation', {}) + orientation_data = { + 'x': orientation.get('x', 0.0), + 'y': orientation.get('y', 0.0), + 'z': orientation.get('z', 0.0), + 'w': orientation.get('w', 0.0) + } + + # Extract angular velocity + angular_velocity = message.get('angular_velocity', {}) + angular_velocity_data = { + 'x': angular_velocity.get('x', 0.0), + 'y': angular_velocity.get('y', 0.0), + 'z': angular_velocity.get('z', 0.0) + } + + # Extract linear acceleration + linear_acceleration = message.get('linear_acceleration', {}) + linear_acceleration_data = { + 'x': linear_acceleration.get('x', 0.0), + 'y': linear_acceleration.get('y', 0.0), + 'z': linear_acceleration.get('z', 0.0) + } + + # Create structured data + imu_data = { + 'orientation': orientation_data, + 'angular_velocity': angular_velocity_data, + 'linear_acceleration': linear_acceleration_data, + 'timestamp': time.time() + } + + return imu_data + + except Exception as e: + logger.error(f"[IMUSubscriber] Error extracting IMU data: {e}") + return None + + def get_last_imu_data(self): + """ + Method that returns the latest IMU data. + + Returns: + dict or None: Latest IMU data or None if no data available + """ + with self.imu_lock: + if self.last_imu_data is None: + logger.debug("[IMUSubscriber] No IMU data received") + return None + return self.last_imu_data.copy() + + def get_last_message(self): + """ + Returns the last received message. + + Returns: + dict or None: Last ROS message or None + """ + with self.imu_lock: + return self.last_message + + def subscribe(self): + """Subscribe to IMU data topic.""" + try: + self.imu_topic.subscribe(self.on_imu_received) + logger.info("[IMUSubscriber] Successfully subscribed to IMU topic") + except Exception as e: + logger.error(f"[IMUSubscriber] Failed to subscribe to IMU topic: {e}") + raise + + +class CameraSubscriber: + """ + Class for subscribing to ROS topic with images and receiving the latest snapshot. + Based on user example for reliable image processing. + """ + + def __init__(self, ros_client, image_topic): + """ + Initialize image subscriber. + + Args: + ros_client: ROS client for connection + image_topic: ROS topic with images + """ + self.last_image = None + self.client = ros_client + self.image_topic = image_topic + self.image_lock = threading.Lock() + self.last_message = None + + def on_image_received(self, message): + """ + Callback that is called when a new message is received. + + Args: + message: ROS message with image + """ + try: + with self.image_lock: + self.last_message = message + # Decode image immediately upon receipt + decoded_image = self._decode_image_from_message(message) + if decoded_image is not None: + self.last_image = decoded_image + else: + logger.warning("[CameraSubscriber] Failed to decode image from message") + except Exception as e: + logger.error(f"[CameraSubscriber] Error in on_image_received: {e}") + + def _decode_image_from_message(self, message): + """ + Decode image from ROS message. + + Args: + message: ROS message with image + + Returns: + np.ndarray or None: Decoded image or None on error + """ + try: + if message is None: + logger.warning("[CameraSubscriber] Message is None") + return None + + # Get image data + img_data = message.get('data') + if img_data is None: + logger.warning("[CameraSubscriber] No 'data' field in message") + return None + + # Handle different data formats + if isinstance(img_data, str): + # If data is a string, try to decode as Base64 + try: + image_bytes = base64.b64decode(img_data) + except Exception as e: + logger.warning(f"[CameraSubscriber] Failed to decode Base64 string: {e}") + # If not Base64, try as regular string + image_bytes = img_data.encode('latin-1') + else: + # If data is already in bytes format + image_bytes = img_data + + # Convert byte array to NumPy array + img_np = np.frombuffer(image_bytes, np.uint8) + + # Decode image from JPEG/PNG using OpenCV + img_cv = cv2.imdecode(img_np, cv2.IMREAD_UNCHANGED) + + if img_cv is None: + logger.warning("[CameraSubscriber] Failed to decode image with OpenCV") + return None + + return img_cv + + except Exception as e: + logger.error(f"[CameraSubscriber] Error decoding image: {e}") + return None + + def get_last_image(self): + """ + Method that returns the last saved snapshot. + + Returns: + np.ndarray or None: Last image or None if no data available + """ + with self.image_lock: + if self.last_image is None: + logger.debug("[CameraSubscriber] No image data received") + return None + return self.last_image.copy() + + def get_last_message(self): + """ + Returns the last received message. + + Returns: + dict or None: Last ROS message or None + """ + with self.image_lock: + return self.last_message + + def subscribe(self): + """Subscribe to image topic.""" + try: + self.image_topic.subscribe(self.on_image_received) + logger.info("[CameraSubscriber] Successfully subscribed to image topic") + except Exception as e: + logger.error(f"[CameraSubscriber] Failed to subscribe to image topic: {e}") + raise + + +class BrewieBase(Robot): + + config_class = BrewieConfig + name = "BrewieBase" + + def __init__(self, config: BrewieConfig): + super().__init__(config) + self.config = config + + # ROS connection + self.ros_client = None + self.position_service = None + self.set_position_topic = None + self.camera_topic = None + self.joy_topic = None + self.imu_topic = None + + # Camera data + self.latest_image = None + self.image_lock = threading.Lock() + self.camera_subscriber = None + + # Additional sensor data + self.joystick_subscriber = None + self.imu_subscriber = None + + # Servo positions cache + self.current_positions = {} + self.position_lock = threading.Lock() + + # Create reverse mapping (joint_name -> servo_id) + self.joint_to_id = {v: k for k, v in config.servo_mapping.items()} + + # Initialize cameras if configured + self.cameras = make_cameras_from_configs(config.cameras) if config.cameras else {} + + @property + def _motors_ft(self) -> dict[str, type]: + # Return features for all joints defined in servo_mapping + return {f"{joint}.pos": float for joint in self.config.servo_mapping.values()} + + @property + def _cameras_ft(self) -> dict[str, tuple]: + if self.cameras: + return { + cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras + } + else: + # Default camera dimensions for ROS camera + return {"camera": (480, 640, 3)} + + @property + def _joystick_ft(self) -> dict[str, type]: + """Joystick features: individual axes and buttons.""" + # Create separate features for each axis and button + features = {} + # 8 joystick axes + for i in range(8): + features[f"joystick.axis_{i}"] = float + # 15 joystick buttons (also float for LeRobot compatibility) + for i in range(15): + features[f"joystick.button_{i}"] = float + return features + + @property + def _imu_ft(self) -> dict[str, type]: + """IMU features: orientation, angular velocity, linear acceleration.""" + return { + "imu.orientation.x": float, + "imu.orientation.y": float, + "imu.orientation.z": float, + "imu.orientation.w": float, + "imu.angular_velocity.x": float, + "imu.angular_velocity.y": float, + "imu.angular_velocity.z": float, + "imu.linear_acceleration.x": float, + "imu.linear_acceleration.y": float, + "imu.linear_acceleration.z": float, + } + + @cached_property + def observation_features(self) -> dict[str, type | tuple]: + return {**self._motors_ft, **self._cameras_ft, **self._joystick_ft, **self._imu_ft} + + @cached_property + def action_features(self) -> dict[str, type]: + return self._motors_ft + + @property + def is_connected(self) -> bool: + ros_connected = self.ros_client is not None and self.ros_client.is_connected + cameras_connected = all(cam.is_connected for cam in self.cameras.values()) if self.cameras else True + return ros_connected and cameras_connected + + def connect(self, calibrate: bool = True) -> None: + """ + Connect to ROS and initialize services and topics. + """ + if self.is_connected: + raise DeviceAlreadyConnectedError(f"{self} already connected") + + # Connect to ROS + self.ros_client = Ros(host=self.config.master_ip, port=self.config.master_port) + self.ros_client.run() + + # Wait for connection + timeout = 10 + start_time = time.time() + while not self.ros_client.is_connected and (time.time() - start_time) < timeout: + time.sleep(0.1) + + if not self.ros_client.is_connected: + raise ConnectionError(f"Failed to connect to ROS at {self.config.master_ip}:{self.config.master_port}") + + # Initialize services and topics + self._setup_ros_services() + self._setup_ros_topics() + + # Initialize default positions for all servos + self._initialize_default_positions() + + # Connect cameras if configured + for cam in self.cameras.values(): + cam.connect() + + logger.info(f"{self} connected to ROS at {self.config.master_ip}:{self.config.master_port}") + + def _setup_ros_services(self): + """Setup ROS services for servo control.""" + self.position_service = Service( + self.ros_client, + self.config.position_service, + 'ros_robot_controller/GetBusServosPosition' + ) + + def _setup_ros_topics(self): + """Setup ROS topics for servo control, camera, joystick and IMU.""" + self.set_position_topic = Topic( + self.ros_client, + self.config.set_position_topic, + 'ros_robot_controller/SetBusServosPosition' + ) + + # Setup camera topic if not using standard cameras + if not self.cameras: + self.camera_topic = Topic( + self.ros_client, + self.config.camera_topic, + 'sensor_msgs/CompressedImage' + ) + # Create CameraSubscriber for reliable image processing + self.camera_subscriber = CameraSubscriber(self.ros_client, self.camera_topic) + self.camera_subscriber.subscribe() + + # Setup joystick topic + self.joy_topic = Topic( + self.ros_client, + self.config.joy_topic, + 'sensor_msgs/Joy' + ) + self.joystick_subscriber = JoystickSubscriber(self.ros_client, self.joy_topic) + self.joystick_subscriber.subscribe() + + # Setup IMU topic + self.imu_topic = Topic( + self.ros_client, + self.config.imu_topic, + 'sensor_msgs/Imu' + ) + self.imu_subscriber = IMUSubscriber(self.ros_client, self.imu_topic) + self.imu_subscriber.subscribe() + + def _initialize_default_positions(self): + """Initialize default positions for all servos.""" + with self.position_lock: + for servo_id in self.config.servo_ids: + joint_name = self.config.servo_mapping.get(servo_id) + if joint_name: + joint_key = f"{joint_name}.pos" + if joint_key not in self.current_positions: + # Use middle position as default (500 out of 0-1000 range) + self.current_positions[joint_key] = 500.0 + logger.debug(f"Initialized default position for {joint_name}: 500.0") + + + @property + def is_calibrated(self) -> bool: + # For ROS-based control, we assume the robot is always "calibrated" + # since calibration is handled by the ROS controller + return True + + def calibrate(self) -> None: + """ + For ROS-based control, calibration is handled by the ROS controller. + This method is kept for compatibility but does nothing. + """ + logger.info("Calibration is handled by the ROS controller. Skipping local calibration.") + + def configure(self) -> None: + """ + For ROS-based control, configuration is handled by the ROS controller. + This method is kept for compatibility but does nothing. + """ + logger.info("Configuration is handled by the ROS controller. Skipping local configuration.") + + def setup_motors(self) -> None: + """ + For ROS-based control, motor setup is handled by the ROS controller. + This method is kept for compatibility but does nothing. + """ + logger.info("Motor setup is handled by the ROS controller. Skipping local setup.") + + def get_observation(self) -> dict[str, Any]: + if not self.is_connected: + raise DeviceNotConnectedError(f"{self} is not connected.") + + obs_dict = {} + total_start = time.perf_counter() + + # Read servo positions via ROS service + start = time.perf_counter() + try: + request = ServiceRequest({'id': self.config.servo_ids}) + result = self.position_service.call(request) + if result.get('success', False): + positions = result.get('position', []) + received_servo_ids = set() + + # Process received positions + for pos_data in positions: + servo_id = pos_data['id'] + position = pos_data['position'] + joint_name = self.config.servo_mapping.get(servo_id) + if joint_name: + # Ensure position is float32 and within valid range (0-1000) + position = float(max(0.0, min(1000.0, float(position)))) + obs_dict[f"{joint_name}.pos"] = position + received_servo_ids.add(servo_id) + + # Fill missing servo positions with previous values + with self.position_lock: + for servo_id in self.config.servo_ids: + if servo_id not in received_servo_ids: + joint_name = self.config.servo_mapping.get(servo_id) + if joint_name: + joint_key = f"{joint_name}.pos" + if joint_key in self.current_positions: + obs_dict[joint_key] = self.current_positions[joint_key] + logger.debug(f"Using previous position for missing servo {servo_id} ({joint_name}): {self.current_positions[joint_key]}") + else: + # If no previous value available, use default position (middle of range) + default_position = 500.0 # Middle of 0-1000 range + obs_dict[joint_key] = default_position + logger.warning(f"No previous position for servo {servo_id} ({joint_name}), using default: {default_position}") + + # Update current positions cache + self.current_positions = obs_dict.copy() + else: + logger.warning("Failed to get servo positions from ROS service") + # Use cached positions if available + with self.position_lock: + obs_dict = self.current_positions.copy() + + except Exception as e: + logger.error(f"Error reading servo positions: {e}") + # Use cached positions if available + with self.position_lock: + obs_dict = self.current_positions.copy() + + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read servo state: {dt_ms:.1f}ms") + + # Get camera data + if self.cameras: + # Use standard cameras + for cam_key, cam in self.cameras.items(): + start = time.perf_counter() + obs_dict[cam_key] = cam.async_read() + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") + else: + # Use ROS camera with improved CameraSubscriber + start = time.perf_counter() + if self.camera_subscriber is not None: + received_img = self.camera_subscriber.get_last_image() + if received_img is None: + logger.warning("[Image] No data received from subscriber") + # Return empty image if no data available + obs_dict["camera"] = np.zeros((480, 640, 3), dtype=np.uint8) + else: + obs_dict["camera"] = received_img + logger.debug(f"[Image] Successfully received image with shape: {received_img.shape}") + else: + logger.warning("[Image] Camera subscriber not initialized") + obs_dict["camera"] = np.zeros((480, 640, 3), dtype=np.uint8) + + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read camera: {dt_ms:.1f}ms") + + # Get joystick data + start = time.perf_counter() + if self.joystick_subscriber is not None: + joy_data = self.joystick_subscriber.get_last_joy_data() + if joy_data is not None: + # Convert joystick data to separate features + axes = joy_data['axes'] + buttons = joy_data['buttons'] + + # Add axis data (up to 8 axes) + for i in range(8): + if i < len(axes): + obs_dict[f"joystick.axis_{i}"] = float(axes[i]) + else: + obs_dict[f"joystick.axis_{i}"] = 0.0 + + # Add button data (up to 15 buttons) + for i in range(15): + if i < len(buttons): + obs_dict[f"joystick.button_{i}"] = float(buttons[i]) + else: + obs_dict[f"joystick.button_{i}"] = 0.0 + + logger.debug(f"[Joystick] Successfully received data: {len(axes)} axes, {len(buttons)} buttons") + else: + # Return zero values if no data available + for i in range(8): + obs_dict[f"joystick.axis_{i}"] = 0.0 + for i in range(15): + obs_dict[f"joystick.button_{i}"] = 0.0 + logger.debug("[Joystick] No data received, using default values") + else: + logger.warning("[Joystick] Joystick subscriber not initialized") + for i in range(8): + obs_dict[f"joystick.axis_{i}"] = 0.0 + for i in range(15): + obs_dict[f"joystick.button_{i}"] = 0.0 + + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read joystick: {dt_ms:.1f}ms") + + # Get IMU data + start = time.perf_counter() + if self.imu_subscriber is not None: + imu_data = self.imu_subscriber.get_last_imu_data() + if imu_data is not None: + # Extract IMU data + orientation = imu_data['orientation'] + angular_velocity = imu_data['angular_velocity'] + linear_acceleration = imu_data['linear_acceleration'] + + obs_dict["imu.orientation.x"] = float(orientation['x']) + obs_dict["imu.orientation.y"] = float(orientation['y']) + obs_dict["imu.orientation.z"] = float(orientation['z']) + obs_dict["imu.orientation.w"] = float(orientation['w']) + + obs_dict["imu.angular_velocity.x"] = float(angular_velocity['x']) + obs_dict["imu.angular_velocity.y"] = float(angular_velocity['y']) + obs_dict["imu.angular_velocity.z"] = float(angular_velocity['z']) + + obs_dict["imu.linear_acceleration.x"] = float(linear_acceleration['x']) + obs_dict["imu.linear_acceleration.y"] = float(linear_acceleration['y']) + obs_dict["imu.linear_acceleration.z"] = float(linear_acceleration['z']) + + logger.debug(f"[IMU] Successfully received data: orientation=({orientation['x']:.3f}, {orientation['y']:.3f}, {orientation['z']:.3f}, {orientation['w']:.3f})") + else: + # Return zero values if no data available + obs_dict["imu.orientation.x"] = 0.0 + obs_dict["imu.orientation.y"] = 0.0 + obs_dict["imu.orientation.z"] = 0.0 + obs_dict["imu.orientation.w"] = 0.0 + obs_dict["imu.angular_velocity.x"] = 0.0 + obs_dict["imu.angular_velocity.y"] = 0.0 + obs_dict["imu.angular_velocity.z"] = 0.0 + obs_dict["imu.linear_acceleration.x"] = 0.0 + obs_dict["imu.linear_acceleration.y"] = 0.0 + obs_dict["imu.linear_acceleration.z"] = 0.0 + logger.debug("[IMU] No data received, using default values") + else: + logger.warning("[IMU] IMU subscriber not initialized") + obs_dict["imu.orientation.x"] = 0.0 + obs_dict["imu.orientation.y"] = 0.0 + obs_dict["imu.orientation.z"] = 0.0 + obs_dict["imu.orientation.w"] = 0.0 + obs_dict["imu.angular_velocity.x"] = 0.0 + obs_dict["imu.angular_velocity.y"] = 0.0 + obs_dict["imu.angular_velocity.z"] = 0.0 + obs_dict["imu.linear_acceleration.x"] = 0.0 + obs_dict["imu.linear_acceleration.y"] = 0.0 + obs_dict["imu.linear_acceleration.z"] = 0.0 + + dt_ms = (time.perf_counter() - start) * 1e3 + logger.debug(f"{self} read IMU: {dt_ms:.1f}ms") + + # Log total observation time + total_dt_ms = (time.perf_counter() - total_start) * 1e3 + logger.info(f"{self} TOTAL observation time: {total_dt_ms:.1f}ms") + + return obs_dict + + def send_action(self, action: dict[str, Any]) -> dict[str, Any]: + """Command arm to move to a target joint configuration via ROS. + + The relative action magnitude may be clipped depending on the configuration parameter + `max_relative_target`. In this case, the action sent differs from original action. + Thus, this function always returns the action actually sent. + + Raises: + RobotDeviceNotConnectedError: if robot is not connected. + + Returns: + the action sent to the servos, potentially clipped. + """ + if not self.is_connected: + raise DeviceNotConnectedError(f"{self} is not connected.") + + goal_pos = {key.removesuffix(".pos"): val for key, val in action.items() if key.endswith(".pos")} + + # Cap goal position when too far away from present position. + if self.config.max_relative_target is not None: + with self.position_lock: + present_pos = {key.removesuffix(".pos"): val for key, val in self.current_positions.items()} + goal_present_pos = {key: (g_pos, present_pos.get(key, 0)) for key, g_pos in goal_pos.items()} + # Convert max_relative_target to float if it's not None + max_relative = float(self.config.max_relative_target) + goal_pos = ensure_safe_goal_position(goal_present_pos, max_relative) + + # Convert joint positions to servo positions and send via ROS + try: + new_positions = [] + for joint_name, position in goal_pos.items(): + servo_id = self.joint_to_id.get(joint_name) + if servo_id is not None: + # Ensure position is within valid range (0-1000) + position = float(max(0.0, min(1000.0, float(position)))) + new_positions.append({'id': servo_id, 'position': position}) + + if new_positions: + servo_msg = Message({ + 'duration': self.config.servo_duration, + 'position': new_positions, + }) + self.set_position_topic.publish(servo_msg) + + except Exception as e: + logger.error(f"Error sending servo positions: {e}") + + return {f"{joint}.pos": val for joint, val in goal_pos.items()} + + def test_camera_connection(self) -> dict[str, Any]: + """ + Test camera connection and return status information. + + Returns: + dict: Camera status information + """ + result = { + "camera_available": False, + "last_image_shape": None, + "last_message_received": False, + "error_message": None + } + + try: + if self.camera_subscriber is not None: + # Check last message + last_msg = self.camera_subscriber.get_last_message() + result["last_message_received"] = last_msg is not None + + # Check last image + last_img = self.camera_subscriber.get_last_image() + if last_img is not None: + result["camera_available"] = True + result["last_image_shape"] = last_img.shape + logger.info(f"[CameraTest] Camera working, image shape: {last_img.shape}") + else: + result["error_message"] = "No image data received" + logger.warning("[CameraTest] No image data available") + else: + result["error_message"] = "Camera subscriber not initialized" + logger.warning("[CameraTest] Camera subscriber not initialized") + + except Exception as e: + result["error_message"] = str(e) + logger.error(f"[CameraTest] Error testing camera: {e}") + + return result + + def disconnect(self): + if not self.is_connected: + raise DeviceNotConnectedError(f"{self} is not connected.") + + # Disconnect cameras + for cam in self.cameras.values(): + cam.disconnect() + + # Disconnect ROS + if self.ros_client and self.ros_client.is_connected: + self.ros_client.close() + self.ros_client = None + + logger.info(f"{self} disconnected from ROS.") + diff --git a/src/lerobot/robots/brewie/__init__.py b/src/lerobot/robots/brewie/__init__.py new file mode 100644 index 00000000000..9ccf3ce9d76 --- /dev/null +++ b/src/lerobot/robots/brewie/__init__.py @@ -0,0 +1,18 @@ +#!/usr/bin/env python + +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .config_Brewie import BrewieConfig +from .Brewie_base import BrewieBase diff --git a/src/lerobot/robots/brewie/brewie.mdx b/src/lerobot/robots/brewie/brewie.mdx new file mode 100644 index 00000000000..83e3542beb9 --- /dev/null +++ b/src/lerobot/robots/brewie/brewie.mdx @@ -0,0 +1 @@ +../../../../docs/source/Brewie.mdx \ No newline at end of file diff --git a/src/lerobot/robots/brewie/config_Brewie.py b/src/lerobot/robots/brewie/config_Brewie.py new file mode 100644 index 00000000000..b79f6aeec84 --- /dev/null +++ b/src/lerobot/robots/brewie/config_Brewie.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python + +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field + +from lerobot.cameras import CameraConfig + +from ..config import RobotConfig + + +@RobotConfig.register_subclass("BrewieConfig") +@dataclass +class BrewieConfig(RobotConfig): + + # ROS connection parameters + master_ip: str = "192.168.20.21" + master_port: int = 9090 + + # Servo configuration + servo_ids: list[int] = field(default_factory=lambda: [13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]) + servo_duration: float = 0.1 # Duration for servo movements + + # Servo mapping (ID -> joint name) + servo_mapping: dict[int, str] = field(default_factory=lambda: { + 13: "left_shoulder_pan", + 14: "right_shoulder_pan", + 15: "left_shoulder_lift", + 16: "right_shoulder_lift", + 17: "left_forearm_roll", + 18: "right_forearm_roll", + 19: "left_forearm_pitch", + 20: "right_forearm_pitch", + 21: "left_gripper", + 22: "right_gripper", + 23: "head_pan", # head rotation + 24: "head_tilt" # head tilt + }) + + # ROS topics and services + position_service: str = "/ros_robot_controller/bus_servo/get_position" + set_position_topic: str = "/ros_robot_controller/bus_servo/set_position" + camera_topic: str = "/camera/image_raw/compressed" + + # Additional sensor topics + joy_topic: str = "/joy" # Joystick data + imu_topic: str = "/imu" # IMU data (gyroscope, accelerometer) + + # Safety parameters + max_relative_target: float | None = None # Maximum relative movement per step + disable_torque_on_disconnect: bool = True # Disable torque when disconnecting + + # cameras Will use? + cameras: dict[str, CameraConfig] = field(default_factory=dict) \ No newline at end of file diff --git a/src/lerobot/robots/utils.py b/src/lerobot/robots/utils.py index 92da597f13c..9b448be5f7b 100644 --- a/src/lerobot/robots/utils.py +++ b/src/lerobot/robots/utils.py @@ -68,6 +68,10 @@ def make_robot_from_config(config: RobotConfig) -> Robot: from .bi_openarm_follower import BiOpenArmFollower return BiOpenArmFollower(config) + elif config.type == "BrewieConfig": + from .brewie import BrewieBase + + return BrewieBase(config) elif config.type == "mock_robot": from tests.mocks.mock_robot import MockRobot diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index 72708ba2356..65d2138b737 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -104,6 +104,7 @@ Robot, RobotConfig, bi_openarm_follower, + brewie, bi_so_follower, earthrover_mini_plus, hope_jr, diff --git a/src/lerobot/scripts/lerobot_teleoperate.py b/src/lerobot/scripts/lerobot_teleoperate.py index dad479b2e77..9a231a104da 100644 --- a/src/lerobot/scripts/lerobot_teleoperate.py +++ b/src/lerobot/scripts/lerobot_teleoperate.py @@ -71,6 +71,7 @@ Robot, RobotConfig, bi_openarm_follower, + brewie, bi_so_follower, earthrover_mini_plus, hope_jr,