diff --git a/src/lerobot/datasets/compute_stats.py b/src/lerobot/datasets/compute_stats.py index c6d12d257..133c59f0e 100644 --- a/src/lerobot/datasets/compute_stats.py +++ b/src/lerobot/datasets/compute_stats.py @@ -520,6 +520,9 @@ def compute_episode_stats( continue if any(d == 0 for d in features[key].get("shape", ())): + logging.warning( + f"Skipping stats for feature '{key}' with a zero-width shape {features[key]['shape']}." + ) continue if features[key]["dtype"] in ["image", "video"]: diff --git a/tests/datasets/test_compute_stats.py b/tests/datasets/test_compute_stats.py index b1cadf0e3..aca5c89e5 100644 --- a/tests/datasets/test_compute_stats.py +++ b/tests/datasets/test_compute_stats.py @@ -13,6 +13,7 @@ # 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 from unittest.mock import patch import numpy as np @@ -687,8 +688,8 @@ def test_compute_episode_stats_string_features_skipped(): assert "q01" in stats["action"] -def test_compute_episode_stats_zero_width_features_skipped(): - """Test that features with a zero-width dim (e.g. shape=(0,)) are skipped.""" +def test_compute_episode_stats_zero_width_features_skipped(caplog): + """Test that features with a zero-width dim (e.g. shape=(0,)) are skipped with a warning.""" episode_data = { "empty": np.zeros((100, 0), dtype=np.float32), # Zero-width feature "action": np.random.normal(0, 1, (100, 5)), @@ -698,13 +699,12 @@ def test_compute_episode_stats_zero_width_features_skipped(): "action": {"dtype": "float32", "shape": (5,)}, } - stats = compute_episode_stats( - episode_data, - features, - ) + with caplog.at_level(logging.WARNING): + stats = compute_episode_stats(episode_data, features) - # Zero-width features should be skipped + # Zero-width features should be skipped with a warning, others computed as usual assert "empty" not in stats + assert "empty" in caplog.text assert "action" in stats assert "q01" in stats["action"] assert stats["action"]["mean"].shape == (5,)