Constrain datasets and fail clearly on unsupported LeRobot dataset formats

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",

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={

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)