diff --git a/pyproject.toml b/pyproject.toml index 5377e3dac5..3f1d0d89cc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,6 +7,7 @@ requires-python = ">=3.11" license = { file = "LICENSE" } dependencies = [ "augmax>=0.3.4", + "datasets>=2.19.0,<4.0.0", "dm-tree>=0.1.8", "einops>=0.8.0", "equinox>=0.11.8", diff --git a/src/openpi/training/data_loader.py b/src/openpi/training/data_loader.py index e2ee7dd06b..fed9881d48 100644 --- a/src/openpi/training/data_loader.py +++ b/src/openpi/training/data_loader.py @@ -17,6 +17,8 @@ import openpi.transforms as _transforms T_co = TypeVar("T_co", covariant=True) +_SUPPORTED_LEROBOT_DATASET_VERSION = "v2.1" +_UNSUPPORTED_LEROBOT_DATASET_VERSION = (3, 0) class Dataset(Protocol[T_co]): @@ -127,6 +129,33 @@ def __len__(self) -> int: return self._num_samples +def _parse_lerobot_dataset_version(version: str | None) -> tuple[int, int] | None: + if version is None: + return None + + version_parts = version.lower().removeprefix("v").split(".") + try: + major = int(version_parts[0]) + minor = int(version_parts[1]) if len(version_parts) > 1 else 0 + except ValueError: + return None + + return major, minor + + +def check_lerobot_dataset_version(version: str | None, repo_id: str) -> None: + parsed_version = _parse_lerobot_dataset_version(version) + if parsed_version is None or parsed_version < _UNSUPPORTED_LEROBOT_DATASET_VERSION: + return + + raise ValueError( + f"Dataset {repo_id} uses LeRobot dataset format {version}. openpi currently supports format " + f"{_SUPPORTED_LEROBOT_DATASET_VERSION} (see the pinned lerobot dependency in pyproject.toml). " + f"Re-create the dataset with lerobot format {_SUPPORTED_LEROBOT_DATASET_VERSION}, or convert it back " + f"to {_SUPPORTED_LEROBOT_DATASET_VERSION}, before training." + ) + + def create_torch_dataset( data_config: _config.DataConfig, action_horizon: int, model_config: _model.BaseModelConfig ) -> Dataset: @@ -138,6 +167,7 @@ def create_torch_dataset( return FakeDataset(model_config, num_samples=1024) dataset_meta = lerobot_dataset.LeRobotDatasetMetadata(repo_id) + check_lerobot_dataset_version(dataset_meta.info.get("codebase_version"), repo_id) dataset = lerobot_dataset.LeRobotDataset( data_config.repo_id, delta_timestamps={ diff --git a/src/openpi/training/data_loader_test.py b/src/openpi/training/data_loader_test.py index d15a73529e..415ed2517a 100644 --- a/src/openpi/training/data_loader_test.py +++ b/src/openpi/training/data_loader_test.py @@ -1,12 +1,24 @@ import dataclasses import jax +import pytest from openpi.models import pi0_config from openpi.training import config as _config from openpi.training import data_loader as _data_loader +def test_lerobot_dataset_version_check_accepts_v21(): + _data_loader.check_lerobot_dataset_version("v2.1", "test/repo") + _data_loader.check_lerobot_dataset_version("2.1", "test/repo") + + +@pytest.mark.parametrize("version", ["v3.0", "v3.1"]) +def test_lerobot_dataset_version_check_rejects_v3(version): + with pytest.raises(ValueError, match="v2.1"): + _data_loader.check_lerobot_dataset_version(version, "test/repo") + + def test_torch_data_loader(): config = pi0_config.Pi0Config(action_dim=24, action_horizon=50, max_token_len=48) dataset = _data_loader.FakeDataset(config, 16) diff --git a/uv.lock b/uv.lock index d7c9e410b1..2191289110 100644 --- a/uv.lock +++ b/uv.lock @@ -3087,6 +3087,7 @@ dependencies = [ { name = "augmax" }, { name = "beartype" }, { name = "chex" }, + { name = "datasets" }, { name = "dm-tree", version = "0.1.8", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, { name = "dm-tree", version = "0.1.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, { name = "einops" }, @@ -3140,6 +3141,7 @@ requires-dist = [ { name = "augmax", specifier = ">=0.3.4" }, { name = "beartype", specifier = "==0.19.0" }, { name = "chex", specifier = "==0.1.90" }, + { name = "datasets", specifier = ">=2.19.0,<4.0.0" }, { name = "dm-tree", specifier = ">=0.1.8" }, { name = "einops", specifier = ">=0.8.0" }, { name = "equinox", specifier = ">=0.11.8" },