op model 16 deep port

2026-07-26 11:52:04 +08:00 · 2026-07-21 19:08:03 +02:00
parent 2cdd8b0551
commit e94257aa9f
17 changed files with 128 additions and 67 deletions
@@ -86,14 +86,16 @@ frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ

 for usbgpu in [False, True] if USBGPU else [False]:
  target_pkl_path = File(modeld_pkl_path(usbgpu)).abspath
-  # BIG_INTO_SMALL=1 builds the default target from the big model, e.g. to test it without a USB GPU
  file_prefix, cmd_flags = ('big_', usbgpu_tg_flags) if usbgpu else ('big_' if os.getenv('BIG_INTO_SMALL') else '', tg_flags)
-  driving_onnx_deps = get_existing_chunks(File(f"models/{file_prefix}driving_supercombo.onnx").abspath)
+  driving_onnx_deps = [p for m in [f'{file_prefix}driving_vision', f'{file_prefix}driving_on_policy', f'{file_prefix}driving_off_policy']
+                       for p in get_existing_chunks(File(f"models/{m}.onnx").abspath)]
  camera_res_args = ' '.join(f'{cw}x{ch}' for cw, ch in CAMERA_CONFIGS)
  cmd = (f'{cmd_flags} {mac_brew_string} python3 {modeld_dir}/compile_modeld.py '
        f'--model-size {model_w}x{model_h} '
        f'--camera-resolutions {camera_res_args} '
-        f'--onnx {File(f"models/{file_prefix}driving_supercombo.onnx").abspath} '
+        f'--vision-onnx {File(f"models/{file_prefix}driving_vision.onnx").abspath} '
+        f'--off-policy-onnx {File(f"models/{file_prefix}driving_off_policy.onnx").abspath} '
+        f'--on-policy-onnx {File(f"models/{file_prefix}driving_on_policy.onnx").abspath} '
        f'--output {target_pkl_path} --frame-skip {frame_skip}')
  onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps)
  chunk_targets = get_chunk_targets(target_pkl_path, estimate_pickle_max_size(onnx_sizes_sum))
@@ -156,19 +156,17 @@ def get_policy_npy_shapes(input_shapes):
  dp = input_shapes['desire_pulse']  # (1, 25, 8)
  tc = input_shapes['traffic_convention']  # (1, 2)
  at = input_shapes['action_t']  # (1, 2)
-  fb = input_shapes['features_buffer']  # (1, 24, 512)
-  # TODO prev_feat shouldn't exist and be handled inside the JIT, but corrupt on QCOM for now
-  shapes = {'desire': (dp[2],), 'traffic_convention': tuple(tc), 'action_t': tuple(at), 'prev_feat': (fb[0], fb[2])}
+  shapes = {'desire': (dp[2],), 'traffic_convention': tuple(tc), 'action_t': tuple(at)}
  return shapes, [math.prod(s) for s in shapes.values()]


-def make_input_queues(input_shapes, frame_skip, device):
-  input_queues, npy = make_warp_input_queues(input_shapes, frame_skip, device)
+def make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, device):
+  input_queues, npy = make_warp_input_queues(vision_input_shapes, frame_skip, device)

-  fb = input_shapes['features_buffer']  # (1, 24, 512), past features only; the model appends the current frame's feature
-  dp = input_shapes['desire_pulse']  # (1, 25, 8)
+  fb = policy_input_shapes['features_buffer']  # (1, 24, 512), past features only; the model appends the current frame's feature
+  dp = policy_input_shapes['desire_pulse']  # (1, 25, 8)

-  shapes, sizes = get_policy_npy_shapes(input_shapes)
+  shapes, sizes = get_policy_npy_shapes(policy_input_shapes)
  packed_npy_inputs = np.zeros(sum(sizes), dtype=np.float32)
  # views into the packed inputs, to be refilled at runtime
  npy.update({k: v.reshape(s) for (k, s), v in zip(shapes.items(), np.split(packed_npy_inputs, np.cumsum(sizes[:-1])), strict=True)})
@@ -208,10 +206,10 @@ def make_warp(nv12, model_w, model_h, frame_skip):
  return warp


-def make_run_policy(model_runner, model_metadata, frame_skip):
+def make_run_policy(model_runners, model_metadata, frame_skip):
  sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)
  sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
-  npy_shapes, npy_sizes = get_policy_npy_shapes(model_metadata['input_shapes'])
+  vision_features_slice = model_metadata['vision']['output_slices']['hidden_state']

  def run_policy(warped, img_q, big_img_q, feat_q, desire_q, packed_npy_inputs):
    packed_npy_inputs = packed_npy_inputs.to(Device.DEFAULT)
@@ -221,20 +219,23 @@ def make_run_policy(model_runner, model_metadata, frame_skip):
    img = shift_and_sample(img_q, warped[0:1], sample_skip_fn)
    big_img = shift_and_sample(big_img_q, warped[1:2], sample_skip_fn)

-    desire, traffic_convention, action_t, prev_feat = (t.reshape(s) for t, s in zip(packed_npy_inputs.split(npy_sizes), npy_shapes.values(), strict=True))
+    npy_shapes, npy_sizes = get_policy_npy_shapes(model_metadata['on_policy']['input_shapes'])
+    desire, traffic_convention, action_t = (t.reshape(s) for t, s in zip(packed_npy_inputs.split(npy_sizes), npy_shapes.values(), strict=True))
    desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn)
-    feat_buf = shift_and_sample(feat_q, prev_feat.reshape(1, 1, -1), sample_skip_fn)

-    inputs = {
-      'img': img,
-      'big_img': big_img,
+    vision_out = next(iter(model_runners['vision']({'img': img, 'big_img': big_img}).values())).cast('float32')
+    new_feat = vision_out[:, vision_features_slice].reshape(1, -1).unsqueeze(0)
+    feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn)
+
+    policy_inputs = {
      'features_buffer': feat_buf,
      'desire_pulse': desire_buf,
      'traffic_convention': traffic_convention,
      'action_t': action_t,
    }
-    out = next(iter(model_runner(inputs).values())).cast('float32')
-    return out,
+    on_policy_out = next(iter(model_runners['on_policy'](policy_inputs).values())).cast('float32')
+    off_policy_out = next(iter(model_runners['off_policy'](policy_inputs).values())).cast('float32')
+    return vision_out, on_policy_out, off_policy_out
  return run_policy


@@ -306,21 +307,30 @@ if __name__ == "__main__":
  p.add_argument('--model-size', type=_parse_size, required=True, help='model input WxH')
  p.add_argument('--camera-resolutions', type=_parse_size, nargs='+', required=True,
                 help='camera resolutions WxH (one or more)')
-  p.add_argument('--onnx', required=True)
+  p.add_argument('--vision-onnx', required=True)
+  p.add_argument('--off-policy-onnx', required=True)
+  p.add_argument('--on-policy-onnx', required=True)
  p.add_argument('--output', required=True)
  p.add_argument('--frame-skip', type=int, required=True)
  args = p.parse_args()

-  model_path = read_file_chunked_to_disk(args.onnx)
+  model_paths = {
+    'vision': read_file_chunked_to_disk(args.vision_onnx),
+    'off_policy': read_file_chunked_to_disk(args.off_policy_onnx),
+    'on_policy': read_file_chunked_to_disk(args.on_policy_onnx),
+  }
  model_w, model_h = args.model_size

-  model_runner = OnnxRunner(model_path)
-  out = {'metadata': make_metadata_dict(model_path)}
+  model_runners = {name: OnnxRunner(path) for name, path in model_paths.items()}
+  out = {'metadata': {name: make_metadata_dict(path) for name, path in model_paths.items()}}

-  run_policy_jit = TinyJit(make_run_policy(model_runner, out['metadata'], args.frame_skip), prune=True)
+  assert out['metadata']['off_policy']['input_shapes'] == out['metadata']['on_policy']['input_shapes']

-  make_policy_queues = partial(make_input_queues, out['metadata']['input_shapes'], args.frame_skip)
-  make_random_model_inputs = partial(make_random_images, keys=['warped'], shape=(2, 6, *out['metadata']['input_shapes']['img'][2:]), device=WARP_DEV)
+  run_policy_jit = TinyJit(make_run_policy(model_runners, out['metadata'], args.frame_skip), prune=True)
+
+  make_policy_queues = partial(make_input_queues, out['metadata']['vision']['input_shapes'],
+                               out['metadata']['on_policy']['input_shapes'], args.frame_skip)
+  make_random_model_inputs = partial(make_random_images, keys=['warped'], shape=(2, 6, *out['metadata']['vision']['input_shapes']['img'][2:]), device=WARP_DEV)
  out['run_policy'] = compile_jit(run_policy_jit, make_random_model_inputs, POLICY_INPUTS,
                                  make_policy_queues)

@@ -328,7 +338,7 @@ if __name__ == "__main__":
    nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
    make_random_warp_inputs = partial(make_random_images, keys=['frame', 'big_frame'], shape=nv12.size, device=WARP_DEV)
    warp = TinyJit(make_warp(nv12, model_w, model_h, args.frame_skip), prune=True)
-    make_warp_queues = partial(make_warp_input_queues, out['metadata']['input_shapes'], args.frame_skip)
+    make_warp_queues = partial(make_warp_input_queues, out['metadata']['vision']['input_shapes'], args.frame_skip)
    out[(cam_w,cam_h)] = compile_jit(warp, make_random_warp_inputs, WARP_INPUTS, make_warp_queues)

  with open(args.output, "wb") as f:
@@ -85,14 +85,22 @@ class ModelState(ModelStateBase):
    self.WARP_DEV, self.QUEUE_DEV = input_devices['WARP_DEV'], input_devices['QUEUE_DEV']
    jits = load_oob(open_file_chunked(modeld_pkl_path(usbgpu)))
    metadata = jits['metadata']
-    self.input_shapes = metadata['input_shapes']
-    self.vision_input_names = [k for k in self.input_shapes if 'img' in k]
-    self.output_slices = metadata['output_slices']
+    vision_metadata = metadata['vision']
+    self.vision_input_shapes = vision_metadata['input_shapes']
+    self.vision_input_names = [k for k in self.vision_input_shapes if 'img' in k]
+    self.vision_output_slices = vision_metadata['output_slices']
+
+    off_policy_metadata = metadata['off_policy']
+    self.off_policy_output_slices = off_policy_metadata['output_slices']
+
+    policy_metadata = metadata['on_policy']
+    self.policy_input_shapes = policy_metadata['input_shapes']
+    self.policy_output_slices = policy_metadata['output_slices']

    self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32)

    self.frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
-    self.input_queues, self.npy = make_input_queues(self.input_shapes, self.frame_skip, device=self.QUEUE_DEV)
+    self.input_queues, self.npy = make_input_queues(self.vision_input_shapes, self.policy_input_shapes, self.frame_skip, device=self.QUEUE_DEV)
    self.full_frames: dict[str, Tensor] = {}
    self._blob_cache: dict[int, Tensor] = {}
    self.parser = Parser()
@@ -126,15 +134,19 @@ class ModelState(ModelStateBase):

    warped = self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames['img'], big_frame=self.full_frames['big_img'])

-    outs, = self.run_policy(
+    vision_output, on_policy_output, off_policy_output = self.run_policy(
      **{k: self.input_queues[k] for k in POLICY_INPUTS if k in self.input_queues}, warped=warped
    )
-    model_output = outs.numpy()[0]
-    outputs_dict = self.parser.parse_outputs(self.slice_outputs(model_output, self.output_slices))
-    self.npy['prev_feat'][:] = model_output[self.output_slices['hidden_state']]
+    vision_output = vision_output.numpy().flatten()
+    off_policy_output = off_policy_output.numpy().flatten()
+    on_policy_output = on_policy_output.numpy().flatten()
+    vision_outputs_dict = self.parser.parse_vision_outputs(self.slice_outputs(vision_output, self.vision_output_slices))
+    off_policy_outputs_dict = self.parser.parse_off_policy_outputs(self.slice_outputs(off_policy_output, self.off_policy_output_slices))
+    policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(on_policy_output, self.policy_output_slices))
+    outputs_dict = {**vision_outputs_dict, **off_policy_outputs_dict, **policy_outputs_dict}

    if SEND_RAW_PRED:
-      outputs_dict['raw_pred'] = model_output.copy()
+      outputs_dict['raw_pred'] = np.concatenate([vision_output.copy(), on_policy_output.copy(), off_policy_output.copy()])
    return outputs_dict


@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8a26866121d1d3a1152bfce024ed7584b8569507d120d4bc8917320093dcd31a
+size 41191256
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94b07ef7a0f65d5c41ac696b4ae7bdc59e2d4c5f504460e2b0d720620892c2e8
+size 33679037
@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a501760a9d1d5fef0eab2b8c5d122d06124fc26dc8e0782e0aa94b82a208f0ff
-size 1757355221
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eda005282417ffa825092ece5c16b5584142044cdbcf15b6d0246136ac6db601
+size 120584466
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6173be8a69b1d9633a09969c80b2a8bd990bfe7d3e76e192a0e537f6fd72222b
+size 41192485
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b66ef783af3fa86190e85a6b4f729cd1443b20be41134aa258f9c376825a45c
+size 33680163
@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:659727c4d4839adc4992a254409a54259a8756a743f2d567bf5fdc6579f8009b
-size 60881999
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bbd0761201b3b161587d097f173c66bf82cd02966e5f0d1edd888c970d6f6d87
+size 21735970
@@ -96,11 +96,17 @@ class Parser:
    self.parse_mdn('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('road_transform', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,))
+    self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN,ModelConstants.DESIRE_PRED_WIDTH))
+    self.parse_binary_crossentropy('meta', outs)
    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_binary_crossentropy('lane_lines_prob', outs)
-    self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN,ModelConstants.DESIRE_PRED_WIDTH))
-    self.parse_binary_crossentropy('meta', outs)
+    return outs
+
+  def parse_off_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
+    plan_mhp = self.is_mhp(outs, 'plan',  ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH)
+    plan_in_N, plan_out_N = (ModelConstants.PLAN_MHP_N, ModelConstants.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=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH))
    self.parse_binary_crossentropy('lead_prob', outs)
    lead_mhp = self.is_mhp(outs, 'lead', ModelConstants.LEAD_MHP_SELECTION * ModelConstants.LEAD_TRAJ_LEN * ModelConstants.LEAD_WIDTH)
    lead_in_N, lead_out_N = (ModelConstants.LEAD_MHP_N, ModelConstants.LEAD_MHP_SELECTION) if lead_mhp else (0, 0)
@@ -110,11 +116,12 @@ class Parser:
    return outs

  def parse_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
-    self.parse_mdn('plan', outs, in_N=0, out_N=0, out_shape=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH))
+    self.parse_mdn('action', outs, in_N=0, out_N=0, out_shape=(ModelConstants.ACTION_WIDTH,))
    self.parse_categorical_crossentropy('desire_state', outs, out_shape=(ModelConstants.DESIRE_PRED_WIDTH,))
    return outs

  def parse_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
    outs = self.parse_vision_outputs(outs)
+    outs = self.parse_off_policy_outputs(outs)
    outs = self.parse_policy_outputs(outs)
    return outs
@@ -60,22 +60,24 @@ export GIT_LFS_SKIP_SMUDGE=1
 pull_lfs() {
  # The big driving model is not used on these devices yet. Keep its pointer in
  # the worktree, but don't download or copy the 1.8 GB LFS object.
-  LFS_EXCLUDE="openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx"
+  LFS_EXCLUDE="openpilot/selfdrive/modeld/models/big_driving_vision.onnx openpilot/selfdrive/modeld/models/big_driving_on_policy.onnx openpilot/selfdrive/modeld/models/big_driving_off_policy.onnx"

  git config --local lfs.fetchexclude "$LFS_EXCLUDE"
  git lfs pull --exclude="$LFS_EXCLUDE"
-  if git cat-file -e "HEAD:$LFS_EXCLUDE"; then
-    rm -f "$LFS_EXCLUDE"
-    git checkout -- "$LFS_EXCLUDE"
+  for f in $LFS_EXCLUDE; do
+    if git cat-file -e "HEAD:$f"; then
+      rm -f "$f"
+      git checkout -- "$f"

-    # `git lfs prune` retains objects referenced by HEAD, even when excluded.
-    # Remove this one explicitly so safe checkout doesn't rsync it either.
-    oid=$(git show "HEAD:$LFS_EXCLUDE" | sed -n 's/^oid sha256://p')
-    lfs_objects=$(git lfs env | sed -n 's/^LocalMediaDir=//p')
-    if [[ "$oid" =~ ^[0-9a-f]{64}$ && -n "$lfs_objects" ]]; then
-      rm -f "$lfs_objects/${oid:0:2}/${oid:2:2}/$oid"
+      # `git lfs prune` retains objects referenced by HEAD, even when excluded.
+      # Remove this one explicitly so safe checkout doesn't rsync it either.
+      oid=$(git show "HEAD:$f" | sed -n 's/^oid sha256://p')
+      lfs_objects=$(git lfs env | sed -n 's/^LocalMediaDir=//p')
+      if [[ "$oid" =~ ^[0-9a-f]{64}$ && -n "$lfs_objects" ]]; then
+        rm -f "$lfs_objects/${oid:0:2}/${oid:2:2}/$oid"
+      fi
    fi
-  fi
+  done
 }

 safe_checkout() {
@@ -8,13 +8,27 @@ from openpilot.sunnypilot.models.model_name import DEFAULT_MODEL

 DEFAULT_MODEL_NAME_PATH = os.path.join(BASEDIR, "openpilot", "sunnypilot", "models", "model_name.py")
 MODEL_HASH_PATH = os.path.join(BASEDIR, "openpilot", "sunnypilot", "models", "tests", "model_hash")
-SUPERCOMBO_ONNX_PATH = os.path.join(BASEDIR, "openpilot", "selfdrive", "modeld", "models", "driving_supercombo.onnx")
+VISION_ONNX_PATH = os.path.join(BASEDIR, "openpilot", "selfdrive", "modeld", "models", "driving_vision.onnx")
+ON_POLICY_ONNX_PATH = os.path.join(BASEDIR, "openpilot", "selfdrive", "modeld", "models", "driving_on_policy.onnx")
+OFF_POLICY_ONNX_PATH = os.path.join(BASEDIR, "openpilot", "selfdrive", "modeld", "models", "driving_off_policy.onnx")
+
+
+def get_model_hash(path: str) -> str:
+  with open(path, "rb") as f:
+    header = f.read(200).decode("ascii", errors="ignore")
+  if header.startswith("version https://git-lfs.github.com/spec/v1"):
+    for line in header.splitlines():
+      if line.startswith("oid sha256:"):
+        return line.removeprefix("oid sha256:")
+  return get_file_hash(path)


 def update_model_hash():
-  supercombo_hash = get_file_hash(SUPERCOMBO_ONNX_PATH)
+  vision_hash = get_model_hash(VISION_ONNX_PATH)
+  on_policy_hash = get_model_hash(ON_POLICY_ONNX_PATH)
+  off_policy_hash = get_model_hash(OFF_POLICY_ONNX_PATH)

-  combined_hash = hashlib.sha256(supercombo_hash.encode()).hexdigest()
+  combined_hash = hashlib.sha256((vision_hash + on_policy_hash + off_policy_hash).encode()).hexdigest()

  with open(MODEL_HASH_PATH, "w") as f:
    f.write(combined_hash)
@@ -1 +1 @@
-DEFAULT_MODEL = "CD210"
+DEFAULT_MODEL = "Op model16 deep #38073"
@@ -1 +1 @@
-49133798d9cd9cacf47085c7ef8122bfee88cd9c6192a8314c81bfb1b37f5809
+21ce55c002a642d200af88a06ddc618261530beda8246a6b193921bbb9989d8a
@@ -5,16 +5,18 @@ This file is part of sunnypilot and is licensed under the MIT License.
 See the LICENSE.md file in the root directory for more details.
 """

-from openpilot.sunnypilot import get_file_hash
-from openpilot.sunnypilot.models.default_model import MODEL_HASH_PATH, SUPERCOMBO_ONNX_PATH
 import hashlib

+from openpilot.sunnypilot.models.default_model import MODEL_HASH_PATH, OFF_POLICY_ONNX_PATH, ON_POLICY_ONNX_PATH, \
+  VISION_ONNX_PATH, get_model_hash
+

 class TestDefaultModel:
  def test_compare_onnx_hashes(self):
-    supercombo_hash = get_file_hash(SUPERCOMBO_ONNX_PATH)
-
-    combined_hash = hashlib.sha256(supercombo_hash.encode()).hexdigest()
+    vision_hash = get_model_hash(VISION_ONNX_PATH)
+    on_policy_hash = get_model_hash(ON_POLICY_ONNX_PATH)
+    off_policy_hash = get_model_hash(OFF_POLICY_ONNX_PATH)
+    combined_hash = hashlib.sha256((vision_hash + on_policy_hash + off_policy_hash).encode()).hexdigest()

    with open(MODEL_HASH_PATH) as f:
      current_hash = f.read().strip()