From 45243dcf7c16009b4a924f5adee4593d5ba5ad2f Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Mon, 27 Jul 2026 18:36:15 +0200 Subject: [PATCH] chore(dataset): add check dataset shape --- src/lerobot/datasets/compute_stats.py | 4 ++-- src/lerobot/datasets/feature_utils.py | 5 +++++ tests/datasets/test_compute_stats.py | 15 +++++---------- tests/datasets/test_dataset_writer.py | 19 +++++++++++-------- 4 files changed, 23 insertions(+), 20 deletions(-) diff --git a/src/lerobot/datasets/compute_stats.py b/src/lerobot/datasets/compute_stats.py index 4d68f3feb..8cfa778d7 100644 --- a/src/lerobot/datasets/compute_stats.py +++ b/src/lerobot/datasets/compute_stats.py @@ -519,8 +519,8 @@ def compute_episode_stats( if features[key]["dtype"] in {"string", "language"}: continue - # Features with a zero-width dimension (e.g. shape=(0,)) carry no - # statistics-bearing values. Skip them like strings instead of letting + # Features with a zero-width dimension contain no statistics-bearing + # values. Skip them like strings instead of letting # get_feature_stats -> RunningQuantileStats.update reshape a size-0 array, # which raises "ValueError: cannot reshape array of size 0". if any(dim == 0 for dim in features[key].get("shape", ())): diff --git a/src/lerobot/datasets/feature_utils.py b/src/lerobot/datasets/feature_utils.py index a6bbfae2c..b60fdb04a 100644 --- a/src/lerobot/datasets/feature_utils.py +++ b/src/lerobot/datasets/feature_utils.py @@ -64,6 +64,11 @@ def get_hf_features_from_features(features: dict) -> datasets.Features: continue elif ft["dtype"] == "image": hf_features[key] = datasets.Image() + elif len(ft["shape"]) > 1 and any(dim == 0 for dim in ft["shape"]): + raise ValueError( + f"Multidimensional features with a zero-width dimension are not supported: " + f"'{key}' has shape {ft['shape']}. Only the one-dimensional shape (0,) is supported." + ) elif ft["shape"] == (1,): hf_features[key] = datasets.Value(dtype=ft["dtype"]) elif len(ft["shape"]) == 1: diff --git a/tests/datasets/test_compute_stats.py b/tests/datasets/test_compute_stats.py index 8a3eb9833..c512bf9bf 100644 --- a/tests/datasets/test_compute_stats.py +++ b/tests/datasets/test_compute_stats.py @@ -687,21 +687,16 @@ def test_compute_episode_stats_string_features_skipped(): assert "q01" in stats["action"] -def test_compute_episode_stats_zero_width_feature_skipped(): - """Features with a zero-width dimension carry no values and must be skipped, not crash. - - Regression test for https://github.com/huggingface/lerobot/issues/3654: - a feature declared with shape=(0,) previously reached RunningQuantileStats.update, - where ``batch.reshape(-1, batch.shape[-1])`` raised - "ValueError: cannot reshape array of size 0 into shape (0)". - """ +@pytest.mark.parametrize("shape", [(0,), (0, 2), (2, 0), (1, 0, 2)]) +def test_compute_episode_stats_zero_width_feature_skipped(shape): + """Features with any zero-width dimension carry no values and are skipped.""" episode_data = { "action": np.random.normal(0, 1, (100, 5)).astype(np.float32), - "target": np.zeros((100, 0), dtype=np.float32), # zero-width feature + "target": np.zeros((100, *shape), dtype=np.float32), } features = { "action": {"dtype": "float32", "shape": (5,)}, - "target": {"dtype": "float32", "shape": (0,)}, + "target": {"dtype": "float32", "shape": shape}, } stats = compute_episode_stats(episode_data, features) diff --git a/tests/datasets/test_dataset_writer.py b/tests/datasets/test_dataset_writer.py index a20064129..3517e7bfc 100644 --- a/tests/datasets/test_dataset_writer.py +++ b/tests/datasets/test_dataset_writer.py @@ -27,6 +27,7 @@ pytest.importorskip("datasets", reason="datasets is required (install lerobot[da from lerobot.configs import VideoEncoderConfig from lerobot.datasets.dataset_writer import _encode_video_worker +from lerobot.datasets.feature_utils import get_hf_features_from_features 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 @@ -190,12 +191,7 @@ def test_save_multiple_episodes(tmp_path): def test_save_episode_with_zero_width_feature(tmp_path): - """save_episode() succeeds when a feature has a zero-width dimension (shape=(0,)). - - Regression test for https://github.com/huggingface/lerobot/issues/3654: such a - feature previously crashed stats computation with - "ValueError: cannot reshape array of size 0 into shape (0)". - """ + """A one-dimensional empty numeric feature round-trips and has no statistics.""" features = { **SIMPLE_FEATURES, "target": {"dtype": "float32", "shape": (0,), "names": None}, @@ -204,19 +200,26 @@ def test_save_episode_with_zero_width_feature(tmp_path): dataset = LeRobotDataset.create(repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=features, root=root) for _ in range(4): dataset.add_frame(_make_frame(features)) - dataset.save_episode() # previously raised ValueError on the zero-width 'target' feature + dataset.save_episode() dataset.finalize() assert dataset.meta.total_episodes == 1 assert dataset.meta.total_frames == 4 - # The zero-width feature round-trips back as an empty vector and is excluded from stats. reloaded = LeRobotDataset(repo_id=DUMMY_REPO_ID, root=root) target = np.asarray(reloaded[0]["target"]) assert target.shape == (0,) assert "target" not in (reloaded.meta.stats or {}) +@pytest.mark.parametrize("shape", [(0, 2), (2, 0), (1, 0, 2)]) +def test_multidimensional_zero_width_feature_rejected(shape): + features = {"target": {"dtype": "float32", "shape": shape, "names": None}} + + with pytest.raises(ValueError, match="Multidimensional features with a zero-width dimension"): + get_hf_features_from_features(features) + + # ── clear / lifecycle ────────────────────────────────────────────────