chore(dataset): add check dataset shape

This commit is contained in:
Steven Palma
2026-07-27 18:36:15 +02:00
parent 034693f724
commit 45243dcf7c
4 changed files with 23 additions and 20 deletions
+2 -2
View File
@@ -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", ())):
+5
View File
@@ -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:
+5 -10
View File
@@ -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)
+11 -8
View File
@@ -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 ────────────────────────────────────────────────