This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new d95fd18b48 [python] Support LeRobot 0.6 dataset imports and training
reads (#10734)
d95fd18b48 is described below
commit d95fd18b48bd26ed394a113b255f369675c7e2d3
Author: XiaoHongbo <[email protected]>
AuthorDate: Fri Oct 23 15:45:25 2026 +0800
[python] Support LeRobot 0.6 dataset imports and training reads (#10734)
---
.github/workflows/ci-python.yml | 36 ++++++++++
docs/docs/pypaimon/lerobot.md | 2 +
paimon-python/pypaimon/multimodal/lerobot/api.py | 6 +-
.../pypaimon/multimodal/lerobot/dataset.py | 5 +-
.../pypaimon/multimodal/lerobot/source.py | 9 ++-
.../pypaimon/multimodal/lerobot/writer.py | 9 +++
.../pypaimon/tests/multimodal_lerobot_test.py | 66 ++++++++++++++++-
.../tests/multimodal_lerobot_writer_test.py | 83 ++++++++++++++++++++++
paimon-python/pyproject.toml | 3 +-
9 files changed, 212 insertions(+), 7 deletions(-)
diff --git a/.github/workflows/ci-python.yml b/.github/workflows/ci-python.yml
index 775c578943..1dd9742be0 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 9fba86ce96..a3c0cfcc78 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 ce9f776eed..7a38a3e720 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.schema import (
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 60386c82f4..929b83d54d 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
@@ -49,6 +49,7 @@ from pypaimon.multimodal.lerobot.schema import (
_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 @@ class _PaimonLeRobotMetadata:
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 7eb8265c73..07e7dfb487 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 adc5dd54d3..311e202556 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 1aea33b2db..b718ca9d7f 100644
--- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
@@ -3500,7 +3500,7 @@ class _FailingCloseDataset:
@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 @@ class LeRobotImportTest(unittest.TestCase):
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 b796e04061..acdbb2f1f6 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 @@ class PaimonLeRobotWriterTest(unittest.TestCase):
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 6241313793..9a9232c82f 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'",
