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

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