[python] Speed up TorchCodec random video frame reads with persisted indexes

.github/workflows/ci-python.yml

@@ -267,19 +267,33 @@ jobs:
           done
 
   lerobot-compatibility:
-    name: LeRobot ${{ matrix.lerobot }} / Python ${{ matrix.python }}
+    name: LeRobot ${{ matrix.lerobot }} / Python ${{ matrix.python }} / TorchCodec ${{ matrix.torchcodec }}
     timeout-minutes: 40
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     strategy:
       fail-fast: false
       matrix:
         include:
           - python: "3.10"
             lerobot: "0.4.4"
+            torch: "2.7.1"
+            torchvision: "0.22.1"
+            torchcodec: "0.3.0"
+          - python: "3.10"
+            lerobot: "0.4.4"
+            torch: "2.6.0"
+            torchvision: "0.21.0"
+            torchcodec: "0.2.1"
           - python: "3.12"
             lerobot: "0.6.0"
+            torch: "2.7.1"
+            torchvision: "0.22.1"
+            torchcodec: "0.3.0"
           - python: "3.12"
             lerobot: "0.6.1"
+            torch: "2.7.1"
+            torchvision: "0.22.1"
+            torchcodec: "0.3.0"
     steps:
       - uses: actions/checkout@v6
       - uses: actions/setup-python@v6
@@ -290,18 +304,31 @@ jobs:
       - 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 torch==${{ matrix.torch }} torchvision==${{ matrix.torchvision }} torchcodec==${{ matrix.torchcodec }} --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
+        if: matrix.torchcodec != '0.2.1'
         working-directory: paimon-python
         run: |
+          python -c 'from lerobot.datasets.lerobot_dataset import LeRobotDataset'
           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
 
+      - name: Test indexed TorchCodec (no skips)
+        working-directory: paimon-python
+        run: |
+          python -c 'import av; from torchcodec.decoders import VideoDecoder'
+          python -m pytest -q pypaimon/tests/indexed_torchcodec_test.py --junitxml=indexed-torchcodec.xml
+          python - <<'PYTEST_REPORT'
+          import xml.etree.ElementTree as ET
+          cases = list(ET.parse("indexed-torchcodec.xml").iter("testcase"))
+          assert cases, "No indexed TorchCodec tests ran"
+          assert not any(case.find("skipped") is not None for case in cases), "Indexed TorchCodec tests must not skip"
+          PYTEST_REPORT
+
   torch:
     name: PyTorch / Python 3.10
     timeout-minutes: 60

docs/docs/pypaimon/lerobot.md

@@ -249,7 +249,10 @@ batches, avoiding a second full-resolution copy during collation.
 Without `tag_name`, the latest snapshots are used. Frame lookups use the BTree
 on `index`; payloads remain lazy. Indexed videos prefetch metadata and target GOPs, then fetch
 uncached PyAV reads on demand. Unindexed videos use the TorchCodec/PyAV scan path. Set
-`video_backend` to force either decoder.
+`video_backend` to force either decoder. With `video_backend="torchcodec"`, indexed
+MP4 videos use timestamp seeks verified against their sample table and returned PTS.
+Decoder replicas share read-only timelines; uniform timing is stored compactly.
+Unsupported timelines or mismatched timestamps fall back to TorchCodec exact mode.
 
 With `pypaimon-rust` installed, LeRobot defaults to Rust frame lookup and table reads
 unless native read/planning options are explicitly configured. Set
@@ -262,9 +265,9 @@ Each open decoder uses memory and a file handle, so usage grows with cameras and
 `read.video.decode-parallelism` limits concurrent video decodes across all cameras in each
 Dataset worker (default: 16). Set it to 1 for serial decoding. Parallel decoding
 applies to video windows without `image_transforms` and uses more CPU and memory.
-`read.video.codec-threads` controls FFmpeg threads per PyAV decoder (default: 1),
+`read.video.codec-threads` controls FFmpeg threads per PyAV or indexed TorchCodec decoder (default: 1),
 avoiding nested codec parallelism with multiple workers and concurrent videos.
-Set it to 0 for automatic selection. It does not affect TorchCodec.
+Set it to 0 for automatic selection. TorchCodec exact-mode fallback uses its own defaults.
 
 `read.video.decoders-per-payload` (default: 4) also splits single-video windows across
 independent cached decoders, within those limits. Set it to 1 to disable this splitting;

paimon-python/pypaimon/common/options/core_options.py

@@ -1185,7 +1185,7 @@ class CoreOptions:
         .int_type()
         .default_value(1)
         .with_description(
-            "FFmpeg threads per PyAV video decoder in a LeRobot worker. "
+            "FFmpeg threads per PyAV or indexed TorchCodec decoder in a LeRobot worker. "
             "0 keeps automatic selection; 1 avoids nested codec parallelism.")
     )
 

paimon-python/pypaimon/multimodal/lerobot/dataset.py

@@ -1789,6 +1789,9 @@ def _open_video_decoder(stream, backend=None, codec_threads=0):
             stream.seek(0)
         else:
             return _FallbackVideoDecoder(stream, decoder)
+    if backend == "torchcodec" and keyframe_index is not None:
+        from pypaimon.multimodal.lerobot.indexed_torchcodec import _IndexedTorchCodecDecoder
+        return _IndexedTorchCodecDecoder(stream, keyframe_index, codec_threads)
     if backend in (None, "torchcodec"):
         try:
             return _open_torchcodec_decoder(stream)