diff --git a/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot.py b/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot.py index 99265a4..b4a4e0b 100644 --- a/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot.py +++ b/src/lerobot_robot_ufactory/robots/uf_robot/uf_robot.py @@ -449,9 +449,9 @@ class UFRobot(Robot, Thread): warn_code = getattr(arm, "warn_code", "unknown") return f"mode={mode}, state={state}, error_code={error_code}, warn_code={warn_code}" - def _check_motion_code(self, command: str, code: int) -> None: - """Fail loudly when the SDK rejects a joint command.""" - if code != 0: + def _check_motion_code(self, command: str, code: int | None) -> None: + """Fail loudly when the SDK explicitly rejects a motion command.""" + if code is not None and code != 0: raise RuntimeError(f"{command} failed, code={code}, {self._motion_status()}") def send_action(self, action: dict) -> np.ndarray: diff --git a/src/lerobot_robot_ufactory/scripts/uf_lerobot_record.py b/src/lerobot_robot_ufactory/scripts/uf_lerobot_record.py index c092a45..3947fab 100644 --- a/src/lerobot_robot_ufactory/scripts/uf_lerobot_record.py +++ b/src/lerobot_robot_ufactory/scripts/uf_lerobot_record.py @@ -510,26 +510,40 @@ def _ask_choice(prompt: str, options: dict[str, str]) -> str: print(f"Invalid choice, please enter {keys}.") +def _missing_dataset_files(root: Path) -> list[str]: + """Return the local files required before a dataset can be resumed.""" + missing = [] + for relative_path in ("meta/info.json", "meta/tasks.parquet"): + if not (root / relative_path).is_file(): + missing.append(relative_path) + if not any((root / "meta" / "episodes").glob("*/*.parquet")): + missing.append("meta/episodes/*/*.parquet") + if not any((root / "data").glob("*/*.parquet")): + missing.append("data/*/*.parquet") + return missing + + def _prepare_dataset_root(cfg: UFRecordConfig) -> None: - """Create the dataset root and ask how to handle an existing dataset.""" + """Prepare an existing dataset root without pre-creating a new one.""" 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: + if not existed: + if cfg.resume: + raise RuntimeError(f"Cannot resume because the dataset directory does not exist: {root}") 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." - ) + missing = _missing_dataset_files(root) + if missing: + missing_text = ", ".join(missing) + message = ( + f"Dataset directory is incomplete and cannot be resumed: {root}\n" + f"Missing: {missing_text}" + ) + if cfg.resume or not sys.stdin.isatty(): + raise RuntimeError(message) choice = _ask_choice( - f"Directory exists but is not a valid LeRobot dataset: {root}", + message, options={ "o": "Overwrite: remove this directory and record a new dataset", "c": "Cancel", @@ -541,6 +555,9 @@ def _prepare_dataset_root(cfg: UFRecordConfig) -> None: raise SystemExit("Recording cancelled.") return + if cfg.resume: + return + # A valid LeRobot dataset already exists. if not sys.stdin.isatty(): # Non-interactive run: keep the previous auto-resume behaviour. diff --git a/start_manual_record.sh b/start_manual_record.sh index 39db6c2..d0543a9 100755 --- a/start_manual_record.sh +++ b/start_manual_record.sh @@ -2,49 +2,7 @@ set -euo pipefail -repo_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" -cd "$repo_root" - -config_path="config/manual_mode/xarm7_manual_record_config.yaml" -dataset_root="$( - sed -n '/^dataset:/,/^[^[:space:]]/s/^[[:space:]]*root:[[:space:]]*//p' "$config_path" \ - | head -n 1 \ - | tr -d "\"'" -)" - -if [[ -z "$dataset_root" ]]; then - printf 'Could not read dataset.root from %s\n' "$config_path" >&2 - exit 1 -fi - -if [[ "$dataset_root" != /* ]]; then - dataset_root="$repo_root/$dataset_root" -fi - -record_args=("$@") -resume_requested=false -for arg in "${record_args[@]}"; do - if [[ "$arg" == "-r" ]]; then - resume_requested=true - break - fi -done - -if [[ -e "$dataset_root" ]]; then - if [[ ! -f "$dataset_root/meta/info.json" ]]; then - printf 'Dataset directory exists but is not a valid LeRobot dataset: %s\n' "$dataset_root" >&2 - printf 'Choose a new dataset.root, or remove this empty/incomplete directory before recording.\n' >&2 - exit 1 - fi - - if [[ "$resume_requested" == false ]]; then - record_args=("-r" "${record_args[@]}") - fi -elif [[ "$resume_requested" == true ]]; then - printf 'Cannot resume because the dataset directory does not exist: %s\n' "$dataset_root" >&2 - exit 1 -fi exec uv run uf-lerobot-record \ - --config_path "$config_path" \ - "${record_args[@]}" + --config_path config/manual_mode/xarm7_manual_record_config.yaml \ + "$@" diff --git a/tests/test_manual_mode.py b/tests/test_manual_mode.py index 9b0a94d..b716312 100644 --- a/tests/test_manual_mode.py +++ b/tests/test_manual_mode.py @@ -315,6 +315,27 @@ def test_prepare_dataset_root_rejects_incomplete_resume(tmp_path): _prepare_dataset_root(cfg) +def test_prepare_dataset_root_rejects_resume_when_root_is_missing(tmp_path): + cfg = SimpleNamespace(dataset=SimpleNamespace(root=tmp_path / "missing"), resume=True) + + with pytest.raises(RuntimeError, match="does not exist"): + _prepare_dataset_root(cfg) + + +def test_prepare_dataset_root_resumes_complete_dataset_without_prompt(tmp_path, monkeypatch): + root = tmp_path / "dataset" + (root / "meta" / "episodes" / "chunk-000").mkdir(parents=True) + (root / "data" / "chunk-000").mkdir(parents=True) + (root / "meta" / "info.json").write_text("{}") + (root / "meta" / "tasks.parquet").write_bytes(b"tasks") + (root / "meta" / "episodes" / "chunk-000" / "file-000.parquet").write_bytes(b"episodes") + (root / "data" / "chunk-000" / "file-000.parquet").write_bytes(b"data") + cfg = SimpleNamespace(dataset=SimpleNamespace(root=root), resume=True) + monkeypatch.setattr(record_module.sys.stdin, "isatty", lambda: True) + + _prepare_dataset_root(cfg) + + def test_manual_record_loop_writes_actual_state_as_action(tmp_path): class FakeRobot: name = "fake_manual_robot"