115 lines
5.8 KiB
Python
Executable File
115 lines
5.8 KiB
Python
Executable File
import os
|
|
import time
|
|
from aloha_mobile import AlohaRobotRos
|
|
from utils import save_data, init_keyboard_listener, load_config, log_say
|
|
|
|
|
|
def main(config_path):
|
|
args = load_config(config_path)
|
|
|
|
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()
|
|
|
|
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
|
|
|
|
# Multi-episode collection mode
|
|
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
|
|
""")
|
|
log_say("欢迎您为 具身智能科学家项目采集数据，您辛苦了。我已经将一切准备就绪，请您按方向左键开始录制数据。", play_sounds=True)
|
|
|
|
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")
|
|
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
|
|
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
|
|
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) # Reduce CPU usage
|
|
|
|
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")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
try:
|
|
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")
|