diff --git a/README.md b/README.md index 27baa16..f66e5d9 100644 --- a/README.md +++ b/README.md @@ -142,6 +142,8 @@ Controls: `Space` start the episode, `→` save it, `←` discard and re-record > During collection the **relative position between the robot arm and the camera must not change**, and the camera setup at inference time must match the one used during collection. If the arm or camera moves, previously collected data becomes invalid. +> If the dataset root already exists and `-r` is not given, the script asks whether to overwrite it, resume, or cancel. + ### 3. Manual drag recording ```bash @@ -149,7 +151,16 @@ Controls: `Space` start the episode, `→` save it, `←` discard and re-record ./start_manual_record.sh -r # force resume; fails if the dataset directory does not exist ``` -The launcher reads `dataset.root` from `config/manual_mode/xarm7_manual_record_config.yaml`: on first run it creates the dataset, and on later runs it automatically resumes an existing valid dataset. If the directory exists but is not a valid LeRobot dataset, choose a new `dataset.root`, or remove that directory after confirming it contains no data. +The record script (`uf-lerobot-record`) reads `dataset.root` from the config and creates the path with `mkdir -p` first, then: + +- Directory does not exist → records a new dataset. +- Directory already exists (valid LeRobot dataset) and no `-r` was given → asks interactively: + - `o` overwrite: delete the existing dataset and record a new one + - `r` resume: keep existing episodes and continue recording + - `c` cancel +- Directory exists but is not a valid LeRobot dataset (missing `meta/info.json`) → asks to overwrite it or cancel; non-interactive runs error out instead. + +Pass `-r` to resume directly without asking. Note that `./start_manual_record.sh` keeps its original launcher behavior — it automatically resumes an existing valid dataset (equivalent to `-r`), so run `uv run uf-lerobot-record --config_path config/manual_mode/xarm7_manual_record_config.yaml` directly if you want to see the overwrite/resume prompt. During recording the arm is in teach mode: the actual joint state is written as both the observation and the action. Hold `C` to slowly close the gripper and `O` to slowly open it. Controls: `Space` start, `→` save, `←` discard & re-record, `Esc` stop. Reset the arm manually between episodes. diff --git a/README_ZH.md b/README_ZH.md index d0d85cb..471c8de 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -141,6 +141,8 @@ uv run uf-lerobot-record --config_path config/gello/xarm7_gello_record_config.ya > 采集过程中**机械臂与相机(D435 / D435i)的相对位置必须保持不变**,推理时的相机位置必须与采集时一致。若机械臂或相机发生变化,此前采集的数据将失效。 +> 如果数据集目录已存在且未加 `-r`,脚本会询问是覆盖、续录还是取消。 + ### 3. 手动拖拽数据采集 ```bash @@ -148,7 +150,16 @@ uv run uf-lerobot-record --config_path config/gello/xarm7_gello_record_config.ya ./start_manual_record.sh -r # 强制续录;数据集目录不存在时会报错 ``` -启动脚本会从 `config/manual_mode/xarm7_manual_record_config.yaml` 读取 `dataset.root`:首次运行创建数据集,之后运行会自动续录已有的有效数据集。如果目录存在但不是有效的 LeRobot 数据集,请更换 `dataset.root`,或确认没有数据后删除该目录。 +录制脚本(`uf-lerobot-record`)会从配置读取 `dataset.root` 并先执行 `mkdir -p` 创建路径,然后: + +- 目录不存在:直接录制新数据集。 +- 目录已存在(有效 LeRobot 数据集)且未加 `-r`:交互询问: + - `o` 覆盖:删除已有数据集,重新录制 + - `r` 续录:保留已有 episode,继续录制 + - `c` 取消 +- 目录存在但不是有效 LeRobot 数据集(缺少 `meta/info.json`):询问覆盖或取消;非交互运行时直接报错。 + +加 `-r` 可跳过询问直接续录。注意 `./start_manual_record.sh` 保持原有启动脚本行为——检测到有效数据集会自动续录(相当于 `-r`),如果想看到覆盖/续录的询问,请直接用 `uv run uf-lerobot-record --config_path config/manual_mode/xarm7_manual_record_config.yaml` 运行。 录制时机械臂处于示教模式,实际关节状态会同时作为 observation 和 action 写入数据集。按住 `C` 缓慢闭合夹爪,按住 `O` 缓慢张开。按键控制:`Space` 开始,`→` 保存,`←` 放弃并重录,`Esc` 停止。episode 之间手动复位机械臂。 diff --git a/src/lerobot_robot_ufactory/scripts/uf_lerobot_record.py b/src/lerobot_robot_ufactory/scripts/uf_lerobot_record.py index f46d58b..c092a45 100644 --- a/src/lerobot_robot_ufactory/scripts/uf_lerobot_record.py +++ b/src/lerobot_robot_ufactory/scripts/uf_lerobot_record.py @@ -493,12 +493,85 @@ def _print_record_controls(is_recorded, manual_mode): print(f'⌨ {controls}') +def _ask_choice(prompt: str, options: dict[str, str]) -> str: + """Prompt the user to pick one of the given options (lowercase keys).""" + keys = "/".join(options) + while True: + print(f"\n{prompt}") + for key, description in options.items(): + print(f" [{key}] {description}") + try: + choice = input(f"Choose [{keys}]: ").strip().lower() + except EOFError: + print("No input available, cancelling.") + raise SystemExit(1) + if choice in options: + return choice + print(f"Invalid choice, please enter {keys}.") + + +def _prepare_dataset_root(cfg: UFRecordConfig) -> None: + """Create the dataset root and ask how to handle an existing dataset.""" + root = Path(cfg.dataset.root) + existed = root.exists() + + # Create the dataset root (and any parent directories) first. + root.mkdir(parents=True, exist_ok=True) + + if not existed or cfg.resume: + return + + if not (root / "meta" / "info.json").is_file(): + # The directory exists but is not a valid LeRobot dataset. + if not sys.stdin.isatty(): + raise RuntimeError( + f"Dataset directory exists but is not a valid LeRobot dataset: {root}\n" + "Choose a new dataset.root, or remove this empty/incomplete directory before recording." + ) + choice = _ask_choice( + f"Directory exists but is not a valid LeRobot dataset: {root}", + options={ + "o": "Overwrite: remove this directory and record a new dataset", + "c": "Cancel", + }, + ) + if choice == "o": + shutil.rmtree(root) + else: + raise SystemExit("Recording cancelled.") + return + + # A valid LeRobot dataset already exists. + if not sys.stdin.isatty(): + # Non-interactive run: keep the previous auto-resume behaviour. + cfg.resume = True + print(f"Existing dataset found, resuming recording (non-interactive): {root}") + return + + choice = _ask_choice( + f"Dataset directory already exists: {root}", + options={ + "o": "Overwrite: delete the existing dataset and record a new one", + "r": "Resume: keep existing episodes and continue recording", + "c": "Cancel", + }, + ) + if choice == "o": + shutil.rmtree(root) + elif choice == "r": + cfg.resume = True + else: + raise SystemExit("Recording cancelled.") + + def record(cfg: UFRecordConfig, async_save: bool = False) -> LeRobotDataset: init_logging() logging.info(pformat(asdict(cfg))) if cfg.display_data: init_rerun(session_name="recording") + _prepare_dataset_root(cfg) + robot = make_robot_from_config(cfg.robot) teleop = make_teleoperator_from_config(cfg.teleop) if cfg.teleop is not None else None manual_mode = bool(getattr(cfg.robot, "manual_mode", False))