modeld_v2: spatial features (PR-1934)

openpilot/sunnypilot/modeld_v2/compile_modeld.py

@@ -7,6 +7,7 @@ See the LICENSE.md file in the root directory for more details.
 """
 
 import argparse
+import math
 import os
 import tempfile
 import time
@@ -66,14 +67,15 @@ def get_policy_npy_shapes(input_shapes: dict, is_supercombo: bool = False) -> tu
   if desire_key:
     shapes['desire'] = (input_shapes[desire_key][2],)
 
-  if is_supercombo and 'features_buffer' in input_shapes:
-    fb = input_shapes['features_buffer']
-    shapes['prev_feat'] = (fb[0], fb[2])
-
   for key, shape in input_shapes.items():
     if key not in (desire_key, 'features_buffer') and 'img' not in key:
       shapes[key] = tuple(shape)
 
+  if is_supercombo and 'features_buffer' in input_shapes:
+    fb = input_shapes['features_buffer']
+    feat_dim = math.prod(fb[2:])
+    shapes['prev_feat'] = (fb[0], feat_dim)
+
   sizes = [int(np.prod(size)) for size in shapes.values()]
   return shapes, sizes
 
@@ -117,8 +119,9 @@ def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = D
   }
 
   if features_buffer:
+    feat_dim = math.prod(features_buffer[2:])
     feat_q_len = frame_skip * features_buffer[1] if is_supercombo else frame_skip * (features_buffer[1] - 1) + 1
-    queues['feat_q'] = Tensor(np.zeros((feat_q_len, features_buffer[0], features_buffer[2]),
+    queues['feat_q'] = Tensor(np.zeros((feat_q_len, features_buffer[0], feat_dim),
                        dtype=np.float32), device=device).contiguous().realize()
 
   queues.update({key: Tensor(value, device='NPY').realize() for key, value in npy_arrays.items() if key in ('tfm', 'big_tfm')})
@@ -183,14 +186,14 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice,
     warped_dev = warped.to(Device.DEFAULT)
     Tensor.realize(packed_npy_inputs_dev, warped_dev)
 
-    img = shift_and_sample(img_q, warped_dev[0:1], sample_skip_fn).realize()
+    img = shift_and_sample(img_q, warped_dev[0:1], sample_skip_fn)
-    big_img = shift_and_sample(big_img_q, warped_dev[1:2], sample_skip_fn).realize()
+    big_img = shift_and_sample(big_img_q, warped_dev[1:2], sample_skip_fn)
 
     unpacked_tensors = [tensor.reshape(shape) for tensor, shape in zip(packed_npy_inputs_dev.split(npy_sizes), npy_shapes.values(), strict=True)]
     unpacked_dict = dict(zip(npy_shapes.keys(), unpacked_tensors, strict=True))
 
     desire_dev = unpacked_dict['desire']
-    desire_buf = shift_and_sample(desire_q, desire_dev.reshape(1, 1, -1), sample_desire_fn).realize()
+    desire_buf = shift_and_sample(desire_q, desire_dev.reshape(1, 1, -1), sample_desire_fn)
 
     inputs = {desire_key: desire_buf}
     for key, tensor_val in unpacked_dict.items():
@@ -199,19 +202,22 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice,
 
     if 'prev_feat' in unpacked_dict:
       prev_feat_dev = unpacked_dict['prev_feat']
-      inputs['features_buffer'] = shift_and_sample(feat_q, prev_feat_dev.reshape(1, 1, -1), sample_skip_fn).realize()
+      feat_buf = shift_and_sample(feat_q, prev_feat_dev.reshape(1, 1, -1), sample_skip_fn)
+      inputs['features_buffer'] = feat_buf if len(fb := input_shapes['features_buffer']) <= 3 else feat_buf.reshape(fb)
 
     if vision_runner:
       vision_out_cast = next(iter(vision_runner({road_key: img, wide_key: big_img}).values())).cast('float32').realize()
       if 'features_buffer' not in inputs:
         new_feat = vision_out_cast[:, features_slice].reshape(1, -1).unsqueeze(0)
-        inputs['features_buffer'] = shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
+        feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
+        inputs['features_buffer'] = feat_buf if len(fb := input_shapes['features_buffer']) <= 3 else feat_buf.reshape(fb)
       policy_outs = [next(iter(pol_runner(inputs).values())).cast('float32').realize() for pol_runner in policy_runners]
       return (vision_out_cast, *policy_outs) if len(policy_outs) > 1 else (vision_out_cast, policy_outs[0])
 
     inputs.update({road_key: img, wide_key: big_img})
     if 'features_buffer' not in inputs:
-      inputs['features_buffer'] = sample_skip_fn(feat_q)
+      feat_buf = sample_skip_fn(feat_q)
+      inputs['features_buffer'] = feat_buf if len(fb := input_shapes['features_buffer']) <= 3 else feat_buf.reshape(fb)
 
     policy_out = next(iter(policy_runners[0](inputs).values())).cast('float32').realize()
     if 'features_buffer' not in inputs and features_slice is not None:

openpilot/sunnypilot/modeld_v2/tests/test_compile_modeld.py

@@ -195,3 +195,85 @@ class TestReadFileChunkedToDisk(OpenpilotTestCase):
 
       assert out.parent == Path(d)
       assert out.read_bytes() == payload
+
+
+class Test4DFeaturesBuffer(OpenpilotTestCase):
+  def test_get_policy_npy_shapes_4d(self):
+    from openpilot.sunnypilot.modeld_v2.compile_modeld import get_policy_npy_shapes
+    input_shapes = {
+      'desire_pulse': (1, 25, 8),
+      'features_buffer': (1, 24, 32, 512),  # compare 4d to 3d for regression
+      'traffic_convention': (1, 2),
+      'action_t': (1, 2)
+    }
+    shapes, sizes = get_policy_npy_shapes(input_shapes, is_supercombo=True)
+    assert shapes['prev_feat'] == (1, 16384)
+    assert sizes == [8, 2, 2, 16384]
+
+  def test_get_policy_npy_shapes_3d(self):
+    from openpilot.sunnypilot.modeld_v2.compile_modeld import get_policy_npy_shapes
+    input_shapes = {
+      'desire_pulse': (1, 25, 8),
+      'features_buffer': (1, 24, 512),
+      'traffic_convention': (1, 2),
+      'action_t': (1, 2)
+    }
+    shapes, sizes = get_policy_npy_shapes(input_shapes, is_supercombo=True)
+    assert shapes['prev_feat'] == (1, 512)
+    assert sizes == [8, 2, 2, 512]
+
+
+class TestStockCompileModeldEquivalence(OpenpilotTestCase):
+  def test_get_policy_npy_shapes_matches_stock(self):
+    from openpilot.selfdrive.modeld.compile_modeld import get_policy_npy_shapes as stock_get_policy_npy_shapes
+    from openpilot.sunnypilot.modeld_v2.compile_modeld import get_policy_npy_shapes as sunny_get_policy_npy_shapes
+
+    stock_input_shapes = {
+      'desire_pulse': (1, 25, 8),
+      'features_buffer': (1, 24, 512), # see below comment
+      'traffic_convention': (1, 2),
+      'action_t': (1, 2),
+    }
+
+    stock_shapes, stock_sizes = stock_get_policy_npy_shapes(stock_input_shapes)
+    sunny_shapes, sunny_sizes = sunny_get_policy_npy_shapes(stock_input_shapes, is_supercombo=True)
+
+    assert sunny_shapes == stock_shapes
+    assert sunny_sizes == stock_sizes
+    assert sunny_shapes['prev_feat'] == (1, 512)
+
+  def test_make_input_queues_full_stock_equivalence(self):
+    from openpilot.selfdrive.modeld.compile_modeld import make_input_queues as stock_make_input_queues
+    from openpilot.sunnypilot.modeld_v2.compile_modeld import make_supercombo_input_queues as sunny_make_supercombo_input_queues
+    input_shapes = {
+      'img': (1, 12, 128, 256),
+      'desire_pulse': (1, 25, 8),
+      'features_buffer': (1, 24, 512),  # when https://github.com/commaai/openpilot/pull/38681 merges, update to 1,24,32,512
+      'traffic_convention': (1, 2),
+      'action_t': (1, 2),
+    }
+    frame_skip = 4
+
+    stock_queues, stock_npy = stock_make_input_queues(input_shapes, frame_skip, device='NPY')
+    sunny_queues, sunny_npy = sunny_make_supercombo_input_queues(input_shapes, frame_skip, device='NPY')
+    assert set(sunny_queues.keys()) == set(stock_queues.keys())
+    for key in stock_queues:
+      assert sunny_queues[key].shape == stock_queues[key].shape, \
+        f"Queue shape mismatch for {key}: sunny {sunny_queues[key].shape} != stock {stock_queues[key].shape}"
+    assert set(sunny_npy.keys()) == set(stock_npy.keys())
+    for key in stock_npy:
+      assert sunny_npy[key].shape == stock_npy[key].shape, \
+        f"Numpy array shape mismatch for {key}: sunny {sunny_npy[key].shape} != stock {stock_npy[key].shape}"
+
+  def test_make_warp_queues_stock_equivalence(self):
+    from openpilot.selfdrive.modeld.compile_modeld import make_warp_input_queues as stock_make_warp_queues
+    from openpilot.sunnypilot.modeld_v2.compile_modeld import make_warp_queues as sunny_make_warp_queues
+    stock_vision_shapes = {'img': (1, 12, 128, 256)}  # for now?
+    stock_queues, stock_npy = stock_make_warp_queues(stock_vision_shapes, frame_skip=4, device='NPY')
+    sunny_queues, sunny_npy = sunny_make_warp_queues(device='NPY')
+
+    assert set(sunny_npy.keys()) == set(stock_npy.keys()) == {'tfm', 'big_tfm'}
+    for key in sunny_npy:
+      assert sunny_npy[key].shape == stock_npy[key].shape == (3, 3)
+
+

openpilot/sunnypilot/models/fetcher.py

@@ -141,7 +141,7 @@ class ModelCache:
 class ModelFetcher:
   """Handles fetching and caching of model data from remote source"""
   MODEL_URL = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_v20.json"
-  MODEL_URL_USBGPU = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_usbgpu_v21.json"
+  MODEL_URL_USBGPU = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_usbgpu_v22.json"
 
   def __init__(self, params: Params):
     self.params = params