refactor: update input handling and buffer logic in modeld

sunnypilot/modeld_v2/model_metadata_lookup.py

@@ -145,27 +145,6 @@ MODEL_METADATA = {
       "pad": slice(-2, None),
     },
   },
-  "driving_vision_steam_powered": {
-    "metadata_path": "selfdrive/modeld/models/driving_vision_metadata.pkl",
-    "model_checkpoint": "2454c1c0-faac-47a7-b0a4-244fa0aeb474/500",
-    "split": True,
-    "input_shapes": {"img": (1, 12, 128, 256), "big_img": (1, 12, 128, 256)},
-    "output_shapes": {"outputs": (1, 1576)},
-    "output_slices": {
-      "meta": slice(0, 55),
-      "desire_pred": slice(55, 87),
-      "pose": slice(87, 99),
-      "wide_from_device_euler": slice(99, 105),
-      "road_transform": slice(105, 117),
-      "lane_lines": slice(117, 645),
-      "lane_lines_prob": slice(645, 653),
-      "road_edges": slice(653, 917),
-      "lead": slice(917, 1061),
-      "lead_prob": slice(1061, 1064),
-      "hidden_state": slice(1064, 1576),
-      "pad": slice(0, None),
-    },
-  },
   "supercombo_op": {
     "metadata_path": "/Users/james/Downloads/model-Optimus Prime (September 21, 2023)-559/supercombo_op_metadata.pkl",
     "model_checkpoint": None,
@@ -262,6 +241,24 @@ MODEL_METADATA = {
     },
     "20hz": True
   },
+  "driving_policy_renamed_desire": {
+    "metadata_path": "/Users/james/Downloads/model-ugh (August 27, 2025)-575/driving_policy_ugh_metadata.pkl",
+    "model_checkpoint": "a8f96b93-bde2-4e28-a732-4df21ebba968/400",
+    "split": True,
+    "input_shapes": {
+      "desire_pulse": (1, 25, 8),
+      "traffic_convention": (1, 2),
+      "features_buffer": (1, 25, 512)
+    },
+    "output_shapes": {
+      "outputs": (1, 1000)
+    },
+    "output_slices": {
+      "plan": slice(0, 990),
+      "desire_state": slice(990, 998),
+      "pad": slice(-2, None)
+    }
+  },
 }
 
 

sunnypilot/modeld_v2/modeld.py

@@ -40,11 +40,84 @@ class FrameMeta:
       self.frame_id, self.timestamp_sof, self.timestamp_eof = vipc.frame_id, vipc.timestamp_sof, vipc.timestamp_eof
 
 
+class InputQueues:
+  def __init__(self, input_shapes: dict, input_dtypes: dict):
+    self.input_shapes = input_shapes
+    self.input_dtypes = input_dtypes
+    self.buffers: dict[str, np.ndarray | None] = {}
+    self.indices: dict[str, np.ndarray | None] = {}
+    for key, shape in input_shapes.items():
+      self._setup_buffer_for_key(key, shape, input_dtypes[key])
+
+  def _setup_buffer_for_key(self, key, shape, dtype):
+    # Temporal input: shape is [batch, history, features]
+    if len(shape) == 3 and shape[1] > 1:
+      buffer_history_len = max(100, shape[1] * 4 if shape[1] < 100 else shape[1])
+      self.buffers[key] = np.zeros((1, buffer_history_len, shape[2]), dtype=dtype)
+      features_buffer_shape = self.input_shapes.get('features_buffer')
+      if shape[1] in (24, 25) and features_buffer_shape and features_buffer_shape[1] == 24:
+        step = int(-buffer_history_len / shape[1])
+        self.indices[key] = np.arange(step, step * (shape[1] + 1), step)[::-1]
+      elif shape[1] == 25:
+        skip = buffer_history_len // shape[1]
+        self.indices[key] = np.arange(buffer_history_len)[-1 - (skip * (shape[1] - 1))::skip]
+      elif shape[1] == buffer_history_len:
+        self.indices[key] = np.arange(buffer_history_len)
+      else:
+        self.indices[key] = None
+    else:
+      self.buffers[key] = self.indices[key] = None
+
+  def update_dtypes_and_shapes(self, input_dtypes: dict, input_shapes: dict) -> None:
+    self.input_dtypes.update(input_dtypes)
+    self.input_shapes.update(input_shapes)
+    for key in input_dtypes:
+      if key in self.buffers and self.buffers[key] is not None:
+        shape = input_shapes[key]
+        self._setup_buffer_for_key(key, shape, input_dtypes[key])
+
+  def reset(self) -> None:
+    for key, shape in self.input_shapes.items():
+      self._setup_buffer_for_key(key, shape, self.input_dtypes[key])
+
+  def enqueue(self, inputs: dict[str, np.ndarray]) -> None:
+    for key, new_val in inputs.items():
+      if key not in self.buffers or self.buffers[key] is None:
+        continue
+      if new_val.dtype != self.input_dtypes[key]:
+        raise ValueError(f'Input {key} has wrong dtype {new_val.dtype}, expected {self.input_dtypes[key]}')
+      buf = self.buffers[key]
+      if buf is not None:
+        if buf.shape[1] == new_val.shape[0]:
+          buf[0, -new_val.shape[0]:] = new_val
+          buf[0, :-new_val.shape[0]] = buf[0, new_val.shape[0]:]
+        else:
+          buf[0, :-1] = buf[0, 1:]
+          buf[0, -1] = new_val
+
+  def get(self, *names) -> dict[str, np.ndarray]:
+    result: dict[str, np.ndarray] = {}
+    for key in names:
+      buf = self.buffers.get(key, None)
+      if buf is not None:
+        out_shape = self.input_shapes.get(key)
+        # Roll buffer and assign based on desire.shape[1] value
+        if out_shape is not None and key.startswith('desire') and buf.shape[1] > out_shape[1]:
+          skip = buf.shape[1] // out_shape[1]
+          result[key] = buf.reshape((out_shape[0], out_shape[1], skip, -1)).max(axis=2)
+        elif self.indices[key] is not None and buf.shape[1] > 1:
+          result[key] = buf[0, self.indices[key]]
+        elif out_shape is not None and buf.shape[1] >= out_shape[1]:
+          result[key] = buf[0, -out_shape[1]:]
+        else:
+          result[key] = buf[0]
+    return result
+
+
 class ModelState(ModelStateBase):
   frames: dict[str, DrivingModelFrame]
   inputs: dict[str, np.ndarray]
   prev_desire: np.ndarray  # for tracking the rising edge of the pulse
-  temporal_idxs: slice | np.ndarray
 
   def __init__(self, context: CLContext):
     ModelStateBase.__init__(self)
@@ -56,39 +129,22 @@ class ModelState(ModelStateBase):
       raise
 
     model_bundle = get_active_bundle()
-    self.generation = model_bundle.generation if model_bundle is not None else None
-    overrides = {override.key: override.value for override in model_bundle.overrides}
+    self.generation = model_bundle.generation if model_bundle else None
+    overrides = {override.key: override.value for override in model_bundle.overrides} if model_bundle else {}
 
     self.LAT_SMOOTH_SECONDS = float(overrides.get('lat', ".0"))
     self.LONG_SMOOTH_SECONDS = float(overrides.get('long', ".0"))
     self.MIN_LAT_CONTROL_SPEED = 0.3
 
-    buffer_length = 5 if self.model_runner.is_20hz else 2
+    buffer_length = 4 if self.model_runner.is_20hz else 2
     self.frames = {name: DrivingModelFrame(context, buffer_length) for name in self.model_runner.vision_input_names}
-    self.prev_desire = np.zeros(self.constants.DESIRE_LEN, dtype=np.float32)
 
-    # img buffers are managed in openCL transform code
-    self.numpy_inputs = {}
-    self.temporal_buffers = {}
-    self.temporal_idxs_map = {}
+    input_dtypes = dict.fromkeys(self.model_runner.input_shapes, np.float32)
+    self.numpy_inputs = {k: np.zeros(shape, dtype=input_dtypes[k]) for k, shape in self.model_runner.input_shapes.items() if k not in self.frames}
 
-    for key, shape in self.model_runner.input_shapes.items():
-      if key not in self.frames: # Managed by opencl
-        self.numpy_inputs[key] = np.zeros(shape, dtype=np.float32)
-        # Temporal input: shape is [batch, history, features]
-        if len(shape) == 3 and shape[1] > 1:
-          buffer_history_len = max(100, (shape[1] * 4 if shape[1] < 100 else shape[1]))   # Allow for higher history buffers in the future
-          feature_len = shape[2]
-          self.temporal_buffers[key] = np.zeros((1, buffer_history_len, feature_len), dtype=np.float32)
-          features_buffer_shape = self.model_runner.input_shapes.get('features_buffer')
-          if shape[1] in (24, 25) and features_buffer_shape is not None and features_buffer_shape[1] == 24:  # 20Hz
-            step = int(-buffer_history_len / shape[1])
-            self.temporal_idxs_map[key] = np.arange(step, step * (shape[1] + 1), step)[::-1]
-          elif shape[1] == 25:  # Split
-            skip = buffer_history_len // shape[1]
-            self.temporal_idxs_map[key] = np.arange(buffer_history_len)[-1 - (skip * (shape[1] - 1))::skip]
-          elif shape[1] == buffer_history_len:  # non20hz
-            self.temporal_idxs_map[key] = np.arange(buffer_history_len)
+    temporal_inputs = {k: v for k, v in self.model_runner.input_shapes.items() if len(v) == 3 and v[1] > 1}
+    self.input_queues = InputQueues(temporal_inputs, dict.fromkeys(temporal_inputs, np.float32))
+    self.prev_desire = np.zeros(self.numpy_inputs[self.desire_key].shape[2], dtype=np.float32)
 
   @property
   def mlsim(self) -> bool:
@@ -104,24 +160,16 @@ class ModelState(ModelStateBase):
     inputs[self.desire_key][0] = 0
     new_desire = np.where(inputs[self.desire_key] - self.prev_desire > .99, inputs[self.desire_key], 0)
     self.prev_desire[:] = inputs[self.desire_key]
-    if self.numpy_inputs[self.desire_key].shape[1] == self.temporal_buffers[self.desire_key].shape[1]:
-      desire_len = inputs[self.desire_key].shape[-1]
-      self.temporal_buffers[self.desire_key][0][:-desire_len] = self.temporal_buffers[self.desire_key][0][desire_len:]
-      self.temporal_buffers[self.desire_key][0][-desire_len:] = new_desire
-    else:
-      self.temporal_buffers[self.desire_key][0,:-1] = self.temporal_buffers[self.desire_key][0,1:]
-      self.temporal_buffers[self.desire_key][0,-1] = new_desire
 
-    # Roll buffer and assign based on desire.shape[1] value
-    if self.temporal_buffers[self.desire_key].shape[1] > self.numpy_inputs[self.desire_key].shape[1]:
-      skip = self.temporal_buffers[self.desire_key].shape[1] // self.numpy_inputs[self.desire_key].shape[1]
-      self.numpy_inputs[self.desire_key][:] = (self.temporal_buffers[self.desire_key][0].reshape(
-                                               self.numpy_inputs[self.desire_key].shape[0], self.numpy_inputs[self.desire_key].shape[1], skip, -1).max(axis=2))
-    else:
-      self.numpy_inputs[self.desire_key][:] = self.temporal_buffers[self.desire_key][0, self.temporal_idxs_map[self.desire_key]]
+    batch_inputs = {key: (new_desire if key == self.desire_key else inputs[key])
+                    for key in self.input_queues.buffers
+                    if not (key == 'features_buffer' and 'hidden_state' in self.numpy_inputs) and (key == self.desire_key or key in inputs)}
+    self.input_queues.enqueue(batch_inputs)
 
     for key in self.numpy_inputs:
-      if key in inputs and key not in [self.desire_key]:
+      if key in self.input_queues.buffers and self.input_queues.buffers[key] is not None:
+        self.numpy_inputs[key][:] = self.input_queues.get(key)[key]
+      elif key in inputs:
         self.numpy_inputs[key][:] = inputs[key]
 
     imgs_cl = {name: self.frames[name].prepare(bufs[name], transforms[name].flatten()) for name in self.model_runner.vision_input_names}
@@ -136,18 +184,14 @@ class ModelState(ModelStateBase):
     outputs = self.model_runner.run_model()
 
     if "lat_planner_solution" in outputs and "lat_planner_state" in self.numpy_inputs:
-      idx_n = outputs['lat_planner_solution'].shape[1]    # Reshaped by parse_mdn from slice(5990, 6254)= 264= 1,264/2 = 1,1,132/4 == 1,33,4
+      idx_n = outputs['lat_planner_solution'].shape[1]
       t_idxs = [10.0 * ((i / (idx_n - 1))**2) for i in range(idx_n)]
       self.numpy_inputs['lat_planner_state'][2] = np.interp(DT_MDL, t_idxs, outputs['lat_planner_solution'][0, :, 2])
       self.numpy_inputs['lat_planner_state'][3] = np.interp(DT_MDL, t_idxs, outputs['lat_planner_solution'][0, :, 3])
 
-    # Update features_buffer
-    self.temporal_buffers['features_buffer'][0, :-1] = self.temporal_buffers['features_buffer'][0, 1:]
-    self.temporal_buffers['features_buffer'][0, -1] = outputs['hidden_state'][0, :]
-    if 'features_buffer' in self.temporal_idxs_map:
-      self.numpy_inputs['features_buffer'][:] = self.temporal_buffers['features_buffer'][0, self.temporal_idxs_map['features_buffer']]
-    else:
-      self.numpy_inputs['features_buffer'][:] = self.temporal_buffers['features_buffer'][0, -self.numpy_inputs['features_buffer'].shape[1]:]
+    # Enqueue features buffer
+    self.input_queues.enqueue({'features_buffer': outputs['hidden_state'][0, :]})
+    self.numpy_inputs['features_buffer'][:] = self.input_queues.get('features_buffer')['features_buffer']
 
     if "desired_curvature" in outputs:
       input_name_prev = None
@@ -155,16 +199,16 @@ class ModelState(ModelStateBase):
         input_name_prev = 'prev_desired_curvs'
       elif "prev_desired_curv" in self.numpy_inputs.keys():
         input_name_prev = 'prev_desired_curv'
-      if input_name_prev and input_name_prev in self.temporal_buffers:
+      if input_name_prev and input_name_prev in self.input_queues.buffers:
         self.process_desired_curvature(outputs, input_name_prev)
     return outputs
 
   def process_desired_curvature(self, outputs, input_name_prev):
-    self.temporal_buffers[input_name_prev][0,:-1] = self.temporal_buffers[input_name_prev][0,1:]
-    self.temporal_buffers[input_name_prev][0,-1,:] = outputs['desired_curvature'][0, :]
-    self.numpy_inputs[input_name_prev][:] = self.temporal_buffers[input_name_prev][0, self.temporal_idxs_map[input_name_prev]]
+    self.input_queues.enqueue({input_name_prev: outputs['desired_curvature'][0, :]})
+    self.numpy_inputs[input_name_prev][:] = self.input_queues.get(input_name_prev)[input_name_prev]
     if self.mlsim:
-      self.numpy_inputs[input_name_prev][:] = 0*self.temporal_buffers[input_name_prev][0, self.temporal_idxs_map[input_name_prev]]
+      self.numpy_inputs[input_name_prev][:] = 0 * self.input_queues.get(input_name_prev)[input_name_prev]
+
 
   def get_action_from_model(self, model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action,
                             lat_action_t: float, long_action_t: float, v_ego: float) -> log.ModelDataV2.Action:
@@ -225,19 +269,13 @@ def main(demo=False):
 
   publish_state = PublishState()
   params = Params()
+  frame_dropped_filter = FirstOrderFilter(0., 10., 1. / model.constants.MODEL_RUN_FREQ)
+  frame_id = last_vipc_frame_id = run_count = 0
 
-  # setup filter to track dropped frames
-  frame_dropped_filter = FirstOrderFilter(0., 10., 1. / model.constants.MODEL_FREQ)
-  frame_id = 0
-  last_vipc_frame_id = 0
-  run_count = 0
-
-  model_transform_main = np.zeros((3, 3), dtype=np.float32)
-  model_transform_extra = np.zeros((3, 3), dtype=np.float32)
+  model_transform_main = model_transform_extra = np.zeros((3, 3), dtype=np.float32)
   live_calib_seen = False
-  buf_main, buf_extra = None, None
-  meta_main = FrameMeta()
-  meta_extra = FrameMeta()
+  buf_main = buf_extra = None
+  meta_main = meta_extra = FrameMeta()
 
 
   if demo:

sunnypilot/modeld_v2/tests/test_buffer_logic_inspect.py

@@ -116,8 +116,8 @@ def test_buffer_shapes_and_indices(shapes, mode, apply_patches):
   state = ModelState(None)
   constants = DummyModelRunner(shapes).constants
   for key in shapes:
-    buf = state.temporal_buffers.get(key, None)
-    idxs = state.temporal_idxs_map.get(key, None)
+    buf = state.input_queues.buffers.get(key, None)
+    idxs = state.input_queues.indices.get(key, None)
     if buf is None:
       continue    # not all shapes are 3d, and the non 3d are not buffered
     # Buffer shape logic
@@ -145,7 +145,7 @@ def test_buffer_shapes_and_indices(shapes, mode, apply_patches):
 def legacy_buffer_update(buf, new_val, mode, key, constants, idxs, input_shape, prev_desire=None):
   # This is what we compare the new dynamic logic to, to ensure it does the same thing
   if mode == 'split':
-    if key == 'desire':
+    if key == 'desire' or key.startswith('desire'):
       buf[0,:-1] = buf[0,1:]
       buf[0,-1] = new_val
       return buf.reshape((1, constants.INPUT_HISTORY_BUFFER_LEN, constants.TEMPORAL_SKIP, -1)).max(axis=2)
@@ -211,14 +211,14 @@ def legacy_buffer_update(buf, new_val, mode, key, constants, idxs, input_shape,
 
 
 def dynamic_buffer_update(state, key, new_val, mode):
-  if key == 'desire':
+  if key == 'desire' or key.startswith('desire'):
     inputs = {k: np.zeros(v[2], dtype=np.float32) if len(v) == 3 else np.zeros(v[1], dtype=np.float32)
-              for k, v in state.model_runner.input_shapes.items() if k != 'desire'}
-    inputs['desire'] = new_val.copy()
+              for k, v in state.model_runner.input_shapes.items() if k != key}
+    inputs[key] = new_val.copy()
     # ModelState.run expects desire as a pulse, so we zero the first element.
-    inputs['desire'][0] = 0
+    inputs[key][0] = 0
     state.run({}, {}, inputs, prepare_only=False)
-    return state.numpy_inputs['desire'][0]
+    return state.numpy_inputs[key]
 
   if key == 'features_buffer':
     inputs = {k: np.zeros(v[2], dtype=np.float32) if len(v) == 3 else np.zeros(v[1], dtype=np.float32)
@@ -249,19 +249,33 @@ def dynamic_buffer_update(state, key, new_val, mode):
 @pytest.mark.parametrize("key", ["desire", "features_buffer", "prev_desired_curv"])
 def test_buffer_update_equivalence(shapes, mode, key, apply_patches):
   state = ModelState(None)
-  if key not in state.numpy_inputs:
-    pytest.skip(f"{key} not in state.numpy_inputs")
+  
+  # Dynamically map desire keys to their actual names in the model
+  if key == "desire":
+    # Find any key in shapes that starts with 'desire'
+    desire_keys = [k for k in shapes.keys() if k.startswith('desire')]
+    if desire_keys:
+      actual_key = desire_keys[0]  # Use the first (and likely only) desire key
+    else:
+      pytest.skip("No desire key found in model inputs")
+  else:
+    actual_key = key
+
+  if actual_key not in state.numpy_inputs:
+    pytest.skip(f"{actual_key} not in state.numpy_inputs")
+  
   constants = DummyModelRunner(shapes).constants
-  buf = state.temporal_buffers.get(key, None)
-  idxs = state.temporal_idxs_map.get(key, None)
-  input_shape = shapes[key]
+  buf = state.input_queues.buffers.get(actual_key, None)
+  idxs = state.input_queues.indices.get(actual_key, None)
+  input_shape = shapes[actual_key]
   prev_desire = np.zeros(constants.DESIRE_LEN, dtype=np.float32) if key == 'desire' else None
+  
   for step in range(20):    # multiple steps to ensure history is built up
     new_val = np.full((input_shape[2],), step, dtype=np.float32)
-    expected = legacy_buffer_update(buf, new_val, mode, key, constants, idxs, input_shape, prev_desire)
-    actual = dynamic_buffer_update(state, key, new_val, mode)
+    expected = legacy_buffer_update(buf, new_val, mode, actual_key, constants, idxs, input_shape, prev_desire)
+    actual = dynamic_buffer_update(state, actual_key, new_val, mode)
     # Model returns the reduced numpy_inputs history, compare the last n entries so the test is checking the same slices.
     if expected is not None and actual is not None and expected.shape != actual.shape:
       if expected.ndim == 2 and actual.ndim == 2 and expected.shape[1] == actual.shape[1]:
         expected = expected[-actual.shape[0]:]
-    assert np.allclose(actual, expected), f"{mode} {key}: dynamic buffer update does not match legacy logic"
+    assert np.allclose(actual, expected), f"{mode} {actual_key}: dynamic buffer update does not match legacy logic"