admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-23 01:41:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into feat/audio_dataset

Browse Source
This commit is contained in:
CarolinePascal
2026-04-01 17:16:58 +02:00
parent 8b9451b585 15934d8d08
commit 9d42de328e
351 changed files with 26219 additions and 7470 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/datasets/test_aggregate.py
+97 -8
View File
@@ -260,8 +260,8 @@ def test_aggregate_datasets(tmp_path, lerobot_dataset_factory):
# Mock the revision to prevent Hub calls during dataset loading
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "test_aggr")
@@ -311,8 +311,8 @@ def test_aggregate_with_low_threshold(tmp_path, lerobot_dataset_factory):
# Mock the revision to prevent Hub calls during dataset loading
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "small_aggr")
@@ -367,8 +367,8 @@ def test_video_timestamps_regression(tmp_path, lerobot_dataset_factory):
)
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "regression_aggr")
@@ -492,8 +492,8 @@ def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory):
# Load the aggregated dataset
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "image_aggr")
@@ -525,3 +525,92 @@ def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory):
assert img.shape[0] == 3, f"Image {image_key} should have 3 channels"
assert_dataset_iteration_works(aggr_ds)
def test_aggregate_already_merged_dataset(tmp_path, lerobot_dataset_factory):
"""Regression test for aggregating a dataset that is itself a result of a previous merge.
This test reproduces the bug where merging datasets with multiple parquet files
(e.g., from a previous merge with file rotation) would cause FileNotFoundError
because metadata file indices were incorrectly preserved instead of being mapped
to their actual destination files.
The fix adds src_to_dst tracking in aggregate_data() to correctly map source
file indices to destination file indices.
"""
# Step 1: Create datasets A and B
ds_a = lerobot_dataset_factory(
root=tmp_path / "ds_a",
repo_id=f"{DUMMY_REPO_ID}_a",
total_episodes=4,
total_frames=200,
)
ds_b = lerobot_dataset_factory(
root=tmp_path / "ds_b",
repo_id=f"{DUMMY_REPO_ID}_b",
total_episodes=4,
total_frames=200,
)
# Step 2: Merge A+B into AB with small file size to force multiple files
aggregate_datasets(
repo_ids=[ds_a.repo_id, ds_b.repo_id],
roots=[ds_a.root, ds_b.root],
aggr_repo_id=f"{DUMMY_REPO_ID}_ab",
aggr_root=tmp_path / "ds_ab",
data_files_size_in_mb=0.01, # Force file rotation
)
with (
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "ds_ab")
ds_ab = LeRobotDataset(f"{DUMMY_REPO_ID}_ab", root=tmp_path / "ds_ab")
# Verify AB has multiple data files (file rotation occurred)
ab_data_files = list((tmp_path / "ds_ab" / "data").rglob("*.parquet"))
assert len(ab_data_files) > 1, "First merge should create multiple parquet files"
# Step 3: Create dataset C
ds_c = lerobot_dataset_factory(
root=tmp_path / "ds_c",
repo_id=f"{DUMMY_REPO_ID}_c",
total_episodes=2,
total_frames=100,
)
# Step 4: Merge AB+C into final - THIS IS WHERE THE BUG OCCURRED
aggregate_datasets(
repo_ids=[ds_ab.repo_id, ds_c.repo_id],
roots=[ds_ab.root, ds_c.root],
aggr_repo_id=f"{DUMMY_REPO_ID}_abc",
aggr_root=tmp_path / "ds_abc",
)
with (
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "ds_abc")
ds_abc = LeRobotDataset(f"{DUMMY_REPO_ID}_abc", root=tmp_path / "ds_abc")
# Step 5: Verify all data files referenced in metadata actually exist
for ep_idx in range(ds_abc.num_episodes):
data_file_path = ds_abc.root / ds_abc.meta.get_data_file_path(ep_idx)
assert data_file_path.exists(), (
f"Episode {ep_idx} references non-existent file: {data_file_path}\n"
"This indicates the src_to_dst mapping fix is not working correctly."
)
# Step 6: Verify we can iterate through the entire dataset without FileNotFoundError
expected_episodes = ds_a.num_episodes + ds_b.num_episodes + ds_c.num_episodes
expected_frames = ds_a.num_frames + ds_b.num_frames + ds_c.num_frames
assert ds_abc.num_episodes == expected_episodes
assert ds_abc.num_frames == expected_frames
# This would raise FileNotFoundError before the fix
assert_dataset_iteration_works(ds_abc)
tests/datasets/test_dataset_metadata.py
+385
View File
@@ -0,0 +1,385 @@
#!/usr/bin/env python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Contract tests for LeRobotDatasetMetadata."""
import json
import numpy as np
import pytest
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
from lerobot.datasets.utils import INFO_PATH
from tests.fixtures.constants import DEFAULT_FPS, DUMMY_ROBOT_TYPE
# ── helpers ──────────────────────────────────────────────────────────
SIMPLE_FEATURES = {
"state": {"dtype": "float32", "shape": (6,), "names": None},
"action": {"dtype": "float32", "shape": (6,), "names": None},
}
VIDEO_FEATURES = {
**SIMPLE_FEATURES,
"observation.images.laptop": {
"dtype": "video",
"shape": (64, 96, 3),
"names": ["height", "width", "channels"],
"info": None,
},
}
IMAGE_FEATURES = {
**SIMPLE_FEATURES,
"observation.images.laptop": {
"dtype": "image",
"shape": (64, 96, 3),
"names": ["height", "width", "channels"],
"info": None,
},
}
def _make_dummy_stats(features: dict) -> dict:
"""Create minimal episode stats matching the given features."""
stats = {}
for key, ft in features.items():
if ft["dtype"] in ("image", "video"):
stats[key] = {
"max": np.ones((3, 1, 1), dtype=np.float32),
"mean": np.full((3, 1, 1), 0.5, dtype=np.float32),
"min": np.zeros((3, 1, 1), dtype=np.float32),
"std": np.full((3, 1, 1), 0.25, dtype=np.float32),
"count": np.array([5]),
}
elif ft["dtype"] in ("float32", "float64", "int64"):
stats[key] = {
"max": np.ones(ft["shape"], dtype=np.float32),
"mean": np.full(ft["shape"], 0.5, dtype=np.float32),
"min": np.zeros(ft["shape"], dtype=np.float32),
"std": np.full(ft["shape"], 0.25, dtype=np.float32),
"count": np.array([5]),
}
return stats
# ── Construction contracts ───────────────────────────────────────────
def test_create_produces_valid_info_on_disk(tmp_path):
"""create() writes info.json and the returned object reflects the provided settings."""
root = tmp_path / "new_ds"
meta = LeRobotDatasetMetadata.create(
repo_id="test/meta",
fps=DEFAULT_FPS,
features=SIMPLE_FEATURES,
robot_type=DUMMY_ROBOT_TYPE,
root=root,
use_videos=False,
)
# info.json was written to disk
assert (root / INFO_PATH).exists()
with open(root / INFO_PATH) as f:
info_on_disk = json.load(f)
assert meta.fps == DEFAULT_FPS
assert meta.robot_type == DUMMY_ROBOT_TYPE
assert "state" in meta.features
assert "action" in meta.features
assert info_on_disk["fps"] == DEFAULT_FPS
def test_create_starts_with_zero_counts(tmp_path):
"""A freshly created metadata has zero episode/frame/task counts."""
root = tmp_path / "empty_ds"
meta = LeRobotDatasetMetadata.create(
repo_id="test/empty", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
)
assert meta.total_episodes == 0
assert meta.total_frames == 0
assert meta.total_tasks == 0
assert meta.tasks is None
assert meta.episodes is None
assert meta.stats is None
def test_create_with_videos_sets_video_path(tmp_path):
"""When features include video-dtype keys, create() produces a non-None video_path."""
root = tmp_path / "video_ds"
meta = LeRobotDatasetMetadata.create(
repo_id="test/video", fps=DEFAULT_FPS, features=VIDEO_FEATURES, root=root, use_videos=True
)
assert meta.video_path is not None
assert len(meta.video_keys) == 1
assert "observation.images.laptop" in meta.video_keys
def test_create_without_videos_has_no_video_path(tmp_path):
"""When use_videos=False and no video features, video_path is None."""
root = tmp_path / "no_video"
meta = LeRobotDatasetMetadata.create(
repo_id="test/novid", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
)
assert meta.video_path is None
assert meta.video_keys == []
def test_create_raises_on_existing_directory(tmp_path):
"""create() raises if root directory already exists."""
root = tmp_path / "existing"
root.mkdir()
with pytest.raises(FileExistsError):
LeRobotDatasetMetadata.create(
repo_id="test/exists", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
)
def test_init_loads_existing_metadata(tmp_path, lerobot_dataset_metadata_factory, info_factory):
"""When metadata files exist on disk, __init__ loads them correctly."""
root = tmp_path / "load_test"
info = info_factory(total_episodes=3, total_frames=150, total_tasks=1, use_videos=False)
meta = lerobot_dataset_metadata_factory(root=root, info=info)
assert meta.total_episodes == 3
assert meta.total_frames == 150
assert meta.fps == info["fps"]
# ── Property accessors ───────────────────────────────────────────────
def test_property_accessors_reflect_info(tmp_path):
"""Properties return values consistent with the info dict."""
root = tmp_path / "props_ds"
meta = LeRobotDatasetMetadata.create(
repo_id="test/props",
fps=DEFAULT_FPS,
features=IMAGE_FEATURES,
robot_type=DUMMY_ROBOT_TYPE,
root=root,
use_videos=False,
)
assert meta.fps == DEFAULT_FPS
assert meta.robot_type == DUMMY_ROBOT_TYPE
# shapes should be tuples
for _key, shape in meta.shapes.items():
assert isinstance(shape, tuple)
# image_keys should contain the image feature
assert "observation.images.laptop" in meta.image_keys
# camera_keys is a superset of image_keys and video_keys
assert set(meta.image_keys + meta.video_keys) == set(meta.camera_keys)
def test_data_path_is_formattable(tmp_path):
"""data_path contains format placeholders that can be .format()-ed."""
root = tmp_path / "fmt_ds"
meta = LeRobotDatasetMetadata.create(
repo_id="test/fmt", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
)
formatted = meta.data_path.format(chunk_index=0, file_index=0)
assert "chunk" in formatted.lower() or "0" in formatted
# ── Task management ──────────────────────────────────────────────────
def test_save_episode_tasks_creates_tasks_dataframe(tmp_path):
"""On a fresh metadata, save_episode_tasks() creates the tasks DataFrame."""
root = tmp_path / "task_ds"
meta = LeRobotDatasetMetadata.create(
repo_id="test/task", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
)
assert meta.tasks is None
meta.save_episode_tasks(["Pick up the cube"])
assert meta.tasks is not None
assert len(meta.tasks) == 1
assert "Pick up the cube" in meta.tasks.index
def test_save_episode_tasks_is_additive(tmp_path):
"""New tasks are added; existing tasks keep their original index."""
root = tmp_path / "additive_ds"
meta = LeRobotDatasetMetadata.create(
repo_id="test/add", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
)
meta.save_episode_tasks(["Task A"])
idx_a = meta.get_task_index("Task A")
meta.save_episode_tasks(["Task A", "Task B"])
assert meta.get_task_index("Task A") == idx_a # unchanged
assert meta.get_task_index("Task B") is not None
assert len(meta.tasks) == 2
def test_get_task_index_returns_none_for_unknown(tmp_path):
"""get_task_index() returns None for an unknown task."""
root = tmp_path / "unknown_ds"
meta = LeRobotDatasetMetadata.create(
repo_id="test/unknown", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
)
meta.save_episode_tasks(["Known task"])
assert meta.get_task_index("Known task") == 0
assert meta.get_task_index("Unknown task") is None
def test_save_episode_tasks_rejects_duplicates(tmp_path):
"""save_episode_tasks() raises ValueError on duplicate task strings."""
root = tmp_path / "dup_ds"
meta = LeRobotDatasetMetadata.create(
repo_id="test/dup", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
)
with pytest.raises(ValueError):
meta.save_episode_tasks(["Same task", "Same task"])
# ── Episode saving ───────────────────────────────────────────────────
def test_save_episode_increments_counters(tmp_path):
"""After save_episode(), total_episodes and total_frames increase."""
root = tmp_path / "ep_ds"
meta = LeRobotDatasetMetadata.create(
repo_id="test/ep", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
)
meta.save_episode_tasks(["Task 1"])
stats = _make_dummy_stats(meta.features)
meta.save_episode(
episode_index=0,
episode_length=10,
episode_tasks=["Task 1"],
episode_stats=stats,
episode_metadata={},
)
assert meta.total_episodes == 1
assert meta.total_frames == 10
def test_save_episode_updates_stats(tmp_path):
"""After save_episode(), .stats is non-None and has feature keys."""
root = tmp_path / "stats_ds"
meta = LeRobotDatasetMetadata.create(
repo_id="test/stats", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
)
meta.save_episode_tasks(["Task 1"])
stats = _make_dummy_stats(meta.features)
meta.save_episode(
episode_index=0,
episode_length=5,
episode_tasks=["Task 1"],
episode_stats=stats,
episode_metadata={},
)
assert meta.stats is not None
# Stats should contain at least the user-defined feature keys
for key in SIMPLE_FEATURES:
assert key in meta.stats
# ── Chunk settings ───────────────────────────────────────────────────
def test_update_chunk_settings_persists(tmp_path):
"""update_chunk_settings() changes values and writes info.json."""
root = tmp_path / "chunk_ds"
meta = LeRobotDatasetMetadata.create(
repo_id="test/chunk", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
)
original = meta.get_chunk_settings()
meta.update_chunk_settings(chunks_size=500)
assert meta.chunks_size == 500
assert meta.chunks_size != original["chunks_size"] or original["chunks_size"] == 500
# Verify persisted
with open(root / INFO_PATH) as f:
info_on_disk = json.load(f)
assert info_on_disk["chunks_size"] == 500
def test_update_chunk_settings_rejects_non_positive(tmp_path):
"""update_chunk_settings() raises ValueError for <= 0 values."""
root = tmp_path / "bad_chunk"
meta = LeRobotDatasetMetadata.create(
repo_id="test/bad", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
)
with pytest.raises(ValueError):
meta.update_chunk_settings(chunks_size=0)
with pytest.raises(ValueError):
meta.update_chunk_settings(data_files_size_in_mb=-1)
# ── Finalization ─────────────────────────────────────────────────────
def test_finalize_is_idempotent(tmp_path):
"""Calling finalize() multiple times does not raise."""
root = tmp_path / "fin_ds"
meta = LeRobotDatasetMetadata.create(
repo_id="test/fin", fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root, use_videos=False
)
meta.finalize()
meta.finalize() # second call should not raise
def test_finalize_flushes_buffered_metadata(tmp_path):
"""Episodes saved before finalize() are written to parquet."""
root = tmp_path / "flush_ds"
meta = LeRobotDatasetMetadata.create(
repo_id="test/flush",
fps=DEFAULT_FPS,
features=SIMPLE_FEATURES,
root=root,
use_videos=False,
metadata_buffer_size=100, # large buffer so nothing auto-flushes
)
meta.save_episode_tasks(["Task 1"])
stats = _make_dummy_stats(meta.features)
# Save a few episodes (won't auto-flush since buffer_size=100)
for i in range(3):
meta.save_episode(
episode_index=i,
episode_length=5,
episode_tasks=["Task 1"],
episode_stats=stats,
episode_metadata={},
)
# Before finalize, the parquet might not exist yet
meta.finalize()
# After finalize, episodes parquet should exist
episodes_dir = root / "meta" / "episodes"
assert episodes_dir.exists()
parquet_files = list(episodes_dir.rglob("*.parquet"))
assert len(parquet_files) > 0
tests/datasets/test_dataset_reader.py
+168
View File
@@ -0,0 +1,168 @@
#!/usr/bin/env python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Contract tests for DatasetReader."""
from lerobot.datasets.dataset_reader import DatasetReader
from lerobot.datasets.video_utils import get_safe_default_codec
# ── Loading ──────────────────────────────────────────────────────────
def test_try_load_returns_true_when_data_exists(tmp_path, lerobot_dataset_factory):
"""Given a fully written dataset, try_load() returns True."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=20, use_videos=False
)
reader = DatasetReader(
meta=dataset.meta,
root=dataset.root,
episodes=None,
tolerance_s=1e-4,
video_backend=get_safe_default_codec(),
delta_timestamps=None,
image_transforms=None,
)
assert reader.try_load() is True
assert reader.hf_dataset is not None
def test_try_load_returns_false_when_no_data(tmp_path):
"""When only metadata exists (no data/ parquets), try_load() returns False."""
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
root = tmp_path / "meta_only"
features = {"state": {"dtype": "float32", "shape": (2,), "names": None}}
meta = LeRobotDatasetMetadata.create(
repo_id="test/meta_only", fps=30, features=features, root=root, use_videos=False
)
reader = DatasetReader(
meta=meta,
root=meta.root,
episodes=None,
tolerance_s=1e-4,
video_backend=get_safe_default_codec(),
delta_timestamps=None,
image_transforms=None,
)
assert reader.try_load() is False
assert reader.hf_dataset is None
# ── Counts ───────────────────────────────────────────────────────────
def test_num_frames_without_filter(tmp_path, lerobot_dataset_factory):
"""With episodes=None, num_frames equals total_frames."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=3, total_frames=60, use_videos=False
)
assert dataset.reader.num_frames == dataset.meta.total_frames
def test_num_episodes_without_filter(tmp_path, lerobot_dataset_factory):
"""With episodes=None, num_episodes equals total_episodes."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=3, total_frames=60, use_videos=False
)
assert dataset.reader.num_episodes == dataset.meta.total_episodes
def test_num_frames_with_episode_filter(tmp_path, lerobot_dataset_factory):
"""When filtering to a subset, only those episodes' frames are counted."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=5, total_frames=100, episodes=[0, 2], use_videos=False
)
# Filtered frames should be less than total
assert dataset.reader.num_frames <= dataset.meta.total_frames
assert dataset.reader.num_episodes == 2
# ── get_item ─────────────────────────────────────────────────────────
def test_get_item_returns_expected_keys(tmp_path, lerobot_dataset_factory):
"""get_item(0) returns a dict with expected keys."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=1, total_frames=10, use_videos=False
)
item = dataset.reader.get_item(0)
# Standard keys that must always be present
for key in ["index", "episode_index", "frame_index", "timestamp", "task_index", "task"]:
assert key in item, f"Missing key: {key}"
def test_get_item_values_are_correct(tmp_path, lerobot_dataset_factory):
"""get_item() returns correct index and episode_index."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=20, use_videos=False
)
item_0 = dataset.reader.get_item(0)
assert item_0["index"].item() == 0
assert item_0["episode_index"].item() == 0
# ── Transforms ───────────────────────────────────────────────────────
def test_image_transforms_are_applied(tmp_path, lerobot_dataset_factory):
"""When image_transforms is provided, get_item() applies it to camera keys."""
transform_called = {"count": 0}
def sentinel_transform(img):
transform_called["count"] += 1
return img
dataset = lerobot_dataset_factory(
root=tmp_path / "ds",
total_episodes=1,
total_frames=5,
use_videos=False,
image_transforms=sentinel_transform,
)
item = dataset[0] # noqa: F841
# Should have been called once per camera key per frame
num_cameras = len(dataset.meta.camera_keys)
if num_cameras > 0:
assert transform_called["count"] >= 1
# ── File paths ───────────────────────────────────────────────────────
def test_get_episodes_file_paths_returns_data_paths(tmp_path, lerobot_dataset_factory):
"""get_episodes_file_paths() returns paths including data/ paths."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=20, use_videos=False
)
paths = dataset.reader.get_episodes_file_paths()
assert len(paths) > 0
assert any("data/" in str(p) for p in paths)
def test_get_episodes_file_paths_includes_video_paths(tmp_path, lerobot_dataset_factory):
"""When dataset has video keys, file paths include video/ paths."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=20, use_videos=True
)
if len(dataset.meta.video_keys) > 0:
paths = dataset.reader.get_episodes_file_paths()
assert any("video" in str(p).lower() for p in paths)
tests/datasets/test_dataset_tools.py
+227 -58
View File
@@ -26,6 +26,7 @@ from lerobot.datasets.dataset_tools import (
delete_episodes,
merge_datasets,
modify_features,
modify_tasks,
remove_feature,
split_dataset,
)
@@ -66,8 +67,8 @@ def test_delete_single_episode(sample_dataset, tmp_path):
output_dir = tmp_path / "filtered"
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(output_dir)
@@ -92,8 +93,8 @@ def test_delete_multiple_episodes(sample_dataset, tmp_path):
output_dir = tmp_path / "filtered"
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(output_dir)
@@ -149,8 +150,8 @@ def test_split_by_episodes(sample_dataset, tmp_path):
}
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
@@ -192,8 +193,8 @@ def test_split_by_fractions(sample_dataset, tmp_path):
}
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
@@ -269,8 +270,8 @@ def test_merge_two_datasets(sample_dataset, tmp_path, empty_lerobot_dataset_fact
dataset2.finalize()
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "merged_dataset")
@@ -309,8 +310,8 @@ def test_add_features_with_values(sample_dataset, tmp_path):
}
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "with_reward")
@@ -345,8 +346,8 @@ def test_add_features_with_callable(sample_dataset, tmp_path):
"reward": (compute_reward, feature_info),
}
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "with_reward")
@@ -400,8 +401,8 @@ def test_modify_features_add_and_remove(sample_dataset, tmp_path):
feature_info = {"dtype": "float32", "shape": (1,), "names": None}
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "modified")
@@ -433,8 +434,8 @@ def test_modify_features_only_add(sample_dataset, tmp_path):
feature_info = {"dtype": "float32", "shape": (1,), "names": None}
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "modified")
@@ -456,8 +457,8 @@ def test_modify_features_only_remove(sample_dataset, tmp_path):
feature_info = {"dtype": "float32", "shape": (1,), "names": None}
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.side_effect = lambda repo_id, **kwargs: str(kwargs.get("local_dir", tmp_path))
@@ -493,8 +494,8 @@ def test_remove_single_feature(sample_dataset, tmp_path):
"reward": (np.random.randn(50, 1).astype(np.float32), feature_info),
}
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.side_effect = lambda repo_id, **kwargs: str(kwargs.get("local_dir", tmp_path))
@@ -520,8 +521,8 @@ def test_remove_single_feature(sample_dataset, tmp_path):
def test_remove_multiple_features(sample_dataset, tmp_path):
"""Test removing multiple features at once."""
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.side_effect = lambda repo_id, **kwargs: str(kwargs.get("local_dir", tmp_path))
@@ -575,8 +576,8 @@ def test_remove_camera_feature(sample_dataset, tmp_path):
camera_to_remove = camera_keys[0]
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "without_camera")
@@ -597,8 +598,8 @@ def test_remove_camera_feature(sample_dataset, tmp_path):
def test_complex_workflow_integration(sample_dataset, tmp_path):
"""Test a complex workflow combining multiple operations."""
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.side_effect = lambda repo_id, **kwargs: str(kwargs.get("local_dir", tmp_path))
@@ -646,8 +647,8 @@ def test_delete_episodes_preserves_stats(sample_dataset, tmp_path):
output_dir = tmp_path / "filtered"
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(output_dir)
@@ -670,8 +671,8 @@ def test_delete_episodes_preserves_tasks(sample_dataset, tmp_path):
output_dir = tmp_path / "filtered"
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(output_dir)
@@ -698,8 +699,8 @@ def test_split_three_ways(sample_dataset, tmp_path):
}
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
@@ -731,8 +732,8 @@ def test_split_preserves_stats(sample_dataset, tmp_path):
splits = {"train": [0, 1, 2], "val": [3, 4]}
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
@@ -789,8 +790,8 @@ def test_merge_three_datasets(sample_dataset, tmp_path, empty_lerobot_dataset_fa
datasets.append(dataset)
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "merged_dataset")
@@ -831,8 +832,8 @@ def test_merge_preserves_stats(sample_dataset, tmp_path, empty_lerobot_dataset_f
dataset2.finalize()
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "merged_dataset")
@@ -865,8 +866,8 @@ def test_add_features_preserves_existing_stats(sample_dataset, tmp_path):
}
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(tmp_path / "with_reward")
@@ -889,8 +890,8 @@ def test_remove_feature_updates_stats(sample_dataset, tmp_path):
feature_info = {"dtype": "float32", "shape": (1,), "names": None}
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.side_effect = lambda repo_id, **kwargs: str(kwargs.get("local_dir", tmp_path))
@@ -918,8 +919,8 @@ def test_delete_consecutive_episodes(sample_dataset, tmp_path):
output_dir = tmp_path / "filtered"
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(output_dir)
@@ -942,8 +943,8 @@ def test_delete_first_and_last_episodes(sample_dataset, tmp_path):
output_dir = tmp_path / "filtered"
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(output_dir)
@@ -970,8 +971,8 @@ def test_split_all_episodes_assigned(sample_dataset, tmp_path):
}
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
@@ -998,8 +999,8 @@ def test_modify_features_preserves_file_structure(sample_dataset, tmp_path):
feature_info = {"dtype": "float32", "shape": (1,), "names": None}
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
@@ -1019,7 +1020,7 @@ def test_modify_features_preserves_file_structure(sample_dataset, tmp_path):
# Get original chunk/file indices from first episode
if train_dataset.meta.episodes is None:
from lerobot.datasets.utils import load_episodes
from lerobot.datasets.io_utils import load_episodes
train_dataset.meta.episodes = load_episodes(train_dataset.meta.root)
original_chunk_indices = [ep["data/chunk_index"] for ep in train_dataset.meta.episodes]
@@ -1039,7 +1040,7 @@ def test_modify_features_preserves_file_structure(sample_dataset, tmp_path):
# Check that chunk/file indices are preserved
if modified_dataset.meta.episodes is None:
from lerobot.datasets.utils import load_episodes
from lerobot.datasets.io_utils import load_episodes
modified_dataset.meta.episodes = load_episodes(modified_dataset.meta.root)
new_chunk_indices = [ep["data/chunk_index"] for ep in modified_dataset.meta.episodes]
@@ -1050,6 +1051,174 @@ def test_modify_features_preserves_file_structure(sample_dataset, tmp_path):
assert "reward" in modified_dataset.meta.features
def test_modify_tasks_single_task_for_all(sample_dataset):
"""Test setting a single task for all episodes."""
new_task = "Pick up the cube and place it"
modified_dataset = modify_tasks(sample_dataset, new_task=new_task)
# Verify all episodes have the new task
assert len(modified_dataset.meta.tasks) == 1
assert new_task in modified_dataset.meta.tasks.index
# Verify task_index is 0 for all frames (only one task)
for i in range(len(modified_dataset)):
item = modified_dataset[i]
assert item["task_index"].item() == 0
assert item["task"] == new_task
def test_modify_tasks_episode_specific(sample_dataset):
"""Test setting different tasks for specific episodes."""
episode_tasks = {
0: "Task A",
1: "Task B",
2: "Task A",
3: "Task C",
4: "Task B",
}
modified_dataset = modify_tasks(sample_dataset, episode_tasks=episode_tasks)
# Verify correct number of unique tasks
unique_tasks = set(episode_tasks.values())
assert len(modified_dataset.meta.tasks) == len(unique_tasks)
# Verify each episode has the correct task
for ep_idx, expected_task in episode_tasks.items():
ep_data = modified_dataset.meta.episodes[ep_idx]
assert ep_data["tasks"][0] == expected_task
def test_modify_tasks_default_with_overrides(sample_dataset):
"""Test setting a default task with specific overrides."""
default_task = "Default task"
override_task = "Special task"
episode_tasks = {2: override_task, 4: override_task}
modified_dataset = modify_tasks(
sample_dataset,
new_task=default_task,
episode_tasks=episode_tasks,
)
# Verify correct number of unique tasks
assert len(modified_dataset.meta.tasks) == 2
assert default_task in modified_dataset.meta.tasks.index
assert override_task in modified_dataset.meta.tasks.index
# Verify episodes have correct tasks
for ep_idx in range(5):
ep_data = modified_dataset.meta.episodes[ep_idx]
if ep_idx in episode_tasks:
assert ep_data["tasks"][0] == override_task
else:
assert ep_data["tasks"][0] == default_task
def test_modify_tasks_no_task_specified(sample_dataset):
"""Test error when no task is specified."""
with pytest.raises(ValueError, match="Must specify at least one of new_task or episode_tasks"):
modify_tasks(sample_dataset)
def test_modify_tasks_invalid_episode_indices(sample_dataset):
"""Test error with invalid episode indices."""
with pytest.raises(ValueError, match="Invalid episode indices"):
modify_tasks(sample_dataset, episode_tasks={10: "Task", 20: "Task"})
def test_modify_tasks_updates_info_json(sample_dataset):
"""Test that total_tasks is updated in info.json."""
episode_tasks = {0: "Task A", 1: "Task B", 2: "Task C", 3: "Task A", 4: "Task B"}
modified_dataset = modify_tasks(sample_dataset, episode_tasks=episode_tasks)
# Verify total_tasks is updated
assert modified_dataset.meta.total_tasks == 3
def test_modify_tasks_preserves_other_metadata(sample_dataset):
"""Test that modifying tasks preserves other metadata."""
original_frames = sample_dataset.meta.total_frames
original_episodes = sample_dataset.meta.total_episodes
original_fps = sample_dataset.meta.fps
modified_dataset = modify_tasks(sample_dataset, new_task="New task")
# Verify other metadata is preserved
assert modified_dataset.meta.total_frames == original_frames
assert modified_dataset.meta.total_episodes == original_episodes
assert modified_dataset.meta.fps == original_fps
def test_modify_tasks_task_index_correct(sample_dataset):
"""Test that task_index values are correct in data files."""
# Create tasks that will have predictable indices (sorted alphabetically)
episode_tasks = {
0: "Alpha task", # Will be index 0
1: "Beta task", # Will be index 1
2: "Alpha task", # Will be index 0
3: "Gamma task", # Will be index 2
4: "Beta task", # Will be index 1
}
modified_dataset = modify_tasks(sample_dataset, episode_tasks=episode_tasks)
# Verify task indices are correct
task_to_expected_idx = {
"Alpha task": 0,
"Beta task": 1,
"Gamma task": 2,
}
for i in range(len(modified_dataset)):
item = modified_dataset[i]
ep_idx = item["episode_index"].item()
expected_task = episode_tasks[ep_idx]
expected_idx = task_to_expected_idx[expected_task]
assert item["task_index"].item() == expected_idx
assert item["task"] == expected_task
def test_modify_tasks_in_place(sample_dataset):
"""Test that modify_tasks modifies the dataset in-place."""
original_root = sample_dataset.root
modified_dataset = modify_tasks(sample_dataset, new_task="New task")
# Verify same instance is returned and root is unchanged
assert modified_dataset is sample_dataset
assert modified_dataset.root == original_root
def test_modify_tasks_keeps_original_when_not_overridden(sample_dataset):
"""Test that original tasks are kept when using episode_tasks without new_task."""
from lerobot.datasets.io_utils import load_episodes
# Ensure episodes metadata is loaded
if sample_dataset.meta.episodes is None:
sample_dataset.meta.episodes = load_episodes(sample_dataset.meta.root)
# Get original tasks for episodes not being overridden
original_task_ep0 = sample_dataset.meta.episodes[0]["tasks"][0]
original_task_ep1 = sample_dataset.meta.episodes[1]["tasks"][0]
# Only override episodes 2, 3, 4
episode_tasks = {2: "New Task A", 3: "New Task B", 4: "New Task A"}
modified_dataset = modify_tasks(sample_dataset, episode_tasks=episode_tasks)
# Verify original tasks are kept for episodes 0 and 1
assert modified_dataset.meta.episodes[0]["tasks"][0] == original_task_ep0
assert modified_dataset.meta.episodes[1]["tasks"][0] == original_task_ep1
# Verify new tasks for overridden episodes
assert modified_dataset.meta.episodes[2]["tasks"][0] == "New Task A"
assert modified_dataset.meta.episodes[3]["tasks"][0] == "New Task B"
assert modified_dataset.meta.episodes[4]["tasks"][0] == "New Task A"
def test_convert_image_to_video_dataset(tmp_path):
"""Test converting lerobot/pusht_image dataset to video format."""
from lerobot.datasets.lerobot_dataset import LeRobotDataset
@@ -1060,8 +1229,8 @@ def test_convert_image_to_video_dataset(tmp_path):
output_dir = tmp_path / "pusht_video"
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(output_dir)
@@ -1123,8 +1292,8 @@ def test_convert_image_to_video_dataset_subset_episodes(tmp_path):
output_dir = tmp_path / "pusht_video_subset"
with (
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
):
mock_get_safe_version.return_value = "v3.0"
mock_snapshot_download.return_value = str(output_dir)
tests/datasets/test_dataset_utils.py
+19 -2
View File
@@ -19,11 +19,28 @@ import torch
from datasets import Dataset
from huggingface_hub import DatasetCard
from lerobot.datasets.push_dataset_to_hub.utils import calculate_episode_data_index
from lerobot.datasets.utils import combine_feature_dicts, create_lerobot_dataset_card, hf_transform_to_torch
from lerobot.datasets.feature_utils import combine_feature_dicts
from lerobot.datasets.io_utils import hf_transform_to_torch
from lerobot.datasets.utils import create_lerobot_dataset_card
from lerobot.utils.constants import ACTION, OBS_IMAGES
def calculate_episode_data_index(hf_dataset: Dataset) -> dict[str, torch.Tensor]:
"""Calculate episode data index for testing. Returns {"from": Tensor, "to": Tensor}."""
episode_data_index: dict[str, list[int]] = {"from": [], "to": []}
current_episode = None
if len(hf_dataset) == 0:
return {"from": torch.tensor([]), "to": torch.tensor([])}
for idx, episode_idx in enumerate(hf_dataset["episode_index"]):
if episode_idx != current_episode:
episode_data_index["from"].append(idx)
if current_episode is not None:
episode_data_index["to"].append(idx)
current_episode = episode_idx
episode_data_index["to"].append(idx + 1)
return {k: torch.tensor(v) for k, v in episode_data_index.items()}
def test_default_parameters():
card = create_lerobot_dataset_card()
assert isinstance(card, DatasetCard)
tests/datasets/test_dataset_writer.py
+226
View File
@@ -0,0 +1,226 @@
#!/usr/bin/env python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Contract tests for DatasetWriter."""
from pathlib import Path
from unittest.mock import patch
import numpy as np
import pytest
import torch
from PIL import Image
from lerobot.datasets.dataset_writer import _encode_video_worker
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.utils import DEFAULT_IMAGE_PATH
from tests.fixtures.constants import DEFAULT_FPS, DUMMY_REPO_ID
SIMPLE_FEATURES = {
"state": {"dtype": "float32", "shape": (6,), "names": None},
"action": {"dtype": "float32", "shape": (6,), "names": None},
}
def _make_frame(features: dict, task: str = "Dummy task") -> dict:
"""Build a valid frame dict for the given features."""
frame = {"task": task}
for key, ft in features.items():
if ft["dtype"] in ("image", "video"):
frame[key] = np.random.randint(0, 256, size=ft["shape"], dtype=np.uint8)
elif ft["dtype"] in ("float32", "float64"):
frame[key] = torch.randn(ft["shape"])
elif ft["dtype"] == "int64":
frame[key] = torch.zeros(ft["shape"], dtype=torch.int64)
return frame
# ── Existing encode_video_worker tests ───────────────────────────────
def test_encode_video_worker_forwards_vcodec(tmp_path):
"""_encode_video_worker correctly forwards the vcodec parameter."""
video_key = "observation.images.laptop"
fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)
img_dir = tmp_path / Path(fpath).parent
img_dir.mkdir(parents=True, exist_ok=True)
Image.new("RGB", (64, 64), color="red").save(img_dir / "frame-000000.png")
captured_kwargs = {}
def mock_encode(imgs_dir, video_path, fps, **kwargs):
captured_kwargs.update(kwargs)
Path(video_path).parent.mkdir(parents=True, exist_ok=True)
Path(video_path).touch()
with patch("lerobot.datasets.dataset_writer.encode_video_frames", side_effect=mock_encode):
_encode_video_worker(video_key, 0, tmp_path, fps=30, vcodec="h264")
assert captured_kwargs["vcodec"] == "h264"
def test_encode_video_worker_default_vcodec(tmp_path):
"""_encode_video_worker uses libsvtav1 as the default codec."""
video_key = "observation.images.laptop"
fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)
img_dir = tmp_path / Path(fpath).parent
img_dir.mkdir(parents=True, exist_ok=True)
Image.new("RGB", (64, 64), color="red").save(img_dir / "frame-000000.png")
captured_kwargs = {}
def mock_encode(imgs_dir, video_path, fps, **kwargs):
captured_kwargs.update(kwargs)
Path(video_path).parent.mkdir(parents=True, exist_ok=True)
Path(video_path).touch()
with patch("lerobot.datasets.dataset_writer.encode_video_frames", side_effect=mock_encode):
_encode_video_worker(video_key, 0, tmp_path, fps=30)
assert captured_kwargs["vcodec"] == "libsvtav1"
# ── add_frame contracts ──────────────────────────────────────────────
def test_add_frame_increments_buffer_size(tmp_path):
"""Each add_frame() call increases episode_buffer['size'] by 1."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
assert dataset.writer.episode_buffer["size"] == 0
dataset.add_frame(_make_frame(SIMPLE_FEATURES))
assert dataset.writer.episode_buffer["size"] == 1
dataset.add_frame(_make_frame(SIMPLE_FEATURES))
assert dataset.writer.episode_buffer["size"] == 2
def test_add_frame_rejects_missing_feature(tmp_path):
"""add_frame() raises ValueError when a required feature is missing."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
with pytest.raises(ValueError, match="Missing features"):
dataset.add_frame({"task": "Dummy task", "state": torch.randn(6)})
# missing 'action'
# ── save_episode contracts ───────────────────────────────────────────
def test_save_episode_writes_parquet(tmp_path):
"""After save_episode(), at least one .parquet file exists under data/."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
for _ in range(3):
dataset.add_frame(_make_frame(SIMPLE_FEATURES))
dataset.save_episode()
parquet_files = list((tmp_path / "ds" / "data").rglob("*.parquet"))
assert len(parquet_files) > 0
def test_save_episode_updates_counters(tmp_path):
"""After save_episode(), metadata counters are updated."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
for _ in range(5):
dataset.add_frame(_make_frame(SIMPLE_FEATURES))
dataset.save_episode()
assert dataset.meta.total_episodes == 1
assert dataset.meta.total_frames == 5
def test_save_episode_resets_buffer(tmp_path):
"""After save_episode(), the episode buffer is reset."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
for _ in range(3):
dataset.add_frame(_make_frame(SIMPLE_FEATURES))
dataset.save_episode()
assert dataset.writer.episode_buffer["size"] == 0
def test_save_multiple_episodes(tmp_path):
"""Recording 3 episodes results in correct total counts."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
total_frames = 0
for ep in range(3):
n_frames = ep + 2 # 2, 3, 4
for _ in range(n_frames):
dataset.add_frame(_make_frame(SIMPLE_FEATURES))
dataset.save_episode()
total_frames += n_frames
assert dataset.meta.total_episodes == 3
assert dataset.meta.total_frames == total_frames
# ── clear / lifecycle ────────────────────────────────────────────────
def test_clear_resets_buffer(tmp_path):
"""clear_episode_buffer() resets the buffer size to 0."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
dataset.add_frame(_make_frame(SIMPLE_FEATURES))
assert dataset.writer.episode_buffer["size"] == 1
dataset.clear_episode_buffer()
assert dataset.writer.episode_buffer["size"] == 0
def test_finalize_is_idempotent(tmp_path):
"""Calling finalize() twice does not raise."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
for _ in range(3):
dataset.add_frame(_make_frame(SIMPLE_FEATURES))
dataset.save_episode()
dataset.finalize()
dataset.finalize() # second call should not raise
def test_finalize_then_read_roundtrip(tmp_path):
"""Write data, finalize, re-open, and verify data matches."""
root = tmp_path / "roundtrip"
features = {"state": {"dtype": "float32", "shape": (2,), "names": None}}
dataset = LeRobotDataset.create(repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=features, root=root)
# Record known values
known_states = []
for i in range(5):
state = torch.tensor([float(i), float(i * 10)])
known_states.append(state)
dataset.add_frame({"task": "Test task", "state": state})
dataset.save_episode()
dataset.finalize()
# Read back
for i in range(5):
item = dataset[i]
assert torch.allclose(item["state"], known_states[i], atol=1e-5)
tests/datasets/test_datasets.py
+63 -127
View File
@@ -30,23 +30,19 @@ import lerobot
from lerobot.configs.default import DatasetConfig
from lerobot.configs.train import TrainPipelineConfig
from lerobot.datasets.factory import make_dataset
from lerobot.datasets.feature_utils import get_hf_features_from_features, hw_to_dataset_features
from lerobot.datasets.image_writer import image_array_to_pil_image
from lerobot.datasets.lerobot_dataset import (
VALID_VIDEO_CODECS,
LeRobotDataset,
MultiLeRobotDataset,
_encode_video_worker,
)
from lerobot.datasets.io_utils import hf_transform_to_torch
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.multi_dataset import MultiLeRobotDataset
from lerobot.datasets.utils import (
DEFAULT_AUDIO_CHUNK_DURATION,
DEFAULT_CHUNK_SIZE,
DEFAULT_DATA_FILE_SIZE_IN_MB,
DEFAULT_VIDEO_FILE_SIZE_IN_MB,
create_branch,
get_hf_features_from_features,
hf_transform_to_torch,
hw_to_dataset_features,
)
from lerobot.datasets.video_utils import VALID_VIDEO_CODECS
from lerobot.envs.factory import make_env_config
from lerobot.policies.factory import make_policy_config
from lerobot.robots import make_robot_from_config
@@ -111,7 +107,7 @@ def audio_dataset(tmp_path, empty_lerobot_dataset_factory):
def test_same_attributes_defined(tmp_path, lerobot_dataset_factory):
"""
Instantiate a LeRobotDataset both ways with '__init__()' and 'create()' and verify that instantiated
objects have the same sets of attributes defined.
objects have the same sets of facade-level attributes defined.
"""
# Instantiate both ways
robot = make_robot_from_config(MockRobotConfig())
@@ -126,6 +122,7 @@ def test_same_attributes_defined(tmp_path, lerobot_dataset_factory):
root_init = tmp_path / "init"
dataset_init = lerobot_dataset_factory(root=root_init, total_episodes=1, total_frames=1)
# Facade-level attributes should match between __init__ and create()
init_attr = set(vars(dataset_init).keys())
create_attr = set(vars(dataset_create).keys())
@@ -253,6 +250,7 @@ def test_add_frame(tmp_path, empty_lerobot_dataset_factory):
dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
dataset.add_frame({"state": torch.randn(1), "task": "Dummy task"})
dataset.save_episode()
dataset.finalize()
assert len(dataset) == 1
assert dataset[0]["task"] == "Dummy task"
@@ -265,6 +263,7 @@ def test_add_frame_state_1d(tmp_path, empty_lerobot_dataset_factory):
dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
dataset.add_frame({"state": torch.randn(2), "task": "Dummy task"})
dataset.save_episode()
dataset.finalize()
assert dataset[0]["state"].shape == torch.Size([2])
@@ -274,6 +273,7 @@ def test_add_frame_state_2d(tmp_path, empty_lerobot_dataset_factory):
dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
dataset.add_frame({"state": torch.randn(2, 4), "task": "Dummy task"})
dataset.save_episode()
dataset.finalize()
assert dataset[0]["state"].shape == torch.Size([2, 4])
@@ -283,6 +283,7 @@ def test_add_frame_state_3d(tmp_path, empty_lerobot_dataset_factory):
dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
dataset.add_frame({"state": torch.randn(2, 4, 3), "task": "Dummy task"})
dataset.save_episode()
dataset.finalize()
assert dataset[0]["state"].shape == torch.Size([2, 4, 3])
@@ -292,6 +293,7 @@ def test_add_frame_state_4d(tmp_path, empty_lerobot_dataset_factory):
dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
dataset.add_frame({"state": torch.randn(2, 4, 3, 5), "task": "Dummy task"})
dataset.save_episode()
dataset.finalize()
assert dataset[0]["state"].shape == torch.Size([2, 4, 3, 5])
@@ -301,6 +303,7 @@ def test_add_frame_state_5d(tmp_path, empty_lerobot_dataset_factory):
dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
dataset.add_frame({"state": torch.randn(2, 4, 3, 5, 1), "task": "Dummy task"})
dataset.save_episode()
dataset.finalize()
assert dataset[0]["state"].shape == torch.Size([2, 4, 3, 5, 1])
@@ -310,6 +313,7 @@ def test_add_frame_state_numpy(tmp_path, empty_lerobot_dataset_factory):
dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
dataset.add_frame({"state": np.array([1], dtype=np.float32), "task": "Dummy task"})
dataset.save_episode()
dataset.finalize()
assert dataset[0]["state"].ndim == 0
@@ -319,6 +323,7 @@ def test_add_frame_string(tmp_path, empty_lerobot_dataset_factory):
dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
dataset.add_frame({"caption": "Dummy caption", "task": "Dummy task"})
dataset.save_episode()
dataset.finalize()
assert dataset[0]["caption"] == "Dummy caption"
@@ -354,6 +359,7 @@ def test_add_frame_image(image_dataset):
dataset = image_dataset
dataset.add_frame({"image": np.random.rand(*DUMMY_CHW), "task": "Dummy task"})
dataset.save_episode()
dataset.finalize()
assert dataset[0]["image"].shape == torch.Size(DUMMY_CHW)
@@ -362,6 +368,7 @@ def test_add_frame_image_h_w_c(image_dataset):
dataset = image_dataset
dataset.add_frame({"image": np.random.rand(*DUMMY_HWC), "task": "Dummy task"})
dataset.save_episode()
dataset.finalize()
assert dataset[0]["image"].shape == torch.Size(DUMMY_CHW)
@@ -371,6 +378,7 @@ def test_add_frame_image_uint8(image_dataset):
image = np.random.randint(0, 256, DUMMY_HWC, dtype=np.uint8)
dataset.add_frame({"image": image, "task": "Dummy task"})
dataset.save_episode()
dataset.finalize()
assert dataset[0]["image"].shape == torch.Size(DUMMY_CHW)
@@ -380,6 +388,7 @@ def test_add_frame_image_pil(image_dataset):
image = np.random.randint(0, 256, DUMMY_HWC, dtype=np.uint8)
dataset.add_frame({"image": Image.fromarray(image), "task": "Dummy task"})
dataset.save_episode()
dataset.finalize()
assert dataset[0]["image"].shape == torch.Size(DUMMY_CHW)
@@ -400,7 +409,7 @@ def test_tmp_image_deletion(tmp_path, empty_lerobot_dataset_factory):
ds_img = empty_lerobot_dataset_factory(root=tmp_path / "img", features=features_image)
ds_img.add_frame({"image": np.random.rand(*DUMMY_CHW), "task": "Dummy task"})
ds_img.save_episode()
img_dir = ds_img._get_image_file_dir(0, image_key)
img_dir = ds_img.writer._get_image_file_dir(0, image_key)
assert not img_dir.exists(), "Temporary image directory should be removed for image features"
@@ -413,10 +422,10 @@ def test_tmp_video_deletion(tmp_path, empty_lerobot_dataset_factory):
}
ds_vid = empty_lerobot_dataset_factory(root=tmp_path / "vid", features=features_video)
ds_vid.batch_encoding_size = 1
ds_vid.writer._batch_encoding_size = 1
ds_vid.add_frame({vid_key: np.random.rand(*DUMMY_CHW), "task": "Dummy task"})
ds_vid.save_episode()
vid_img_dir = ds_vid._get_image_file_dir(0, vid_key)
vid_img_dir = ds_vid.writer._get_image_file_dir(0, vid_key)
assert not vid_img_dir.exists(), (
"Temporary image directory should be removed when batch_encoding_size == 1"
)
@@ -431,7 +440,7 @@ def test_tmp_mixed_deletion(tmp_path, empty_lerobot_dataset_factory):
vid_key: {"dtype": "video", "shape": DUMMY_HWC, "names": ["height", "width", "channels"]},
}
ds_mixed = empty_lerobot_dataset_factory(
root=tmp_path / "mixed", features=features_mixed, batch_encoding_size=2
root=tmp_path / "mixed", features=features_mixed, batch_encoding_size=2, streaming_encoding=False
)
ds_mixed.add_frame(
{
@@ -441,8 +450,8 @@ def test_tmp_mixed_deletion(tmp_path, empty_lerobot_dataset_factory):
}
)
ds_mixed.save_episode()
img_dir = ds_mixed._get_image_file_dir(0, image_key)
vid_img_dir = ds_mixed._get_image_file_dir(0, vid_key)
img_dir = ds_mixed.writer._get_image_file_dir(0, image_key)
vid_img_dir = ds_mixed.writer._get_image_file_dir(0, vid_key)
assert not img_dir.exists(), "Temporary image directory should be removed for image features"
assert vid_img_dir.exists(), (
"Temporary image directory should not be removed for video features when batch_encoding_size == 2"
@@ -503,7 +512,7 @@ def test_add_frame_audio_file(audio_dataset):
)
# Create the audio file that should be created in the background by the Microphone class
for audio_key in dataset.meta.audio_keys:
fpath = dataset._get_raw_audio_file_path(0, audio_key)
fpath = dataset.writer._get_raw_audio_file_path(0, audio_key)
fpath.parent.mkdir(parents=True, exist_ok=True)
write(
fpath,
@@ -748,29 +757,29 @@ def test_check_cached_episodes_sufficient(tmp_path, lerobot_dataset_factory):
)
# Test hf_dataset is None
dataset.hf_dataset = None
assert dataset._check_cached_episodes_sufficient() is False
dataset.reader.hf_dataset = None
assert dataset.reader._check_cached_episodes_sufficient() is False
# Test hf_dataset is empty
import datasets
empty_features = get_hf_features_from_features(dataset.features)
dataset.hf_dataset = datasets.Dataset.from_dict(
dataset.reader.hf_dataset = datasets.Dataset.from_dict(
{key: [] for key in empty_features}, features=empty_features
)
dataset.hf_dataset.set_transform(hf_transform_to_torch)
assert dataset._check_cached_episodes_sufficient() is False
dataset.reader.hf_dataset.set_transform(hf_transform_to_torch)
assert dataset.reader._check_cached_episodes_sufficient() is False
# Restore the original dataset for remaining tests
dataset.hf_dataset = dataset.load_hf_dataset()
dataset.reader.hf_dataset = dataset.reader._load_hf_dataset()
# Test all episodes requested (self.episodes = None) and all are available
dataset.episodes = None
assert dataset._check_cached_episodes_sufficient() is True
dataset.reader.episodes = None
assert dataset.reader._check_cached_episodes_sufficient() is True
# Test specific episodes requested that are all available
dataset.episodes = [0, 2, 4]
assert dataset._check_cached_episodes_sufficient() is True
dataset.reader.episodes = [0, 2, 4]
assert dataset.reader._check_cached_episodes_sufficient() is True
# Test request episodes that don't exist in the cached dataset
# Create a dataset with only episodes 0, 1, 2
@@ -782,8 +791,8 @@ def test_check_cached_episodes_sufficient(tmp_path, lerobot_dataset_factory):
)
# Request episodes that include non-existent ones
limited_dataset.episodes = [0, 1, 2, 3, 4]
assert limited_dataset._check_cached_episodes_sufficient() is False
limited_dataset.reader.episodes = [0, 1, 2, 3, 4]
assert limited_dataset.reader._check_cached_episodes_sufficient() is False
# Test create a dataset with sparse episodes (e.g., only episodes 0, 2, 4)
# First create the full dataset structure
@@ -819,22 +828,22 @@ def test_check_cached_episodes_sufficient(tmp_path, lerobot_dataset_factory):
filtered_data[key] = filtered_values
sparse_dataset.hf_dataset = datasets.Dataset.from_dict(
sparse_dataset.reader.hf_dataset = datasets.Dataset.from_dict(
filtered_data, features=get_hf_features_from_features(sparse_dataset.features)
)
sparse_dataset.hf_dataset.set_transform(hf_transform_to_torch)
sparse_dataset.reader.hf_dataset.set_transform(hf_transform_to_torch)
# Test requesting all episodes when only some are cached
sparse_dataset.episodes = None
assert sparse_dataset._check_cached_episodes_sufficient() is False
sparse_dataset.reader.episodes = None
assert sparse_dataset.reader._check_cached_episodes_sufficient() is False
# Test requesting only the available episodes
sparse_dataset.episodes = [0, 2, 4]
assert sparse_dataset._check_cached_episodes_sufficient() is True
sparse_dataset.reader.episodes = [0, 2, 4]
assert sparse_dataset.reader._check_cached_episodes_sufficient() is True
# Test requesting a mix of available and unavailable episodes
sparse_dataset.episodes = [0, 1, 2]
assert sparse_dataset._check_cached_episodes_sufficient() is False
sparse_dataset.reader.episodes = [0, 1, 2]
assert sparse_dataset.reader._check_cached_episodes_sufficient() is False
def test_update_chunk_settings(tmp_path, empty_lerobot_dataset_factory):
@@ -1306,13 +1315,13 @@ def test_dataset_resume_recording(tmp_path, empty_lerobot_dataset_factory):
del dataset_verify
# Phase 3: Resume recording - add more episodes
dataset_resumed = LeRobotDataset(initial_repo_id, root=initial_root, revision="v3.0")
dataset_resumed = LeRobotDataset.resume(initial_repo_id, root=initial_root, revision="v3.0")
assert dataset_resumed.meta.total_episodes == initial_episodes
assert dataset_resumed.meta.total_frames == initial_episodes * frames_per_episode
assert dataset_resumed.latest_episode is None # Not recording yet
assert dataset_resumed.writer is None
assert dataset_resumed.meta.writer is None
assert dataset_resumed.writer._latest_episode is None # Not recording yet
assert dataset_resumed.writer._pq_writer is None
assert dataset_resumed.meta._pq_writer is None
additional_episodes = 2
for ep_idx in range(initial_episodes, initial_episodes + additional_episodes):
@@ -1388,7 +1397,7 @@ def test_frames_in_current_file_calculation(tmp_path, empty_lerobot_dataset_fact
dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features, use_videos=False)
dataset.meta.update_chunk_settings(data_files_size_in_mb=100)
assert dataset._current_file_start_frame is None
assert dataset.writer._current_file_start_frame is None
frames_per_episode = 10
for _ in range(frames_per_episode):
@@ -1401,7 +1410,7 @@ def test_frames_in_current_file_calculation(tmp_path, empty_lerobot_dataset_fact
)
dataset.save_episode()
assert dataset._current_file_start_frame == 0
assert dataset.writer._current_file_start_frame == 0
assert dataset.meta.total_episodes == 1
assert dataset.meta.total_frames == frames_per_episode
@@ -1415,12 +1424,12 @@ def test_frames_in_current_file_calculation(tmp_path, empty_lerobot_dataset_fact
)
dataset.save_episode()
assert dataset._current_file_start_frame == 0
assert dataset.writer._current_file_start_frame == 0
assert dataset.meta.total_episodes == 2
assert dataset.meta.total_frames == 2 * frames_per_episode
ep1_chunk = dataset.latest_episode["data/chunk_index"]
ep1_file = dataset.latest_episode["data/file_index"]
ep1_chunk = dataset.writer._latest_episode["data/chunk_index"]
ep1_file = dataset.writer._latest_episode["data/file_index"]
assert ep1_chunk == 0
assert ep1_file == 0
@@ -1434,18 +1443,18 @@ def test_frames_in_current_file_calculation(tmp_path, empty_lerobot_dataset_fact
)
dataset.save_episode()
assert dataset._current_file_start_frame == 0
assert dataset.writer._current_file_start_frame == 0
assert dataset.meta.total_episodes == 3
assert dataset.meta.total_frames == 3 * frames_per_episode
ep2_chunk = dataset.latest_episode["data/chunk_index"]
ep2_file = dataset.latest_episode["data/file_index"]
ep2_chunk = dataset.writer._latest_episode["data/chunk_index"]
ep2_file = dataset.writer._latest_episode["data/file_index"]
assert ep2_chunk == 0
assert ep2_file == 0
dataset.finalize()
from lerobot.datasets.utils import load_episodes
from lerobot.datasets.io_utils import load_episodes
dataset.meta.episodes = load_episodes(dataset.root)
assert dataset.meta.episodes is not None
@@ -1471,82 +1480,6 @@ def test_frames_in_current_file_calculation(tmp_path, empty_lerobot_dataset_fact
assert frame["episode_index"].item() == expected_ep
def test_encode_video_worker_forwards_vcodec(tmp_path):
"""Test that _encode_video_worker correctly forwards the vcodec parameter to encode_video_frames."""
from unittest.mock import patch
from lerobot.datasets.utils import DEFAULT_IMAGE_PATH
# Create the expected directory structure
video_key = "observation.images.laptop"
episode_index = 0
frame_index = 0
fpath = DEFAULT_IMAGE_PATH.format(
image_key=video_key, episode_index=episode_index, frame_index=frame_index
)
img_dir = tmp_path / Path(fpath).parent
img_dir.mkdir(parents=True, exist_ok=True)
# Create a dummy image file
dummy_img = Image.new("RGB", (64, 64), color="red")
dummy_img.save(img_dir / "frame-000000.png")
# Track what vcodec was passed to encode_video_frames
captured_kwargs = {}
def mock_encode_video_frames(imgs_dir, video_path, fps, **kwargs):
captured_kwargs.update(kwargs)
# Create a dummy output file so the worker doesn't fail
Path(video_path).parent.mkdir(parents=True, exist_ok=True)
Path(video_path).touch()
with patch("lerobot.datasets.lerobot_dataset.encode_video_frames", side_effect=mock_encode_video_frames):
# Test with h264 codec
_encode_video_worker(video_key, episode_index, tmp_path, fps=30, vcodec="h264")
assert "vcodec" in captured_kwargs
assert captured_kwargs["vcodec"] == "h264"
def test_encode_video_worker_default_vcodec(tmp_path):
"""Test that _encode_video_worker uses libsvtav1 as the default codec."""
from unittest.mock import patch
from lerobot.datasets.utils import DEFAULT_IMAGE_PATH
# Create the expected directory structure
video_key = "observation.images.laptop"
episode_index = 0
frame_index = 0
fpath = DEFAULT_IMAGE_PATH.format(
image_key=video_key, episode_index=episode_index, frame_index=frame_index
)
img_dir = tmp_path / Path(fpath).parent
img_dir.mkdir(parents=True, exist_ok=True)
# Create a dummy image file
dummy_img = Image.new("RGB", (64, 64), color="red")
dummy_img.save(img_dir / "frame-000000.png")
# Track what vcodec was passed to encode_video_frames
captured_kwargs = {}
def mock_encode_video_frames(imgs_dir, video_path, fps, **kwargs):
captured_kwargs.update(kwargs)
# Create a dummy output file so the worker doesn't fail
Path(video_path).parent.mkdir(parents=True, exist_ok=True)
Path(video_path).touch()
with patch("lerobot.datasets.lerobot_dataset.encode_video_frames", side_effect=mock_encode_video_frames):
# Test with default codec (no vcodec specified)
_encode_video_worker(video_key, episode_index, tmp_path, fps=30)
assert "vcodec" in captured_kwargs
assert captured_kwargs["vcodec"] == "libsvtav1"
def test_lerobot_dataset_vcodec_validation():
"""Test that LeRobotDataset validates the vcodec parameter."""
# Test that invalid vcodec raises ValueError
@@ -1566,7 +1499,10 @@ def test_valid_video_codecs_constant():
assert "h264" in VALID_VIDEO_CODECS
assert "hevc" in VALID_VIDEO_CODECS
assert "libsvtav1" in VALID_VIDEO_CODECS
assert len(VALID_VIDEO_CODECS) == 3
assert "auto" in VALID_VIDEO_CODECS
assert "h264_videotoolbox" in VALID_VIDEO_CODECS
assert "h264_nvenc" in VALID_VIDEO_CODECS
assert len(VALID_VIDEO_CODECS) == 10
def test_delta_timestamps_with_episodes_filter(tmp_path, empty_lerobot_dataset_factory):
tests/datasets/test_delta_timestamps.py
+1 -1
View File
@@ -13,7 +13,7 @@
# limitations under the License.
import pytest
from lerobot.datasets.utils import (
from lerobot.datasets.feature_utils import (
check_delta_timestamps,
get_delta_indices,
)
tests/datasets/test_image_transforms.py
+24
View File
@@ -390,6 +390,30 @@ def test_sharpness_jitter_invalid_range_max_smaller():
SharpnessJitter((2.0, 0.1))
def test_make_transform_from_config_with_v2_resize(img_tensor_factory):
img_tensor = img_tensor_factory()
tf_cfg = ImageTransformConfig(type="Resize", kwargs={"size": (32, 32)})
tf = make_transform_from_config(tf_cfg)
assert isinstance(tf, v2.Resize)
output = tf(img_tensor)
assert output.shape[-2:] == (32, 32)
def test_make_transform_from_config_with_v2_identity(img_tensor_factory):
img_tensor = img_tensor_factory()
tf_cfg = ImageTransformConfig(type="Identity", kwargs={})
tf = make_transform_from_config(tf_cfg)
assert isinstance(tf, v2.Identity)
output = tf(img_tensor)
assert output.shape == img_tensor.shape
def test_make_transform_from_config_invalid_type():
tf_cfg = ImageTransformConfig(type="NotARealTransform", kwargs={})
with pytest.raises(ValueError, match="not valid"):
make_transform_from_config(tf_cfg)
def test_save_all_transforms(img_tensor_factory, tmp_path):
img_tensor = img_tensor_factory()
tf_cfg = ImageTransformsConfig(enable=True)
tests/datasets/test_image_writer.py
+10 -6
View File
@@ -142,9 +142,9 @@ def test_write_image_image(tmp_path, img_factory):
def test_write_image_exception(tmp_path):
image_array = "invalid data"
fpath = tmp_path / DUMMY_IMAGE
with patch("builtins.print") as mock_print:
with patch("lerobot.datasets.image_writer.logger") as mock_logger:
write_image(image_array, fpath)
mock_print.assert_called()
mock_logger.error.assert_called()
assert not fpath.exists()
@@ -243,10 +243,10 @@ def test_save_image_invalid_data(tmp_path):
image_array = "invalid data"
fpath = tmp_path / DUMMY_IMAGE
fpath.parent.mkdir(parents=True, exist_ok=True)
with patch("builtins.print") as mock_print:
with patch("lerobot.datasets.image_writer.logger") as mock_logger:
writer.save_image(image_array, fpath)
writer.wait_until_done()
mock_print.assert_called()
mock_logger.error.assert_called()
assert not fpath.exists()
finally:
writer.stop()
@@ -352,10 +352,14 @@ def test_with_different_image_formats(tmp_path, img_array_factory):
def test_safe_stop_image_writer_decorator():
class MockDataset:
class MockWriter:
def __init__(self):
self.image_writer = MagicMock(spec=AsyncImageWriter)
class MockDataset:
def __init__(self):
self.writer = MockWriter()
@safe_stop_image_writer
def function_that_raises_exception(dataset=None):
raise Exception("Test exception")
@@ -366,7 +370,7 @@ def test_safe_stop_image_writer_decorator():
function_that_raises_exception(dataset=dataset)
assert str(exc_info.value) == "Test exception"
dataset.image_writer.stop.assert_called_once()
dataset.writer.image_writer.stop.assert_called_once()
def test_main_process_time(tmp_path, img_tensor_factory):
tests/datasets/test_lerobot_dataset.py
+632
View File
@@ -0,0 +1,632 @@
#!/usr/bin/env python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Contract tests for the LeRobotDataset facade.
Tests focus on mode contracts (read-only, write-only, resume), guards,
property delegation, and the full create-record-finalize-read lifecycle.
"""
from pathlib import Path
from unittest.mock import Mock
import pytest
import torch
import lerobot.datasets.dataset_metadata as dataset_metadata_module
import lerobot.datasets.lerobot_dataset as lerobot_dataset_module
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
from lerobot.datasets.dataset_reader import DatasetReader
from lerobot.datasets.dataset_writer import DatasetWriter
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from tests.fixtures.constants import DEFAULT_FPS, DUMMY_REPO_ID
SIMPLE_FEATURES = {
"state": {"dtype": "float32", "shape": (2,), "names": None},
}
SNAPSHOT_MAIN_FEATURES = {
**SIMPLE_FEATURES,
"test": {"dtype": "float32", "shape": (2,), "names": None},
}
def _make_frame(task: str = "Dummy task") -> dict:
return {"task": task, "state": torch.randn(2)}
def _set_default_cache_root(monkeypatch: pytest.MonkeyPatch, cache_root: Path) -> None:
monkeypatch.setattr(dataset_metadata_module, "HF_LEROBOT_HOME", cache_root)
monkeypatch.setattr(dataset_metadata_module, "HF_LEROBOT_HUB_CACHE", cache_root / "hub")
monkeypatch.setattr(lerobot_dataset_module, "HF_LEROBOT_HUB_CACHE", cache_root / "hub")
def _write_dataset_tree(
root: Path,
*,
motor_features: dict[str, dict],
info_factory,
stats_factory,
tasks_factory,
episodes_factory,
hf_dataset_factory,
create_info,
create_stats,
create_tasks,
create_episodes,
create_hf_dataset,
) -> None:
root.mkdir(parents=True, exist_ok=True)
info = info_factory(
total_episodes=1,
total_frames=3,
total_tasks=1,
use_videos=False,
motor_features=motor_features,
camera_features={},
)
tasks = tasks_factory(total_tasks=1)
episodes = episodes_factory(
features=info["features"],
fps=info["fps"],
total_episodes=1,
total_frames=3,
tasks=tasks,
)
stats = stats_factory(features=info["features"])
hf_dataset = hf_dataset_factory(
features=info["features"],
tasks=tasks,
episodes=episodes,
fps=info["fps"],
)
create_info(root, info)
create_stats(root, stats)
create_tasks(root, tasks)
create_episodes(root, episodes)
create_hf_dataset(root, hf_dataset)
# ── Read-only mode (via __init__) ────────────────────────────────────
def test_init_creates_reader_no_writer(tmp_path, lerobot_dataset_factory):
"""__init__() sets reader to a DatasetReader and writer to None."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=1, total_frames=10, use_videos=False
)
assert isinstance(dataset.reader, DatasetReader)
assert dataset.writer is None
def test_init_loads_data(tmp_path, lerobot_dataset_factory):
"""After __init__(), the dataset has data and len > 0."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=1, total_frames=10, use_videos=False
)
assert len(dataset) > 0
def test_getitem_works_in_read_mode(tmp_path, lerobot_dataset_factory):
"""dataset[0] returns a dict with expected keys in read-only mode."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=1, total_frames=10, use_videos=False
)
item = dataset[0]
assert isinstance(item, dict)
assert "index" in item
assert "task" in item
def test_len_matches_num_frames(tmp_path, lerobot_dataset_factory):
"""len(dataset) equals dataset.num_frames."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=30, use_videos=False
)
assert len(dataset) == dataset.num_frames
def test_metadata_without_root_uses_hub_cache_snapshot_download(
tmp_path,
info_factory,
stats_factory,
tasks_factory,
episodes_factory,
hf_dataset_factory,
create_info,
create_stats,
create_tasks,
create_episodes,
create_hf_dataset,
monkeypatch,
):
"""Metadata refresh uses the dedicated Hub cache instead of a shared local_dir mirror."""
repo_id = DUMMY_REPO_ID
cache_root = tmp_path / "lerobot_cache"
snapshot_root = cache_root / "hub" / "datasets--dummy--repo" / "snapshots" / "commit-main"
_write_dataset_tree(
snapshot_root,
motor_features=SNAPSHOT_MAIN_FEATURES,
info_factory=info_factory,
stats_factory=stats_factory,
tasks_factory=tasks_factory,
episodes_factory=episodes_factory,
hf_dataset_factory=hf_dataset_factory,
create_info=create_info,
create_stats=create_stats,
create_tasks=create_tasks,
create_episodes=create_episodes,
create_hf_dataset=create_hf_dataset,
)
_set_default_cache_root(monkeypatch, cache_root)
snapshot_download = Mock(return_value=str(snapshot_root))
monkeypatch.setattr(dataset_metadata_module, "snapshot_download", snapshot_download)
meta = LeRobotDatasetMetadata(repo_id=repo_id, revision="main", force_cache_sync=True)
assert meta.root == snapshot_root
assert snapshot_download.call_count == 1
assert snapshot_download.call_args.args == (repo_id,)
assert snapshot_download.call_args.kwargs == {
"repo_type": "dataset",
"revision": "main",
"cache_dir": cache_root / "hub",
"allow_patterns": "meta/",
"ignore_patterns": None,
}
def test_without_root_reads_different_revisions_from_distinct_snapshot_roots(
tmp_path,
info_factory,
stats_factory,
tasks_factory,
episodes_factory,
hf_dataset_factory,
create_info,
create_stats,
create_tasks,
create_episodes,
create_hf_dataset,
monkeypatch,
):
"""Different revisions resolve to different on-disk snapshot roots."""
repo_id = DUMMY_REPO_ID
old_revision = "b59010db93eb6cc3cf06ef2f7cae1bbe62b726d9"
cache_root = tmp_path / "lerobot_cache"
main_root = cache_root / "hub" / "datasets--dummy--repo" / "snapshots" / "commit-main"
old_root = cache_root / "hub" / "datasets--dummy--repo" / "snapshots" / "commit-old"
_write_dataset_tree(
main_root,
motor_features=SNAPSHOT_MAIN_FEATURES,
info_factory=info_factory,
stats_factory=stats_factory,
tasks_factory=tasks_factory,
episodes_factory=episodes_factory,
hf_dataset_factory=hf_dataset_factory,
create_info=create_info,
create_stats=create_stats,
create_tasks=create_tasks,
create_episodes=create_episodes,
create_hf_dataset=create_hf_dataset,
)
_write_dataset_tree(
old_root,
motor_features=SIMPLE_FEATURES,
info_factory=info_factory,
stats_factory=stats_factory,
tasks_factory=tasks_factory,
episodes_factory=episodes_factory,
hf_dataset_factory=hf_dataset_factory,
create_info=create_info,
create_stats=create_stats,
create_tasks=create_tasks,
create_episodes=create_episodes,
create_hf_dataset=create_hf_dataset,
)
_set_default_cache_root(monkeypatch, cache_root)
snapshot_roots = {
"main": main_root,
old_revision: old_root,
}
meta_snapshot_download = Mock(
side_effect=lambda repo_id, **kwargs: str(snapshot_roots[kwargs["revision"]])
)
data_snapshot_download = Mock(
side_effect=lambda repo_id, **kwargs: str(snapshot_roots[kwargs["revision"]])
)
monkeypatch.setattr(dataset_metadata_module, "snapshot_download", meta_snapshot_download)
monkeypatch.setattr(lerobot_dataset_module, "snapshot_download", data_snapshot_download)
main_dataset = LeRobotDataset(
repo_id=repo_id, revision="main", download_videos=False, force_cache_sync=True
)
old_dataset = LeRobotDataset(
repo_id=repo_id, revision=old_revision, download_videos=False, force_cache_sync=True
)
assert main_dataset.root == main_root
assert old_dataset.root == old_root
assert "test" in main_dataset.hf_dataset.column_names
assert "test" not in old_dataset.hf_dataset.column_names
# Metadata downloads use cache_dir, not local_dir
assert meta_snapshot_download.call_count == 2
for download_call in meta_snapshot_download.call_args_list:
assert download_call.kwargs["cache_dir"] == cache_root / "hub"
assert "local_dir" not in download_call.kwargs
# Data downloads also use cache_dir, not local_dir
assert data_snapshot_download.call_count == 2
for download_call in data_snapshot_download.call_args_list:
assert download_call.kwargs["cache_dir"] == cache_root / "hub"
assert "local_dir" not in download_call.kwargs
def test_metadata_without_root_ignores_legacy_local_dir_cache(
tmp_path,
info_factory,
stats_factory,
tasks_factory,
episodes_factory,
hf_dataset_factory,
create_info,
create_stats,
create_tasks,
create_episodes,
create_hf_dataset,
monkeypatch,
):
"""Legacy local-dir mirrors are bypassed in favor of revision-safe snapshots."""
repo_id = DUMMY_REPO_ID
cache_root = tmp_path / "lerobot_cache"
legacy_root = cache_root / repo_id
snapshot_root = cache_root / "hub" / "datasets--dummy--repo" / "snapshots" / "commit-main"
_write_dataset_tree(
legacy_root,
motor_features=SIMPLE_FEATURES,
info_factory=info_factory,
stats_factory=stats_factory,
tasks_factory=tasks_factory,
episodes_factory=episodes_factory,
hf_dataset_factory=hf_dataset_factory,
create_info=create_info,
create_stats=create_stats,
create_tasks=create_tasks,
create_episodes=create_episodes,
create_hf_dataset=create_hf_dataset,
)
(legacy_root / ".cache" / "huggingface" / "download").mkdir(parents=True, exist_ok=True)
_write_dataset_tree(
snapshot_root,
motor_features=SNAPSHOT_MAIN_FEATURES,
info_factory=info_factory,
stats_factory=stats_factory,
tasks_factory=tasks_factory,
episodes_factory=episodes_factory,
hf_dataset_factory=hf_dataset_factory,
create_info=create_info,
create_stats=create_stats,
create_tasks=create_tasks,
create_episodes=create_episodes,
create_hf_dataset=create_hf_dataset,
)
_set_default_cache_root(monkeypatch, cache_root)
snapshot_download = Mock(return_value=str(snapshot_root))
monkeypatch.setattr(dataset_metadata_module, "snapshot_download", snapshot_download)
meta = LeRobotDatasetMetadata(repo_id=repo_id, revision="main")
assert meta.root == snapshot_root
assert "test" in meta.features
assert snapshot_download.call_count == 1
def test_download_without_root_uses_hub_cache(
tmp_path,
info_factory,
stats_factory,
tasks_factory,
episodes_factory,
hf_dataset_factory,
create_info,
create_stats,
create_tasks,
create_episodes,
create_hf_dataset,
monkeypatch,
):
"""LeRobotDataset._download() uses cache_dir (not local_dir) when root is not provided."""
repo_id = DUMMY_REPO_ID
cache_root = tmp_path / "lerobot_cache"
snapshot_root = cache_root / "hub" / "datasets--dummy--repo" / "snapshots" / "commit-main"
# Pre-populate snapshot directory so metadata loads succeed, but leave
# data absent so that _download() is triggered.
_write_dataset_tree(
snapshot_root,
motor_features=SIMPLE_FEATURES,
info_factory=info_factory,
stats_factory=stats_factory,
tasks_factory=tasks_factory,
episodes_factory=episodes_factory,
hf_dataset_factory=hf_dataset_factory,
create_info=create_info,
create_stats=create_stats,
create_tasks=create_tasks,
create_episodes=create_episodes,
create_hf_dataset=create_hf_dataset,
)
_set_default_cache_root(monkeypatch, cache_root)
meta_snapshot_download = Mock(return_value=str(snapshot_root))
monkeypatch.setattr(dataset_metadata_module, "snapshot_download", meta_snapshot_download)
# Mock the data snapshot_download to return the same root (data already
# exists there from _write_dataset_tree).
data_snapshot_download = Mock(return_value=str(snapshot_root))
monkeypatch.setattr(lerobot_dataset_module, "snapshot_download", data_snapshot_download)
LeRobotDataset(repo_id=repo_id, revision="main", force_cache_sync=True)
# _download() should have called snapshot_download with cache_dir
assert data_snapshot_download.call_count == 1
call_kwargs = data_snapshot_download.call_args.kwargs
assert call_kwargs["cache_dir"] == cache_root / "hub"
assert "local_dir" not in call_kwargs
# ── Write-only mode (via create()) ──────────────────────────────────
def test_create_sets_writer_no_reader(tmp_path):
"""create() sets writer to a DatasetWriter and reader to None."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
assert isinstance(dataset.writer, DatasetWriter)
assert dataset.reader is None
def test_create_initial_counts_zero(tmp_path):
"""After create(), num_episodes == 0 and num_frames == 0."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
assert dataset.num_episodes == 0
assert dataset.num_frames == 0
def test_add_frame_works_in_write_mode(tmp_path):
"""add_frame() succeeds on a dataset created via create()."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
dataset.add_frame(_make_frame()) # should not raise
# ── Resume mode ──────────────────────────────────────────────────────
def test_resume_creates_writer(tmp_path):
"""After resume(), writer is a DatasetWriter."""
root = tmp_path / "resume_ds"
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root
)
for _ in range(3):
dataset.add_frame(_make_frame())
dataset.save_episode()
dataset.finalize()
resumed = LeRobotDataset.resume(repo_id=DUMMY_REPO_ID, root=root)
assert isinstance(resumed.writer, DatasetWriter)
def test_resume_preserves_episode_count(tmp_path):
"""After resume(), existing episodes are counted."""
root = tmp_path / "resume_ds"
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root
)
for _ in range(3):
dataset.add_frame(_make_frame())
dataset.save_episode()
dataset.finalize()
resumed = LeRobotDataset.resume(repo_id=DUMMY_REPO_ID, root=root)
assert resumed.meta.total_episodes == 1
def test_resume_can_add_more_episodes(tmp_path):
"""After resume(), new episodes can be added."""
root = tmp_path / "resume_ds"
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root
)
for _ in range(3):
dataset.add_frame(_make_frame())
dataset.save_episode()
dataset.finalize()
resumed = LeRobotDataset.resume(repo_id=DUMMY_REPO_ID, root=root)
for _ in range(2):
resumed.add_frame(_make_frame())
resumed.save_episode()
assert resumed.meta.total_episodes == 2
# ── Writer guard ─────────────────────────────────────────────────────
def test_add_frame_raises_without_writer(tmp_path, lerobot_dataset_factory):
"""add_frame() raises RuntimeError on a read-only dataset."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=1, total_frames=5, use_videos=False
)
with pytest.raises(RuntimeError, match="read-only"):
dataset.add_frame(_make_frame())
def test_save_episode_raises_without_writer(tmp_path, lerobot_dataset_factory):
"""save_episode() raises RuntimeError on a read-only dataset."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=1, total_frames=5, use_videos=False
)
with pytest.raises(RuntimeError, match="read-only"):
dataset.save_episode()
def test_clear_episode_buffer_raises_without_writer(tmp_path, lerobot_dataset_factory):
"""clear_episode_buffer() raises RuntimeError on a read-only dataset."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=1, total_frames=5, use_videos=False
)
with pytest.raises(RuntimeError, match="read-only"):
dataset.clear_episode_buffer()
# ── Reader guard ─────────────────────────────────────────────────────
def test_getitem_raises_before_finalize(tmp_path):
"""dataset[0] raises RuntimeError while recording (before finalize)."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
for _ in range(3):
dataset.add_frame(_make_frame())
dataset.save_episode()
with pytest.raises(RuntimeError, match="finalize"):
dataset[0]
def test_getitem_works_after_finalize(tmp_path):
"""After finalize(), dataset[0] returns data."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
for _ in range(3):
dataset.add_frame(_make_frame())
dataset.save_episode()
dataset.finalize()
item = dataset[0]
assert "state" in item
assert "task" in item
# ── Property delegation ──────────────────────────────────────────────
def test_fps_delegates_to_meta(tmp_path, lerobot_dataset_factory):
"""dataset.fps == dataset.meta.fps."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=1, total_frames=5, use_videos=False
)
assert dataset.fps == dataset.meta.fps
def test_features_delegates_to_meta(tmp_path, lerobot_dataset_factory):
"""dataset.features is dataset.meta.features."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=1, total_frames=5, use_videos=False
)
assert dataset.features is dataset.meta.features
def test_num_frames_uses_meta_in_write_mode(tmp_path):
"""In write-only mode (reader=None), num_frames comes from metadata."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
assert dataset.reader is None
assert dataset.num_frames == dataset.meta.total_frames
# ── Lifecycle ────────────────────────────────────────────────────────
def test_finalize_is_idempotent(tmp_path):
"""Calling finalize() twice does not raise."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
dataset.finalize()
dataset.finalize()
def test_has_pending_frames_lifecycle(tmp_path):
"""has_pending_frames: False -> True (add_frame) -> False (save_episode)."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
assert dataset.has_pending_frames() is False
dataset.add_frame(_make_frame())
assert dataset.has_pending_frames() is True
dataset.save_episode()
assert dataset.has_pending_frames() is False
def test_create_record_finalize_read_roundtrip(tmp_path):
"""End-to-end: create, record 2 episodes, finalize, re-open, verify data."""
root = tmp_path / "roundtrip"
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root
)
# Episode 0: 3 frames with known values
ep0_states = []
for i in range(3):
state = torch.tensor([float(i), float(i * 2)])
ep0_states.append(state)
dataset.add_frame({"task": "Task A", "state": state})
dataset.save_episode()
# Episode 1: 2 frames
ep1_states = []
for i in range(2):
state = torch.tensor([float(i + 100), float(i + 200)])
ep1_states.append(state)
dataset.add_frame({"task": "Task B", "state": state})
dataset.save_episode()
dataset.finalize()
# Re-open as read-only
reopened = LeRobotDataset(repo_id=DUMMY_REPO_ID, root=root)
assert len(reopened) == 5
assert reopened.num_episodes == 2
# Verify episode 0
for i in range(3):
item = reopened[i]
assert torch.allclose(item["state"], ep0_states[i], atol=1e-5)
assert item["episode_index"].item() == 0
# Verify episode 1
for i in range(2):
item = reopened[3 + i]
assert torch.allclose(item["state"], ep1_states[i], atol=1e-5)
assert item["episode_index"].item() == 1
tests/datasets/test_online_buffer.py
-282
View File
@@ -1,282 +0,0 @@
#!/usr/bin/env python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.d
from copy import deepcopy
from uuid import uuid4
import numpy as np
import pytest
import torch
from lerobot.datasets.online_buffer import OnlineBuffer, compute_sampler_weights
# Some constants for OnlineBuffer tests.
data_key = "data"
data_shape = (2, 3) # just some arbitrary > 1D shape
buffer_capacity = 100
fps = 10
def make_new_buffer(
write_dir: str | None = None, delta_timestamps: dict[str, list[float]] | None = None
) -> tuple[OnlineBuffer, str]:
if write_dir is None:
write_dir = f"/tmp/online_buffer_{uuid4().hex}"
buffer = OnlineBuffer(
write_dir,
data_spec={data_key: {"shape": data_shape, "dtype": np.dtype("float32")}},
buffer_capacity=buffer_capacity,
fps=fps,
delta_timestamps=delta_timestamps,
)
return buffer, write_dir
def make_spoof_data_frames(n_episodes: int, n_frames_per_episode: int) -> dict[str, np.ndarray]:
new_data = {
data_key: np.arange(n_frames_per_episode * n_episodes * np.prod(data_shape)).reshape(-1, *data_shape),
OnlineBuffer.INDEX_KEY: np.arange(n_frames_per_episode * n_episodes),
OnlineBuffer.EPISODE_INDEX_KEY: np.repeat(np.arange(n_episodes), n_frames_per_episode),
OnlineBuffer.FRAME_INDEX_KEY: np.tile(np.arange(n_frames_per_episode), n_episodes),
OnlineBuffer.TIMESTAMP_KEY: np.tile(np.arange(n_frames_per_episode) / fps, n_episodes),
}
return new_data
def test_non_mutate():
"""Checks that the data provided to the add_data method is copied rather than passed by reference.
This means that mutating the data in the buffer does not mutate the original data.
NOTE: If this test fails, it means some of the other tests may be compromised. For example, we can't trust
a success case for `test_write_read`.
"""
buffer, _ = make_new_buffer()
new_data = make_spoof_data_frames(2, buffer_capacity // 4)
new_data_copy = deepcopy(new_data)
buffer.add_data(new_data)
buffer._data[data_key][:] += 1
assert all(np.array_equal(new_data[k], new_data_copy[k]) for k in new_data)
def test_index_error_no_data():
buffer, _ = make_new_buffer()
with pytest.raises(IndexError):
buffer[0]
def test_index_error_with_data():
buffer, _ = make_new_buffer()
n_frames = buffer_capacity // 2
new_data = make_spoof_data_frames(1, n_frames)
buffer.add_data(new_data)
with pytest.raises(IndexError):
buffer[n_frames]
with pytest.raises(IndexError):
buffer[-n_frames - 1]
@pytest.mark.parametrize("do_reload", [False, True])
def test_write_read(do_reload: bool):
"""Checks that data can be added to the buffer and read back.
If do_reload we delete the buffer object and load the buffer back from disk before reading.
"""
buffer, write_dir = make_new_buffer()
n_episodes = 2
n_frames_per_episode = buffer_capacity // 4
new_data = make_spoof_data_frames(n_episodes, n_frames_per_episode)
buffer.add_data(new_data)
if do_reload:
del buffer
buffer, _ = make_new_buffer(write_dir)
assert len(buffer) == n_frames_per_episode * n_episodes
for i, item in enumerate(buffer):
assert all(isinstance(item[k], torch.Tensor) for k in item)
assert np.array_equal(item[data_key].numpy(), new_data[data_key][i])
def test_read_data_key():
"""Tests that data can be added to a buffer and all data for a. specific key can be read back."""
buffer, _ = make_new_buffer()
n_episodes = 2
n_frames_per_episode = buffer_capacity // 4
new_data = make_spoof_data_frames(n_episodes, n_frames_per_episode)
buffer.add_data(new_data)
data_from_buffer = buffer.get_data_by_key(data_key)
assert isinstance(data_from_buffer, torch.Tensor)
assert np.array_equal(data_from_buffer.numpy(), new_data[data_key])
def test_fifo():
"""Checks that if data is added beyond the buffer capacity, we discard the oldest data first."""
buffer, _ = make_new_buffer()
n_frames_per_episode = buffer_capacity // 4
n_episodes = 3
new_data = make_spoof_data_frames(n_episodes, n_frames_per_episode)
buffer.add_data(new_data)
n_more_episodes = 2
# Developer sanity check (in case someone changes the global `buffer_capacity`).
assert (n_episodes + n_more_episodes) * n_frames_per_episode > buffer_capacity, (
"Something went wrong with the test code."
)
more_new_data = make_spoof_data_frames(n_more_episodes, n_frames_per_episode)
buffer.add_data(more_new_data)
assert len(buffer) == buffer_capacity, "The buffer should be full."
expected_data = {}
for k in new_data:
# Concatenate, left-truncate, then roll, to imitate the cyclical FIFO pattern in OnlineBuffer.
expected_data[k] = np.roll(
np.concatenate([new_data[k], more_new_data[k]])[-buffer_capacity:],
shift=len(new_data[k]) + len(more_new_data[k]) - buffer_capacity,
axis=0,
)
for i, item in enumerate(buffer):
assert all(isinstance(item[k], torch.Tensor) for k in item)
assert np.array_equal(item[data_key].numpy(), expected_data[data_key][i])
def test_delta_timestamps_within_tolerance():
"""Check that getting an item with delta_timestamps within tolerance succeeds.
Note: Copied from `test_datasets.py::test_load_previous_and_future_frames_within_tolerance`.
"""
# Sanity check on global fps as we are assuming it is 10 here.
assert fps == 10, "This test assumes fps==10"
buffer, _ = make_new_buffer(delta_timestamps={"index": [-0.2, 0, 0.139]})
new_data = make_spoof_data_frames(n_episodes=1, n_frames_per_episode=5)
buffer.add_data(new_data)
buffer.tolerance_s = 0.04
item = buffer[2]
data, is_pad = item["index"], item[f"index{OnlineBuffer.IS_PAD_POSTFIX}"]
torch.testing.assert_close(data, torch.tensor([0, 2, 3]), msg="Data does not match expected values")
assert not is_pad.any(), "Unexpected padding detected"
def test_delta_timestamps_outside_tolerance_inside_episode_range():
"""Check that getting an item with delta_timestamps outside of tolerance fails.
We expect it to fail if and only if the requested timestamps are within the episode range.
Note: Copied from
`test_datasets.py::test_load_previous_and_future_frames_outside_tolerance_inside_episode_range`
"""
# Sanity check on global fps as we are assuming it is 10 here.
assert fps == 10, "This test assumes fps==10"
buffer, _ = make_new_buffer(delta_timestamps={"index": [-0.2, 0, 0.141]})
new_data = make_spoof_data_frames(n_episodes=1, n_frames_per_episode=5)
buffer.add_data(new_data)
buffer.tolerance_s = 0.04
with pytest.raises(AssertionError):
buffer[2]
def test_delta_timestamps_outside_tolerance_outside_episode_range():
"""Check that copy-padding of timestamps outside of the episode range works.
Note: Copied from
`test_datasets.py::test_load_previous_and_future_frames_outside_tolerance_outside_episode_range`
"""
# Sanity check on global fps as we are assuming it is 10 here.
assert fps == 10, "This test assumes fps==10"
buffer, _ = make_new_buffer(delta_timestamps={"index": [-0.3, -0.24, 0, 0.26, 0.3]})
new_data = make_spoof_data_frames(n_episodes=1, n_frames_per_episode=5)
buffer.add_data(new_data)
buffer.tolerance_s = 0.04
item = buffer[2]
data, is_pad = item["index"], item["index_is_pad"]
assert torch.equal(data, torch.tensor([0, 0, 2, 4, 4])), "Data does not match expected values"
assert torch.equal(is_pad, torch.tensor([True, False, False, True, True])), (
"Padding does not match expected values"
)
# Arbitrarily set small dataset sizes, making sure to have uneven sizes.
@pytest.mark.parametrize("offline_dataset_size", [1, 6])
@pytest.mark.parametrize("online_dataset_size", [0, 4])
@pytest.mark.parametrize("online_sampling_ratio", [0.0, 1.0])
def test_compute_sampler_weights_trivial(
lerobot_dataset_factory,
tmp_path,
offline_dataset_size: int,
online_dataset_size: int,
online_sampling_ratio: float,
):
offline_dataset = lerobot_dataset_factory(tmp_path, total_episodes=1, total_frames=offline_dataset_size)
online_dataset, _ = make_new_buffer()
if online_dataset_size > 0:
online_dataset.add_data(
make_spoof_data_frames(n_episodes=2, n_frames_per_episode=online_dataset_size // 2)
)
weights = compute_sampler_weights(
offline_dataset, online_dataset=online_dataset, online_sampling_ratio=online_sampling_ratio
)
if offline_dataset_size == 0 or online_dataset_size == 0:
expected_weights = torch.ones(offline_dataset_size + online_dataset_size)
elif online_sampling_ratio == 0:
expected_weights = torch.cat([torch.ones(offline_dataset_size), torch.zeros(online_dataset_size)])
elif online_sampling_ratio == 1:
expected_weights = torch.cat([torch.zeros(offline_dataset_size), torch.ones(online_dataset_size)])
expected_weights /= expected_weights.sum()
torch.testing.assert_close(weights, expected_weights)
def test_compute_sampler_weights_nontrivial_ratio(lerobot_dataset_factory, tmp_path):
# Arbitrarily set small dataset sizes, making sure to have uneven sizes.
offline_dataset = lerobot_dataset_factory(tmp_path, total_episodes=1, total_frames=4)
online_dataset, _ = make_new_buffer()
online_dataset.add_data(make_spoof_data_frames(n_episodes=4, n_frames_per_episode=2))
online_sampling_ratio = 0.8
weights = compute_sampler_weights(
offline_dataset, online_dataset=online_dataset, online_sampling_ratio=online_sampling_ratio
)
torch.testing.assert_close(
weights, torch.tensor([0.05, 0.05, 0.05, 0.05, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1])
)
def test_compute_sampler_weights_nontrivial_ratio_and_drop_last_n(lerobot_dataset_factory, tmp_path):
# Arbitrarily set small dataset sizes, making sure to have uneven sizes.
offline_dataset = lerobot_dataset_factory(tmp_path, total_episodes=1, total_frames=4)
online_dataset, _ = make_new_buffer()
online_dataset.add_data(make_spoof_data_frames(n_episodes=4, n_frames_per_episode=2))
weights = compute_sampler_weights(
offline_dataset, online_dataset=online_dataset, online_sampling_ratio=0.8, online_drop_n_last_frames=1
)
torch.testing.assert_close(
weights, torch.tensor([0.05, 0.05, 0.05, 0.05, 0.2, 0.0, 0.2, 0.0, 0.2, 0.0, 0.2, 0.0])
)
def test_compute_sampler_weights_drop_n_last_frames(lerobot_dataset_factory, tmp_path):
"""Note: test copied from test_sampler."""
offline_dataset = lerobot_dataset_factory(tmp_path, total_episodes=1, total_frames=2)
online_dataset, _ = make_new_buffer()
online_dataset.add_data(make_spoof_data_frames(n_episodes=4, n_frames_per_episode=2))
weights = compute_sampler_weights(
offline_dataset,
offline_drop_n_last_frames=1,
online_dataset=online_dataset,
online_sampling_ratio=0.5,
online_drop_n_last_frames=1,
)
torch.testing.assert_close(weights, torch.tensor([0.5, 0, 0.125, 0, 0.125, 0, 0.125, 0, 0.125, 0]))
tests/datasets/test_sampler.py
+47 -3
View File
@@ -13,13 +13,32 @@
# 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 logging
import pytest
import torch
from datasets import Dataset
from lerobot.datasets.push_dataset_to_hub.utils import calculate_episode_data_index
from lerobot.datasets.sampler import EpisodeAwareSampler
from lerobot.datasets.utils import (
from lerobot.datasets.io_utils import (
hf_transform_to_torch,
)
from lerobot.datasets.sampler import EpisodeAwareSampler
def calculate_episode_data_index(hf_dataset: Dataset) -> dict[str, torch.Tensor]:
"""Calculate episode data index for testing. Returns {"from": Tensor, "to": Tensor}."""
episode_data_index: dict[str, list[int]] = {"from": [], "to": []}
current_episode = None
if len(hf_dataset) == 0:
return {"from": torch.tensor([]), "to": torch.tensor([])}
for idx, episode_idx in enumerate(hf_dataset["episode_index"]):
if episode_idx != current_episode:
episode_data_index["from"].append(idx)
if current_episode is not None:
episode_data_index["to"].append(idx)
current_episode = episode_idx
episode_data_index["to"].append(idx + 1)
return {k: torch.tensor(v) for k, v in episode_data_index.items()}
def test_drop_n_first_frames():
@@ -90,3 +109,28 @@ def test_shuffle():
assert sampler.indices == [0, 1, 2, 3, 4, 5]
assert len(sampler) == 6
assert set(sampler) == {0, 1, 2, 3, 4, 5}
def test_negative_drop_first_frames_raises():
with pytest.raises(ValueError, match="drop_n_first_frames must be >= 0"):
EpisodeAwareSampler([0], [10], drop_n_first_frames=-1)
def test_negative_drop_last_frames_raises():
with pytest.raises(ValueError, match="drop_n_last_frames must be >= 0"):
EpisodeAwareSampler([0], [10], drop_n_last_frames=-1)
def test_all_episodes_dropped_raises():
# All episodes have 1 frame, drop_n_first_frames=1 removes all
with pytest.raises(ValueError, match="No valid frames remain"):
EpisodeAwareSampler([0, 1, 2], [1, 2, 3], drop_n_first_frames=1)
def test_partial_episode_drop_warns(caplog):
# Episode 0: 1 frame (dropped), Episode 1: 5 frames (kept)
with caplog.at_level(logging.WARNING, logger="lerobot.datasets.sampler"):
sampler = EpisodeAwareSampler([0, 1], [1, 6], drop_n_first_frames=1)
# Episode 0 is skipped (1 frame, drop 1), Episode 1 keeps frames 2-5
assert sampler.indices == [2, 3, 4, 5]
assert "Episode 0" in caplog.text
tests/datasets/test_streaming_video_encoder.py
+730
View File
@@ -0,0 +1,730 @@
#!/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 streaming video encoding and hardware-accelerated encoding."""
import queue
import threading
from unittest.mock import patch
import av
import numpy as np
import pytest
from lerobot.datasets.video_utils import (
VALID_VIDEO_CODECS,
StreamingVideoEncoder,
_CameraEncoderThread,
_get_codec_options,
detect_available_hw_encoders,
resolve_vcodec,
)
from lerobot.utils.constants import OBS_IMAGES
# ─── _get_codec_options tests ───
class TestGetCodecOptions:
def test_libsvtav1_defaults(self):
opts = _get_codec_options("libsvtav1")
assert opts["g"] == "2"
assert opts["crf"] == "30"
assert opts["preset"] == "12"
def test_libsvtav1_custom_preset(self):
opts = _get_codec_options("libsvtav1", preset=8)
assert opts["preset"] == "8"
def test_h264_options(self):
opts = _get_codec_options("h264", g=10, crf=23)
assert opts["g"] == "10"
assert opts["crf"] == "23"
assert "preset" not in opts
def test_videotoolbox_options(self):
opts = _get_codec_options("h264_videotoolbox", g=2, crf=30)
assert opts["g"] == "2"
# CRF 30 maps to quality = max(1, min(100, 100 - 30*2)) = 40
assert opts["q:v"] == "40"
assert "crf" not in opts
def test_nvenc_options(self):
opts = _get_codec_options("h264_nvenc", g=2, crf=25)
assert opts["rc"] == "constqp"
assert opts["qp"] == "25"
assert "crf" not in opts
# NVENC doesn't support g
assert "g" not in opts
def test_vaapi_options(self):
opts = _get_codec_options("h264_vaapi", crf=28)
assert opts["qp"] == "28"
def test_qsv_options(self):
opts = _get_codec_options("h264_qsv", crf=25)
assert opts["global_quality"] == "25"
def test_no_g_no_crf(self):
opts = _get_codec_options("h264", g=None, crf=None)
assert "g" not in opts
assert "crf" not in opts
# ─── HW encoder detection tests ───
class TestHWEncoderDetection:
def test_detect_available_hw_encoders_returns_list(self):
result = detect_available_hw_encoders()
assert isinstance(result, list)
def test_detect_available_hw_encoders_only_valid(self):
from lerobot.datasets.video_utils import HW_ENCODERS
result = detect_available_hw_encoders()
for encoder in result:
assert encoder in HW_ENCODERS
def test_resolve_vcodec_passthrough(self):
assert resolve_vcodec("libsvtav1") == "libsvtav1"
assert resolve_vcodec("h264") == "h264"
def test_resolve_vcodec_auto_fallback(self):
"""When no HW encoders are available, auto should fall back to libsvtav1."""
with patch("lerobot.datasets.video_utils.detect_available_hw_encoders", return_value=[]):
assert resolve_vcodec("auto") == "libsvtav1"
def test_resolve_vcodec_auto_picks_hw(self):
"""When a HW encoder is available, auto should pick it."""
with patch(
"lerobot.datasets.video_utils.detect_available_hw_encoders",
return_value=["h264_videotoolbox"],
):
assert resolve_vcodec("auto") == "h264_videotoolbox"
def test_resolve_vcodec_auto_returns_valid(self):
"""Test that resolve_vcodec('auto') returns a known valid codec."""
result = resolve_vcodec("auto")
assert result in VALID_VIDEO_CODECS
def test_hw_encoder_names_accepted_in_validation(self):
"""Test that HW encoder names pass validation in VALID_VIDEO_CODECS."""
assert "auto" in VALID_VIDEO_CODECS
assert "h264_videotoolbox" in VALID_VIDEO_CODECS
assert "h264_nvenc" in VALID_VIDEO_CODECS
def test_resolve_vcodec_invalid_raises(self):
"""Test that resolve_vcodec raises ValueError for invalid codecs."""
with pytest.raises(ValueError, match="Invalid vcodec"):
resolve_vcodec("not_a_real_codec")
# ─── _CameraEncoderThread tests ───
class TestCameraEncoderThread:
def test_encodes_valid_mp4(self, tmp_path):
"""Test that the encoder thread creates a valid MP4 file with correct frame count."""
num_frames = 30
height, width = 64, 96
fps = 30
video_path = tmp_path / "test_output" / "test.mp4"
frame_queue: queue.Queue = queue.Queue(maxsize=60)
result_queue: queue.Queue = queue.Queue(maxsize=1)
stop_event = threading.Event()
encoder_thread = _CameraEncoderThread(
video_path=video_path,
fps=fps,
vcodec="libsvtav1",
pix_fmt="yuv420p",
g=2,
crf=30,
preset=13,
frame_queue=frame_queue,
result_queue=result_queue,
stop_event=stop_event,
)
encoder_thread.start()
# Feed frames (HWC uint8)
for _ in range(num_frames):
frame = np.random.randint(0, 255, (height, width, 3), dtype=np.uint8)
frame_queue.put(frame)
# Send sentinel
frame_queue.put(None)
encoder_thread.join(timeout=60)
assert not encoder_thread.is_alive()
# Check result
status, data = result_queue.get(timeout=5)
assert status == "ok"
assert data is not None # Stats should be returned
assert "mean" in data
assert "std" in data
assert "min" in data
assert "max" in data
assert "count" in data
# Verify the MP4 file is valid
assert video_path.exists()
with av.open(str(video_path)) as container:
stream = container.streams.video[0]
# The frame count should match
total_frames = sum(1 for _ in container.decode(stream))
assert total_frames == num_frames
def test_handles_chw_input(self, tmp_path):
"""Test that CHW format input is handled correctly."""
num_frames = 5
fps = 30
video_path = tmp_path / "test_chw" / "test.mp4"
frame_queue: queue.Queue = queue.Queue(maxsize=60)
result_queue: queue.Queue = queue.Queue(maxsize=1)
stop_event = threading.Event()
encoder_thread = _CameraEncoderThread(
video_path=video_path,
fps=fps,
vcodec="libsvtav1",
pix_fmt="yuv420p",
g=2,
crf=30,
preset=13,
frame_queue=frame_queue,
result_queue=result_queue,
stop_event=stop_event,
)
encoder_thread.start()
# Feed CHW frames
for _ in range(num_frames):
frame = np.random.randint(0, 255, (3, 64, 96), dtype=np.uint8)
frame_queue.put(frame)
frame_queue.put(None)
encoder_thread.join(timeout=60)
status, _ = result_queue.get(timeout=5)
assert status == "ok"
assert video_path.exists()
def test_stop_event_cancellation(self, tmp_path):
"""Test that setting the stop event causes the thread to exit."""
fps = 30
video_path = tmp_path / "test_cancel" / "test.mp4"
frame_queue: queue.Queue = queue.Queue(maxsize=60)
result_queue: queue.Queue = queue.Queue(maxsize=1)
stop_event = threading.Event()
encoder_thread = _CameraEncoderThread(
video_path=video_path,
fps=fps,
vcodec="libsvtav1",
pix_fmt="yuv420p",
g=2,
crf=30,
preset=13,
frame_queue=frame_queue,
result_queue=result_queue,
stop_event=stop_event,
)
encoder_thread.start()
# Feed a few frames
for _ in range(3):
frame = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
frame_queue.put(frame)
# Signal stop instead of sending sentinel
stop_event.set()
encoder_thread.join(timeout=10)
assert not encoder_thread.is_alive()
# ─── StreamingVideoEncoder tests ───
class TestStreamingVideoEncoder:
def test_single_camera_episode(self, tmp_path):
"""Test encoding a single camera episode."""
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13)
video_keys = [f"{OBS_IMAGES}.laptop"]
encoder.start_episode(video_keys, tmp_path)
num_frames = 20
for _ in range(num_frames):
frame = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
encoder.feed_frame(f"{OBS_IMAGES}.laptop", frame)
results = encoder.finish_episode()
assert f"{OBS_IMAGES}.laptop" in results
mp4_path, stats = results[f"{OBS_IMAGES}.laptop"]
assert mp4_path.exists()
assert stats is not None
# Verify frame count
with av.open(str(mp4_path)) as container:
stream = container.streams.video[0]
total_frames = sum(1 for _ in container.decode(stream))
assert total_frames == num_frames
encoder.close()
def test_multi_camera_episode(self, tmp_path):
"""Test encoding multiple cameras simultaneously."""
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30)
video_keys = [f"{OBS_IMAGES}.laptop", f"{OBS_IMAGES}.phone"]
encoder.start_episode(video_keys, tmp_path)
num_frames = 15
for _ in range(num_frames):
frame0 = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
frame1 = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
encoder.feed_frame(video_keys[0], frame0)
encoder.feed_frame(video_keys[1], frame1)
results = encoder.finish_episode()
for key in video_keys:
assert key in results
mp4_path, stats = results[key]
assert mp4_path.exists()
assert stats is not None
encoder.close()
def test_sequential_episodes(self, tmp_path):
"""Test that multiple sequential episodes work correctly."""
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30)
video_keys = [f"{OBS_IMAGES}.cam"]
for ep in range(3):
encoder.start_episode(video_keys, tmp_path)
num_frames = 10 + ep * 5
for _ in range(num_frames):
frame = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
encoder.feed_frame(f"{OBS_IMAGES}.cam", frame)
results = encoder.finish_episode()
mp4_path, stats = results[f"{OBS_IMAGES}.cam"]
assert mp4_path.exists()
with av.open(str(mp4_path)) as container:
stream = container.streams.video[0]
total_frames = sum(1 for _ in container.decode(stream))
assert total_frames == num_frames
encoder.close()
def test_cancel_episode(self, tmp_path):
"""Test that canceling an episode cleans up properly."""
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30)
video_keys = [f"{OBS_IMAGES}.cam"]
encoder.start_episode(video_keys, tmp_path)
for _ in range(5):
frame = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
encoder.feed_frame(f"{OBS_IMAGES}.cam", frame)
encoder.cancel_episode()
# Should be able to start a new episode after cancel
encoder.start_episode(video_keys, tmp_path)
for _ in range(5):
frame = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
encoder.feed_frame(f"{OBS_IMAGES}.cam", frame)
results = encoder.finish_episode()
assert f"{OBS_IMAGES}.cam" in results
encoder.close()
def test_feed_without_start_raises(self, tmp_path):
"""Test that feeding frames without starting an episode raises."""
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p")
with pytest.raises(RuntimeError, match="No active episode"):
encoder.feed_frame("cam", np.zeros((64, 96, 3), dtype=np.uint8))
encoder.close()
def test_finish_without_start_raises(self, tmp_path):
"""Test that finishing without starting raises."""
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p")
with pytest.raises(RuntimeError, match="No active episode"):
encoder.finish_episode()
encoder.close()
def test_close_is_idempotent(self, tmp_path):
"""Test that close() can be called multiple times safely."""
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p")
encoder.close()
encoder.close() # Should not raise
def test_video_duration_matches_frame_count(self, tmp_path):
"""Test that encoded video duration matches num_frames / fps."""
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13)
video_keys = [f"{OBS_IMAGES}.cam"]
encoder.start_episode(video_keys, tmp_path)
num_frames = 90 # 3 seconds at 30fps
for _ in range(num_frames):
frame = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
encoder.feed_frame(f"{OBS_IMAGES}.cam", frame)
results = encoder.finish_episode()
mp4_path, _ = results[f"{OBS_IMAGES}.cam"]
expected_duration = num_frames / 30.0 # 3.0 seconds
with av.open(str(mp4_path)) as container:
stream = container.streams.video[0]
total_frames = sum(1 for _ in container.decode(stream))
if stream.duration is not None:
actual_duration = float(stream.duration * stream.time_base)
else:
actual_duration = float(container.duration / av.time_base)
assert total_frames == num_frames
# Allow small tolerance for duration due to codec framing
assert abs(actual_duration - expected_duration) < 0.5, (
f"Video duration {actual_duration:.2f}s != expected {expected_duration:.2f}s"
)
encoder.close()
def test_multi_camera_start_episode_called_once(self, tmp_path):
"""Test that with multiple cameras, no frames are lost due to double start_episode."""
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30)
video_keys = [f"{OBS_IMAGES}.cam1", f"{OBS_IMAGES}.cam2"]
encoder.start_episode(video_keys, tmp_path)
num_frames = 30
for _ in range(num_frames):
frame0 = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
frame1 = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
encoder.feed_frame(video_keys[0], frame0)
encoder.feed_frame(video_keys[1], frame1)
results = encoder.finish_episode()
# Both cameras should have all frames
for key in video_keys:
mp4_path, stats = results[key]
assert mp4_path.exists()
with av.open(str(mp4_path)) as container:
stream = container.streams.video[0]
total_frames = sum(1 for _ in container.decode(stream))
assert total_frames == num_frames, (
f"Camera {key}: expected {num_frames} frames, got {total_frames}"
)
encoder.close()
def test_encoder_threads_passed_to_thread(self, tmp_path):
"""Test that encoder_threads is stored and passed through to encoder threads."""
encoder = StreamingVideoEncoder(
fps=30, vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, encoder_threads=2
)
assert encoder.encoder_threads == 2
video_keys = [f"{OBS_IMAGES}.cam"]
encoder.start_episode(video_keys, tmp_path)
# Verify the thread received the encoder_threads value
thread = encoder._threads[f"{OBS_IMAGES}.cam"]
assert thread.encoder_threads == 2
# Feed some frames and finish to ensure it works end-to-end
num_frames = 10
for _ in range(num_frames):
frame = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
encoder.feed_frame(f"{OBS_IMAGES}.cam", frame)
results = encoder.finish_episode()
mp4_path, stats = results[f"{OBS_IMAGES}.cam"]
assert mp4_path.exists()
assert stats is not None
with av.open(str(mp4_path)) as container:
stream = container.streams.video[0]
total_frames = sum(1 for _ in container.decode(stream))
assert total_frames == num_frames
encoder.close()
def test_encoder_threads_none_by_default(self, tmp_path):
"""Test that encoder_threads defaults to None (codec auto-detect)."""
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p")
assert encoder.encoder_threads is None
encoder.close()
def test_graceful_frame_dropping(self, tmp_path):
"""Test that full queue drops frames instead of crashing."""
encoder = StreamingVideoEncoder(
fps=30, vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13, queue_maxsize=1
)
video_keys = [f"{OBS_IMAGES}.cam"]
encoder.start_episode(video_keys, tmp_path)
# Feed many frames quickly - with queue_maxsize=1, some will be dropped
num_frames = 50
for _ in range(num_frames):
frame = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
encoder.feed_frame(f"{OBS_IMAGES}.cam", frame)
# Should not raise - frames are dropped gracefully
results = encoder.finish_episode()
assert f"{OBS_IMAGES}.cam" in results
mp4_path, _ = results[f"{OBS_IMAGES}.cam"]
assert mp4_path.exists()
# Some frames should have been dropped (queue was tiny)
dropped = encoder._dropped_frames.get(f"{OBS_IMAGES}.cam", 0)
# We can't guarantee drops but can verify no crash occurred
assert dropped >= 0
encoder.close()
# ─── Integration tests with LeRobotDataset ───
class TestStreamingEncoderIntegration:
def test_add_frame_save_episode_streaming(self, tmp_path):
"""Full integration test: add_frame -> save_episode with streaming encoding."""
from lerobot.datasets.lerobot_dataset import LeRobotDataset
features = {
"observation.images.cam": {
"dtype": "video",
"shape": (64, 96, 3),
"names": ["height", "width", "channels"],
},
"action": {"dtype": "float32", "shape": (6,), "names": ["j1", "j2", "j3", "j4", "j5", "j6"]},
}
dataset = LeRobotDataset.create(
repo_id="test/streaming",
fps=30,
features=features,
root=tmp_path / "streaming_test",
use_videos=True,
streaming_encoding=True,
)
assert dataset.writer._streaming_encoder is not None
num_frames = 20
for _ in range(num_frames):
frame = {
"observation.images.cam": np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8),
"action": np.random.randn(6).astype(np.float32),
"task": "test task",
}
dataset.add_frame(frame)
dataset.save_episode()
# Verify dataset metadata
assert dataset.meta.total_episodes == 1
assert dataset.meta.total_frames == num_frames
# Verify stats exist for the video key
assert dataset.meta.stats is not None
assert "observation.images.cam" in dataset.meta.stats
assert "action" in dataset.meta.stats
dataset.finalize()
def test_streaming_disabled_creates_pngs(self, tmp_path):
"""Test that disabling streaming encoding falls back to PNG path."""
from lerobot.datasets.lerobot_dataset import LeRobotDataset
features = {
"observation.images.cam": {
"dtype": "video",
"shape": (64, 96, 3),
"names": ["height", "width", "channels"],
},
"action": {"dtype": "float32", "shape": (6,), "names": ["j1", "j2", "j3", "j4", "j5", "j6"]},
}
dataset = LeRobotDataset.create(
repo_id="test/no_streaming",
fps=30,
features=features,
root=tmp_path / "no_streaming_test",
use_videos=True,
streaming_encoding=False,
)
assert dataset.writer._streaming_encoder is None
num_frames = 5
for _ in range(num_frames):
frame = {
"observation.images.cam": np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8),
"action": np.random.randn(6).astype(np.float32),
"task": "test task",
}
dataset.add_frame(frame)
# With streaming disabled, PNG files should be written
images_dir = dataset.root / "images"
assert images_dir.exists()
dataset.save_episode()
dataset.finalize()
def test_multi_episode_streaming(self, tmp_path):
"""Test recording multiple episodes with streaming encoding."""
from lerobot.datasets.lerobot_dataset import LeRobotDataset
features = {
"observation.images.cam": {
"dtype": "video",
"shape": (64, 96, 3),
"names": ["height", "width", "channels"],
},
"action": {"dtype": "float32", "shape": (2,), "names": ["j1", "j2"]},
}
dataset = LeRobotDataset.create(
repo_id="test/multi_ep",
fps=30,
features=features,
root=tmp_path / "multi_ep_test",
use_videos=True,
streaming_encoding=True,
)
for ep in range(3):
num_frames = 10 + ep * 5
for _ in range(num_frames):
frame = {
"observation.images.cam": np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8),
"action": np.random.randn(2).astype(np.float32),
"task": f"task_{ep}",
}
dataset.add_frame(frame)
dataset.save_episode()
assert dataset.meta.total_episodes == 3
assert dataset.meta.total_frames == 10 + 15 + 20
dataset.finalize()
def test_clear_episode_buffer_cancels_streaming(self, tmp_path):
"""Test that clearing episode buffer cancels streaming encoding."""
from lerobot.datasets.lerobot_dataset import LeRobotDataset
features = {
"observation.images.cam": {
"dtype": "video",
"shape": (64, 96, 3),
"names": ["height", "width", "channels"],
},
"action": {"dtype": "float32", "shape": (2,), "names": ["j1", "j2"]},
}
dataset = LeRobotDataset.create(
repo_id="test/cancel",
fps=30,
features=features,
root=tmp_path / "cancel_test",
use_videos=True,
streaming_encoding=True,
)
# Add some frames
for _ in range(5):
frame = {
"observation.images.cam": np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8),
"action": np.random.randn(2).astype(np.float32),
"task": "task",
}
dataset.add_frame(frame)
# Cancel and re-record
dataset.clear_episode_buffer()
# Record a new episode
for _ in range(10):
frame = {
"observation.images.cam": np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8),
"action": np.random.randn(2).astype(np.float32),
"task": "task",
}
dataset.add_frame(frame)
dataset.save_episode()
assert dataset.meta.total_episodes == 1
assert dataset.meta.total_frames == 10
dataset.finalize()
def test_multi_camera_add_frame_streaming(self, tmp_path):
"""Test that start_episode is called once with multiple video keys."""
from lerobot.datasets.lerobot_dataset import LeRobotDataset
features = {
"observation.images.cam1": {
"dtype": "video",
"shape": (64, 96, 3),
"names": ["height", "width", "channels"],
},
"observation.images.cam2": {
"dtype": "video",
"shape": (64, 96, 3),
"names": ["height", "width", "channels"],
},
"action": {"dtype": "float32", "shape": (2,), "names": ["j1", "j2"]},
}
dataset = LeRobotDataset.create(
repo_id="test/multi_cam",
fps=30,
features=features,
root=tmp_path / "multi_cam_test",
use_videos=True,
streaming_encoding=True,
)
num_frames = 15
for _ in range(num_frames):
frame = {
"observation.images.cam1": np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8),
"observation.images.cam2": np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8),
"action": np.random.randn(2).astype(np.float32),
"task": "test task",
}
dataset.add_frame(frame)
dataset.save_episode()
assert dataset.meta.total_episodes == 1
assert dataset.meta.total_frames == num_frames
dataset.finalize()
tests/datasets/test_subtask_dataset.py
+190
View File
@@ -0,0 +1,190 @@
#!/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 subtask functionality in LeRobotDataset.
These tests verify that:
- Subtask information is correctly loaded from datasets that have subtask data
- The __getitem__ method correctly adds subtask strings to returned items
- Subtask handling gracefully handles missing data
"""
import pandas as pd
import pytest
import torch
from lerobot.datasets.lerobot_dataset import LeRobotDataset
class TestSubtaskDataset:
"""Tests for subtask handling in LeRobotDataset."""
@pytest.fixture
def subtask_dataset(self):
"""Load the test subtask dataset from the hub."""
# Use lerobot/pusht-subtask dataset with episode 1
return LeRobotDataset(
repo_id="lerobot/pusht-subtask",
episodes=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
)
def test_subtask_dataset_loads(self, subtask_dataset):
"""Test that the subtask dataset loads successfully."""
assert subtask_dataset is not None
assert len(subtask_dataset) > 0
def test_subtask_metadata_loaded(self, subtask_dataset):
"""Test that subtask metadata is loaded when present in dataset."""
# The dataset should have subtasks metadata loaded
assert subtask_dataset.meta.subtasks is not None
assert isinstance(subtask_dataset.meta.subtasks, pd.DataFrame)
def test_subtask_index_in_features(self, subtask_dataset):
"""Test that subtask_index is a feature when dataset has subtasks."""
assert "subtask_index" in subtask_dataset.features
def test_getitem_returns_subtask_string(self, subtask_dataset):
"""Test that __getitem__ correctly adds subtask string to returned item."""
item = subtask_dataset[0]
# Subtask should be present in the returned item
assert "subtask" in item
assert isinstance(item["subtask"], str)
assert len(item["subtask"]) > 0 # Should not be empty
def test_getitem_has_subtask_index(self, subtask_dataset):
"""Test that __getitem__ includes subtask_index."""
item = subtask_dataset[0]
assert "subtask_index" in item
assert isinstance(item["subtask_index"], torch.Tensor)
def test_subtask_index_maps_to_valid_subtask(self, subtask_dataset):
"""Test that subtask_index correctly maps to a subtask in metadata."""
item = subtask_dataset[0]
subtask_idx = item["subtask_index"].item()
subtask_from_metadata = subtask_dataset.meta.subtasks.iloc[subtask_idx].name
assert item["subtask"] == subtask_from_metadata
def test_all_items_have_subtask(self, subtask_dataset):
"""Test that all items in the dataset have subtask information."""
for i in range(min(len(subtask_dataset), 5)): # Check first 5 items
item = subtask_dataset[i]
assert "subtask" in item
assert isinstance(item["subtask"], str)
def test_task_and_subtask_coexist(self, subtask_dataset):
"""Test that both task and subtask are present in returned items."""
item = subtask_dataset[0]
# Both task and subtask should be present
assert "task" in item
assert "subtask" in item
assert isinstance(item["task"], str)
assert isinstance(item["subtask"], str)
class TestSubtaskDatasetMissing:
"""Tests for graceful handling when subtask data is missing."""
@pytest.fixture
def dataset_without_subtasks(self, tmp_path, empty_lerobot_dataset_factory):
"""Create a dataset without subtask information."""
features = {"state": {"dtype": "float32", "shape": (2,), "names": None}}
dataset = empty_lerobot_dataset_factory(root=tmp_path / "no_subtask", features=features)
# Add some frames and save
for _ in range(5):
dataset.add_frame({"state": torch.randn(2), "task": "Test task"})
dataset.save_episode()
dataset.finalize()
# Reload the dataset
return LeRobotDataset(dataset.repo_id, root=dataset.root)
def test_no_subtask_in_features(self, dataset_without_subtasks):
"""Test that subtask_index is not in features when not provided."""
assert "subtask_index" not in dataset_without_subtasks.features
def test_getitem_without_subtask(self, dataset_without_subtasks):
"""Test that __getitem__ works when subtask is not present."""
item = dataset_without_subtasks[0]
# Item should still be retrievable
assert item is not None
assert "state" in item
assert "task" in item
# Subtask should NOT be present
assert "subtask" not in item
def test_subtasks_metadata_is_none(self, dataset_without_subtasks):
"""Test that subtasks metadata is None when not present."""
assert dataset_without_subtasks.meta.subtasks is None
class TestSubtaskEdgeCases:
"""Edge case tests for subtask handling."""
def test_subtask_with_multiple_episodes(self):
"""Test subtask handling with multiple episodes if available."""
try:
dataset = LeRobotDataset(
repo_id="lerobot/pusht-subtask",
episodes=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
)
except Exception:
pytest.skip("Could not load test-subtask dataset")
# Check first and last items have valid subtasks
first_item = dataset[0]
last_item = dataset[len(dataset) - 1]
assert "subtask" in first_item
assert "subtask" in last_item
assert isinstance(first_item["subtask"], str)
assert isinstance(last_item["subtask"], str)
def test_subtask_index_consistency(self):
"""Test that same subtask_index returns same subtask string."""
try:
dataset = LeRobotDataset(
repo_id="lerobot/pusht-subtask",
episodes=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
)
except Exception:
pytest.skip("Could not load test-subtask dataset")
if len(dataset) < 2:
pytest.skip("Dataset too small for this test")
# Collect subtask_index to subtask mappings
subtask_map = {}
for i in range(min(len(dataset), 10)):
item = dataset[i]
idx = item["subtask_index"].item()
subtask = item["subtask"]
if idx in subtask_map:
# Same index should always return same subtask
assert subtask_map[idx] == subtask, (
f"Inconsistent subtask for index {idx}: '{subtask_map[idx]}' vs '{subtask}'"
)
else:
subtask_map[idx] = subtask