add convert aloha 2 lerobot

.gitignore

@@ -1,4 +1,5 @@
 cobot_magic/
 librealsense/
 data*/
-outputs/
+outputs/
+lerobot_datasets/

README.MD

@@ -0,0 +1 @@
+export LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtiff.so.5

collect_data/aloha.yaml

@@ -0,0 +1,23 @@
+camera_names:
+- cam_high
+- cam_left_wrist
+- cam_right_wrist
+dataset_dir: /home/ubuntu/LYT/lerobot_aloha/datasets/3camera
+episode_idx: 0
+frame_rate: 30
+img_front_depth_topic: /camera_f/depth/image_raw
+img_front_topic: /camera_f/color/image_raw
+img_left_depth_topic: /camera_l/depth/image_raw
+img_left_topic: /camera_l/color/image_raw
+img_right_depth_topic: /camera_r/depth/image_raw
+img_right_topic: /camera_r/color/image_raw
+master_arm_left_topic: /master/joint_left
+master_arm_right_topic: /master/joint_right
+max_timesteps: 500
+num_episodes: 50
+puppet_arm_left_topic: /puppet/joint_left
+puppet_arm_right_topic: /puppet/joint_right
+robot_base_topic: /odom
+task_name: aloha_mobile_dummy
+use_depth_image: false
+use_robot_base: false

collect_data/collect_data.py

@@ -1,77 +1,18 @@
 import os
 import time
-import argparse
 from aloha_mobile import AlohaRobotRos
-from utils import save_data, init_keyboard_listener
+from utils import save_data, init_keyboard_listener, load_config, log_say
 
 
-def get_arguments():
+def main(config_path):
-    parser = argparse.ArgumentParser()
+    args = load_config(config_path)
-    parser.add_argument('--dataset_dir', action='store', type=str, help='Dataset_dir.',
-                        default="./data", required=False)
-    parser.add_argument('--task_name', action='store', type=str, help='Task name.',
-                        default="aloha_mobile_dummy", required=False)
-    parser.add_argument('--episode_idx', action='store', type=int, help='Episode index.',
-                        default=0, required=False)
-    parser.add_argument('--max_timesteps', action='store', type=int, help='Max_timesteps.',
-                        default=500, required=False)
-    parser.add_argument('--camera_names', action='store', type=str, help='camera_names',
-                        default=['cam_high', 'cam_left_wrist', 'cam_right_wrist'], required=False)
-    parser.add_argument('--num_episodes', action='store', type=int, help='Num_episodes.',
-                        default=1, required=False)
     
-
-    #  topic name of color image
-    parser.add_argument('--img_front_topic', action='store', type=str, help='img_front_topic',
-                        default='/camera_f/color/image_raw', required=False)
-    parser.add_argument('--img_left_topic', action='store', type=str, help='img_left_topic',
-                        default='/camera_l/color/image_raw', required=False)
-    parser.add_argument('--img_right_topic', action='store', type=str, help='img_right_topic',
-                        default='/camera_r/color/image_raw', required=False)
-    
-    # topic name of depth image
-    parser.add_argument('--img_front_depth_topic', action='store', type=str, help='img_front_depth_topic',
-                        default='/camera_f/depth/image_raw', required=False)
-    parser.add_argument('--img_left_depth_topic', action='store', type=str, help='img_left_depth_topic',
-                        default='/camera_l/depth/image_raw', required=False)
-    parser.add_argument('--img_right_depth_topic', action='store', type=str, help='img_right_depth_topic',
-                        default='/camera_r/depth/image_raw', required=False)
-    
-    # topic name of arm
-    parser.add_argument('--master_arm_left_topic', action='store', type=str, help='master_arm_left_topic',
-                        default='/master/joint_left', required=False)
-    parser.add_argument('--master_arm_right_topic', action='store', type=str, help='master_arm_right_topic',
-                        default='/master/joint_right', required=False)
-    parser.add_argument('--puppet_arm_left_topic', action='store', type=str, help='puppet_arm_left_topic',
-                        default='/puppet/joint_left', required=False)
-    parser.add_argument('--puppet_arm_right_topic', action='store', type=str, help='puppet_arm_right_topic',
-                        default='/puppet/joint_right', required=False)
-    
-    # topic name of robot_base
-    parser.add_argument('--robot_base_topic', action='store', type=str, help='robot_base_topic',
-                        default='/odom', required=False)
-    
-    parser.add_argument('--use_robot_base', action='store', type=bool, help='use_robot_base',
-                        default=False, required=False)
-    
-    # collect depth image
-    parser.add_argument('--use_depth_image', action='store', type=bool, help='use_depth_image',
-                        default=False, required=False)
-    
-    parser.add_argument('--frame_rate', action='store', type=int, help='frame_rate',
-                        default=30, required=False)
-    
-    args = parser.parse_args()
-    return args
-
-
-def main():
-    args = get_arguments()
     ros_operator = AlohaRobotRos(args)
     dataset_dir = os.path.join(args.dataset_dir, args.task_name)
-    # 确保数据集目录存在
+    # Ensure dataset directory exists
     os.makedirs(dataset_dir, exist_ok=True)
-    # 单集收集模式
+    
+    # Single episode collection mode
     if args.num_episodes == 1:
         print(f"Recording single episode {args.episode_idx}...")
         timesteps, actions = ros_operator.process()
@@ -84,7 +25,8 @@ def main():
         save_data(args, timesteps, actions, dataset_path)
         print(f"\033[32mEpisode {args.episode_idx} saved successfully at {dataset_path}\033[0m")
         return 0
-    # 多集收集模式
+    
+    # Multi-episode collection mode
     print("""
 \033[1;36mKeyboard Controls:\033[0m
 ← \033[1mLeft Arrow\033[0m:  Start Recording
@@ -93,7 +35,8 @@ def main():
 ↑ \033[1mUp Arrow\033[0m:    Replay Data (if implemented)
 \033[1mESC\033[0m:           Exit Program
 """)
-    # 初始化键盘监听器
+    log_say("欢迎您为 具身智能科学家项目采集数据，您辛苦了。我已经将一切准备就绪，请您按方向左键开始录制数据。", play_sounds=True)
+
     listener, events = init_keyboard_listener()
     episode_idx = args.episode_idx
     collected_episodes = 0
@@ -102,27 +45,31 @@ def main():
         while collected_episodes < args.num_episodes:
             if events["exit_early"]:
                 print("\033[33mOperation terminated by user\033[0m")
+                log_say("操作被你停止了，如果这是个误操作，请重新开始。", play_sounds=True)
                 break
                 
             if events["record_start"]:
-                # 重置事件状态，开始新的录制
+                # Reset event states for new recording
                 events["record_start"] = False
                 events["save_data"] = False
                 events["discard_data"] = False
-                
+                log_say(f"开始录制第{episode_idx}条轨迹，请开始操作机械臂。", play_sounds=True)
                 print(f"\n\033[1;32mRecording episode {episode_idx}...\033[0m")
                 timesteps, actions = ros_operator.process()
                 print(f"\033[1;33mRecorded {len(actions)} timesteps. (→ to save, ↓ to discard)\033[0m")
+                log_say(f"第{episode_idx}条轨迹的录制已经到达最大时间步并录制结束，请选择是否保留该条轨迹。按方向右键保留，方向下键丢弃。", play_sounds=True)
                 
-                # 等待用户决定保存或丢弃
+                # Wait for user decision to save or discard
                 while True:
                     if events["save_data"]:
                         events["save_data"] = False
                         
                         if len(actions) < args.max_timesteps:
                             print(f"\033[31mSave failure: Recorded only {len(actions)}/{args.max_timesteps} timesteps.\033[0m")
+                            log_say(f"由于当前轨迹的实际时间步数小于最大时间步数，因此无法保存该条轨迹。该条轨迹将被丢弃，请重新录制。", play_sounds=True)
                         else:
                             dataset_path = os.path.join(dataset_dir, f"episode_{episode_idx}")
+                            log_say(f"你选择了保留该轨迹作为第{episode_idx}条轨迹数据。接下来请你按方向左键开始录制下一条轨迹。", play_sounds=True)
                             save_data(args, timesteps, actions, dataset_path)
                             print(f"\033[32mEpisode {episode_idx} saved successfully at {dataset_path}\033[0m")
                             episode_idx += 1
@@ -132,22 +79,25 @@ def main():
                         
                     if events["discard_data"]:
                         events["discard_data"] = False
+                        log_say(f"你选择了丢弃该轨迹作为第{episode_idx}条轨迹数据。接下来请你按方向左键开始录制下一条轨迹。", play_sounds=True)
                         print("\033[33mData discarded. Press ← to start a new recording.\033[0m")
                         break
                         
                     if events["exit_early"]:
                         print("\033[33mOperation terminated by user\033[0m")
+                        log_say("操作被你停止了，如果这是个误操作，请重新开始。", play_sounds=True)
                         return 0
                         
-                    time.sleep(0.1)  # 减少CPU使用率
+                    time.sleep(0.1)  # Reduce CPU usage
             
-            time.sleep(0.1)  # 减少CPU使用率
+            time.sleep(0.1)  # Reduce CPU usage
             
         if collected_episodes == args.num_episodes:
+            log_say("恭喜你，本次数据已经全部录制完成。您辛苦了~", play_sounds=True)
             print(f"\n\033[1;32mData collection complete! All {args.num_episodes} episodes collected.\033[0m")
             
     finally:
-        # 确保监听器被清理
+        # Ensure listener is cleaned up
         if listener:
             listener.stop()
             print("Keyboard listener stopped")
@@ -155,12 +105,10 @@ def main():
             
 if __name__ == '__main__':
     try:
-        exit_code = main()
+        exit_code = main("/home/ubuntu/LYT/lerobot_aloha/collect_data/aloha.yaml")
         exit(exit_code if exit_code is not None else 0)
     except KeyboardInterrupt:
         print("\n\033[33mProgram interrupted by user\033[0m")
         exit(0)
     except Exception as e:
         print(f"\n\033[31mError: {e}\033[0m")
-
-# python collect_data.py --dataset_dir ~/data --max_timesteps 500 --episode_idx 0

collect_data/collect_data_gui.py

@@ -1,416 +1,403 @@
-import os
 import sys
-import time
+import os
-import argparse
+import yaml
-from aloha_mobile import AlohaRobotRos
+from collect_data import main
-from utils import save_data, init_keyboard_listener
+from types import SimpleNamespace
 from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
-                           QLabel, QLineEdit, QPushButton, QCheckBox, QSpinBox,
+                             QLabel, QLineEdit, QPushButton, QSpinBox, QCheckBox, 
-                           QGroupBox, QFormLayout, QTabWidget, QTextEdit, QFileDialog,
+                             QGroupBox, QTabWidget, QMessageBox, QFileDialog)
-                           QMessageBox, QProgressBar, QComboBox)
+from PyQt5.QtCore import Qt
-from PyQt5.QtCore import Qt, QThread, pyqtSignal, pyqtSlot
+
-from PyQt5.QtGui import QFont, QIcon, QTextCursor
 
-class DataCollectionThread(QThread):
-    """处理数据收集的线程"""
-    update_signal = pyqtSignal(str)
-    progress_signal = pyqtSignal(int)
-    finish_signal = pyqtSignal(bool, str)
-    
-    def __init__(self, args, parent=None):
-        super(DataCollectionThread, self).__init__(parent)
-        self.args = args
-        self.is_running = True
-        self.ros_operator = None
-        
-    def run(self):
-        try:
-            self.update_signal.emit("正在初始化ROS操作...\n")
-            self.ros_operator = AlohaRobotRos(self.args)
-            dataset_dir = os.path.join(self.args.dataset_dir, self.args.task_name)
-            os.makedirs(dataset_dir, exist_ok=True)
-            
-            # 单集收集模式
-            if self.args.num_episodes == 1:
-                self.update_signal.emit(f"开始录制第 {self.args.episode_idx} 集...\n")
-                timesteps, actions = self.ros_operator.process()
-                
-                if len(actions) < self.args.max_timesteps:
-                    self.update_signal.emit(f"保存失败: 只录制了 {len(actions)}/{self.args.max_timesteps} 个时间步.\n")
-                    self.finish_signal.emit(False, f"只录制了 {len(actions)}/{self.args.max_timesteps} 个时间步")
-                    return
-                    
-                dataset_path = os.path.join(dataset_dir, f"episode_{self.args.episode_idx}")
-                save_data(self.args, timesteps, actions, dataset_path)
-                self.update_signal.emit(f"第 {self.args.episode_idx} 集成功保存到 {dataset_path}.\n")
-                self.finish_signal.emit(True, "数据收集完成")
-                
-            # 多集收集模式
-            else:
-                self.update_signal.emit("""
-键盘控制:
-← 左箭头: 开始录制
-→ 右箭头: 保存当前数据
-↓ 下箭头: 丢弃当前数据
-ESC:      退出程序
-""")
-                # 初始化键盘监听器
-                listener, events = init_keyboard_listener()
-                episode_idx = self.args.episode_idx
-                collected_episodes = 0
-                
-                try:
-                    while collected_episodes < self.args.num_episodes and self.is_running:
-                        if events["exit_early"]:
-                            self.update_signal.emit("操作被用户终止.\n")
-                            break
-                            
-                        if events["record_start"]:
-                            # 重置事件状态，开始新的录制
-                            events["record_start"] = False
-                            events["save_data"] = False
-                            events["discard_data"] = False
-                            
-                            self.update_signal.emit(f"\n正在录制第 {episode_idx} 集...\n")
-                            timesteps, actions = self.ros_operator.process()
-                            self.update_signal.emit(f"已录制 {len(actions)} 个时间步. (→ 保存, ↓ 丢弃)\n")
-                            
-                            # 等待用户决定保存或丢弃
-                            while self.is_running:
-                                if events["save_data"]:
-                                    events["save_data"] = False
-                                    
-                                    if len(actions) < self.args.max_timesteps:
-                                        self.update_signal.emit(f"保存失败: 只录制了 {len(actions)}/{self.args.max_timesteps} 个时间步.\n")
-                                    else:
-                                        dataset_path = os.path.join(dataset_dir, f"episode_{episode_idx}")
-                                        save_data(self.args, timesteps, actions, dataset_path)
-                                        self.update_signal.emit(f"第 {episode_idx} 集成功保存到 {dataset_path}.\n")
-                                        episode_idx += 1
-                                        collected_episodes += 1
-                                        progress_percentage = int(collected_episodes * 100 / self.args.num_episodes)
-                                        self.progress_signal.emit(progress_percentage)
-                                        self.update_signal.emit(f"进度: {collected_episodes}/{self.args.num_episodes} 集已收集. (← 开始新一集)\n")
-                                    break
-                                    
-                                if events["discard_data"]:
-                                    events["discard_data"] = False
-                                    self.update_signal.emit("数据已丢弃. 请按 ← 开始新的录制.\n")
-                                    break
-                                    
-                                if events["exit_early"]:
-                                    self.update_signal.emit("操作被用户终止.\n")
-                                    self.is_running = False
-                                    break
-                                    
-                                time.sleep(0.1)  # 减少CPU使用率
-                        
-                        time.sleep(0.1)  # 减少CPU使用率
-                        
-                    if collected_episodes == self.args.num_episodes:
-                        self.update_signal.emit(f"\n数据收集完成! 所有 {self.args.num_episodes} 集已收集.\n")
-                        self.finish_signal.emit(True, "全部数据集收集完成")
-                    else:
-                        self.finish_signal.emit(False, "数据收集未完成")
-                        
-                finally:
-                    # 确保监听器被清理
-                    if listener:
-                        listener.stop()
-                        self.update_signal.emit("键盘监听器已停止\n")
-                
-        except Exception as e:
-            self.update_signal.emit(f"错误: {str(e)}\n")
-            self.finish_signal.emit(False, str(e))
-            
-    def stop(self):
-        self.is_running = False
-        self.wait()
 
 class AlohaDataCollectionGUI(QMainWindow):
     def __init__(self):
         super().__init__()
-        self.setWindowTitle("ALOHA 数据收集工具")
+        self.setWindowTitle("MindRobot-V1 Data Collection")
-        self.setGeometry(100, 100, 800, 700)
+        self.setGeometry(100, 100, 800, 600)
         
-        # 主组件
         self.central_widget = QWidget()
         self.setCentralWidget(self.central_widget)
+        
         self.main_layout = QVBoxLayout(self.central_widget)
         
-        # 创建选项卡
+        self.config_path = os.path.expanduser("/home/ubuntu/LYT/lerobot_aloha/collect_data/aloha.yaml")
-        self.tab_widget = QTabWidget()
+        self.create_ui()
-        self.main_layout.addWidget(self.tab_widget)
+        self.setup_connections()
+        self.load_default_config()
         
-        # 创建配置选项卡
+    def create_ui(self):
-        self.config_tab = QWidget()
+        # Create tabs
-        self.tab_widget.addTab(self.config_tab, "配置")
+        self.tabs = QTabWidget()
+        self.main_layout.addWidget(self.tabs)
         
-        # 创建数据收集选项卡
+        # General Settings Tab
-        self.collection_tab = QWidget()
+        self.general_tab = QWidget()
-        self.tab_widget.addTab(self.collection_tab, "数据收集")
+        self.tabs.addTab(self.general_tab, "General Settings")
+        self.create_general_tab()
         
-        self.setup_config_tab()
+        # Camera Settings Tab
-        self.setup_collection_tab()
+        self.camera_tab = QWidget()
+        self.tabs.addTab(self.camera_tab, "Camera Settings")
+        self.create_camera_tab()
         
-        # 初始化数据收集线程
+        # Arm Settings Tab
-        self.collection_thread = None
+        self.arm_tab = QWidget()
+        self.tabs.addTab(self.arm_tab, "Arm Settings")
+        self.create_arm_tab()
         
-    def setup_config_tab(self):
+        # Control Buttons
-        config_layout = QVBoxLayout(self.config_tab)
+        self.control_group = QGroupBox("Control")
+        control_layout = QHBoxLayout()
         
-        # 基本配置组
+        self.load_config_button = QPushButton("Load Config")
-        basic_group = QGroupBox("基本配置")
+        self.save_config_button = QPushButton("Save Config")
-        basic_form = QFormLayout()
+        self.start_button = QPushButton("Start Recording")
+        self.stop_button = QPushButton("Stop Recording")
+        self.exit_button = QPushButton("Exit")
         
-        self.dataset_dir = QLineEdit("./data")
+        control_layout.addWidget(self.load_config_button)
-        self.browse_button = QPushButton("浏览")
+        control_layout.addWidget(self.save_config_button)
-        self.browse_button.clicked.connect(self.browse_dataset_dir)
+        control_layout.addWidget(self.start_button)
+        control_layout.addWidget(self.stop_button)
+        control_layout.addWidget(self.exit_button)
         
-        dir_layout = QHBoxLayout()
+        self.control_group.setLayout(control_layout)
-        dir_layout.addWidget(self.dataset_dir)
+        self.main_layout.addWidget(self.control_group)
-        dir_layout.addWidget(self.browse_button)
         
-        self.task_name = QLineEdit("aloha_mobile_dummy")
+    def create_general_tab(self):
-        self.episode_idx = QSpinBox()
+        layout = QVBoxLayout(self.general_tab)
-        self.episode_idx.setRange(0, 1000)
-        self.max_timesteps = QSpinBox()
-        self.max_timesteps.setRange(1, 10000)
-        self.max_timesteps.setValue(500)
-        self.num_episodes = QSpinBox()
-        self.num_episodes.setRange(1, 100)
-        self.num_episodes.setValue(1)
-        self.frame_rate = QSpinBox()
-        self.frame_rate.setRange(1, 120)
-        self.frame_rate.setValue(30)
         
-        basic_form.addRow("数据集目录:", dir_layout)
+        # Config File Path
-        basic_form.addRow("任务名称:", self.task_name)
+        config_group = QGroupBox("Configuration File")
-        basic_form.addRow("起始集索引:", self.episode_idx)
+        config_layout = QHBoxLayout()
-        basic_form.addRow("最大时间步:", self.max_timesteps)
-        basic_form.addRow("集数:", self.num_episodes)
-        basic_form.addRow("帧率:", self.frame_rate)
         
-        basic_group.setLayout(basic_form)
+        self.config_path_edit = QLineEdit(self.config_path)
-        config_layout.addWidget(basic_group)
+        self.browse_config_button = QPushButton("Browse...")
         
-        # 相机话题组
+        config_layout.addWidget(QLabel("Config File:"))
-        camera_group = QGroupBox("相机话题")
+        config_layout.addWidget(self.config_path_edit)
-        camera_form = QFormLayout()
+        config_layout.addWidget(self.browse_config_button)
         
-        self.img_front_topic = QLineEdit('/camera_f/color/image_raw')
-        self.img_left_topic = QLineEdit('/camera_l/color/image_raw')
-        self.img_right_topic = QLineEdit('/camera_r/color/image_raw')
-        
-        self.img_front_depth_topic = QLineEdit('/camera_f/depth/image_raw')
-        self.img_left_depth_topic = QLineEdit('/camera_l/depth/image_raw')
-        self.img_right_depth_topic = QLineEdit('/camera_r/depth/image_raw')
-        
-        self.use_depth_image = QCheckBox("使用深度图像")
-        
-        camera_form.addRow("前置相机:", self.img_front_topic)
-        camera_form.addRow("左腕相机:", self.img_left_topic)
-        camera_form.addRow("右腕相机:", self.img_right_topic)
-        camera_form.addRow("前置深度:", self.img_front_depth_topic)
-        camera_form.addRow("左腕深度:", self.img_left_depth_topic)
-        camera_form.addRow("右腕深度:", self.img_right_depth_topic)
-        camera_form.addRow("", self.use_depth_image)
-        
-        camera_group.setLayout(camera_form)
-        config_layout.addWidget(camera_group)
-        
-        # 机器人话题组
-        robot_group = QGroupBox("机器人话题")
-        robot_form = QFormLayout()
-        
-        self.master_arm_left_topic = QLineEdit('/master/joint_left')
-        self.master_arm_right_topic = QLineEdit('/master/joint_right')
-        self.puppet_arm_left_topic = QLineEdit('/puppet/joint_left')
-        self.puppet_arm_right_topic = QLineEdit('/puppet/joint_right')
-        self.robot_base_topic = QLineEdit('/odom')
-        self.use_robot_base = QCheckBox("使用机器人底盘")
-        
-        robot_form.addRow("主左臂:", self.master_arm_left_topic)
-        robot_form.addRow("主右臂:", self.master_arm_right_topic)
-        robot_form.addRow("从左臂:", self.puppet_arm_left_topic)
-        robot_form.addRow("从右臂:", self.puppet_arm_right_topic)
-        robot_form.addRow("底盘:", self.robot_base_topic)
-        robot_form.addRow("", self.use_robot_base)
-        
-        robot_group.setLayout(robot_form)
-        config_layout.addWidget(robot_group)
-        
-        # 相机名称配置
-        camera_names_group = QGroupBox("相机名称")
-        camera_names_layout = QVBoxLayout()
-        
-        self.camera_names = ['cam_high', 'cam_left_wrist', 'cam_right_wrist']
-        self.camera_checkboxes = {}
-        
-        for cam_name in self.camera_names:
-            self.camera_checkboxes[cam_name] = QCheckBox(cam_name)
-            self.camera_checkboxes[cam_name].setChecked(True)
-            camera_names_layout.addWidget(self.camera_checkboxes[cam_name])
-        
-        camera_names_group.setLayout(camera_names_layout)
-        config_layout.addWidget(camera_names_group)
-        
-        # 保存配置按钮
-        self.save_config_button = QPushButton("保存配置")
-        self.save_config_button.clicked.connect(self.save_config)
-        config_layout.addWidget(self.save_config_button)
-        
-    def setup_collection_tab(self):
-        collection_layout = QVBoxLayout(self.collection_tab)
-        
-        # 当前配置展示
-        config_group = QGroupBox("当前配置")
-        self.config_text = QTextEdit()
-        self.config_text.setReadOnly(True)
-        config_layout = QVBoxLayout()
-        config_layout.addWidget(self.config_text)
         config_group.setLayout(config_layout)
-        collection_layout.addWidget(config_group)
+        layout.addWidget(config_group)
         
-        # 操作按钮
+        # Dataset Directory
-        buttons_layout = QHBoxLayout()
+        dir_group = QGroupBox("Dataset Directory")
+        dir_layout = QHBoxLayout()
         
-        self.start_button = QPushButton("开始收集")
+        self.dataset_dir_edit = QLineEdit()
-        self.start_button.setIcon(QIcon.fromTheme("media-playback-start"))
+        self.browse_dir_button = QPushButton("Browse...")
-        self.start_button.clicked.connect(self.start_collection)
         
-        self.stop_button = QPushButton("停止")
+        dir_layout.addWidget(QLabel("Dataset Directory:"))
-        self.stop_button.setIcon(QIcon.fromTheme("media-playback-stop"))
+        dir_layout.addWidget(self.dataset_dir_edit)
-        self.stop_button.setEnabled(False)
+        dir_layout.addWidget(self.browse_dir_button)
-        self.stop_button.clicked.connect(self.stop_collection)
         
-        buttons_layout.addWidget(self.start_button)
+        dir_group.setLayout(dir_layout)
-        buttons_layout.addWidget(self.stop_button)
+        layout.addWidget(dir_group)
-        collection_layout.addLayout(buttons_layout)
         
-        # 进度条
+        # Task Settings
-        self.progress_bar = QProgressBar()
+        task_group = QGroupBox("Task Settings")
-        self.progress_bar.setValue(0)
+        task_layout = QVBoxLayout()
-        collection_layout.addWidget(self.progress_bar)
         
-        # 日志输出
+        self.task_name_edit = QLineEdit()
-        log_group = QGroupBox("操作日志")
+        self.episode_idx_spin = QSpinBox()
-        self.log_text = QTextEdit()
+        self.episode_idx_spin.setRange(0, 9999)
-        self.log_text.setReadOnly(True)
+        self.max_timesteps_spin = QSpinBox()
-        log_layout = QVBoxLayout()
+        self.max_timesteps_spin.setRange(1, 10000)
-        log_layout.addWidget(self.log_text)
+        self.num_episodes_spin = QSpinBox()
-        log_group.setLayout(log_layout)
+        self.num_episodes_spin.setRange(1, 1000)
-        collection_layout.addWidget(log_group)
+        self.frame_rate_spin = QSpinBox()
+        self.frame_rate_spin.setRange(1, 60)
         
+        task_layout.addWidget(QLabel("Task Name:"))
+        task_layout.addWidget(self.task_name_edit)
+        task_layout.addWidget(QLabel("Episode Index:"))
+        task_layout.addWidget(self.episode_idx_spin)
+        task_layout.addWidget(QLabel("Max Timesteps:"))
+        task_layout.addWidget(self.max_timesteps_spin)
+        task_layout.addWidget(QLabel("Number of Episodes:"))
+        task_layout.addWidget(self.num_episodes_spin)
+        task_layout.addWidget(QLabel("Frame Rate:"))
+        task_layout.addWidget(self.frame_rate_spin)
+        
+        task_group.setLayout(task_layout)
+        layout.addWidget(task_group)
+        
+        # Options
+        options_group = QGroupBox("Options")
+        options_layout = QVBoxLayout()
+        
+        self.use_robot_base_check = QCheckBox("Use Robot Base")
+        self.use_depth_image_check = QCheckBox("Use Depth Image")
+        
+        options_layout.addWidget(self.use_robot_base_check)
+        options_layout.addWidget(self.use_depth_image_check)
+        
+        options_group.setLayout(options_layout)
+        layout.addWidget(options_group)
+        
+        layout.addStretch()
+        
+    def create_camera_tab(self):
+        layout = QVBoxLayout(self.camera_tab)
+        
+        # Color Image Topics
+        color_group = QGroupBox("Color Image Topics")
+        color_layout = QVBoxLayout()
+        
+        self.img_front_topic_edit = QLineEdit()
+        self.img_left_topic_edit = QLineEdit()
+        self.img_right_topic_edit = QLineEdit()
+        
+        color_layout.addWidget(QLabel("Front Camera Topic:"))
+        color_layout.addWidget(self.img_front_topic_edit)
+        color_layout.addWidget(QLabel("Left Camera Topic:"))
+        color_layout.addWidget(self.img_left_topic_edit)
+        color_layout.addWidget(QLabel("Right Camera Topic:"))
+        color_layout.addWidget(self.img_right_topic_edit)
+        
+        color_group.setLayout(color_layout)
+        layout.addWidget(color_group)
+        
+        # Depth Image Topics
+        depth_group = QGroupBox("Depth Image Topics")
+        depth_layout = QVBoxLayout()
+        
+        self.img_front_depth_topic_edit = QLineEdit()
+        self.img_left_depth_topic_edit = QLineEdit()
+        self.img_right_depth_topic_edit = QLineEdit()
+        
+        depth_layout.addWidget(QLabel("Front Depth Topic:"))
+        depth_layout.addWidget(self.img_front_depth_topic_edit)
+        depth_layout.addWidget(QLabel("Left Depth Topic:"))
+        depth_layout.addWidget(self.img_left_depth_topic_edit)
+        depth_layout.addWidget(QLabel("Right Depth Topic:"))
+        depth_layout.addWidget(self.img_right_depth_topic_edit)
+        
+        depth_group.setLayout(depth_layout)
+        layout.addWidget(depth_group)
+        
+        layout.addStretch()
+        
+    def create_arm_tab(self):
+        layout = QVBoxLayout(self.arm_tab)
+        
+        # Master Arm Topics
+        master_group = QGroupBox("Master Arm Topics")
+        master_layout = QVBoxLayout()
+        
+        self.master_arm_left_topic_edit = QLineEdit()
+        self.master_arm_right_topic_edit = QLineEdit()
+        
+        master_layout.addWidget(QLabel("Master Left Arm Topic:"))
+        master_layout.addWidget(self.master_arm_left_topic_edit)
+        master_layout.addWidget(QLabel("Master Right Arm Topic:"))
+        master_layout.addWidget(self.master_arm_right_topic_edit)
+        
+        master_group.setLayout(master_layout)
+        layout.addWidget(master_group)
+        
+        # Puppet Arm Topics
+        puppet_group = QGroupBox("Puppet Arm Topics")
+        puppet_layout = QVBoxLayout()
+        
+        self.puppet_arm_left_topic_edit = QLineEdit()
+        self.puppet_arm_right_topic_edit = QLineEdit()
+        
+        puppet_layout.addWidget(QLabel("Puppet Left Arm Topic:"))
+        puppet_layout.addWidget(self.puppet_arm_left_topic_edit)
+        puppet_layout.addWidget(QLabel("Puppet Right Arm Topic:"))
+        puppet_layout.addWidget(self.puppet_arm_right_topic_edit)
+        
+        puppet_group.setLayout(puppet_layout)
+        layout.addWidget(puppet_group)
+        
+        # Robot Base Topic
+        base_group = QGroupBox("Robot Base Topic")
+        base_layout = QVBoxLayout()
+        
+        self.robot_base_topic_edit = QLineEdit()
+        
+        base_layout.addWidget(QLabel("Robot Base Topic:"))
+        base_layout.addWidget(self.robot_base_topic_edit)
+        
+        base_group.setLayout(base_layout)
+        layout.addWidget(base_group)
+        
+        layout.addStretch()
+        
+    def setup_connections(self):
+        self.load_config_button.clicked.connect(self.load_config)
+        self.save_config_button.clicked.connect(self.save_config)
+        self.browse_config_button.clicked.connect(self.browse_config_file)
+        self.browse_dir_button.clicked.connect(self.browse_dataset_dir)
+        self.start_button.clicked.connect(self.start_recording)
+        self.stop_button.clicked.connect(self.stop_recording)
+        self.exit_button.clicked.connect(self.close)
+        
+    def load_default_config(self):
+        default_config = {
+            'dataset_dir': '/home/ubuntu/LYT/lerobot_aloha/datasets/3camera',
+            'task_name': 'aloha_mobile_dummy',
+            'episode_idx': 0,
+            'max_timesteps': 500,
+            'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist'],
+            'num_episodes': 50,
+            'img_front_topic': '/camera_f/color/image_raw',
+            'img_left_topic': '/camera_l/color/image_raw',
+            'img_right_topic': '/camera_r/color/image_raw',
+            'img_front_depth_topic': '/camera_f/depth/image_raw',
+            'img_left_depth_topic': '/camera_l/depth/image_raw',
+            'img_right_depth_topic': '/camera_r/depth/image_raw',
+            'master_arm_left_topic': '/master/joint_left',
+            'master_arm_right_topic': '/master/joint_right',
+            'puppet_arm_left_topic': '/puppet/joint_left',
+            'puppet_arm_right_topic': '/puppet/joint_right',
+            'robot_base_topic': '/odom',
+            'use_robot_base': False,
+            'use_depth_image': False,
+            'frame_rate': 30
+        }
+        
+        # Update UI with default values
+        self.update_ui_from_config(default_config)
+        
+    def update_ui_from_config(self, config):
+        """Update UI elements from a config dictionary"""
+        self.dataset_dir_edit.setText(config.get('dataset_dir', ''))
+        self.task_name_edit.setText(config.get('task_name', ''))
+        self.episode_idx_spin.setValue(config.get('episode_idx', 0))
+        self.max_timesteps_spin.setValue(config.get('max_timesteps', 500))
+        self.num_episodes_spin.setValue(config.get('num_episodes', 1))
+        self.frame_rate_spin.setValue(config.get('frame_rate', 30))
+        
+        self.img_front_topic_edit.setText(config.get('img_front_topic', ''))
+        self.img_left_topic_edit.setText(config.get('img_left_topic', ''))
+        self.img_right_topic_edit.setText(config.get('img_right_topic', ''))
+        
+        self.img_front_depth_topic_edit.setText(config.get('img_front_depth_topic', ''))
+        self.img_left_depth_topic_edit.setText(config.get('img_left_depth_topic', ''))
+        self.img_right_depth_topic_edit.setText(config.get('img_right_depth_topic', ''))
+        
+        self.master_arm_left_topic_edit.setText(config.get('master_arm_left_topic', ''))
+        self.master_arm_right_topic_edit.setText(config.get('master_arm_right_topic', ''))
+        self.puppet_arm_left_topic_edit.setText(config.get('puppet_arm_left_topic', ''))
+        self.puppet_arm_right_topic_edit.setText(config.get('puppet_arm_right_topic', ''))
+        
+        self.robot_base_topic_edit.setText(config.get('robot_base_topic', ''))
+        self.use_robot_base_check.setChecked(config.get('use_robot_base', False))
+        self.use_depth_image_check.setChecked(config.get('use_depth_image', False))
+        
+    def get_config_from_ui(self):
+        """Get current UI values as a config dictionary"""
+        config = {
+            'dataset_dir': self.dataset_dir_edit.text(),
+            'task_name': self.task_name_edit.text(),
+            'episode_idx': self.episode_idx_spin.value(),
+            'max_timesteps': self.max_timesteps_spin.value(),
+            'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist'],
+            'num_episodes': self.num_episodes_spin.value(),
+            'img_front_topic': self.img_front_topic_edit.text(),
+            'img_left_topic': self.img_left_topic_edit.text(),
+            'img_right_topic': self.img_right_topic_edit.text(),
+            'img_front_depth_topic': self.img_front_depth_topic_edit.text(),
+            'img_left_depth_topic': self.img_left_depth_topic_edit.text(),
+            'img_right_depth_topic': self.img_right_depth_topic_edit.text(),
+            'master_arm_left_topic': self.master_arm_left_topic_edit.text(),
+            'master_arm_right_topic': self.master_arm_right_topic_edit.text(),
+            'puppet_arm_left_topic': self.puppet_arm_left_topic_edit.text(),
+            'puppet_arm_right_topic': self.puppet_arm_right_topic_edit.text(),
+            'robot_base_topic': self.robot_base_topic_edit.text(),
+            'use_robot_base': self.use_robot_base_check.isChecked(),
+            'use_depth_image': self.use_depth_image_check.isChecked(),
+            'frame_rate': self.frame_rate_spin.value()
+        }
+        return config
+        
+    def browse_config_file(self):
+        file_path, _ = QFileDialog.getOpenFileName(
+            self, "Select Config File", "", "YAML Files (*.yaml *.yml)"
+        )
+        if file_path:
+            self.config_path_edit.setText(file_path)
+            self.load_config()
+            
     def browse_dataset_dir(self):
-        directory = QFileDialog.getExistingDirectory(self, "选择数据集目录", self.dataset_dir.text())
+        dir_path = QFileDialog.getExistingDirectory(
-        if directory:
+            self, "Select Dataset Directory"
-            self.dataset_dir.setText(directory)
+        )
-    
+        if dir_path:
-    def save_config(self):
+            self.dataset_dir_edit.setText(dir_path)
-        # 更新配置显示
+            
-        selected_cameras = [cam for cam, checkbox in self.camera_checkboxes.items() if checkbox.isChecked()]
+    def load_config(self):
-        
+        config_path = self.config_path_edit.text()
-        config_text = f"""
+        if not os.path.exists(config_path):
-任务名称: {self.task_name.text()}
+            QMessageBox.warning(self, "Warning", f"Config file not found: {config_path}")
-数据集目录: {self.dataset_dir.text()}
-起始集索引: {self.episode_idx.value()}
-最大时间步: {self.max_timesteps.value()}
-集数: {self.num_episodes.value()}
-帧率: {self.frame_rate.value()}
-使用深度图像: {"是" if self.use_depth_image.isChecked() else "否"}
-使用机器人底盘: {"是" if self.use_robot_base.isChecked() else "否"}
-相机: {', '.join(selected_cameras)}
-        """
-        
-        self.config_text.setText(config_text)
-        self.tab_widget.setCurrentIndex(1)  # 切换到收集选项卡
-        
-        QMessageBox.information(self, "配置已保存", "配置已更新，可以开始数据收集")
-    
-    def start_collection(self):
-        if not self.task_name.text():
-            QMessageBox.warning(self, "配置错误", "请输入有效的任务名称")
             return
             
-        # 构建参数
+        try:
-        args = argparse.Namespace(
+            with open(config_path, 'r') as f:
-            dataset_dir=self.dataset_dir.text(),
+                config = yaml.safe_load(f)
-            task_name=self.task_name.text(),
+                self.update_ui_from_config(config)
-            episode_idx=self.episode_idx.value(),
+                self.statusBar().showMessage(f"Config loaded from {config_path}", 3000)
-            max_timesteps=self.max_timesteps.value(),
+        except Exception as e:
-            num_episodes=self.num_episodes.value(),
+            QMessageBox.critical(self, "Error", f"Failed to load config: {str(e)}")
-            camera_names=[cam for cam, checkbox in self.camera_checkboxes.items() if checkbox.isChecked()],
-            img_front_topic=self.img_front_topic.text(),
-            img_left_topic=self.img_left_topic.text(),
-            img_right_topic=self.img_right_topic.text(),
-            img_front_depth_topic=self.img_front_depth_topic.text(),
-            img_left_depth_topic=self.img_left_depth_topic.text(),
-            img_right_depth_topic=self.img_right_depth_topic.text(),
-            master_arm_left_topic=self.master_arm_left_topic.text(),
-            master_arm_right_topic=self.master_arm_right_topic.text(),
-            puppet_arm_left_topic=self.puppet_arm_left_topic.text(),
-            puppet_arm_right_topic=self.puppet_arm_right_topic.text(),
-            robot_base_topic=self.robot_base_topic.text(),
-            use_robot_base=self.use_robot_base.isChecked(),
-            use_depth_image=self.use_depth_image.isChecked(),
-            frame_rate=self.frame_rate.value()
-        )
-        
-        # 更新UI状态
-        self.start_button.setEnabled(False)
-        self.stop_button.setEnabled(True)
-        self.progress_bar.setValue(0)
-        self.log_text.clear()
-        self.log_text.append("正在初始化数据收集...\n")
-        
-        # 创建并启动线程
-        self.collection_thread = DataCollectionThread(args)
-        self.collection_thread.update_signal.connect(self.update_log)
-        self.collection_thread.progress_signal.connect(self.update_progress)
-        self.collection_thread.finish_signal.connect(self.collection_finished)
-        self.collection_thread.start()
-        
-    def stop_collection(self):
-        if self.collection_thread and self.collection_thread.isRunning():
-            self.log_text.append("正在停止数据收集...\n")
-            self.collection_thread.stop()
-            self.collection_thread = None
             
-        self.start_button.setEnabled(True)
+    def save_config(self):
-        self.stop_button.setEnabled(False)
+        config_path = self.config_path_edit.text()
-    
+        if not config_path:
-    @pyqtSlot(str)
+            QMessageBox.warning(self, "Warning", "Please specify a config file path")
-    def update_log(self, message):
+            return
-        self.log_text.append(message)
-        # 自动滚动到底部
-        cursor = self.log_text.textCursor()
-        cursor.movePosition(QTextCursor.End)
-        self.log_text.setTextCursor(cursor)
-        
-    @pyqtSlot(int)
-    def update_progress(self, value):
-        self.progress_bar.setValue(value)
-        
-    @pyqtSlot(bool, str)
-    def collection_finished(self, success, message):
-        self.start_button.setEnabled(True)
-        self.stop_button.setEnabled(False)
-        
-        if success:
-            QMessageBox.information(self, "完成", message)
-        else:
-            QMessageBox.warning(self, "出错", f"数据收集失败: {message}")
             
-        # 更新episode_idx值
+        try:
-        if success and self.num_episodes.value() > 0:
+            config = self.get_config_from_ui()
-            self.episode_idx.setValue(self.episode_idx.value() + self.num_episodes.value())
+            with open(config_path, 'w') as f:
+                yaml.dump(config, f, default_flow_style=False)
+            self.statusBar().showMessage(f"Config saved to {config_path}", 3000)
+        except Exception as e:
+            QMessageBox.critical(self, "Error", f"Failed to save config: {str(e)}")
+            
+    def start_recording(self):
+        try:
+            # Save current config to a temporary file
+            temp_config_path = "/tmp/aloha_temp_config.yaml"
+            config = self.get_config_from_ui()
+            
+            # Validate inputs
+            if not config['dataset_dir']:
+                QMessageBox.warning(self, "Warning", "Dataset directory cannot be empty!")
+                return
+                
+            if not config['task_name']:
+                QMessageBox.warning(self, "Warning", "Task name cannot be empty!")
+                return
+                
+            with open(temp_config_path, 'w') as f:
+                yaml.dump(config, f, default_flow_style=False)
+                
+            self.statusBar().showMessage("Recording started...")
+            
+            # Start recording with the temporary config file
+            exit_code = main(temp_config_path)
+            
+            if exit_code == 0:
+                self.statusBar().showMessage("Recording completed successfully!", 5000)
+            else:
+                self.statusBar().showMessage("Recording completed with errors", 5000)
+                
+        except Exception as e:
+            QMessageBox.critical(self, "Error", f"An error occurred: {str(e)}")
+            self.statusBar().showMessage("Recording failed", 5000)
+            
+    def stop_recording(self):
+        # In a real application, this would signal the recording thread to stop
+        self.statusBar().showMessage("Recording stopped", 5000)
+        QMessageBox.information(self, "Info", "Stop recording requested. This would stop the recording in a real implementation.")
 
-def main():
+
+if __name__ == "__main__":
     app = QApplication(sys.argv)
     window = AlohaDataCollectionGUI()
     window.show()
     sys.exit(app.exec_())
-
-if __name__ == "__main__":
-    main()

collect_data/convert_aloha_data_to_lerobot.py

@@ -0,0 +1,292 @@
+"""
+Script to convert Aloha hdf5 data to the LeRobot dataset v2.0 format.
+
+Example usage: uv run examples/aloha_real/convert_aloha_data_to_lerobot.py --raw-dir /path/to/raw/data --repo-id <org>/<dataset-name>
+"""
+
+import dataclasses
+from pathlib import Path
+import shutil
+from typing import Literal
+import os
+import h5py
+from lerobot.common.datasets.lerobot_dataset import LEROBOT_HOME
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.common.datasets.push_dataset_to_hub._download_raw import download_raw
+import numpy as np
+import torch
+import tqdm
+import tyro
+
+# 使用自定义路径覆盖
+LEROBOT_HOME = Path("/home/ubuntu/hdd0/lerobot_datasets/3camera")
+
+@dataclasses.dataclass(frozen=True)
+class DatasetConfig:
+    use_videos: bool = True
+    tolerance_s: float = 0.0001
+    image_writer_processes: int = 10
+    image_writer_threads: int = 5
+    video_backend: str | None = None
+
+
+DEFAULT_DATASET_CONFIG = DatasetConfig()
+
+
+def create_empty_dataset(
+    repo_id: str,
+    robot_type: str,
+    mode: Literal["video", "image"] = "video",
+    *,
+    has_velocity: bool = False,
+    has_effort: bool = False,
+    dataset_config: DatasetConfig = DEFAULT_DATASET_CONFIG,
+) -> LeRobotDataset:
+    motors = [
+        "right_waist",
+        "right_shoulder",
+        "right_elbow",
+        "right_forearm_roll",
+        "right_wrist_angle",
+        "right_wrist_rotate",
+        "right_gripper",
+        "left_waist",
+        "left_shoulder",
+        "left_elbow",
+        "left_forearm_roll",
+        "left_wrist_angle",
+        "left_wrist_rotate",
+        "left_gripper",
+    ]
+    # 确定camera的情况
+    # cameras = [
+    #     "cam_high",
+    #     "cam_low",
+    #     "cam_left_wrist",
+    #     "cam_right_wrist",
+    # ]
+    cameras = [
+        "cam_high",
+        "cam_left_wrist",
+        "cam_right_wrist",
+    ]
+
+    features = {
+        "observation.state": {
+            "dtype": "float32",
+            "shape": (len(motors),),
+            "names": [
+                motors,
+            ],
+        },
+        "action": {
+            "dtype": "float32",
+            "shape": (len(motors),),
+            "names": [
+                motors,
+            ],
+        },
+    }
+
+    if has_velocity:
+        features["observation.velocity"] = {
+            "dtype": "float32",
+            "shape": (len(motors),),
+            "names": [
+                motors,
+            ],
+        }
+
+    if has_effort:
+        features["observation.effort"] = {
+            "dtype": "float32",
+            "shape": (len(motors),),
+            "names": [
+                motors,
+            ],
+        }
+
+    for cam in cameras:
+        features[f"observation.images.{cam}"] = {
+            "dtype": mode,
+            "shape": (3, 480, 640),
+            "names": [
+                "channels",
+                "height",
+                "width",
+            ],
+        }
+
+    if Path(LEROBOT_HOME / repo_id).exists():
+        shutil.rmtree(LEROBOT_HOME / repo_id)
+
+    return LeRobotDataset.create(
+        repo_id=repo_id,
+        fps=30,
+        root=Path(LEROBOT_HOME / repo_id),
+        robot_type=robot_type,
+        features=features,
+        use_videos=dataset_config.use_videos,
+        tolerance_s=dataset_config.tolerance_s,
+        image_writer_processes=dataset_config.image_writer_processes,
+        image_writer_threads=dataset_config.image_writer_threads,
+        video_backend=dataset_config.video_backend,
+    )
+
+
+def get_cameras(hdf5_files: list[Path]) -> list[str]:
+    with h5py.File(hdf5_files[0], "r") as ep:
+        # ignore depth channel, not currently handled
+        return [key for key in ep["/observations/images"].keys() if "depth" not in key]  # noqa: SIM118
+
+
+def has_velocity(hdf5_files: list[Path]) -> bool:
+    with h5py.File(hdf5_files[0], "r") as ep:
+        return "/observations/qvel" in ep
+
+
+def has_effort(hdf5_files: list[Path]) -> bool:
+    with h5py.File(hdf5_files[0], "r") as ep:
+        return "/observations/effort" in ep
+
+
+def load_raw_images_per_camera(ep: h5py.File, cameras: list[str]) -> dict[str, np.ndarray]:
+    imgs_per_cam = {}
+    for camera in cameras:
+        uncompressed = ep[f"/observations/images/{camera}"].ndim == 4
+
+        if uncompressed:
+            # load all images in RAM
+            imgs_array = ep[f"/observations/images/{camera}"][:]
+        else:
+            import cv2
+
+            # load one compressed image after the other in RAM and uncompress
+            imgs_array = []
+            for data in ep[f"/observations/images/{camera}"]:
+                imgs_array.append(cv2.cvtColor(cv2.imdecode(data, 1), cv2.COLOR_BGR2RGB))
+            imgs_array = np.array(imgs_array)
+
+        imgs_per_cam[camera] = imgs_array
+    return imgs_per_cam
+
+
+def load_raw_episode_data(
+    ep_path: Path,
+) -> tuple[dict[str, np.ndarray], torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+    with h5py.File(ep_path, "r") as ep:
+        state = torch.from_numpy(ep["/observations/qpos"][:])
+        action = torch.from_numpy(ep["/action"][:])
+
+        velocity = None
+        if "/observations/qvel" in ep:
+            velocity = torch.from_numpy(ep["/observations/qvel"][:])
+
+        effort = None
+        if "/observations/effort" in ep:
+            effort = torch.from_numpy(ep["/observations/effort"][:])
+
+        # 确定camera的情况
+        # imgs_per_cam = load_raw_images_per_camera(
+        #     ep,
+        #     [
+        #         "cam_high",
+        #         "cam_low",
+        #         "cam_left_wrist",
+        #         "cam_right_wrist",
+        #     ],
+        # )
+        imgs_per_cam = load_raw_images_per_camera(
+            ep,
+            [
+                "cam_high",
+                "cam_left_wrist",
+                "cam_right_wrist",
+            ],
+        )
+
+    return imgs_per_cam, state, action, velocity, effort
+
+
+def populate_dataset(
+    dataset: LeRobotDataset,
+    hdf5_files: list[Path],
+    task: str,
+    episodes: list[int] | None = None,
+) -> LeRobotDataset:
+    if episodes is None:
+        episodes = range(len(hdf5_files))
+
+    for ep_idx in tqdm.tqdm(episodes):
+        ep_path = hdf5_files[ep_idx]
+        # import pdb
+        # pdb.set_trace()
+        imgs_per_cam, state, action, velocity, effort = load_raw_episode_data(ep_path)
+        num_frames = state.shape[0]
+
+        for i in range(num_frames):
+            frame = {
+                "observation.state": state[i],
+                "action": action[i],
+            }
+
+            for camera, img_array in imgs_per_cam.items():
+                frame[f"observation.images.{camera}"] = img_array[i]
+
+            if velocity is not None:
+                frame["observation.velocity"] = velocity[i]
+            if effort is not None:
+                frame["observation.effort"] = effort[i]
+
+            dataset.add_frame(frame)
+
+        dataset.save_episode(task=task)
+
+    return dataset
+
+
+def port_aloha(
+    raw_dir: Path,
+    repo_id: str,
+    raw_repo_id: str | None = None,
+    task: str = "DEBUG",
+    *,
+    episodes: list[int] | None = None,
+    push_to_hub: bool = False,
+    is_mobile: bool = False,
+    mode: Literal["video", "image"] = "image",
+    dataset_config: DatasetConfig = DEFAULT_DATASET_CONFIG,
+):
+    print(LEROBOT_HOME)
+    if (LEROBOT_HOME / repo_id).exists():
+        shutil.rmtree(LEROBOT_HOME / repo_id)
+
+    if not raw_dir.exists():
+        if raw_repo_id is None:
+            raise ValueError("raw_repo_id must be provided if raw_dir does not exist")
+        download_raw(raw_dir, repo_id=raw_repo_id)
+
+    hdf5_files = sorted(raw_dir.glob("episode_*.hdf5"))
+
+    dataset = create_empty_dataset(
+        repo_id,
+        robot_type="mobile_aloha" if is_mobile else "aloha",
+        mode=mode,
+        has_effort=has_effort(hdf5_files),
+        has_velocity=has_velocity(hdf5_files),
+        dataset_config=dataset_config,
+    )
+    dataset = populate_dataset(
+        dataset,
+        hdf5_files,
+        task=task,
+        episodes=episodes,
+    )
+    dataset.consolidate()
+
+    if push_to_hub:
+        dataset.push_to_hub()
+
+
+if __name__ == "__main__":
+    tyro.cli(port_aloha)

collect_data/utils.py

@@ -270,3 +270,53 @@ def init_keyboard_listener():
     listener.start()
 
     return listener, events
+
+
+import yaml
+from argparse import Namespace
+def load_config(yaml_path):
+    """Load configuration from YAML file and return as Namespace object"""
+    with open(yaml_path, 'r') as f:
+        config_dict = yaml.safe_load(f)
+    
+    # Convert dict to Namespace (similar to argparse.Namespace)
+    return Namespace(**config_dict)
+
+
+import platform
+import subprocess
+# import pyttsx3
+def say(text, blocking=False):
+    system = platform.system()
+
+    if system == "Darwin":
+        cmd = ["say", text]
+
+    elif system == "Linux":
+        # cmd = ["spd-say", text]
+        # if blocking:
+        #     cmd.append("--wait")
+        cmd = ["edge-playback", "--text", text]
+
+    elif system == "Windows":
+        cmd = [
+            "PowerShell",
+            "-Command",
+            "Add-Type -AssemblyName System.Speech; "
+            f"(New-Object System.Speech.Synthesis.SpeechSynthesizer).Speak('{text}')",
+        ]
+
+    else:
+        raise RuntimeError("Unsupported operating system for text-to-speech.")
+
+    if blocking:
+        subprocess.run(cmd, check=True)
+    else:
+        subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
+
+
+def log_say(text, play_sounds, blocking=False):
+    print(text)
+
+    if play_sounds:
+        say(text, blocking)

init_robot.bash

@@ -1,2 +0,0 @@
-source ~/ros_noetic/devel_isolated/setup.bash
-cd cobot_magic/remote_control-x86-can-v2 && ./tools/can.sh && ./tools/jgl_2follower.sh