From c2cf1edda324610ed6310f578cd7fe1ccf17f0f4 Mon Sep 17 00:00:00 2001 From: Advait Jayant Date: Fri, 12 Jun 2026 01:28:14 +0100 Subject: [PATCH] Constrain datasets and fail clearly on unsupported LeRobot dataset formats datasets 4.0 returns lazy Column objects for column access, which breaks the pinned lerobot with 'TypeError: stack(): ... not Column' in any environment that does not use the uv lockfile, since neither openpi nor the pinned lerobot bound datasets from above. Datasets stored in the newer LeRobot v3 format fail much deeper with misleading tracebacks because the pinned lerobot can only read format v2.1. This adds an explicit datasets>=2.19.0,<4.0.0 dependency and an early codebase version check that raises an actionable error before constructing the dataset. Fixes #753 --- pyproject.toml | 1 + src/openpi/training/data_loader.py | 30 +++++++++++++++++++++++++ src/openpi/training/data_loader_test.py | 12 ++++++++++ uv.lock | 2 ++ 4 files changed, 45 insertions(+) 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" },