diff --git a/src/lerobot/datasets/dataset_tools.py b/src/lerobot/datasets/dataset_tools.py index 31e075d7c..ce3055441 100644 --- a/src/lerobot/datasets/dataset_tools.py +++ b/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 diff --git a/tests/datasets/test_dataset_tools.py b/tests/datasets/test_dataset_tools.py index c19e7f41f..f9365e7da 100644 --- a/tests/datasets/test_dataset_tools.py +++ b/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,)}