Support task replacement mappings in modify_tasks

Part of #2326.

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

src/lerobot/datasets/dataset_tools.py

@@ -1435,15 +1435,18 @@ def modify_tasks(
     dataset: LeRobotDataset,
     new_task: str | None = None,
     episode_tasks: dict[int, str] | None = None,
+    task_replacements: dict[str, str] | None = None,
 ) -> LeRobotDataset:
     """Modify tasks in a LeRobotDataset.
 
     This function allows you to either:
     1. Set a single task for the entire dataset (using `new_task`)
     2. Set specific tasks for specific episodes (using `episode_tasks`)
+    3. Replace existing task strings wherever they appear (using `task_replacements`)
 
     You can combine both: `new_task` sets the default, and `episode_tasks` overrides
-    specific episodes.
+    specific episodes. `task_replacements` can be combined with `episode_tasks`, with
+    episode-specific overrides taking precedence.
 
     The dataset is modified in-place, updating only the task-related files:
     - meta/tasks.parquet
@@ -1475,9 +1478,21 @@ def modify_tasks(
                 new_task="Default task",
                 episode_tasks={5: "Special task for episode 5"}
             )
+
+        Replace existing task strings in-place:
+            dataset = modify_tasks(
+                dataset,
+                task_replacements={"Pick up the cube": "Lift the cube"}
+            )
     """
-    if new_task is None and episode_tasks is None:
-        raise ValueError("Must specify at least one of new_task or episode_tasks")
+    episode_tasks = episode_tasks or None
+    task_replacements = task_replacements or None
+
+    if new_task is not None and task_replacements is not None:
+        raise ValueError("Cannot combine new_task with task_replacements")
+
+    if new_task is None and episode_tasks is None and task_replacements is None:
+        raise ValueError("Must specify at least one of new_task, episode_tasks, or task_replacements")
 
     if episode_tasks is not None:
         valid_indices = set(range(dataset.meta.total_episodes))
@@ -1489,19 +1504,33 @@ def modify_tasks(
     if dataset.meta.episodes is None:
         dataset.meta.episodes = load_episodes(dataset.root)
 
+    if task_replacements is not None:
+        current_tasks = {
+            dataset.meta.episodes[ep_idx]["tasks"][0]
+            for ep_idx in range(dataset.meta.total_episodes)
+            if dataset.meta.episodes[ep_idx]["tasks"]
+        }
+        invalid_tasks = set(task_replacements) - current_tasks
+        if invalid_tasks:
+            raise ValueError(f"Task replacements reference unknown tasks: {sorted(invalid_tasks)}")
+
     # Build the mapping from episode index to task string
     episode_to_task: dict[int, str] = {}
     for ep_idx in range(dataset.meta.total_episodes):
+        original_tasks = dataset.meta.episodes[ep_idx]["tasks"]
+        if not original_tasks:
+            raise ValueError(f"Episode {ep_idx} has no tasks and no default task was provided")
+        original_task = original_tasks[0]
+
         if episode_tasks and ep_idx in episode_tasks:
             episode_to_task[ep_idx] = episode_tasks[ep_idx]
         elif new_task is not None:
             episode_to_task[ep_idx] = new_task
+        elif task_replacements and original_task in task_replacements:
+            episode_to_task[ep_idx] = task_replacements[original_task]
         else:
             # Keep original task if not overridden and no default provided
-            original_tasks = dataset.meta.episodes[ep_idx]["tasks"]
-            if not original_tasks:
-                raise ValueError(f"Episode {ep_idx} has no tasks and no default task was provided")
-            episode_to_task[ep_idx] = original_tasks[0]
+            episode_to_task[ep_idx] = original_task
 
     # Collect all unique tasks and create new task mapping
     unique_tasks = sorted(set(episode_to_task.values()))

src/lerobot/scripts/lerobot_edit_dataset.py

@@ -127,6 +127,12 @@ Modify tasks - set default task with overrides for specific episodes (WARNING: m
         --operation.new_task "Default task" \
         --operation.episode_tasks '{"5": "Special task for episode 5"}'
 
+Modify tasks - replace existing task strings in-place (WARNING: modifies in-place):
+    lerobot-edit-dataset \
+        --repo_id lerobot/pusht \
+        --operation.type modify_tasks \
+        --operation.task_replacements '{"Pick up the red cube": "Lift the red cube"}'
+
 Convert image dataset to video format and save locally:
     lerobot-edit-dataset \
         --repo_id lerobot/pusht_image \
@@ -303,6 +309,7 @@ class RemoveFeatureConfig(OperationConfig):
 class ModifyTasksConfig(OperationConfig):
     new_task: str | None = None
     episode_tasks: dict[str, str] | None = None
+    task_replacements: dict[str, str] | None = None
 
 
 @OperationConfig.register_subclass("convert_image_to_video")
@@ -551,9 +558,15 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
 
     new_task = cfg.operation.new_task
     episode_tasks_raw = cfg.operation.episode_tasks
+    task_replacements = cfg.operation.task_replacements
 
-    if new_task is None and episode_tasks_raw is None:
-        raise ValueError("Must specify at least one of new_task or episode_tasks for modify_tasks operation")
+    if new_task is not None and task_replacements is not None:
+        raise ValueError("Cannot combine new_task with task_replacements for modify_tasks operation")
+
+    if new_task is None and episode_tasks_raw is None and task_replacements is None:
+        raise ValueError(
+            "Must specify at least one of new_task, episode_tasks, or task_replacements for modify_tasks operation"
+        )
 
     if cfg.new_repo_id is not None or cfg.new_root is not None:
         logging.warning(
@@ -573,11 +586,14 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
         logging.info(f"  Default task: '{new_task}'")
     if episode_tasks:
         logging.info(f"  Episode-specific tasks: {episode_tasks}")
+    if task_replacements:
+        logging.info(f"  Task replacements: {task_replacements}")
 
     modified_dataset = modify_tasks(
         dataset,
         new_task=new_task,
         episode_tasks=episode_tasks,
+        task_replacements=task_replacements,
     )
 
     logging.info(f"Dataset modified at {dataset.root}")

tests/datasets/test_dataset_tools.py

@@ -1125,9 +1125,45 @@ def test_modify_tasks_default_with_overrides(sample_dataset):
             assert ep_data["tasks"][0] == default_task
 
 
+def test_modify_tasks_replacements(sample_dataset):
+    """Test replacing task strings based on their current values."""
+    modified_dataset = modify_tasks(
+        sample_dataset,
+        task_replacements={
+            "task_0": "Pick the cube",
+            "task_1": "Place the cube",
+        },
+    )
+
+    assert len(modified_dataset.meta.tasks) == 2
+    assert "Pick the cube" in modified_dataset.meta.tasks.index
+    assert "Place the cube" in modified_dataset.meta.tasks.index
+
+    for ep_idx in range(5):
+        expected_task = "Pick the cube" if ep_idx % 2 == 0 else "Place the cube"
+        assert modified_dataset.meta.episodes[ep_idx]["tasks"][0] == expected_task
+
+
+def test_modify_tasks_replacements_with_episode_overrides(sample_dataset):
+    """Test that explicit episode overrides take precedence over replacements."""
+    modified_dataset = modify_tasks(
+        sample_dataset,
+        task_replacements={
+            "task_0": "Pick the cube",
+            "task_1": "Place the cube",
+        },
+        episode_tasks={1: "Inspect the cube"},
+    )
+
+    assert modified_dataset.meta.episodes[0]["tasks"][0] == "Pick the cube"
+    assert modified_dataset.meta.episodes[1]["tasks"][0] == "Inspect the cube"
+    assert modified_dataset.meta.episodes[3]["tasks"][0] == "Place the cube"
+    assert len(modified_dataset.meta.tasks) == 3
+
+
 def test_modify_tasks_no_task_specified(sample_dataset):
     """Test error when no task is specified."""
-    with pytest.raises(ValueError, match="Must specify at least one of new_task or episode_tasks"):
+    with pytest.raises(ValueError, match="Must specify at least one of new_task, episode_tasks, or task_replacements"):
         modify_tasks(sample_dataset)
 
 
@@ -1137,6 +1173,22 @@ def test_modify_tasks_invalid_episode_indices(sample_dataset):
         modify_tasks(sample_dataset, episode_tasks={10: "Task", 20: "Task"})
 
 
+def test_modify_tasks_invalid_task_replacements(sample_dataset):
+    """Test error when task replacements refer to unknown task strings."""
+    with pytest.raises(ValueError, match="Task replacements reference unknown tasks"):
+        modify_tasks(sample_dataset, task_replacements={"missing_task": "New task"})
+
+
+def test_modify_tasks_rejects_default_task_and_replacements(sample_dataset):
+    """Test that default-task assignment cannot be combined with find-and-replace."""
+    with pytest.raises(ValueError, match="Cannot combine new_task with task_replacements"):
+        modify_tasks(
+            sample_dataset,
+            new_task="Default task",
+            task_replacements={"task_0": "Pick the cube"},
+        )
+
+
 def test_modify_tasks_updates_info_json(sample_dataset):
     """Test that total_tasks is updated in info.json."""
     episode_tasks = {0: "Task A", 1: "Task B", 2: "Task C", 3: "Task A", 4: "Task B"}

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."""