from __future__ import annotations from dataclasses import dataclass import numpy as np from openpilot.iqpilot.selfdrive.iqmodeld.models.split_model_constants import SplitModelConstants from openpilot.iqpilot.selfdrive.iqmodeld.config import ModelConstants def _bounded_exp(values, out=None): return np.exp(np.clip(values, -np.inf, 11), out=out) def _sigmoid(values): return 1.0 / (1.0 + _bounded_exp(-values)) def _softmax_last(values, axis=-1): values -= np.max(values, axis=axis, keepdims=True) if values.dtype in (np.float32, np.float64): _bounded_exp(values, out=values) else: values = _bounded_exp(values) values /= np.sum(values, axis=axis, keepdims=True) return values @dataclass(frozen=True) class _MixtureRecipe: input_heads: int output_heads: int final_shape: tuple[int, ...] class _TensorKitchen: def __init__(self, ignore_missing: bool = False): self.ignore_missing = ignore_missing def _grab(self, outputs: dict[str, np.ndarray], tensor_name: str) -> np.ndarray | None: if tensor_name not in outputs: if not self.ignore_missing: raise ValueError(f"Missing output {tensor_name}") return return outputs[tensor_name] def categorical(self, outputs: dict[str, np.ndarray], tensor_name: str, shape=None) -> None: raw = self._grab(outputs, tensor_name) if raw is None: return if shape is not None: raw = raw.reshape((raw.shape[0],) + shape) outputs[tensor_name] = _softmax_last(raw, axis=-1) def binary(self, outputs: dict[str, np.ndarray], tensor_name: str) -> None: raw = self._grab(outputs, tensor_name) if raw is None: return outputs[tensor_name] = _sigmoid(raw) def mixture(self, outputs: dict[str, np.ndarray], tensor_name: str, recipe: _MixtureRecipe) -> None: raw = self._grab(outputs, tensor_name) if raw is None: return reshaped = raw.reshape((raw.shape[0], max(recipe.input_heads, 1), -1)) value_count = (reshaped.shape[2] - recipe.output_heads) // 2 means = reshaped[:, :, :value_count] stds = _bounded_exp(reshaped[:, :, value_count:2 * value_count]) if recipe.input_heads > 1: weights = np.zeros((reshaped.shape[0], recipe.input_heads, recipe.output_heads), dtype=reshaped.dtype) for output_idx in range(recipe.output_heads): weights[:, :, output_idx - recipe.output_heads] = _softmax_last( reshaped[:, :, output_idx - recipe.output_heads], axis=-1 ) if recipe.output_heads == 1: for batch_idx in range(weights.shape[0]): order = np.argsort(weights[batch_idx][:, 0])[::-1] weights[batch_idx] = weights[batch_idx][order] means[batch_idx] = means[batch_idx][order] stds[batch_idx] = stds[batch_idx][order] hypothesis_shape = (reshaped.shape[0], recipe.input_heads, *recipe.final_shape) outputs[f"{tensor_name}_weights"] = weights outputs[f"{tensor_name}_hypotheses"] = means.reshape(hypothesis_shape) outputs[f"{tensor_name}_stds_hypotheses"] = stds.reshape(hypothesis_shape) picked_means = np.zeros((reshaped.shape[0], recipe.output_heads, value_count), dtype=reshaped.dtype) picked_stds = np.zeros((reshaped.shape[0], recipe.output_heads, value_count), dtype=reshaped.dtype) for batch_idx in range(weights.shape[0]): for output_idx in range(recipe.output_heads): order = np.argsort(weights[batch_idx, :, output_idx])[::-1] picked_means[batch_idx, output_idx] = means[batch_idx, order[0]] picked_stds[batch_idx, output_idx] = stds[batch_idx, order[0]] else: picked_means = means picked_stds = stds final_shape = ((reshaped.shape[0], recipe.output_heads, *recipe.final_shape) if recipe.output_heads > 1 else (reshaped.shape[0], *recipe.final_shape)) outputs[tensor_name] = picked_means.reshape(final_shape) outputs[f"{tensor_name}_stds"] = picked_stds.reshape(final_shape) class ArchiveParser(_TensorKitchen): def __init__(self, ignore_missing: bool = False): super().__init__(ignore_missing=ignore_missing) self._c = ModelConstants def _recipes(self) -> list[tuple[str, _MixtureRecipe]]: c = self._c return [ ("plan", _MixtureRecipe(c.PLAN_MHP_N, c.PLAN_MHP_SELECTION, (c.IDX_N, c.PLAN_WIDTH))), ("lane_lines", _MixtureRecipe(0, 0, (c.NUM_LANE_LINES, c.IDX_N, c.LANE_LINES_WIDTH))), ("road_edges", _MixtureRecipe(0, 0, (c.NUM_ROAD_EDGES, c.IDX_N, c.LANE_LINES_WIDTH))), ("pose", _MixtureRecipe(0, 0, (c.POSE_WIDTH,))), ("road_transform", _MixtureRecipe(0, 0, (c.POSE_WIDTH,))), ("wide_from_device_euler", _MixtureRecipe(0, 0, (c.WIDE_FROM_DEVICE_WIDTH,))), ("lead", _MixtureRecipe(c.LEAD_MHP_N, c.LEAD_MHP_SELECTION, (c.LEAD_TRAJ_LEN, c.LEAD_WIDTH))), ] def parse_outputs(self, outputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: c = self._c for tensor_name, recipe in self._recipes(): self.mixture(outputs, tensor_name, recipe) if "sim_pose" in outputs: self.mixture(outputs, "sim_pose", _MixtureRecipe(0, 0, (c.POSE_WIDTH,))) if "lat_planner_solution" in outputs: self.mixture(outputs, "lat_planner_solution", _MixtureRecipe(0, 0, (c.IDX_N, c.LAT_PLANNER_SOLUTION_WIDTH))) if "desired_curvature" in outputs: self.mixture(outputs, "desired_curvature", _MixtureRecipe(0, 0, (c.DESIRED_CURV_WIDTH,))) for name in ("lead_prob", "lane_lines_prob", "meta"): self.binary(outputs, name) self.categorical(outputs, "desire_state", shape=(c.DESIRE_PRED_WIDTH,)) self.categorical(outputs, "desire_pred", shape=(c.DESIRE_PRED_LEN, c.DESIRE_PRED_WIDTH)) return outputs class PhaseParser(_TensorKitchen): def __init__(self, ignore_missing: bool = False): super().__init__(ignore_missing=ignore_missing) self._c = SplitModelConstants def _has_mixture_heads(self, outputs: dict[str, np.ndarray], tensor_name: str, flat_width: int) -> bool: raw = self._grab(outputs, tensor_name) if raw is None: return False return raw.shape[1] != 2 * flat_width def _decode_dynamic_family(self, outputs: dict[str, np.ndarray]) -> None: c = self._c if "lead" in outputs: uses_heads = self._has_mixture_heads(outputs, "lead", c.LEAD_MHP_SELECTION * c.LEAD_TRAJ_LEN * c.LEAD_WIDTH) self.mixture(outputs, "lead", _MixtureRecipe( c.LEAD_MHP_N if uses_heads else 0, c.LEAD_MHP_SELECTION if uses_heads else 0, (c.LEAD_TRAJ_LEN, c.LEAD_WIDTH) if uses_heads else (c.LEAD_MHP_SELECTION, c.LEAD_TRAJ_LEN, c.LEAD_WIDTH), )) if "plan" in outputs: uses_heads = self._has_mixture_heads(outputs, "plan", c.IDX_N * c.PLAN_WIDTH) self.mixture(outputs, "plan", _MixtureRecipe( c.PLAN_MHP_N if uses_heads else 0, c.PLAN_MHP_SELECTION if uses_heads else 0, (c.IDX_N, c.PLAN_WIDTH), )) if "planplus" in outputs: self.mixture(outputs, "planplus", _MixtureRecipe(0, 0, (c.IDX_N, c.PLAN_WIDTH))) def _decode_policy_family(self, outputs: dict[str, np.ndarray]) -> None: c = self._c if "action" in outputs: self.mixture(outputs, "action", _MixtureRecipe(0, 0, (c.ACTION_WIDTH,))) if "desired_curvature" in outputs: self.mixture(outputs, "desired_curvature", _MixtureRecipe(0, 0, (c.DESIRED_CURV_WIDTH,))) if "desire_pred" in outputs: self.categorical(outputs, "desire_pred", shape=(c.DESIRE_PRED_LEN, c.DESIRE_PRED_WIDTH)) if "desire_state" in outputs: self.categorical(outputs, "desire_state", shape=(c.DESIRE_PRED_WIDTH,)) if "lane_lines" in outputs: self.mixture(outputs, "lane_lines", _MixtureRecipe(0, 0, (c.NUM_LANE_LINES, c.IDX_N, c.LANE_LINES_WIDTH))) if "lane_lines_prob" in outputs: self.binary(outputs, "lane_lines_prob") if "lead_prob" in outputs: self.binary(outputs, "lead_prob") if "lat_planner_solution" in outputs: self.mixture(outputs, "lat_planner_solution", _MixtureRecipe(0, 0, (c.IDX_N, c.LAT_PLANNER_SOLUTION_WIDTH))) if "meta" in outputs: self.binary(outputs, "meta") if "road_edges" in outputs: self.mixture(outputs, "road_edges", _MixtureRecipe(0, 0, (c.NUM_ROAD_EDGES, c.IDX_N, c.LANE_LINES_WIDTH))) if "sim_pose" in outputs: self.mixture(outputs, "sim_pose", _MixtureRecipe(0, 0, (c.POSE_WIDTH,))) def parse_vision_outputs(self, outputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: c = self._c self.mixture(outputs, "pose", _MixtureRecipe(0, 0, (c.POSE_WIDTH,))) self.mixture(outputs, "wide_from_device_euler", _MixtureRecipe(0, 0, (c.WIDE_FROM_DEVICE_WIDTH,))) self.mixture(outputs, "road_transform", _MixtureRecipe(0, 0, (c.POSE_WIDTH,))) self._decode_dynamic_family(outputs) self._decode_policy_family(outputs) return outputs def parse_policy_outputs(self, outputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: self._decode_dynamic_family(outputs) self._decode_policy_family(outputs) return outputs def parse_outputs(self, outputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: return self.parse_policy_outputs(self.parse_vision_outputs(outputs)) __all__ = [ "ArchiveParser", "PhaseParser", ]