Support task replacement mappings in modify_tasks

Part of #2326.

Signed-off-by: 陈伟 <woshei0a0a0a@qq.com>

tests/scripts/test_edit_dataset_modify_tasks.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python
+
+# Copyright 2026 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.
+
+import numpy as np
+import pytest
+
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.scripts.lerobot_edit_dataset import EditDatasetConfig, ModifyTasksConfig, handle_modify_tasks
+
+
+@pytest.fixture
+def sample_dataset(tmp_path, empty_lerobot_dataset_factory):
+    features = {
+        "action": {"dtype": "float32", "shape": (6,), "names": None},
+        "observation.state": {"dtype": "float32", "shape": (4,), "names": None},
+        "observation.images.top": {"dtype": "image", "shape": (224, 224, 3), "names": None},
+    }
+
+    dataset = empty_lerobot_dataset_factory(
+        root=tmp_path / "test_dataset",
+        features=features,
+    )
+
+    for ep_idx in range(5):
+        for _ in range(10):
+            frame = {
+                "action": np.random.randn(6).astype(np.float32),
+                "observation.state": np.random.randn(4).astype(np.float32),
+                "observation.images.top": np.random.randint(0, 255, size=(224, 224, 3), dtype=np.uint8),
+                "task": f"task_{ep_idx % 2}",
+            }
+            dataset.add_frame(frame)
+        dataset.save_episode()
+
+    dataset.finalize()
+    return dataset
+
+
+def test_handle_modify_tasks_with_replacements(sample_dataset):
+    cfg = EditDatasetConfig(
+        repo_id=sample_dataset.repo_id,
+        root=str(sample_dataset.root),
+        operation=ModifyTasksConfig(
+            task_replacements={
+                "task_0": "Pick the cube",
+                "task_1": "Place the cube",
+            }
+        ),
+    )
+
+    handle_modify_tasks(cfg)
+
+    modified_dataset = LeRobotDataset(cfg.repo_id, root=sample_dataset.root)
+    assert modified_dataset.meta.episodes[0]["tasks"][0] == "Pick the cube"
+    assert modified_dataset.meta.episodes[1]["tasks"][0] == "Place the cube"
+    assert len(modified_dataset.meta.tasks) == 2
+
+
+def test_handle_modify_tasks_rejects_default_task_and_replacements(sample_dataset):
+    cfg = EditDatasetConfig(
+        repo_id=sample_dataset.repo_id,
+        root=str(sample_dataset.root),
+        operation=ModifyTasksConfig(
+            new_task="Default task",
+            task_replacements={"task_0": "Pick the cube"},
+        ),
+    )
+
+    with pytest.raises(ValueError, match="Cannot combine new_task with task_replacements"):
+        handle_modify_tasks(cfg)

tests/scripts/test_edit_dataset_parsing.py

@@ -105,6 +105,23 @@ class TestOperationTypeParsing:
         resolved_name = OperationConfig.get_choice_name(type(cfg.operation))
         assert resolved_name == type_name
 
+    def test_modify_tasks_replacements_args_parse(self):
+        cfg = parse_cfg(
+            [
+                "--repo_id",
+                "test/repo",
+                "--operation.type",
+                "modify_tasks",
+                "--operation.task_replacements",
+                '{"task_0": "pick cube", "task_1": "place cube"}',
+            ]
+        )
+        assert isinstance(cfg.operation, ModifyTasksConfig)
+        assert cfg.operation.task_replacements == {
+            "task_0": "pick cube",
+            "task_1": "place cube",
+        }
+
 
 class TestDepthEncoderParsing:
     """Test that the depth encoder is exposed and parsed for video operations."""