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