From ad516e619beabbf7b03c4255a774e7b7bc306e82 Mon Sep 17 00:00:00 2001 From: nayan Date: Mon, 27 Jul 2026 16:25:05 -0400 Subject: [PATCH] simplify everything. --- .../sunnypilot/modeld_v2/mhp_inference.py | 128 ---------- openpilot/sunnypilot/modeld_v2/modeld.py | 20 +- .../modeld_v2/parse_model_outputs.py | 151 +++++++----- .../modeld_v2/parse_model_outputs_split.py | 93 +++---- .../modeld_v2/tests/test_dynamic_mhp.py | 228 +++++++----------- 5 files changed, 211 insertions(+), 409 deletions(-) delete mode 100644 openpilot/sunnypilot/modeld_v2/mhp_inference.py diff --git a/openpilot/sunnypilot/modeld_v2/mhp_inference.py b/openpilot/sunnypilot/modeld_v2/mhp_inference.py deleted file mode 100644 index e686c18c0e..0000000000 --- a/openpilot/sunnypilot/modeld_v2/mhp_inference.py +++ /dev/null @@ -1,128 +0,0 @@ -"""Infer MHP (mixture-density-hypothesis) parameter values from output slice sizes. - -The legacy supercombo encoded its outputs using fixed values for -``PLAN_MHP_N``, ``PLAN_MHP_SELECTION``, ``LEAD_MHP_N``, and ``LEAD_MHP_SELECTION``. -Newer supercombo architectures emit differently-sized slices for the same heads, -yet the parser downstream still expects to be told the layout. - -This module figures those numbers out at runtime from the model's output slice, -so the existing ``parse_mdn`` path can handle any supercombo flavor without code -changes for each variant. ``infer_mhp`` always tries the legacy values first so -existing compiled pkls behave identically (full backwards compatibility). - -Schema reminder (``parse_mdn`` packs hypotheses contiguously in the order -``in_N × (mu | std | weights)`` along the channel axis): - - slice_size == in_N × (2·n_values + out_N) - -where ``n_values`` is the per-hypothesis value width (``IDX_N × PLAN_WIDTH`` -for plan, ``LEAD_TRAJ_LEN × LEAD_WIDTH`` for lead, etc.). -""" - -def infer_mhp( - slice_size: int, - n_values: int, - legacy_in_n: int, - legacy_out_n: int, - max_in_n: int = 16, -) -> tuple[int, int]: - """Infer ``(in_N, out_N)`` for an MDN-encoded output slice. - - Tries values in this order, returning on the first match: - - 1. ``(legacy_in_n, legacy_out_n)`` exactly — preserves exact backwards - compatibility for existing supercombo pkls. - 2. ``out_N ∈ (0, 1, 3)`` (no weights, single weight, three-way selection) at - increasing in_N — covers the common architectural patterns. - 3. Brute-force any valid ``out_N`` that divides ``slice_size``. - - Args: - slice_size: Number of floats in the output slice (typically - ``slices[name].stop - slices[name].start``). - n_values: Per-hypothesis mu/std width (e.g. ``IDX_N × PLAN_WIDTH`` for plan). - legacy_in_n: The legacy in_N value (highest priority for backwards compat). - legacy_out_n: The legacy out_N value. - max_in_n: Upper bound on accepted hypothesis counts (filters silly parses). - - Returns: - ``(in_N, out_N)``. If nothing fits the formulas, returns ``(1, 0)`` — - single hypothesis with no weights — which is the gentlest fallback. - """ - if slice_size <= 0 or n_values <= 0: - return 1, 0 - - # Priority 1: exact legacy match (BC-preserving). - per_hyp_legacy = 2 * n_values + legacy_out_n - if per_hyp_legacy > 0 and legacy_in_n * per_hyp_legacy == slice_size: - return legacy_in_n, legacy_out_n - - # Priority 2: common weight layouts across supercombo variants. - for out_n in (0, 1, 3): - per_hyp = 2 * n_values + out_n - if per_hyp <= 0: - continue - if slice_size % per_hyp == 0: - in_n = slice_size // per_hyp - if 1 <= in_n <= max_in_n: - return in_n, out_n - - # Priority 3: brute-force any divisor that yields a sensible in_N. - # Bound out_n by max_in_n to keep the search tiny (3 hypotheses - # of weights is already a lot). - for out_n in range(0, max_in_n + 1): - per_hyp = 2 * n_values + out_n - if per_hyp <= 0: - continue - if slice_size % per_hyp == 0: - in_n = slice_size // per_hyp - if 1 <= in_n <= max_in_n: - return in_n, out_n - - # Last resort: best-effort single hypothesis with no weights. - return 1, 0 - - -def slice_size(sl) -> int: - """Return the float-width of a ``slice``/``None`` from ``output_slices``. - - ``None`` and ``slice(None, None, None)`` are treated as 0 (output absent). - Negative-end slices (ONNX-style "from end") aren't supported by this helper - because parser expects single-tensor packed outputs. - """ - if sl is None: - return 0 - start = 0 if sl.start is None else sl.start - stop = sl.stop - if stop is None or stop < 0: - return 0 - return max(0, stop - start) - - -def infer_mhp_for_outputs( - output_slices: dict, - constants, - max_in_n: int = 16, -) -> dict: - """Build a dict of MHP values keyed by head name from a model's output_slices. - - Reads ``output_slices['plan']`` and ``output_slices['lead']`` (if present) - and infers their ``in_N``/``out_N``. Other heads aren't MDN-encoded the same - way and remain driven by the ``constants`` module. - """ - config: dict[str, int] = {} - - plan_size = slice_size(output_slices.get('plan')) - if plan_size > 0: - n = constants.IDX_N * constants.PLAN_WIDTH - in_n, out_n = infer_mhp(plan_size, n, constants.PLAN_MHP_N, constants.PLAN_MHP_SELECTION, max_in_n) - config['plan_mhp_n'] = in_n - config['plan_mhp_selection'] = out_n - - lead_size = slice_size(output_slices.get('lead')) - if lead_size > 0: - n = constants.LEAD_TRAJ_LEN * constants.LEAD_WIDTH - in_n, out_n = infer_mhp(lead_size, n, constants.LEAD_MHP_N, constants.LEAD_MHP_SELECTION, max_in_n) - config['lead_mhp_n'] = in_n - config['lead_mhp_selection'] = out_n - - return config diff --git a/openpilot/sunnypilot/modeld_v2/modeld.py b/openpilot/sunnypilot/modeld_v2/modeld.py index 5182a191ab..17c1851c54 100755 --- a/openpilot/sunnypilot/modeld_v2/modeld.py +++ b/openpilot/sunnypilot/modeld_v2/modeld.py @@ -156,25 +156,15 @@ class ModelState(ModelStateBase): from openpilot.sunnypilot.modeld_v2.constants import ModelConstants self.constants = ModelConstants() - # Derive the parser's per-head MHP values from the appropriate output - # slices. Legacy pkls fall back to the constants values (Priority 1 in - # ``mhp_inference.infer_mhp``), so this is fully backwards compatible. - # Supercombo pkls carry plan/lead in the vision tensor; split/multi-policy - # pkls carry them on the policy tensor (and even there we use the first - # policy's slices -- the existing code only tracks one ``policy_output_slices``). - from openpilot.sunnypilot.modeld_v2.mhp_inference import infer_mhp_for_outputs - - if self._combined_model_type == 'supercombo': - mhp_config = infer_mhp_for_outputs(self.vision_output_slices, self.constants) - else: - mhp_config = infer_mhp_for_outputs(self.policy_output_slices, self.constants) - + # Combined parsers auto-detect ``(in_N, out_N)`` for plan/lead from the raw + # slice size, so they transparently support both legacy and newer + # supercombo ONNXes without any per-model configuration. from openpilot.sunnypilot.modeld_v2.parse_model_outputs_split import Parser as SplitParser from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser as CombinedParser if self._combined_model_type != 'supercombo': - self.parser = SplitParser(mhp_config=mhp_config) + self.parser = SplitParser() else: - self.parser = CombinedParser(mhp_config=mhp_config) + self.parser = CombinedParser() self.prev_desire = np.zeros(self.constants.DESIRE_LEN, dtype=np.float32) self.full_frames: dict = {} diff --git a/openpilot/sunnypilot/modeld_v2/parse_model_outputs.py b/openpilot/sunnypilot/modeld_v2/parse_model_outputs.py index 4d56ae5475..3ac2ad88b9 100644 --- a/openpilot/sunnypilot/modeld_v2/parse_model_outputs.py +++ b/openpilot/sunnypilot/modeld_v2/parse_model_outputs.py @@ -1,13 +1,16 @@ import numpy as np from openpilot.sunnypilot.modeld_v2.constants import ModelConstants + def safe_exp(x, out=None): # -11 is around 10**14, more causes float16 overflow return np.exp(np.clip(x, -np.inf, 11), out=out) + def sigmoid(x): return 1. / (1. + safe_exp(-x)) + def softmax(x, axis=-1): x -= np.max(x, axis=axis, keepdims=True) if x.dtype == np.float32 or x.dtype == np.float64: @@ -17,17 +20,34 @@ def softmax(x, axis=-1): x /= np.sum(x, axis=axis, keepdims=True) return x -class Parser: - def __init__(self, ignore_missing=False, mhp_config=None): - self.ignore_missing = ignore_missing - # Optional MHP overrides keyed by head: 'plan_mhp_n', 'plan_mhp_selection', - # 'lead_mhp_n', 'lead_mhp_selection'. ``None`` (or missing keys) keeps the - # legacy ``ModelConstants`` values so existing models behave identically. - self.mhp = mhp_config or {} - def _mhp(self, head, default_in, default_out): - return (self.mhp.get(f'{head}_mhp_n', default_in), - self.mhp.get(f'{head}_mhp_selection', default_out)) +def _infer_mhp(slice_size: int, prod_out_shape: int, max_in_n: int = 16, max_out_n: int = 6) -> tuple[int, int]: + """Derive ``(in_N, out_N)`` from a packed MDN slice. + + Layout (combined supercombo): for each hypothesis we have ``mu``, ``std``, + and an optional scalar ``weight`` block. So: + + slice_size = in_N * (2 * prod_out_shape + out_N) + + We scan small ``out_N`` values (no weights is most common in modern models, + one or three weights in the legacy) and accept the first division that + yields an integer ``in_N`` in ``[1, max_in_n]``. The candidates are spaced + wide enough that there's no ambiguity for the bands we care about. + """ + for out_n in range(max_out_n + 1): + per = 2 * prod_out_shape + out_n + if per <= 0: + continue + if slice_size % per == 0: + in_n = slice_size // per + if 1 <= in_n <= max_in_n: + return in_n, out_n + return 1, 0 # single hypothesis, no weights — matches a non-MDN output + + +class Parser: + def __init__(self, ignore_missing=False): + self.ignore_missing = ignore_missing def check_missing(self, outs, name): if name not in outs and not self.ignore_missing: @@ -48,85 +68,92 @@ class Parser: raw = outs[name] outs[name] = sigmoid(raw) - def parse_mdn(self, name, outs, in_N=0, out_N=1, out_shape=None): + def parse_mdn(self, name, outs, out_shape, in_N=0, out_N=0): + """Parse a packed MDN output. Pass ``in_N``/``out_N`` explicitly for the + legacy layout; pass neither (defaults of 0) to auto-detect from the + slice size, which transparently supports newer supercombos that drop the + per-hypothesis weight block or change the hypothesis count.""" if self.check_missing(outs, name): return raw = outs[name] - raw = raw.reshape((raw.shape[0], max(in_N, 1), -1)) + + if in_N == 0 and out_N == 0: + prod = int(np.prod(out_shape)) + in_N, out_N = _infer_mhp(raw.shape[1], prod) + + raw = raw.reshape((raw.shape[0], in_N, -1)) n_values = (raw.shape[2] - out_N)//2 pred_mu = raw[:,:,:n_values] pred_std = safe_exp(raw[:,:,n_values: 2*n_values]) - if in_N > 1: - if out_N > 0: - weights = np.zeros((raw.shape[0], in_N, out_N), dtype=raw.dtype) - for i in range(out_N): - weights[:,:,i - out_N] = softmax(raw[:,:,i - out_N], axis=-1) + if in_N > 1 and out_N > 0: + weights = np.zeros((raw.shape[0], in_N, out_N), dtype=raw.dtype) + for i in range(out_N): + weights[:,:,i - out_N] = softmax(raw[:,:,i - out_N], axis=-1) - if out_N == 1: - for fidx in range(weights.shape[0]): - idxs = np.argsort(weights[fidx][:,0])[::-1] - weights[fidx] = weights[fidx][idxs] - pred_mu[fidx] = pred_mu[fidx][idxs] - pred_std[fidx] = pred_std[fidx][idxs] - assert out_shape is not None - full_shape = tuple([raw.shape[0], in_N] + list(out_shape)) - outs[name + '_weights'] = weights - outs[name + '_hypotheses'] = pred_mu.reshape(full_shape) - outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape) - - pred_mu_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype) - pred_std_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype) + if out_N == 1: for fidx in range(weights.shape[0]): - for hidx in range(out_N): - idxs = np.argsort(weights[fidx,:,hidx])[::-1] - pred_mu_final[fidx, hidx] = pred_mu[fidx, idxs[0]] - pred_std_final[fidx, hidx] = pred_std[fidx, idxs[0]] - else: - # MHP without weights: keep every hypothesis intact, surface them as - # ``*_hypotheses`` outputs and use the same shape for the primary - # output so downstream consumers can iterate over the full set. - assert out_shape is not None - full_shape = tuple([raw.shape[0], in_N] + list(out_shape)) - outs[name + '_hypotheses'] = pred_mu.reshape(full_shape) - outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape) - pred_mu_final = pred_mu - pred_std_final = pred_std + idxs = np.argsort(weights[fidx][:,0])[::-1] + weights[fidx] = weights[fidx][idxs] + pred_mu[fidx] = pred_mu[fidx][idxs] + pred_std[fidx] = pred_std[fidx][idxs] + full_shape = tuple([raw.shape[0], in_N] + list(out_shape)) + outs[name + '_weights'] = weights + outs[name + '_hypotheses'] = pred_mu.reshape(full_shape) + outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape) + + pred_mu_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype) + pred_std_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype) + for fidx in range(weights.shape[0]): + for hidx in range(out_N): + idxs = np.argsort(weights[fidx,:,hidx])[::-1] + pred_mu_final[fidx, hidx] = pred_mu[fidx, idxs[0]] + pred_std_final[fidx, hidx] = pred_std[fidx, idxs[0]] + elif in_N > 1 and out_N == 0: + # MHP without weights: keep every hypothesis intact, surface them as + # ``*_hypotheses`` and propagate the full multi-hypothesis tensor. + full_shape = tuple([raw.shape[0], in_N] + list(out_shape)) + outs[name + '_hypotheses'] = pred_mu.reshape(full_shape) + outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape) + pred_mu_final = pred_mu + pred_std_final = pred_std else: pred_mu_final = pred_mu pred_std_final = pred_std + # Final-shape selector: keep an extra hypothesis axis only when + # multi-hypothesis data must survive — i.e. when out_N > 1 (legacy + # multiple-selection) OR when in_N > 1 with no weights (newer MHP). For + # single-hypothesis / collapsed cases (lane_lines, pose, etc.) drop the + # extra axis so the consumer sees the historical ``(batch, *out_shape)``. if out_N > 1 or (in_N > 1 and out_N == 0): - assert out_shape is not None n_selections = out_N if out_N > 1 else in_N final_shape = tuple([raw.shape[0], n_selections] + list(out_shape)) else: - assert out_shape is not None final_shape = tuple([raw.shape[0],] + list(out_shape)) outs[name] = pred_mu_final.reshape(final_shape) outs[name + '_stds'] = pred_std_final.reshape(final_shape) def parse_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: - plan_in, plan_out = self._mhp('plan', ModelConstants.PLAN_MHP_N, ModelConstants.PLAN_MHP_SELECTION) - lead_in, lead_out = self._mhp('lead', ModelConstants.LEAD_MHP_N, ModelConstants.LEAD_MHP_SELECTION) - self.parse_mdn('plan', outs, in_N=plan_in, out_N=plan_out, - out_shape=(ModelConstants.IDX_N,ModelConstants.PLAN_WIDTH)) - self.parse_mdn('lane_lines', outs, in_N=0, out_N=0, out_shape=(ModelConstants.NUM_LANE_LINES,ModelConstants.IDX_N,ModelConstants.LANE_LINES_WIDTH)) - self.parse_mdn('road_edges', outs, in_N=0, out_N=0, out_shape=(ModelConstants.NUM_ROAD_EDGES,ModelConstants.IDX_N,ModelConstants.LANE_LINES_WIDTH)) - self.parse_mdn('pose', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,)) - self.parse_mdn('road_transform', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,)) + # Pass no explicit ``in_N``/``out_N`` for plan/lead — ``parse_mdn`` infers + # them from the raw slice size, which naturally handles both the legacy + # supercombo (4955 / 102) and newer variants (e.g. 990 / 144). + self.parse_mdn('plan', outs, out_shape=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) + self.parse_mdn('lane_lines', outs, out_shape=(ModelConstants.NUM_LANE_LINES, ModelConstants.IDX_N, ModelConstants.LANE_LINES_WIDTH)) + self.parse_mdn('road_edges', outs, out_shape=(ModelConstants.NUM_ROAD_EDGES, ModelConstants.IDX_N, ModelConstants.LANE_LINES_WIDTH)) + self.parse_mdn('pose', outs, out_shape=(ModelConstants.POSE_WIDTH,)) + self.parse_mdn('road_transform', outs, out_shape=(ModelConstants.POSE_WIDTH,)) if 'sim_pose' in outs: - self.parse_mdn('sim_pose', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,)) - self.parse_mdn('wide_from_device_euler', outs, in_N=0, out_N=0, out_shape=(ModelConstants.WIDE_FROM_DEVICE_WIDTH,)) - self.parse_mdn('lead', outs, in_N=lead_in, out_N=lead_out, - out_shape=(ModelConstants.LEAD_TRAJ_LEN,ModelConstants.LEAD_WIDTH)) + self.parse_mdn('sim_pose', outs, out_shape=(ModelConstants.POSE_WIDTH,)) + self.parse_mdn('wide_from_device_euler', outs, out_shape=(ModelConstants.WIDE_FROM_DEVICE_WIDTH,)) + self.parse_mdn('lead', outs, out_shape=(ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH)) if 'lat_planner_solution' in outs: - self.parse_mdn('lat_planner_solution', outs, in_N=0, out_N=0, out_shape=(ModelConstants.IDX_N,ModelConstants.LAT_PLANNER_SOLUTION_WIDTH)) + self.parse_mdn('lat_planner_solution', outs, out_shape=(ModelConstants.IDX_N, ModelConstants.LAT_PLANNER_SOLUTION_WIDTH)) if 'desired_curvature' in outs: - self.parse_mdn('desired_curvature', outs, in_N=0, out_N=0, out_shape=(ModelConstants.DESIRED_CURV_WIDTH,)) + self.parse_mdn('desired_curvature', outs, out_shape=(ModelConstants.DESIRED_CURV_WIDTH,)) for k in ['lead_prob', 'lane_lines_prob', 'meta']: self.parse_binary_crossentropy(k, outs) self.parse_categorical_crossentropy('desire_state', outs, out_shape=(ModelConstants.DESIRE_PRED_WIDTH,)) - self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN,ModelConstants.DESIRE_PRED_WIDTH)) + self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN, ModelConstants.DESIRE_PRED_WIDTH)) return outs diff --git a/openpilot/sunnypilot/modeld_v2/parse_model_outputs_split.py b/openpilot/sunnypilot/modeld_v2/parse_model_outputs_split.py index 8c152ad220..895a42ed4b 100644 --- a/openpilot/sunnypilot/modeld_v2/parse_model_outputs_split.py +++ b/openpilot/sunnypilot/modeld_v2/parse_model_outputs_split.py @@ -22,16 +22,8 @@ def softmax(x, axis=-1): class Parser: - def __init__(self, ignore_missing=False, mhp_config=None): + def __init__(self, ignore_missing=False): self.ignore_missing = ignore_missing - # Optional overrides for plan/lead ``in_N`` / ``out_N``. ``None`` or any - # missing key falls back to ``SplitModelConstants`` so previously compiled - # pkls continue to behave identically. - self.mhp = mhp_config or {} - - def _mhp(self, head, default_in, default_out): - return (self.mhp.get(f'{head}_mhp_n', default_in), - self.mhp.get(f'{head}_mhp_selection', default_out)) def check_missing(self, outs, name): if name not in outs and not self.ignore_missing: @@ -63,46 +55,36 @@ class Parser: pred_std = safe_exp(raw[:,:,n_values: 2*n_values]) if in_N > 1: - if out_N > 0: - weights = np.zeros((raw.shape[0], in_N, out_N), dtype=raw.dtype) - for i in range(out_N): - weights[:,:,i - out_N] = softmax(raw[:,:,i - out_N], axis=-1) + weights = np.zeros((raw.shape[0], in_N, out_N), dtype=raw.dtype) + for i in range(out_N): + weights[:,:,i - out_N] = softmax(raw[:,:,i - out_N], axis=-1) - if out_N == 1: - for fidx in range(weights.shape[0]): - idxs = np.argsort(weights[fidx][:,0])[::-1] - weights[fidx] = weights[fidx][idxs] - pred_mu[fidx] = pred_mu[fidx][idxs] - pred_std[fidx] = pred_std[fidx][idxs] - assert out_shape is not None - full_shape = tuple([raw.shape[0], in_N] + list(out_shape)) - outs[name + '_weights'] = weights - outs[name + '_hypotheses'] = pred_mu.reshape(full_shape) - outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape) - - pred_mu_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype) - pred_std_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype) + if out_N == 1: for fidx in range(weights.shape[0]): - for hidx in range(out_N): - idxs = np.argsort(weights[fidx,:,hidx])[::-1] - pred_mu_final[fidx, hidx] = pred_mu[fidx, idxs[0]] - pred_std_final[fidx, hidx] = pred_std[fidx, idxs[0]] - else: - # MHP without weights: keep every hypothesis intact. - assert out_shape is not None - full_shape = tuple([raw.shape[0], in_N] + list(out_shape)) - outs[name + '_hypotheses'] = pred_mu.reshape(full_shape) - outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape) - pred_mu_final = pred_mu - pred_std_final = pred_std + idxs = np.argsort(weights[fidx][:,0])[::-1] + weights[fidx] = weights[fidx][idxs] + pred_mu[fidx] = pred_mu[fidx][idxs] + pred_std[fidx] = pred_std[fidx][idxs] + assert out_shape is not None + full_shape = tuple([raw.shape[0], in_N] + list(out_shape)) + outs[name + '_weights'] = weights + outs[name + '_hypotheses'] = pred_mu.reshape(full_shape) + outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape) + + pred_mu_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype) + pred_std_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype) + for fidx in range(weights.shape[0]): + for hidx in range(out_N): + idxs = np.argsort(weights[fidx,:,hidx])[::-1] + pred_mu_final[fidx, hidx] = pred_mu[fidx, idxs[0]] + pred_std_final[fidx, hidx] = pred_std[fidx, idxs[0]] else: pred_mu_final = pred_mu pred_std_final = pred_std - if out_N > 1 or (in_N > 1 and out_N == 0): + if out_N > 1: assert out_shape is not None - n_selections = out_N if out_N > 1 else in_N - final_shape = tuple([raw.shape[0], n_selections] + list(out_shape)) + final_shape = tuple([raw.shape[0], out_N] + list(out_shape)) else: assert out_shape is not None final_shape = tuple([raw.shape[0],] + list(out_shape)) @@ -118,26 +100,15 @@ class Parser: def parse_dynamic_outputs(self, outs: dict[str, np.ndarray]) -> None: if 'lead' in outs: - # Prefer explicit overrides in ``mhp``; otherwise fall back to the - # legacy `is_mhp` heuristic that inspects the raw tensor's last axis. - if self.mhp.get('lead_mhp_n') is not None or self.mhp.get('lead_mhp_selection') is not None: - lead_in_N, lead_out_N = (SplitModelConstants.LEAD_MHP_N, SplitModelConstants.LEAD_MHP_SELECTION) - lead_in_N = self.mhp.get('lead_mhp_n', lead_in_N) - lead_out_N = self.mhp.get('lead_mhp_selection', lead_out_N) - lead_out_shape = (SplitModelConstants.LEAD_TRAJ_LEN, SplitModelConstants.LEAD_WIDTH) - else: - lead_mhp = self.is_mhp(outs, 'lead', - SplitModelConstants.LEAD_MHP_SELECTION * SplitModelConstants.LEAD_TRAJ_LEN * SplitModelConstants.LEAD_WIDTH) - lead_in_N, lead_out_N = (SplitModelConstants.LEAD_MHP_N, SplitModelConstants.LEAD_MHP_SELECTION) if lead_mhp else (0, 0) - lead_out_shape = (SplitModelConstants.LEAD_TRAJ_LEN, SplitModelConstants.LEAD_WIDTH) if lead_mhp else \ - (SplitModelConstants.LEAD_MHP_SELECTION, SplitModelConstants.LEAD_TRAJ_LEN, SplitModelConstants.LEAD_WIDTH) + lead_mhp = self.is_mhp(outs, 'lead', + SplitModelConstants.LEAD_MHP_SELECTION * SplitModelConstants.LEAD_TRAJ_LEN * SplitModelConstants.LEAD_WIDTH) + lead_in_N, lead_out_N = (SplitModelConstants.LEAD_MHP_N, SplitModelConstants.LEAD_MHP_SELECTION) if lead_mhp else (0, 0) + lead_out_shape = (SplitModelConstants.LEAD_TRAJ_LEN, SplitModelConstants.LEAD_WIDTH) if lead_mhp else \ + (SplitModelConstants.LEAD_MHP_SELECTION, SplitModelConstants.LEAD_TRAJ_LEN, SplitModelConstants.LEAD_WIDTH) self.parse_mdn('lead', outs, in_N=lead_in_N, out_N=lead_out_N, out_shape=lead_out_shape) if 'plan' in outs: - if self.mhp.get('plan_mhp_n') is not None or self.mhp.get('plan_mhp_selection') is not None: - plan_in_N, plan_out_N = self._mhp('plan', SplitModelConstants.PLAN_MHP_N, SplitModelConstants.PLAN_MHP_SELECTION) - else: - plan_mhp = self.is_mhp(outs, 'plan', SplitModelConstants.IDX_N * SplitModelConstants.PLAN_WIDTH) - plan_in_N, plan_out_N = (SplitModelConstants.PLAN_MHP_N, SplitModelConstants.PLAN_MHP_SELECTION) if plan_mhp else (0, 0) + plan_mhp = self.is_mhp(outs, 'plan', SplitModelConstants.IDX_N * SplitModelConstants.PLAN_WIDTH) + plan_in_N, plan_out_N = (SplitModelConstants.PLAN_MHP_N, SplitModelConstants.PLAN_MHP_SELECTION) if plan_mhp else (0, 0) self.parse_mdn('plan', outs, in_N=plan_in_N, out_N=plan_out_N, out_shape=(SplitModelConstants.IDX_N, SplitModelConstants.PLAN_WIDTH)) if 'planplus' in outs: @@ -165,7 +136,7 @@ class Parser: self.parse_mdn('road_edges', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.NUM_ROAD_EDGES,SplitModelConstants.IDX_N,SplitModelConstants.LANE_LINES_WIDTH)) if 'sim_pose' in outs: - self.parse_mdn('sim_pose', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.POSE_WIDTH,)) + self.parse_mdn('sim_pose', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.POSE_WIDTH,)) if 'action' in outs: self.parse_mdn('action', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.ACTION_WIDTH,)) diff --git a/openpilot/sunnypilot/modeld_v2/tests/test_dynamic_mhp.py b/openpilot/sunnypilot/modeld_v2/tests/test_dynamic_mhp.py index efe6421d43..17307fc8f8 100644 --- a/openpilot/sunnypilot/modeld_v2/tests/test_dynamic_mhp.py +++ b/openpilot/sunnypilot/modeld_v2/tests/test_dynamic_mhp.py @@ -1,171 +1,113 @@ -"""Tests for the dynamic MDN inference in ``mhp_inference`` and the parser -changes that read from ``output_slices`` instead of hardcoded constants. -""" +"""Tests for the dynamic MDN-hypothesis inference in ``parse_model_outputs``.""" import numpy as np import pytest -from openpilot.sunnypilot.modeld_v2.mhp_inference import ( - infer_mhp, - slice_size, - infer_mhp_for_outputs, -) -from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser as CombinedParser -from openpilot.sunnypilot.modeld_v2.parse_model_outputs_split import Parser as SplitParser + +from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser, _infer_mhp from openpilot.sunnypilot.modeld_v2.constants import ModelConstants -# -- infer_mhp -------------------------------------------------------------- +# -- _infer_mhp ------------------------------------------------------------- class TestInferMhp: - N_PLAN = ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH # 495 - N_LEAD = ModelConstants.LEAD_TRAJ_LEN * ModelConstants.LEAD_WIDTH # 24 + P_PLAN = ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH # 495 + P_LEAD = ModelConstants.LEAD_TRAJ_LEN * ModelConstants.LEAD_WIDTH # 24 - def test_legacy_plan_preserved(self): - # Legacy: 5 hypotheses x (2*495 + 1) = 4955 - assert infer_mhp(4955, self.N_PLAN, legacy_in_n=5, legacy_out_n=1) == (5, 1) + def test_legacy_plan(self): + # 5 hypotheses * (2*495 + 1) = 4955 + assert _infer_mhp(4955, self.P_PLAN) == (5, 1) - def test_legacy_lead_preserved(self): - # Legacy: 2 hypotheses x (2*24 + 3) = 102 - assert infer_mhp(102, self.N_LEAD, legacy_in_n=2, legacy_out_n=3) == (2, 3) + def test_legacy_lead(self): + # 2 hypotheses * (2*24 + 3) = 102 + assert _infer_mhp(102, self.P_LEAD) == (2, 3) - def test_new_supercombo_plan_single_hypothesis_no_weights(self): - # New combined supercombo: 1 hypothesis x (2*495 + 0) = 990 - assert infer_mhp(990, self.N_PLAN, legacy_in_n=5, legacy_out_n=1) == (1, 0) + def test_new_plan(self): + # 1 hypothesis * (2*495 + 0) = 990 + assert _infer_mhp(990, self.P_PLAN) == (1, 0) - def test_new_supercombo_lead_three_hypotheses_no_weights(self): - # New combined supercombo: 3 hypotheses x (2*24 + 0) = 144 - assert infer_mhp(144, self.N_LEAD, legacy_in_n=2, legacy_out_n=3) == (3, 0) + def test_new_lead(self): + # 3 hypotheses * (2*24 + 0) = 144 + assert _infer_mhp(144, self.P_LEAD) == (3, 0) - def test_out_n_one(self): - # 4 hypotheses x (2*495 + 1) = 3964 - assert infer_mhp(3964, self.N_PLAN, legacy_in_n=5, legacy_out_n=1) == (4, 1) + def test_no_mhp_output(self): + # lane_lines (528 = 4*132): 1 hypothesis * (2*528 + 0) = 1056; 2 hypotheses * 528+0 = 1056 too — same. + # In practice a single-hypothesis MDN with ``out_N == 0`` yields 1*2*prod = 528 * 2 = ... no, lane_lines is + # ``in_N=0`` because it's laid out as 2*prod directly (1056 = 4 * 132 * 2 simply). The auto-detect prefers + # the smallest valid ``in_N``, so 1056 = 2*1*528 + 0 -> (1, 0). + assert _infer_mhp(1056, 528) == (1, 0) - def test_out_n_three(self): - # 2 hypotheses x (2*24 + 3) = 102 (matches legacy_lead as well) - assert infer_mhp(102, self.N_LEAD, legacy_in_n=4, legacy_out_n=99) == (2, 3) - - def test_zero_or_invalid_returns_fallback(self): - assert infer_mhp(0, self.N_PLAN, 5, 1) == (1, 0) - assert infer_mhp(-1, self.N_PLAN, 5, 1) == (1, 0) - assert infer_mhp(990, 0, 5, 1) == (1, 0) - - def test_no_match_returns_single_hypothesis(self): - # 987 doesn't cleanly factor under the constraints we care about. - assert infer_mhp(987, self.N_PLAN, 5, 1) == (1, 0) - - -class TestSliceSize: - def test_none_returns_zero(self): - assert slice_size(None) == 0 - - def test_basic_slice(self): - assert slice_size(slice(10, 50)) == 40 - - def test_negative_stop_returns_zero(self): - assert slice_size(slice(10, -2)) == 0 - - def test_none_bounds(self): - assert slice_size(slice(None, 100)) == 100 - - -class TestInferMhpForOutputs: - def test_infers_for_plan_and_lead(self): - slices = { - 'plan': slice(1576, 2566), # 990 - 'lead': slice(917, 1061), # 144 - } - cfg = infer_mhp_for_outputs(slices, ModelConstants) - assert cfg == {'plan_mhp_n': 1, 'plan_mhp_selection': 0, - 'lead_mhp_n': 3, 'lead_mhp_selection': 0} - - def test_legacy_falls_back_to_constants(self): - # Legacy sizes: 4955 plan, 102 lead -> both match Priority 1. - slices = { - 'plan': slice(0, 4955), - 'lead': slice(4955, 5057), - } - cfg = infer_mhp_for_outputs(slices, ModelConstants) - assert cfg == {'plan_mhp_n': ModelConstants.PLAN_MHP_N, - 'plan_mhp_selection': ModelConstants.PLAN_MHP_SELECTION, - 'lead_mhp_n': ModelConstants.LEAD_MHP_N, - 'lead_mhp_selection': ModelConstants.LEAD_MHP_SELECTION} - - def test_missing_outputs_are_skipped(self): - cfg = infer_mhp_for_outputs({}, ModelConstants) - assert cfg == {} + def test_unknown_size_keeps_single_hypothesis(self): + # 989 doesn't divide cleanly under any out_N ∈ {0..6} for P = 495, so we + # fall back to the safe single-hypothesis default. + assert _infer_mhp(989, self.P_PLAN) == (1, 0) # -- CombinedParser --------------------------------------------------------- -def _synth_outputs(in_n_plan=5, out_n_plan=1, in_n_lead=2, out_n_lead=3, - n_plan=ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH, - n_lead=ModelConstants.LEAD_TRAJ_LEN * ModelConstants.LEAD_WIDTH): - """Synthesize a flat-output dict with the right per-head sizes.""" - plan_size = in_n_plan * (2 * n_plan + out_n_plan) - lead_size = in_n_lead * (2 * n_lead + out_n_lead) - return { - 'plan': np.random.RandomState(0).randn(1, plan_size).astype(np.float32), - 'lead': np.random.RandomState(1).randn(1, lead_size).astype(np.float32), - # Other outputs that parse_outputs() consumes: - 'lane_lines': np.random.RandomState(2).randn(1, 528).astype(np.float32), - 'road_edges': np.random.RandomState(3).randn(1, 264).astype(np.float32), - 'pose': np.random.RandomState(4).randn(1, 12).astype(np.float32), - 'road_transform': np.random.RandomState(5).randn(1, 12).astype(np.float32), - 'wide_from_device_euler': np.random.RandomState(6).randn(1, 6).astype(np.float32), - 'lead_prob': np.random.RandomState(7).randn(1, 3).astype(np.float32), - 'lane_lines_prob':np.random.RandomState(8).randn(1, 8).astype(np.float32), - 'meta': np.random.RandomState(9).randn(1, 55).astype(np.float32), - 'desire_state': np.random.RandomState(10).randn(1, 8).astype(np.float32), - 'desire_pred': np.random.RandomState(11).randn(1, 32).astype(np.float32), +def _synth_outs( + plan_in_n: int = 5, plan_out_n: int = 1, + lead_in_n: int = 2, lead_out_n: int = 3, + extras: bool = True, +) -> dict[str, np.ndarray]: + plan_size = plan_in_n * (2 * ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH + plan_out_n) + lead_size = lead_in_n * (2 * ModelConstants.LEAD_TRAJ_LEN * ModelConstants.LEAD_WIDTH + lead_out_n) + rng = np.random.RandomState(0) + d = { + 'plan': rng.randn(1, plan_size).astype(np.float32), + 'lead': rng.randn(1, lead_size).astype(np.float32), } + if extras: + d.update({ + 'lane_lines': rng.randn(1, 528).astype(np.float32), + 'road_edges': rng.randn(1, 264).astype(np.float32), + 'pose': rng.randn(1, 12).astype(np.float32), + 'road_transform': rng.randn(1, 12).astype(np.float32), + 'wide_from_device_euler': rng.randn(1, 6).astype(np.float32), + 'lead_prob': rng.randn(1, 3).astype(np.float32), + 'lane_lines_prob':rng.randn(1, 8).astype(np.float32), + 'meta': rng.randn(1, 55).astype(np.float32), + 'desire_state': rng.randn(1, 8).astype(np.float32), + 'desire_pred': rng.randn(1, 32).astype(np.float32), + }) + return d class TestCombinedParser: - def test_legacy_keeps_existing_shape(self): - p = CombinedParser() # empty mhp -> legacy constants - out = p.parse_outputs(_synth_outputs(5, 1, 2, 3)) - assert out['plan'].shape == (1, 33, 15) - assert out['plan_stds'].shape == (1, 33, 15) - # Lead primary output collapses to LEAD_MHP_SELECTION=3 selections per - # ``parse_mdn``; raw hypotheses survive as ``lead_hypotheses``. - assert out['lead'].shape == (1, 3, 6, 4) - assert out['lead_stds'].shape == (1, 3, 6, 4) - assert out['plan_hypotheses'].shape == (1, 5, 33, 15) - assert out['lead_hypotheses'].shape == (1, 2, 6, 4) + def test_legacy_supercombo_shapes(self): + p = Parser() + out = p.parse_outputs(_synth_outs(plan_in_n=5, plan_out_n=1, lead_in_n=2, lead_out_n=3)) + assert out['plan'].shape == (1, ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH) + assert out['plan_stds'].shape == (1, ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH) + assert out['lead'].shape == (1, ModelConstants.LEAD_MHP_SELECTION, ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH) + assert out['lead_stds'].shape == (1, ModelConstants.LEAD_MHP_SELECTION, ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH) + # Per-hypothesis outputs preserved by legacy code path + assert out['plan_hypotheses'].shape == (1, 5, ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH) + assert out['lead_hypotheses'].shape == (1, 2, ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH) - def test_new_supercombo_plan_and_lead_parse(self): - p = CombinedParser(mhp_config={ - 'plan_mhp_n': 1, 'plan_mhp_selection': 0, - 'lead_mhp_n': 3, 'lead_mhp_selection': 0, - }) - out = p.parse_outputs(_synth_outputs(1, 0, 3, 0)) - assert out['plan'].shape == (1, 33, 15) - assert out['plan_stds'].shape == (1, 33, 15) - assert out['lead'].shape == (1, 3, 6, 4) - assert out['lead_stds'].shape == (1, 3, 6, 4) - # MHP-without-weights keeps every hypothesis as ``*_hypotheses`` - assert out['lead_hypotheses'].shape == (1, 3, 6, 4) - # Plan with a single hypothesis takes the in_N<=1 branch, which - # (matching legacy behavior) does not emit ``plan_hypotheses``. - assert 'plan_hypotheses' not in out + def test_new_supercombo_shapes(self): + p = Parser() + out = p.parse_outputs(_synth_outs(plan_in_n=1, plan_out_n=0, lead_in_n=3, lead_out_n=0)) + # Plan single hypothesis collapses straight to (1, IDX_N, PLAN_WIDTH) + assert out['plan'].shape == (1, ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH) + assert out['plan_stds'].shape == (1, ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH) + # Lead with 3 hypotheses and no weights keeps all hypotheses + assert out['lead'].shape == (1, 3, ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH) + assert out['lead_stds'].shape == (1, 3, ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH) + assert out['lead_hypotheses'].shape == (1, 3, ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH) - -class TestSplitParser: - def test_default_uses_is_mhp_heuristic(self): - # No mhp_config -> falls back to inspecting the raw tensor's last axis. - # The legacy 102-element lead pack fits the "MHP branch" path. - n = SplitParser() - outs = {'lead': np.zeros((1, 102), dtype=np.float32)} - n.parse_dynamic_outputs(outs) - assert outs['lead'].shape == (1, 3, 6, 4) - assert outs['lead_hypotheses'].shape == (1, 2, 6, 4) - - def test_explicit_mhp_overrides_is_mhp(self): - n = SplitParser(mhp_config={'lead_mhp_n': 3, 'lead_mhp_selection': 0}) - outs = {'lead': np.zeros((1, 144), dtype=np.float32)} - n.parse_dynamic_outputs(outs) - assert outs['lead'].shape == (1, 3, 6, 4) + def test_unknown_size_does_not_crash_legacy_layout(self): + # Provide EVERY output parse_outputs expects so the parser can auto-detect + # each head from the actual raw slice size, including the new-supercombo + # mixture (legacy-shaped lane_lines/etc. + new-shaped plan/lead). + p = Parser() + out = p.parse_outputs(_synth_outs(plan_in_n=1, plan_out_n=0, lead_in_n=3, lead_out_n=0)) + assert out['plan'].shape == (1, ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH) + assert out['lead'].shape == (1, 3, ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH) + # Non-MHP outputs must keep the historical 3D shape (no spurious leading 1). + assert out['lane_lines'].shape == (1, ModelConstants.NUM_LANE_LINES, ModelConstants.IDX_N, ModelConstants.LANE_LINES_WIDTH) + assert out['pose'].shape == (1, ModelConstants.POSE_WIDTH) + assert out['road_transform'].shape == (1, ModelConstants.POSE_WIDTH) if __name__ == '__main__':