Fix bug in dataset root resolution
This commit is contained in:
parent
ed0063730d
commit
1b775a4be2
@ -449,9 +449,9 @@ class UFRobot(Robot, Thread):
|
|||||||
warn_code = getattr(arm, "warn_code", "unknown")
|
warn_code = getattr(arm, "warn_code", "unknown")
|
||||||
return f"mode={mode}, state={state}, error_code={error_code}, warn_code={warn_code}"
|
return f"mode={mode}, state={state}, error_code={error_code}, warn_code={warn_code}"
|
||||||
|
|
||||||
def _check_motion_code(self, command: str, code: int) -> None:
|
def _check_motion_code(self, command: str, code: int | None) -> None:
|
||||||
"""Fail loudly when the SDK rejects a joint command."""
|
"""Fail loudly when the SDK explicitly rejects a motion command."""
|
||||||
if code != 0:
|
if code is not None and code != 0:
|
||||||
raise RuntimeError(f"{command} failed, code={code}, {self._motion_status()}")
|
raise RuntimeError(f"{command} failed, code={code}, {self._motion_status()}")
|
||||||
|
|
||||||
def send_action(self, action: dict) -> np.ndarray:
|
def send_action(self, action: dict) -> np.ndarray:
|
||||||
|
|||||||
@ -510,26 +510,40 @@ def _ask_choice(prompt: str, options: dict[str, str]) -> str:
|
|||||||
print(f"Invalid choice, please enter {keys}.")
|
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:
|
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)
|
root = Path(cfg.dataset.root)
|
||||||
existed = root.exists()
|
existed = root.exists()
|
||||||
|
|
||||||
# Create the dataset root (and any parent directories) first.
|
if not existed:
|
||||||
root.mkdir(parents=True, exist_ok=True)
|
if cfg.resume:
|
||||||
|
raise RuntimeError(f"Cannot resume because the dataset directory does not exist: {root}")
|
||||||
if not existed or cfg.resume:
|
|
||||||
return
|
return
|
||||||
|
|
||||||
if not (root / "meta" / "info.json").is_file():
|
missing = _missing_dataset_files(root)
|
||||||
# The directory exists but is not a valid LeRobot dataset.
|
if missing:
|
||||||
if not sys.stdin.isatty():
|
missing_text = ", ".join(missing)
|
||||||
raise RuntimeError(
|
message = (
|
||||||
f"Dataset directory exists but is not a valid LeRobot dataset: {root}\n"
|
f"Dataset directory is incomplete and cannot be resumed: {root}\n"
|
||||||
"Choose a new dataset.root, or remove this empty/incomplete directory before recording."
|
f"Missing: {missing_text}"
|
||||||
)
|
)
|
||||||
|
if cfg.resume or not sys.stdin.isatty():
|
||||||
|
raise RuntimeError(message)
|
||||||
choice = _ask_choice(
|
choice = _ask_choice(
|
||||||
f"Directory exists but is not a valid LeRobot dataset: {root}",
|
message,
|
||||||
options={
|
options={
|
||||||
"o": "Overwrite: remove this directory and record a new dataset",
|
"o": "Overwrite: remove this directory and record a new dataset",
|
||||||
"c": "Cancel",
|
"c": "Cancel",
|
||||||
@ -541,6 +555,9 @@ def _prepare_dataset_root(cfg: UFRecordConfig) -> None:
|
|||||||
raise SystemExit("Recording cancelled.")
|
raise SystemExit("Recording cancelled.")
|
||||||
return
|
return
|
||||||
|
|
||||||
|
if cfg.resume:
|
||||||
|
return
|
||||||
|
|
||||||
# A valid LeRobot dataset already exists.
|
# A valid LeRobot dataset already exists.
|
||||||
if not sys.stdin.isatty():
|
if not sys.stdin.isatty():
|
||||||
# Non-interactive run: keep the previous auto-resume behaviour.
|
# Non-interactive run: keep the previous auto-resume behaviour.
|
||||||
|
|||||||
@ -2,49 +2,7 @@
|
|||||||
|
|
||||||
set -euo pipefail
|
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 \
|
exec uv run uf-lerobot-record \
|
||||||
--config_path "$config_path" \
|
--config_path config/manual_mode/xarm7_manual_record_config.yaml \
|
||||||
"${record_args[@]}"
|
"$@"
|
||||||
|
|||||||
@ -315,6 +315,27 @@ def test_prepare_dataset_root_rejects_incomplete_resume(tmp_path):
|
|||||||
_prepare_dataset_root(cfg)
|
_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):
|
def test_manual_record_loop_writes_actual_state_as_action(tmp_path):
|
||||||
class FakeRobot:
|
class FakeRobot:
|
||||||
name = "fake_manual_robot"
|
name = "fake_manual_robot"
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user