Fix add_features for multi-dimensional per-frame features

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

@@ -378,6 +378,40 @@ def test_add_features_with_callable(sample_dataset, tmp_path):
     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
+    latent_values = np.random.randn(num_frames, 1, 4).astype(np.float32)
+
+    feature_info = {
+        "dtype": "float32",
+        "shape": (1, 4),
+        "names": None,
+    }
+    features = {
+        "latent_vectors": (latent_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_latents")
+
+        new_dataset = add_features(
+            dataset=sample_dataset,
+            features=features,
+            output_dir=tmp_path / "with_latents",
+        )
+
+    assert "latent_vectors" in new_dataset.meta.features
+    sample_item = new_dataset[0]
+    assert "latent_vectors" in sample_item
+    assert isinstance(sample_item["latent_vectors"], torch.Tensor)
+    assert tuple(sample_item["latent_vectors"].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,)}