admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-30 21:19:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

tests(all shapes): enhancing tests to cover all possible features shapes

Browse Source
This commit is contained in:
CarolinePascal
2026-07-30 18:23:45 +02:00
parent 910aac4c03
commit ba169ce524
1 changed files with 22 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/datasets/test_dataset_tools.py
+22 -52
View File
@@ -304,40 +304,45 @@ def test_merge_empty_list(tmp_path):
merge_datasets([], output_repo_id="merged", output_dir=tmp_path)
def test_add_features_with_values(sample_dataset, tmp_path):
"""Test adding a feature with pre-computed values."""
@pytest.mark.parametrize(
"values, feature_shape, expected_item_shape",
[
(np.random.randn(50).astype(np.float32), (1,), ()),
(np.random.randn(50, 1).astype(np.float32), (1,), ()),
(np.random.randn(50, 7).astype(np.float32), (7,), (7,)),
(np.random.randn(50, 1, 4).astype(np.float32), (1, 4), (1, 4)),
],
ids=["scalar_1d", "scalar_2d", "vector", "matrix"],
)
def test_add_features_with_values(sample_dataset, tmp_path, values, feature_shape, expected_item_shape):
"""Test adding a pre-computed feature across supported per-frame shapes."""
num_frames = sample_dataset.meta.total_frames
reward_values = np.random.randn(num_frames, 1).astype(np.float32)
assert len(values) == num_frames
feature_info = {
"dtype": "float32",
"shape": (1,),
"names": None,
}
features = {
"reward": (reward_values, feature_info),
}
feature_info = {"dtype": "float32", "shape": feature_shape, "names": None}
features = {"new_feature": (values, feature_info)}
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 / "with_reward")
mock_snapshot_download.return_value = str(tmp_path / "with_feature")
new_dataset = add_features(
dataset=sample_dataset,
features=features,
output_dir=tmp_path / "with_reward",
output_dir=tmp_path / "with_feature",
)
assert "reward" in new_dataset.meta.features
assert new_dataset.meta.features["reward"] == feature_info
assert "new_feature" in new_dataset.meta.features
assert new_dataset.meta.features["new_feature"] == feature_info
assert len(new_dataset) == num_frames
sample_item = new_dataset[0]
assert "reward" in sample_item
assert isinstance(sample_item["reward"], torch.Tensor)
assert "new_feature" in sample_item
assert isinstance(sample_item["new_feature"], torch.Tensor)
assert tuple(sample_item["new_feature"].shape) == expected_item_shape
def test_add_features_with_callable(sample_dataset, tmp_path):
@@ -377,41 +382,6 @@ def test_add_features_with_callable(sample_dataset, tmp_path):
assert first_frame["frame_index"] == 0
assert float(first_frame["reward"]) == 0.0
def test_add_features_with_multidimensional_values(sample_dataset, tmp_path):
"""Test adding a multi-dimensional per-frame feature."""
num_frames = sample_dataset.meta.total_frames
feature_values = np.random.randn(num_frames, 1, 4).astype(np.float32)
feature_info = {
"dtype": "float32",
"shape": (1, 4),
"names": None,
}
features = {
"multi_dim_feature": (feature_values, feature_info),
}
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 / "with_multi_dim_feature")
new_dataset = add_features(
dataset=sample_dataset,
features=features,
output_dir=tmp_path / "with_multi_dim_feature",
)
assert "multi_dim_feature" in new_dataset.meta.features
sample_item = new_dataset[0]
assert "multi_dim_feature" in sample_item
assert isinstance(sample_item["multi_dim_feature"], torch.Tensor)
assert tuple(sample_item["multi_dim_feature"].shape) == (1, 4)
def test_add_existing_feature(sample_dataset, tmp_path):
"""Test error when adding an existing feature."""
feature_info = {"dtype": "float32", "shape": (1,)}