chore: dataloader tests (nuscenes, mimicgen, yaak) (#16)

.gitattributes

@@ -0,0 +1 @@
+tests/data/** filter=lfs diff=lfs merge=lfs -text

.github/workflows/ci.yaml

@@ -10,9 +10,8 @@ on:
   workflow_dispatch:
 
 jobs:
-  pre-commit:
+  ci:
     runs-on: ubuntu-latest
-
     steps:
       - name: setup ssh
         uses: webfactory/ssh-agent@v0.9.0
@@ -24,15 +23,19 @@ jobs:
         with:
           submodules: recursive
           persist-credentials: false
+          lfs: true
+
+      - name: git lfs checkout
+        run: git lfs checkout
+
+      - name: setup just
+        uses: extractions/setup-just@v2
 
       - name: setup python
         uses: actions/setup-python@v5
         with:
           python-version-file: "pyproject.toml"
 
-      - name: setup just
-        uses: extractions/setup-just@v2
-
       - name: setup ytt
         uses: carvel-dev/setup-action@v2
         with:
@@ -48,14 +51,17 @@ jobs:
       - name: sync
         run: just sync
 
-      - name: install tools
-        run: just install-tools
+      - name: build protos
+        run: just build-protos
 
-      - name: restore pre-commit cache
-        uses: actions/cache@v4
-        with:
-          path: ~/.cache/pre-commit
-          key: pre-commit-${{ hashFiles('.pre-commit-config.yaml') }}
+      - name: format
+        run: just format --check
+
+      - name: lint
+        run: just lint
+
+      - name: typecheck
+        run: just typecheck
 
-      - name: pre-commit
-        run: just pre-commit
+      - name: test
+        run: just test

.gitignore

@@ -10,7 +10,9 @@ wheels/
 .venv
 
 # generated
-examples/config
+config/*
+!config/_templates
+
 **/*_pb2.py
 **/*_pb2.pyi
 **/outputs/*

.pre-commit-config.yaml

@@ -1,6 +1,4 @@
 ---
-fail_fast: true
-
 repos:
   - repo: https://github.com/abravalheri/validate-pyproject
     rev: v0.21
@@ -17,21 +15,15 @@ repos:
     hooks:
       - id: pyupgrade
 
-  - repo: https://github.com/google/yamlfmt
-    rev: v0.13.0
-    hooks:
-      - id: yamlfmt
-        exclude: examples/config
-
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.6.9
+    rev: v0.7.0
     hooks:
       - id: ruff
         args: [--fix]
       - id: ruff-format
 
   - repo: https://github.com/DetachHead/basedpyright-pre-commit-mirror
-    rev: 1.18.4
+    rev: 1.19.0
     hooks:
       - id: basedpyright
 
@@ -44,11 +36,3 @@ repos:
         entry: just --fmt --unstable
         pass_filenames: false
         always_run: true
-
-      - id: generate-example-config
-        name: generate-example-config
-        language: system
-        stages: [pre-commit]
-        entry: just generate-example-config
-        pass_filenames: false
-        always_run: true

README.md

@@ -20,7 +20,7 @@ uv add https://github.com/yaak-ai/rbyte/releases/latest/download/rbyte-X.Y.Z-py3
 
 ## Examples
 
-See [examples/config_templates](examples/config_templates) ([`ytt`](https://carvel.dev/ytt/) templates) and [justfile](justfile) for usage examples.
+See [config/_templates](./config/_templates) ([`ytt`](https://carvel.dev/ytt/) templates) and [justfile](./justfile) for usage examples.
 
 <details>
 <summary><a href=https://nuscenes.org> nuScenes </a> x <a href=https://mcap.dev> mcap </a></summary>

config/_templates/dataloader/torch.yaml

@@ -0,0 +1,13 @@
+---
+_target_: torch.utils.data.DataLoader
+dataset: ${dataset}
+shuffle: false
+batch_size: 1
+collate_fn:
+  _target_: rbyte.utils.dataloader.collate_identity
+  _partial_: true
+
+num_workers: 1
+pin_memory: false
+persistent_workers: true
+multiprocessing_context: forkserver

examples/config_templates/dataset/carla.yaml → config/_templates/dataset/carla.yaml

@@ -105,7 +105,7 @@ inputs:
         filter: |
           `control.throttle` > 0.5
 
-    #@ end
+  #@ end
 
 sample_builder:
   _target_: rbyte.sample.builder.GreedySampleTableBuilder

examples/config_templates/dataset/hdf5.yaml → config/_templates/dataset/mimicgen.yaml

@@ -1,21 +1,16 @@
-#! https://sites.google.com/view/il-for-mm/datasets#h.cq0r3rd5nr9m
+#! https://huggingface.co/datasets/amandlek/mimicgen_datasets/blob/main/source/coffee.hdf5
 
-#@yaml/map-key-override
 #@yaml/text-templated-strings
 
 #@ inputs = {
-#@     "table_setup_from_dishwasher_sample": [
+#@     "coffee": [
 #@         "/data/demo_0",
 #@         "/data/demo_1",
-#@         "/data/demo_10",
-#@         "/data/demo_101",
-#@         "/data/demo_102",
 #@     ]
 #@ }
 
 #@ frame_keys = [
-#@     '/obs/rgb',
-#@     '/obs/depth',
+#@     'obs/agentview_image',
 #@ ]
 ---
 _target_: rbyte.Dataset
@@ -30,9 +25,9 @@ inputs:
       (@=frame_key@):
         index_column: _idx_
         reader:
-          _target_: rbyte.io.frame.hdf5.Hdf5FrameReader
+          _target_: rbyte.io.frame.Hdf5FrameReader
           path: "${data_dir}/(@=input_id@).hdf5"
-          key: (@=input_key@)(@=frame_key@)
+          key: (@=input_key@)/(@=frame_key@)
       #@ end
 
     table:
@@ -42,13 +37,12 @@ inputs:
         readers:
           - path: "${data_dir}/(@=input_id@).hdf5"
             reader:
-              _target_: rbyte.io.table.hdf5.Hdf5TableReader
+              _target_: rbyte.io.table.Hdf5TableReader
               _recursive_: false
               fields:
                 (@=input_key@):
                   _idx_:
-                  obs/object:
-                  task_successes:
+                  obs/robot0_eef_pos:
 
         merger:
           _target_: rbyte.io.table.TableConcater

examples/config_templates/dataset/mcap.yaml → config/_templates/dataset/nuscenes.yaml

@@ -3,7 +3,7 @@
 #@yaml/text-templated-strings
 
 #@ inputs = [
-#@     'NuScenes-v1.0-mini-scene-0103',
+#@     'nuScenes-v1.0-mini-scene-0061-cut',
 #@ ]
 
 #@ camera_topics = [
@@ -23,7 +23,7 @@ inputs:
       (@=topic@):
         index_column: (@=topic@)/_idx_
         reader:
-          _target_: rbyte.io.frame.mcap.McapFrameReader
+          _target_: rbyte.io.frame.McapFrameReader
           path: "${data_dir}/(@=input_id@).mcap"
           topic: (@=topic@)
           decoder_factory: mcap_protobuf.decoder.DecoderFactory
@@ -42,7 +42,7 @@ inputs:
         readers:
           - path: "${data_dir}/(@=input_id@).mcap"
             reader:
-              _target_: rbyte.io.table.mcap.McapTableReader
+              _target_: rbyte.io.table.McapTableReader
               _recursive_: false
               decoder_factories:
                 - rbyte.utils.mcap.ProtobufDecoderFactory
@@ -78,7 +78,7 @@ inputs:
                 method: ref
               _idx_:
                 method: asof
-                tolerance: 10ms
+                tolerance: 40ms
                 strategy: nearest
             #@ end
 
@@ -89,7 +89,7 @@ inputs:
                 method: interp
 
         filter: |
-          `/odom/vel.x` >= 8.6
+          `/odom/vel.x` >= 8
 
         cache:
           _target_: rbyte.utils.dataframe.DataframeDiskCache

examples/config_templates/dataset/yaak.yaml → config/_templates/dataset/yaak.yaml

@@ -1,7 +1,7 @@
 #@yaml/text-templated-strings
 
 #@ drives = [
-#@     'Niro098-HQ/2024-01-22--09-03-16',
+#@     'Niro098-HQ/2024-06-18--13-39-54',
 #@ ]
 
 #@ cameras = [
@@ -21,16 +21,11 @@ inputs:
       (@=source_id@):
         index_column: "ImageMetadata.(@=source_id@).frame_idx"
         reader:
-          _target_: rbyte.io.frame.DirectoryFrameReader
+          _target_: rbyte.io.frame.FfmpegFrameReader
           _recursive_: true
-          path: "${data_dir}/(@=input_id@)/frames/(@=source_id@).pii.mp4/576x324/{:09d}.jpg"
-          frame_decoder:
-            _target_: simplejpeg.decode_jpeg
-            _partial_: true
-            colorspace: rgb
-            fastdct: true
-            fastupsample: true
-        #@ end
+          path: "${data_dir}/(@=input_id@)/(@=source_id@).pii.mp4"
+          resize_shorter_side: 324
+      #@ end
 
     table:
       builder:
@@ -72,7 +67,7 @@ inputs:
 
           - path: ${data_dir}/(@=input_id@)/ai.mcap
             reader:
-              _target_: rbyte.io.table.mcap.McapTableReader
+              _target_: rbyte.io.table.McapTableReader
               _recursive_: false
               decoder_factories: [rbyte.utils.mcap.ProtobufDecoderFactory]
               fields:
@@ -110,14 +105,15 @@ inputs:
               gear:
                 method: asof
                 tolerance: 100ms
+                strategy: nearest
 
             /ai/safety_score:
               clip.end_timestamp:
                 method: ref
 
               score:
                 method: asof
-                tolerance: 100ms
+                tolerance: 500ms
                 strategy: nearest
 
         filter: |
@@ -132,8 +128,8 @@ inputs:
 sample_builder:
   _target_: rbyte.sample.builder.GreedySampleTableBuilder
   index_column: ImageMetadata.(@=cameras[0]@).frame_idx
-  length: 6
-  stride: 10
-  min_step: 6
+  length: 1
+  stride: 1
+  min_step: 1
   filter: |
-    array_lower(`VehicleMotion.speed`) > 50
+    array_mean(`VehicleMotion.speed`) > 40

examples/config_templates/frame_reader/mcap.yaml → config/_templates/frame_reader/mcap.yaml

@@ -1,5 +1,5 @@
 ---
-_target_: rbyte.io.frame.mcap.McapFrameReader
+_target_: rbyte.io.frame.McapFrameReader
 _recursive_: true
 path: ???
 topic: ???

config/_templates/frame_reader/video/vali.yaml

@@ -0,0 +1,5 @@
+---
+_target_: rbyte.io.frame.video.vali_reader.ValiGpuFrameReader
+_convert_: all
+path: ???
+pixel_format_chain: [NV12]

examples/config_templates/logger/rerun/carla.yaml → config/_templates/logger/rerun/carla.yaml

@@ -3,7 +3,6 @@
 #@ cameras = [
 #@     'cam_front_left',
 #@ ]
-
 ---
 _target_: rbyte.viz.loggers.RerunLogger
 _recursive_: true

config/_templates/logger/rerun/mimicgen.yaml

@@ -0,0 +1,11 @@
+---
+_target_: rbyte.viz.loggers.RerunLogger
+schema:
+  frame:
+    obs/agentview_image:
+      Image:
+        color_model: RGB
+
+  table:
+    _idx_: TimeSequenceColumn
+    obs/robot0_eef_pos: Points3D

examples/config_templates/logger/rerun/yaak.yaml → config/_templates/logger/rerun/yaak.yaml

@@ -5,7 +5,6 @@
 #@     'cam_left_backward',
 #@     'cam_right_backward',
 #@ ]
-
 ---
 _target_: rbyte.viz.loggers.RerunLogger
 schema:

examples/config_templates/table_builder/carla.yaml → config/_templates/table_builder/carla.yaml

@@ -24,5 +24,5 @@ merger:
   _target_: rbyte.io.table.TableConcater
   method: vertical
 
-filter: |
+filter: |-
   `control.throttle` > 0.5

examples/config_templates/table_builder/hdf5.yaml → config/_templates/table_builder/hdf5.yaml

@@ -4,21 +4,16 @@ _convert_: all
 readers:
   - path: ???
     reader:
-      _target_: rbyte.io.table.hdf5.Hdf5TableReader
+      _target_: rbyte.io.table.Hdf5TableReader
       _recursive_: false
       fields:
         /data/demo_0:
           _idx_:
           actions:
           dones:
-          obs/gt_nav:
-          obs/object:
-          obs/proprio:
-          obs/proprio_nav:
-          obs/scan:
+          obs/robot0_eef_pos:
           rewards:
           states:
-          task_successes:
 
 merger:
   _target_: rbyte.io.table.TableConcater