diff --git a/src/lerobot/datasets/dataset_tools.py b/src/lerobot/datasets/dataset_tools.py index 772f8b7e5..317841a65 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 95c94f22c..992766e41 100644 --- a/tests/datasets/test_dataset_tools.py +++ b/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):