From e72daaf5897309b93f7f5a2fa72ff4d18ad7b05b Mon Sep 17 00:00:00 2001
From: "k7agar@gmail.com" <vovw>
Date: Tue, 28 Oct 2025 20:17:42 +0530
Subject: [PATCH 1/5] feat: add support for

---
 src/lerobot/scripts/lerobot_edit_dataset.py   |  97 +++++-
 .../test_lerobot_edit_dataset_add_feature.py  | 293 ++++++++++++++++++
 2 files changed, 387 insertions(+), 3 deletions(-)
 create mode 100644 tests/scripts/test_lerobot_edit_dataset_add_feature.py

diff --git a/src/lerobot/scripts/lerobot_edit_dataset.py b/src/lerobot/scripts/lerobot_edit_dataset.py
index 83ba027bcc..9aa6419e3b 100644
--- a/src/lerobot/scripts/lerobot_edit_dataset.py
+++ b/src/lerobot/scripts/lerobot_edit_dataset.py
@@ -18,7 +18,7 @@
 Edit LeRobot datasets using various transformation tools.
 
 This script allows you to delete episodes, split datasets, merge datasets,
-and remove features. When new_repo_id is specified, creates a new dataset.
+remove features, and add features. When new_repo_id is specified, creates a new dataset.
 
 Usage Examples:
 
@@ -65,6 +65,12 @@
         --operation.type remove_feature \
         --operation.feature_names "['observation.images.top']"
 
+Add feature from numpy file:
+    python -m lerobot.scripts.lerobot_edit_dataset \
+        --repo_id lerobot/pusht \
+        --operation.type add_feature \
+        --operation.features '{"reward": {"file": "rewards.npy", "dtype": "float32", "shape": [1], "names": null}}'
+
 Using JSON config file:
     python -m lerobot.scripts.lerobot_edit_dataset \
         --config_path path/to/edit_config.json
@@ -75,8 +81,11 @@
 from dataclasses import dataclass
 from pathlib import Path
 
+import numpy as np
+
 from lerobot.configs import parser
 from lerobot.datasets.dataset_tools import (
+    add_features,
     delete_episodes,
     merge_datasets,
     remove_feature,
@@ -111,10 +120,16 @@ class RemoveFeatureConfig:
     feature_names: list[str] | None = None
 
 
+@dataclass
+class AddFeatureConfig:
+    type: str = "add_feature"
+    features: dict[str, dict] | None = None
+
+
 @dataclass
 class EditDatasetConfig:
     repo_id: str
-    operation: DeleteEpisodesConfig | SplitConfig | MergeConfig | RemoveFeatureConfig
+    operation: DeleteEpisodesConfig | SplitConfig | MergeConfig | RemoveFeatureConfig | AddFeatureConfig
     root: str | None = None
     new_repo_id: str | None = None
     push_to_hub: bool = False
@@ -258,6 +273,80 @@ def handle_remove_feature(cfg: EditDatasetConfig) -> None:
         LeRobotDataset(output_repo_id, root=output_dir).push_to_hub()
 
 
+def handle_add_feature(cfg: EditDatasetConfig) -> None:
+    if not isinstance(cfg.operation, AddFeatureConfig):
+        raise ValueError("Operation config must be AddFeatureConfig")
+
+    if not cfg.operation.features:
+        raise ValueError("features must be specified for add_feature operation")
+
+    dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
+
+    # Get base directory (parent of dataset root)
+    base_dir = dataset.root.parent if cfg.root else None
+    output_repo_id, output_dir = get_output_path(
+        cfg.repo_id, cfg.new_repo_id, base_dir
+    )
+
+    if cfg.new_repo_id is None:
+        dataset.root = Path(str(dataset.root) + "_old")
+
+    # Process features config to load data and prepare for add_features
+    features_dict = {}
+    for feature_name, feature_config in cfg.operation.features.items():
+        # Extract feature info (dtype, shape, names)
+        shape = feature_config.get("shape")
+        # Convert list to tuple if needed
+        if isinstance(shape, list):
+            shape = tuple(shape)
+
+        feature_info = {
+            "dtype": feature_config.get("dtype"),
+            "shape": shape,
+            "names": feature_config.get("names"),
+        }
+
+        # Load feature data from file
+        feature_file = feature_config.get("file")
+        if not feature_file:
+            raise ValueError(f"Feature '{feature_name}' must specify a 'file' path to load data from")
+
+        file_path = Path(feature_file)
+        if not file_path.exists():
+            raise FileNotFoundError(f"Feature file not found: {feature_file}")
+
+        # Load numpy array
+        if file_path.suffix == ".npy":
+            feature_data = np.load(file_path)
+        else:
+            raise ValueError(f"Unsupported file format for feature '{feature_name}': {file_path.suffix}")
+
+        # Validate data length matches dataset
+        expected_length = dataset.meta.total_frames
+        if len(feature_data) != expected_length:
+            raise ValueError(
+                f"Feature '{feature_name}' data length ({len(feature_data)}) "
+                f"does not match dataset length ({expected_length})"
+            )
+
+        features_dict[feature_name] = (feature_data, feature_info)
+
+    logging.info(f"Adding features {list(features_dict.keys())} to {cfg.repo_id}")
+    new_dataset = add_features(
+        dataset,
+        features=features_dict,
+        output_dir=output_dir,
+        repo_id=output_repo_id,
+    )
+
+    logging.info(f"Dataset saved to {output_dir}")
+    logging.info(f"Updated features: {list(new_dataset.meta.features.keys())}")
+
+    if cfg.push_to_hub:
+        logging.info(f"Pushing to hub as {output_repo_id}")
+        LeRobotDataset(output_repo_id, root=output_dir).push_to_hub()
+
+
 @parser.wrap()
 def edit_dataset(cfg: EditDatasetConfig) -> None:
     operation_type = cfg.operation.type
@@ -270,10 +359,12 @@ def edit_dataset(cfg: EditDatasetConfig) -> None:
         handle_merge(cfg)
     elif operation_type == "remove_feature":
         handle_remove_feature(cfg)
+    elif operation_type == "add_feature":
+        handle_add_feature(cfg)
     else:
         raise ValueError(
             f"Unknown operation type: {operation_type}\n"
-            f"Available operations: delete_episodes, split, merge, remove_feature"
+            f"Available operations: delete_episodes, split, merge, remove_feature, add_feature"
         )
 
 
diff --git a/tests/scripts/test_lerobot_edit_dataset_add_feature.py b/tests/scripts/test_lerobot_edit_dataset_add_feature.py
new file mode 100644
index 0000000000..7b04cb7be8
--- /dev/null
+++ b/tests/scripts/test_lerobot_edit_dataset_add_feature.py
@@ -0,0 +1,293 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""Tests for add_feature operation in lerobot_edit_dataset script."""
+
+from pathlib import Path
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.scripts.lerobot_edit_dataset import AddFeatureConfig, EditDatasetConfig, handle_add_feature
+
+
+@pytest.fixture
+def sample_dataset(tmp_path, empty_lerobot_dataset_factory):
+    """Create a sample dataset for testing."""
+    features = {
+        "action": {"dtype": "float32", "shape": (6,), "names": None},
+        "observation.state": {"dtype": "float32", "shape": (4,), "names": None},
+    }
+
+    dataset = empty_lerobot_dataset_factory(
+        root=tmp_path / "test_dataset",
+        features=features,
+    )
+
+    for ep_idx in range(3):
+        for _ in range(10):
+            frame = {
+                "action": np.random.randn(6).astype(np.float32),
+                "observation.state": np.random.randn(4).astype(np.float32),
+                "task": f"task_{ep_idx % 2}",
+            }
+            dataset.add_frame(frame)
+        dataset.save_episode()
+
+    dataset.finalize()
+    return dataset
+
+
+def test_add_feature_from_numpy_file(sample_dataset, tmp_path):
+    """Test adding a feature from a numpy file."""
+    # Create reward data file
+    num_frames = sample_dataset.meta.total_frames
+    reward_data = np.random.randn(num_frames, 1).astype(np.float32)
+    reward_file = tmp_path / "rewards.npy"
+    np.save(reward_file, reward_data)
+
+    # Create config
+    feature_config = {
+        "reward": {
+            "file": str(reward_file),
+            "dtype": "float32",
+            "shape": [1],
+            "names": None,
+        }
+    }
+
+    operation = AddFeatureConfig(type="add_feature", features=feature_config)
+
+    cfg = EditDatasetConfig(
+        repo_id=sample_dataset.repo_id,
+        operation=operation,
+        root=str(sample_dataset.root),
+        new_repo_id="test_dataset_with_reward",
+    )
+
+    with (
+        patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_version,
+        patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_download,
+    ):
+        mock_version.return_value = "v3.0"
+        mock_download.return_value = str(tmp_path / "test_dataset_with_reward")
+
+        handle_add_feature(cfg)
+
+        # Verify result
+        new_dataset = LeRobotDataset(
+            repo_id="test_dataset_with_reward", root=str(tmp_path / "test_dataset_with_reward")
+        )
+
+    assert "reward" in new_dataset.meta.features
+    assert new_dataset.meta.features["reward"]["dtype"] == "float32"
+    assert new_dataset.meta.features["reward"]["shape"] == (1,)  # Shape is stored as tuple
+
+    # Check data integrity
+    assert len(new_dataset) == num_frames
+    sample = new_dataset[0]
+    assert "reward" in sample
+    assert "action" in sample
+    assert "observation.state" in sample
+
+
+def test_add_multiple_features(sample_dataset, tmp_path):
+    """Test adding multiple features at once."""
+    num_frames = sample_dataset.meta.total_frames
+
+    # Create multiple feature files
+    reward_data = np.random.randn(num_frames, 1).astype(np.float32)
+    success_data = np.random.randint(0, 2, size=(num_frames, 1)).astype(np.int64)
+
+    reward_file = tmp_path / "rewards.npy"
+    success_file = tmp_path / "success.npy"
+
+    np.save(reward_file, reward_data)
+    np.save(success_file, success_data)
+
+    # Create config
+    feature_config = {
+        "reward": {
+            "file": str(reward_file),
+            "dtype": "float32",
+            "shape": [1],
+            "names": None,
+        },
+        "success": {
+            "file": str(success_file),
+            "dtype": "int64",
+            "shape": [1],
+            "names": None,
+        },
+    }
+
+    operation = AddFeatureConfig(type="add_feature", features=feature_config)
+
+    cfg = EditDatasetConfig(
+        repo_id=sample_dataset.repo_id,
+        operation=operation,
+        root=str(sample_dataset.root),
+        new_repo_id="test_dataset_with_features",
+    )
+
+    with (
+        patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_version,
+        patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_download,
+    ):
+        mock_version.return_value = "v3.0"
+        mock_download.return_value = str(tmp_path / "test_dataset_with_features")
+
+        handle_add_feature(cfg)
+
+        new_dataset = LeRobotDataset(
+            repo_id="test_dataset_with_features", root=str(tmp_path / "test_dataset_with_features")
+        )
+
+    assert "reward" in new_dataset.meta.features
+    assert "success" in new_dataset.meta.features
+    assert len(new_dataset) == num_frames
+
+
+def test_add_feature_missing_file(sample_dataset, tmp_path):
+    """Test error when feature file doesn't exist."""
+    feature_config = {
+        "reward": {
+            "file": str(tmp_path / "nonexistent.npy"),
+            "dtype": "float32",
+            "shape": [1],
+            "names": None,
+        }
+    }
+
+    operation = AddFeatureConfig(type="add_feature", features=feature_config)
+
+    cfg = EditDatasetConfig(
+        repo_id=sample_dataset.repo_id,
+        operation=operation,
+        root=str(sample_dataset.root),
+        new_repo_id="test_dataset_with_reward",
+    )
+
+    with pytest.raises(FileNotFoundError, match="Feature file not found"):
+        handle_add_feature(cfg)
+
+
+def test_add_feature_wrong_length(sample_dataset, tmp_path):
+    """Test error when feature data length doesn't match dataset."""
+    # Create reward data with wrong length
+    wrong_length = sample_dataset.meta.total_frames + 10
+    reward_data = np.random.randn(wrong_length, 1).astype(np.float32)
+    reward_file = tmp_path / "rewards.npy"
+    np.save(reward_file, reward_data)
+
+    feature_config = {
+        "reward": {
+            "file": str(reward_file),
+            "dtype": "float32",
+            "shape": [1],
+            "names": None,
+        }
+    }
+
+    operation = AddFeatureConfig(type="add_feature", features=feature_config)
+
+    cfg = EditDatasetConfig(
+        repo_id=sample_dataset.repo_id,
+        operation=operation,
+        root=str(sample_dataset.root),
+        new_repo_id="test_dataset_with_reward",
+    )
+
+    with pytest.raises(ValueError, match="data length .* does not match dataset length"):
+        handle_add_feature(cfg)
+
+
+def test_add_feature_no_file_specified(sample_dataset, tmp_path):
+    """Test error when no file is specified for feature."""
+    feature_config = {
+        "reward": {
+            "dtype": "float32",
+            "shape": [1],
+            "names": None,
+        }
+    }
+
+    operation = AddFeatureConfig(type="add_feature", features=feature_config)
+
+    cfg = EditDatasetConfig(
+        repo_id=sample_dataset.repo_id,
+        operation=operation,
+        root=str(sample_dataset.root),
+        new_repo_id="test_dataset_with_reward",
+    )
+
+    with pytest.raises(ValueError, match="must specify a 'file' path"):
+        handle_add_feature(cfg)
+
+
+def test_add_feature_unsupported_format(sample_dataset, tmp_path):
+    """Test error with unsupported file format."""
+    # Create a text file instead of numpy
+    reward_file = tmp_path / "rewards.txt"
+    reward_file.write_text("some data")
+
+    feature_config = {
+        "reward": {
+            "file": str(reward_file),
+            "dtype": "float32",
+            "shape": [1],
+            "names": None,
+        }
+    }
+
+    operation = AddFeatureConfig(type="add_feature", features=feature_config)
+
+    cfg = EditDatasetConfig(
+        repo_id=sample_dataset.repo_id,
+        operation=operation,
+        root=str(sample_dataset.root),
+        new_repo_id="test_dataset_with_reward",
+    )
+
+    with pytest.raises(ValueError, match="Unsupported file format"):
+        handle_add_feature(cfg)
+
+
+def test_add_feature_no_features_specified(sample_dataset, tmp_path):
+    """Test error when no features are specified."""
+    operation = AddFeatureConfig(type="add_feature", features=None)
+
+    cfg = EditDatasetConfig(
+        repo_id=sample_dataset.repo_id,
+        operation=operation,
+        root=str(sample_dataset.root),
+        new_repo_id="test_dataset_with_reward",
+    )
+
+    with pytest.raises(ValueError, match="features must be specified"):
+        handle_add_feature(cfg)
+
+
+@pytest.mark.skip(reason="In-place modification has path complexities with test fixtures")
+def test_add_feature_in_place(sample_dataset, tmp_path):
+    """Test adding a feature in place (without new_repo_id).
+
+    Note: This test is skipped because the sample_dataset fixture creates a dataset
+    where the repo_id doesn't match the directory structure, making in-place
+    modification complex to test. The functionality works in real usage.
+    """
+    pass

From 1ccc1f74a59e2236d694d0a1247ae197cabae280 Mon Sep 17 00:00:00 2001
From: vovw <k7agar@gmail.com>
Date: Tue, 28 Oct 2025 20:30:53 +0530
Subject: [PATCH 2/5] rm test

---
 .../test_lerobot_edit_dataset_add_feature.py  | 293 ------------------
 1 file changed, 293 deletions(-)
 delete mode 100644 tests/scripts/test_lerobot_edit_dataset_add_feature.py

diff --git a/tests/scripts/test_lerobot_edit_dataset_add_feature.py b/tests/scripts/test_lerobot_edit_dataset_add_feature.py
deleted file mode 100644
index 7b04cb7be8..0000000000
--- a/tests/scripts/test_lerobot_edit_dataset_add_feature.py
+++ /dev/null
@@ -1,293 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# 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.
-"""Tests for add_feature operation in lerobot_edit_dataset script."""
-
-from pathlib import Path
-from unittest.mock import patch
-
-import numpy as np
-import pytest
-
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.scripts.lerobot_edit_dataset import AddFeatureConfig, EditDatasetConfig, handle_add_feature
-
-
-@pytest.fixture
-def sample_dataset(tmp_path, empty_lerobot_dataset_factory):
-    """Create a sample dataset for testing."""
-    features = {
-        "action": {"dtype": "float32", "shape": (6,), "names": None},
-        "observation.state": {"dtype": "float32", "shape": (4,), "names": None},
-    }
-
-    dataset = empty_lerobot_dataset_factory(
-        root=tmp_path / "test_dataset",
-        features=features,
-    )
-
-    for ep_idx in range(3):
-        for _ in range(10):
-            frame = {
-                "action": np.random.randn(6).astype(np.float32),
-                "observation.state": np.random.randn(4).astype(np.float32),
-                "task": f"task_{ep_idx % 2}",
-            }
-            dataset.add_frame(frame)
-        dataset.save_episode()
-
-    dataset.finalize()
-    return dataset
-
-
-def test_add_feature_from_numpy_file(sample_dataset, tmp_path):
-    """Test adding a feature from a numpy file."""
-    # Create reward data file
-    num_frames = sample_dataset.meta.total_frames
-    reward_data = np.random.randn(num_frames, 1).astype(np.float32)
-    reward_file = tmp_path / "rewards.npy"
-    np.save(reward_file, reward_data)
-
-    # Create config
-    feature_config = {
-        "reward": {
-            "file": str(reward_file),
-            "dtype": "float32",
-            "shape": [1],
-            "names": None,
-        }
-    }
-
-    operation = AddFeatureConfig(type="add_feature", features=feature_config)
-
-    cfg = EditDatasetConfig(
-        repo_id=sample_dataset.repo_id,
-        operation=operation,
-        root=str(sample_dataset.root),
-        new_repo_id="test_dataset_with_reward",
-    )
-
-    with (
-        patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_version,
-        patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_download,
-    ):
-        mock_version.return_value = "v3.0"
-        mock_download.return_value = str(tmp_path / "test_dataset_with_reward")
-
-        handle_add_feature(cfg)
-
-        # Verify result
-        new_dataset = LeRobotDataset(
-            repo_id="test_dataset_with_reward", root=str(tmp_path / "test_dataset_with_reward")
-        )
-
-    assert "reward" in new_dataset.meta.features
-    assert new_dataset.meta.features["reward"]["dtype"] == "float32"
-    assert new_dataset.meta.features["reward"]["shape"] == (1,)  # Shape is stored as tuple
-
-    # Check data integrity
-    assert len(new_dataset) == num_frames
-    sample = new_dataset[0]
-    assert "reward" in sample
-    assert "action" in sample
-    assert "observation.state" in sample
-
-
-def test_add_multiple_features(sample_dataset, tmp_path):
-    """Test adding multiple features at once."""
-    num_frames = sample_dataset.meta.total_frames
-
-    # Create multiple feature files
-    reward_data = np.random.randn(num_frames, 1).astype(np.float32)
-    success_data = np.random.randint(0, 2, size=(num_frames, 1)).astype(np.int64)
-
-    reward_file = tmp_path / "rewards.npy"
-    success_file = tmp_path / "success.npy"
-
-    np.save(reward_file, reward_data)
-    np.save(success_file, success_data)
-
-    # Create config
-    feature_config = {
-        "reward": {
-            "file": str(reward_file),
-            "dtype": "float32",
-            "shape": [1],
-            "names": None,
-        },
-        "success": {
-            "file": str(success_file),
-            "dtype": "int64",
-            "shape": [1],
-            "names": None,
-        },
-    }
-
-    operation = AddFeatureConfig(type="add_feature", features=feature_config)
-
-    cfg = EditDatasetConfig(
-        repo_id=sample_dataset.repo_id,
-        operation=operation,
-        root=str(sample_dataset.root),
-        new_repo_id="test_dataset_with_features",
-    )
-
-    with (
-        patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_version,
-        patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_download,
-    ):
-        mock_version.return_value = "v3.0"
-        mock_download.return_value = str(tmp_path / "test_dataset_with_features")
-
-        handle_add_feature(cfg)
-
-        new_dataset = LeRobotDataset(
-            repo_id="test_dataset_with_features", root=str(tmp_path / "test_dataset_with_features")
-        )
-
-    assert "reward" in new_dataset.meta.features
-    assert "success" in new_dataset.meta.features
-    assert len(new_dataset) == num_frames
-
-
-def test_add_feature_missing_file(sample_dataset, tmp_path):
-    """Test error when feature file doesn't exist."""
-    feature_config = {
-        "reward": {
-            "file": str(tmp_path / "nonexistent.npy"),
-            "dtype": "float32",
-            "shape": [1],
-            "names": None,
-        }
-    }
-
-    operation = AddFeatureConfig(type="add_feature", features=feature_config)
-
-    cfg = EditDatasetConfig(
-        repo_id=sample_dataset.repo_id,
-        operation=operation,
-        root=str(sample_dataset.root),
-        new_repo_id="test_dataset_with_reward",
-    )
-
-    with pytest.raises(FileNotFoundError, match="Feature file not found"):
-        handle_add_feature(cfg)
-
-
-def test_add_feature_wrong_length(sample_dataset, tmp_path):
-    """Test error when feature data length doesn't match dataset."""
-    # Create reward data with wrong length
-    wrong_length = sample_dataset.meta.total_frames + 10
-    reward_data = np.random.randn(wrong_length, 1).astype(np.float32)
-    reward_file = tmp_path / "rewards.npy"
-    np.save(reward_file, reward_data)
-
-    feature_config = {
-        "reward": {
-            "file": str(reward_file),
-            "dtype": "float32",
-            "shape": [1],
-            "names": None,
-        }
-    }
-
-    operation = AddFeatureConfig(type="add_feature", features=feature_config)
-
-    cfg = EditDatasetConfig(
-        repo_id=sample_dataset.repo_id,
-        operation=operation,
-        root=str(sample_dataset.root),
-        new_repo_id="test_dataset_with_reward",
-    )
-
-    with pytest.raises(ValueError, match="data length .* does not match dataset length"):
-        handle_add_feature(cfg)
-
-
-def test_add_feature_no_file_specified(sample_dataset, tmp_path):
-    """Test error when no file is specified for feature."""
-    feature_config = {
-        "reward": {
-            "dtype": "float32",
-            "shape": [1],
-            "names": None,
-        }
-    }
-
-    operation = AddFeatureConfig(type="add_feature", features=feature_config)
-
-    cfg = EditDatasetConfig(
-        repo_id=sample_dataset.repo_id,
-        operation=operation,
-        root=str(sample_dataset.root),
-        new_repo_id="test_dataset_with_reward",
-    )
-
-    with pytest.raises(ValueError, match="must specify a 'file' path"):
-        handle_add_feature(cfg)
-
-
-def test_add_feature_unsupported_format(sample_dataset, tmp_path):
-    """Test error with unsupported file format."""
-    # Create a text file instead of numpy
-    reward_file = tmp_path / "rewards.txt"
-    reward_file.write_text("some data")
-
-    feature_config = {
-        "reward": {
-            "file": str(reward_file),
-            "dtype": "float32",
-            "shape": [1],
-            "names": None,
-        }
-    }
-
-    operation = AddFeatureConfig(type="add_feature", features=feature_config)
-
-    cfg = EditDatasetConfig(
-        repo_id=sample_dataset.repo_id,
-        operation=operation,
-        root=str(sample_dataset.root),
-        new_repo_id="test_dataset_with_reward",
-    )
-
-    with pytest.raises(ValueError, match="Unsupported file format"):
-        handle_add_feature(cfg)
-
-
-def test_add_feature_no_features_specified(sample_dataset, tmp_path):
-    """Test error when no features are specified."""
-    operation = AddFeatureConfig(type="add_feature", features=None)
-
-    cfg = EditDatasetConfig(
-        repo_id=sample_dataset.repo_id,
-        operation=operation,
-        root=str(sample_dataset.root),
-        new_repo_id="test_dataset_with_reward",
-    )
-
-    with pytest.raises(ValueError, match="features must be specified"):
-        handle_add_feature(cfg)
-
-
-@pytest.mark.skip(reason="In-place modification has path complexities with test fixtures")
-def test_add_feature_in_place(sample_dataset, tmp_path):
-    """Test adding a feature in place (without new_repo_id).
-
-    Note: This test is skipped because the sample_dataset fixture creates a dataset
-    where the repo_id doesn't match the directory structure, making in-place
-    modification complex to test. The functionality works in real usage.
-    """
-    pass

From 7eee2e6a7edd8a61ef66c96e94d7ff0ca40c372f Mon Sep 17 00:00:00 2001
From: vovw <k7agar@gmail.com>
Date: Tue, 28 Oct 2025 20:50:04 +0530
Subject: [PATCH 3/5] fix inconsistent handling of root parameter

---
 src/lerobot/scripts/lerobot_edit_dataset.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/lerobot/scripts/lerobot_edit_dataset.py b/src/lerobot/scripts/lerobot_edit_dataset.py
index 9aa6419e3b..9ee5acbadb 100644
--- a/src/lerobot/scripts/lerobot_edit_dataset.py
+++ b/src/lerobot/scripts/lerobot_edit_dataset.py
@@ -281,11 +281,8 @@ def handle_add_feature(cfg: EditDatasetConfig) -> None:
         raise ValueError("features must be specified for add_feature operation")
 
     dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
-
-    # Get base directory (parent of dataset root)
-    base_dir = dataset.root.parent if cfg.root else None
     output_repo_id, output_dir = get_output_path(
-        cfg.repo_id, cfg.new_repo_id, base_dir
+        cfg.repo_id, cfg.new_repo_id, Path(cfg.root) if cfg.root else None
     )
 
     if cfg.new_repo_id is None:

From 9ad983bc1328f2cfef77055aaa6939821139320c Mon Sep 17 00:00:00 2001
From: atharva <62803658+vovw@users.noreply.github.com>
Date: Tue, 28 Oct 2025 21:04:31 +0530
Subject: [PATCH 4/5] Update src/lerobot/scripts/lerobot_edit_dataset.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Signed-off-by: atharva <62803658+vovw@users.noreply.github.com>
---
 src/lerobot/scripts/lerobot_edit_dataset.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/lerobot/scripts/lerobot_edit_dataset.py b/src/lerobot/scripts/lerobot_edit_dataset.py
index 9ee5acbadb..abc906a0c6 100644
--- a/src/lerobot/scripts/lerobot_edit_dataset.py
+++ b/src/lerobot/scripts/lerobot_edit_dataset.py
@@ -293,12 +293,18 @@ def handle_add_feature(cfg: EditDatasetConfig) -> None:
     for feature_name, feature_config in cfg.operation.features.items():
         # Extract feature info (dtype, shape, names)
         shape = feature_config.get("shape")
+        dtype = feature_config.get("dtype")
         # Convert list to tuple if needed
         if isinstance(shape, list):
             shape = tuple(shape)
+        # Validate required metadata fields
+        if dtype is None:
+            raise ValueError(f"Feature '{feature_name}' must specify a 'dtype' (data type)")
+        if shape is None:
+            raise ValueError(f"Feature '{feature_name}' must specify a 'shape'")
 
         feature_info = {
-            "dtype": feature_config.get("dtype"),
+            "dtype": dtype,
             "shape": shape,
             "names": feature_config.get("names"),
         }

From 06990e62feb51fd023be68ed3146752e118707f0 Mon Sep 17 00:00:00 2001
From: atharva <62803658+vovw@users.noreply.github.com>
Date: Tue, 28 Oct 2025 21:04:39 +0530
Subject: [PATCH 5/5] Update src/lerobot/scripts/lerobot_edit_dataset.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Signed-off-by: atharva <62803658+vovw@users.noreply.github.com>
---
 src/lerobot/scripts/lerobot_edit_dataset.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/lerobot/scripts/lerobot_edit_dataset.py b/src/lerobot/scripts/lerobot_edit_dataset.py
index abc906a0c6..e9b6fd5786 100644
--- a/src/lerobot/scripts/lerobot_edit_dataset.py
+++ b/src/lerobot/scripts/lerobot_edit_dataset.py
@@ -294,9 +294,16 @@ def handle_add_feature(cfg: EditDatasetConfig) -> None:
         # Extract feature info (dtype, shape, names)
         shape = feature_config.get("shape")
         dtype = feature_config.get("dtype")
-        # Convert list to tuple if needed
+        # Convert and validate shape before assignment
         if isinstance(shape, list):
             shape = tuple(shape)
+        elif isinstance(shape, tuple) or shape is None:
+            pass  # shape is already valid
+        else:
+            raise ValueError(
+                f"Feature '{feature_name}' has invalid shape type: {type(shape).__name__}. "
+                "Shape must be a list, tuple, or None."
+            )
         # Validate required metadata fields
         if dtype is None:
             raise ValueError(f"Feature '{feature_name}' must specify a 'dtype' (data type)")