Compare commits

..

1 Commits

20 changed files with 1690 additions and 317 deletions
+50
View File
@@ -239,6 +239,56 @@ Every module is on by default and can be toggled independently (set to
| `--vqa.restrict_to_default_camera` | `false` | Ground VQA only on `--vlm.camera_key` (else every camera). |
| `--executor.episode_parallelism` | `16` | Episodes processed concurrently within each phase. |
## Camera-view curation
`lerobot-curate-cameras` is a separate, lightweight command that uses the same
VLM backend for a **dataset-filtering / curation** pass. It downloads only the
**first episode**, then for each camera view asks the VLM to:
1. **flag** whether the view is blurry / unusable, and
2. **label** the view with a canonical name from a closed vocabulary
(`top`, `wrist`, `front`, `bottom`, `left`, `right`, plus two-word combos
like `left_wrist`).
It runs in one of two modes:
- `--mode=report` (default) — write the labels + verdicts into `meta/`
(`meta/camera_curation.json` and a `curation` block on each camera in
`meta/info.json`). Nothing is moved; this is the cheap triage pass and works
for any dataset.
- `--mode=rename` — apply the labels by renaming each camera key to
`observation.images.<label>`. For **video** datasets this is a
**download-free, server-side Hub commit**: the `videos/<key>/` files are moved
with the Hub's LFS copy/delete (no video is downloaded or re-encoded), and only
the small `meta/` files are edited.
```bash
# Cheap, mutation-free triage (writes meta/camera_curation.json):
uv run lerobot-curate-cameras --repo_id=user/dataset --mode=report
# Apply the labels by renaming camera keys on a new branch (keeps `main` intact):
uv run lerobot-curate-cameras --repo_id=user/dataset --mode=rename --branch=curated
# Run the VLM decision on a GPU via HF Jobs (same --job.* flags as above):
uv run lerobot-curate-cameras --repo_id=user/dataset --mode=rename --job.target=h200
```
Notes:
- The Hub rename is **in place** on the source repo — the Hub does not support
cross-repo LFS copies. Use `--branch` to commit to a branch so `main` is
preserved.
- **Image** datasets store frames inside the data parquet, so their rename can't
avoid touching the data; the rename falls back to a local rewrite (via
[`rename_features`](./using_dataset_tools#rename-features)). Prefer `--mode=report`
for image datasets.
- Views judged unusable are only flagged by default (still renamed). Pass
`--drop_unusable=true` (local path) to remove them.
Key options: `--mode`, `--branch`, `--n_frames`, `--view_vocabulary`,
`--allow_combos`, `--on_collision`, `--drop_unusable`, and the shared
`--vlm.*` / `--job.*` flags documented above.
## Contributing new modules
The pipeline is built to grow, and **contributions are very welcome** —
+22
View File
@@ -89,6 +89,28 @@ lerobot-edit-dataset \
--operation.feature_names "['observation.images.top']"
```
#### Rename Features
Rename feature keys — typically to canonicalize camera views (e.g.
`observation.images.cam_0` → `observation.images.left_wrist`). A rename changes
no pixel data, so it is a cheap key-remap: it rewrites `meta/` (info features,
episode `videos/*` and `stats/*` columns, `stats.json`), moves the
`videos/<key>/` directory, and — for image datasets — renames the embedded
image column. Videos are **not** re-encoded.
```bash
# Rename one or more camera keys
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type rename_features \
--operation.name_mapping '{"observation.images.cam_0": "observation.images.left_wrist"}'
```
If two targets collide (e.g. two cameras both labeled `top`), the operation
raises by default; pass `--operation.on_collision suffix` to disambiguate
deterministically (`top`, `top_2`, …). To label camera views automatically with
a VLM, see the [Annotation Pipeline](./annotation_pipeline#camera-view-curation).
#### Convert to Video
Convert an image-based dataset to video format, creating a new LeRobotDataset where images are stored as videos. This is useful for reducing storage requirements and improving data loading performance. The new dataset will have the exact same structure as the original, but with images encoded as MP4 videos in the proper LeRobot format.
+1
View File
@@ -356,6 +356,7 @@ lerobot-imgtransform-viz="lerobot.scripts.lerobot_imgtransform_viz:main"
lerobot-edit-dataset="lerobot.scripts.lerobot_edit_dataset:main"
lerobot-setup-can="lerobot.scripts.lerobot_setup_can:main"
lerobot-annotate="lerobot.scripts.lerobot_annotate:main"
lerobot-curate-cameras="lerobot.scripts.lerobot_curate_cameras:main"
lerobot-rollout="lerobot.scripts.lerobot_rollout:main"
# ---------------- Tool Configurations ----------------
@@ -0,0 +1,46 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""VLM camera-view curation for LeRobot datasets.
For each dataset, the first episode is inspected by a vision-language model to
(1) judge whether each camera view is blurry/unusable and (2) assign a canonical
view label (``top``/``wrist``/``front``/…). The labels can then be applied by
renaming the camera keys — for video datasets via a download-free, server-side
Hub commit. Exposed as the ``lerobot-curate-cameras`` CLI.
"""
from .config import DEFAULT_VIEW_VOCABULARY, CameraCurationConfig
from .curator import (
CameraVerdict,
build_name_mapping,
build_report,
curate_cameras,
is_valid_view_label,
rename_camera_keys_on_hub,
write_report,
)
__all__ = [
"DEFAULT_VIEW_VOCABULARY",
"CameraCurationConfig",
"CameraVerdict",
"build_name_mapping",
"build_report",
"curate_cameras",
"is_valid_view_label",
"rename_camera_keys_on_hub",
"write_report",
]
@@ -0,0 +1,80 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Config for ``lerobot-curate-cameras`` (VLM camera-view curation)."""
from __future__ import annotations
from dataclasses import dataclass, field
from pathlib import Path
from lerobot.annotations.steerable_pipeline.config import AnnotationJobConfig, VlmConfig
# The closed vocabulary of canonical camera-view labels. Combos are formed by
# joining two of these with ``_`` (e.g. ``left_wrist``).
DEFAULT_VIEW_VOCABULARY: tuple[str, ...] = ("top", "wrist", "front", "bottom", "left", "right")
@dataclass
class CameraCurationConfig:
"""Top-level config for ``lerobot-curate-cameras``.
The VLM decision only ever reads the first episode (a cheap partial
download). ``mode="report"`` writes the labels + quality verdicts into
``meta/`` and moves nothing (works for any dataset). ``mode="rename"``
additionally renames the camera keys to ``observation.images.<label>`` —
for video datasets this is a server-side, download-free Hub commit.
"""
# Hub dataset id (downloaded when ``root`` is unset) — also the rename target.
repo_id: str | None = None
# Local dataset directory (skips the Hub download).
root: Path | None = None
# "report": write mapping + verdicts into meta/, no file moves.
# "rename": physically rename camera keys to observation.images.<label>.
mode: str = "report"
# Commit target branch for the Hub rename; keeps ``main`` intact when set.
# None commits to the default branch.
branch: str | None = None
# Episode inspected by the VLM (first episode by default).
episode_index: int = 0
# Frames sampled from that episode per camera and shown to the VLM.
n_frames: int = 4
# Closed label vocabulary and whether two-token combos (left_wrist) are allowed.
view_vocabulary: tuple[str, ...] = DEFAULT_VIEW_VOCABULARY
allow_combos: bool = True
# "error" raises on colliding target labels; "suffix" disambiguates (top -> top_2).
on_collision: str = "error"
# Remove cameras judged unusable (default: only flag them, still rename).
drop_unusable: bool = False
# Where to write the machine-readable report (default <root>/meta/camera_curation.json).
report_path: Path | None = None
vlm: VlmConfig = field(default_factory=VlmConfig)
job: AnnotationJobConfig = field(default_factory=AnnotationJobConfig)
seed: int = 1729
# Keyframe decode backend forwarded to ``decode_video_frames`` (None = default).
video_backend: str | None = None
# Upload the result (rename mode). Kept off by default so runs are dry.
push_to_hub: bool = False
push_commit_message: str | None = None
@@ -0,0 +1,349 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Camera-view curation: per-camera VLM quality + label judgments and the
lightweight (download-free) Hub rename that applies the chosen labels.
The decision (:func:`curate_cameras`) is a pure function of a
``{camera_key: [frames]}`` map and a VLM client, so it unit-tests with a stub
VLM and no dataset. The orchestrating CLI (``lerobot-curate-cameras``) samples
those frames from the dataset's first episode.
"""
from __future__ import annotations
import logging
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any
from lerobot.annotations.steerable_pipeline.frames import to_image_blocks
from lerobot.datasets.dataset_tools import _remap_camera_key_in_meta, _resolve_rename_collisions
from lerobot.datasets.io_utils import load_info, write_info
from lerobot.utils.io_utils import write_json
from .config import CameraCurationConfig
logger = logging.getLogger(__name__)
_PROMPT_PATH = Path(__file__).parent / "prompts" / "camera_curation.txt"
# The canonical prefix every curated camera key gets.
OBS_IMAGE_PREFIX = "observation.images."
@dataclass
class CameraVerdict:
"""One camera's VLM verdict."""
camera_key: str
usable: bool
view_label: str | None
blur_reason: str | None = None
confidence: float | None = None
# Populated by ``build_name_mapping`` once collisions are resolved.
proposed_new_key: str | None = None
def _load_prompt() -> str:
return _PROMPT_PATH.read_text(encoding="utf-8")
def is_valid_view_label(label: str, vocabulary: tuple[str, ...], allow_combos: bool) -> bool:
"""True if ``label`` is a single vocab word, or (when allowed) an underscore
combo of at most two distinct vocab words."""
if not label:
return False
tokens = label.split("_")
if not allow_combos:
return len(tokens) == 1 and tokens[0] in vocabulary
if not (1 <= len(tokens) <= 2):
return False
return all(tok in vocabulary for tok in tokens) and len(set(tokens)) == len(tokens)
def _build_messages(frames: list[Any], cfg: CameraCurationConfig) -> list[dict[str, Any]]:
if cfg.allow_combos:
combo_rule = (
"You may combine at most two of these words with an underscore when "
"one word is not precise enough (e.g. \"left_wrist\"). "
)
else:
combo_rule = "Use exactly one of these words (no combinations). "
prompt = _load_prompt().format(
vocabulary=", ".join(cfg.view_vocabulary),
combo_rule=combo_rule,
)
content = [*to_image_blocks(frames), {"type": "text", "text": prompt}]
return [{"role": "user", "content": content}]
def _parse_verdict(camera_key: str, result: Any, cfg: CameraCurationConfig) -> CameraVerdict:
"""Turn a parsed VLM JSON object into a :class:`CameraVerdict` (defensively)."""
if not isinstance(result, dict):
return CameraVerdict(camera_key=camera_key, usable=True, view_label=None, blur_reason=None)
usable = bool(result.get("usable", True))
blur_reason = result.get("blur_reason")
blur_reason = str(blur_reason) if blur_reason else None
raw_label = result.get("view_label")
label = str(raw_label).strip().lower().replace(" ", "_") if raw_label else ""
view_label = label if is_valid_view_label(label, cfg.view_vocabulary, cfg.allow_combos) else None
if raw_label and view_label is None:
logger.warning(
"camera %s: VLM returned view_label=%r which is not in the vocabulary %s; leaving unlabeled",
camera_key,
raw_label,
cfg.view_vocabulary,
)
confidence = result.get("confidence")
try:
confidence = float(confidence) if confidence is not None else None
except (TypeError, ValueError):
confidence = None
return CameraVerdict(
camera_key=camera_key,
usable=usable,
view_label=view_label,
blur_reason=blur_reason,
confidence=confidence,
)
def curate_cameras(
frames_by_camera: dict[str, list[Any]],
cfg: CameraCurationConfig,
vlm: Any,
) -> list[CameraVerdict]:
"""Judge each camera's quality + view label from a few sampled frames.
``frames_by_camera`` maps a camera key to a list of decoded frames (torch
tensors or PIL images). One batched ``generate_json`` call is issued across
all cameras. Cameras with no frames are still reported (usable, unlabeled)
so the caller sees the full camera set.
"""
ordered_keys = list(frames_by_camera)
callable_keys = [k for k in ordered_keys if frames_by_camera[k]]
verdicts: dict[str, CameraVerdict] = {
k: CameraVerdict(camera_key=k, usable=True, view_label=None) for k in ordered_keys
}
if callable_keys:
messages_batch = [_build_messages(frames_by_camera[k], cfg) for k in callable_keys]
results = vlm.generate_json(messages_batch)
for key, result in zip(callable_keys, results, strict=True):
verdicts[key] = _parse_verdict(key, result, cfg)
return [verdicts[k] for k in ordered_keys]
def build_name_mapping(
verdicts: list[CameraVerdict],
existing_features: dict[str, dict],
cfg: CameraCurationConfig,
) -> dict[str, str]:
"""Compute ``{old_key: observation.images.<label>}`` for labeled cameras.
Cameras without a valid label (or already at their canonical name) are
skipped. Collisions are resolved with ``cfg.on_collision`` and the resolved
target is written back onto each verdict's ``proposed_new_key``.
"""
desired: dict[str, str] = {}
for v in verdicts:
if v.view_label is None:
continue
target = f"{OBS_IMAGE_PREFIX}{v.view_label}"
if target != v.camera_key:
desired[v.camera_key] = target
if not desired:
return {}
resolved = _resolve_rename_collisions(desired, existing_features, cfg.on_collision)
by_key = {v.camera_key: v for v in verdicts}
for old, new in resolved.items():
by_key[old].proposed_new_key = new
return resolved
def build_report(
verdicts: list[CameraVerdict],
mapping: dict[str, str],
cfg: CameraCurationConfig,
) -> dict[str, Any]:
"""Assemble the machine-readable curation report."""
return {
"repo_id": cfg.repo_id,
"episode_index": cfg.episode_index,
"view_vocabulary": list(cfg.view_vocabulary),
"cameras": {
v.camera_key: {
"view_label": v.view_label,
"usable": v.usable,
"blur_reason": v.blur_reason,
"confidence": v.confidence,
"proposed_new_key": mapping.get(v.camera_key),
}
for v in verdicts
},
}
def write_report(
root: Path,
verdicts: list[CameraVerdict],
mapping: dict[str, str],
cfg: CameraCurationConfig,
) -> Path:
"""Write ``meta/camera_curation.json`` and stamp verdicts into ``info.json``.
Stamping goes into each camera's ``features[key]["info"]["curation"]`` so the
verdict travels with the dataset. Returns the report path.
"""
report = build_report(verdicts, mapping, cfg)
default_report_path = root / "meta" / "camera_curation.json"
report_path = Path(cfg.report_path) if cfg.report_path is not None else default_report_path
report_path.parent.mkdir(parents=True, exist_ok=True)
write_json(report, report_path)
info = load_info(root)
changed = False
for v in verdicts:
feature = info.features.get(v.camera_key)
if feature is None:
continue
feature.setdefault("info", {})
if feature["info"] is None:
feature["info"] = {}
feature["info"]["curation"] = {
"view_label": v.view_label,
"usable": v.usable,
"blur_reason": v.blur_reason,
"confidence": v.confidence,
}
changed = True
if changed:
write_info(info, root)
return report_path
def _swap_key_in_path(path: str, old_key: str, new_key: str) -> str:
"""Rewrite the ``<old_key>`` path segment of a ``videos/<key>/...`` repo path."""
prefix = f"videos/{old_key}/"
return f"videos/{new_key}/{path[len(prefix):]}" if path.startswith(prefix) else path
def rename_camera_keys_on_hub(
repo_id: str,
name_mapping: dict[str, str],
local_root: Path,
*,
revision: str | None = None,
branch: str | None = None,
token: str | None = None,
commit_message: str | None = None,
) -> Any:
"""Rename camera keys on the Hub without downloading video data.
Edits the small ``meta/`` files locally (under ``local_root``, which must be
a writable dataset root whose ``meta/`` is already present), then commits, in
one atomic ``create_commit``: ``CommitOperationCopy`` + ``CommitOperationDelete``
to move each ``videos/<old>/*`` LFS file server-side, and ``CommitOperationAdd``
for the edited meta files. Renames in place on ``repo_id`` (cross-repo copies
are unsupported); pass ``branch`` to commit to a branch and keep ``main`` intact.
Only video keys can be moved this way — reject swaps/cycles and image keys
(handled by the local ``rename_features`` path instead).
"""
from huggingface_hub import CommitOperationAdd, CommitOperationCopy, CommitOperationDelete, HfApi
# A swap/cycle (a target that is also a source) cannot be expressed in a
# single base-revision commit; defer to the local rename path.
swaps = set(name_mapping.values()) & set(name_mapping)
if swaps:
raise NotImplementedError(
f"Hub rename cannot swap keys in one commit (offending: {sorted(swaps)}); "
"use the local rename_features path for swaps/cycles."
)
# Determine which OLD keys are video-stored (only those have a videos/ tree)
# BEFORE remapping the metadata.
info = load_info(local_root)
video_old_keys = {
old for old in name_mapping if info.features.get(old, {}).get("dtype") == "video"
}
image_old_keys = {
old for old in name_mapping if info.features.get(old, {}).get("dtype") == "image"
}
if image_old_keys:
raise NotImplementedError(
f"Hub rename cannot move image data stored in the data parquet (keys: {sorted(image_old_keys)}); "
"use --mode report (metadata mapping) or the local rename_features path for image datasets."
)
# 1. Rewrite meta/ locally (info features, episodes columns, stats keys).
_remap_camera_key_in_meta(local_root, name_mapping)
api = HfApi(token=token)
operations: list[Any] = []
# 2. Add the (small) meta files we just edited.
meta_dir = local_root / "meta"
meta_files = [meta_dir / "info.json"]
stats_file = meta_dir / "stats.json"
if stats_file.exists():
meta_files.append(stats_file)
meta_files.extend(sorted((meta_dir / "episodes").glob("*/*.parquet")))
for fpath in meta_files:
rel = fpath.relative_to(local_root).as_posix()
operations.append(CommitOperationAdd(path_in_repo=rel, path_or_fileobj=str(fpath)))
# 3. Move video LFS files server-side (copy + delete), no download.
repo_files = api.list_repo_files(repo_id, repo_type="dataset", revision=revision)
n_moved = 0
for old in video_old_keys:
new = name_mapping[old]
prefix = f"videos/{old}/"
for f in repo_files:
if f.startswith(prefix):
operations.append(
CommitOperationCopy(src_path_in_repo=f, path_in_repo=_swap_key_in_path(f, old, new))
)
operations.append(CommitOperationDelete(path_in_repo=f))
n_moved += 1
logger.info(
"hub rename: moving %d video file(s) server-side across %d camera(s)",
n_moved,
len(video_old_keys),
)
commit_info = api.create_commit(
repo_id=repo_id,
repo_type="dataset",
operations=operations,
revision=branch or revision,
commit_message=commit_message or "curate: rename camera views (lerobot-curate-cameras)",
)
return commit_info
def as_report_dict(verdicts: list[CameraVerdict]) -> list[dict[str, Any]]:
"""Convenience: verdicts as plain dicts (for logging/JSON)."""
return [asdict(v) for v in verdicts]
@@ -0,0 +1,28 @@
You are inspecting frames from ONE camera of a robot manipulation dataset. All
frames come from the same fixed camera during a single episode; use them
together to judge the camera, not any single moment.
Do two things and return them as one JSON object.
1. QUALITY. Decide whether this camera view is usable for training a policy.
Mark it UNUSABLE if it is blurry / out of focus, badly over- or
under-exposed, mostly occluded, static/frozen, corrupted, or otherwise does
not clearly show the scene. Otherwise it is usable.
2. VIEW LABEL. Choose the single best label for where this camera is mounted /
what it looks at, using ONLY this closed vocabulary:
{vocabulary}
{combo_rule}Pick the label that best matches the viewpoint (e.g. a
downward overhead shot is "top"; a camera on the robot's gripper/hand that
moves with the arm is "wrist"). Do not invent words outside the vocabulary.
Output strictly valid JSON, no prose, no code fences, with exactly these keys:
{{
"usable": true or false,
"blur_reason": "<short reason if unusable, else null>",
"view_label": "<one label from the vocabulary, combos joined by '_'>",
"confidence": <number between 0 and 1>
}}
+2
View File
@@ -33,6 +33,7 @@ from .dataset_tools import (
recompute_stats,
reencode_dataset,
remove_feature,
rename_features,
split_dataset,
)
from .factory import make_dataset, make_train_eval_datasets, resolve_delta_timestamps
@@ -96,6 +97,7 @@ __all__ = [
"recompute_stats",
"reencode_dataset",
"remove_feature",
"rename_features",
"resolve_delta_timestamps",
"safe_stop_image_writer",
"split_dataset",
+249
View File
@@ -47,6 +47,7 @@ from lerobot.configs import (
)
from lerobot.configs.video import DEPTH_ENCODER_INFO_FIELD_NAMES
from lerobot.utils.constants import ACTION, HF_LEROBOT_HOME, OBS_IMAGE, OBS_STATE
from lerobot.utils.io_utils import load_json, write_json
from lerobot.utils.utils import flatten_dict
from .aggregate import aggregate_datasets
@@ -60,6 +61,8 @@ from .image_writer import write_image
from .io_utils import (
get_parquet_file_size_in_mb,
load_episodes,
load_info,
to_parquet_one_row_group_per_episode,
write_info,
write_stats,
write_tasks,
@@ -72,7 +75,9 @@ from .utils import (
DEFAULT_DATA_PATH,
DEFAULT_EPISODES_PATH,
DEPTH_FILE_PATTERN,
EPISODES_DIR,
IMAGE_FILE_PATTERN,
STATS_PATH,
VIDEO_DIR,
update_chunk_file_indices,
)
@@ -484,6 +489,250 @@ def remove_feature(
)
# Columns in ``meta/episodes/*.parquet`` are namespaced by feature key under
# these prefixes (e.g. ``videos/observation.images.top/from_timestamp`` and
# ``stats/observation.images.top/mean``). Renaming a feature means rewriting the
# middle ``<key>`` segment of every such column. Note ``stats/*`` columns are
# invisible via ``meta.episodes`` (``load_episodes`` drops them), so we operate
# on the raw parquet.
_EPISODE_KEY_PREFIXES = ("videos", "stats")
# Features that must never be renamed (or become a rename target): the dataset
# indexing/bookkeeping columns.
_REQUIRED_FEATURES = frozenset({"timestamp", "frame_index", "episode_index", "index", "task_index"})
def _resolve_rename_collisions(
name_mapping: dict[str, str],
existing_features: dict[str, dict],
on_collision: str,
) -> dict[str, str]:
"""Validate/disambiguate a ``{old_key: new_key}`` mapping against collisions.
The post-rename key set is ``(features \\ sources) targets``. A collision is
either two sources mapping to the same target, or a target equal to an
untouched existing key. Swaps/cycles between sources are *not* collisions
(handled downstream). ``on_collision="error"`` raises listing every offending
pair; ``"suffix"`` disambiguates deterministically (``top`` ``top_2`` )
in sorted-source order.
"""
if on_collision not in ("error", "suffix"):
raise ValueError(f"on_collision must be 'error' or 'suffix', got {on_collision!r}")
sources = set(name_mapping)
untouched = set(existing_features) - sources
targets = list(name_mapping.values())
duplicate_targets = {t for t in targets if targets.count(t) > 1}
untouched_collisions = set(targets) & untouched
if on_collision == "error":
problems = []
if duplicate_targets:
problems.append(f"multiple cameras map to the same target(s): {sorted(duplicate_targets)}")
if untouched_collisions:
problems.append(
f"target(s) collide with existing feature(s) not being renamed: "
f"{sorted(untouched_collisions)}"
)
if problems:
raise ValueError(
"rename_features collision(s): "
+ "; ".join(problems)
+ ". Resolve the labels (e.g. use combos like 'left_wrist') or pass "
"on_collision='suffix'."
)
return dict(name_mapping)
# suffix mode: greedily de-collide in a deterministic (sorted) order.
used = set(untouched)
resolved: dict[str, str] = {}
for src in sorted(name_mapping):
target = name_mapping[src]
if target in used:
base, i = target, 2
while target in used:
target = f"{base}_{i}"
i += 1
resolved[src] = target
used.add(target)
return resolved
def _remap_camera_key_in_meta(root: Path, name_mapping: dict[str, str]) -> None:
"""Rename feature keys across the dataset's ``meta/`` files (no file moves).
Touches: ``meta/info.json`` ``features`` (key renamed, feature dict carried
verbatim so codec ``info`` / depth params survive), every
``meta/episodes/*/*.parquet`` (``videos/<old>/*`` and ``stats/<old>/*``
columns), and ``meta/stats.json`` (top-level ``<old>`` key). All three are
simultaneous relabels, so swaps/cycles are safe here.
"""
# info.json — rebuild features preserving insertion order.
info = load_info(root)
info.features = {name_mapping.get(key, key): ft for key, ft in info.features.items()}
write_info(info, root)
# episodes parquet — rename namespaced columns by prefix.
def _rename_column(col: str) -> str:
for prefix in _EPISODE_KEY_PREFIXES:
head = f"{prefix}/"
if col.startswith(head):
rest = col[len(head) :]
for old, new in name_mapping.items():
if rest == old or rest.startswith(f"{old}/"):
return f"{head}{new}{rest[len(old) :]}"
return col
for path in sorted((root / EPISODES_DIR).glob("*/*.parquet")):
df = pd.read_parquet(path)
col_map = {c: _rename_column(c) for c in df.columns if _rename_column(c) != c}
if col_map:
df = df.rename(columns=col_map)
to_parquet_one_row_group_per_episode(df, path)
# stats.json — remap top-level feature keys.
stats_path = root / STATS_PATH
if stats_path.exists():
stats = load_json(stats_path)
if isinstance(stats, dict):
stats = {name_mapping.get(key, key): value for key, value in stats.items()}
write_json(stats, stats_path)
def _move_camera_key_dirs(root: Path, name_mapping: dict[str, str]) -> None:
"""Move ``videos/<old>`` and ``images/<old>`` trees to their new key names.
Two-phase (source sentinel target) so a swap like ``{a: b, b: a}`` cannot
clobber. Missing source dirs are skipped (a key may be stored one way only).
"""
for subdir in (VIDEO_DIR, "images"):
base = root / subdir
if not base.exists():
continue
# Phase 1: move every source to a unique sentinel.
sentinels: dict[str, Path] = {}
for i, old in enumerate(name_mapping):
src = base / old
if src.exists():
sentinel = base / f".__rename_tmp_{i}__"
shutil.move(str(src), str(sentinel))
sentinels[old] = sentinel
# Phase 2: sentinel → final target.
for old, sentinel in sentinels.items():
shutil.move(str(sentinel), str(base / name_mapping[old]))
def _rename_image_data_columns(root: Path, name_mapping: dict[str, str]) -> None:
"""Rename image-feature columns inside ``data/*.parquet`` at the Arrow level.
Image datasets embed frames as HF ``Image()`` columns in the data parquet.
We rename the Arrow field *and* the matching key in the schema-level
``huggingface`` metadata (which references columns by name), so no pixel
bytes are decoded or re-embedded and ``datasets`` still types the column as
an image after the rename.
"""
import json
data_dir = root / DATA_DIR
if not data_dir.exists():
return
for path in sorted(data_dir.glob("*/*.parquet")):
table = pq.read_table(path)
col_map = {c: name_mapping[c] for c in table.column_names if c in name_mapping}
if not col_map:
continue
table = table.rename_columns([col_map.get(c, c) for c in table.column_names])
metadata = dict(table.schema.metadata or {})
hf_key = b"huggingface"
if hf_key in metadata:
hf_meta = json.loads(metadata[hf_key])
features = hf_meta.get("info", {}).get("features")
if isinstance(features, dict):
for old, new in col_map.items():
if old in features:
features[new] = features.pop(old)
metadata[hf_key] = json.dumps(hf_meta).encode()
table = table.replace_schema_metadata(metadata)
pq.write_table(table, str(path))
def rename_features(
dataset: LeRobotDataset,
name_mapping: dict[str, str],
output_dir: str | Path | None = None,
repo_id: str | None = None,
*,
on_collision: str = "error",
) -> LeRobotDataset:
"""Rename dataset feature keys without re-encoding any pixel data.
A rename changes zero frame content, so this does a cheap key-remap rather
than the full-copy ``modify_features`` path (which would re-embed images and
byte-copy videos). It rewrites ``meta/`` (info features, episodes
``videos/*``+``stats/*`` columns, stats.json keys), moves the physical
``videos/<key>/`` (and ``images/<key>/``) directories, and for image
datasets renames the embedded ``data/*.parquet`` image column at the Arrow
level. Feature ``info`` dicts (video codec params, depth ``is_depth_map``) are
carried verbatim.
Args:
dataset: The source LeRobotDataset.
name_mapping: ``{old_feature_key: new_feature_key}``. Identity pairs are
ignored. Typically used to canonicalize camera keys, e.g.
``{"observation.images.cam_0": "observation.images.left_wrist"}``.
output_dir: Where the renamed dataset is written. Defaults to
``$HF_LEROBOT_HOME/repo_id``. When it equals ``dataset.root`` the
rename is applied in place.
repo_id: Identifier for the renamed dataset (default ``<repo_id>_renamed``).
on_collision: ``"error"`` (default) raises on colliding targets;
``"suffix"`` disambiguates deterministically (``top`` ``top_2``).
Returns:
The renamed LeRobotDataset.
"""
if not name_mapping:
raise ValueError("name_mapping must be a non-empty {old_key: new_key} dict")
features = dataset.meta.features
mapping = {old: new for old, new in name_mapping.items() if old != new}
if not mapping:
raise ValueError("name_mapping only contains identity renames (old == new); nothing to do")
missing = [old for old in mapping if old not in features]
if missing:
raise ValueError(f"Feature(s) not found in dataset: {missing}")
bad_required = sorted(
{name for pair in mapping.items() for name in pair if name in _REQUIRED_FEATURES}
)
if bad_required:
raise ValueError(f"Cannot rename to/from required features: {bad_required}")
bad_names = [new for new in mapping.values() if "/" in new]
if bad_names:
raise ValueError(f"Target feature name(s) cannot contain '/': {bad_names}")
mapping = _resolve_rename_collisions(mapping, features, on_collision)
if repo_id is None:
repo_id = f"{dataset.repo_id}_renamed"
output_dir = Path(output_dir) if output_dir is not None else HF_LEROBOT_HOME / repo_id
in_place = output_dir.resolve() == Path(dataset.root).resolve()
if not in_place:
shutil.copytree(dataset.root, output_dir)
image_keys = set(dataset.meta.image_keys)
_remap_camera_key_in_meta(output_dir, mapping)
_move_camera_key_dirs(output_dir, mapping)
image_mapping = {old: new for old, new in mapping.items() if old in image_keys}
if image_mapping:
_rename_image_data_columns(output_dir, image_mapping)
return LeRobotDataset(repo_id=repo_id, root=output_dir)
def _fractions_to_episode_indices(
total_episodes: int,
splits: dict[str, float],
+112
View File
@@ -0,0 +1,112 @@
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Run ``lerobot-curate-cameras`` on HF Jobs (HuggingFace GPUs).
Same shape as the annotation submitter (``lerobot.jobs.annotate``): the VLM
decision needs a GPU, so the pod boots the ``vllm/vllm-openai`` image, installs
lerobot on top, and replays the user's CLI with ``lerobot-curate-cameras``. The
``--mode=rename`` commit runs from the pod (which holds ``HF_TOKEN``); a bare
``--mode=report`` run leaves its output only on the pod, so we warn.
"""
from __future__ import annotations
import shlex
import sys
from dataclasses import is_dataclass
from typing import TYPE_CHECKING
from huggingface_hub import HfApi, get_token, run_job
from .annotate import build_pod_setup
from .dataset import ensure_dataset_available
from .hf import _pod_forwarded_args, follow_job, resolve_job_tags
if TYPE_CHECKING:
from lerobot.annotations.camera_curation.config import CameraCurationConfig
# Same rationale as the annotate submitter: --root is host-local, --repo_id is
# re-emitted, config files can't be read on the pod, and --job could smuggle a
# remote target back onto the pod.
_SUBMITTER_OWNED_ARGS = ("--root", "--repo_id", "--config_path", "--job")
def _local_config_file_args(cfg: CameraCurationConfig) -> list[str]:
return ["--config_path", *(f"--{name}" for name in vars(cfg) if is_dataclass(getattr(cfg, name)))]
def build_pod_command(repo_id: str, lerobot_ref: str, argv: list[str]) -> list[str]:
"""``bash -c`` command the pod runs: setup prelude, then curate-cameras."""
forwarded = _pod_forwarded_args(argv, drop_names=_SUBMITTER_OWNED_ARGS, drop_prefixes=("--job.",))
curate = shlex.join(
["lerobot-curate-cameras", f"--repo_id={repo_id}", *forwarded, "--job.target=local"]
)
return ["bash", "-c", f"{build_pod_setup(lerobot_ref)} && {curate}"]
def submit_curate_to_hf(cfg: CameraCurationConfig) -> None:
"""Submit a camera-curation run to HF Jobs and tail its logs."""
token = get_token()
if not token:
raise RuntimeError("Not logged in to Hugging Face. Run `hf auth login` first.")
if cfg.repo_id is None:
raise ValueError(
"Remote curation requires --repo_id: the pod downloads the dataset from the Hub, "
"and --root only names a directory on this machine."
)
argv = sys.argv[1:]
passed = {tok.split("=", 1)[0] for tok in argv}
used_config_files = sorted(passed.intersection(_local_config_file_args(cfg)))
if used_config_files:
raise ValueError(
f"{', '.join(used_config_files)} cannot be used with a remote --job.target: the pod "
"cannot read config files from this machine. Pass the settings as CLI flags instead."
)
if cfg.mode == "report":
print(
"WARNING: --mode=report writes its result into the pod's local copy, which is discarded "
"when the job ends. Use --mode=rename to commit the result to the Hub."
)
api = HfApi(token=token)
tags = resolve_job_tags(cfg.job.tags)
ensure_dataset_available(cfg.repo_id, api=api, tags=tags)
command = build_pod_command(cfg.repo_id, cfg.job.lerobot_ref, argv)
print(f"Submitting job to HF Jobs (flavor={cfg.job.target}, image={cfg.job.image}) ...")
job_info = run_job(
image=cfg.job.image,
command=command,
flavor=cfg.job.target,
secrets={"HF_TOKEN": token},
timeout=cfg.job.timeout,
labels=dict.fromkeys(tags, "true"),
)
job_id = job_info.id
job_url = getattr(job_info, "url", None)
print(f"Job submitted: {job_id}")
if job_url:
print(f" Job page: {job_url}")
print(f" Dataset repo: https://huggingface.co/datasets/{cfg.repo_id}")
print(f" Monitor: hf jobs logs {job_id}")
print(f" Cancel: hf jobs cancel {job_id}")
if not follow_job(job_id, detach=cfg.job.detach):
return
print("\nCuration complete.")
@@ -0,0 +1,244 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""``lerobot-curate-cameras`` — VLM camera-view curation for a LeRobot dataset.
Downloads only the first episode, asks a VLM to (1) flag blurry/unusable views
and (2) label each view (``top``/``wrist``/``front``/), then either records the
result in ``meta/`` (``--mode=report``) or renames the camera keys to
``observation.images.<label>`` (``--mode=rename``). For video datasets the
rename is a download-free, server-side Hub commit.
Examples:
# Cheap, mutation-free triage (writes meta/camera_curation.json):
uv run lerobot-curate-cameras --repo_id=user/dataset --mode=report
# Apply the labels by renaming camera keys on a new branch (video datasets):
uv run lerobot-curate-cameras --repo_id=user/dataset --mode=rename --branch=curated
# Run the VLM decision on a GPU via HF Jobs:
uv run lerobot-curate-cameras --repo_id=user/dataset --mode=rename --job.target=h200
"""
import logging
import shutil
import tempfile
from pathlib import Path
from typing import TYPE_CHECKING, Any
from lerobot.annotations.camera_curation import curator
from lerobot.annotations.camera_curation.config import CameraCurationConfig
from lerobot.annotations.steerable_pipeline.frames import make_frame_provider
from lerobot.annotations.steerable_pipeline.reader import iter_episodes
from lerobot.annotations.steerable_pipeline.vlm_client import make_vlm_client
from lerobot.configs import parser
from lerobot.utils.constants import HF_LEROBOT_HOME
from lerobot.utils.import_utils import _datasets_available, require_package
if TYPE_CHECKING or _datasets_available:
from lerobot.datasets.lerobot_dataset import LeRobotDataset
logger = logging.getLogger(__name__)
def _resolve_root(cfg: CameraCurationConfig) -> Path:
"""Concrete, writable root for the dataset (never the symlinked snapshot cache)."""
if cfg.root is not None:
return Path(cfg.root)
if cfg.repo_id is not None:
return HF_LEROBOT_HOME / cfg.repo_id
raise ValueError("Either --repo_id or --root must be provided.")
def _uniform_indices(n: int, k: int) -> list[int]:
if n <= 0 or k <= 0:
return []
if k >= n:
return list(range(n))
step = (n - 1) / (k - 1) if k > 1 else 0.0
return sorted({round(i * step) for i in range(k)})
def _to_uint8_frame(frame: Any) -> Any:
"""Scale a float [0,1] image tensor to uint8; pass uint8/PIL through."""
import torch
if isinstance(frame, torch.Tensor) and torch.is_floating_point(frame):
return (frame.clamp(0, 1) * 255).to(torch.uint8)
return frame
def _sample_frames(dataset: "LeRobotDataset", cfg: CameraCurationConfig) -> dict[str, list[Any]]:
"""Sample ``n_frames`` from the inspected episode for each (non-depth) camera.
Video cameras go through the annotation frame provider (uint8 frames); image
cameras are read straight from the dataset rows and scaled to uint8.
"""
meta = dataset.meta
depth_keys = set(meta.depth_keys)
video_keys = set(meta.video_keys)
image_keys = set(meta.image_keys)
cameras = [k for k in meta.camera_keys if k not in depth_keys]
frames: dict[str, list[Any]] = {k: [] for k in cameras}
video_cameras = [k for k in cameras if k in video_keys]
if video_cameras:
provider = make_frame_provider(dataset.root, video_backend=cfg.video_backend)
records = list(iter_episodes(dataset.root, only_episodes=(cfg.episode_index,)))
record = records[0] if records else None
if record is not None:
for key in video_cameras:
frames[key] = provider.video_for_episode(record, cfg.n_frames, camera_key=key)
image_cameras = [k for k in cameras if k in image_keys]
if image_cameras:
n = len(dataset)
for i in _uniform_indices(n, cfg.n_frames):
item = dataset[i]
for key in image_cameras:
if key in item:
frames[key].append(_to_uint8_frame(item[key]))
return frames
def _apply_rename(
root: Path,
dataset: "LeRobotDataset",
cfg: CameraCurationConfig,
mapping: dict[str, str],
verdicts: list["curator.CameraVerdict"],
) -> None:
"""Apply the computed ``{old: new}`` camera-key mapping.
Video datasets on the Hub download-free server-side rename commit.
Otherwise (image dataset, local-only, or a swap/cycle) local
``rename_features`` over a full copy of the dataset.
"""
video_keys = set(dataset.meta.video_keys)
all_video = set(mapping) <= video_keys
has_swap = bool(set(mapping.values()) & set(mapping))
if cfg.repo_id is not None and all_video and not has_swap:
if cfg.drop_unusable:
logger.warning(
"--drop_unusable is only applied via the local rename path; the Hub rename keeps "
"flagged views (they are still recorded in meta/). Re-run with a local --root to drop."
)
# Edit a throwaway copy of meta/ so the local cached copy stays pristine
# and only the intended files land in the commit.
work = Path(tempfile.mkdtemp(prefix="lerobot_curate_"))
try:
shutil.copytree(root / "meta", work / "meta")
commit = curator.rename_camera_keys_on_hub(
cfg.repo_id,
mapping,
work,
branch=cfg.branch,
commit_message=cfg.push_commit_message,
)
oid = getattr(commit, "oid", None)
ref = cfg.branch or "main"
logger.info("Hub rename committed to %s@%s (%s)", cfg.repo_id, ref, oid)
finally:
shutil.rmtree(work, ignore_errors=True)
return
# Local path: needs the full dataset, so re-load without the episode filter.
require_package("datasets", "dataset")
from lerobot.datasets import LeRobotDataset as _LeRobotDataset
from lerobot.datasets import remove_feature, rename_features
logger.info("Local rename path (image/local/swap): loading the full dataset from %s", root)
full = _LeRobotDataset(cfg.repo_id or "local", root=root)
renamed = rename_features(
full, mapping, output_dir=root, repo_id=full.repo_id, on_collision=cfg.on_collision
)
if cfg.drop_unusable:
unusable_new_keys = [
mapping[v.camera_key] for v in verdicts if not v.usable and v.camera_key in mapping
]
if unusable_new_keys:
logger.info("Dropping unusable views: %s", unusable_new_keys)
renamed = remove_feature(renamed, unusable_new_keys, output_dir=root, repo_id=renamed.repo_id)
if cfg.push_to_hub:
logger.info("Pushing renamed dataset to %s", renamed.repo_id)
renamed.push_to_hub()
@parser.wrap()
def curate_cameras(cfg: CameraCurationConfig) -> None:
"""Run the camera-view curation pipeline over a dataset's first episode."""
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
if cfg.mode not in ("report", "rename"):
raise ValueError(f"--mode must be 'report' or 'rename', got {cfg.mode!r}")
if cfg.job.is_remote:
from lerobot.jobs.curate import submit_curate_to_hf
return submit_curate_to_hf(cfg)
require_package("datasets", "dataset")
from lerobot.datasets.lerobot_dataset import LeRobotDataset
root = _resolve_root(cfg)
logger.info("curate-cameras: repo_id=%s root=%s mode=%s", cfg.repo_id, root, cfg.mode)
# Only episode ``cfg.episode_index`` is fetched (a cheap partial download).
dataset = LeRobotDataset(
cfg.repo_id or "local",
root=root,
episodes=[cfg.episode_index],
download_videos=True,
)
frames = _sample_frames(dataset, cfg)
n_with_frames = sum(1 for v in frames.values() if v)
logger.info("curate-cameras: %d camera(s), %d with sampled frames", len(frames), n_with_frames)
vlm = make_vlm_client(cfg.vlm)
verdicts = curator.curate_cameras(frames, cfg, vlm)
for v in verdicts:
logger.info(
" %s -> label=%s usable=%s%s",
v.camera_key,
v.view_label,
v.usable,
"" if v.usable else f" (blur_reason={v.blur_reason!r})",
)
mapping = curator.build_name_mapping(verdicts, dataset.meta.features, cfg)
report_path = curator.write_report(dataset.root, verdicts, mapping, cfg)
logger.info("curate-cameras: report written to %s", report_path)
logger.info("curate-cameras: proposed rename mapping: %s", mapping or "(none)")
if cfg.mode == "rename":
if not mapping:
logger.info("curate-cameras: nothing to rename (no confident labels differ from current keys)")
return
_apply_rename(dataset.root, dataset, cfg, mapping, verdicts)
def main() -> None:
curate_cameras()
if __name__ == "__main__":
main()
@@ -87,11 +87,6 @@ import tqdm
from lerobot.configs import DEPTH_MILLIMETER_UNIT
from lerobot.datasets import LeRobotDataset
from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
from lerobot.utils.dataset_visualization_utils import (
get_extra_scalar_keys,
is_scalar_like,
scalar_to_float,
)
from lerobot.utils.utils import init_logging
logger = logging.getLogger(__name__)
@@ -158,8 +153,6 @@ def build_blueprint_from_dataset(dataset: LeRobotDataset):
for key in (DONE, REWARD, SUCCESS):
if key in dataset.features:
views.append(rrb.TimeSeriesView(origin=key, name=key))
for key in get_extra_scalar_keys(dataset):
views.append(rrb.TimeSeriesView(origin=key, name=key))
return rrb.Blueprint(rrb.Grid(*views))
@@ -251,8 +244,6 @@ def visualize_dataset(
hi = stats["q99"] if "q99" in stats else stats["max"]
depth_ranges[key] = (float(np.asarray(lo).item()), float(np.asarray(hi).item()))
extra_scalar_keys = get_extra_scalar_keys(dataset)
first_index = None
for batch in tqdm.tqdm(dataloader, total=len(dataloader)):
if first_index is None:
@@ -296,10 +287,6 @@ def visualize_dataset(
if SUCCESS in batch:
rr.log(SUCCESS, rr.Scalars(batch[SUCCESS][i].item()))
for key in extra_scalar_keys:
if key in batch and is_scalar_like(batch[key][i]):
rr.log(key, rr.Scalars(scalar_to_float(batch[key][i])))
# save .rrd locally
if mode == "local" and save:
output_dir.mkdir(parents=True, exist_ok=True)
@@ -108,6 +108,12 @@ Remove camera feature:
--operation.type remove_feature \
--operation.feature_names "['observation.image']"
Rename features/camera keys (no pixel data is re-encoded):
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type rename_features \
--operation.name_mapping '{"observation.images.cam_0": "observation.images.left_wrist"}'
Modify tasks - set a single task for all episodes (WARNING: modifies in-place):
lerobot-edit-dataset \
--repo_id lerobot/pusht \
@@ -257,6 +263,7 @@ from lerobot.datasets import (
recompute_stats,
reencode_dataset,
remove_feature,
rename_features,
split_dataset,
)
from lerobot.utils.constants import HF_LEROBOT_HOME
@@ -298,6 +305,16 @@ class RemoveFeatureConfig(OperationConfig):
feature_names: list[str] | None = None
@OperationConfig.register_subclass("rename_features")
@dataclass
class RenameFeaturesConfig(OperationConfig):
# Mapping of {old_feature_key: new_feature_key}, e.g.
# {"observation.images.cam_0": "observation.images.left_wrist"}.
name_mapping: dict[str, str] | None = None
# "error" raises on colliding targets; "suffix" disambiguates (top -> top_2).
on_collision: str = "error"
@OperationConfig.register_subclass("modify_tasks")
@dataclass
class ModifyTasksConfig(OperationConfig):
@@ -545,6 +562,42 @@ def handle_remove_feature(cfg: EditDatasetConfig) -> None:
LeRobotDataset(output_repo_id, root=output_dir).push_to_hub()
def handle_rename_features(cfg: EditDatasetConfig) -> None:
if not isinstance(cfg.operation, RenameFeaturesConfig):
raise ValueError("Operation config must be RenameFeaturesConfig")
if not cfg.operation.name_mapping:
raise ValueError("name_mapping must be specified for rename_features operation")
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
output_repo_id, output_dir = get_output_path(
cfg.repo_id,
new_repo_id=cfg.new_repo_id,
root=cfg.root,
new_root=cfg.new_root,
)
# In case of in-place modification, make the dataset point to the backup directory
if output_dir == dataset.root:
dataset.root = dataset.root.with_name(dataset.root.name + "_old")
logging.info(f"Renaming features {cfg.operation.name_mapping} in {cfg.repo_id}")
new_dataset = rename_features(
dataset,
name_mapping=cfg.operation.name_mapping,
output_dir=output_dir,
repo_id=output_repo_id,
on_collision=cfg.operation.on_collision,
)
logging.info(f"Dataset saved to {output_dir}")
logging.info(f"Features: {list(new_dataset.meta.features.keys())}")
if cfg.push_to_hub:
logging.info(f"Pushing to hub as {output_repo_id}")
LeRobotDataset(output_repo_id, root=output_dir).push_to_hub()
def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
if not isinstance(cfg.operation, ModifyTasksConfig):
raise ValueError("Operation config must be ModifyTasksConfig")
@@ -830,6 +883,8 @@ def edit_dataset(cfg: EditDatasetConfig) -> None:
handle_merge(cfg)
elif operation_type == "remove_feature":
handle_remove_feature(cfg)
elif operation_type == "rename_features":
handle_rename_features(cfg)
elif operation_type == "modify_tasks":
handle_modify_tasks(cfg)
elif operation_type == "convert_image_to_video":
@@ -1,99 +0,0 @@
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Shared helpers for visualizing scalar features from a LeRobot dataset."""
from collections.abc import Iterable, Mapping
from typing import TYPE_CHECKING
import numpy as np
import torch
from .constants import ACTION, DEFAULT_FEATURES, DONE, OBS_STATE, REWARD, SUCCESS
if TYPE_CHECKING:
from lerobot.datasets import LeRobotDataset
METADATA_KEYS = {*DEFAULT_FEATURES, "task"}
KNOWN_SCALAR_KEYS = {DONE, REWARD, SUCCESS}
SCALAR_DTYPE_KINDS = {"b", "i", "u", "f"}
def is_scalar_feature(feature: Mapping) -> bool:
"""Return whether a feature schema describes a numeric or boolean scalar."""
dtype = feature.get("dtype")
if not isinstance(dtype, str):
return False
try:
dtype_kind = np.dtype(dtype).kind
except (TypeError, ValueError):
return False
if dtype_kind not in SCALAR_DTYPE_KINDS:
return False
shape = feature.get("shape")
if shape is None:
return True
if isinstance(shape, int):
return shape == 1
if not isinstance(shape, (list, tuple)):
return False
return len(shape) == 0 or (len(shape) == 1 and shape[0] == 1)
def get_extra_scalar_keys(dataset: "LeRobotDataset", additional_known_keys: Iterable[str] = ()) -> list[str]:
"""Return scalar feature keys not handled by the visualizer's standard paths."""
known_keys = {
ACTION,
OBS_STATE,
*KNOWN_SCALAR_KEYS,
*METADATA_KEYS,
*additional_known_keys,
*dataset.meta.camera_keys,
}
return [
key
for key, feature in dataset.features.items()
if key not in known_keys and is_scalar_feature(feature)
]
def is_scalar_like(value: object) -> bool:
"""Return whether a runtime value contains exactly one numeric or boolean scalar."""
if isinstance(value, torch.Tensor):
return value.numel() == 1 and not value.is_complex()
if isinstance(value, np.ndarray):
return value.size == 1 and value.dtype.kind in SCALAR_DTYPE_KINDS
return np.isscalar(value) and np.asarray(value).dtype.kind in SCALAR_DTYPE_KINDS
def scalar_to_float(value: object) -> float:
"""Convert a scalar-like tensor, array, or Python value to ``float``."""
return float(value.item() if hasattr(value, "item") else value)
def get_scalar_values(sample: Mapping, keys: Iterable[str]) -> dict[str, float]:
"""Select and convert scalar-like values from ``sample`` for the requested keys."""
values = {}
for key in keys:
value = sample.get(key)
if value is not None and is_scalar_like(value):
values[key] = scalar_to_float(value)
return values
+10 -28
View File
@@ -23,7 +23,6 @@ importing from here directly. Requires the ``viz`` extra (``pip install 'lerobot
import logging
import numbers
import time
from collections.abc import Iterable, Mapping
import cv2
import numpy as np
@@ -42,7 +41,6 @@ from .constants import (
SUCCESS,
TRUNCATED,
)
from .dataset_visualization_utils import get_extra_scalar_keys, get_scalar_values
from .import_utils import require_package
# Static schema shared by all scalar topics. Each message carries a flat list of ``{label, value}``
@@ -407,24 +405,6 @@ def _frame_to_scalars(sample: dict, key: str, labels: list[str] | None = None) -
return _labeled_scalars(name, arr.flatten(), labels)
def _dataset_frame_scalar_groups(
sample: Mapping, extra_scalar_keys: Iterable[str]
) -> tuple[dict[str, float], dict[str, float]]:
"""Return standard episode scalars and custom scalar features for one dataset frame."""
episode_scalars = {}
for feature, label in (
(DONE, "done"),
(TRUNCATED, "truncated"),
(REWARD, "reward"),
(SUCCESS, "success"),
):
value = sample.get(feature)
if value is not None:
episode_scalars[label] = float(value)
return episode_scalars, get_scalar_values(sample, extra_scalar_keys)
def serve_foxglove_dataset_playback(
dataset,
episode_index: int,
@@ -472,7 +452,6 @@ def serve_foxglove_dataset_playback(
raise ValueError("Cannot visualize an empty episode.")
first_ns, last_ns = times_ns[0], times_ns[-1]
camera_keys = list(dataset.meta.camera_keys)
extra_scalar_keys = get_extra_scalar_keys(dataset, additional_known_keys=(TRUNCATED,))
# Dataset-wide q01/q99 depth bounds (fallback min/max) used to normalize depth to [0, 1].
depth_ranges: dict[str, tuple[float, float]] = {}
for key in dataset.meta.depth_keys:
@@ -521,14 +500,17 @@ def serve_foxglove_dataset_playback(
channels=channels,
log_time=log_time,
)
episode_scalars, extra_scalars = _dataset_frame_scalar_groups(sample, extra_scalar_keys)
episode_scalars = {}
for feat, label in (
(DONE, "done"),
(TRUNCATED, "truncated"),
(REWARD, "reward"),
(SUCCESS, "success"),
):
v = sample.get(feat)
if v is not None:
episode_scalars[label] = float(v)
_log_foxglove_scalars("/episode/state", episode_scalars, channels=channels, log_time=log_time)
_log_foxglove_scalars(
"/episode/extras",
extra_scalars,
channels=channels,
log_time=log_time,
)
lock = threading.Lock()
stop_event = threading.Event()
+226
View File
@@ -0,0 +1,226 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for the camera-view curation pipeline (stubbed VLM, mocked Hub)."""
from __future__ import annotations
from pathlib import Path
from unittest.mock import MagicMock, patch
import PIL.Image
import pytest
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
pytest.importorskip("pandas", reason="pandas is required (install lerobot[dataset])")
import pandas as pd # noqa: E402
from lerobot.annotations.camera_curation.config import CameraCurationConfig # noqa: E402
from lerobot.annotations.camera_curation.curator import ( # noqa: E402
CameraVerdict,
build_name_mapping,
curate_cameras,
is_valid_view_label,
rename_camera_keys_on_hub,
write_report,
)
from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient # noqa: E402
from lerobot.datasets.io_utils import load_info, write_info # noqa: E402
from lerobot.datasets.utils import DatasetInfo # noqa: E402
from lerobot.utils.io_utils import load_json, write_json # noqa: E402
VOCAB = ("top", "wrist", "front", "bottom", "left", "right")
def _queued_vlm(responses: list) -> StubVlmClient:
"""Stub VLM that returns queued responses in batch order."""
state = {"i": 0}
def responder(_messages):
r = responses[state["i"]]
state["i"] += 1
return r
return StubVlmClient(responder=responder)
def _tiny_image() -> PIL.Image.Image:
return PIL.Image.new("RGB", (16, 12))
def _make_min_meta(root: Path, camera_key: str, dtype: str = "video") -> None:
"""Write a minimal ``meta/`` tree with one camera + one action feature."""
(root / "meta" / "episodes" / "chunk-000").mkdir(parents=True, exist_ok=True)
features = {
camera_key: {
"dtype": dtype,
"shape": (64, 96, 3),
"names": ["height", "width", "channels"],
"info": {"video.fps": 10.0} if dtype == "video" else None,
},
"action": {"dtype": "float32", "shape": (2,), "names": None},
}
write_info(DatasetInfo(codebase_version="v3.0", fps=10, features=features), root)
write_json({camera_key: {"mean": [0.0]}, "action": {"mean": [0.0]}}, root / "meta" / "stats.json")
df = pd.DataFrame(
{
"episode_index": [0],
f"videos/{camera_key}/from_timestamp": [0.0],
f"videos/{camera_key}/to_timestamp": [1.0],
f"videos/{camera_key}/chunk_index": [0],
f"videos/{camera_key}/file_index": [0],
f"stats/{camera_key}/mean": [[0.0]],
}
)
df.to_parquet(root / "meta" / "episodes" / "chunk-000" / "file-000.parquet")
# ------------------------------ pure logic ------------------------------
def test_is_valid_view_label():
assert is_valid_view_label("top", VOCAB, allow_combos=True)
assert is_valid_view_label("left_wrist", VOCAB, allow_combos=True)
assert not is_valid_view_label("left_wrist", VOCAB, allow_combos=False)
assert not is_valid_view_label("banana", VOCAB, allow_combos=True)
assert not is_valid_view_label("left_left", VOCAB, allow_combos=True) # duplicate token
assert not is_valid_view_label("top_wrist_front", VOCAB, allow_combos=True) # 3 tokens
assert not is_valid_view_label("", VOCAB, allow_combos=True)
def test_curate_cameras_parses_and_validates(tmp_path):
cfg = CameraCurationConfig(view_vocabulary=VOCAB)
frames = {
"observation.images.a": [_tiny_image()],
"observation.images.b": [_tiny_image()],
"observation.images.c": [], # no frames -> reported, not sent to the VLM
}
vlm = _queued_vlm(
[
{"usable": True, "blur_reason": None, "view_label": "Left Wrist", "confidence": 0.9},
{"usable": False, "blur_reason": "out of focus", "view_label": "banana", "confidence": 0.2},
]
)
verdicts = {v.camera_key: v for v in curate_cameras(frames, cfg, vlm)}
assert verdicts["observation.images.a"].view_label == "left_wrist" # normalized
assert verdicts["observation.images.a"].usable is True
assert verdicts["observation.images.b"].usable is False
assert verdicts["observation.images.b"].blur_reason == "out of focus"
assert verdicts["observation.images.b"].view_label is None # invalid label dropped
assert verdicts["observation.images.c"].view_label is None # no frames
def test_build_name_mapping_and_collision():
cfg = CameraCurationConfig(view_vocabulary=VOCAB)
existing = {"observation.images.cam_0": {}, "observation.images.cam_1": {}, "observation.images.top": {}}
verdicts = [
CameraVerdict("observation.images.cam_0", usable=True, view_label="left_wrist"),
CameraVerdict("observation.images.cam_1", usable=True, view_label="front"),
# already canonical -> skipped by build_name_mapping
CameraVerdict("observation.images.top", usable=True, view_label="top"),
]
mapping = build_name_mapping(verdicts, existing, cfg)
assert mapping == {
"observation.images.cam_0": "observation.images.left_wrist",
"observation.images.cam_1": "observation.images.front",
}
# proposed_new_key stamped back onto the verdicts
assert verdicts[0].proposed_new_key == "observation.images.left_wrist"
# two cameras wanting the same label collide under the default policy
clash = [
CameraVerdict("observation.images.cam_0", usable=True, view_label="top"),
CameraVerdict("observation.images.cam_1", usable=True, view_label="top"),
]
with pytest.raises(ValueError, match="collision"):
build_name_mapping(clash, {"observation.images.cam_0": {}, "observation.images.cam_1": {}}, cfg)
def test_write_report(tmp_path):
_make_min_meta(tmp_path, "observation.images.cam_0", dtype="video")
cfg = CameraCurationConfig(repo_id="user/ds", view_vocabulary=VOCAB)
verdicts = [
CameraVerdict("observation.images.cam_0", usable=True, view_label="left_wrist", confidence=0.9)
]
mapping = {"observation.images.cam_0": "observation.images.left_wrist"}
report_path = write_report(tmp_path, verdicts, mapping, cfg)
report = load_json(report_path)
cam = report["cameras"]["observation.images.cam_0"]
assert cam["view_label"] == "left_wrist"
assert cam["proposed_new_key"] == "observation.images.left_wrist"
# verdict stamped into info.json so it travels with the dataset
info = load_info(tmp_path)
assert info.features["observation.images.cam_0"]["info"]["curation"]["view_label"] == "left_wrist"
# ------------------------- lightweight Hub rename -------------------------
def test_rename_camera_keys_on_hub_builds_ops(tmp_path):
from huggingface_hub import CommitOperationAdd, CommitOperationCopy, CommitOperationDelete
camera_key = "observation.images.cam_0"
new_key = "observation.images.left_wrist"
_make_min_meta(tmp_path, camera_key, dtype="video")
old_mp4 = f"videos/{camera_key}/chunk-000/file-000.mp4"
fake_api = MagicMock()
fake_api.list_repo_files.return_value = [old_mp4, "meta/info.json", "data/chunk-000/file-000.parquet"]
fake_api.create_commit.return_value = MagicMock(oid="deadbeef")
with patch("huggingface_hub.HfApi", return_value=fake_api):
rename_camera_keys_on_hub("user/ds", {camera_key: new_key}, tmp_path, branch="curated")
kwargs = fake_api.create_commit.call_args.kwargs
ops = kwargs["operations"]
copies = [o for o in ops if isinstance(o, CommitOperationCopy)]
deletes = [o for o in ops if isinstance(o, CommitOperationDelete)]
adds = [o for o in ops if isinstance(o, CommitOperationAdd)]
new_mp4 = f"videos/{new_key}/chunk-000/file-000.mp4"
assert any(o.src_path_in_repo == old_mp4 and o.path_in_repo == new_mp4 for o in copies)
assert any(o.path_in_repo == old_mp4 for o in deletes)
assert any(o.path_in_repo == "meta/info.json" for o in adds)
assert kwargs["revision"] == "curated"
# meta on disk was actually remapped
info = load_info(tmp_path)
assert new_key in info.features and camera_key not in info.features
def test_rename_camera_keys_on_hub_rejects_image_keys(tmp_path):
camera_key = "observation.images.cam_0"
_make_min_meta(tmp_path, camera_key, dtype="image")
with patch("huggingface_hub.HfApi", return_value=MagicMock()):
with pytest.raises(NotImplementedError, match="image data"):
rename_camera_keys_on_hub("user/ds", {camera_key: "observation.images.top"}, tmp_path)
def test_rename_camera_keys_on_hub_rejects_swaps(tmp_path):
_make_min_meta(tmp_path, "observation.images.a", dtype="video")
with patch("huggingface_hub.HfApi", return_value=MagicMock()):
with pytest.raises(NotImplementedError, match="swap"):
rename_camera_keys_on_hub(
"user/ds",
{
"observation.images.a": "observation.images.b",
"observation.images.b": "observation.images.a",
},
tmp_path,
)
+163 -1
View File
@@ -23,6 +23,7 @@ import torch
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
import pandas as pd # noqa: E402
from lerobot.configs import DepthEncoderConfig, RGBEncoderConfig
from lerobot.datasets.dataset_tools import (
@@ -34,9 +35,11 @@ from lerobot.datasets.dataset_tools import (
modify_tasks,
reencode_dataset,
remove_feature,
rename_features,
split_dataset,
)
from lerobot.datasets.io_utils import load_info
from lerobot.datasets.dataset_tools import _resolve_rename_collisions
from lerobot.datasets.io_utils import load_info, load_stats
from tests.datasets.test_video_encoding import require_h264, require_hevc, require_libsvtav1
from tests.fixtures.constants import DUMMY_DEPTH_FEATURES, DUMMY_DEPTH_KEY
from tests.fixtures.dataset_factories import add_frames
@@ -1492,3 +1495,162 @@ def test_reencode_dataset_multi_key_multiprocessing(
for vk in dataset.meta.video_keys:
persisted_encoder = RGBEncoderConfig.from_video_info(persisted_info.features[vk].get("info", {}))
assert persisted_encoder == target_cfg
# ----------------------------- rename_features -----------------------------
def _mock_hub(tmp_path):
"""Context managers that stop dataset reload from hitting the Hub."""
return (
patch("lerobot.datasets.dataset_metadata.get_safe_version", return_value="v3.0"),
patch(
"lerobot.datasets.dataset_metadata.snapshot_download",
side_effect=lambda repo_id, **kwargs: str(kwargs.get("local_dir", tmp_path)),
),
)
@pytest.fixture
def two_camera_image_dataset(tmp_path, empty_lerobot_dataset_factory):
"""An image dataset with two camera views (for collision tests)."""
features = {
"action": {"dtype": "float32", "shape": (6,), "names": None},
"observation.images.cam_0": {"dtype": "image", "shape": (32, 32, 3), "names": None},
"observation.images.cam_1": {"dtype": "image", "shape": (32, 32, 3), "names": None},
}
dataset = empty_lerobot_dataset_factory(root=tmp_path / "two_cam", features=features)
for _ in range(2):
for _ in range(4):
dataset.add_frame(
{
"action": np.random.randn(6).astype(np.float32),
"observation.images.cam_0": np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8),
"observation.images.cam_1": np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8),
"task": "t",
}
)
dataset.save_episode()
dataset.finalize()
return dataset
def test_resolve_rename_collisions_error_and_suffix():
features = {"a": {}, "b": {}, "c": {}}
# many-to-one
with pytest.raises(ValueError, match="same target"):
_resolve_rename_collisions({"a": "top", "b": "top"}, features, "error")
# target collides with an untouched key
with pytest.raises(ValueError, match="existing feature"):
_resolve_rename_collisions({"a": "c"}, features, "error")
# suffix disambiguates deterministically
resolved = _resolve_rename_collisions({"a": "top", "b": "top"}, features, "suffix")
assert set(resolved.values()) == {"top", "top_2"}
assert resolved["a"] == "top" # sorted-source order keeps the first
def test_rename_image_feature(sample_dataset, tmp_path):
old, new = "observation.images.top", "observation.images.wrist"
m1, m2 = _mock_hub(tmp_path)
with m1, m2:
renamed = rename_features(sample_dataset, {old: new}, output_dir=tmp_path / "renamed")
assert new in renamed.meta.features
assert old not in renamed.meta.features
assert renamed.meta.features[new]["dtype"] == "image"
# the frame still decodes under the new key
item = renamed[0]
assert new in item and old not in item
# stats moved to the new key
stats = load_stats(renamed.root)
assert new in stats and old not in stats
@require_h264
def test_rename_video_feature_no_reencode(tmp_path, empty_lerobot_dataset_factory, features_factory):
features = features_factory(use_videos=True) # observation.images.{laptop,phone}
dataset = empty_lerobot_dataset_factory(root=tmp_path / "vid", features=features, use_videos=True)
add_frames(dataset, num_frames=4)
dataset.save_episode()
dataset.finalize()
old, new = "laptop", "observation.images.top" # features_factory uses bare camera keys
old_bytes = (dataset.root / dataset.meta.get_video_file_path(0, old)).read_bytes()
m1, m2 = _mock_hub(tmp_path)
with m1, m2:
renamed = rename_features(dataset, {old: new}, output_dir=tmp_path / "renamed")
assert new in renamed.meta.features and old not in renamed.meta.features
new_mp4 = renamed.root / renamed.meta.get_video_file_path(0, new)
assert new_mp4.exists()
# a rename must not re-encode: the mp4 is byte-identical.
assert new_mp4.read_bytes() == old_bytes
# episodes metadata columns were remapped.
ep_parquet = next((renamed.root / "meta" / "episodes").glob("*/*.parquet"))
cols = pd.read_parquet(ep_parquet).columns
assert f"videos/{new}/from_timestamp" in cols
assert f"videos/{old}/from_timestamp" not in cols
# the video still decodes under the new key.
assert new in renamed[0]
def test_rename_collision_raises(two_camera_image_dataset, tmp_path):
m1, m2 = _mock_hub(tmp_path)
with m1, m2, pytest.raises(ValueError, match="collision"):
rename_features(
two_camera_image_dataset,
{
"observation.images.cam_0": "observation.images.top",
"observation.images.cam_1": "observation.images.top",
},
output_dir=tmp_path / "out",
)
def test_rename_collision_suffix(two_camera_image_dataset, tmp_path):
m1, m2 = _mock_hub(tmp_path)
with m1, m2:
renamed = rename_features(
two_camera_image_dataset,
{
"observation.images.cam_0": "observation.images.top",
"observation.images.cam_1": "observation.images.top",
},
output_dir=tmp_path / "out",
on_collision="suffix",
)
keys = set(renamed.meta.features)
assert {"observation.images.top", "observation.images.top_2"} <= keys
def test_rename_identity_only_raises(sample_dataset, tmp_path):
with pytest.raises(ValueError, match="identity"):
rename_features(
sample_dataset,
{"observation.images.top": "observation.images.top"},
output_dir=tmp_path / "out",
)
def test_rename_missing_key_raises(sample_dataset, tmp_path):
with pytest.raises(ValueError, match="not found"):
rename_features(
sample_dataset,
{"observation.images.nope": "observation.images.top"},
output_dir=tmp_path / "out",
)
def test_rename_required_feature_raises(sample_dataset, tmp_path):
with pytest.raises(ValueError, match="required"):
rename_features(sample_dataset, {"timestamp": "t2"}, output_dir=tmp_path / "out")
def test_rename_slash_in_target_raises(sample_dataset, tmp_path):
with pytest.raises(ValueError, match="'/'"):
rename_features(
sample_dataset,
{"observation.images.top": "observation/images/top"},
output_dir=tmp_path / "out",
)
-157
View File
@@ -13,78 +13,11 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import sys
from types import SimpleNamespace
import numpy as np
import pytest
import torch
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
from lerobot.scripts.lerobot_dataset_viz import visualize_dataset
from lerobot.utils import import_utils
from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS, TRUNCATED
from lerobot.utils.dataset_visualization_utils import (
get_extra_scalar_keys,
is_scalar_feature,
is_scalar_like,
scalar_to_float,
)
class DummyMeta:
camera_keys = ["observation.images.front"]
class DummyFeatureDataset:
meta = DummyMeta()
features = {
"index": {"dtype": "int64", "shape": [1]},
"timestamp": {"dtype": "float32", "shape": [1]},
"episode_index": {"dtype": "int64", "shape": [1]},
"frame_index": {"dtype": "int64", "shape": [1]},
"task_index": {"dtype": "int64", "shape": [1]},
ACTION: {"dtype": "float32", "shape": [6]},
OBS_STATE: {"dtype": "float32", "shape": [6]},
DONE: {"dtype": "bool", "shape": [1]},
REWARD: {"dtype": "float32", "shape": [1]},
SUCCESS: {"dtype": "bool", "shape": [1]},
TRUNCATED: {"dtype": "bool", "shape": [1]},
"observation.images.front": {"dtype": "video", "shape": [3, 480, 640]},
"q_target": {"dtype": "float32", "shape": [1]},
"quality": {"dtype": "double", "shape": [1]},
"intervention": {"dtype": "bool", "shape": []},
"embedding": {"dtype": "float32", "shape": [32]},
"comment": {"dtype": "string", "shape": [1]},
}
class DummyVizDataset:
repo_id = "dummy/custom-scalars"
depth_output_unit = "m"
meta = SimpleNamespace(camera_keys=[], depth_keys=[], stats=None)
features = {
"index": {"dtype": "int64", "shape": [1]},
"timestamp": {"dtype": "float32", "shape": [1]},
ACTION: {"dtype": "float32", "shape": [2], "names": ["x", "y"]},
"q_target": {"dtype": "float32", "shape": [1]},
"intervention": {"dtype": "bool", "shape": [1]},
"embedding": {"dtype": "float32", "shape": [2]},
}
def __len__(self):
return 2
def __getitem__(self, index):
return {
"index": torch.tensor(index),
"timestamp": torch.tensor(index / 10),
ACTION: torch.tensor([index, index + 1], dtype=torch.float32),
"q_target": torch.tensor([index + 0.5]),
"intervention": torch.tensor([index == 0]),
"embedding": torch.tensor([index, index + 1], dtype=torch.float32),
}
@pytest.mark.skip("TODO: add dummy videos")
@@ -100,93 +33,3 @@ def test_visualize_local_dataset(tmp_path, lerobot_dataset_factory):
output_dir=output_dir,
)
assert rrd_path.exists()
def test_get_extra_scalar_keys_skips_known_metadata_and_non_scalars():
dataset = DummyFeatureDataset()
assert get_extra_scalar_keys(dataset) == [TRUNCATED, "q_target", "quality", "intervention"]
assert get_extra_scalar_keys(dataset, additional_known_keys=(TRUNCATED,)) == [
"q_target",
"quality",
"intervention",
]
@pytest.mark.parametrize(
("feature", "expected"),
[
({"dtype": "float32", "shape": (1,)}, True),
({"dtype": "double", "shape": [1]}, True),
({"dtype": "bool", "shape": []}, True),
({"dtype": "int64", "shape": 1}, True),
({"dtype": "float32", "shape": (3,)}, False),
({"dtype": "complex64", "shape": [1]}, False),
({"dtype": "datetime64[ns]", "shape": [1]}, False),
({"dtype": "string", "shape": [1]}, False),
({"shape": [1]}, False),
],
)
def test_is_scalar_feature(feature, expected):
assert is_scalar_feature(feature) is expected
@pytest.mark.parametrize(
("value", "expected"),
[
(torch.tensor(1.5), True),
(torch.tensor([1.5]), True),
(torch.tensor([1.5, 2.5]), False),
(np.array(2.0), True),
(np.array([2.0]), True),
(np.array([2.0, 3.0]), False),
(True, True),
("not numeric", False),
],
)
def test_is_scalar_like(value, expected):
assert is_scalar_like(value) is expected
def test_scalar_to_float():
assert scalar_to_float(torch.tensor(3.0)) == 3.0
assert scalar_to_float(np.array([4.0])) == 4.0
def test_visualize_dataset_logs_extra_scalars(monkeypatch):
logged = []
initialized = []
dummy_rrb = SimpleNamespace(
Spatial2DView=lambda origin=None, name=None: SimpleNamespace(
kind="spatial", origin=origin, name=name
),
TimeSeriesView=lambda origin=None, name=None, overrides=None: SimpleNamespace(
kind="time_series", origin=origin, name=name, overrides=overrides
),
Grid=lambda *views: SimpleNamespace(views=views),
Blueprint=lambda root: SimpleNamespace(root=root),
)
dummy_rr = SimpleNamespace(
SeriesLines=lambda names=None: SimpleNamespace(names=names),
Scalars=lambda value: SimpleNamespace(value=value),
init=lambda *args, **kwargs: initialized.append((args, kwargs)),
log=lambda key, entity: logged.append((key, entity)),
set_time=lambda *args, **kwargs: None,
blueprint=dummy_rrb,
)
monkeypatch.setitem(sys.modules, "rerun", dummy_rr)
monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb)
monkeypatch.setattr(import_utils, "require_package", lambda *args, **kwargs: None)
visualize_dataset(DummyVizDataset(), episode_index=0, batch_size=2)
logged_keys = [key for key, _ in logged]
assert logged_keys.count("q_target") == 2
assert logged_keys.count("intervention") == 2
assert "embedding" not in logged_keys
blueprint = initialized[0][1]["default_blueprint"]
time_series_origins = {view.origin for view in blueprint.root.views if view.kind == "time_series"}
assert {"q_target", "intervention"} <= time_series_origins
assert "embedding" not in time_series_origins
@@ -0,0 +1,52 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import pytest
import torch
# The script imports ``lerobot.datasets`` (via the annotation frame provider),
# which only ships under the ``dataset`` extra.
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
from lerobot.scripts.lerobot_curate_cameras import _to_uint8_frame, _uniform_indices # noqa: E402
@pytest.mark.parametrize(
"n,k,expected",
[
(0, 4, []),
(5, 0, []),
(3, 5, [0, 1, 2]), # k >= n -> all frames
(1, 4, [0]),
(10, 1, [0]),
(10, 4, [0, 3, 6, 9]), # evenly spaced, endpoints included
],
)
def test_uniform_indices(n, k, expected):
assert _uniform_indices(n, k) == expected
def test_to_uint8_frame_scales_floats():
frame = torch.ones(3, 4, 4, dtype=torch.float32) # [0,1] float
out = _to_uint8_frame(frame)
assert out.dtype == torch.uint8
assert int(out.max()) == 255
def test_to_uint8_frame_passthrough_uint8():
frame = torch.zeros(3, 4, 4, dtype=torch.uint8)
out = _to_uint8_frame(frame)
assert out is frame # uint8 passes through untouched
+1 -19
View File
@@ -24,7 +24,7 @@ the functions that talk to the server, so the helpers below run in the base test
import numpy as np
from lerobot.utils import foxglove_visualization as fv
from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
from lerobot.utils.constants import ACTION, OBS_STATE
def test_foxglove_safe_name_collapses_dots():
@@ -93,24 +93,6 @@ def test_feature_dim_names_formats():
assert fv._feature_dim_names({"shape": [2]}) is None
def test_dataset_frame_scalar_groups_include_custom_scalars():
sample = {
DONE: np.array(True),
REWARD: np.array(0.5),
SUCCESS: np.array(False),
"q_target": np.array([0.75]),
"intervention": np.array([True]),
"embedding": np.array([1.0, 2.0]),
}
episode_scalars, extra_scalars = fv._dataset_frame_scalar_groups(
sample, ("q_target", "intervention", "embedding")
)
assert episode_scalars == {"done": 1.0, "reward": 0.5, "success": 0.0}
assert extra_scalars == {"q_target": 0.75, "intervention": 1.0}
def test_is_scalar():
assert fv._is_scalar(1.0)
assert fv._is_scalar(np.float32(2.0))