diff --git a/.github/workflows/ci-python.yml b/.github/workflows/ci-python.yml index 775c578943fb..1dd9742be001 100644 --- a/.github/workflows/ci-python.yml +++ b/.github/workflows/ci-python.yml @@ -265,6 +265,42 @@ jobs: -k 'TIMESTAMP or temporal' done + lerobot-compatibility: + name: LeRobot ${{ matrix.lerobot }} / Python ${{ matrix.python }} + timeout-minutes: 40 + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + include: + - python: "3.10" + lerobot: "0.4.4" + - python: "3.12" + lerobot: "0.6.0" + - python: "3.12" + lerobot: "0.6.1" + steps: + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.python }} + - name: Install FFmpeg + run: sudo apt-get update && sudo apt-get install -y ffmpeg + - name: Install LeRobot dependencies + run: | + python -m pip install --upgrade pip + python -m pip install torch==2.7.1 torchvision==0.22.1 torchcodec==0.3.0 --index-url https://download.pytorch.org/whl/cpu + python -m pip install -e './paimon-python[lerobot]' 'lerobot==${{ matrix.lerobot }}' --group ./paimon-python/pyproject.toml:test + python -m pip check + python -c 'from lerobot.datasets.lerobot_dataset import LeRobotDataset' + - name: Test import, training reads and writer statistics + working-directory: paimon-python + run: | + python -m pytest -q \ + pypaimon/tests/multimodal_lerobot_test.py \ + pypaimon/tests/multimodal_lerobot_writer_test.py::PaimonLeRobotWriterTest::test_multiframe_boolean_statistics_preserve_extrema \ + pypaimon/tests/multimodal_lerobot_writer_test.py::PaimonLeRobotWriterTest::test_multiframe_int64_statistics_preserve_extrema + torch: name: PyTorch / Python 3.10 timeout-minutes: 60 diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 9fba86ce96d9..a3c0cfcc78c9 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -46,6 +46,8 @@ Missing or empty stats create no stats table. The frame table's pip install 'pypaimon[lerobot]' ``` +Python 3.10–3.11 uses LeRobot 0.4.x; Python 3.12+ uses LeRobot 0.6.x. + ```python conn.load_from_lerobot( "robot_data", diff --git a/paimon-python/pypaimon/multimodal/lerobot/api.py b/paimon-python/pypaimon/multimodal/lerobot/api.py index ce9f776eed14..7a38a3e720fa 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/api.py +++ b/paimon-python/pypaimon/multimodal/lerobot/api.py @@ -43,6 +43,7 @@ from pypaimon.multimodal.lerobot.source import ( _close_quietly, _import_lerobot_dataset, + _info_dict, _load_hub_info, _open_resolved_dataset, _resolved_source, @@ -122,15 +123,16 @@ def load_from_lerobot( download_videos=bool(video_fields), ) try: - info = dict(dataset.meta.info) + info = _info_dict(dataset.meta.info) _require_v3(info, resolved_source.path) _validated_counts(info, resolved_source.path) _validate_v3_required_features(info) video_fields = _video_feature_names(info) lerobot_schema = _schema_from_info(info) + # Preserve source extensions omitted by LeRobot DatasetInfo.to_dict(). metadata = _load_dataset_metadata( - dataset, info, resolved_source) + dataset, local_info, resolved_source) _import_dataset( connection, table_name, diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 60386c82f473..929b83d54da4 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -49,6 +49,7 @@ _schema_from_info, _validate_lerobot_schema, ) +from pypaimon.multimodal.lerobot.source import _info_dict from pypaimon.multimodal.table import _target_schema, _time_travel_table from pypaimon.multimodal.video import VideoFrameCollator, _VideoDecodePool from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex @@ -885,10 +886,10 @@ def get_task_index(self, task): def _reader_metadata(metadata): info = _metadata_member(metadata, "info") - if not isinstance(info, Mapping): + if not isinstance(info, Mapping) and not callable(getattr(info, "to_dict", None)): raise TypeError( "PaimonDatasetReader metadata must contain an info map.") - info = dict(info) + info = _info_dict(info) features = info.get("features") if isinstance(features, Mapping): info["features"] = { diff --git a/paimon-python/pypaimon/multimodal/lerobot/source.py b/paimon-python/pypaimon/multimodal/lerobot/source.py index 7eb8265c7399..07e7dfb487d4 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/source.py +++ b/paimon-python/pypaimon/multimodal/lerobot/source.py @@ -214,6 +214,11 @@ def _import_lerobot_dataset(): return LeRobotDataset +def _info_dict(info): + # LeRobot 0.6 exposes DatasetInfo instead of a dict. + return dict(info.to_dict() if hasattr(info, "to_dict") else info) + + def _load_hub_info(source): try: from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata @@ -222,7 +227,9 @@ def _load_hub_info(source): "load_from_lerobot requires LeRobot; install " "'pypaimon[lerobot]'.") from error try: - return dict(LeRobotDatasetMetadata(repo_id=source.repo_id).info) + metadata = LeRobotDatasetMetadata(repo_id=source.repo_id) + # DatasetInfo omits unknown fields; preserve the downloaded source JSON. + return _local_source(Path(metadata.root), source.path)[1] except Exception as error: raise ValueError( "Cannot open LeRobot Dataset v3 metadata %s: %s" diff --git a/paimon-python/pypaimon/multimodal/lerobot/writer.py b/paimon-python/pypaimon/multimodal/lerobot/writer.py index adc5dd54d36b..311e20255695 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/writer.py +++ b/paimon-python/pypaimon/multimodal/lerobot/writer.py @@ -208,6 +208,15 @@ def _compute_stats(episode, features): array = array.reshape(array.shape[0], -1) data[name] = array numeric_stats = compute_episode_stats(data, numeric_features) + # Preserve exact extrema when LeRobot promotes integer and bool samples. + for name, feature in numeric_features.items(): + dtype = feature.get("dtype") + if dtype in _INTEGER_DTYPES or dtype in ("bool", "boolean"): + values = data[name] + numeric_stats[name]["min"] = np.min( + values, axis=0, keepdims=values.ndim == 1) + numeric_stats[name]["max"] = np.max( + values, axis=0, keepdims=values.ndim == 1) for name, shape in reshaped_features.items(): numeric_stats[name] = { stat: value if stat == "count" else value.reshape(shape) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index 7d48264a3a88..a9b28f8e98bd 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -3500,7 +3500,7 @@ def close(self): @unittest.skipUnless( sys.version_info >= (3, 10) and LeRobotDataset is not None, - "LeRobot 0.4.x requires Python 3.10+ and the lerobot extra", + "LeRobot requires Python 3.10+ and the lerobot extra", ) class LeRobotImportTest(unittest.TestCase): @@ -3583,6 +3583,70 @@ def _create_image_dataset(root): dataset.save_episode() dataset.finalize() + def test_hub_metadata_accepts_installed_lerobot_info(self): + from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata + from pypaimon.multimodal.lerobot.source import _load_hub_info + + metadata = LeRobotDatasetMetadata( + repo_id="pypaimon/local-image-test", root=self.image_source) + with patch("lerobot.datasets.lerobot_dataset.LeRobotDatasetMetadata", + return_value=metadata): + info = _load_hub_info(_LeRobotSource( + path="pypaimon/local-image-test", root=None, + repo_id="pypaimon/local-image-test")) + self.assertIsInstance(info, dict) + self.assertEqual(info["total_frames"], 5) + self.assertEqual(info["total_episodes"], 2) + self.assertEqual(info["features"]["observation.state"]["dtype"], "float32") + + def test_custom_reader_accepts_installed_lerobot_metadata(self): + from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata + + metadata = LeRobotDatasetMetadata( + repo_id="pypaimon/local-image-test", root=self.image_source) + reader = _ManualDatasetReader(metadata) + try: + self.assertEqual(reader.meta.info["total_frames"], 5) + self.assertEqual(reader.meta.info["total_episodes"], 2) + self.assertEqual( + reader.meta.features["observation.state"]["shape"], (3,)) + finally: + reader.close() + + def test_local_import_preserves_extended_info(self): + source = self.temp_dir / "extended-source" + shutil.copytree(self.image_source, source) + info_path = source / "meta" / "info.json" + info = json.loads(info_path.read_text()) + info["vendor_extension"] = {"labels": ["pick"], "enabled": True} + info_path.write_text(json.dumps(info)) + + self.connection.load_from_lerobot("extended_info", source) + + self.assertEqual( + info, _catalog_metadata(self.connection, "extended_info__info")) + + def test_hub_import_preserves_extended_info(self): + from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata + + source = self.temp_dir / "hub-source" + shutil.copytree(self.image_source, source) + info_path = source / "meta" / "info.json" + info = json.loads(info_path.read_text()) + info["vendor_extension"] = {"labels": ["pick"], "enabled": True} + info_path.write_text(json.dumps(info)) + repo_id = "pypaimon/hub-image-test" + metadata = LeRobotDatasetMetadata(repo_id=repo_id, root=source) + dataset = LeRobotDataset(repo_id=repo_id, root=source) + with patch("lerobot.datasets.lerobot_dataset.LeRobotDatasetMetadata", + return_value=metadata), patch( + "pypaimon.multimodal.lerobot.source._open_dataset", + return_value=dataset): + self.connection.load_from_lerobot("hub_info", repo_id) + + self.assertEqual( + info, _catalog_metadata(self.connection, "hub_info__info")) + def test_table_dataset_pickle_preserves_episode_metadata_and_reads(self): import torch diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py index b796e0406134..acdbb2f1f61e 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_writer_test.py @@ -817,6 +817,89 @@ def test_writes_multidimensional_and_boolean_numpy_features(self): rows[0]["observation.matrix"]) self.assertEqual([True, False], rows[0]["observation.flags"]) + def test_multiframe_boolean_statistics_preserve_extrema(self): + for shape in ((2,), (2, 2)): + with self.subTest(shape=shape): + name = "bool_stats_%d" % len(shape) + writer = PaimonLeRobotWriter( + self.connection, name, fps=10, features={ + "observation.flags": { + "dtype": "bool", "shape": shape, "names": None, + }, + }) + first = np.array([True, False] * (int(np.prod(shape)) // 2), + dtype=np.bool_).reshape(shape) + for flags in (first, ~first): + writer.add_frame({ + "observation.flags": flags, "task": "inspect", + }) + writer.save_episode() + writer.finalize() + + rows = self.connection.get_table(name).scan().select([ + "index", "observation.flags" + ]).to_arrow().sort_by("index").to_pylist() + self.assertEqual( + [first.tolist(), (~first).tolist()], + [row["observation.flags"] for row in rows]) + episode = _catalog_rows( + self.connection, name + "__episodes")[0] + stats = { + row["key"]: json.loads(row["value"]) + for row in _catalog_rows(self.connection, name + "__stats") + }["observation.flags"] + for stat, expected in (("min", np.zeros(shape, dtype=bool)), + ("max", np.ones(shape, dtype=bool))): + value = episode["stats/observation.flags/" + stat] + self.assertEqual(expected.tolist(), value) + self.assertEqual(np.dtype(bool), np.asarray(value).dtype) + self.assertEqual(expected.tolist(), stats[stat]) + self.assertEqual( + np.dtype(bool), np.asarray(stats[stat]).dtype) + + def test_multiframe_int64_statistics_preserve_extrema(self): + cases = ( + ("precision", (2 ** 53 + 1, 2 ** 53 + 3)), + ("upper_bound", (2 ** 63 - 2, 2 ** 63 - 1)), + ) + for name, values in cases: + with self.subTest(name=name): + table_name = "int64_stats_" + name + writer = PaimonLeRobotWriter( + self.connection, table_name, fps=10, features={ + "observation.counter": { + "dtype": "int64", "shape": (1,), "names": None, + }, + }) + for value in values: + writer.add_frame({ + "observation.counter": np.array([value], dtype=np.int64), + "task": "inspect", + }) + writer.save_episode() + writer.finalize() + + rows = self.connection.get_table(table_name).scan().select([ + "index", "observation.counter" + ]).to_arrow().sort_by("index").to_pylist() + self.assertEqual( + list(values), + [row["observation.counter"] for row in rows]) + episode = _catalog_rows( + self.connection, table_name + "__episodes")[0] + stats = { + row["key"]: json.loads(row["value"]) + for row in _catalog_rows(self.connection, table_name + "__stats") + }["observation.counter"] + for stat, expected in (("min", min(values)), + ("max", max(values))): + episode_value = episode["stats/observation.counter/" + stat][0] + stats_value = stats[stat][0] + self.assertEqual(expected, episode_value) + self.assertEqual(expected, stats_value) + self.assertIsInstance(episode_value, int) + self.assertIsInstance(stats_value, int) + def test_discards_rerecorded_episode_and_finalizes_tail_batch(self): writer = PaimonLeRobotWriter( self.connection, diff --git a/paimon-python/pyproject.toml b/paimon-python/pyproject.toml index 624131379361..9a9232c82f36 100644 --- a/paimon-python/pyproject.toml +++ b/paimon-python/pyproject.toml @@ -102,7 +102,8 @@ lerobot = [ "datasets>=4,<5; python_version >= '3.10'", "pandas>=2.2.2,<3; python_version >= '3.10'", "torch>=2.3; python_version >= '3.10'", - "lerobot>=0.4.4,<0.5; python_version >= '3.10'", + "lerobot>=0.4.4,<0.5; python_version >= '3.10' and python_version < '3.12'", + "lerobot[dataset]>=0.6,<0.7; python_version >= '3.12'", ] ray = [ "ray>=2.10,<3; python_version >= '3.8'",