diff --git a/sunnypilot/modeld_v2/model_metadata_lookup.py b/sunnypilot/modeld_v2/model_metadata_lookup.py index 41340b6268..400ab5d3a7 100644 --- a/sunnypilot/modeld_v2/model_metadata_lookup.py +++ b/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) + } + }, } diff --git a/sunnypilot/modeld_v2/modeld.py b/sunnypilot/modeld_v2/modeld.py index f4b728844a..0d5ab00ae4 100755 --- a/sunnypilot/modeld_v2/modeld.py +++ b/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: diff --git a/sunnypilot/modeld_v2/tests/test_buffer_logic_inspect.py b/sunnypilot/modeld_v2/tests/test_buffer_logic_inspect.py index d710484461..0750c8b29c 100644 --- a/sunnypilot/modeld_v2/tests/test_buffer_logic_inspect.py +++ b/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" \ No newline at end of file