diff --git a/src/lerobot/datasets/dataset_tools.py b/src/lerobot/datasets/dataset_tools.py index 31e075d7c..13b65f43f 100644 --- a/src/lerobot/datasets/dataset_tools.py +++ b/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())) diff --git a/src/lerobot/scripts/lerobot_edit_dataset.py b/src/lerobot/scripts/lerobot_edit_dataset.py index 42dce438f..775b85196 100644 --- a/src/lerobot/scripts/lerobot_edit_dataset.py +++ b/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}") diff --git a/tests/datasets/test_dataset_tools.py b/tests/datasets/test_dataset_tools.py index c19e7f41f..ec62ceb28 100644 --- a/tests/datasets/test_dataset_tools.py +++ b/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"} diff --git a/tests/scripts/test_edit_dataset_modify_tasks.py b/tests/scripts/test_edit_dataset_modify_tasks.py new file mode 100644 index 000000000..545b074f5 --- /dev/null +++ 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) diff --git a/tests/scripts/test_edit_dataset_parsing.py b/tests/scripts/test_edit_dataset_parsing.py index 22a3c1be2..c9bfb33be 100644 --- a/tests/scripts/test_edit_dataset_parsing.py +++ b/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."""