feat(task modification precedence): Improving task modification precedence so that all modes can be used in a single run. Adapting tests accordinginly.

src/lerobot/datasets/dataset_tools.py

@@ -1444,9 +1444,8 @@ def modify_tasks(
     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. `task_replacements` can be combined with `episode_tasks`, with
-    episode-specific overrides taking precedence.
+    You can combine all of them: `new_task` sets the default, while `episode_tasks` and 
+    `task_replacements` override specific episodes, with `episode_tasks` taking precedence.
 
     The dataset is modified in-place, updating only the task-related files:
     - meta/tasks.parquet
@@ -1485,16 +1484,10 @@ def modify_tasks(
                 task_replacements={"Pick up the cube": "Lift the cube"}
             )
     """
-    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:
+    if not new_task and not episode_tasks and not task_replacements:
         raise ValueError("Must specify at least one of new_task, episode_tasks, or task_replacements")
 
-    if episode_tasks is not None:
+    if episode_tasks:
         valid_indices = set(range(dataset.meta.total_episodes))
         invalid = set(episode_tasks.keys()) - valid_indices
         if invalid:
@@ -1504,12 +1497,8 @@ 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"]
-        }
+    if task_replacements:
+        current_tasks = set(dataset.meta.tasks.index)
         invalid_tasks = set(task_replacements) - current_tasks
         if invalid_tasks:
             raise ValueError(f"Task replacements reference unknown tasks: {sorted(invalid_tasks)}")
@@ -1518,19 +1507,19 @@ def modify_tasks(
     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]
+        original_task = original_tasks[0] if original_tasks else None
 
         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:
+        elif new_task:
+            episode_to_task[ep_idx] = new_task
+        elif original_task:
             # Keep original task if not overridden and no default provided
             episode_to_task[ep_idx] = original_task
+        else:
+            raise ValueError(f"Episode {ep_idx} has no task; provide new_task or episode_tasks")
 
     # Collect all unique tasks and create new task mapping
     unique_tasks = sorted(set(episode_to_task.values()))

src/lerobot/scripts/lerobot_edit_dataset.py

@@ -560,9 +560,6 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
     episode_tasks_raw = cfg.operation.episode_tasks
     task_replacements = cfg.operation.task_replacements
 
-    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"

tests/datasets/test_dataset_tools.py

@@ -1161,6 +1161,20 @@ def test_modify_tasks_replacements_with_episode_overrides(sample_dataset):
     assert len(modified_dataset.meta.tasks) == 3
 
 
+def test_modify_tasks_default_task_and_replacements(sample_dataset):
+    """Test that new_task acts as the default for episodes not matched by task_replacements."""
+    modified_dataset = modify_tasks(
+        sample_dataset,
+        new_task="Default task",
+        task_replacements={"task_0": "Pick the cube"},
+    )
+
+    for ep_idx in range(5):
+        expected_task = "Pick the cube" if ep_idx % 2 == 0 else "Default task"
+        assert modified_dataset.meta.episodes[ep_idx]["tasks"][0] == expected_task
+    assert len(modified_dataset.meta.tasks) == 2
+
+
 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, episode_tasks, or task_replacements"):
@@ -1179,16 +1193,6 @@ def test_modify_tasks_invalid_task_replacements(sample_dataset):
         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

@@ -1,83 +0,0 @@
-#!/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)