lerobot_aloha/collect_data/collect_data.py

import os
import time
import argparse
from aloha_mobile import AlohaRobotRos
from utils import save_data, init_keyboard_listener


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)
    

    #  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)
    # 确保数据集目录存在
    os.makedirs(dataset_dir, exist_ok=True)
    # 单集收集模式
    if args.num_episodes == 1:
        print(f"Recording single episode {args.episode_idx}...")
        timesteps, actions = ros_operator.process()
        
        if len(actions) < args.max_timesteps:
            print(f"\033[31m\nSave failure: Recorded only {len(actions)}/{args.max_timesteps} timesteps.\033[0m\n")
            return -1
            
        dataset_path = os.path.join(dataset_dir, f"episode_{args.episode_idx}")
        save_data(args, timesteps, actions, dataset_path)
        print(f"\033[32mEpisode {args.episode_idx} saved successfully at {dataset_path}\033[0m")
        return 0
    # 多集收集模式
    print("""
\033[1;36mKeyboard Controls:\033[0m
← \033[1mLeft Arrow\033[0m:  Start Recording
→ \033[1mRight Arrow\033[0m: Save Current Data
↓ \033[1mDown Arrow\033[0m:  Discard Current Data
↑ \033[1mUp Arrow\033[0m:    Replay Data (if implemented)
\033[1mESC\033[0m:           Exit Program
""")
    # 初始化键盘监听器
    listener, events = init_keyboard_listener()
    episode_idx = args.episode_idx
    collected_episodes = 0
    
    try:
        while collected_episodes < args.num_episodes:
            if events["exit_early"]:
                print("\033[33mOperation terminated by user\033[0m")
                break
                
            if events["record_start"]:
                # 重置事件状态，开始新的录制
                events["record_start"] = False
                events["save_data"] = False
                events["discard_data"] = False
                
                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")
                
                # 等待用户决定保存或丢弃
                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")
                        else:
                            dataset_path = os.path.join(dataset_dir, f"episode_{episode_idx}")
                            save_data(args, timesteps, actions, dataset_path)
                            print(f"\033[32mEpisode {episode_idx} saved successfully at {dataset_path}\033[0m")
                            episode_idx += 1
                            collected_episodes += 1
                            print(f"\033[1mProgress: {collected_episodes}/{args.num_episodes} episodes collected. (← to start new episode)\033[0m")
                        break
                        
                    if events["discard_data"]:
                        events["discard_data"] = False
                        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")
                        return 0
                        
                    time.sleep(0.1)  # 减少CPU使用率
            
            time.sleep(0.1)  # 减少CPU使用率
            
        if collected_episodes == args.num_episodes:
            print(f"\n\033[1;32mData collection complete! All {args.num_episodes} episodes collected.\033[0m")
            
    finally:
        # 确保监听器被清理
        if listener:
            listener.stop()
            print("Keyboard listener stopped")

            
if __name__ == '__main__':
    try:
        exit_code = main()
        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