diff --git a/src/lerobot/datasets/dataset_tools.py b/src/lerobot/datasets/dataset_tools.py index 13b65f43f..cc9fea20c 100644 --- a/src/lerobot/datasets/dataset_tools.py +++ b/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())) diff --git a/src/lerobot/scripts/lerobot_edit_dataset.py b/src/lerobot/scripts/lerobot_edit_dataset.py index 775b85196..6d0473f5b 100644 --- a/src/lerobot/scripts/lerobot_edit_dataset.py +++ b/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" diff --git a/tests/datasets/test_dataset_tools.py b/tests/datasets/test_dataset_tools.py index ec62ceb28..72928b1f0 100644 --- a/tests/datasets/test_dataset_tools.py +++ b/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"} diff --git a/tests/scripts/test_edit_dataset_modify_tasks.py b/tests/scripts/test_edit_dataset_modify_tasks.py deleted file mode 100644 index 545b074f5..000000000 --- a/tests/scripts/test_edit_dataset_modify_tasks.py +++ /dev/null @@ -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)