[python] Support LeRobot 0.6 dataset imports and training reads

.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

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

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,

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

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"

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)

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
 

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,

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