Merge remote-tracking branch 'origin/main' into user/rcadene/2025_04_11_dataset_v3

tests/datasets/test_compute_stats.py

@@ -18,7 +18,7 @@ from unittest.mock import patch
 import numpy as np
 import pytest
 
-from lerobot.common.datasets.compute_stats import (
+from lerobot.datasets.compute_stats import (
     _assert_type_and_shape,
     aggregate_feature_stats,
     aggregate_stats,
@@ -61,7 +61,7 @@ def test_sample_indices():
     assert len(indices) == estimate_num_samples(10)
 
 
-@patch("lerobot.common.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy)
+@patch("lerobot.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy)
 def test_sample_images(mock_load):
     image_paths = [f"image_{i}.jpg" for i in range(100)]
     images = sample_images(image_paths)
@@ -144,9 +144,7 @@ def test_compute_episode_stats():
         "observation.state": {"dtype": "numeric"},
     }
 
-    with patch(
-        "lerobot.common.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy
-    ):
+    with patch("lerobot.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy):
         stats = compute_episode_stats(episode_data, features)
 
     assert "observation.image" in stats and "observation.state" in stats

tests/datasets/test_delta_timestamps.py

@@ -19,7 +19,7 @@ import pyarrow.compute as pc
 import pytest
 import torch
 
-from lerobot.common.datasets.utils import (
+from lerobot.datasets.utils import (
     check_delta_timestamps,
     check_timestamps_sync,
     get_delta_indices,

tests/datasets/test_image_transforms.py

@@ -16,6 +16,7 @@
 
 import pytest
 import torch
+from packaging import version
 from safetensors.torch import load_file
 from torchvision.transforms import v2
 from torchvision.transforms.v2 import functional as F  # noqa: N812
@@ -253,7 +254,14 @@ def test_backward_compatibility_single_transforms(
 
 
 @require_x86_64_kernel
+@pytest.mark.skipif(
+    version.parse(torch.__version__) < version.parse("2.7.0"),
+    reason="Test artifacts were generated with PyTorch >= 2.7.0 which has different multinomial behavior",
+)
 def test_backward_compatibility_default_config(img_tensor, default_transforms):
+    # NOTE: PyTorch versions have different randomness, it might break this test.
+    # See this PR: https://github.com/huggingface/lerobot/pull/1127.
+
     cfg = ImageTransformsConfig(enable=True)
     default_tf = ImageTransforms(cfg)
 

tests/datasets/test_image_writer.py

@@ -20,7 +20,7 @@ import numpy as np
 import pytest
 from PIL import Image
 
-from lerobot.common.datasets.image_writer import (
+from lerobot.datasets.image_writer import (
     AsyncImageWriter,
     image_array_to_pil_image,
     safe_stop_image_writer,

tests/datasets/test_online_buffer.py

@@ -20,7 +20,7 @@ import numpy as np
 import pytest
 import torch
 
-from lerobot.common.datasets.online_buffer import OnlineBuffer, compute_sampler_weights
+from lerobot.datasets.online_buffer import OnlineBuffer, compute_sampler_weights
 
 # Some constants for OnlineBuffer tests.
 data_key = "data"

tests/datasets/test_sampler.py

@@ -15,9 +15,9 @@
 # limitations under the License.
 from datasets import Dataset
 
-from lerobot.common.datasets.push_dataset_to_hub.utils import calculate_episode_data_index
-from lerobot.common.datasets.sampler import EpisodeAwareSampler
-from lerobot.common.datasets.utils import (
+from lerobot.datasets.push_dataset_to_hub.utils import calculate_episode_data_index
+from lerobot.datasets.sampler import EpisodeAwareSampler
+from lerobot.datasets.utils import (
     hf_transform_to_torch,
 )