tangger/lerobot_aloha
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add convert aloha 2 lerobot

Browse Source
This commit is contained in:
tangger
2025-04-20 21:50:20 +08:00
parent 722de584d2
commit 25fb9c0d33
13 changed files with 753 additions and 452 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
.gitignore vendored
View File

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

1
README.MD
View File

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

BIN
collect_data/__pycache__/aloha_mobile.cpython-310.pyc Normal file
View File

Binary file not shown.

BIN
collect_data/__pycache__/collect_data.cpython-310.pyc Normal file
View File

Binary file not shown.

BIN
collect_data/__pycache__/utils.cpython-310.pyc Normal file
View File

Binary file not shown.

23
collect_data/aloha.yaml Normal file
View File

@@ -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

100
collect_data/collect_data.py
View File

@@ -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():
parser = argparse.ArgumentParser()
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)
def main(config_path):
args = load_config(config_path)
# 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

733
collect_data/collect_data_gui.py
View File

@@ -1,416 +1,403 @@
import os
import sys
import time
import argparse
from aloha_mobile import AlohaRobotRos
from utils import save_data, init_keyboard_listener
import os
import yaml
from collect_data import main
from types import SimpleNamespace
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
QLabel, QLineEdit, QPushButton, QCheckBox, QSpinBox,
QGroupBox, QFormLayout, QTabWidget, QTextEdit, QFileDialog,
QMessageBox, QProgressBar, QComboBox)
from PyQt5.QtCore import Qt, QThread, pyqtSignal, pyqtSlot
from PyQt5.QtGui import QFont, QIcon, QTextCursor
QLabel, QLineEdit, QPushButton, QSpinBox, QCheckBox,
QGroupBox, QTabWidget, QMessageBox, QFileDialog)
from PyQt5.QtCore import Qt
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.setGeometry(100, 100, 800, 700)
self.setWindowTitle("MindRobot-V1 Data Collection")
self.setGeometry(100, 100, 800, 600)
# 主组件
self.central_widget = QWidget()
self.setCentralWidget(self.central_widget)
self.main_layout = QVBoxLayout(self.central_widget)
# 创建选项卡
self.tab_widget = QTabWidget()
self.main_layout.addWidget(self.tab_widget)
self.config_path = os.path.expanduser("/home/ubuntu/LYT/lerobot_aloha/collect_data/aloha.yaml")
self.create_ui()
self.setup_connections()
self.load_default_config()
# 创建配置选项卡
self.config_tab = QWidget()
self.tab_widget.addTab(self.config_tab, "配置")
def create_ui(self):
# Create tabs
self.tabs = QTabWidget()
self.main_layout.addWidget(self.tabs)
# 创建数据收集选项卡
self.collection_tab = QWidget()
self.tab_widget.addTab(self.collection_tab, "数据收集")
# General Settings Tab
self.general_tab = QWidget()
self.tabs.addTab(self.general_tab, "General Settings")
self.create_general_tab()
self.setup_config_tab()
self.setup_collection_tab()
# Camera Settings Tab
self.camera_tab = QWidget()
self.tabs.addTab(self.camera_tab, "Camera Settings")
self.create_camera_tab()
# 初始化数据收集线程
self.collection_thread = None
# Arm Settings Tab
self.arm_tab = QWidget()
self.tabs.addTab(self.arm_tab, "Arm Settings")
self.create_arm_tab()
def setup_config_tab(self):
config_layout = QVBoxLayout(self.config_tab)
# Control Buttons
self.control_group = QGroupBox("Control")
control_layout = QHBoxLayout()
# 基本配置组
basic_group = QGroupBox("基本配置")
basic_form = QFormLayout()
self.load_config_button = QPushButton("Load Config")
self.save_config_button = QPushButton("Save Config")
self.start_button = QPushButton("Start Recording")
self.stop_button = QPushButton("Stop Recording")
self.exit_button = QPushButton("Exit")
self.dataset_dir = QLineEdit("./data")
self.browse_button = QPushButton("浏览")
self.browse_button.clicked.connect(self.browse_dataset_dir)
control_layout.addWidget(self.load_config_button)
control_layout.addWidget(self.save_config_button)
control_layout.addWidget(self.start_button)
control_layout.addWidget(self.stop_button)
control_layout.addWidget(self.exit_button)
dir_layout = QHBoxLayout()
dir_layout.addWidget(self.dataset_dir)
dir_layout.addWidget(self.browse_button)
self.control_group.setLayout(control_layout)
self.main_layout.addWidget(self.control_group)
self.task_name = QLineEdit("aloha_mobile_dummy")
self.episode_idx = QSpinBox()
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)
def create_general_tab(self):
layout = QVBoxLayout(self.general_tab)
basic_form.addRow("数据集目录:", dir_layout)
basic_form.addRow("任务名称:", self.task_name)
basic_form.addRow("起始集索引:", self.episode_idx)
basic_form.addRow("最大时间步:", self.max_timesteps)
basic_form.addRow("集数:", self.num_episodes)
basic_form.addRow("帧率:", self.frame_rate)
# Config File Path
config_group = QGroupBox("Configuration File")
config_layout = QHBoxLayout()
basic_group.setLayout(basic_form)
config_layout.addWidget(basic_group)
self.config_path_edit = QLineEdit(self.config_path)
self.browse_config_button = QPushButton("Browse...")
# 相机话题组
camera_group = QGroupBox("相机话题")
camera_form = QFormLayout()
config_layout.addWidget(QLabel("Config File:"))
config_layout.addWidget(self.config_path_edit)
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)
# 操作按钮
buttons_layout = QHBoxLayout()
# Dataset Directory
dir_group = QGroupBox("Dataset Directory")
dir_layout = QHBoxLayout()
self.start_button = QPushButton("开始收集")
self.start_button.setIcon(QIcon.fromTheme("media-playback-start"))
self.start_button.clicked.connect(self.start_collection)
self.dataset_dir_edit = QLineEdit()
self.browse_dir_button = QPushButton("Browse...")
self.stop_button = QPushButton("停止")
self.stop_button.setIcon(QIcon.fromTheme("media-playback-stop"))
self.stop_button.setEnabled(False)
self.stop_button.clicked.connect(self.stop_collection)
dir_layout.addWidget(QLabel("Dataset Directory:"))
dir_layout.addWidget(self.dataset_dir_edit)
dir_layout.addWidget(self.browse_dir_button)
buttons_layout.addWidget(self.start_button)
buttons_layout.addWidget(self.stop_button)
collection_layout.addLayout(buttons_layout)
dir_group.setLayout(dir_layout)
layout.addWidget(dir_group)
# 进度条
self.progress_bar = QProgressBar()
self.progress_bar.setValue(0)
collection_layout.addWidget(self.progress_bar)
# Task Settings
task_group = QGroupBox("Task Settings")
task_layout = QVBoxLayout()
# 日志输出
log_group = QGroupBox("操作日志")
self.log_text = QTextEdit()
self.log_text.setReadOnly(True)
log_layout = QVBoxLayout()
log_layout.addWidget(self.log_text)
log_group.setLayout(log_layout)
collection_layout.addWidget(log_group)
self.task_name_edit = QLineEdit()
self.episode_idx_spin = QSpinBox()
self.episode_idx_spin.setRange(0, 9999)
self.max_timesteps_spin = QSpinBox()
self.max_timesteps_spin.setRange(1, 10000)
self.num_episodes_spin = QSpinBox()
self.num_episodes_spin.setRange(1, 1000)
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())
if directory:
self.dataset_dir.setText(directory)
def save_config(self):
# 更新配置显示
selected_cameras = [cam for cam, checkbox in self.camera_checkboxes.items() if checkbox.isChecked()]
config_text = f"""
任务名称: {self.task_name.text()}
数据集目录: {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, "配置错误", "请输入有效的任务名称")
dir_path = QFileDialog.getExistingDirectory(
self, "Select Dataset Directory"
)
if dir_path:
self.dataset_dir_edit.setText(dir_path)
def load_config(self):
config_path = self.config_path_edit.text()
if not os.path.exists(config_path):
QMessageBox.warning(self, "Warning", f"Config file not found: {config_path}")
return
# 构建参数
args = argparse.Namespace(
dataset_dir=self.dataset_dir.text(),
task_name=self.task_name.text(),
episode_idx=self.episode_idx.value(),
max_timesteps=self.max_timesteps.value(),
num_episodes=self.num_episodes.value(),
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
try:
with open(config_path, 'r') as f:
config = yaml.safe_load(f)
self.update_ui_from_config(config)
self.statusBar().showMessage(f"Config loaded from {config_path}", 3000)
except Exception as e:
QMessageBox.critical(self, "Error", f"Failed to load config: {str(e)}")
self.start_button.setEnabled(True)
self.stop_button.setEnabled(False)
@pyqtSlot(str)
def update_log(self, message):
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}")
def save_config(self):
config_path = self.config_path_edit.text()
if not config_path:
QMessageBox.warning(self, "Warning", "Please specify a config file path")
return
# 更新episode_idx值
if success and self.num_episodes.value() > 0:
self.episode_idx.setValue(self.episode_idx.value() + self.num_episodes.value())
try:
config = self.get_config_from_ui()
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()

292
collect_data/convert_aloha_data_to_lerobot.py Normal file
View File

@@ -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)

50
collect_data/utils.py
View File

@@ -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)

2
init_robot.bash
View File

@@ -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

1
openpi Submodule

Submodule openpi added at 36dc3c037e

BIN
test.pt
View File

Binary file not shown.