├── scripts ├── __init__.py ├── backend.py ├── frontend.py └── fullstack.py ├── app ├── __init__.py ├── replaying.py ├── config.py ├── teleoperating.py ├── training.py ├── main.py ├── calibrating.py └── recording.py ├── pyproject.toml ├── .gitignore └── README.md /scripts/__init__.py: -------------------------------------------------------------------------------- 1 | # Scripts package for LeLab commands 2 | -------------------------------------------------------------------------------- /app/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | LeLab - A FastAPI application for audio recording and calibration. 3 | """ 4 | 5 | __version__ = "0.1.0" 6 | -------------------------------------------------------------------------------- /scripts/backend.py: -------------------------------------------------------------------------------- 1 | """ 2 | Backend-only script for LeLab 3 | Runs just the FastAPI server with uvicorn 4 | """ 5 | 6 | import logging 7 | import uvicorn 8 | 9 | # Set up logging 10 | logging.basicConfig(level=logging.INFO) 11 | logger = logging.getLogger(__name__) 12 | 13 | 14 | def main(): 15 | """Start the FastAPI backend server only""" 16 | logger.info("🚀 Starting LeLab FastAPI backend server...") 17 | uvicorn.run( 18 | "app.main:app", host="0.0.0.0", port=8000, reload=True, log_level="info" 19 | ) 20 | 21 | 22 | if __name__ == "__main__": 23 | main() 24 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools>=45", "wheel", "setuptools_scm[toml]>=6.2"] 3 | build-backend = "setuptools.build_meta" 4 | 5 | [project] 6 | name = "LeLab" 7 | version = "0.1.0" 8 | description = "LeRobot Lab - A web interface for robotics with lerobot integration" 9 | readme = "README.md" 10 | requires-python = ">=3.10" 11 | dependencies = [ 12 | "fastapi[standard]>=0.115.12", 13 | "websockets>=15.0.1", 14 | "uvicorn>=0.24.0", 15 | "lerobot @ git+https://github.com/huggingface/lerobot.git", 16 | ] 17 | 18 | [project.scripts] 19 | lelab = "scripts.backend:main" 20 | "lelab-fullstack" = "scripts.fullstack:main" 21 | "lelab-frontend" = "scripts.frontend:main" 22 | 23 | [tool.setuptools.packages.find] 24 | where = ["."] 25 | include = ["app*", "scripts*"] 26 | 27 | [tool.setuptools.package-data] 28 | app = ["static/**/*"] -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Python-generated files 2 | __pycache__/ 3 | *.py[oc] 4 | build/ 5 | dist/ 6 | wheels/ 7 | *.egg-info 8 | 9 | # Virtual environments 10 | .venv 11 | venv/ 12 | env/ 13 | ENV/ 14 | 15 | # IDE and Editor files 16 | .vscode/ 17 | .idea/ 18 | *.swp 19 | *.swo 20 | *~ 21 | 22 | # System files 23 | .DS_Store 24 | .DS_Store? 25 | ._* 26 | .Spotlight-V100 27 | .Trashes 28 | ehthumbs.db 29 | Thumbs.db 30 | 31 | # Logs 32 | *.log 33 | logs/ 34 | 35 | # Environment variables 36 | .env 37 | .env.local 38 | .env.*.local 39 | 40 | # Testing 41 | .coverage 42 | .pytest_cache/ 43 | .tox/ 44 | htmlcov/ 45 | 46 | # Frontend build outputs (if built to root) 47 | frontend/dist/ 48 | frontend/build/ 49 | frontend/.next/ 50 | frontend/out/ 51 | 52 | # Node.js (in case frontend files leak to root) 53 | node_modules/ 54 | npm-debug.log* 55 | yarn-debug.log* 56 | yarn-error.log* 57 | package-lock.json 58 | 59 | # Temporary files 60 | *.tmp 61 | *.temp 62 | 63 | # Database files 64 | *.db 65 | *.sqlite 66 | *.sqlite3 67 | -------------------------------------------------------------------------------- /scripts/frontend.py: -------------------------------------------------------------------------------- 1 | """ 2 | Frontend-only script for LeLab 3 | Clones and runs the frontend development server 4 | """ 5 | 6 | import os 7 | import subprocess 8 | import logging 9 | import webbrowser 10 | import time 11 | from pathlib import Path 12 | 13 | # Set up logging 14 | logging.basicConfig(level=logging.INFO) 15 | logger = logging.getLogger(__name__) 16 | 17 | FRONTEND_REPO_URL = "https://github.com/jurmy24/leLab-space.git" 18 | FRONTEND_DIR_NAME = "leLab-space" 19 | 20 | 21 | def get_frontend_path(): 22 | """Get the path to the frontend directory""" 23 | # Check if frontend exists in parent directory (same level as leLab) 24 | parent_dir = Path(__file__).parent.parent.parent 25 | frontend_path = parent_dir / FRONTEND_DIR_NAME 26 | 27 | if frontend_path.exists(): 28 | logger.info(f"✅ Found existing frontend at: {frontend_path}") 29 | return frontend_path 30 | 31 | return None 32 | 33 | 34 | def clone_frontend(): 35 | """Clone the frontend repository""" 36 | parent_dir = Path(__file__).parent.parent.parent 37 | frontend_path = parent_dir / FRONTEND_DIR_NAME 38 | 39 | logger.info(f"📥 Cloning frontend repository to: {frontend_path}") 40 | 41 | try: 42 | subprocess.run( 43 | ["git", "clone", FRONTEND_REPO_URL, str(frontend_path)], 44 | check=True, 45 | cwd=parent_dir, 46 | ) 47 | logger.info("✅ Frontend repository cloned successfully") 48 | return frontend_path 49 | except subprocess.CalledProcessError as e: 50 | logger.error(f"❌ Failed to clone frontend repository: {e}") 51 | return None 52 | except FileNotFoundError: 53 | logger.error("❌ git not found. Please install git") 54 | return None 55 | 56 | 57 | def install_frontend_deps(frontend_path): 58 | """Install frontend dependencies""" 59 | logger.info("📦 Installing frontend dependencies...") 60 | 61 | try: 62 | subprocess.run(["npm", "install"], check=True, cwd=frontend_path) 63 | logger.info("✅ Frontend dependencies installed successfully") 64 | return True 65 | except subprocess.CalledProcessError as e: 66 | logger.error(f"❌ Failed to install frontend dependencies: {e}") 67 | return False 68 | except FileNotFoundError: 69 | logger.error("❌ npm not found. Please install Node.js and npm") 70 | return False 71 | 72 | 73 | def start_frontend_dev_server(frontend_path): 74 | """Start the frontend development server""" 75 | logger.info("🎨 Starting Vite frontend development server...") 76 | 77 | try: 78 | # Start the dev server 79 | process = subprocess.Popen(["npm", "run", "dev"], cwd=frontend_path) 80 | 81 | # Wait a moment for server to start 82 | time.sleep(3) 83 | 84 | # Auto-open browser 85 | logger.info("🌐 Opening browser...") 86 | webbrowser.open("http://localhost:8080") 87 | 88 | # Wait for the process 89 | process.wait() 90 | 91 | except subprocess.CalledProcessError as e: 92 | logger.error(f"❌ Failed to start frontend server: {e}") 93 | return False 94 | except FileNotFoundError: 95 | logger.error("❌ npm not found. Please install Node.js and npm") 96 | return False 97 | except KeyboardInterrupt: 98 | logger.info("🛑 Frontend server stopped by user") 99 | if process: 100 | process.terminate() 101 | return True 102 | 103 | 104 | def main(): 105 | """Main function to run frontend only""" 106 | logger.info("🎨 Starting LeLab frontend development server...") 107 | 108 | # Get or clone frontend 109 | frontend_path = get_frontend_path() 110 | if not frontend_path: 111 | frontend_path = clone_frontend() 112 | if not frontend_path: 113 | logger.error("❌ Failed to get frontend repository") 114 | return 115 | 116 | # Install dependencies 117 | if not install_frontend_deps(frontend_path): 118 | logger.error("❌ Failed to install frontend dependencies") 119 | return 120 | 121 | # Start dev server 122 | start_frontend_dev_server(frontend_path) 123 | 124 | 125 | if __name__ == "__main__": 126 | main() 127 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # LeLab - Web Interface for LeRobot 2 | 3 | A modern web-based interface for controlling and monitoring robots using the [LeRobot](https://github.com/huggingface/lerobot) framework. This application provides an intuitive dashboard for robot teleoperation, data recording, and calibration management. 4 | 5 | ## 🤖 About 6 | 7 | LeLab bridges the gap between LeRobot's powerful robotics capabilities and user-friendly web interfaces. It offers: 8 | 9 | - **Real-time robot control** through an intuitive web dashboard 10 | - **Dataset recording** for training machine learning models 11 | - **Live teleoperation** with WebSocket-based real-time feedback 12 | - **Configuration management** for leader/follower robot setups 13 | - **Joint position monitoring** and visualization 14 | 15 | ## 🏗️ Architecture 16 | 17 | ``` 18 | ┌─────────────────┐ ┌──────────────────┐ ┌─────────────────┐ 19 | │ Frontend │ │ FastAPI │ │ LeRobot │ 20 | │ (React/TS) │◄──►│ Backend │◄──►│ Framework │ 21 | │ │ │ │ │ │ 22 | │ • Dashboard │ │ • REST APIs │ │ • Robot │ 23 | │ • Controls │ │ • WebSockets │ │ Control │ 24 | │ • Monitoring │ │ • Recording │ │ • Sensors │ 25 | └─────────────────┘ └──────────────────┘ └─────────────────┘ 26 | ``` 27 | 28 | ## ✨ Features 29 | 30 | ### 🎮 Robot Control 31 | 32 | - **Teleoperation**: Direct robot arm control through web interface 33 | - **Joint monitoring**: Real-time joint position feedback via WebSocket 34 | - **Safety controls**: Start/stop teleoperation with status monitoring 35 | 36 | ### 📹 Data Recording 37 | 38 | - **Dataset creation**: Record episodes for training ML models 39 | - **Session management**: Start, stop, and manage recording sessions 40 | - **Episode controls**: Skip to next episode or re-record current one 41 | - **Real-time status**: Monitor recording progress and status 42 | 43 | ### ⚙️ Configuration 44 | 45 | - **Config management**: Handle leader and follower robot configurations 46 | - **Calibration support**: Load and manage calibration settings 47 | - **Health monitoring**: System health checks and diagnostics 48 | 49 | ### 🌐 Web Interface 50 | 51 | - **Modern UI**: Built with React, TypeScript, and Tailwind CSS 52 | - **Real-time updates**: WebSocket integration for live data 53 | - **Responsive design**: Works on desktop and mobile devices 54 | 55 | ## 🚀 Quick Start 56 | 57 | ### Prerequisites 58 | 59 | - Python 3.8+ 60 | - Node.js 16+ (for frontend development) 61 | - LeRobot framework installed and configured 62 | - Compatible robot hardware 63 | 64 | ### Installation 65 | 66 | 1. **Clone the repository** 67 | 68 | ```bash 69 | git clone 70 | cd leLab 71 | ``` 72 | 73 | 2. **Install the Python backend** 74 | 75 | ```bash 76 | # If installing in virtual environment 77 | python -m venv .venv 78 | source .venv/bin/activate 79 | # If installing globally 80 | # Note: Git-LFS required: brew install git-lfs 81 | pip install -e . 82 | ``` 83 | 84 | ### Running the Application 85 | 86 | After installation, you can use the `lelab` command-line tools: 87 | 88 | ```bash 89 | # Start only the backend server (default) 90 | lelab 91 | 92 | # Start both backend and frontend servers 93 | lelab-fullstack 94 | 95 | # Start only the frontend development server 96 | lelab-frontend 97 | ``` 98 | 99 | **Command Options:** 100 | 101 | - `lelab` - Starts only the FastAPI backend server on `http://localhost:8000` 102 | - `lelab-fullstack` - Starts both FastAPI backend (port 8000) and Vite frontend (port 8080) with auto-browser opening 103 | - `lelab-frontend` - Starts only the frontend development server with auto-browser opening 104 | 105 | **Frontend Repository:** 106 | 107 | The frontend is automatically cloned from [leLab-space](https://github.com/jurmy24/leLab-space.git) when you run `lelab-frontend` or `lelab-fullstack`. The system will: 108 | 109 | 1. Check if the frontend already exists in the parent directory 110 | 2. Clone the repository if it doesn't exist 111 | 3. Install dependencies with `npm install` 112 | 4. Start the development server and auto-open your browser 113 | 114 | ### Key Endpoints 115 | 116 | - `POST /move-arm` - Start robot teleoperation 117 | - `POST /stop-teleoperation` - Stop current teleoperation 118 | - `GET /joint-positions` - Get current joint positions 119 | - `POST /start-recording` - Begin dataset recording 120 | - `POST /stop-recording` - End recording session 121 | - `GET /get-configs` - Retrieve available configurations 122 | - `WS /ws/joint-data` - WebSocket for real-time joint data 123 | 124 | ## 🏗️ Project Structure 125 | 126 | ``` 127 | leLab/ 128 | ├── app/ # FastAPI backend 129 | │ ├── main.py # Main FastAPI application 130 | │ ├── recording.py # Dataset recording logic 131 | │ ├── teleoperating.py # Robot teleoperation logic 132 | │ ├── calibrating.py # Robot calibration logic 133 | │ ├── training.py # ML training logic 134 | │ ├── config.py # Configuration management 135 | │ ├── scripts/ # Command-line scripts 136 | │ │ ├── backend.py # Backend-only startup 137 | │ │ ├── frontend.py # Frontend-only startup 138 | │ │ └── fullstack.py # Full-stack startup 139 | │ └── static/ # Static web files 140 | ├── ../leLab-space/ # React frontend (auto-cloned) 141 | │ ├── src/ 142 | │ │ ├── components/ # React components 143 | │ │ ├── pages/ # Page components 144 | │ │ ├── hooks/ # Custom React hooks 145 | │ │ └── contexts/ # React contexts 146 | │ ├── public/ # Static assets 147 | │ └── package.json # Frontend dependencies 148 | ├── pyproject.toml # Python project configuration 149 | └── README.md # This file 150 | ``` 151 | 152 | ## 🔧 Development 153 | 154 | ### Backend Development 155 | 156 | ```bash 157 | # Install in editable mode 158 | pip install -e . 159 | 160 | # Run backend only with auto-reload 161 | lelab 162 | ``` 163 | 164 | ### Frontend Development 165 | 166 | ```bash 167 | # Automatically clones, installs deps, and starts dev server 168 | lelab-frontend 169 | ``` 170 | 171 | ### Full-Stack Development 172 | 173 | ```bash 174 | # Start both backend and frontend with auto-reload 175 | lelab-fullstack 176 | ``` 177 | 178 | **Development Notes:** 179 | 180 | - The frontend repository is automatically managed (cloned/updated) 181 | - Both commands auto-open your browser to the appropriate URL 182 | - Backend runs on `http://localhost:8000` 183 | - Frontend runs on `http://localhost:8080` with API proxying 184 | 185 | ## 🤝 Contributing 186 | 187 | 1. Fork the repository 188 | 2. Create a feature branch (`git checkout -b feature/amazing-feature`) 189 | 3. Commit your changes (`git commit -m 'Add amazing feature'`) 190 | 4. Push to the branch (`git push origin feature/amazing-feature`) 191 | 5. Open a Pull Request 192 | 193 | ## 📄 License 194 | 195 | This project is licensed under the MIT License - see the LICENSE file for details. 196 | 197 | ## 🙏 Acknowledgments 198 | 199 | - [LeRobot](https://github.com/huggingface/lerobot) - The underlying robotics framework 200 | - [FastAPI](https://fastapi.tiangolo.com/) - Modern web framework for building APIs 201 | - [React](https://reactjs.org/) - Frontend user interface library 202 | 203 | --- 204 | 205 | **Note**: Make sure your LeRobot environment is properly configured and your robot hardware is connected before running the application. 206 | -------------------------------------------------------------------------------- /app/replaying.py: -------------------------------------------------------------------------------- 1 | from pydantic import BaseModel 2 | from typing import Dict, Any 3 | import threading 4 | import logging 5 | 6 | logger = logging.getLogger(__name__) 7 | 8 | # Import lerobot replay functionality directly 9 | from lerobot.replay import replay, ReplayConfig, DatasetReplayConfig 10 | from lerobot.common.robots import make_robot_from_config 11 | from lerobot.common.robots.so101_follower import SO101FollowerConfig 12 | from lerobot.common.robots.so100_follower import SO100FollowerConfig 13 | 14 | # Import calibration setup from config 15 | from .config import setup_follower_calibration_file 16 | 17 | # Simple global state 18 | replay_active = False 19 | replay_thread = None 20 | replay_status = { 21 | "replay_active": False, 22 | "status": "idle", 23 | "error_message": None, 24 | } 25 | 26 | class ReplayRequest(BaseModel): 27 | robot_type: str = "so101_follower" 28 | robot_port: str = "/dev/tty.usbmodem58760431541" 29 | robot_id: str = "my_awesome_follower_arm" 30 | dataset_repo_id: str 31 | episode: int = 0 32 | 33 | def run_replay_directly(request: ReplayRequest): 34 | """Run the lerobot replay function directly""" 35 | global replay_active, replay_status 36 | 37 | try: 38 | # Debug logging for all parameters 39 | logger.info(f"🔍 DEBUG: Replay request received") 40 | logger.info(f"🔍 DEBUG: robot_type='{request.robot_type}'") 41 | logger.info(f"🔍 DEBUG: robot_port='{request.robot_port}'") 42 | logger.info(f"🔍 DEBUG: robot_id='{request.robot_id}'") 43 | logger.info(f"🔍 DEBUG: dataset_repo_id='{request.dataset_repo_id}'") 44 | logger.info(f"🔍 DEBUG: episode={request.episode}") 45 | 46 | replay_status.update({ 47 | "replay_active": True, 48 | "status": "running", 49 | "error_message": None, 50 | }) 51 | 52 | logger.info(f"🎬 Starting replay: {request.robot_type} on {request.robot_port}") 53 | logger.info(f"📁 Dataset: {request.dataset_repo_id}, Episode: {request.episode}") 54 | 55 | # Setup calibration file and get the proper config name 56 | logger.info(f"🔧 Setting up calibration file for robot_id: {request.robot_id}") 57 | try: 58 | follower_config_name = setup_follower_calibration_file(request.robot_id) 59 | logger.info(f"✅ Using follower config name: {follower_config_name}") 60 | except Exception as calib_error: 61 | logger.error(f"❌ Calibration setup failed: {calib_error}") 62 | raise 63 | 64 | # Create robot config based on robot type 65 | logger.info(f"🤖 Creating robot config for type: {request.robot_type}") 66 | try: 67 | if request.robot_type == "so101_follower": 68 | robot_config = SO101FollowerConfig( 69 | port=request.robot_port, 70 | id=follower_config_name, # Use processed config name 71 | ) 72 | elif request.robot_type == "so100_follower": 73 | robot_config = SO100FollowerConfig( 74 | port=request.robot_port, 75 | id=follower_config_name, # Use processed config name 76 | ) 77 | else: 78 | raise ValueError(f"Unsupported robot type: {request.robot_type}") 79 | logger.info(f"✅ Robot config created successfully") 80 | except Exception as robot_error: 81 | logger.error(f"❌ Robot config creation failed: {robot_error}") 82 | raise 83 | 84 | # Create dataset config 85 | logger.info(f"📊 Creating dataset config") 86 | try: 87 | dataset_config = DatasetReplayConfig( 88 | repo_id=request.dataset_repo_id, 89 | episode=request.episode, 90 | fps=30 91 | ) 92 | logger.info(f"✅ Dataset config created successfully") 93 | except Exception as dataset_error: 94 | logger.error(f"❌ Dataset config creation failed: {dataset_error}") 95 | raise 96 | 97 | # Create complete replay config 98 | logger.info(f"⚙️ Creating complete replay config") 99 | try: 100 | cfg = ReplayConfig( 101 | robot=robot_config, 102 | dataset=dataset_config, 103 | play_sounds=False # Disable sounds for web interface 104 | ) 105 | logger.info(f"✅ Complete replay config created successfully") 106 | except Exception as config_error: 107 | logger.error(f"❌ Replay config creation failed: {config_error}") 108 | raise 109 | 110 | # Validate robot connection before replay 111 | logger.info(f"🔍 Validating robot connection before replay") 112 | try: 113 | # Create and test robot connection 114 | test_robot = make_robot_from_config(robot_config) 115 | logger.info(f"🤖 Robot created, attempting connection...") 116 | logger.info(f"🔍 Robot config - Port: {robot_config.port}, ID: {robot_config.id}") 117 | 118 | test_robot.connect() 119 | logger.info(f"✅ Robot connected successfully") 120 | logger.info(f"🔍 Robot bus info: {type(test_robot.bus).__name__}") 121 | 122 | # Check if robot has motors configured 123 | if hasattr(test_robot, 'bus') and hasattr(test_robot.bus, 'models'): 124 | if not test_robot.bus.models: 125 | logger.error(f"❌ No motors detected on robot bus") 126 | raise ValueError("No motors detected on robot bus. Check robot connection and calibration.") 127 | else: 128 | logger.info(f"✅ Robot has {len(test_robot.bus.models)} motors configured") 129 | 130 | # Test a simple motor operation to ensure they're working 131 | try: 132 | # Try to read current positions to verify motors are responsive 133 | test_robot.bus.read("Present_Position") 134 | logger.info(f"✅ Motors are responsive") 135 | except Exception as motor_test_error: 136 | logger.error(f"❌ Motors not responsive: {motor_test_error}") 137 | raise ValueError(f"Motors not responsive: {motor_test_error}") 138 | 139 | # Disconnect test connection 140 | test_robot.disconnect() 141 | logger.info(f"✅ Robot validation completed successfully") 142 | 143 | except ValueError as ve: 144 | # Re-raise ValueError with original message 145 | raise ve 146 | except Exception as validation_error: 147 | logger.error(f"❌ Robot validation failed: {validation_error}") 148 | # Provide more specific error message for common issues 149 | error_msg = str(validation_error) 150 | if "No such file or directory" in error_msg or "Permission denied" in error_msg: 151 | raise ValueError(f"Robot port '{request.robot_port}' is not accessible. Check USB connection and permissions.") 152 | elif "models" in error_msg.lower() or "motor" in error_msg.lower(): 153 | raise ValueError(f"Motor configuration error: {validation_error}") 154 | else: 155 | raise ValueError(f"Robot validation failed: {validation_error}") 156 | 157 | # Run the replay directly 158 | logger.info(f"🚀 Starting lerobot replay function") 159 | try: 160 | replay(cfg) 161 | logger.info(f"✅ Lerobot replay function completed successfully") 162 | except StopIteration as stop_error: 163 | logger.error(f"❌ StopIteration error during replay - no motors detected") 164 | raise ValueError("No motors detected during replay. The robot may have been disconnected or the motors are not properly configured. Please check your robot hardware setup and calibration files.") 165 | except Exception as replay_error: 166 | logger.error(f"❌ Lerobot replay function failed: {replay_error}") 167 | raise 168 | 169 | replay_status.update({ 170 | "status": "completed", 171 | "replay_active": False 172 | }) 173 | logger.info("✅ Replay completed successfully") 174 | 175 | except Exception as e: 176 | import traceback 177 | error_msg = str(e) 178 | full_traceback = traceback.format_exc() 179 | 180 | # Enhanced error logging 181 | logger.error(f"❌ Replay failed with exception type: {type(e).__name__}") 182 | logger.error(f"❌ Replay error message: '{error_msg}'") 183 | logger.error(f"❌ Full traceback:\n{full_traceback}") 184 | 185 | # Use more descriptive error message if original is empty 186 | if not error_msg or error_msg.strip() == "": 187 | error_msg = f"Unknown error of type {type(e).__name__}" 188 | 189 | replay_status.update({ 190 | "status": "error", 191 | "replay_active": False, 192 | "error_message": error_msg 193 | }) 194 | finally: 195 | replay_active = False 196 | 197 | def handle_start_replay(request: ReplayRequest) -> Dict[str, Any]: 198 | """Handle starting a replay session""" 199 | global replay_thread, replay_active 200 | 201 | if replay_active: 202 | return { 203 | "success": False, 204 | "message": "Replay is already active" 205 | } 206 | 207 | try: 208 | replay_active = True 209 | replay_thread = threading.Thread( 210 | target=run_replay_directly, 211 | args=(request,), 212 | daemon=True 213 | ) 214 | replay_thread.start() 215 | 216 | return { 217 | "success": True, 218 | "message": "Replay started successfully" 219 | } 220 | 221 | except Exception as e: 222 | replay_active = False 223 | logger.error(f"❌ Failed to start replay: {e}") 224 | return { 225 | "success": False, 226 | "message": str(e) 227 | } 228 | 229 | def handle_stop_replay() -> Dict[str, Any]: 230 | """Handle stopping the current replay session""" 231 | global replay_active 232 | 233 | if not replay_active: 234 | return { 235 | "success": False, 236 | "message": "No active replay session" 237 | } 238 | 239 | # Note: The lerobot replay function doesn't have a built-in stop mechanism 240 | # since it's designed to run to completion. We can only stop between episodes. 241 | replay_active = False 242 | replay_status.update({ 243 | "replay_active": False, 244 | "status": "stopped" 245 | }) 246 | 247 | return { 248 | "success": True, 249 | "message": "Replay stop requested (will complete current episode)" 250 | } 251 | 252 | def handle_replay_status() -> Dict[str, Any]: 253 | """Handle getting the current replay status""" 254 | return { 255 | "success": True, 256 | "status": replay_status.copy() 257 | } 258 | 259 | def handle_replay_logs() -> Dict[str, Any]: 260 | """Handle getting recent replay logs""" 261 | return { 262 | "success": True, 263 | "logs": [] # Logs are handled by lerobot's logging system 264 | } 265 | 266 | def cleanup(): 267 | """Clean up replay resources""" 268 | global replay_active 269 | replay_active = False 270 | -------------------------------------------------------------------------------- /scripts/fullstack.py: -------------------------------------------------------------------------------- 1 | """ 2 | Fullstack script for LeLab 3 | Runs both backend and frontend development servers 4 | Frontend starts detached first, then backend 5 | """ 6 | 7 | import os 8 | import subprocess 9 | import logging 10 | import webbrowser 11 | import time 12 | import signal 13 | import sys 14 | import threading 15 | from pathlib import Path 16 | 17 | # Set up logging 18 | logging.basicConfig(level=logging.INFO) 19 | logger = logging.getLogger(__name__) 20 | 21 | FRONTEND_REPO_URL = "https://github.com/jurmy24/leLab-space.git" 22 | FRONTEND_DIR_NAME = "leLab-space" 23 | 24 | # Global variables to track processes 25 | frontend_process = None 26 | backend_process = None 27 | 28 | 29 | def get_frontend_path(): 30 | """Get the path to the frontend directory""" 31 | # Check if frontend exists in parent directory (same level as leLab) 32 | parent_dir = Path(__file__).parent.parent.parent 33 | frontend_path = parent_dir / FRONTEND_DIR_NAME 34 | 35 | if frontend_path.exists(): 36 | logger.info(f"✅ Found existing frontend at: {frontend_path}") 37 | return frontend_path 38 | 39 | return None 40 | 41 | 42 | def clone_frontend(): 43 | """Clone the frontend repository""" 44 | parent_dir = Path(__file__).parent.parent.parent 45 | frontend_path = parent_dir / FRONTEND_DIR_NAME 46 | 47 | logger.info(f"📥 Cloning frontend repository to: {frontend_path}") 48 | 49 | try: 50 | subprocess.run( 51 | ["git", "clone", FRONTEND_REPO_URL, str(frontend_path)], 52 | check=True, 53 | cwd=parent_dir, 54 | ) 55 | logger.info("✅ Frontend repository cloned successfully") 56 | return frontend_path 57 | except subprocess.CalledProcessError as e: 58 | logger.error(f"❌ Failed to clone frontend repository: {e}") 59 | return None 60 | except FileNotFoundError: 61 | logger.error("❌ git not found. Please install git") 62 | return None 63 | 64 | 65 | def install_frontend_deps(frontend_path): 66 | """Install frontend dependencies""" 67 | logger.info("📦 Installing frontend dependencies...") 68 | 69 | try: 70 | subprocess.run(["npm", "install"], check=True, cwd=frontend_path) 71 | logger.info("✅ Frontend dependencies installed successfully") 72 | return True 73 | except subprocess.CalledProcessError as e: 74 | logger.error(f"❌ Failed to install frontend dependencies: {e}") 75 | return False 76 | except FileNotFoundError: 77 | logger.error("❌ npm not found. Please install Node.js and npm") 78 | return False 79 | 80 | 81 | def start_frontend_detached(frontend_path): 82 | """Start the frontend development server detached""" 83 | global frontend_process 84 | logger.info("🎨 Starting Vite frontend development server (detached)...") 85 | 86 | try: 87 | # Start frontend detached 88 | frontend_process = subprocess.Popen( 89 | ["npm", "run", "dev"], 90 | cwd=frontend_path, 91 | stdout=subprocess.DEVNULL, 92 | stderr=subprocess.DEVNULL, 93 | start_new_session=True, # Detach from parent 94 | ) 95 | 96 | logger.info(f"✅ Frontend server started (PID: {frontend_process.pid})") 97 | return True 98 | except Exception as e: 99 | logger.error(f"❌ Failed to start frontend server: {e}") 100 | return False 101 | 102 | 103 | def wait_for_frontend_ready(): 104 | """Wait for frontend server to be ready""" 105 | logger.info("⏳ Waiting for frontend server to be ready...") 106 | 107 | import socket 108 | import errno 109 | 110 | max_attempts = 30 # 30 seconds max 111 | for attempt in range(max_attempts): 112 | try: 113 | # Try to connect to the frontend port 114 | sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 115 | sock.settimeout(1) 116 | result = sock.connect_ex(("localhost", 8080)) # Vite default port 117 | sock.close() 118 | 119 | if result == 0: 120 | logger.info("✅ Frontend server is ready!") 121 | return True 122 | 123 | except Exception: 124 | pass 125 | 126 | time.sleep(1) 127 | if attempt % 5 == 0: 128 | logger.info(f"⏳ Still waiting for frontend... ({attempt}s)") 129 | 130 | logger.warning("⚠️ Frontend server didn't respond within 30 seconds") 131 | return False 132 | 133 | 134 | def start_backend_detached(): 135 | """Start the backend server detached""" 136 | global backend_process 137 | logger.info("🚀 Starting FastAPI backend server (detached)...") 138 | 139 | try: 140 | # Get the project root directory (where pyproject.toml is located) 141 | project_root = Path(__file__).parent.parent 142 | 143 | # Preserve the current environment variables 144 | env = os.environ.copy() 145 | 146 | backend_process = subprocess.Popen( 147 | [ 148 | sys.executable, 149 | "-m", 150 | "uvicorn", 151 | "app.main:app", 152 | "--host", 153 | "0.0.0.0", 154 | "--port", 155 | "8000", 156 | "--reload", 157 | ], 158 | cwd=project_root, # Set working directory to project root 159 | env=env, # Preserve environment variables 160 | stdout=subprocess.DEVNULL, 161 | stderr=subprocess.DEVNULL, 162 | start_new_session=True, # Detach from parent 163 | ) 164 | 165 | logger.info(f"✅ Backend server started (PID: {backend_process.pid})") 166 | logger.info(f"📁 Backend working directory: {project_root}") 167 | return True 168 | except Exception as e: 169 | logger.error(f"❌ Failed to start backend server: {e}") 170 | return False 171 | 172 | 173 | def wait_for_backend_ready(): 174 | """Wait for backend server to be ready""" 175 | logger.info("⏳ Waiting for backend server to be ready...") 176 | 177 | import socket 178 | 179 | max_attempts = 15 # 15 seconds max 180 | for attempt in range(max_attempts): 181 | try: 182 | sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 183 | sock.settimeout(1) 184 | result = sock.connect_ex(("localhost", 8000)) 185 | sock.close() 186 | 187 | if result == 0: 188 | logger.info("✅ Backend server is ready!") 189 | return True 190 | 191 | except Exception: 192 | pass 193 | 194 | time.sleep(1) 195 | if attempt % 5 == 0: 196 | logger.info(f"⏳ Still waiting for backend... ({attempt}s)") 197 | 198 | logger.warning("⚠️ Backend server didn't respond within 15 seconds") 199 | return False 200 | 201 | 202 | def is_process_running(process): 203 | """Check if a process is still running""" 204 | if process is None: 205 | return False 206 | try: 207 | return process.poll() is None 208 | except: 209 | return False 210 | 211 | 212 | def cleanup_processes(): 213 | """Clean up all processes""" 214 | global backend_process, frontend_process 215 | 216 | logger.info("🛑 Shutting down servers...") 217 | 218 | processes_to_kill = [] 219 | 220 | if backend_process: 221 | processes_to_kill.append(("Backend", backend_process)) 222 | 223 | if frontend_process: 224 | processes_to_kill.append(("Frontend", frontend_process)) 225 | 226 | for name, process in processes_to_kill: 227 | try: 228 | if is_process_running(process): 229 | logger.info(f"🔄 Stopping {name} server (PID: {process.pid})...") 230 | 231 | # Try to kill the process group (handles child processes) 232 | try: 233 | os.killpg(os.getpgid(process.pid), signal.SIGTERM) 234 | except: 235 | process.terminate() 236 | 237 | # Wait for graceful shutdown 238 | try: 239 | process.wait(timeout=5) 240 | logger.info(f"✅ {name} server stopped gracefully") 241 | except subprocess.TimeoutExpired: 242 | logger.warning( 243 | f"⚠️ {name} server didn't stop gracefully, force killing..." 244 | ) 245 | try: 246 | os.killpg(os.getpgid(process.pid), signal.SIGKILL) 247 | except: 248 | process.kill() 249 | logger.info(f"✅ {name} server force stopped") 250 | except Exception as e: 251 | logger.error(f"❌ Error stopping {name}: {e}") 252 | 253 | logger.info("✅ All servers stopped") 254 | 255 | 256 | def signal_handler(signum, frame): 257 | """Handle shutdown signals gracefully""" 258 | logger.info("\n🛑 Received shutdown signal...") 259 | cleanup_processes() 260 | sys.exit(0) 261 | 262 | 263 | def main(): 264 | """Main function to run both backend and frontend""" 265 | logger.info("🚀 Starting LeLab fullstack development servers...") 266 | 267 | # Set up signal handlers for graceful shutdown 268 | signal.signal(signal.SIGINT, signal_handler) 269 | signal.signal(signal.SIGTERM, signal_handler) 270 | 271 | try: 272 | # Get or clone frontend 273 | frontend_path = get_frontend_path() 274 | if not frontend_path: 275 | frontend_path = clone_frontend() 276 | if not frontend_path: 277 | logger.error("❌ Failed to get frontend repository") 278 | return 279 | 280 | # Install frontend dependencies 281 | if not install_frontend_deps(frontend_path): 282 | logger.error("❌ Failed to install frontend dependencies") 283 | return 284 | 285 | # Step 1: Start frontend detached 286 | if not start_frontend_detached(frontend_path): 287 | logger.error("❌ Failed to start frontend server") 288 | return 289 | 290 | # Step 2: Wait for frontend to be ready 291 | if not wait_for_frontend_ready(): 292 | logger.error("❌ Frontend server not ready") 293 | cleanup_processes() 294 | return 295 | 296 | # Step 3: Start backend detached 297 | if not start_backend_detached(): 298 | logger.error("❌ Failed to start backend server") 299 | cleanup_processes() 300 | return 301 | 302 | # Step 4: Wait for backend to be ready 303 | if not wait_for_backend_ready(): 304 | logger.error("❌ Backend server not ready") 305 | cleanup_processes() 306 | return 307 | 308 | # Step 5: Open browser with auto-reset to localhost 309 | logger.info("🌐 Opening browser...") 310 | # Open with a URL that will automatically reset to localhost mode 311 | webbrowser.open("http://localhost:8080?reset_to_localhost=true") 312 | 313 | # Success! 314 | logger.info("✅ Both servers are running!") 315 | logger.info("📱 Backend: http://localhost:8000") 316 | logger.info("🌐 Frontend: http://localhost:8080") 317 | logger.info("🛑 Press Ctrl+C to stop both servers") 318 | 319 | # Keep the script running and monitor processes 320 | while True: 321 | time.sleep(5) 322 | 323 | # Check if processes are still running 324 | if not is_process_running(frontend_process): 325 | logger.error("❌ Frontend process died") 326 | break 327 | if not is_process_running(backend_process): 328 | logger.error("❌ Backend process died") 329 | break 330 | 331 | except KeyboardInterrupt: 332 | logger.info("\n🛑 Received interrupt signal") 333 | except Exception as e: 334 | logger.error(f"❌ Unexpected error: {e}") 335 | finally: 336 | cleanup_processes() 337 | 338 | 339 | if __name__ == "__main__": 340 | main() 341 | -------------------------------------------------------------------------------- /app/config.py: -------------------------------------------------------------------------------- 1 | import os 2 | import shutil 3 | import logging 4 | import platform 5 | import time 6 | from pathlib import Path 7 | 8 | logger = logging.getLogger(__name__) 9 | 10 | # Define the calibration config paths (shared between features) 11 | CALIBRATION_BASE_PATH_TELEOP = os.path.expanduser( 12 | "~/.cache/huggingface/lerobot/calibration/teleoperators" 13 | ) 14 | CALIBRATION_BASE_PATH_ROBOTS = os.path.expanduser( 15 | "~/.cache/huggingface/lerobot/calibration/robots" 16 | ) 17 | LEADER_CONFIG_PATH = os.path.join(CALIBRATION_BASE_PATH_TELEOP, "so101_leader") 18 | FOLLOWER_CONFIG_PATH = os.path.join(CALIBRATION_BASE_PATH_ROBOTS, "so101_follower") 19 | 20 | # Define port storage path 21 | PORT_CONFIG_PATH = os.path.expanduser("~/.cache/huggingface/lerobot/ports") 22 | LEADER_PORT_FILE = os.path.join(PORT_CONFIG_PATH, "leader_port.txt") 23 | FOLLOWER_PORT_FILE = os.path.join(PORT_CONFIG_PATH, "follower_port.txt") 24 | 25 | # Define configuration storage path 26 | CONFIG_STORAGE_PATH = os.path.expanduser("~/.cache/huggingface/lerobot/saved_configs") 27 | LEADER_CONFIG_FILE = os.path.join(CONFIG_STORAGE_PATH, "leader_config.txt") 28 | FOLLOWER_CONFIG_FILE = os.path.join(CONFIG_STORAGE_PATH, "follower_config.txt") 29 | 30 | def setup_calibration_files(leader_config: str, follower_config: str): 31 | """Setup calibration files in the correct locations for teleoperation and recording""" 32 | # Extract config names from file paths (remove .json extension) 33 | leader_config_name = os.path.splitext(leader_config)[0] 34 | follower_config_name = os.path.splitext(follower_config)[0] 35 | 36 | # Log the full paths to check if files exist 37 | leader_config_full_path = os.path.join(LEADER_CONFIG_PATH, leader_config) 38 | follower_config_full_path = os.path.join(FOLLOWER_CONFIG_PATH, follower_config) 39 | 40 | logger.info(f"Checking calibration files:") 41 | logger.info(f"Leader config path: {leader_config_full_path}") 42 | logger.info(f"Follower config path: {follower_config_full_path}") 43 | logger.info(f"Leader config exists: {os.path.exists(leader_config_full_path)}") 44 | logger.info(f"Follower config exists: {os.path.exists(follower_config_full_path)}") 45 | 46 | # Create calibration directories if they don't exist 47 | leader_calibration_dir = os.path.join(CALIBRATION_BASE_PATH_TELEOP, "so101_leader") 48 | follower_calibration_dir = os.path.join(CALIBRATION_BASE_PATH_ROBOTS, "so101_follower") 49 | os.makedirs(leader_calibration_dir, exist_ok=True) 50 | os.makedirs(follower_calibration_dir, exist_ok=True) 51 | 52 | # Copy calibration files to the correct locations if they're not already there 53 | leader_target_path = os.path.join(leader_calibration_dir, f"{leader_config_name}.json") 54 | follower_target_path = os.path.join(follower_calibration_dir, f"{follower_config_name}.json") 55 | 56 | if not os.path.exists(leader_target_path): 57 | if os.path.exists(leader_config_full_path): 58 | shutil.copy2(leader_config_full_path, leader_target_path) 59 | logger.info(f"Copied leader calibration to {leader_target_path}") 60 | else: 61 | raise FileNotFoundError(f"Leader calibration file not found: {leader_config_full_path}") 62 | else: 63 | logger.info(f"Leader calibration already exists at {leader_target_path}") 64 | 65 | if not os.path.exists(follower_target_path): 66 | if os.path.exists(follower_config_full_path): 67 | shutil.copy2(follower_config_full_path, follower_target_path) 68 | logger.info(f"Copied follower calibration to {follower_target_path}") 69 | else: 70 | raise FileNotFoundError(f"Follower calibration file not found: {follower_config_full_path}") 71 | else: 72 | logger.info(f"Follower calibration already exists at {follower_target_path}") 73 | 74 | return leader_config_name, follower_config_name 75 | 76 | 77 | def setup_follower_calibration_file(follower_config: str): 78 | """Setup follower calibration file in the correct location for replay functionality""" 79 | # Extract config name from file path (remove .json extension) 80 | follower_config_name = os.path.splitext(follower_config)[0] 81 | 82 | # Log the full path to check if file exists 83 | follower_config_full_path = os.path.join(FOLLOWER_CONFIG_PATH, follower_config) 84 | 85 | logger.info(f"Checking follower calibration file:") 86 | logger.info(f"Follower config path: {follower_config_full_path}") 87 | logger.info(f"Follower config exists: {os.path.exists(follower_config_full_path)}") 88 | 89 | # Create calibration directory if it doesn't exist 90 | follower_calibration_dir = os.path.join(CALIBRATION_BASE_PATH_ROBOTS, "so101_follower") 91 | os.makedirs(follower_calibration_dir, exist_ok=True) 92 | 93 | # Copy calibration file to the correct location if it's not already there 94 | follower_target_path = os.path.join(follower_calibration_dir, f"{follower_config_name}.json") 95 | 96 | if not os.path.exists(follower_target_path): 97 | if os.path.exists(follower_config_full_path): 98 | shutil.copy2(follower_config_full_path, follower_target_path) 99 | logger.info(f"Copied follower calibration to {follower_target_path}") 100 | else: 101 | raise FileNotFoundError(f"Follower calibration file not found: {follower_config_full_path}") 102 | else: 103 | logger.info(f"Follower calibration already exists at {follower_target_path}") 104 | 105 | return follower_config_name 106 | 107 | 108 | def find_available_ports(): 109 | """Find all available serial ports on the system""" 110 | try: 111 | from serial.tools import list_ports # Part of pyserial library 112 | except ImportError: 113 | raise ImportError("pyserial library is required. Install it with: pip install pyserial") 114 | 115 | if platform.system() == "Windows": 116 | # List COM ports using pyserial 117 | ports = [port.device for port in list_ports.comports()] 118 | else: # Linux/macOS 119 | # List /dev/tty* ports for Unix-based systems 120 | ports = [str(path) for path in Path("/dev").glob("tty*")] 121 | return sorted(ports) 122 | 123 | 124 | def find_robot_port(robot_type="robot"): 125 | """ 126 | Find the port for a robot by detecting the difference when disconnecting/reconnecting 127 | 128 | Args: 129 | robot_type (str): Type of robot ("leader" or "follower" or generic "robot") 130 | 131 | Returns: 132 | str: The detected port 133 | """ 134 | logger.info(f"Finding port for {robot_type}") 135 | 136 | # Get initial ports 137 | ports_before = find_available_ports() 138 | logger.info(f"Ports before disconnecting: {ports_before}") 139 | 140 | # This function returns the port detection logic, but the actual user interaction 141 | # should be handled by the frontend 142 | return { 143 | "ports_before": ports_before, 144 | "robot_type": robot_type 145 | } 146 | 147 | 148 | def detect_port_after_disconnect(ports_before): 149 | """ 150 | Detect the port after disconnection by comparing with ports before 151 | 152 | Args: 153 | ports_before (list): List of ports before disconnection 154 | 155 | Returns: 156 | str: The detected port 157 | """ 158 | time.sleep(0.5) # Allow some time for port to be released 159 | ports_after = find_available_ports() 160 | ports_diff = list(set(ports_before) - set(ports_after)) 161 | 162 | logger.info(f"Ports after disconnecting: {ports_after}") 163 | logger.info(f"Port difference: {ports_diff}") 164 | 165 | if len(ports_diff) == 1: 166 | port = ports_diff[0] 167 | logger.info(f"Detected port: {port}") 168 | return port 169 | elif len(ports_diff) == 0: 170 | raise OSError("Could not detect the port. No difference was found.") 171 | else: 172 | raise OSError(f"Could not detect the port. More than one port was found ({ports_diff}).") 173 | 174 | 175 | def save_robot_port(robot_type, port): 176 | """ 177 | Save the robot port to a file for future use 178 | 179 | Args: 180 | robot_type (str): "leader" or "follower" 181 | port (str): The port to save 182 | """ 183 | # Create port config directory if it doesn't exist 184 | os.makedirs(PORT_CONFIG_PATH, exist_ok=True) 185 | 186 | port_file = LEADER_PORT_FILE if robot_type == "leader" else FOLLOWER_PORT_FILE 187 | 188 | with open(port_file, 'w') as f: 189 | f.write(port) 190 | 191 | logger.info(f"Saved {robot_type} port: {port}") 192 | 193 | 194 | def get_saved_robot_port(robot_type): 195 | """ 196 | Get the saved robot port from file 197 | 198 | Args: 199 | robot_type (str): "leader" or "follower" 200 | 201 | Returns: 202 | str or None: The saved port, or None if not found 203 | """ 204 | port_file = LEADER_PORT_FILE if robot_type == "leader" else FOLLOWER_PORT_FILE 205 | 206 | if os.path.exists(port_file): 207 | with open(port_file, 'r') as f: 208 | port = f.read().strip() 209 | logger.info(f"Retrieved saved {robot_type} port: {port}") 210 | return port 211 | 212 | logger.info(f"No saved port found for {robot_type}") 213 | return None 214 | 215 | 216 | def get_default_robot_port(robot_type): 217 | """ 218 | Get the default port for a robot, checking saved ports first 219 | 220 | Args: 221 | robot_type (str): "leader" or "follower" 222 | 223 | Returns: 224 | str: The default port to use 225 | """ 226 | saved_port = get_saved_robot_port(robot_type) 227 | if saved_port: 228 | return saved_port 229 | 230 | # Fallback to common default ports 231 | if platform.system() == "Windows": 232 | return "COM3" # Common Windows default 233 | else: 234 | return "/dev/ttyUSB0" # Common Linux/macOS default 235 | 236 | 237 | def save_robot_config(robot_type: str, config_name: str): 238 | """Save the robot configuration to a file for future use""" 239 | try: 240 | # Create the config storage directory if it doesn't exist 241 | os.makedirs(CONFIG_STORAGE_PATH, exist_ok=True) 242 | 243 | # Determine the config file path 244 | if robot_type.lower() == "leader": 245 | config_file_path = LEADER_CONFIG_FILE 246 | elif robot_type.lower() == "follower": 247 | config_file_path = FOLLOWER_CONFIG_FILE 248 | else: 249 | logger.error(f"Unknown robot type: {robot_type}") 250 | return False 251 | 252 | # Write the config name to file 253 | with open(config_file_path, 'w') as f: 254 | f.write(config_name.strip()) 255 | 256 | logger.info(f"Saved {robot_type} configuration: {config_name}") 257 | return True 258 | 259 | except Exception as e: 260 | logger.error(f"Error saving {robot_type} configuration: {e}") 261 | return False 262 | 263 | 264 | def get_saved_robot_config(robot_type: str): 265 | """Get the saved robot configuration from file""" 266 | try: 267 | # Determine the config file path 268 | if robot_type.lower() == "leader": 269 | config_file_path = LEADER_CONFIG_FILE 270 | elif robot_type.lower() == "follower": 271 | config_file_path = FOLLOWER_CONFIG_FILE 272 | else: 273 | logger.error(f"Unknown robot type: {robot_type}") 274 | return None 275 | 276 | # Read the config name from file 277 | if os.path.exists(config_file_path): 278 | with open(config_file_path, 'r') as f: 279 | config_name = f.read().strip() 280 | if config_name: 281 | logger.info(f"Found saved {robot_type} configuration: {config_name}") 282 | return config_name 283 | 284 | logger.info(f"No saved {robot_type} configuration found") 285 | return None 286 | 287 | except Exception as e: 288 | logger.error(f"Error reading saved {robot_type} configuration: {e}") 289 | return None 290 | 291 | 292 | def get_default_robot_config(robot_type: str, available_configs: list): 293 | """Get the default configuration for a robot, checking saved configs first""" 294 | saved_config = get_saved_robot_config(robot_type) 295 | if saved_config and saved_config in available_configs: 296 | return saved_config 297 | 298 | # Return first available config as fallback 299 | if available_configs: 300 | return available_configs[0] 301 | 302 | return None 303 | -------------------------------------------------------------------------------- /app/teleoperating.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | import shutil 4 | import select 5 | import sys 6 | import termios 7 | import tty 8 | import time 9 | from typing import Dict, Any, List 10 | from concurrent.futures import ThreadPoolExecutor 11 | from pydantic import BaseModel 12 | 13 | from lerobot.common.teleoperators.so101_leader import SO101LeaderConfig, SO101Leader 14 | from lerobot.common.robots.so101_follower import SO101FollowerConfig, SO101Follower 15 | from lerobot.teleoperate import teleoperate, TeleoperateConfig 16 | 17 | # Import calibration paths and functions from config (shared constants) 18 | from .config import ( 19 | CALIBRATION_BASE_PATH_TELEOP, 20 | CALIBRATION_BASE_PATH_ROBOTS, 21 | LEADER_CONFIG_PATH, 22 | FOLLOWER_CONFIG_PATH, 23 | setup_calibration_files 24 | ) 25 | 26 | logger = logging.getLogger(__name__) 27 | 28 | # Global variables for teleoperation state 29 | teleoperation_active = False 30 | teleoperation_thread = None 31 | current_robot = None 32 | current_teleop = None 33 | 34 | 35 | class TeleoperateRequest(BaseModel): 36 | leader_port: str 37 | follower_port: str 38 | leader_config: str 39 | follower_config: str 40 | 41 | 42 | def setup_keyboard(): 43 | """Set up keyboard for non-blocking input""" 44 | old_settings = termios.tcgetattr(sys.stdin) 45 | tty.setraw(sys.stdin.fileno()) 46 | return old_settings 47 | 48 | 49 | def restore_keyboard(old_settings): 50 | """Restore keyboard settings""" 51 | termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings) 52 | 53 | 54 | def check_quit_key(): 55 | """Check if 'q' key was pressed""" 56 | if select.select([sys.stdin], [], [], 0) == ([sys.stdin], [], []): 57 | key = sys.stdin.read(1) 58 | return key.lower() == "q" 59 | return False 60 | 61 | 62 | def get_joint_positions_from_robot(robot) -> Dict[str, float]: 63 | """ 64 | Extract current joint positions from the robot and convert to URDF joint format. 65 | 66 | Args: 67 | robot: The robot instance (SO101Follower) 68 | 69 | Returns: 70 | Dictionary mapping URDF joint names to radian values 71 | """ 72 | try: 73 | # Get the current observation from the robot 74 | observation = robot.get_observation() 75 | 76 | # Map robot motor names to URDF joint names 77 | # Based on the motor configuration in SO101Follower and URDF joint names 78 | motor_to_urdf_mapping = { 79 | "shoulder_pan": "Rotation", # Base rotation 80 | "shoulder_lift": "Pitch", # Shoulder pitch 81 | "elbow_flex": "Elbow", # Elbow flexion 82 | "wrist_flex": "Wrist_Pitch", # Wrist pitch 83 | "wrist_roll": "Wrist_Roll", # Wrist roll 84 | "gripper": "Jaw", # Gripper/jaw 85 | } 86 | 87 | joint_positions = {} 88 | 89 | # Extract joint positions and convert degrees to radians 90 | for motor_name, urdf_joint_name in motor_to_urdf_mapping.items(): 91 | motor_key = f"{motor_name}.pos" 92 | if motor_key in observation: 93 | # Convert degrees to radians for the URDF viewer 94 | angle_degrees = observation[motor_key] 95 | angle_radians = angle_degrees * (3.14159 / 180.0) 96 | joint_positions[urdf_joint_name] = angle_radians 97 | else: 98 | logger.warning(f"Motor {motor_key} not found in observation") 99 | joint_positions[urdf_joint_name] = 0.0 100 | 101 | return joint_positions 102 | 103 | except Exception as e: 104 | logger.error(f"Error getting joint positions: {e}") 105 | return { 106 | "Rotation": 0.0, 107 | "Pitch": 0.0, 108 | "Elbow": 0.0, 109 | "Wrist_Pitch": 0.0, 110 | "Wrist_Roll": 0.0, 111 | "Jaw": 0.0, 112 | } 113 | 114 | 115 | 116 | def handle_start_teleoperation(request: TeleoperateRequest, websocket_manager=None) -> Dict[str, Any]: 117 | """Handle start teleoperation request""" 118 | global teleoperation_active, teleoperation_thread, current_robot, current_teleop 119 | 120 | if teleoperation_active: 121 | return {"success": False, "message": "Teleoperation is already active"} 122 | 123 | try: 124 | logger.info(f"Starting teleoperation with leader port: {request.leader_port}, follower port: {request.follower_port}") 125 | 126 | # Setup calibration files 127 | leader_config_name, follower_config_name = setup_calibration_files( 128 | request.leader_config, request.follower_config 129 | ) 130 | 131 | # Create robot and teleop configs 132 | robot_config = SO101FollowerConfig( 133 | port=request.follower_port, 134 | id=follower_config_name, 135 | ) 136 | 137 | teleop_config = SO101LeaderConfig( 138 | port=request.leader_port, 139 | id=leader_config_name, 140 | ) 141 | 142 | # Start teleoperation in a separate thread 143 | def teleoperation_worker(): 144 | global teleoperation_active, current_robot, current_teleop 145 | teleoperation_active = True 146 | 147 | try: 148 | logger.info("Initializing robot and teleop device...") 149 | robot = SO101Follower(robot_config) 150 | teleop_device = SO101Leader(teleop_config) 151 | 152 | current_robot = robot 153 | current_teleop = teleop_device 154 | 155 | logger.info("Connecting to devices...") 156 | robot.bus.connect() 157 | teleop_device.bus.connect() 158 | 159 | # Write calibration to motors' memory 160 | logger.info("Writing calibration to motors...") 161 | robot.bus.write_calibration(robot.calibration) 162 | teleop_device.bus.write_calibration(teleop_device.calibration) 163 | 164 | # Connect cameras and configure motors 165 | logger.info("Connecting cameras and configuring motors...") 166 | for cam in robot.cameras.values(): 167 | cam.connect() 168 | robot.configure() 169 | teleop_device.configure() 170 | logger.info("Successfully connected to both devices") 171 | 172 | logger.info("Starting teleoperation loop...") 173 | logger.info("Press 'q' to quit teleoperation") 174 | 175 | # Set up keyboard for non-blocking input 176 | old_settings = setup_keyboard() 177 | 178 | try: 179 | want_to_disconnect = False 180 | last_broadcast_time = 0 181 | broadcast_interval = 0.05 # Broadcast every 50ms (20 FPS) 182 | 183 | while not want_to_disconnect and teleoperation_active: 184 | # Check teleoperation_active flag first (for web stop requests) 185 | if not teleoperation_active: 186 | logger.info("Teleoperation stopped via web interface") 187 | break 188 | 189 | action = teleop_device.get_action() 190 | robot.send_action(action) 191 | 192 | # Broadcast joint positions to connected WebSocket clients 193 | current_time = time.time() 194 | if current_time - last_broadcast_time >= broadcast_interval: 195 | try: 196 | joint_positions = get_joint_positions_from_robot(robot) 197 | joint_data = { 198 | "type": "joint_update", 199 | "joints": joint_positions, 200 | "timestamp": current_time, 201 | } 202 | 203 | # Use websocket manager to broadcast the data 204 | if websocket_manager and websocket_manager.active_connections: 205 | websocket_manager.broadcast_joint_data_sync(joint_data) 206 | 207 | last_broadcast_time = current_time 208 | except Exception as e: 209 | logger.error(f"Error broadcasting joint data: {e}") 210 | 211 | # Check for keyboard input 212 | if check_quit_key(): 213 | want_to_disconnect = True 214 | logger.info("Quit key pressed, stopping teleoperation...") 215 | 216 | # Small delay to prevent excessive CPU usage and allow for responsive stopping 217 | time.sleep(0.001) # 1ms delay 218 | finally: 219 | # Always restore keyboard settings 220 | restore_keyboard(old_settings) 221 | robot.disconnect() 222 | teleop_device.disconnect() 223 | logger.info("Teleoperation stopped") 224 | 225 | return {"success": True, "message": "Teleoperation completed successfully"} 226 | 227 | except Exception as e: 228 | logger.error(f"Error during teleoperation: {e}") 229 | return {"success": False, "error": str(e)} 230 | finally: 231 | teleoperation_active = False 232 | current_robot = None 233 | current_teleop = None 234 | 235 | teleoperation_thread = ThreadPoolExecutor(max_workers=1) 236 | future = teleoperation_thread.submit(teleoperation_worker) 237 | 238 | return { 239 | "success": True, 240 | "message": "Teleoperation started successfully", 241 | "leader_port": request.leader_port, 242 | "follower_port": request.follower_port 243 | } 244 | 245 | except Exception as e: 246 | teleoperation_active = False 247 | logger.error(f"Failed to start teleoperation: {e}") 248 | return {"success": False, "message": f"Failed to start teleoperation: {str(e)}"} 249 | 250 | 251 | def handle_stop_teleoperation() -> Dict[str, Any]: 252 | """Handle stop teleoperation request""" 253 | global teleoperation_active, teleoperation_thread, current_robot, current_teleop 254 | 255 | if not teleoperation_active: 256 | return {"success": False, "message": "No teleoperation session is active"} 257 | 258 | try: 259 | logger.info("Stop teleoperation triggered from web interface") 260 | 261 | # Stop the teleoperation flag 262 | teleoperation_active = False 263 | 264 | # Force cleanup of robot connections if they exist 265 | try: 266 | if current_robot: 267 | logger.info("Disconnecting robot...") 268 | current_robot.disconnect() 269 | 270 | if current_teleop: 271 | logger.info("Disconnecting teleop device...") 272 | current_teleop.disconnect() 273 | except Exception as cleanup_error: 274 | logger.warning(f"Error during device cleanup: {cleanup_error}") 275 | 276 | # Wait for the thread to finish (with timeout) 277 | if teleoperation_thread: 278 | try: 279 | logger.info("Waiting for teleoperation thread to finish...") 280 | # Give the thread a moment to finish gracefully 281 | import time 282 | time.sleep(0.5) 283 | 284 | # Shutdown the thread pool 285 | teleoperation_thread.shutdown(wait=True, timeout=5.0) 286 | logger.info("Teleoperation thread stopped") 287 | except Exception as thread_error: 288 | logger.warning(f"Error stopping teleoperation thread: {thread_error}") 289 | 290 | # Clean up global variables 291 | current_robot = None 292 | current_teleop = None 293 | teleoperation_thread = None 294 | 295 | logger.info("Teleoperation stopped successfully") 296 | 297 | return { 298 | "success": True, 299 | "message": "Teleoperation stopped successfully", 300 | } 301 | 302 | except Exception as e: 303 | logger.error(f"Error stopping teleoperation: {e}") 304 | return {"success": False, "message": f"Failed to stop teleoperation: {str(e)}"} 305 | 306 | 307 | def handle_teleoperation_status() -> Dict[str, Any]: 308 | """Handle teleoperation status request""" 309 | return { 310 | "teleoperation_active": teleoperation_active, 311 | "available_controls": { 312 | "stop_teleoperation": teleoperation_active, 313 | }, 314 | "message": "Teleoperation status retrieved successfully" 315 | } 316 | 317 | 318 | def handle_get_joint_positions() -> Dict[str, Any]: 319 | """Handle get current robot joint positions request""" 320 | global current_robot 321 | 322 | if not teleoperation_active or current_robot is None: 323 | return {"success": False, "message": "No active teleoperation session"} 324 | 325 | try: 326 | joint_positions = get_joint_positions_from_robot(current_robot) 327 | return { 328 | "success": True, 329 | "joint_positions": joint_positions, 330 | "timestamp": time.time() 331 | } 332 | except Exception as e: 333 | logger.error(f"Error getting joint positions: {e}") 334 | return {"success": False, "message": f"Failed to get joint positions: {str(e)}"} 335 | -------------------------------------------------------------------------------- /app/training.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import subprocess 3 | import threading 4 | import time 5 | from pathlib import Path 6 | from typing import Dict, Any, Optional, List 7 | from pydantic import BaseModel 8 | import json 9 | import queue 10 | import os 11 | import signal 12 | import psutil 13 | 14 | logger = logging.getLogger(__name__) 15 | 16 | class TrainingRequest(BaseModel): 17 | # Dataset configuration - exact matches from CLI 18 | dataset_repo_id: str # --dataset.repo_id 19 | dataset_revision: Optional[str] = None # --dataset.revision 20 | dataset_root: Optional[str] = None # --dataset.root 21 | dataset_episodes: Optional[List[int]] = None # --dataset.episodes 22 | 23 | # Policy configuration - only type is configurable at top level 24 | policy_type: str = "act" # --policy.type (act, diffusion, pi0, smolvla, tdmpc, vqbet, pi0fast, sac, reward_classifier) 25 | 26 | # Core training parameters - exact matches from CLI 27 | steps: int = 10000 # --steps 28 | batch_size: int = 8 # --batch_size 29 | seed: Optional[int] = 1000 # --seed 30 | num_workers: int = 4 # --num_workers 31 | 32 | # Logging and checkpointing - exact matches from CLI 33 | log_freq: int = 250 # --log_freq 34 | save_freq: int = 1000 # --save_freq 35 | eval_freq: int = 0 # --eval_freq 36 | save_checkpoint: bool = True # --save_checkpoint 37 | 38 | # Output configuration - exact matches from CLI 39 | output_dir: str = "outputs/train" # --output_dir 40 | resume: bool = False # --resume 41 | job_name: Optional[str] = None # --job_name 42 | 43 | # Weights & Biases - exact matches from CLI 44 | wandb_enable: bool = False # --wandb.enable 45 | wandb_project: Optional[str] = None # --wandb.project 46 | wandb_entity: Optional[str] = None # --wandb.entity 47 | wandb_notes: Optional[str] = None # --wandb.notes 48 | wandb_run_id: Optional[str] = None # --wandb.run_id 49 | wandb_mode: Optional[str] = "online" # --wandb.mode (online, offline, disabled) 50 | wandb_disable_artifact: bool = False # --wandb.disable_artifact 51 | 52 | # Environment and evaluation - exact matches from CLI 53 | env_type: Optional[str] = None # --env.type (aloha, pusht, xarm, gym_manipulator, hil) 54 | env_task: Optional[str] = None # --env.task 55 | eval_n_episodes: int = 10 # --eval.n_episodes 56 | eval_batch_size: int = 50 # --eval.batch_size 57 | eval_use_async_envs: bool = False # --eval.use_async_envs 58 | 59 | # Policy-specific parameters that are commonly used 60 | policy_device: Optional[str] = "cuda" # --policy.device 61 | policy_use_amp: bool = False # --policy.use_amp 62 | 63 | # Optimizer parameters - exact matches from CLI 64 | optimizer_type: Optional[str] = "adam" # --optimizer.type (adam, adamw, sgd, multi_adam) 65 | optimizer_lr: Optional[float] = None # --optimizer.lr (will use policy default if not set) 66 | optimizer_weight_decay: Optional[float] = None # --optimizer.weight_decay 67 | optimizer_grad_clip_norm: Optional[float] = None # --optimizer.grad_clip_norm 68 | 69 | # Advanced configuration 70 | use_policy_training_preset: bool = True # --use_policy_training_preset 71 | config_path: Optional[str] = None # --config_path 72 | 73 | class TrainingStatus(BaseModel): 74 | training_active: bool = False 75 | current_step: int = 0 76 | total_steps: int = 0 77 | current_loss: Optional[float] = None 78 | current_lr: Optional[float] = None 79 | grad_norm: Optional[float] = None 80 | epoch_time: Optional[float] = None 81 | eta_seconds: Optional[float] = None 82 | available_controls: Dict[str, bool] = { 83 | "stop_training": False, 84 | "pause_training": False, 85 | "resume_training": False 86 | } 87 | 88 | class TrainingManager: 89 | def __init__(self): 90 | self.process: Optional[subprocess.Popen] = None 91 | self.status = TrainingStatus() 92 | self.log_queue = queue.Queue() 93 | self.log_thread: Optional[threading.Thread] = None 94 | self.monitor_thread: Optional[threading.Thread] = None 95 | self._stop_monitoring = threading.Event() 96 | 97 | def start_training(self, request: TrainingRequest) -> Dict[str, Any]: 98 | """Start a training session""" 99 | if self.status.training_active: 100 | return {"success": False, "message": "Training is already active"} 101 | 102 | try: 103 | # Create output directory 104 | output_path = Path(request.output_dir) 105 | output_path.mkdir(parents=True, exist_ok=True) 106 | 107 | # Build the training command 108 | cmd = self._build_training_command(request) 109 | logger.info(f"Starting training with command: {' '.join(cmd)}") 110 | 111 | # Start the training process 112 | self.process = subprocess.Popen( 113 | cmd, 114 | stdout=subprocess.PIPE, 115 | stderr=subprocess.STDOUT, 116 | universal_newlines=True, 117 | bufsize=1, 118 | env=os.environ.copy() 119 | ) 120 | 121 | # Update status 122 | self.status.training_active = True 123 | self.status.total_steps = request.steps 124 | self.status.current_step = 0 125 | self.status.available_controls = { 126 | "stop_training": True, 127 | "pause_training": False, # Not implemented in LeRobot 128 | "resume_training": False 129 | } 130 | 131 | # Start monitoring threads 132 | self._start_monitoring() 133 | 134 | return {"success": True, "message": "Training started successfully"} 135 | 136 | except Exception as e: 137 | logger.error(f"Failed to start training: {e}") 138 | return {"success": False, "message": f"Failed to start training: {str(e)}"} 139 | 140 | def stop_training(self) -> Dict[str, Any]: 141 | """Stop the current training session""" 142 | if not self.status.training_active: 143 | return {"success": False, "message": "No training session is active"} 144 | 145 | try: 146 | if self.process and self.process.poll() is None: 147 | # Try graceful shutdown first 148 | self.process.terminate() 149 | 150 | # Wait for a bit, then force kill if necessary 151 | try: 152 | self.process.wait(timeout=10) 153 | except subprocess.TimeoutExpired: 154 | logger.warning("Training process didn't terminate gracefully, forcing kill") 155 | self.process.kill() 156 | self.process.wait() 157 | 158 | self._stop_monitoring_threads() 159 | self._reset_status() 160 | 161 | return {"success": True, "message": "Training stopped successfully"} 162 | 163 | except Exception as e: 164 | logger.error(f"Failed to stop training: {e}") 165 | return {"success": False, "message": f"Failed to stop training: {str(e)}"} 166 | 167 | def get_status(self) -> TrainingStatus: 168 | """Get current training status""" 169 | # Check if process is still running 170 | if self.process and self.process.poll() is not None: 171 | # Process has ended 172 | if self.status.training_active: 173 | self._stop_monitoring_threads() 174 | self._reset_status() 175 | 176 | return self.status 177 | 178 | def get_logs(self) -> list: 179 | """Get recent training logs""" 180 | logs = [] 181 | try: 182 | while not self.log_queue.empty(): 183 | logs.append(self.log_queue.get_nowait()) 184 | except queue.Empty: 185 | pass 186 | return logs 187 | 188 | def _build_training_command(self, request: TrainingRequest) -> list: 189 | """Build the training command from the request parameters - only using actual CLI parameters""" 190 | cmd = ["python", "-m", "lerobot.scripts.train"] 191 | 192 | # Dataset configuration 193 | cmd.extend(["--dataset.repo_id", request.dataset_repo_id]) 194 | if request.dataset_revision: 195 | cmd.extend(["--dataset.revision", request.dataset_revision]) 196 | if request.dataset_root: 197 | cmd.extend(["--dataset.root", request.dataset_root]) 198 | if request.dataset_episodes: 199 | cmd.extend(["--dataset.episodes"] + [str(ep) for ep in request.dataset_episodes]) 200 | 201 | # Policy type 202 | cmd.extend(["--policy.type", request.policy_type]) 203 | 204 | # Core training parameters 205 | cmd.extend(["--steps", str(request.steps)]) 206 | cmd.extend(["--batch-size", str(request.batch_size)]) 207 | cmd.extend(["--num-workers", str(request.num_workers)]) 208 | 209 | if request.seed is not None: 210 | cmd.extend(["--seed", str(request.seed)]) 211 | 212 | # Policy device and AMP 213 | if request.policy_device: 214 | cmd.extend(["--policy.device", request.policy_device]) 215 | if request.policy_use_amp: 216 | cmd.append("--policy.use_amp") 217 | 218 | # Logging and checkpointing 219 | cmd.extend(["--log-freq", str(request.log_freq)]) 220 | cmd.extend(["--save-freq", str(request.save_freq)]) 221 | cmd.extend(["--eval-freq", str(request.eval_freq)]) 222 | 223 | if request.save_checkpoint: 224 | cmd.append("--save_checkpoint") 225 | 226 | # Output configuration 227 | cmd.extend(["--output-dir", request.output_dir]) 228 | if request.resume: 229 | cmd.append("--resume") 230 | if request.job_name: 231 | cmd.extend(["--job-name", request.job_name]) 232 | 233 | # Weights & Biases 234 | if request.wandb_enable: 235 | cmd.append("--wandb.enable") 236 | if request.wandb_project: 237 | cmd.extend(["--wandb.project", request.wandb_project]) 238 | if request.wandb_entity: 239 | cmd.extend(["--wandb.entity", request.wandb_entity]) 240 | if request.wandb_notes: 241 | cmd.extend(["--wandb.notes", request.wandb_notes]) 242 | if request.wandb_run_id: 243 | cmd.extend(["--wandb.run_id", request.wandb_run_id]) 244 | if request.wandb_mode: 245 | cmd.extend(["--wandb.mode", request.wandb_mode]) 246 | if request.wandb_disable_artifact: 247 | cmd.append("--wandb.disable_artifact") 248 | 249 | # Environment configuration 250 | if request.env_type: 251 | cmd.extend(["--env.type", request.env_type]) 252 | if request.env_task: 253 | cmd.extend(["--env.task", request.env_task]) 254 | 255 | # Evaluation configuration 256 | cmd.extend(["--eval.n_episodes", str(request.eval_n_episodes)]) 257 | cmd.extend(["--eval.batch_size", str(request.eval_batch_size)]) 258 | if request.eval_use_async_envs: 259 | cmd.append("--eval.use_async_envs") 260 | 261 | # Optimizer configuration 262 | if request.optimizer_type: 263 | cmd.extend(["--optimizer.type", request.optimizer_type]) 264 | if request.optimizer_lr is not None: 265 | cmd.extend(["--optimizer.lr", str(request.optimizer_lr)]) 266 | if request.optimizer_weight_decay is not None: 267 | cmd.extend(["--optimizer.weight_decay", str(request.optimizer_weight_decay)]) 268 | if request.optimizer_grad_clip_norm is not None: 269 | cmd.extend(["--optimizer.grad_clip_norm", str(request.optimizer_grad_clip_norm)]) 270 | 271 | # Advanced options 272 | if request.use_policy_training_preset: 273 | cmd.append("--use_policy_training_preset") 274 | if request.config_path: 275 | cmd.extend(["--config_path", request.config_path]) 276 | 277 | return cmd 278 | 279 | def _start_monitoring(self): 280 | """Start monitoring threads for process output and status""" 281 | self._stop_monitoring.clear() 282 | 283 | # Start log monitoring thread 284 | self.log_thread = threading.Thread( 285 | target=self._monitor_logs, 286 | daemon=True 287 | ) 288 | self.log_thread.start() 289 | 290 | # Start status monitoring thread 291 | self.monitor_thread = threading.Thread( 292 | target=self._monitor_status, 293 | daemon=True 294 | ) 295 | self.monitor_thread.start() 296 | 297 | def _stop_monitoring_threads(self): 298 | """Stop all monitoring threads""" 299 | self._stop_monitoring.set() 300 | 301 | if self.log_thread and self.log_thread.is_alive(): 302 | self.log_thread.join(timeout=2) 303 | 304 | if self.monitor_thread and self.monitor_thread.is_alive(): 305 | self.monitor_thread.join(timeout=2) 306 | 307 | def _monitor_logs(self): 308 | """Monitor training process logs""" 309 | if not self.process: 310 | return 311 | 312 | try: 313 | for line in iter(self.process.stdout.readline, ''): 314 | if self._stop_monitoring.is_set(): 315 | break 316 | 317 | if line.strip(): 318 | # Parse training information from log line 319 | self._parse_log_line(line.strip()) 320 | 321 | # Add to log queue (keep last 1000 lines) 322 | if self.log_queue.qsize() >= 1000: 323 | try: 324 | self.log_queue.get_nowait() 325 | except queue.Empty: 326 | pass 327 | 328 | self.log_queue.put({ 329 | "timestamp": time.time(), 330 | "message": line.strip() 331 | }) 332 | 333 | except Exception as e: 334 | logger.error(f"Error monitoring logs: {e}") 335 | 336 | def _monitor_status(self): 337 | """Monitor training process status""" 338 | while not self._stop_monitoring.is_set() and self.process: 339 | try: 340 | # Check if process is still running 341 | if self.process.poll() is not None: 342 | # Process has ended 343 | break 344 | 345 | # Update available controls 346 | self.status.available_controls["stop_training"] = True 347 | 348 | time.sleep(1) 349 | 350 | except Exception as e: 351 | logger.error(f"Error monitoring status: {e}") 352 | break 353 | 354 | def _parse_log_line(self, line: str): 355 | """Parse training metrics from log line""" 356 | try: 357 | # Look for training metrics in the log line 358 | if "step:" in line.lower() and "loss:" in line.lower(): 359 | # Extract step number 360 | if "step:" in line: 361 | step_part = line.split("step:")[1].split()[0] 362 | try: 363 | self.status.current_step = int(step_part.replace(",", "")) 364 | except ValueError: 365 | pass 366 | 367 | # Extract loss 368 | if "loss:" in line: 369 | loss_part = line.split("loss:")[1].split()[0] 370 | try: 371 | self.status.current_loss = float(loss_part) 372 | except ValueError: 373 | pass 374 | 375 | # Extract learning rate 376 | if "lr:" in line: 377 | lr_part = line.split("lr:")[1].split()[0] 378 | try: 379 | self.status.current_lr = float(lr_part) 380 | except ValueError: 381 | pass 382 | 383 | # Extract gradient norm 384 | if "grdn:" in line: 385 | grdn_part = line.split("grdn:")[1].split()[0] 386 | try: 387 | self.status.grad_norm = float(grdn_part) 388 | except ValueError: 389 | pass 390 | 391 | # Calculate ETA 392 | if self.status.current_step > 0 and self.status.total_steps > 0: 393 | progress = self.status.current_step / self.status.total_steps 394 | if progress > 0: 395 | # Rough estimate based on current progress 396 | remaining_steps = self.status.total_steps - self.status.current_step 397 | # This is a very rough estimate - would need more sophisticated timing 398 | self.status.eta_seconds = remaining_steps * 0.5 # Assume 0.5s per step 399 | 400 | except Exception as e: 401 | logger.debug(f"Error parsing log line '{line}': {e}") 402 | 403 | def _reset_status(self): 404 | """Reset training status""" 405 | self.status = TrainingStatus() 406 | self.process = None 407 | 408 | # Global training manager instance 409 | training_manager = TrainingManager() 410 | 411 | # Handler functions for FastAPI endpoints 412 | def handle_start_training(request: TrainingRequest) -> Dict[str, Any]: 413 | """Handle start training request""" 414 | return training_manager.start_training(request) 415 | 416 | def handle_stop_training() -> Dict[str, Any]: 417 | """Handle stop training request""" 418 | return training_manager.stop_training() 419 | 420 | def handle_training_status() -> TrainingStatus: 421 | """Handle training status request""" 422 | return training_manager.get_status() 423 | 424 | def handle_training_logs() -> Dict[str, Any]: 425 | """Handle training logs request""" 426 | logs = training_manager.get_logs() 427 | return {"logs": logs} 428 | -------------------------------------------------------------------------------- /app/main.py: -------------------------------------------------------------------------------- 1 | from fastapi import FastAPI, WebSocket, WebSocketDisconnect 2 | from fastapi.staticfiles import StaticFiles 3 | from fastapi.responses import FileResponse 4 | from fastapi.middleware.cors import CORSMiddleware 5 | import os 6 | import logging 7 | import glob 8 | import asyncio 9 | from typing import List, Dict, Any 10 | import threading 11 | import queue 12 | from pathlib import Path 13 | from . import config 14 | 15 | # Import our custom recording functionality 16 | from .recording import ( 17 | RecordingRequest, 18 | UploadRequest, 19 | DatasetInfoRequest, 20 | handle_start_recording, 21 | handle_stop_recording, 22 | handle_exit_early, 23 | handle_rerecord_episode, 24 | handle_recording_status, 25 | handle_upload_dataset, 26 | handle_get_dataset_info, 27 | ) 28 | 29 | # Import our custom teleoperation functionality 30 | from .teleoperating import ( 31 | TeleoperateRequest, 32 | handle_start_teleoperation, 33 | handle_stop_teleoperation, 34 | handle_teleoperation_status, 35 | handle_get_joint_positions, 36 | ) 37 | 38 | # Import our custom calibration functionality 39 | from .calibrating import CalibrationRequest, calibration_manager 40 | 41 | # Import our custom training functionality 42 | from .training import ( 43 | TrainingRequest, 44 | handle_start_training, 45 | handle_stop_training, 46 | handle_training_status, 47 | handle_training_logs, 48 | ) 49 | 50 | # Import our custom replay functionality 51 | from .replaying import ( 52 | ReplayRequest, 53 | handle_start_replay, 54 | handle_stop_replay, 55 | handle_replay_status, 56 | handle_replay_logs, 57 | ) 58 | 59 | 60 | # Set up logging 61 | logging.basicConfig(level=logging.INFO) 62 | logger = logging.getLogger(__name__) 63 | 64 | # Global variables for WebSocket connections 65 | connected_websockets: List[WebSocket] = [] 66 | 67 | 68 | app = FastAPI() 69 | 70 | # Add CORS middleware 71 | app.add_middleware( 72 | CORSMiddleware, 73 | allow_origins=["*"], # Allows all origins 74 | allow_credentials=True, 75 | allow_methods=["*"], # Allows all methods 76 | allow_headers=["*"], # Allows all headers 77 | ) 78 | 79 | # Create static directory if it doesn't exist 80 | os.makedirs("app/static", exist_ok=True) 81 | 82 | # Mount the static directory 83 | app.mount("/static", StaticFiles(directory="app/static"), name="static") 84 | 85 | # Get the path to the lerobot root directory (3 levels up from this script) 86 | LEROBOT_PATH = str(Path(__file__).parent.parent.parent.parent) 87 | logger.info(f"LeRobot path: {LEROBOT_PATH}") 88 | 89 | # Import shared configuration constants 90 | from .config import ( 91 | CALIBRATION_BASE_PATH_TELEOP, 92 | CALIBRATION_BASE_PATH_ROBOTS, 93 | LEADER_CONFIG_PATH, 94 | FOLLOWER_CONFIG_PATH, 95 | find_available_ports, 96 | find_robot_port, 97 | detect_port_after_disconnect, 98 | save_robot_port, 99 | get_saved_robot_port, 100 | get_default_robot_port, 101 | ) 102 | 103 | 104 | class ConnectionManager: 105 | def __init__(self): 106 | self.active_connections: List[WebSocket] = [] 107 | self.broadcast_queue = queue.Queue() 108 | self.broadcast_thread = None 109 | self.is_running = False 110 | 111 | async def connect(self, websocket: WebSocket): 112 | await websocket.accept() 113 | self.active_connections.append(websocket) 114 | logger.info( 115 | f"WebSocket connected. Total connections: {len(self.active_connections)}" 116 | ) 117 | 118 | # Start broadcast thread if not running 119 | if not self.is_running: 120 | self.start_broadcast_thread() 121 | 122 | def disconnect(self, websocket: WebSocket): 123 | if websocket in self.active_connections: 124 | self.active_connections.remove(websocket) 125 | logger.info( 126 | f"WebSocket disconnected. Total connections: {len(self.active_connections)}" 127 | ) 128 | 129 | # Stop broadcast thread if no connections 130 | if not self.active_connections and self.is_running: 131 | self.stop_broadcast_thread() 132 | 133 | def start_broadcast_thread(self): 134 | """Start the background thread for broadcasting data""" 135 | if self.is_running: 136 | return 137 | 138 | self.is_running = True 139 | self.broadcast_thread = threading.Thread( 140 | target=self._broadcast_worker, daemon=True 141 | ) 142 | self.broadcast_thread.start() 143 | logger.info("📡 Broadcast thread started") 144 | 145 | def stop_broadcast_thread(self): 146 | """Stop the background thread""" 147 | self.is_running = False 148 | if self.broadcast_thread: 149 | self.broadcast_thread.join(timeout=1.0) 150 | logger.info("📡 Broadcast thread stopped") 151 | 152 | def _broadcast_worker(self): 153 | """Background worker thread for broadcasting WebSocket data""" 154 | import asyncio 155 | 156 | # Create a new event loop for this thread 157 | loop = asyncio.new_event_loop() 158 | asyncio.set_event_loop(loop) 159 | 160 | try: 161 | while self.is_running: 162 | try: 163 | # Get data from queue with timeout 164 | data = self.broadcast_queue.get(timeout=0.1) 165 | if data is None: # Poison pill to stop 166 | break 167 | 168 | # Broadcast to all connections 169 | if self.active_connections: 170 | loop.run_until_complete(self._send_to_all_connections(data)) 171 | 172 | except queue.Empty: 173 | continue 174 | except Exception as e: 175 | logger.error(f"Error in broadcast worker: {e}") 176 | 177 | finally: 178 | loop.close() 179 | 180 | async def _send_to_all_connections(self, data: Dict[str, Any]): 181 | """Send data to all active WebSocket connections""" 182 | if not self.active_connections: 183 | return 184 | 185 | disconnected = [] 186 | for connection in self.active_connections: 187 | try: 188 | await connection.send_json(data) 189 | except Exception as e: 190 | logger.error(f"Error sending data to WebSocket: {e}") 191 | disconnected.append(connection) 192 | 193 | # Remove disconnected connections 194 | for connection in disconnected: 195 | self.disconnect(connection) 196 | 197 | def broadcast_joint_data_sync(self, data: Dict[str, Any]): 198 | """Thread-safe method to queue data for broadcasting""" 199 | if self.is_running and self.active_connections: 200 | try: 201 | self.broadcast_queue.put_nowait(data) 202 | except queue.Full: 203 | logger.warning("Broadcast queue is full, dropping data") 204 | 205 | 206 | manager = ConnectionManager() 207 | 208 | 209 | @app.get("/") 210 | def read_root(): 211 | return FileResponse("app/static/index.html") 212 | 213 | 214 | @app.get("/get-configs") 215 | def get_configs(): 216 | # Get all available calibration configs 217 | leader_configs = [ 218 | os.path.basename(f) 219 | for f in glob.glob(os.path.join(LEADER_CONFIG_PATH, "*.json")) 220 | ] 221 | follower_configs = [ 222 | os.path.basename(f) 223 | for f in glob.glob(os.path.join(FOLLOWER_CONFIG_PATH, "*.json")) 224 | ] 225 | 226 | return {"leader_configs": leader_configs, "follower_configs": follower_configs} 227 | 228 | 229 | @app.post("/move-arm") 230 | def teleoperate_arm(request: TeleoperateRequest): 231 | """Start teleoperation of the robot arm""" 232 | return handle_start_teleoperation(request, manager) 233 | 234 | 235 | @app.post("/stop-teleoperation") 236 | def stop_teleoperation(): 237 | """Stop the current teleoperation session""" 238 | return handle_stop_teleoperation() 239 | 240 | 241 | @app.get("/teleoperation-status") 242 | def teleoperation_status(): 243 | """Get the current teleoperation status""" 244 | return handle_teleoperation_status() 245 | 246 | 247 | @app.get("/joint-positions") 248 | def get_joint_positions(): 249 | """Get current robot joint positions""" 250 | return handle_get_joint_positions() 251 | 252 | 253 | @app.get("/health") 254 | def health_check(): 255 | """Simple health check endpoint to verify server is running""" 256 | return {"status": "ok", "message": "FastAPI server is running"} 257 | 258 | 259 | @app.get("/ws-test") 260 | def websocket_test(): 261 | """Test endpoint to verify WebSocket support""" 262 | return {"websocket_endpoint": "/ws/joint-data", "status": "available"} 263 | 264 | 265 | @app.websocket("/ws/joint-data") 266 | async def websocket_endpoint(websocket: WebSocket): 267 | logger.info("🔗 New WebSocket connection attempt") 268 | try: 269 | await manager.connect(websocket) 270 | logger.info("✅ WebSocket connection established") 271 | 272 | while True: 273 | # Keep the connection alive and wait for messages 274 | try: 275 | data = await asyncio.wait_for(websocket.receive_text(), timeout=1.0) 276 | # Handle any incoming messages if needed 277 | logger.debug(f"Received WebSocket message: {data}") 278 | except asyncio.TimeoutError: 279 | # No message received, continue 280 | pass 281 | except WebSocketDisconnect: 282 | logger.info("🔌 WebSocket client disconnected") 283 | break 284 | 285 | # Small delay to prevent excessive CPU usage 286 | await asyncio.sleep(0.01) 287 | 288 | except WebSocketDisconnect: 289 | logger.info("🔌 WebSocket disconnected normally") 290 | except Exception as e: 291 | logger.error(f"❌ WebSocket error: {e}") 292 | finally: 293 | manager.disconnect(websocket) 294 | logger.info("🧹 WebSocket connection cleaned up") 295 | 296 | 297 | @app.post("/start-recording") 298 | def start_recording(request: RecordingRequest): 299 | """Start a dataset recording session""" 300 | return handle_start_recording(request, manager) 301 | 302 | 303 | @app.post("/stop-recording") 304 | def stop_recording(): 305 | """Stop the current recording session""" 306 | return handle_stop_recording() 307 | 308 | 309 | @app.get("/recording-status") 310 | def recording_status(): 311 | """Get the current recording status""" 312 | return handle_recording_status() 313 | 314 | 315 | @app.post("/recording-exit-early") 316 | def recording_exit_early(): 317 | """Skip to next episode (replaces right arrow key)""" 318 | return handle_exit_early() 319 | 320 | 321 | @app.post("/recording-rerecord-episode") 322 | def recording_rerecord_episode(): 323 | """Re-record current episode (replaces left arrow key)""" 324 | return handle_rerecord_episode() 325 | 326 | 327 | @app.post("/upload-dataset") 328 | def upload_dataset(request: UploadRequest): 329 | """Upload dataset to HuggingFace Hub""" 330 | return handle_upload_dataset(request) 331 | 332 | 333 | @app.post("/dataset-info") 334 | def get_dataset_info(request: DatasetInfoRequest): 335 | """Get information about a saved dataset""" 336 | return handle_get_dataset_info(request) 337 | 338 | 339 | # ============================================================================ 340 | # TRAINING ENDPOINTS 341 | # ============================================================================ 342 | 343 | 344 | @app.post("/start-training") 345 | def start_training(request: TrainingRequest): 346 | """Start a training session""" 347 | return handle_start_training(request) 348 | 349 | 350 | @app.post("/stop-training") 351 | def stop_training(): 352 | """Stop the current training session""" 353 | return handle_stop_training() 354 | 355 | 356 | @app.get("/training-status") 357 | def training_status(): 358 | """Get the current training status""" 359 | return handle_training_status() 360 | 361 | 362 | @app.get("/training-logs") 363 | def training_logs(): 364 | """Get recent training logs""" 365 | return handle_training_logs() 366 | 367 | 368 | # ============================================================================ 369 | # REPLAY ENDPOINTS 370 | # ============================================================================ 371 | 372 | 373 | @app.post("/start-replay") 374 | def start_replay(request: ReplayRequest): 375 | """Start a replay session""" 376 | return handle_start_replay(request) 377 | 378 | 379 | @app.post("/stop-replay") 380 | def stop_replay(): 381 | """Stop the current replay session""" 382 | return handle_stop_replay() 383 | 384 | 385 | @app.get("/replay-status") 386 | def replay_status(): 387 | """Get the current replay status""" 388 | return handle_replay_status() 389 | 390 | 391 | @app.get("/replay-logs") 392 | def replay_logs(): 393 | """Get recent replay logs""" 394 | return handle_replay_logs() 395 | 396 | 397 | # ============================================================================ 398 | # Calibration endpoints 399 | @app.post("/start-calibration") 400 | def start_calibration(request: CalibrationRequest): 401 | """Start calibration process""" 402 | return calibration_manager.start_calibration(request) 403 | 404 | 405 | @app.post("/stop-calibration") 406 | def stop_calibration(): 407 | """Stop calibration process""" 408 | return calibration_manager.stop_calibration_process() 409 | 410 | 411 | @app.get("/calibration-status") 412 | def calibration_status(): 413 | """Get current calibration status""" 414 | from dataclasses import asdict 415 | 416 | status = calibration_manager.get_status() 417 | return asdict(status) 418 | 419 | 420 | @app.post("/complete-calibration-step") 421 | def complete_calibration_step(): 422 | """Complete the current calibration step""" 423 | return calibration_manager.complete_step() 424 | 425 | 426 | @app.get("/calibration-configs/{device_type}") 427 | def get_calibration_configs(device_type: str): 428 | """Get all calibration config files for a specific device type""" 429 | try: 430 | if device_type == "robot": 431 | config_path = FOLLOWER_CONFIG_PATH 432 | elif device_type == "teleop": 433 | config_path = LEADER_CONFIG_PATH 434 | else: 435 | return {"success": False, "message": "Invalid device type"} 436 | 437 | # Get all JSON files in the config directory 438 | configs = [] 439 | if os.path.exists(config_path): 440 | for file in os.listdir(config_path): 441 | if file.endswith(".json"): 442 | config_name = os.path.splitext(file)[0] 443 | file_path = os.path.join(config_path, file) 444 | file_size = os.path.getsize(file_path) 445 | modified_time = os.path.getmtime(file_path) 446 | 447 | configs.append( 448 | { 449 | "name": config_name, 450 | "filename": file, 451 | "size": file_size, 452 | "modified": modified_time, 453 | } 454 | ) 455 | 456 | return {"success": True, "configs": configs, "device_type": device_type} 457 | 458 | except Exception as e: 459 | logger.error(f"Error getting calibration configs: {e}") 460 | return {"success": False, "message": str(e)} 461 | 462 | 463 | @app.delete("/calibration-configs/{device_type}/{config_name}") 464 | def delete_calibration_config(device_type: str, config_name: str): 465 | """Delete a calibration config file""" 466 | try: 467 | if device_type == "robot": 468 | config_path = FOLLOWER_CONFIG_PATH 469 | elif device_type == "teleop": 470 | config_path = LEADER_CONFIG_PATH 471 | else: 472 | return {"success": False, "message": "Invalid device type"} 473 | 474 | # Construct the file path 475 | filename = f"{config_name}.json" 476 | file_path = os.path.join(config_path, filename) 477 | 478 | # Check if file exists 479 | if not os.path.exists(file_path): 480 | return {"success": False, "message": "Configuration file not found"} 481 | 482 | # Delete the file 483 | os.remove(file_path) 484 | logger.info(f"Deleted calibration config: {file_path}") 485 | 486 | return { 487 | "success": True, 488 | "message": f"Configuration '{config_name}' deleted successfully", 489 | } 490 | 491 | except Exception as e: 492 | logger.error(f"Error deleting calibration config: {e}") 493 | return {"success": False, "message": str(e)} 494 | 495 | 496 | # ============================================================================ 497 | # PORT DETECTION ENDPOINTS 498 | # ============================================================================ 499 | 500 | @app.get("/available-ports") 501 | def get_available_ports(): 502 | """Get all available serial ports""" 503 | try: 504 | ports = find_available_ports() 505 | return {"status": "success", "ports": ports} 506 | except Exception as e: 507 | logger.error(f"Error getting available ports: {e}") 508 | return {"status": "error", "message": str(e)} 509 | 510 | 511 | @app.get("/available-cameras") 512 | def get_available_cameras(): 513 | """Get all available cameras""" 514 | try: 515 | # Try to detect cameras using OpenCV 516 | import cv2 517 | cameras = [] 518 | 519 | # Test up to 10 camera indices 520 | for i in range(10): 521 | cap = cv2.VideoCapture(i) 522 | if cap.isOpened(): 523 | ret, frame = cap.read() 524 | if ret: 525 | cameras.append({ 526 | "index": i, 527 | "name": f"Camera {i}", 528 | "available": True, 529 | "width": int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), 530 | "height": int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), 531 | "fps": int(cap.get(cv2.CAP_PROP_FPS)), 532 | }) 533 | cap.release() 534 | 535 | return {"status": "success", "cameras": cameras} 536 | except ImportError: 537 | # OpenCV not available, return empty list 538 | logger.warning("OpenCV not available for camera detection") 539 | return {"status": "success", "cameras": []} 540 | except Exception as e: 541 | logger.error(f"Error detecting cameras: {e}") 542 | return {"status": "error", "message": str(e), "cameras": []} 543 | 544 | 545 | @app.post("/start-port-detection") 546 | def start_port_detection(data: dict): 547 | """Start port detection process for a robot""" 548 | try: 549 | robot_type = data.get("robot_type", "robot") 550 | result = find_robot_port(robot_type) 551 | return {"status": "success", "data": result} 552 | except Exception as e: 553 | logger.error(f"Error starting port detection: {e}") 554 | return {"status": "error", "message": str(e)} 555 | 556 | 557 | @app.post("/detect-port-after-disconnect") 558 | def detect_port_after_disconnect_endpoint(data: dict): 559 | """Detect port after disconnection""" 560 | try: 561 | ports_before = data.get("ports_before", []) 562 | detected_port = detect_port_after_disconnect(ports_before) 563 | return {"status": "success", "port": detected_port} 564 | except Exception as e: 565 | logger.error(f"Error detecting port: {e}") 566 | return {"status": "error", "message": str(e)} 567 | 568 | 569 | @app.post("/save-robot-port") 570 | def save_robot_port_endpoint(data: dict): 571 | """Save a robot port for future use""" 572 | try: 573 | robot_type = data.get("robot_type") 574 | port = data.get("port") 575 | 576 | if not robot_type or not port: 577 | return {"status": "error", "message": "robot_type and port are required"} 578 | 579 | save_robot_port(robot_type, port) 580 | return {"status": "success", "message": f"Port {port} saved for {robot_type}"} 581 | except Exception as e: 582 | logger.error(f"Error saving robot port: {e}") 583 | return {"status": "error", "message": str(e)} 584 | 585 | 586 | @app.get("/robot-port/{robot_type}") 587 | def get_robot_port(robot_type: str): 588 | """Get the saved port for a robot type""" 589 | try: 590 | saved_port = get_saved_robot_port(robot_type) 591 | default_port = get_default_robot_port(robot_type) 592 | return { 593 | "status": "success", 594 | "saved_port": saved_port, 595 | "default_port": default_port 596 | } 597 | except Exception as e: 598 | logger.error(f"Error getting robot port: {e}") 599 | return {"status": "error", "message": str(e)} 600 | 601 | 602 | @app.post("/save-robot-config") 603 | def save_robot_config_endpoint(data: dict): 604 | """Save a robot configuration for future use""" 605 | try: 606 | robot_type = data.get("robot_type") 607 | config_name = data.get("config_name") 608 | 609 | if not robot_type or not config_name: 610 | return {"status": "error", "message": "Missing robot_type or config_name"} 611 | 612 | success = config.save_robot_config(robot_type, config_name) 613 | 614 | if success: 615 | return {"status": "success", "message": f"Configuration saved for {robot_type}"} 616 | else: 617 | return {"status": "error", "message": "Failed to save configuration"} 618 | 619 | except Exception as e: 620 | logger.error(f"Error saving robot configuration: {e}") 621 | return {"status": "error", "message": str(e)} 622 | 623 | 624 | @app.get("/robot-config/{robot_type}") 625 | def get_robot_config(robot_type: str, available_configs: str = ""): 626 | """Get the saved configuration for a robot type""" 627 | try: 628 | # Parse available configs from query parameter 629 | available_configs_list = [] 630 | if available_configs: 631 | available_configs_list = [cfg.strip() for cfg in available_configs.split(",") if cfg.strip()] 632 | 633 | saved_config = config.get_saved_robot_config(robot_type) 634 | default_config = config.get_default_robot_config(robot_type, available_configs_list) 635 | 636 | return { 637 | "status": "success", 638 | "saved_config": saved_config, 639 | "default_config": default_config 640 | } 641 | except Exception as e: 642 | logger.error(f"Error getting robot configuration: {e}") 643 | return {"status": "error", "message": str(e)} 644 | 645 | 646 | @app.on_event("shutdown") 647 | async def shutdown_event(): 648 | """Clean up resources when FastAPI shuts down""" 649 | logger.info("🔄 FastAPI shutting down, cleaning up...") 650 | 651 | # Stop any active recording - handled by recording module cleanup 652 | 653 | # Clean up replay resources 654 | from .replaying import cleanup as replay_cleanup 655 | 656 | replay_cleanup() 657 | 658 | if manager: 659 | manager.stop_broadcast_thread() 660 | logger.info("✅ Cleanup completed") 661 | -------------------------------------------------------------------------------- /app/calibrating.py: -------------------------------------------------------------------------------- 1 | """ 2 | Calibration module for the web interface. 3 | 4 | This module provides calibration functionality similar to the CLI calibrate.py, 5 | but adapted for the web interface with step-by-step guidance. 6 | """ 7 | 8 | import logging 9 | import threading 10 | import time 11 | import traceback 12 | from dataclasses import asdict, dataclass 13 | from typing import Optional, Dict, Any 14 | 15 | from lerobot.common.robots import ( 16 | Robot, 17 | RobotConfig, 18 | make_robot_from_config, 19 | ) 20 | from lerobot.common.teleoperators import ( 21 | Teleoperator, 22 | TeleoperatorConfig, 23 | make_teleoperator_from_config, 24 | ) 25 | from lerobot.common.motors import MotorCalibration 26 | from lerobot.common.motors.feetech import OperatingMode 27 | from lerobot.common.utils.utils import init_logging 28 | 29 | logger = logging.getLogger(__name__) 30 | 31 | 32 | @dataclass 33 | class CalibrationStatus: 34 | """Status information for calibration process""" 35 | calibration_active: bool = False 36 | status: str = "idle" # "idle", "connecting", "homing", "recording", "completed", "error", "stopping" 37 | device_type: Optional[str] = None 38 | error: Optional[str] = None 39 | message: str = "" 40 | step: int = 0 # Current calibration step 41 | total_steps: int = 2 # Total number of calibration steps 42 | current_positions: Dict[str, float] = None 43 | recorded_ranges: Dict[str, Dict[str, float]] = None # {motor: {min: val, max: val, current: val}} 44 | 45 | 46 | @dataclass 47 | class CalibrationRequest: 48 | """Request parameters for starting calibration""" 49 | device_type: str # "robot" or "teleop" 50 | port: str 51 | config_file: str 52 | 53 | 54 | class CalibrationManager: 55 | """Manages calibration process for the web interface""" 56 | 57 | def __init__(self): 58 | self.status = CalibrationStatus() 59 | self.device: Optional[Robot | Teleoperator] = None 60 | self.calibration_thread: Optional[threading.Thread] = None 61 | self.stop_calibration = False 62 | self._status_lock = threading.Lock() 63 | self._step_complete = threading.Event() 64 | self._recording_active = False 65 | self._start_positions = {} 66 | self._mins = {} 67 | self._maxes = {} 68 | self._homing_offsets = {} 69 | 70 | # Initialize logging 71 | init_logging() 72 | 73 | def get_status(self) -> CalibrationStatus: 74 | """Get current calibration status""" 75 | with self._status_lock: 76 | # Update current positions if we're recording and device is connected 77 | if self.status.status == "recording" and self.device and self.device.is_connected: 78 | try: 79 | # Try reading positions with quick retry on port contention 80 | positions = None 81 | for attempt in range(2): # Quick retry for status updates 82 | try: 83 | positions = self.device.bus.sync_read("Present_Position", normalize=False) 84 | break 85 | except Exception as read_error: 86 | if "Port is in use" in str(read_error) and attempt < 1: 87 | time.sleep(0.005) # Very short delay 88 | continue 89 | else: 90 | raise read_error 91 | 92 | if positions: 93 | # Update recorded ranges 94 | if not self.status.recorded_ranges: 95 | self.status.recorded_ranges = {} 96 | 97 | for motor, pos in positions.items(): 98 | # Filter out invalid readings (0, negative, or extreme values) 99 | if pos <= 0 or pos >= 5000: 100 | continue # Skip invalid readings 101 | 102 | if motor not in self.status.recorded_ranges: 103 | self.status.recorded_ranges[motor] = { 104 | "min": pos, "max": pos, "current": pos 105 | } 106 | else: 107 | self.status.recorded_ranges[motor]["current"] = pos 108 | self.status.recorded_ranges[motor]["min"] = min( 109 | self.status.recorded_ranges[motor]["min"], pos 110 | ) 111 | self.status.recorded_ranges[motor]["max"] = max( 112 | self.status.recorded_ranges[motor]["max"], pos 113 | ) 114 | except Exception as e: 115 | # Reduce log spam by using debug level for expected port contention 116 | if "Port is in use" in str(e): 117 | logger.debug(f"Port busy during position read: {e}") 118 | else: 119 | logger.warning(f"Failed to read positions: {e}") 120 | 121 | return self.status 122 | 123 | def _update_status(self, **kwargs): 124 | """Update calibration status thread-safely""" 125 | with self._status_lock: 126 | for key, value in kwargs.items(): 127 | if hasattr(self.status, key): 128 | setattr(self.status, key, value) 129 | 130 | def start_calibration(self, request: CalibrationRequest) -> Dict[str, Any]: 131 | """Start calibration process""" 132 | try: 133 | if self.status.calibration_active: 134 | return {"success": False, "message": "Calibration already active"} 135 | 136 | # Reset status and clear any previous calibration data 137 | self._start_positions = {} 138 | self._mins = {} 139 | self._maxes = {} 140 | self._homing_offsets = {} 141 | 142 | self._update_status( 143 | calibration_active=True, 144 | status="connecting", 145 | device_type=request.device_type, 146 | error=None, 147 | message=f"Starting calibration for {request.device_type}", 148 | step=0, 149 | current_positions=None, 150 | recorded_ranges=None 151 | ) 152 | 153 | # Start calibration in a separate thread 154 | self.calibration_thread = threading.Thread( 155 | target=self._calibration_worker, 156 | args=(request,), 157 | daemon=True 158 | ) 159 | self.stop_calibration = False 160 | self._step_complete.clear() 161 | self.calibration_thread.start() 162 | 163 | return {"success": True, "message": "Calibration started"} 164 | 165 | except Exception as e: 166 | logger.error(f"Error starting calibration: {e}") 167 | self._update_status( 168 | calibration_active=False, 169 | status="error", 170 | error=str(e), 171 | message="Failed to start calibration" 172 | ) 173 | return {"success": False, "message": str(e)} 174 | 175 | def complete_step(self) -> Dict[str, Any]: 176 | """Complete the current calibration step""" 177 | try: 178 | if not self.status.calibration_active: 179 | return {"success": False, "message": "No calibration active"} 180 | 181 | if self.status.status == "homing": 182 | # Complete homing step 183 | self._step_complete.set() 184 | return {"success": True, "message": "Homing position set"} 185 | 186 | elif self.status.status == "recording": 187 | # Complete recording step 188 | self._recording_active = False 189 | self._step_complete.set() 190 | return {"success": True, "message": "Range recording completed"} 191 | 192 | else: 193 | return {"success": False, "message": f"Cannot complete step in status: {self.status.status}"} 194 | 195 | except Exception as e: 196 | logger.error(f"Error completing step: {e}") 197 | return {"success": False, "message": str(e)} 198 | 199 | def stop_calibration_process(self) -> Dict[str, Any]: 200 | """Stop calibration process""" 201 | try: 202 | if not self.status.calibration_active: 203 | return {"success": False, "message": "No calibration active"} 204 | 205 | logger.info("Stopping calibration process...") 206 | self.stop_calibration = True 207 | self._recording_active = False 208 | self._step_complete.set() # Unblock any waiting step 209 | 210 | self._update_status( 211 | status="stopping", 212 | message="Stopping calibration..." 213 | ) 214 | 215 | # Wait for thread to finish 216 | if self.calibration_thread and self.calibration_thread.is_alive(): 217 | self.calibration_thread.join(timeout=5.0) 218 | 219 | # Ensure cleanup is called if thread didn't finish properly 220 | if self.calibration_thread and self.calibration_thread.is_alive(): 221 | logger.warning("Calibration thread did not finish within timeout, forcing cleanup") 222 | 223 | # Force cleanup and finish 224 | self._cleanup_and_finish("Calibration stopped", status="idle") 225 | 226 | logger.info("Calibration stop completed") 227 | return {"success": True, "message": "Calibration stopped"} 228 | 229 | except Exception as e: 230 | logger.error(f"Error stopping calibration: {e}") 231 | # Force cleanup on error too 232 | self._cleanup_and_finish("Calibration stopped with error", status="error") 233 | return {"success": False, "message": str(e)} 234 | 235 | def _calibration_worker(self, request: CalibrationRequest): 236 | """Worker thread for calibration process""" 237 | try: 238 | logger.info(f"Starting calibration worker for {request.device_type}") 239 | 240 | # Create device configuration 241 | if request.device_type == "robot": 242 | from lerobot.common.robots.so101_follower import SO101FollowerConfig 243 | config = SO101FollowerConfig( 244 | port=request.port, 245 | id=request.config_file 246 | ) 247 | elif request.device_type == "teleop": 248 | from lerobot.common.teleoperators.so101_leader import SO101LeaderConfig 249 | config = SO101LeaderConfig( 250 | port=request.port, 251 | id=request.config_file 252 | ) 253 | else: 254 | raise ValueError(f"Unknown device type: {request.device_type}") 255 | 256 | self._update_status( 257 | status="connecting", 258 | message="Connecting to device..." 259 | ) 260 | 261 | # Create and connect device 262 | if request.device_type == "robot": 263 | self.device = make_robot_from_config(config) 264 | else: 265 | self.device = make_teleoperator_from_config(config) 266 | 267 | logger.info("Connecting to device...") 268 | self.device.connect(calibrate=False) 269 | 270 | if self.stop_calibration: 271 | logger.info("Calibration stopped after device connection") 272 | self._cleanup_and_finish("Calibration cancelled") 273 | return 274 | 275 | # Start Step 1: Homing 276 | self._step_homing() 277 | 278 | if self.stop_calibration: 279 | logger.info("Calibration stopped after homing step") 280 | self._cleanup_and_finish("Calibration cancelled") 281 | return 282 | 283 | # Start Step 2: Range Recording 284 | self._step_range_recording() 285 | 286 | if self.stop_calibration: 287 | logger.info("Calibration stopped after recording step") 288 | self._cleanup_and_finish("Calibration cancelled") 289 | return 290 | 291 | # Complete calibration 292 | self._complete_calibration() 293 | 294 | logger.info("Calibration completed successfully") 295 | self._cleanup_and_finish("Calibration completed successfully", status="completed") 296 | 297 | except Exception as e: 298 | logger.error(f"Calibration error: {e}") 299 | logger.error(traceback.format_exc()) 300 | # Ensure cleanup happens even on error 301 | self._cleanup_and_finish(f"Calibration failed: {e}", status="error") 302 | finally: 303 | # Ensure we always clean up and reset the active flag 304 | logger.info("Calibration worker thread finishing") 305 | if self.status.calibration_active: 306 | logger.warning("Worker thread ending but calibration still marked as active - forcing cleanup") 307 | self._cleanup_and_finish("Calibration stopped", status="idle") 308 | 309 | def _step_homing(self): 310 | """Step 1: Set homing position""" 311 | logger.info("Starting homing step") 312 | 313 | # Disable torque to allow manual movement 314 | self.device.bus.disable_torque() 315 | for motor in self.device.bus.motors: 316 | self.device.bus.write("Operating_Mode", motor, OperatingMode.POSITION.value) 317 | 318 | self._update_status( 319 | status="homing", 320 | step=1, 321 | message="Move the device to the middle of its range of motion, then click 'Complete Step'" 322 | ) 323 | 324 | # Wait for user to complete step 325 | while not self._step_complete.is_set() and not self.stop_calibration: 326 | time.sleep(0.1) 327 | 328 | if self.stop_calibration: 329 | logger.info("Homing step cancelled due to stop request") 330 | return 331 | 332 | # Set homing offsets 333 | logger.info("Setting homing offsets...") 334 | self.device.bus.reset_calibration() 335 | actual_positions = self.device.bus.sync_read("Present_Position", normalize=False) 336 | logger.info(f"Current positions for homing: {actual_positions}") 337 | 338 | self._homing_offsets = self.device.bus._get_half_turn_homings(actual_positions) 339 | logger.info(f"Calculated homing offsets: {self._homing_offsets}") 340 | 341 | for motor, offset in self._homing_offsets.items(): 342 | self.device.bus.write("Homing_Offset", motor, offset) 343 | 344 | self._step_complete.clear() 345 | logger.info("Homing step completed") 346 | 347 | def _step_range_recording(self): 348 | """Step 2: Record range of motion""" 349 | logger.info("Starting range recording step") 350 | 351 | # Initialize range tracking with retry and validation 352 | self._start_positions = {} 353 | for attempt in range(5): # Try multiple times to get valid initial positions 354 | try: 355 | positions = self.device.bus.sync_read("Present_Position", normalize=False) 356 | # Validate initial positions 357 | valid_positions = {} 358 | for motor, pos in positions.items(): 359 | if pos > 0 and pos < 5000: # Valid range 360 | valid_positions[motor] = pos 361 | 362 | if len(valid_positions) == len(positions): # All positions are valid 363 | self._start_positions = valid_positions 364 | break 365 | else: 366 | logger.warning(f"Attempt {attempt + 1}: Got invalid initial positions, retrying...") 367 | time.sleep(0.1) 368 | except Exception as e: 369 | logger.warning(f"Attempt {attempt + 1}: Failed to read initial positions: {e}") 370 | time.sleep(0.1) 371 | 372 | if not self._start_positions: 373 | raise RuntimeError("Could not get valid initial positions after multiple attempts") 374 | 375 | logger.info(f"Starting positions for range recording: {self._start_positions}") 376 | 377 | self._mins = self._start_positions.copy() 378 | self._maxes = self._start_positions.copy() 379 | logger.info(f"Initialized mins: {self._mins}") 380 | logger.info(f"Initialized maxes: {self._maxes}") 381 | 382 | self._update_status( 383 | status="recording", 384 | step=2, 385 | message="Move ALL joints through their FULL ranges of motion - from minimum to maximum positions. Ensure each joint moves significantly from its starting position.", 386 | recorded_ranges={motor: {"min": pos, "max": pos, "current": pos} 387 | for motor, pos in self._start_positions.items()} 388 | ) 389 | 390 | self._recording_active = True 391 | 392 | # Record positions until user completes step 393 | while not self._step_complete.is_set() and not self.stop_calibration: 394 | try: 395 | # Try reading positions with retry on port contention 396 | positions = None 397 | for attempt in range(3): # Try up to 3 times 398 | try: 399 | positions = self.device.bus.sync_read("Present_Position", normalize=False) 400 | break # Success, exit retry loop 401 | except Exception as read_error: 402 | if "Port is in use" in str(read_error) and attempt < 2: 403 | time.sleep(0.01) # Short delay before retry 404 | continue 405 | else: 406 | raise read_error # Re-raise if not port contention or final attempt 407 | 408 | if positions: 409 | # Validate the readings - filter out invalid/zero values 410 | valid_positions = {} 411 | for motor, pos in positions.items(): 412 | # Filter out clearly invalid readings (0, negative, or extreme values) 413 | if pos > 0 and pos < 5000: # Reasonable range for motor positions 414 | valid_positions[motor] = pos 415 | else: 416 | logger.debug(f"Filtered invalid position for {motor}: {pos}") 417 | 418 | # Only update if we have valid readings 419 | if valid_positions: 420 | for motor, pos in valid_positions.items(): 421 | if motor in self._mins: 422 | self._mins[motor] = min(self._mins[motor], pos) 423 | self._maxes[motor] = max(self._maxes[motor], pos) 424 | 425 | time.sleep(0.05) # 20Hz update rate 426 | except Exception as e: 427 | if "Port is in use" in str(e): 428 | logger.debug(f"Port busy during position read: {e}") 429 | else: 430 | logger.warning(f"Error reading positions during recording: {e}") 431 | # Increase sleep time on error to reduce port contention 432 | time.sleep(0.2) 433 | 434 | if self.stop_calibration: 435 | logger.info("Range recording step cancelled due to stop request") 436 | return 437 | 438 | # Log the final recorded ranges for debugging 439 | logger.info("Final recorded ranges:") 440 | for motor in self._mins.keys(): 441 | logger.info(f" {motor}: min={self._mins[motor]}, max={self._maxes[motor]}, range={self._maxes[motor] - self._mins[motor]}") 442 | 443 | # Validate ranges 444 | same_min_max = [motor for motor in self._mins.keys() if self._mins[motor] == self._maxes[motor]] 445 | if same_min_max: 446 | raise ValueError(f"Some motors have the same min and max values: {same_min_max}") 447 | 448 | # Check for insufficient range movement (less than 100 motor steps) 449 | insufficient_range = [] 450 | for motor in self._mins.keys(): 451 | range_diff = self._maxes[motor] - self._mins[motor] 452 | if range_diff < 100: # Less than 100 motor steps seems insufficient 453 | insufficient_range.append(f"{motor}: {range_diff}") 454 | 455 | if insufficient_range: 456 | logger.warning(f"Some motors may not have been moved through sufficient range: {insufficient_range}") 457 | logger.warning("Consider moving all joints through their full range of motion during calibration") 458 | 459 | self._step_complete.clear() 460 | logger.info("Range recording step completed") 461 | 462 | def _complete_calibration(self): 463 | """Complete the calibration and save results""" 464 | logger.info("Completing calibration...") 465 | 466 | # Log motor information for debugging 467 | logger.info("Motor configuration:") 468 | for motor, m in self.device.bus.motors.items(): 469 | logger.info(f" {motor}: ID={m.id}, Model={m.model}") 470 | 471 | # Create calibration dict 472 | calibration = {} 473 | for motor, m in self.device.bus.motors.items(): 474 | calibration[motor] = MotorCalibration( 475 | id=m.id, 476 | drive_mode=0, 477 | homing_offset=self._homing_offsets[motor], 478 | range_min=self._mins[motor], 479 | range_max=self._maxes[motor], 480 | ) 481 | logger.info(f"Calibration for {motor}: " 482 | f"ID={m.id}, " 483 | f"homing_offset={self._homing_offsets[motor]}, " 484 | f"range_min={self._mins[motor]}, " 485 | f"range_max={self._maxes[motor]}") 486 | 487 | # Write and save calibration 488 | self.device.calibration = calibration 489 | self.device.bus.write_calibration(calibration) 490 | self.device._save_calibration() 491 | 492 | logger.info(f"Calibration saved to {self.device.calibration_fpath}") 493 | 494 | def _cleanup_and_finish(self, message: str, status: str = "completed"): 495 | """Clean up and finish calibration""" 496 | self._cleanup_device() 497 | self._recording_active = False 498 | self._update_status( 499 | calibration_active=False, 500 | status=status, 501 | message=message 502 | ) 503 | 504 | def _cleanup_device(self): 505 | """Clean up device connection""" 506 | try: 507 | if self.device: 508 | logger.info("Disconnecting device...") 509 | self.device.disconnect() 510 | self.device = None 511 | except Exception as e: 512 | logger.error(f"Error disconnecting device: {e}") 513 | 514 | 515 | # Global calibration manager instance 516 | calibration_manager = CalibrationManager() 517 | -------------------------------------------------------------------------------- /app/recording.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | import shutil 4 | from typing import Dict, Any, Callable 5 | from concurrent.futures import ThreadPoolExecutor 6 | from pydantic import BaseModel 7 | 8 | # Import the main record functionality to reuse it 9 | from lerobot.record import record, RecordConfig, DatasetRecordConfig 10 | from lerobot.common.robots.so101_follower import SO101FollowerConfig 11 | from lerobot.common.teleoperators.so101_leader import SO101LeaderConfig 12 | from lerobot.common.datasets.lerobot_dataset import LeRobotDataset 13 | 14 | # Import for patching the keyboard listener 15 | from lerobot.common.utils import control_utils 16 | import functools 17 | 18 | logger = logging.getLogger(__name__) 19 | 20 | # Import calibration paths from config (shared constants) 21 | from .config import ( 22 | CALIBRATION_BASE_PATH_TELEOP, 23 | CALIBRATION_BASE_PATH_ROBOTS, 24 | LEADER_CONFIG_PATH, 25 | FOLLOWER_CONFIG_PATH, 26 | setup_calibration_files 27 | ) 28 | 29 | # Global variables for recording state 30 | recording_active = False 31 | recording_thread = None 32 | recording_events = None # Events dict for controlling recording session 33 | recording_config = None # Store recording configuration 34 | recording_start_time = None # Track when recording started 35 | current_episode = 1 # Track current episode number 36 | saved_episodes = 0 # Track how many episodes have been saved 37 | current_phase = "preparing" # Track current phase: "preparing", "recording", "resetting", "completed" 38 | phase_start_time = None # Track when current phase started 39 | 40 | 41 | class RecordingRequest(BaseModel): 42 | leader_port: str 43 | follower_port: str 44 | leader_config: str 45 | follower_config: str 46 | dataset_repo_id: str 47 | single_task: str 48 | num_episodes: int = 5 49 | episode_time_s: int = 30 50 | reset_time_s: int = 10 51 | fps: int = 30 52 | video: bool = True 53 | push_to_hub: bool = False 54 | resume: bool = False 55 | cameras: dict = {} 56 | test_mode: bool = False # Skip robot connection for testing 57 | 58 | 59 | class UploadRequest(BaseModel): 60 | dataset_repo_id: str 61 | tags: list[str] = [] 62 | private: bool = False 63 | 64 | 65 | class DatasetInfoRequest(BaseModel): 66 | dataset_repo_id: str 67 | 68 | 69 | 70 | 71 | def create_record_config(request: RecordingRequest) -> RecordConfig: 72 | """Create a RecordConfig from the recording request""" 73 | from lerobot.common.cameras.opencv.configuration_opencv import OpenCVCameraConfig 74 | 75 | # Setup calibration files 76 | leader_config_name, follower_config_name = setup_calibration_files( 77 | request.leader_config, request.follower_config 78 | ) 79 | 80 | # 🔧 CAMERA CONFIG CONVERSION: Convert frontend camera dict to proper CameraConfig objects 81 | camera_configs = {} 82 | for camera_name, camera_data in request.cameras.items(): 83 | if camera_data.get("type") == "opencv": 84 | # Convert frontend format to OpenCVCameraConfig 85 | camera_configs[camera_name] = OpenCVCameraConfig( 86 | index_or_path=camera_data.get("camera_index", 0), 87 | fps=camera_data.get("fps"), 88 | width=camera_data.get("width"), 89 | height=camera_data.get("height"), 90 | ) 91 | logger.info(f"✅ CAMERA CONFIG: Converted {camera_name} -> OpenCVCameraConfig(index={camera_data.get('camera_index')}, {camera_data.get('width')}x{camera_data.get('height')}@{camera_data.get('fps')}fps)") 92 | else: 93 | logger.warning(f"⚠️ CAMERA CONFIG: Unsupported camera type '{camera_data.get('type')}' for {camera_name}") 94 | 95 | # Create robot config 96 | robot_config = SO101FollowerConfig( 97 | port=request.follower_port, 98 | id=follower_config_name, 99 | cameras=camera_configs, 100 | ) 101 | 102 | # Create teleop config 103 | teleop_config = SO101LeaderConfig( 104 | port=request.leader_port, 105 | id=leader_config_name, 106 | ) 107 | 108 | # Create dataset config 109 | dataset_config = DatasetRecordConfig( 110 | repo_id=request.dataset_repo_id, 111 | single_task=request.single_task, 112 | num_episodes=request.num_episodes, 113 | episode_time_s=request.episode_time_s, 114 | reset_time_s=request.reset_time_s, 115 | fps=request.fps, 116 | video=request.video, 117 | push_to_hub=request.push_to_hub, 118 | ) 119 | 120 | # Create the main record config 121 | record_config = RecordConfig( 122 | robot=robot_config, 123 | teleop=teleop_config, 124 | dataset=dataset_config, 125 | resume=request.resume, 126 | display_data=False, # Don't display data in API mode 127 | play_sounds=False, # Don't play sounds in API mode 128 | ) 129 | 130 | return record_config 131 | 132 | 133 | def handle_start_recording(request: RecordingRequest, websocket_manager=None) -> Dict[str, Any]: 134 | """Handle start recording request by using the existing record() function""" 135 | global recording_active, recording_thread, recording_events, recording_config, recording_start_time, current_episode, saved_episodes, current_phase, phase_start_time 136 | 137 | if recording_active: 138 | return {"success": False, "message": "Recording is already active"} 139 | 140 | # 🧹 CLEANUP: Reset all global state from previous sessions 141 | logger.info("🧹 CLEANUP: Resetting all recording state variables") 142 | recording_active = False 143 | recording_thread = None 144 | recording_events = None 145 | recording_config = None 146 | recording_start_time = None 147 | current_episode = 1 148 | saved_episodes = 0 149 | current_phase = "preparing" 150 | phase_start_time = None 151 | 152 | try: 153 | import time 154 | logger.info(f"Starting recording for dataset: {request.dataset_repo_id}") 155 | logger.info(f"Task: {request.single_task}") 156 | 157 | # Store recording configuration and reset episode counter 158 | recording_config = request 159 | recording_start_time = None # Will be set when recording actually starts 160 | current_episode = 1 161 | current_phase = "preparing" 162 | phase_start_time = None 163 | 164 | # Initialize recording events for web control (replaces keyboard controls) 165 | recording_events = { 166 | "exit_early": False, # Right arrow key -> "Skip to next episode" button 167 | "stop_recording": False, # ESC key -> "Stop recording" button 168 | "rerecord_episode": False # Left arrow key -> "Re-record episode" button 169 | } 170 | 171 | # Create the record configuration 172 | record_config = create_record_config(request) 173 | 174 | # Start recording in a separate thread 175 | def recording_worker(): 176 | global recording_active, recording_start_time, current_phase, phase_start_time, current_episode, saved_episodes 177 | recording_active = True 178 | recording_start_time = time.time() # Set start time when recording actually begins 179 | 180 | # Initialize episode counters 181 | current_episode = 1 182 | saved_episodes = 0 183 | 184 | try: 185 | logger.info(f"Starting recording worker with events: {recording_events}") 186 | print(f"🚀 STATUS CHANGE: Recording session started for dataset '{request.dataset_repo_id}'") 187 | print(f"📋 STATUS CHANGE: Task: '{request.single_task}' - {request.num_episodes} episodes planned") 188 | 189 | # 🔓 CRITICAL: Wait for camera streams to be fully released by frontend 190 | if request.cameras: 191 | logger.info(f"🔓 BACKEND: Waiting for camera resources to be released (cameras configured: {list(request.cameras.keys())})") 192 | print(f"🔓 STATUS CHANGE: Waiting for camera resources to be released...") 193 | time.sleep(2.0) # Give cameras more time to be fully released 194 | logger.info(f"✅ BACKEND: Camera wait period complete, proceeding with robot initialization") 195 | 196 | # Use the original record() function but with web-controlled events 197 | dataset = record_with_web_events(record_config, recording_events) 198 | logger.info(f"Recording completed successfully. Dataset has {dataset.num_episodes} episodes") 199 | print(f"🎉 STATUS CHANGE: Recording session completed successfully with {dataset.num_episodes} episodes") 200 | return {"success": True, "episodes": dataset.num_episodes} 201 | except Exception as e: 202 | logger.error(f"❌ CRITICAL ERROR during recording: {e}") 203 | print(f"❌ STATUS CHANGE: Recording session failed with error: {str(e)}") 204 | import traceback 205 | logger.error(f"Full traceback: {traceback.format_exc()}") 206 | 207 | # 🚨 CRITICAL: Set phase to "error" instead of "completed" to distinguish failures 208 | current_phase = "error" 209 | recording_active = False 210 | recording_start_time = None 211 | phase_start_time = None 212 | 213 | return {"success": False, "error": str(e)} 214 | finally: 215 | # Only set to completed if no error occurred 216 | if current_phase != "error": 217 | current_phase = "completed" 218 | 219 | recording_active = False 220 | recording_start_time = None 221 | phase_start_time = None 222 | current_episode = 1 # Reset for next session 223 | saved_episodes = 0 # Reset for next session 224 | logger.info("🔚 RECORDING SESSION: Setting state to completed - frontend should stop polling") 225 | print(f"🔚 STATUS CHANGE: Recording session ended") 226 | 227 | recording_thread = ThreadPoolExecutor(max_workers=1) 228 | future = recording_thread.submit(recording_worker) 229 | 230 | return { 231 | "success": True, 232 | "message": "Recording started successfully", 233 | "dataset_id": request.dataset_repo_id, 234 | "num_episodes": request.num_episodes 235 | } 236 | 237 | except Exception as e: 238 | recording_active = False 239 | logger.error(f"Failed to start recording: {e}") 240 | return {"success": False, "message": f"Failed to start recording: {str(e)}"} 241 | 242 | 243 | def handle_stop_recording() -> Dict[str, Any]: 244 | """Handle stop recording request - replaces ESC key""" 245 | global recording_active, recording_thread, recording_events, current_phase, phase_start_time 246 | 247 | if not recording_active or recording_events is None: 248 | return {"success": False, "message": "No recording session is active"} 249 | 250 | try: 251 | # Trigger the stop recording event (replaces ESC key) 252 | recording_events["stop_recording"] = True 253 | recording_events["exit_early"] = True 254 | 255 | # Update phase to indicate stopping 256 | current_phase = "stopping" 257 | phase_start_time = None 258 | 259 | logger.info("Stop recording triggered from web interface") 260 | print("🛑 STATUS CHANGE: Stop recording requested - session will end soon") 261 | 262 | return { 263 | "success": True, 264 | "message": "Recording stop requested successfully", 265 | "session_ending": True # Signal that session is ending 266 | } 267 | 268 | except Exception as e: 269 | logger.error(f"Error stopping recording: {e}") 270 | return {"success": False, "message": f"Failed to stop recording: {str(e)}"} 271 | 272 | 273 | def handle_exit_early() -> Dict[str, Any]: 274 | """Handle exit early request - replaces right arrow key""" 275 | global recording_events, current_phase 276 | 277 | if not recording_active or recording_events is None: 278 | return {"success": False, "message": "No recording session is active"} 279 | 280 | try: 281 | # Log the current state before setting the flag 282 | logger.info(f"Exit early requested - Current phase: {current_phase}") 283 | logger.info(f"Events before setting exit_early: {recording_events}") 284 | 285 | # Trigger the exit early event (replaces right arrow key) 286 | recording_events["exit_early"] = True 287 | # Also set our tracking flag that won't be reset by record_loop 288 | recording_events["_exit_early_triggered"] = True 289 | 290 | # Log the state after setting the flag 291 | logger.info(f"Exit early flag set - Events after: {recording_events}") 292 | logger.info(f"Exit early triggered from web interface (current phase: {current_phase})") 293 | 294 | phase_name = "recording phase" if current_phase == "recording" else "reset phase" 295 | return { 296 | "success": True, 297 | "message": f"Exit early triggered successfully for {phase_name}", 298 | "current_phase": current_phase, 299 | "events_state": dict(recording_events) # Include events state in response 300 | } 301 | 302 | except Exception as e: 303 | logger.error(f"Error triggering exit early: {e}") 304 | import traceback 305 | logger.error(f"Full traceback: {traceback.format_exc()}") 306 | return {"success": False, "message": f"Failed to trigger exit early: {str(e)}"} 307 | 308 | 309 | def handle_rerecord_episode() -> Dict[str, Any]: 310 | """Handle rerecord episode request - replaces left arrow key""" 311 | global recording_events 312 | 313 | if not recording_active or recording_events is None: 314 | return {"success": False, "message": "No recording session is active"} 315 | 316 | try: 317 | # Log the current state before setting the flags 318 | logger.info(f"Re-record episode requested - Events before: {recording_events}") 319 | 320 | # Trigger the rerecord episode event (replaces left arrow key) 321 | recording_events["rerecord_episode"] = True 322 | recording_events["exit_early"] = True # Also need to exit current loop 323 | 324 | # Log the state after setting the flags 325 | logger.info(f"Re-record flags set - Events after: {recording_events}") 326 | logger.info("Re-record episode triggered from web interface") 327 | 328 | return { 329 | "success": True, 330 | "message": "Re-record episode requested successfully", 331 | "events_state": dict(recording_events) # Include events state in response 332 | } 333 | 334 | except Exception as e: 335 | logger.error(f"Error triggering rerecord episode: {e}") 336 | import traceback 337 | logger.error(f"Full traceback: {traceback.format_exc()}") 338 | return {"success": False, "message": f"Failed to trigger rerecord episode: {str(e)}"} 339 | 340 | 341 | def handle_recording_status() -> Dict[str, Any]: 342 | """Handle recording status request""" 343 | import time 344 | 345 | # If recording is not active and phase is completed or error, indicate session has ended 346 | session_ended = not recording_active and current_phase in ["completed", "error"] 347 | 348 | # Log when session has ended to help debug frontend polling 349 | if session_ended: 350 | if current_phase == "error": 351 | logger.info("📡 RECORDING STATUS REQUEST: Session failed with error - frontend should stop polling") 352 | print("📡 STATUS CHANGE: Frontend is still polling after session error - should stop now") 353 | else: 354 | logger.info("📡 RECORDING STATUS REQUEST: Session has ended - frontend should stop polling") 355 | print("📡 STATUS CHANGE: Frontend is still polling after session end - should stop now") 356 | 357 | status = { 358 | "recording_active": recording_active, 359 | "current_phase": current_phase, # "preparing", "recording", "resetting", "completed" 360 | "session_ended": session_ended, # New field to indicate session completion 361 | "available_controls": { 362 | "stop_recording": recording_active, # ESC key replacement 363 | "exit_early": recording_active, # Right arrow key replacement 364 | "rerecord_episode": recording_active and current_phase == "recording" # Only during recording phase 365 | }, 366 | "message": "Recording session failed with error - check logs" if current_phase == "error" else ("Recording session has ended - stop polling" if session_ended else "Recording status retrieved successfully") 367 | } 368 | 369 | # Add episode information if recording is active 370 | if recording_active and recording_config: 371 | status["current_episode"] = current_episode 372 | status["total_episodes"] = recording_config.num_episodes 373 | status["saved_episodes"] = saved_episodes # Track completed episodes 374 | 375 | # Add session start time if available 376 | if recording_start_time: 377 | status["session_start_time"] = recording_start_time 378 | status["session_elapsed_seconds"] = int(time.time() - recording_start_time) 379 | 380 | # Add phase timing information 381 | if phase_start_time: 382 | status["phase_start_time"] = phase_start_time 383 | status["phase_elapsed_seconds"] = int(time.time() - phase_start_time) 384 | 385 | # Add phase time limits 386 | if current_phase == "recording": 387 | status["phase_time_limit_s"] = recording_config.episode_time_s 388 | elif current_phase == "resetting": 389 | status["phase_time_limit_s"] = recording_config.reset_time_s 390 | 391 | return status 392 | 393 | 394 | # For backward compatibility, in case we want to add frame modifications later 395 | def add_custom_frame_modifier(modifier_func: Callable[[Dict[str, Any]], Dict[str, Any]]): 396 | """Placeholder for future custom frame modifications""" 397 | logger.info("Custom frame modifier registered (not yet implemented in simplified version)") 398 | 399 | 400 | def add_timestamp_modifier(): 401 | """Placeholder for timestamp modifier""" 402 | logger.info("Timestamp modifier registered (not yet implemented in simplified version)") 403 | 404 | 405 | def add_debug_info_modifier(): 406 | """Placeholder for debug info modifier""" 407 | logger.info("Debug info modifier registered (not yet implemented in simplified version)") 408 | 409 | 410 | def handle_get_dataset_info(request: DatasetInfoRequest) -> Dict[str, Any]: 411 | """Get information about a saved dataset""" 412 | try: 413 | # Import LeRobotDataset to load the dataset 414 | from lerobot.common.datasets.lerobot_dataset import LeRobotDataset 415 | import time 416 | 417 | logger.info(f"Loading dataset {request.dataset_repo_id} to get info") 418 | 419 | # Load the dataset from local storage 420 | dataset = LeRobotDataset(request.dataset_repo_id) 421 | 422 | logger.info(f"Dataset loaded with {dataset.num_episodes} episodes") 423 | 424 | # Get dataset metadata 425 | dataset_info = { 426 | "success": True, 427 | "dataset_repo_id": request.dataset_repo_id, 428 | "num_episodes": dataset.num_episodes, 429 | "single_task": getattr(dataset.meta, 'single_task', 'Unknown task'), 430 | "fps": dataset.fps, 431 | "features": list(dataset.features.keys()), 432 | "total_frames": dataset.total_frames, 433 | "robot_type": getattr(dataset.meta, 'robot_type', 'Unknown robot'), 434 | } 435 | 436 | # Try to get task information from the first episode if available 437 | if dataset.num_episodes > 0: 438 | try: 439 | first_episode = dataset[0] 440 | if 'task' in first_episode: 441 | dataset_info["single_task"] = first_episode['task'][0] if len(first_episode['task']) > 0 else dataset_info["single_task"] 442 | except Exception as e: 443 | logger.warning(f"Could not extract task from first episode: {e}") 444 | 445 | logger.info(f"Dataset info retrieved: {dataset_info}") 446 | return dataset_info 447 | 448 | except Exception as e: 449 | logger.error(f"Error loading dataset {request.dataset_repo_id}: {e}") 450 | import traceback 451 | logger.error(f"Full traceback: {traceback.format_exc()}") 452 | return { 453 | "success": False, 454 | "message": f"Failed to load dataset: {str(e)}" 455 | } 456 | 457 | 458 | def handle_upload_dataset(request: UploadRequest) -> Dict[str, Any]: 459 | """Handle dataset upload to HuggingFace Hub""" 460 | try: 461 | # Import LeRobotDataset to load and upload the dataset 462 | from lerobot.common.datasets.lerobot_dataset import LeRobotDataset 463 | 464 | logger.info(f"Loading dataset {request.dataset_repo_id} for upload") 465 | 466 | # Load the dataset from local storage 467 | dataset = LeRobotDataset(request.dataset_repo_id) 468 | 469 | logger.info(f"Dataset loaded with {dataset.num_episodes} episodes") 470 | logger.info(f"Uploading to HuggingFace Hub with tags: {request.tags}, private: {request.private}") 471 | 472 | # Upload dataset to HuggingFace Hub 473 | dataset.push_to_hub( 474 | tags=request.tags, 475 | private=request.private 476 | ) 477 | 478 | logger.info(f"Dataset {request.dataset_repo_id} uploaded successfully to HuggingFace Hub") 479 | 480 | return { 481 | "success": True, 482 | "message": f"Dataset {request.dataset_repo_id} uploaded successfully to HuggingFace Hub", 483 | "dataset_url": f"https://huggingface.co/datasets/{request.dataset_repo_id}", 484 | "num_episodes": dataset.num_episodes 485 | } 486 | 487 | except Exception as e: 488 | logger.error(f"Error uploading dataset {request.dataset_repo_id}: {e}") 489 | import traceback 490 | logger.error(f"Full traceback: {traceback.format_exc()}") 491 | return { 492 | "success": False, 493 | "message": f"Failed to upload dataset: {str(e)}" 494 | } 495 | 496 | 497 | def record_with_web_events(cfg: RecordConfig, web_events: dict) -> LeRobotDataset: 498 | """ 499 | Implement recording with phase tracking - exactly mirrors original record() function behavior 500 | """ 501 | import time 502 | from lerobot.common.utils.utils import log_say 503 | from lerobot.common.datasets.lerobot_dataset import LeRobotDataset 504 | from lerobot.common.datasets.utils import hw_to_dataset_features 505 | from lerobot.common.robots import make_robot_from_config 506 | from lerobot.common.teleoperators import make_teleoperator_from_config 507 | from lerobot.common.utils.control_utils import sanity_check_dataset_name, sanity_check_dataset_robot_compatibility 508 | from lerobot.common.policies.factory import make_policy 509 | from lerobot.common.datasets.image_writer import safe_stop_image_writer 510 | 511 | global current_phase, phase_start_time, current_episode, saved_episodes 512 | 513 | # Import the record_loop function from lerobot.record 514 | from lerobot.record import record_loop 515 | 516 | robot = make_robot_from_config(cfg.robot) 517 | teleop = make_teleoperator_from_config(cfg.teleop) if cfg.teleop is not None else None 518 | 519 | action_features = hw_to_dataset_features(robot.action_features, "action", cfg.dataset.video) 520 | obs_features = hw_to_dataset_features(robot.observation_features, "observation", cfg.dataset.video) 521 | dataset_features = {**action_features, **obs_features} 522 | 523 | if cfg.resume: 524 | dataset = LeRobotDataset( 525 | cfg.dataset.repo_id, 526 | root=cfg.dataset.root, 527 | ) 528 | if hasattr(robot, "cameras") and len(robot.cameras) > 0: 529 | dataset.start_image_writer( 530 | num_processes=cfg.dataset.num_image_writer_processes, 531 | num_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras), 532 | ) 533 | sanity_check_dataset_robot_compatibility(dataset, robot, cfg.dataset.fps, dataset_features) 534 | else: 535 | sanity_check_dataset_name(cfg.dataset.repo_id, cfg.policy) 536 | dataset = LeRobotDataset.create( 537 | cfg.dataset.repo_id, 538 | cfg.dataset.fps, 539 | root=cfg.dataset.root, 540 | robot_type=robot.name, 541 | features=dataset_features, 542 | use_videos=cfg.dataset.video, 543 | image_writer_processes=cfg.dataset.num_image_writer_processes, 544 | image_writer_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras), 545 | ) 546 | 547 | # Load pretrained policy 548 | policy = None if cfg.policy is None else make_policy(cfg.policy, ds_meta=dataset.meta) 549 | 550 | # 🔧 ROBOT CONNECTION: Connect with enhanced error handling for camera conflicts 551 | try: 552 | logger.info("🔧 ROBOT CONNECTION: Attempting to connect robot...") 553 | robot.connect() 554 | logger.info("✅ ROBOT CONNECTION: Robot connected successfully") 555 | except Exception as e: 556 | logger.error(f"❌ ROBOT CONNECTION: Failed to connect robot: {e}") 557 | # If robot connection fails due to camera conflict, provide clear error 558 | if "camera" in str(e).lower() or "device" in str(e).lower() or "busy" in str(e).lower(): 559 | logger.error("💡 ROBOT CONNECTION: Camera connection failure - likely camera resource conflict") 560 | logger.error("💡 ROBOT CONNECTION: Make sure frontend camera streams are released before recording") 561 | raise 562 | 563 | if teleop is not None: 564 | try: 565 | logger.info("🔧 TELEOP CONNECTION: Attempting to connect teleoperator...") 566 | teleop.connect() 567 | logger.info("✅ TELEOP CONNECTION: Teleoperator connected successfully") 568 | except Exception as e: 569 | logger.error(f"❌ TELEOP CONNECTION: Failed to connect teleoperator: {e}") 570 | raise 571 | 572 | # Ensure calibration is properly loaded and applied to the devices 573 | logger.info(f"Applying calibration to devices") 574 | 575 | # Write calibration to motors' memory (similar to teleoperation code) 576 | if hasattr(robot, 'bus') and robot.calibration is not None: 577 | try: 578 | logger.info(f"Writing robot calibration to motors...") 579 | robot.bus.write_calibration(robot.calibration) 580 | logger.info(f"Robot calibration applied successfully") 581 | except Exception as e: 582 | logger.error(f"Error writing robot calibration: {e}") 583 | else: 584 | logger.warning(f"Robot bus or calibration not available - calibration may not be applied") 585 | 586 | if teleop is not None and hasattr(teleop, 'bus') and teleop.calibration is not None: 587 | try: 588 | logger.info(f"Writing teleop calibration to motors...") 589 | teleop.bus.write_calibration(teleop.calibration) 590 | logger.info(f"Teleop calibration applied successfully") 591 | except Exception as e: 592 | logger.error(f"Error writing teleop calibration: {e}") 593 | else: 594 | logger.warning(f"Teleop bus or calibration not available - calibration may not be applied") 595 | 596 | # Start with episode 1 - but track it properly 597 | current_episode = 1 598 | saved_episodes = 0 # Track how many episodes we've actually saved 599 | 600 | try: 601 | while saved_episodes < cfg.dataset.num_episodes: 602 | # RECORDING PHASE - with dataset (matches original record.py exactly) 603 | current_phase = "recording" 604 | phase_start_time = time.time() 605 | logger.info(f"Starting recording phase for episode {current_episode}") 606 | logger.info(f"Events state at start of recording phase: {web_events}") 607 | print(f"🎬 STATUS CHANGE: Starting recording phase for episode {current_episode}/{cfg.dataset.num_episodes}") 608 | 609 | log_say(f"Recording episode {current_episode}", cfg.play_sounds) 610 | 611 | # Add a tracking flag that won't be reset by record_loop 612 | web_events["_exit_early_triggered"] = False 613 | logger.info(f"Recording phase - calling record_loop with events: {web_events}") 614 | 615 | record_loop( 616 | robot=robot, 617 | events=web_events, 618 | fps=cfg.dataset.fps, 619 | teleop=teleop, 620 | policy=policy, 621 | dataset=dataset, 622 | control_time_s=cfg.dataset.episode_time_s, 623 | single_task=cfg.dataset.single_task, 624 | display_data=cfg.display_data, 625 | ) 626 | 627 | logger.info(f"Recording phase completed - events state: {web_events}") 628 | 629 | # Check if exit_early was triggered (use our tracking flag) 630 | recording_interrupted_by_exit_early = web_events.get("_exit_early_triggered", False) 631 | if recording_interrupted_by_exit_early: 632 | logger.info("🟡 RECORDING PHASE INTERRUPTED BY EXIT_EARLY - proceeding to save episode") 633 | print(f"🟡 STATUS CHANGE: Recording phase interrupted by user - episode {current_episode} data collected") 634 | # Reset our tracking flag 635 | web_events["_exit_early_triggered"] = False 636 | else: 637 | # Recording completed due to timeout - trigger re-record behavior 638 | logger.info("⏰ RECORDING PHASE COMPLETED DUE TO TIMEOUT - triggering re-record") 639 | print(f"⏰ STATUS CHANGE: Recording timeout reached for episode {current_episode} - re-recording") 640 | web_events["rerecord_episode"] = True 641 | 642 | # Handle rerecord logic first (before saving) 643 | if web_events["rerecord_episode"]: 644 | log_say("Re-record episode", cfg.play_sounds) 645 | print(f"🔄 STATUS CHANGE: Re-recording episode {current_episode} (episode number stays the same)") 646 | web_events["rerecord_episode"] = False 647 | web_events["exit_early"] = False 648 | dataset.clear_episode_buffer() 649 | 650 | # Go through reset phase before re-recording (don't increment episode counters) 651 | # RESET PHASE - without dataset (matches original record.py exactly) 652 | current_phase = "resetting" 653 | phase_start_time = time.time() 654 | logger.info(f"Starting reset phase for re-record of episode {current_episode}") 655 | logger.info(f"Events state at start of reset phase: {web_events}") 656 | print(f"🔄 STATUS CHANGE: Starting reset phase for episode {current_episode}") 657 | 658 | log_say("Reset the environment", cfg.play_sounds) 659 | 660 | # Reset exit_early flag at the start of each phase 661 | web_events["exit_early"] = False 662 | logger.info(f"Reset phase - calling record_loop with events: {web_events}") 663 | 664 | record_loop( 665 | robot=robot, 666 | events=web_events, 667 | fps=cfg.dataset.fps, 668 | teleop=teleop, 669 | # NOTE: NO dataset parameter here - matches LeRobot CLI exactly 670 | # This means NO recording happens during reset phase 671 | control_time_s=cfg.dataset.reset_time_s, 672 | single_task=cfg.dataset.single_task, 673 | display_data=cfg.display_data, 674 | ) 675 | 676 | logger.info(f"Reset phase completed - events state: {web_events}") 677 | 678 | # Check if reset was interrupted by exit_early 679 | if web_events["exit_early"]: 680 | logger.info("🟡 RESET PHASE INTERRUPTED BY EXIT_EARLY during re-record") 681 | print(f"🟡 STATUS CHANGE: Reset phase interrupted by user during re-record") 682 | web_events["exit_early"] = False 683 | 684 | # Check if stop recording was requested during re-record reset phase 685 | if web_events["stop_recording"]: 686 | logger.info("🛑 STOP RECORDING requested during re-record reset phase - ending session") 687 | print(f"🛑 STATUS CHANGE: Stop recording requested during re-record reset - ending session") 688 | break 689 | 690 | # Don't increment current_episode or saved_episodes - we're re-recording the same episode 691 | continue 692 | 693 | # Save episode immediately after recording phase (matches expected flow) 694 | logger.info(f"💾 Saving episode {current_episode}...") 695 | print(f"💾 STATUS CHANGE: Saving episode {current_episode}") 696 | dataset.save_episode() 697 | logger.info(f"✅ Episode {current_episode} saved successfully") 698 | print(f"✅ STATUS CHANGE: Episode {current_episode} saved successfully") 699 | 700 | # Increment episode counters after successful save 701 | saved_episodes += 1 702 | current_episode += 1 703 | 704 | # Check if we should stop recording 705 | if web_events["stop_recording"]: 706 | print(f"🛑 STATUS CHANGE: Recording manually stopped by user") 707 | break 708 | 709 | # Check if we've completed all episodes 710 | if saved_episodes >= cfg.dataset.num_episodes: 711 | break 712 | 713 | # Execute reset phase to prepare for next episode 714 | # Skip reset for the last episode that was just saved 715 | if saved_episodes < cfg.dataset.num_episodes: 716 | # RESET PHASE - without dataset (matches original record.py exactly) 717 | current_phase = "resetting" 718 | phase_start_time = time.time() 719 | logger.info(f"Starting reset phase for next episode {current_episode}") 720 | logger.info(f"Events state at start of reset phase: {web_events}") 721 | print(f"🔄 STATUS CHANGE: Starting reset phase for episode {current_episode}") 722 | 723 | log_say("Reset the environment", cfg.play_sounds) 724 | 725 | # Reset exit_early flag at the start of each phase 726 | web_events["exit_early"] = False 727 | logger.info(f"Reset phase - calling record_loop with events: {web_events}") 728 | 729 | record_loop( 730 | robot=robot, 731 | events=web_events, 732 | fps=cfg.dataset.fps, 733 | teleop=teleop, 734 | # NOTE: NO dataset parameter here - matches LeRobot CLI exactly 735 | # This means NO recording happens during reset phase 736 | control_time_s=cfg.dataset.reset_time_s, 737 | single_task=cfg.dataset.single_task, 738 | display_data=cfg.display_data, 739 | ) 740 | 741 | logger.info(f"Reset phase completed - events state: {web_events}") 742 | 743 | # Check if reset was interrupted by exit_early 744 | if web_events["exit_early"]: 745 | logger.info("🟡 RESET PHASE INTERRUPTED BY EXIT_EARLY - proceeding to next episode") 746 | print(f"🟡 STATUS CHANGE: Reset phase interrupted by user - proceeding to next episode") 747 | web_events["exit_early"] = False 748 | 749 | # Check if stop recording was requested during reset phase 750 | if web_events["stop_recording"]: 751 | logger.info("🛑 STOP RECORDING requested during reset phase - ending session") 752 | print(f"🛑 STATUS CHANGE: Stop recording requested during reset - ending session") 753 | break 754 | 755 | # Recording completed 756 | current_phase = "completed" 757 | phase_start_time = None 758 | print(f"🏁 STATUS CHANGE: Recording session completed - all episodes finished") 759 | log_say("Stop recording", cfg.play_sounds, blocking=True) 760 | 761 | finally: 762 | robot.disconnect() 763 | if teleop: 764 | teleop.disconnect() 765 | 766 | if cfg.dataset.push_to_hub: 767 | dataset.push_to_hub(tags=cfg.dataset.tags, private=cfg.dataset.private) 768 | 769 | log_say("Exiting", cfg.play_sounds) 770 | return dataset 771 | --------------------------------------------------------------------------------