fix(features copy): adding deepcopy on LeRobot dataset features to avoid shallow copy leaks (#3826)

* fix(features copy): adding deepcopy on LeRobot dataset features to avoid shallow copy leaks

* tests(test): adding new test
This commit is contained in:
Caroline Pascal
2026-06-16 17:58:59 +02:00
committed by GitHub
parent 58ccc01508
commit 287c823f13
3 changed files with 22 additions and 3 deletions
+2 -1
View File
@@ -15,6 +15,7 @@
# limitations under the License. # limitations under the License.
import contextlib import contextlib
from collections.abc import Callable from collections.abc import Callable
from copy import deepcopy
from pathlib import Path from pathlib import Path
import numpy as np import numpy as np
@@ -709,7 +710,7 @@ class LeRobotDatasetMetadata:
obj.root.mkdir(parents=True, exist_ok=False) obj.root.mkdir(parents=True, exist_ok=False)
features = {**features, **DEFAULT_FEATURES} features = {**deepcopy(features), **DEFAULT_FEATURES}
_validate_feature_names(features) _validate_feature_names(features)
obj.tasks = None obj.tasks = None
+4 -1
View File
@@ -27,6 +27,7 @@ import logging
import shutil import shutil
from collections.abc import Callable from collections.abc import Callable
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
from copy import deepcopy
from pathlib import Path from pathlib import Path
import datasets import datasets
@@ -1101,7 +1102,9 @@ def _copy_episodes_metadata_and_stats(
if dst_meta.video_keys and src_dataset.meta.video_keys: if dst_meta.video_keys and src_dataset.meta.video_keys:
for key in dst_meta.video_keys: for key in dst_meta.video_keys:
if key in src_dataset.meta.features: if key in src_dataset.meta.features:
dst_meta.info.features[key]["info"] = src_dataset.meta.info.features[key].get("info", {}) dst_meta.info.features[key]["info"] = deepcopy(
src_dataset.meta.info.features[key].get("info", {})
)
write_info(dst_meta.info, dst_meta.root) write_info(dst_meta.info, dst_meta.root)
+16 -1
View File
@@ -51,7 +51,7 @@ from lerobot.robots import make_robot_from_config
from lerobot.transforms import ImageTransforms, ImageTransformsConfig from lerobot.transforms import ImageTransforms, ImageTransformsConfig
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGES, OBS_STATE, OBS_STR, REWARD from lerobot.utils.constants import ACTION, DONE, OBS_IMAGES, OBS_STATE, OBS_STR, REWARD
from lerobot.utils.feature_utils import hw_to_dataset_features from lerobot.utils.feature_utils import hw_to_dataset_features
from tests.fixtures.constants import DUMMY_CHW, DUMMY_HWC, DUMMY_REPO_ID from tests.fixtures.constants import DUMMY_CHW, DUMMY_HWC, DUMMY_MOTOR_FEATURES, DUMMY_REPO_ID
from tests.mocks.mock_robot import MockRobotConfig from tests.mocks.mock_robot import MockRobotConfig
from tests.utils import require_x86_64_kernel from tests.utils import require_x86_64_kernel
@@ -133,6 +133,21 @@ def test_dataset_feature_with_forward_slash_raises_error():
) )
def test_create_does_not_mutate_input_features(tmp_path, empty_lerobot_dataset_factory):
# ``create`` must deep-copy features so a dataset built from another's features stays independent.
dataset = empty_lerobot_dataset_factory(
root=tmp_path / "ds1", features=DUMMY_MOTOR_FEATURES, use_videos=False
)
dataset_copy = empty_lerobot_dataset_factory(
root=tmp_path / "ds2", features=dataset.meta.features, use_videos=False
)
original_shape = dataset.meta.info.features["state"]["shape"]
dataset_copy.meta.info.features["state"]["shape"] = (999,)
assert dataset.meta.info.features["state"]["shape"] == original_shape
def test_add_frame_missing_task(tmp_path, empty_lerobot_dataset_factory): def test_add_frame_missing_task(tmp_path, empty_lerobot_dataset_factory):
features = {"state": {"dtype": "float32", "shape": (1,), "names": None}} features = {"state": {"dtype": "float32", "shape": (1,), "names": None}}
dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features) dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)