fix(datasets): support multi-dimensional per-frame features in add_features (#4246)

* Fix add_features for multi-dimensional per-frame features

* Use generic names in multidimensional add_features test

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

* chore(format): formatting code

---------

Co-authored-by: felixmin <felix.minze@live.de>

src/lerobot/datasets/dataset_tools.py

@@ -1045,10 +1045,12 @@ def _copy_data_with_feature_changes(
                     df[feature_name] = feature_values
                 else:
                     feature_slice = values[frame_idx:end_idx]
-                    if len(feature_slice.shape) > 1 and feature_slice.shape[1] == 1:
+                    if feature_slice.ndim == 1:
+                        df[feature_name] = feature_slice
+                    elif feature_slice.ndim == 2 and feature_slice.shape[1] == 1:
                         df[feature_name] = feature_slice.flatten()
                     else:
-                        df[feature_name] = feature_slice
+                        df[feature_name] = list(feature_slice)
             frame_idx = end_idx
 
         # Write using the same chunk/file structure as source

tests/datasets/test_dataset_tools.py

@@ -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):