refactor(dataset): update imports across the codebase

tests/datasets/test_dataset_tools.py

@@ -1020,7 +1020,7 @@ def test_modify_features_preserves_file_structure(sample_dataset, tmp_path):
 
         # Get original chunk/file indices from first episode
         if train_dataset.meta.episodes is None:
-            from lerobot.datasets.utils import load_episodes
+            from lerobot.datasets.io_utils import load_episodes
 
             train_dataset.meta.episodes = load_episodes(train_dataset.meta.root)
         original_chunk_indices = [ep["data/chunk_index"] for ep in train_dataset.meta.episodes]
@@ -1040,7 +1040,7 @@ def test_modify_features_preserves_file_structure(sample_dataset, tmp_path):
 
         # Check that chunk/file indices are preserved
         if modified_dataset.meta.episodes is None:
-            from lerobot.datasets.utils import load_episodes
+            from lerobot.datasets.io_utils import load_episodes
 
             modified_dataset.meta.episodes = load_episodes(modified_dataset.meta.root)
         new_chunk_indices = [ep["data/chunk_index"] for ep in modified_dataset.meta.episodes]
@@ -1194,7 +1194,7 @@ def test_modify_tasks_in_place(sample_dataset):
 
 def test_modify_tasks_keeps_original_when_not_overridden(sample_dataset):
     """Test that original tasks are kept when using episode_tasks without new_task."""
-    from lerobot.datasets.utils import load_episodes
+    from lerobot.datasets.io_utils import load_episodes
 
     # Ensure episodes metadata is loaded
     if sample_dataset.meta.episodes is None:

tests/datasets/test_dataset_utils.py

@@ -19,7 +19,9 @@ import torch
 from datasets import Dataset
 from huggingface_hub import DatasetCard
 
-from lerobot.datasets.utils import combine_feature_dicts, create_lerobot_dataset_card, hf_transform_to_torch
+from lerobot.datasets.feature_utils import combine_feature_dicts
+from lerobot.datasets.io_utils import hf_transform_to_torch
+from lerobot.datasets.utils import create_lerobot_dataset_card
 from lerobot.utils.constants import ACTION, OBS_IMAGES
 
 

tests/datasets/test_datasets.py

@@ -29,20 +29,19 @@ import lerobot
 from lerobot.configs.default import DatasetConfig
 from lerobot.configs.train import TrainPipelineConfig
 from lerobot.datasets.factory import make_dataset
+from lerobot.datasets.feature_utils import get_hf_features_from_features, hw_to_dataset_features
 from lerobot.datasets.image_writer import image_array_to_pil_image
+from lerobot.datasets.io_utils import hf_transform_to_torch
 from lerobot.datasets.lerobot_dataset import (
     LeRobotDataset,
-    MultiLeRobotDataset,
     _encode_video_worker,
 )
+from lerobot.datasets.multi_dataset import MultiLeRobotDataset
 from lerobot.datasets.utils import (
     DEFAULT_CHUNK_SIZE,
     DEFAULT_DATA_FILE_SIZE_IN_MB,
     DEFAULT_VIDEO_FILE_SIZE_IN_MB,
     create_branch,
-    get_hf_features_from_features,
-    hf_transform_to_torch,
-    hw_to_dataset_features,
 )
 from lerobot.datasets.video_utils import VALID_VIDEO_CODECS
 from lerobot.envs.factory import make_env_config
@@ -1329,7 +1328,7 @@ def test_frames_in_current_file_calculation(tmp_path, empty_lerobot_dataset_fact
 
     dataset.finalize()
 
-    from lerobot.datasets.utils import load_episodes
+    from lerobot.datasets.io_utils import load_episodes
 
     dataset.meta.episodes = load_episodes(dataset.root)
     assert dataset.meta.episodes is not None

tests/datasets/test_delta_timestamps.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import pytest
 
-from lerobot.datasets.utils import (
+from lerobot.datasets.feature_utils import (
     check_delta_timestamps,
     get_delta_indices,
 )

tests/datasets/test_sampler.py

@@ -19,10 +19,10 @@ import pytest
 import torch
 from datasets import Dataset
 
-from lerobot.datasets.sampler import EpisodeAwareSampler
-from lerobot.datasets.utils import (
+from lerobot.datasets.io_utils import (
     hf_transform_to_torch,
 )
+from lerobot.datasets.sampler import EpisodeAwareSampler
 
 
 def calculate_episode_data_index(hf_dataset: Dataset) -> dict[str, torch.Tensor]:

tests/fixtures/dataset_factories.py

@@ -26,7 +26,10 @@ import pytest
 import torch
 from datasets import Dataset
 
-from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.datasets.dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
+from lerobot.datasets.feature_utils import get_hf_features_from_features
+from lerobot.datasets.io_utils import hf_transform_to_torch
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.datasets.utils import (
     DEFAULT_CHUNK_SIZE,
     DEFAULT_DATA_FILE_SIZE_IN_MB,
@@ -35,8 +38,6 @@ from lerobot.datasets.utils import (
     DEFAULT_VIDEO_FILE_SIZE_IN_MB,
     DEFAULT_VIDEO_PATH,
     flatten_dict,
-    get_hf_features_from_features,
-    hf_transform_to_torch,
 )
 from lerobot.datasets.video_utils import encode_video_frames
 from tests.fixtures.constants import (

tests/fixtures/files.py

@@ -20,17 +20,19 @@ import pandas as pd
 import pytest
 from datasets import Dataset
 
-from lerobot.datasets.utils import (
-    DEFAULT_CHUNK_SIZE,
-    DEFAULT_DATA_FILE_SIZE_IN_MB,
-    DEFAULT_DATA_PATH,
+from lerobot.datasets.io_utils import (
     get_hf_dataset_size_in_mb,
-    update_chunk_file_indices,
     write_episodes,
     write_info,
     write_stats,
     write_tasks,
 )
+from lerobot.datasets.utils import (
+    DEFAULT_CHUNK_SIZE,
+    DEFAULT_DATA_FILE_SIZE_IN_MB,
+    DEFAULT_DATA_PATH,
+    update_chunk_file_indices,
+)
 
 
 def write_hf_dataset(

tests/policies/test_policies.py

@@ -28,7 +28,8 @@ from lerobot.configs.default import DatasetConfig
 from lerobot.configs.train import TrainPipelineConfig
 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.datasets.factory import make_dataset
-from lerobot.datasets.utils import cycle, dataset_to_policy_features
+from lerobot.datasets.feature_utils import dataset_to_policy_features
+from lerobot.datasets.utils import cycle
 from lerobot.envs.factory import make_env, make_env_config
 from lerobot.envs.utils import preprocess_observation
 from lerobot.optim.factory import make_optimizer_and_scheduler