wip

yaak-ai · Nov 26, 2024 · e0c731f · e0c731f
1 parent a7db0a9
commit e0c731f
Show file tree

Hide file tree

Showing 56 changed files with 964 additions and 1,054 deletions.
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -42,7 +42,7 @@ jobs:
           only: ytt
 
       - name: setup uv
-        uses: astral-sh/setup-uv@v3
+        uses: astral-sh/setup-uv@v4
         with:
           version: "latest"
           enable-cache: true

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -11,14 +11,14 @@ repos:
       - id: pyupgrade
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.7.4
+    rev: v0.8.0
     hooks:
       - id: ruff
         args: [--fix]
       - id: ruff-format
 
   - repo: https://github.com/DetachHead/basedpyright-pre-commit-mirror
-    rev: 1.21.1
+    rev: 1.22.0
     hooks:
       - id: basedpyright
 

diff --git a/config/_templates/dataset/carla.yaml b/config/_templates/dataset/carla.yaml
@@ -1,52 +1,52 @@
 #@yaml/text-templated-strings
 
 #@ drives = [
-#@     'Accident',
-#@     'Accident_2',
-#@     'AccidentTwoWays',
-#@     'BlockedIntersection',
-#@     'ConstructionObstacle',
-#@     'ConstructionObstacleTwoWays',
-#@     'ControlLoss',
-#@     'DynamicObjectCrossing_right_container',
-#@     'DynamicObjectCrossing_right_haybalelb',
-#@     'EnterActorFlow_fast',
-#@     'EnterActorFlow_slow',
-#@     'HardBreakRoute',
-#@     'HazardAtSideLane',
-#@     'HazardAtSideLaneTwoWays',
-#@     'HighwayCutIn',
-#@     'HighwayExit',
-#@     'InterurbanActorFlow_fast',
-#@     'InterurbanActorFlow_slow',
-#@     'InterurbanAdvancedActorFlow_fast',
-#@     'InterurbanAdvancedActorFlow_slow',
-#@     'InterurbanAdvancedActorFlow_slow_2',
-#@     'InvadingTurn_less_curve',
-#@     'InvadingTurn_much_curve',
-#@     'MergerIntoSlowTraffic',
-#@     'MergerIntoSlowTrafficV2',
-#@     'NonSignalizedJunctionLeftTurn_fast',
-#@     'NonSignalizedJunctionLeftTurn_slow',
-#@     'NonSignalizedJunctionRightTurn_fast',
-#@     'NonSignalizedJunctionRightTurn_slow',
-#@     'OppositeVehicleRunningRedLight',
-#@     'OppositeVehicleTakingPriority',
-#@     'ParkedObstacle',
-#@     'ParkedObstacleTwoWays',
-#@     'ParkingCrossingPedestrian',
-#@     'ParkingCutIn',
-#@     'ParkingExit',
-#@     'PedestrianCrossing',
-#@     'SignalizedJunctionLeftTurn_fast',
-#@     'SignalizedJunctionLeftTurn_slow',
-#@     'SignalizedJunctionRightTurn_fast',
-#@     'SignalizedJunctionRightTurn_slow',
-#@     'StaticCutIn',
-#@     'VehicleOpensDoorTwoWays',
-#@     'VehicleTurningRoutePedestrian',
+#!     'Accident',
+#!     'Accident_2',
+#!     'AccidentTwoWays',
+#!     'BlockedIntersection',
+#!     'ConstructionObstacle',
+#!     'ConstructionObstacleTwoWays',
+#!     'ControlLoss',
+#!     'DynamicObjectCrossing_right_container',
+#!     'DynamicObjectCrossing_right_haybalelb',
+#!     'EnterActorFlow_fast',
+#!     'EnterActorFlow_slow',
+#!     'HardBreakRoute',
+#!     'HazardAtSideLane',
+#!     'HazardAtSideLaneTwoWays',
+#!     'HighwayCutIn',
+#!     'HighwayExit',
+#!     'InterurbanActorFlow_fast',
+#!     'InterurbanActorFlow_slow',
+#!     'InterurbanAdvancedActorFlow_fast',
+#!     'InterurbanAdvancedActorFlow_slow',
+#!     'InterurbanAdvancedActorFlow_slow_2',
+#!     'InvadingTurn_less_curve',
+#!     'InvadingTurn_much_curve',
+#!     'MergerIntoSlowTraffic',
+#!     'MergerIntoSlowTrafficV2',
+#!     'NonSignalizedJunctionLeftTurn_fast',
+#!     'NonSignalizedJunctionLeftTurn_slow',
+#!     'NonSignalizedJunctionRightTurn_fast',
+#!     'NonSignalizedJunctionRightTurn_slow',
+#!     'OppositeVehicleRunningRedLight',
+#!     'OppositeVehicleTakingPriority',
+#!     'ParkedObstacle',
+#!     'ParkedObstacleTwoWays',
+#!     'ParkingCrossingPedestrian',
+#!     'ParkingCutIn',
+#!     'ParkingExit',
+#!     'PedestrianCrossing',
+#!     'SignalizedJunctionLeftTurn_fast',
+#!     'SignalizedJunctionLeftTurn_slow',
+#!     'SignalizedJunctionRightTurn_fast',
+#!     'SignalizedJunctionRightTurn_slow',
+#!     'StaticCutIn',
+#!     'VehicleOpensDoorTwoWays',
+#!     'VehicleTurningRoutePedestrian',
 #@     'VehicleTurningRoute_right',
-#@     'YieldToEmergencyVehicle',
+#!     'YieldToEmergencyVehicle',
 #@ ]
 
 #@ cameras = [
@@ -75,39 +75,58 @@ inputs:
 
       #@ end
 
-    table_builder:
-      _target_: rbyte.io.table.TableBuilder
-      _convert_: all
-      readers:
-        ego_logs:
-          path: ${data_dir}/(@=input_id@)/ego_logs.json
-          reader:
-            _target_: rbyte.io.JsonTableReader
-            _recursive_: false
-            fields:
-              records:
-                _idx_:
-                control.brake:
-                control.throttle:
-                control.steer:
-                state.velocity.value:
-                state.acceleration.value:
+    samples:
+      pipeline:
+        _target_: pipefunc.Pipeline
+        _convert_: all
+        validate_type_annotations: false
+        functions:
+          - _target_: pipefunc.PipeFunc
+            bound:
+              path: ${data_dir}/(@=input_id@)/ego_logs.json
+            output_name: ego_logs
+            func:
+              _target_: rbyte.io.JsonDataFrameBuilder
+              fields:
+                records:
+                  control.brake:
+                  control.throttle:
+                  control.steer:
+                  state.velocity.value:
+                  state.acceleration.value:
 
-            transforms:
-              - _target_: rbyte.io.FpsResampler
-                source_fps: 20
-                target_fps: 30
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: ego_logs
+            output_name: data
+            func:
+              _target_: rbyte.io.DataFrameConcater
+              method: vertical
 
-      merger:
-        _target_: rbyte.io.TableConcater
-        method: vertical
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: data
+            output_name: data_resampled
+            func:
+              _target_: rbyte.io.DataFrameFpsResampler
+              fps_in: 20
+              fps_out: 30
 
-      filter: |
-        `control.throttle` > 0.5
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: data_resampled
+            output_name: data_filtered
+            func:
+              _target_: rbyte.io.DataFrameFilter
+              predicate: |
+                `control.throttle` > 0.5
 
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: data_filtered
+            output_name: samples
+            func:
+              _target_: rbyte.RollingWindowSampleBuilder
+              index_column: _idx_
+              period: 1i
   #@ end
-
-sample_builder:
-  _target_: rbyte.RollingWindowSampleBuilder
-  index_column: _idx_
-  period: 1i
diff --git a/config/_templates/dataset/mimicgen.yaml b/config/_templates/dataset/mimicgen.yaml
@@ -12,8 +12,8 @@
 #@ ]
 ---
 _target_: rbyte.Dataset
-_convert_: all
 _recursive_: false
+_convert_: all
 inputs:
   #@ for input_id, input_keys in inputs.items():
   #@ for input_key in input_keys:
@@ -28,27 +28,39 @@ inputs:
           key: (@=input_key@)/(@=frame_key@)
       #@ end
 
-    table_builder:
-      _target_: rbyte.io.TableBuilder
-      _convert_: all
-      readers:
-        hdf5:
-          path: "${data_dir}/(@=input_id@).hdf5"
-          reader:
-            _target_: rbyte.io.Hdf5TableReader
-            _recursive_: false
-            fields:
-              (@=input_key@):
-                _idx_:
-                obs/robot0_eef_pos:
+    samples:
+      pipeline:
+        _target_: pipefunc.Pipeline
+        _convert_: all
+        validate_type_annotations: false
+        functions:
+          - _target_: pipefunc.PipeFunc
+            bound:
+              path: "${data_dir}/(@=input_id@).hdf5"
+            output_name: data
+            func:
+              _target_: rbyte.io.Hdf5DataFrameBuilder
+              fields:
+                (@=input_key@):
+                  _idx_:
+                  obs/robot0_eef_pos:
+
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: data
+            output_name: data_concated
+            func:
+              _target_: rbyte.io.DataFrameConcater
+              method: vertical
+
+          - _target_: pipefunc.PipeFunc
+            renames:
+              input: data_concated
+            output_name: samples
+            func:
+              _target_: rbyte.RollingWindowSampleBuilder
+              index_column: _idx_
+              period: 1i
 
-      merger:
-        _target_: rbyte.io.TableConcater
-        method: vertical
   #@ end
   #@ end
-
-sample_builder:
-  _target_: rbyte.RollingWindowSampleBuilder
-  index_column: _idx_
-  period: 1i