mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-30 13:09:40 +00:00
Compare commits
16 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 51ea892a4f | |||
| a6b06eac38 | |||
| 36b8face98 | |||
| cd8984cc0a | |||
| b9ded9e761 | |||
| 185f3e1708 | |||
| e36783253a | |||
| 289e577fc7 | |||
| 9c32722eb9 | |||
| b49cb50e01 | |||
| dd08d4eb53 | |||
| 6e5f6df6e7 | |||
| 265abe6c79 | |||
| b4e2d0b610 | |||
| 5594eba06a | |||
| 207183c2f8 |
+12
-8
@@ -61,15 +61,19 @@ Full details in [`docs/source/so101.mdx`](./docs/source/so101.mdx) and [`docs/so
|
||||
**4.1 Install**
|
||||
|
||||
```bash
|
||||
pip install 'lerobot[feetech]' # SO-100/SO-101 motor stack
|
||||
# pip install 'lerobot[all]' # everything
|
||||
# pip install 'lerobot[aloha,pusht]' # specific features
|
||||
# pip install 'lerobot[smolvla]' # add SmolVLA deps
|
||||
git lfs install && git lfs pull
|
||||
hf auth login # required to push datasets/policies
|
||||
```
|
||||
# uv (recommended — see AGENTS.md and CLAUDE.md)
|
||||
uv sync --locked --extra feetech # SO-100/SO-101 motor stack
|
||||
# uv sync --locked --extra all # everything
|
||||
# uv sync --locked --extra smolvla # add SmolVLA deps
|
||||
|
||||
Contributors can alternatively use `uv sync --locked --extra feetech` (see `AGENTS.md`).
|
||||
# pip (alternative, e.g. when not working from source)
|
||||
# pip install 'lerobot[feetech]'
|
||||
# pip install 'lerobot[all]'
|
||||
# pip install 'lerobot[smolvla]'
|
||||
|
||||
git lfs install && git lfs pull
|
||||
hf auth login # required to push datasets/policies
|
||||
```
|
||||
|
||||
**4.2 Find USB ports** — run once per arm, unplug when prompted.
|
||||
|
||||
|
||||
@@ -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** —
|
||||
|
||||
@@ -58,7 +58,7 @@ final_action = postprocessor(action)
|
||||
|
||||
## Hardware API redesign
|
||||
|
||||
PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is a overview of what changed and how you can continue to work with datasets created before this pull request.
|
||||
PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is an overview of what changed and how you can continue to work with datasets created before this pull request.
|
||||
|
||||
### What changed?
|
||||
|
||||
@@ -129,8 +129,8 @@ python examples/backward_compatibility/replay.py \
|
||||
|
||||
Policies output actions in the same format as the datasets (`torch.Tensors`). Therefore, the same transformations should be applied.
|
||||
|
||||
To find these transformations, we recommend to first try and and replay an episode of the dataset your policy was trained on using the section above.
|
||||
Then, add these same transformations on your inference script (shown here in the `record.py` script):
|
||||
To find these transformations, we recommend first replaying an episode of the dataset your policy was trained on using the section above.
|
||||
Then, add these same transformations to your inference script (shown here in the `record.py` script):
|
||||
|
||||
```diff
|
||||
action_values = predict_action(
|
||||
|
||||
@@ -88,20 +88,6 @@ policy_preprocessor = NormalizerProcessorStep(stats=dataset_stats)
|
||||
|
||||
The same policy can work with different environment processors, and the same environment processor can work with different policies:
|
||||
|
||||
````python
|
||||
# Use SmolVLA policy with LIBERO environment
|
||||
# Use SmolVLA policy with LIBERO environment
|
||||
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
||||
env_cfg=libero_cfg,
|
||||
policy_cfg=smolvla_cfg,
|
||||
)
|
||||
smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
|
||||
# Or use ACT policy with the same LIBERO environment
|
||||
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
||||
env_cfg=libero_cfg,
|
||||
policy_cfg=act_cfg,
|
||||
)
|
||||
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
|
||||
```python
|
||||
# Use SmolVLA policy with LIBERO environment
|
||||
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
||||
@@ -116,6 +102,7 @@ libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
||||
policy_cfg=act_cfg,
|
||||
)
|
||||
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
|
||||
```
|
||||
|
||||
### 3. **Easier Experimentation**
|
||||
|
||||
@@ -145,7 +132,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
|
||||
state = torch.cat([eef_pos, eef_axisangle, eef_vel,
|
||||
gripper_pos, gripper_vel], dim=-1) # 14D
|
||||
return state
|
||||
````
|
||||
```
|
||||
|
||||
### 4. **Cleaner Environment Code**
|
||||
|
||||
|
||||
@@ -40,10 +40,10 @@ This tutorial guides you through updating the firmware of Feetech motors using t
|
||||
For each motor you want to update:
|
||||
|
||||
1. **Select the motor** from the list by clicking on it
|
||||
2. **Click on Upgrade tab**:
|
||||
3. **Click on Online button**:
|
||||
- If an potential firmware update is found, it will be displayed in the box
|
||||
4. **Click on Upgrade button**:
|
||||
2. **Click the Upgrade tab**:
|
||||
3. **Click the Online button**:
|
||||
- If a potential firmware update is found, it will be displayed in the box
|
||||
4. **Click the Upgrade button**:
|
||||
- The update progress will be displayed
|
||||
|
||||
## Step 6: Verify Update
|
||||
|
||||
@@ -211,7 +211,7 @@ Record, Replay and Train with Hope-JR is still experimental.
|
||||
|
||||
### Record
|
||||
|
||||
This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings).
|
||||
This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data).
|
||||
|
||||
```bash
|
||||
lerobot-record \
|
||||
|
||||
@@ -18,7 +18,7 @@ If you're using Feetech or Dynamixel motors, LeRobot provides built-in bus inter
|
||||
- [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) – for controlling Dynamixel servos
|
||||
|
||||
Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API.
|
||||
For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so101_follower/so101_follower.py)
|
||||
For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so_follower.py)
|
||||
|
||||
Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial):
|
||||
|
||||
|
||||
@@ -51,7 +51,7 @@ In addition to these instructions, you need to install the Feetech SDK & ZeroMQ
|
||||
pip install -e ".[lekiwi]"
|
||||
```
|
||||
|
||||
Great :hugs:! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base :robot:.
|
||||
Great 🤗! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base 🤖.
|
||||
Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.
|
||||
|
||||
# Step-by-Step Assembly Instructions
|
||||
|
||||
@@ -174,7 +174,7 @@ The model takes images, text instructions, and robot state as input, and outputs
|
||||
|
||||
## Reproducing π₀Fast results
|
||||
|
||||
We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40kk steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
|
||||
We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40k steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
|
||||
|
||||
The finetuned model can be found here:
|
||||
|
||||
|
||||
@@ -22,7 +22,7 @@ With processors, you choose the learning features you want to use for your polic
|
||||
## Three pipelines
|
||||
|
||||
We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match.
|
||||
Each of these pipelines handle different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
|
||||
Each of these pipelines handles different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
|
||||
|
||||
1. Pipeline 1: Teleop action space → dataset action space (phone pose → EE targets)
|
||||
2. Pipeline 2: Dataset action space → robot command space (EE targets → joints)
|
||||
@@ -74,15 +74,15 @@ In the phone to SO-100 follower examples we use the following adapters:
|
||||
- `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition.
|
||||
- `transition_to_robot_action`: transforms the pipeline transition to a robot action dict.
|
||||
- `observation_to_transition`: transforms the robot observation dict to a pipeline transition.
|
||||
- `transition_to_observation`: transforms the pipeline transition to a observation dict.
|
||||
- `transition_to_observation`: transforms the pipeline transition to an observation dict.
|
||||
|
||||
Checkout [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
|
||||
Check out [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
|
||||
|
||||
## Dataset feature contracts
|
||||
|
||||
Dataset features are determined by the keys saved in the dataset. Each step can declare what features it modifies in a contract called `transform_features(...)`. Once you build a processor, the processor can then aggregate all of these features with `aggregate_pipeline_dataset_features()` and merge multiple feature dicts with `combine_feature_dicts(...)`.
|
||||
|
||||
Below is and example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
|
||||
Below is an example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
|
||||
|
||||
```python
|
||||
def transform_features(
|
||||
|
||||
+2
-2
@@ -57,7 +57,7 @@ policy_cfg.rtc_config = RTCConfig(
|
||||
policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda")
|
||||
|
||||
# Now use predict_action_chunk with RTC parameters
|
||||
inference_delay = 4 # How many steps of inference latency, this values should be calculated based on the inference latency of the policy
|
||||
inference_delay = 4 # How many steps of inference latency, this value should be calculated based on the inference latency of the policy
|
||||
|
||||
# Initialize the action queue
|
||||
action_queue = ActionQueue(policy_cfg.rtc_config)
|
||||
@@ -100,7 +100,7 @@ Typical values: 8-12 steps
|
||||
RTCConfig(execution_horizon=10)
|
||||
```
|
||||
|
||||
**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is a optimal value.
|
||||
**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is an optimal value.
|
||||
|
||||
**`prefix_attention_schedule`**: How to weight consistency across the overlap region.
|
||||
|
||||
|
||||
@@ -93,7 +93,7 @@ lerobot-train --help
|
||||
|
||||
## Evaluate the finetuned model and run it in real-time
|
||||
|
||||
Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots).
|
||||
Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots#record-a-dataset).
|
||||
Once you are logged in, you can run inference in your setup by doing:
|
||||
|
||||
```bash
|
||||
|
||||
@@ -50,11 +50,11 @@ lerobot-edit-dataset \
|
||||
Divide a dataset into multiple subsets.
|
||||
|
||||
```bash
|
||||
# Split by fractions (e.g. 80% train, 20% test, 20% val)
|
||||
# Split by fractions (e.g. 60% train, 20% val, 20% test)
|
||||
lerobot-edit-dataset \
|
||||
--repo_id lerobot/pusht \
|
||||
--operation.type split \
|
||||
--operation.splits '{"train": 0.8, "test": 0.2, "val": 0.2}'
|
||||
--operation.splits '{"train": 0.6, "val": 0.2, "test": 0.2}'
|
||||
|
||||
# Split by specific episode indices
|
||||
lerobot-edit-dataset \
|
||||
@@ -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.
|
||||
|
||||
@@ -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>
|
||||
}}
|
||||
@@ -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",
|
||||
|
||||
@@ -19,6 +19,7 @@ import copy
|
||||
import logging
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Any, NotRequired, TypedDict
|
||||
|
||||
import datasets
|
||||
import pandas as pd
|
||||
@@ -49,8 +50,32 @@ from .utils import (
|
||||
)
|
||||
from .video_utils import concatenate_video_files, get_video_duration_in_s
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> dict[str, dict]:
|
||||
type FeatureDict = dict[str, dict[str, Any]]
|
||||
type ChunkFile = tuple[int, int]
|
||||
|
||||
|
||||
class IndexState(TypedDict):
|
||||
chunk: int
|
||||
file: int
|
||||
src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
|
||||
|
||||
|
||||
class VideoIndex(TypedDict):
|
||||
chunk: int
|
||||
file: int
|
||||
latest_duration: float
|
||||
episode_duration: float
|
||||
src_to_offset: NotRequired[dict[ChunkFile, float]]
|
||||
src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
|
||||
dst_file_durations: NotRequired[dict[ChunkFile, float]]
|
||||
|
||||
|
||||
type VideoIndexState = dict[str, VideoIndex]
|
||||
|
||||
|
||||
def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> FeatureDict:
|
||||
"""Create a merged video feature info dictionary for aggregation. The video encoder info is merged field-by-field: each key is kept only when every source agrees; otherwise that key is set to ``null`` (or ``{}`` for ``video.extra_options``) and a warning is logged.
|
||||
|
||||
Args:
|
||||
@@ -59,14 +84,14 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
|
||||
Returns:
|
||||
dict: A dictionary of merged video feature info.
|
||||
"""
|
||||
merged_info = copy.deepcopy(all_metadata[0].features)
|
||||
merged_info: FeatureDict = copy.deepcopy(all_metadata[0].features)
|
||||
video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"]
|
||||
|
||||
for vk in video_keys:
|
||||
video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata]
|
||||
base_video_info = video_infos[0]
|
||||
|
||||
merged_encoder_info: dict = {}
|
||||
merged_encoder_info: dict[str, Any] = {}
|
||||
fallback_keys: list[str] = []
|
||||
for info_key in VIDEO_ENCODER_INFO_KEYS:
|
||||
values = [info.get(info_key, None) for info in video_infos]
|
||||
@@ -80,7 +105,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
|
||||
merged_encoder_info[info_key] = {} if info_key == "video.extra_options" else None
|
||||
|
||||
if fallback_keys:
|
||||
logging.warning(
|
||||
logger.warning(
|
||||
f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. "
|
||||
f"Setting these keys to null: {fallback_keys}.",
|
||||
)
|
||||
@@ -92,7 +117,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
|
||||
return merged_info
|
||||
|
||||
|
||||
def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
|
||||
def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[int, str | None, FeatureDict]:
|
||||
"""Validates that all dataset metadata have consistent properties.
|
||||
|
||||
Ensures all datasets have the same fps, robot_type, and features to guarantee
|
||||
@@ -129,7 +154,9 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
|
||||
return fps, robot_type, features
|
||||
|
||||
|
||||
def update_data_df(df, src_meta, dst_meta):
|
||||
def update_data_df(
|
||||
df: pd.DataFrame, src_meta: LeRobotDatasetMetadata, dst_meta: LeRobotDatasetMetadata
|
||||
) -> pd.DataFrame:
|
||||
"""Updates a data DataFrame with new indices and task mappings for aggregation.
|
||||
|
||||
Adjusts episode indices, frame indices, and task indices to account for
|
||||
@@ -154,12 +181,12 @@ def update_data_df(df, src_meta, dst_meta):
|
||||
|
||||
|
||||
def update_meta_data(
|
||||
df,
|
||||
dst_meta,
|
||||
meta_idx,
|
||||
data_idx,
|
||||
videos_idx,
|
||||
):
|
||||
df: pd.DataFrame,
|
||||
dst_meta: LeRobotDatasetMetadata,
|
||||
meta_idx: IndexState,
|
||||
data_idx: IndexState,
|
||||
videos_idx: VideoIndexState,
|
||||
) -> pd.DataFrame:
|
||||
"""Updates metadata DataFrame with new chunk, file, and timestamp indices.
|
||||
|
||||
Adjusts all indices and timestamps to account for previously aggregated
|
||||
@@ -289,7 +316,7 @@ def aggregate_datasets(
|
||||
chunk_size: int | None = None,
|
||||
concatenate_videos: bool = True,
|
||||
concatenate_data: bool = True,
|
||||
):
|
||||
) -> None:
|
||||
"""Aggregates multiple LeRobot datasets into a single unified dataset.
|
||||
|
||||
This is the main function that orchestrates the aggregation process by:
|
||||
@@ -309,7 +336,7 @@ def aggregate_datasets(
|
||||
concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
|
||||
concatenate_data: When False, keep one parquet per source file instead of packing into shards.
|
||||
"""
|
||||
logging.info("Start aggregate_datasets")
|
||||
logger.info("Start aggregate_datasets")
|
||||
|
||||
if data_files_size_in_mb is None:
|
||||
data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
|
||||
@@ -341,15 +368,15 @@ def aggregate_datasets(
|
||||
video_files_size_in_mb=video_files_size_in_mb,
|
||||
)
|
||||
|
||||
logging.info("Find all tasks")
|
||||
logger.info("Find all tasks")
|
||||
unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
|
||||
dst_meta.tasks = pd.DataFrame(
|
||||
{"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task")
|
||||
)
|
||||
|
||||
meta_idx = {"chunk": 0, "file": 0}
|
||||
data_idx = {"chunk": 0, "file": 0}
|
||||
videos_idx = {
|
||||
meta_idx: IndexState = {"chunk": 0, "file": 0}
|
||||
data_idx: IndexState = {"chunk": 0, "file": 0}
|
||||
videos_idx: VideoIndexState = {
|
||||
key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys
|
||||
}
|
||||
|
||||
@@ -373,12 +400,17 @@ def aggregate_datasets(
|
||||
dst_meta.info.total_frames += src_meta.total_frames
|
||||
|
||||
finalize_aggregation(dst_meta, all_metadata)
|
||||
logging.info("Aggregation complete.")
|
||||
logger.info("Aggregation complete.")
|
||||
|
||||
|
||||
def aggregate_videos(
|
||||
src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size, concatenate_videos=True
|
||||
):
|
||||
src_meta: LeRobotDatasetMetadata,
|
||||
dst_meta: LeRobotDatasetMetadata,
|
||||
videos_idx: VideoIndexState,
|
||||
video_files_size_in_mb: float,
|
||||
chunk_size: int,
|
||||
concatenate_videos: bool = True,
|
||||
) -> VideoIndexState:
|
||||
"""Aggregates video chunks from a source dataset into the destination dataset.
|
||||
|
||||
Handles video file concatenation and rotation based on file size limits.
|
||||
@@ -406,15 +438,16 @@ def aggregate_videos(
|
||||
videos_idx[key]["dst_file_durations"] = {}
|
||||
|
||||
for key, video_idx in videos_idx.items():
|
||||
unique_chunk_file_pairs = {
|
||||
(chunk, file)
|
||||
for chunk, file in zip(
|
||||
src_meta.episodes[f"videos/{key}/chunk_index"],
|
||||
src_meta.episodes[f"videos/{key}/file_index"],
|
||||
strict=False,
|
||||
)
|
||||
}
|
||||
unique_chunk_file_pairs = sorted(unique_chunk_file_pairs)
|
||||
unique_chunk_file_pairs: list[ChunkFile] = sorted(
|
||||
{
|
||||
(chunk, file)
|
||||
for chunk, file in zip(
|
||||
src_meta.episodes[f"videos/{key}/chunk_index"],
|
||||
src_meta.episodes[f"videos/{key}/file_index"],
|
||||
strict=False,
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
chunk_idx = video_idx["chunk"]
|
||||
file_idx = video_idx["file"]
|
||||
@@ -489,7 +522,14 @@ def aggregate_videos(
|
||||
return videos_idx
|
||||
|
||||
|
||||
def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size, concatenate_data=True):
|
||||
def aggregate_data(
|
||||
src_meta: LeRobotDatasetMetadata,
|
||||
dst_meta: LeRobotDatasetMetadata,
|
||||
data_idx: IndexState,
|
||||
data_files_size_in_mb: float,
|
||||
chunk_size: int,
|
||||
concatenate_data: bool = True,
|
||||
) -> IndexState:
|
||||
"""Aggregates data chunks from a source dataset into the destination dataset.
|
||||
|
||||
Reads source data files, updates indices to match the aggregated dataset,
|
||||
@@ -510,14 +550,16 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
|
||||
Returns:
|
||||
dict: Updated data_idx with current chunk and file indices.
|
||||
"""
|
||||
unique_chunk_file_ids = {
|
||||
(c, f)
|
||||
for c, f in zip(
|
||||
src_meta.episodes["data/chunk_index"], src_meta.episodes["data/file_index"], strict=False
|
||||
)
|
||||
}
|
||||
|
||||
unique_chunk_file_ids = sorted(unique_chunk_file_ids)
|
||||
unique_chunk_file_ids: list[ChunkFile] = sorted(
|
||||
{
|
||||
(c, f)
|
||||
for c, f in zip(
|
||||
src_meta.episodes["data/chunk_index"],
|
||||
src_meta.episodes["data/file_index"],
|
||||
strict=False,
|
||||
)
|
||||
}
|
||||
)
|
||||
contains_images = len(dst_meta.image_keys) > 0
|
||||
|
||||
# retrieve features schema for proper image typing in parquet
|
||||
@@ -525,7 +567,7 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
|
||||
|
||||
# Track source to destination file mapping for metadata update
|
||||
# This is critical for handling datasets that are already results of a merge
|
||||
src_to_dst: dict[tuple[int, int], tuple[int, int]] = {}
|
||||
src_to_dst: dict[ChunkFile, ChunkFile] = {}
|
||||
|
||||
for src_chunk_idx, src_file_idx in unique_chunk_file_ids:
|
||||
src_path = src_meta.root / DEFAULT_DATA_PATH.format(
|
||||
@@ -564,7 +606,13 @@ def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_si
|
||||
return data_idx
|
||||
|
||||
|
||||
def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
|
||||
def aggregate_metadata(
|
||||
src_meta: LeRobotDatasetMetadata,
|
||||
dst_meta: LeRobotDatasetMetadata,
|
||||
meta_idx: IndexState,
|
||||
data_idx: IndexState,
|
||||
videos_idx: VideoIndexState,
|
||||
) -> IndexState:
|
||||
"""Aggregates metadata from a source dataset into the destination dataset.
|
||||
|
||||
Reads source metadata files, updates all indices and timestamps,
|
||||
@@ -580,16 +628,16 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
|
||||
Returns:
|
||||
dict: Updated meta_idx with current chunk and file indices.
|
||||
"""
|
||||
chunk_file_ids = {
|
||||
(c, f)
|
||||
for c, f in zip(
|
||||
src_meta.episodes["meta/episodes/chunk_index"],
|
||||
src_meta.episodes["meta/episodes/file_index"],
|
||||
strict=False,
|
||||
)
|
||||
}
|
||||
|
||||
chunk_file_ids = sorted(chunk_file_ids)
|
||||
chunk_file_ids: list[ChunkFile] = sorted(
|
||||
{
|
||||
(c, f)
|
||||
for c, f in zip(
|
||||
src_meta.episodes["meta/episodes/chunk_index"],
|
||||
src_meta.episodes["meta/episodes/file_index"],
|
||||
strict=False,
|
||||
)
|
||||
}
|
||||
)
|
||||
for chunk_idx, file_idx in chunk_file_ids:
|
||||
src_path = src_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
|
||||
df = pd.read_parquet(src_path)
|
||||
@@ -622,16 +670,16 @@ def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
|
||||
def append_or_create_parquet_file(
|
||||
df: pd.DataFrame,
|
||||
src_path: Path,
|
||||
idx: dict[str, int],
|
||||
idx: IndexState,
|
||||
max_mb: float,
|
||||
chunk_size: int,
|
||||
default_path: str,
|
||||
contains_images: bool = False,
|
||||
aggr_root: Path = None,
|
||||
aggr_root: Path | None = None,
|
||||
hf_features: datasets.Features | None = None,
|
||||
concatenate: bool = True,
|
||||
one_row_group_per_episode: bool = False,
|
||||
) -> tuple[dict[str, int], tuple[int, int]]:
|
||||
) -> tuple[IndexState, ChunkFile]:
|
||||
"""Appends data to an existing parquet file or creates a new one based on size constraints.
|
||||
|
||||
Manages file rotation when size limits are exceeded to prevent individual files
|
||||
@@ -654,7 +702,13 @@ def append_or_create_parquet_file(
|
||||
Returns:
|
||||
tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict
|
||||
and (dst_chunk, dst_file) is the actual destination file the data was written to.
|
||||
|
||||
Raises:
|
||||
ValueError: If aggr_root is not provided.
|
||||
"""
|
||||
if aggr_root is None:
|
||||
raise ValueError("aggr_root must be provided.")
|
||||
|
||||
dst_chunk, dst_file = idx["chunk"], idx["file"]
|
||||
dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)
|
||||
|
||||
@@ -698,7 +752,9 @@ def append_or_create_parquet_file(
|
||||
return idx, (dst_chunk, dst_file)
|
||||
|
||||
|
||||
def finalize_aggregation(aggr_meta, all_metadata):
|
||||
def finalize_aggregation(
|
||||
aggr_meta: LeRobotDatasetMetadata, all_metadata: list[LeRobotDatasetMetadata]
|
||||
) -> None:
|
||||
"""Finalizes the dataset aggregation by writing summary files and statistics.
|
||||
|
||||
Writes the tasks file, info file with total counts and splits, and
|
||||
@@ -708,16 +764,16 @@ def finalize_aggregation(aggr_meta, all_metadata):
|
||||
aggr_meta: Aggregated dataset metadata.
|
||||
all_metadata: List of all source dataset metadata objects.
|
||||
"""
|
||||
logging.info("write tasks")
|
||||
logger.info("write tasks")
|
||||
write_tasks(aggr_meta.tasks, aggr_meta.root)
|
||||
|
||||
logging.info("write info")
|
||||
logger.info("write info")
|
||||
aggr_meta.info.total_tasks = len(aggr_meta.tasks)
|
||||
aggr_meta.info.total_episodes = sum(m.total_episodes for m in all_metadata)
|
||||
aggr_meta.info.total_frames = sum(m.total_frames for m in all_metadata)
|
||||
aggr_meta.info.splits = {"train": f"0:{sum(m.total_episodes for m in all_metadata)}"}
|
||||
write_info(aggr_meta.info, aggr_meta.root)
|
||||
|
||||
logging.info("write stats")
|
||||
logger.info("write stats")
|
||||
aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata])
|
||||
write_stats(aggr_meta.stats, aggr_meta.root)
|
||||
|
||||
@@ -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],
|
||||
|
||||
@@ -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.")
|
||||
@@ -302,6 +302,33 @@ def _pad_evo1_stats(
|
||||
return padded_stats
|
||||
|
||||
|
||||
def _refresh_evo1_normalization_steps(
|
||||
config: Evo1Config,
|
||||
preprocessor: PolicyProcessorPipeline,
|
||||
postprocessor: PolicyProcessorPipeline,
|
||||
) -> None:
|
||||
"""Re-pad checkpoint-loaded (un)normalizer stats/features to EVO1's fixed widths.
|
||||
|
||||
Loading a checkpoint injects the raw dataset stats (unpadded to max_state_dim/max_action_dim)
|
||||
into the (un)normalizer via the generic override path in make_pre_post_processors. Those stats
|
||||
and their declared features must be re-padded/reshaped to EVO1's fixed widths, otherwise
|
||||
normalization fails against the padded state/action tensors (e.g. state padded to 24 vs. 8-dim
|
||||
LIBERO stats). Padding is a no-op when stats are already at the target width.
|
||||
"""
|
||||
normalization_features = _evo1_normalization_features(config)
|
||||
action_features = _evo1_action_features(config)
|
||||
for step in preprocessor.steps:
|
||||
if isinstance(step, NormalizerProcessorStep):
|
||||
step.features = normalization_features
|
||||
step.stats = _pad_evo1_stats(config, step.stats)
|
||||
step.to(device=step.device, dtype=step.dtype)
|
||||
for step in postprocessor.steps:
|
||||
if isinstance(step, UnnormalizerProcessorStep):
|
||||
step.features = action_features
|
||||
step.stats = _pad_evo1_stats(config, step.stats)
|
||||
step.to(device=step.device, dtype=step.dtype)
|
||||
|
||||
|
||||
def reconcile_evo1_processors(
|
||||
config: Evo1Config,
|
||||
preprocessor: PolicyProcessorPipeline,
|
||||
@@ -309,16 +336,19 @@ def reconcile_evo1_processors(
|
||||
) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
|
||||
"""Reconcile checkpoint-loaded pipelines with the current EVO1 config.
|
||||
|
||||
Two things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
|
||||
(converters are plain functions and are never serialized), and eval-time CLI overrides of the
|
||||
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This
|
||||
restores the converter and rebuilds the action step from the current config so those overrides
|
||||
take effect.
|
||||
Three things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
|
||||
(converters are plain functions and are never serialized), eval-time CLI overrides of the
|
||||
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`), and the
|
||||
(un)normalizer stats/features when the generic override path injects raw, unpadded dataset
|
||||
stats. This restores the converter, re-pads the normalization stats to EVO1's fixed widths, and
|
||||
rebuilds the action step from the current config so those overrides take effect.
|
||||
"""
|
||||
# Pipelines reloaded from a checkpoint come back with the default batch converter, which drops
|
||||
# non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
|
||||
preprocessor.to_transition = evo1_batch_to_transition
|
||||
|
||||
_refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
|
||||
|
||||
action_step = Evo1ActionProcessorStep(
|
||||
action_dim=_evo1_action_dim(config),
|
||||
binarize_gripper=config.binarize_gripper,
|
||||
|
||||
@@ -18,7 +18,7 @@ import functools
|
||||
import threading
|
||||
from collections.abc import Callable, Sequence
|
||||
from contextlib import suppress
|
||||
from typing import TypedDict
|
||||
from typing import NotRequired, TypedDict
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F # noqa: N812
|
||||
@@ -36,7 +36,7 @@ class BatchTransition(TypedDict):
|
||||
next_state: dict[str, torch.Tensor]
|
||||
done: torch.Tensor
|
||||
truncated: torch.Tensor
|
||||
complementary_info: dict[str, torch.Tensor | float | int] | None = None
|
||||
complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
|
||||
|
||||
|
||||
def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:
|
||||
|
||||
@@ -510,10 +510,10 @@ class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep):
|
||||
# We only use the ee pose in the dataset, so we don't need the joint positions
|
||||
for n in self.motor_names:
|
||||
features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
|
||||
# We specify the dataset features of this step that we want to be stored in the dataset
|
||||
# Store end-effector features as actions in the dataset schema
|
||||
for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
|
||||
features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature(
|
||||
type=FeatureType.STATE, shape=(1,)
|
||||
type=FeatureType.ACTION, shape=(1,)
|
||||
)
|
||||
return features
|
||||
|
||||
|
||||
@@ -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()
|
||||
@@ -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":
|
||||
|
||||
@@ -28,7 +28,6 @@ lerobot-find-cameras
|
||||
# NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful.
|
||||
|
||||
import argparse
|
||||
import concurrent.futures
|
||||
import logging
|
||||
import time
|
||||
from pathlib import Path
|
||||
@@ -133,7 +132,7 @@ def save_image(
|
||||
camera_identifier: str | int,
|
||||
images_dir: Path,
|
||||
camera_type: str,
|
||||
):
|
||||
) -> None:
|
||||
"""
|
||||
Saves a single image to disk using Pillow. Handles color conversion if necessary.
|
||||
"""
|
||||
@@ -152,7 +151,7 @@ def save_image(
|
||||
logger.error(f"Failed to save image for camera {camera_identifier} (type {camera_type}): {e}")
|
||||
|
||||
|
||||
def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
|
||||
def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> dict[str, Any] | None:
|
||||
"""Create and connect to a camera instance based on metadata."""
|
||||
cam_type = cam_meta.get("type")
|
||||
cam_id = cam_meta.get("id")
|
||||
@@ -165,12 +164,14 @@ def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
|
||||
cv_config = OpenCVCameraConfig(
|
||||
index_or_path=cam_id,
|
||||
color_mode=ColorMode.RGB,
|
||||
warmup_s=warmup_s,
|
||||
)
|
||||
instance = OpenCVCamera(cv_config)
|
||||
elif cam_type == "RealSense":
|
||||
rs_config = RealSenseCameraConfig(
|
||||
serial_number_or_name=cam_id,
|
||||
color_mode=ColorMode.RGB,
|
||||
warmup_s=warmup_s,
|
||||
)
|
||||
instance = RealSenseCamera(rs_config)
|
||||
else:
|
||||
@@ -188,9 +189,7 @@ def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
|
||||
return None
|
||||
|
||||
|
||||
def process_camera_image(
|
||||
cam_dict: dict[str, Any], output_dir: Path, current_time: float
|
||||
) -> concurrent.futures.Future | None:
|
||||
def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_time: float) -> None:
|
||||
"""Capture and process an image from a single camera."""
|
||||
cam = cam_dict["instance"]
|
||||
meta = cam_dict["meta"]
|
||||
@@ -200,7 +199,7 @@ def process_camera_image(
|
||||
try:
|
||||
image_data = cam.read()
|
||||
|
||||
return save_image(
|
||||
save_image(
|
||||
image_data,
|
||||
cam_id_str,
|
||||
output_dir,
|
||||
@@ -215,21 +214,21 @@ def process_camera_image(
|
||||
return None
|
||||
|
||||
|
||||
def cleanup_cameras(cameras_to_use: list[dict[str, Any]]):
|
||||
def cleanup_camera(cam_dict: dict[str, Any]) -> None:
|
||||
"""Disconnect all cameras."""
|
||||
logger.info(f"Disconnecting {len(cameras_to_use)} cameras...")
|
||||
for cam_dict in cameras_to_use:
|
||||
try:
|
||||
if cam_dict["instance"] and cam_dict["instance"].is_connected:
|
||||
cam_dict["instance"].disconnect()
|
||||
except Exception as e:
|
||||
logger.error(f"Error disconnecting camera {cam_dict['meta'].get('id')}: {e}")
|
||||
logger.info(f"Disconnecting camera with ID {cam_dict['meta'].get('id')}...")
|
||||
try:
|
||||
if cam_dict["instance"] and cam_dict["instance"].is_connected:
|
||||
cam_dict["instance"].disconnect()
|
||||
except Exception as e:
|
||||
logger.error(f"Error disconnecting camera {cam_dict['meta'].get('id')}: {e}")
|
||||
|
||||
|
||||
def save_images_from_all_cameras(
|
||||
output_dir: Path,
|
||||
record_time_s: float = 2.0,
|
||||
camera_type: str | None = None,
|
||||
warmup_s: int = 1,
|
||||
):
|
||||
"""
|
||||
Connects to detected cameras (optionally filtered by type) and saves images from each.
|
||||
@@ -240,6 +239,7 @@ def save_images_from_all_cameras(
|
||||
record_time_s: Duration in seconds to record images.
|
||||
camera_type: Optional string to filter cameras ("realsense" or "opencv").
|
||||
If None, uses all detected cameras.
|
||||
warmup_s: Duration in seconds to warmup camera before recording images.
|
||||
"""
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
logger.info(f"Saving images to {output_dir}")
|
||||
@@ -249,40 +249,24 @@ def save_images_from_all_cameras(
|
||||
logger.warning("No cameras detected matching the criteria. Cannot save images.")
|
||||
return
|
||||
|
||||
cameras_to_use = []
|
||||
for cam_meta in all_camera_metadata:
|
||||
camera_instance = create_camera_instance(cam_meta)
|
||||
if camera_instance:
|
||||
cameras_to_use.append(camera_instance)
|
||||
logger.info(
|
||||
f"Starting image capture for {record_time_s} seconds from {len(all_camera_metadata)} cameras."
|
||||
)
|
||||
|
||||
if not cameras_to_use:
|
||||
logger.warning("No cameras could be connected. Aborting image save.")
|
||||
return
|
||||
|
||||
logger.info(f"Starting image capture for {record_time_s} seconds from {len(cameras_to_use)} cameras.")
|
||||
start_time = time.perf_counter()
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=len(cameras_to_use) * 2) as executor:
|
||||
try:
|
||||
try:
|
||||
for cam_meta in all_camera_metadata:
|
||||
cam_dict = create_camera_instance(cam_meta, warmup_s=warmup_s)
|
||||
if cam_dict is None:
|
||||
continue
|
||||
start_time = time.perf_counter()
|
||||
while time.perf_counter() - start_time < record_time_s:
|
||||
futures = []
|
||||
current_capture_time = time.perf_counter()
|
||||
|
||||
for cam_dict in cameras_to_use:
|
||||
future = process_camera_image(cam_dict, output_dir, current_capture_time)
|
||||
if future:
|
||||
futures.append(future)
|
||||
|
||||
if futures:
|
||||
concurrent.futures.wait(futures)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Capture interrupted by user.")
|
||||
finally:
|
||||
print("\nFinalizing image saving...")
|
||||
executor.shutdown(wait=True)
|
||||
cleanup_cameras(cameras_to_use)
|
||||
print(f"Image capture finished. Images saved to {output_dir}")
|
||||
process_camera_image(cam_dict, output_dir, current_capture_time)
|
||||
cleanup_camera(cam_dict)
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Capture interrupted by user.")
|
||||
finally:
|
||||
print(f"Image capture finished. Images saved to {output_dir}")
|
||||
|
||||
|
||||
def main():
|
||||
@@ -291,7 +275,6 @@ def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Unified camera utility script for listing cameras and capturing images."
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"camera_type",
|
||||
type=str,
|
||||
@@ -309,8 +292,14 @@ def main():
|
||||
parser.add_argument(
|
||||
"--record-time-s",
|
||||
type=float,
|
||||
default=6.0,
|
||||
help="Time duration to attempt capturing frames. Default: 6 seconds.",
|
||||
default=2.0,
|
||||
help="Time duration to attempt capturing frames. Default: 2 seconds.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup-s",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Time duration to warmup camera before attempting to capture frames. Default: 1 second.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
save_images_from_all_cameras(**vars(args))
|
||||
|
||||
@@ -171,7 +171,13 @@ class IOSPhone(BasePhone, Teleoperator):
|
||||
# HEBI provides orientation in w, x, y, z format.
|
||||
# Scipy's Rotation expects x, y, z, w.
|
||||
quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]])) # wxyz to xyzw
|
||||
rot = Rotation.from_quat(quat_xyzw)
|
||||
# ARKit can emit zero/NaN quaternions before tracking is ready or on a
|
||||
# dropped packet. Rotation.from_quat now rejects those; degrade the same
|
||||
# way as a missing pose so teleop stays alive mid-session.
|
||||
try:
|
||||
rot = Rotation.from_quat(quat_xyzw)
|
||||
except ValueError:
|
||||
return False, None, None, None
|
||||
pos = ar_pos - rot.apply(self.config.camera_offset)
|
||||
return True, pos, rot, pose
|
||||
|
||||
|
||||
@@ -37,16 +37,25 @@ def auto_select_torch_device() -> torch.device:
|
||||
|
||||
# TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
|
||||
def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
|
||||
"""Given a string, return a torch.device with checks on whether the device is available."""
|
||||
"""Given a string, return a torch.device with checks on whether the device is available.
|
||||
|
||||
Raises:
|
||||
ValueError: If the requested device family is known but not available on
|
||||
this machine (``AssertionError`` was previously used and is easy to
|
||||
mistake for a programmer bug under ``python -O`` where asserts vanish).
|
||||
"""
|
||||
try_device = str(try_device)
|
||||
if try_device.startswith("cuda"):
|
||||
assert torch.cuda.is_available()
|
||||
if not torch.cuda.is_available():
|
||||
raise ValueError(f"Requested device {try_device!r} but CUDA is not available.")
|
||||
device = torch.device(try_device)
|
||||
elif try_device == "mps":
|
||||
assert torch.backends.mps.is_available()
|
||||
if not torch.backends.mps.is_available():
|
||||
raise ValueError("Requested device 'mps' but MPS is not available.")
|
||||
device = torch.device("mps")
|
||||
elif try_device == "xpu":
|
||||
assert torch.xpu.is_available()
|
||||
if not torch.xpu.is_available():
|
||||
raise ValueError("Requested device 'xpu' but XPU is not available.")
|
||||
device = torch.device("xpu")
|
||||
elif try_device == "cpu":
|
||||
device = torch.device("cpu")
|
||||
|
||||
@@ -32,21 +32,21 @@ def load_json(fpath: Path) -> Any:
|
||||
Returns:
|
||||
Any: The data loaded from the JSON file.
|
||||
"""
|
||||
with open(fpath) as f:
|
||||
with open(fpath, encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def write_json(data: dict, fpath: Path) -> None:
|
||||
"""Write data to a JSON file.
|
||||
def write_json(data: JsonLike, fpath: Path) -> None:
|
||||
"""Write JSON-serializable data to a file.
|
||||
|
||||
Creates parent directories if they don't exist.
|
||||
|
||||
Args:
|
||||
data (dict): The dictionary to write.
|
||||
data: JSON-serializable data to write.
|
||||
fpath (Path): The path to the output JSON file.
|
||||
"""
|
||||
fpath.parent.mkdir(exist_ok=True, parents=True)
|
||||
with open(fpath, "w") as f:
|
||||
with open(fpath, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=4, ensure_ascii=False)
|
||||
|
||||
|
||||
|
||||
@@ -30,6 +30,10 @@ def precise_sleep(seconds: float, spin_threshold: float = 0.010, sleep_margin: f
|
||||
"""
|
||||
if seconds <= 0:
|
||||
return
|
||||
if spin_threshold < 0:
|
||||
raise ValueError(f"spin_threshold must be >= 0, got {spin_threshold}")
|
||||
if sleep_margin < 0:
|
||||
raise ValueError(f"sleep_margin must be >= 0, got {sleep_margin}")
|
||||
|
||||
system = platform.system()
|
||||
# On macOS and Windows the scheduler / sleep granularity can make
|
||||
|
||||
@@ -29,10 +29,13 @@ class Rotation:
|
||||
def __init__(self, quat: np.ndarray) -> None:
|
||||
"""Initialize rotation from quaternion [x, y, z, w]."""
|
||||
self._quat = np.asarray(quat, dtype=float)
|
||||
# Normalize quaternion
|
||||
if self._quat.shape != (4,):
|
||||
raise ValueError(f"Quaternion must have shape (4,), got {self._quat.shape}")
|
||||
# Normalize quaternion. Reject the zero vector — it has no orientation.
|
||||
norm = np.linalg.norm(self._quat)
|
||||
if norm > 0:
|
||||
self._quat = self._quat / norm
|
||||
if norm <= 0.0 or not np.isfinite(norm):
|
||||
raise ValueError(f"Quaternion must be a non-zero finite vector; got {self._quat} (norm={norm})")
|
||||
self._quat = self._quat / norm
|
||||
|
||||
@classmethod
|
||||
def from_rotvec(cls, rotvec: np.ndarray) -> "Rotation":
|
||||
|
||||
@@ -14,7 +14,7 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from typing import TypedDict
|
||||
from typing import NotRequired, TypedDict
|
||||
|
||||
import torch
|
||||
|
||||
@@ -28,7 +28,7 @@ class Transition(TypedDict):
|
||||
next_state: dict[str, torch.Tensor]
|
||||
done: bool
|
||||
truncated: bool
|
||||
complementary_info: dict[str, torch.Tensor | float | int] | None = None
|
||||
complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
|
||||
|
||||
|
||||
def move_transition_to_device(transition: Transition, device: str = "cpu") -> Transition:
|
||||
|
||||
@@ -24,7 +24,6 @@ import sys
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from copy import copy, deepcopy
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from statistics import mean
|
||||
from typing import TYPE_CHECKING, Any
|
||||
@@ -61,14 +60,16 @@ def init_logging(
|
||||
accelerator: Optional Accelerator instance (for multi-GPU detection)
|
||||
"""
|
||||
|
||||
def custom_format(record: logging.LogRecord) -> str:
|
||||
dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
fnameline = f"{record.pathname}:{record.lineno}"
|
||||
pid_str = f"[PID: {os.getpid()}] " if display_pid else ""
|
||||
return f"{record.levelname} {pid_str}{dt} {fnameline[-15:]:>15} {record.getMessage()}"
|
||||
class LeRobotFormatter(logging.Formatter):
|
||||
def format(self, record: logging.LogRecord) -> str:
|
||||
record.lerobot_location = f"{record.pathname}:{record.lineno}"[-15:]
|
||||
record.lerobot_pid = f"[PID: {os.getpid()}] " if display_pid else ""
|
||||
return super().format(record)
|
||||
|
||||
formatter = logging.Formatter()
|
||||
formatter.format = custom_format
|
||||
formatter = LeRobotFormatter(
|
||||
"%(levelname)s %(lerobot_pid)s%(asctime)s %(lerobot_location)15s %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
)
|
||||
|
||||
logger = logging.getLogger()
|
||||
logger.setLevel(logging.NOTSET)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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",
|
||||
)
|
||||
|
||||
@@ -496,6 +496,60 @@ def test_evo1_processor_save_load_round_trip_applies_config_overrides(tmp_path):
|
||||
assert "embodiment_id" in processed
|
||||
|
||||
|
||||
def test_reconcile_evo1_processors_repads_overridden_stats(tmp_path):
|
||||
"""Loading a checkpoint and injecting raw (unpadded) dataset stats must be re-padded.
|
||||
|
||||
Regression test: lerobot-train passes the raw dataset stats as normalizer/unnormalizer
|
||||
overrides when resuming from a checkpoint (e.g. stage2 from a stage1 checkpoint). Those stats
|
||||
are at the dataset dims (e.g. LIBERO state=8/action=7), but EVO1 pads state/action to
|
||||
max_state_dim/max_action_dim before normalization, so reconcile_evo1_processors must re-pad the
|
||||
stats or normalization crashes with a shape mismatch.
|
||||
"""
|
||||
config = make_config()
|
||||
preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=make_stats())
|
||||
preprocessor.save_pretrained(tmp_path)
|
||||
postprocessor.save_pretrained(tmp_path)
|
||||
|
||||
# Reload with the generic override path injecting raw, unpadded dataset stats.
|
||||
raw_stats = make_stats()
|
||||
loaded_pre = PolicyProcessorPipeline.from_pretrained(
|
||||
tmp_path,
|
||||
config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json",
|
||||
overrides={"normalizer_processor": {"stats": raw_stats}},
|
||||
to_transition=batch_to_transition,
|
||||
to_output=transition_to_batch,
|
||||
)
|
||||
loaded_post = PolicyProcessorPipeline.from_pretrained(
|
||||
tmp_path,
|
||||
config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json",
|
||||
overrides={"unnormalizer_processor": {"stats": raw_stats}},
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
)
|
||||
|
||||
# Sanity: the override really injected unpadded stats before reconciliation.
|
||||
normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
|
||||
assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (STATE_DIM,)
|
||||
|
||||
loaded_pre, loaded_post = reconcile_evo1_processors(config, loaded_pre, loaded_post)
|
||||
|
||||
normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
|
||||
unnormalizer = next(step for step in loaded_post.steps if isinstance(step, UnnormalizerProcessorStep))
|
||||
assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,)
|
||||
assert normalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
|
||||
assert unnormalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
|
||||
|
||||
# Normalizing a padded state must not raise (this is the exact runtime path that crashed).
|
||||
processed = loaded_pre(
|
||||
{
|
||||
"task": "pick the block",
|
||||
OBS_STATE: torch.zeros(STATE_DIM),
|
||||
f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
|
||||
}
|
||||
)
|
||||
assert processed[OBS_STATE].shape == (1, MAX_STATE_DIM)
|
||||
|
||||
|
||||
def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
|
||||
monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
|
||||
policy = modeling_evo1.Evo1Policy(make_config())
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 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
|
||||
|
||||
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.robots.so_follower.robot_kinematic_processor import (
|
||||
ForwardKinematicsJointsToEEAction,
|
||||
ForwardKinematicsJointsToEEObservation,
|
||||
)
|
||||
|
||||
MOTOR_NAMES = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
|
||||
EE_KEYS = {f"ee.{k}" for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]}
|
||||
|
||||
|
||||
def _joint_bucket(feature_type: FeatureType) -> dict[str, PolicyFeature]:
|
||||
return {f"{n}.pos": PolicyFeature(type=feature_type, shape=(1,)) for n in MOTOR_NAMES}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("step_cls", "bucket", "feature_type"),
|
||||
[
|
||||
(ForwardKinematicsJointsToEEAction, PipelineFeatureType.ACTION, FeatureType.ACTION),
|
||||
(ForwardKinematicsJointsToEEObservation, PipelineFeatureType.OBSERVATION, FeatureType.STATE),
|
||||
],
|
||||
)
|
||||
def test_fk_feature_schema(step_cls, bucket, feature_type):
|
||||
features = {PipelineFeatureType.ACTION: {}, PipelineFeatureType.OBSERVATION: {}}
|
||||
features[bucket] = _joint_bucket(feature_type)
|
||||
out = step_cls(kinematics=None, motor_names=MOTOR_NAMES).transform_features(features)[bucket]
|
||||
assert set(out) == EE_KEYS
|
||||
assert {feature.type for feature in out.values()} == {feature_type}
|
||||
@@ -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
|
||||
@@ -0,0 +1,36 @@
|
||||
#!/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.
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.utils.device_utils import get_safe_torch_device, is_torch_device_available
|
||||
|
||||
|
||||
def test_cpu_always_available():
|
||||
assert get_safe_torch_device("cpu") == torch.device("cpu")
|
||||
assert is_torch_device_available("cpu")
|
||||
|
||||
|
||||
def test_missing_cuda_raises_valueerror():
|
||||
with patch("torch.cuda.is_available", return_value=False), pytest.raises(ValueError, match="CUDA"):
|
||||
get_safe_torch_device("cuda")
|
||||
|
||||
|
||||
def test_missing_mps_raises_valueerror():
|
||||
with patch("torch.backends.mps.is_available", return_value=False), pytest.raises(ValueError, match="MPS"):
|
||||
get_safe_torch_device("mps")
|
||||
@@ -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.
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from lerobot.utils.rotation import Rotation
|
||||
|
||||
|
||||
def test_zero_quaternion_rejected():
|
||||
with pytest.raises(ValueError, match="non-zero"):
|
||||
Rotation(np.zeros(4))
|
||||
|
||||
|
||||
def test_non_finite_quaternion_rejected():
|
||||
with pytest.raises(ValueError, match="non-zero|finite"):
|
||||
Rotation(np.array([np.nan, 0.0, 0.0, 1.0]))
|
||||
|
||||
|
||||
def test_wrong_shape_rejected():
|
||||
with pytest.raises(ValueError, match="shape"):
|
||||
Rotation(np.array([1.0, 0.0, 0.0]))
|
||||
|
||||
|
||||
def test_identity_roundtrip():
|
||||
r = Rotation.from_rotvec(np.zeros(3))
|
||||
assert np.allclose(r.as_rotvec(), 0.0)
|
||||
assert np.allclose(r.as_matrix(), np.eye(3))
|
||||
|
||||
|
||||
def test_rotvec_roundtrip():
|
||||
rotvec = np.array([0.1, -0.2, 0.3])
|
||||
r = Rotation.from_rotvec(rotvec)
|
||||
assert np.allclose(r.as_rotvec(), rotvec, atol=1e-6)
|
||||
Reference in New Issue
Block a user