OP model (#37740)

* Off policy model

* 2f70b996-c604-4a46-9ac9-13ce7534605b/100

* misc fixes

* 1cc1791b-4555-41ce-a5cb-ce046967075a/100

* fix model

* 6ab6fae5-fbbd-4ad0-928a-b33794f60dba/100

* recomp

* update models

* qxfinally correct

* b8b96ac6-7918-401a-a862-eaf1fdbba88d/100

* wrong plan

* wrong plan

* Vf9b3fb5f-4d0d-4dcb-bc3a-5e94d1fdcdaa/200

* bump dbc

* ready to merge

* rename to on-policy

* Just cleanup big models for now

---------

Co-authored-by: Kacper Rączy <gfw.kra@gmail.com>

Author: haraschax

scripts/reporter.py

@@ -38,6 +38,11 @@ def get_checkpoint(f):
       continue
 
     fn = os.path.basename(f)
-    master = get_checkpoint(MASTER_PATH + MODEL_PATH + fn)
+    master_path = MASTER_PATH + MODEL_PATH + fn
+    if os.path.exists(master_path):
+      master = get_checkpoint(master_path)
+      master_col = f"[{master}](https://reporter.comma.life/experiment/{master})"
+    else:
+      master_col = "N/A (new model)"
     pr = get_checkpoint(BASEDIR + MODEL_PATH + fn)
-    print("|", fn, "|", f"[{master}](https://reporter.comma.life/experiment/{master})", "|", f"[{pr}](https://reporter.comma.life/experiment/{pr})", "|")
+    print("|", fn, "|", master_col, "|", f"[{pr}](https://reporter.comma.life/experiment/{pr})", "|")

selfdrive/modeld/SConscript

@@ -21,7 +21,7 @@ tg_flags = {
 }.get(arch, 'DEV=CPU CPU_LLVM=1 THREADS=0')
 
 # Get model metadata
-for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']:
+for model_name in ['driving_vision', 'driving_off_policy', 'driving_on_policy', 'dmonitoring_model']:
   fn = File(f"models/{model_name}").abspath
   script_files = [File(Dir("#selfdrive/modeld").File("get_model_metadata.py").abspath)]
   cmd = f'{tg_flags} python3 {Dir("#selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx'
@@ -59,19 +59,5 @@ def tg_compile(flags, model_name):
   )
 
 # Compile small models
-for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']:
+for model_name in ['driving_vision', 'driving_off_policy', 'driving_on_policy', 'dmonitoring_model']:
   tg_compile(tg_flags, model_name)
-
-# Compile BIG model if USB GPU is available
-if "USBGPU" in os.environ:
-  import subprocess
-  # because tg doesn't support multi-process
-  devs = subprocess.check_output('python3 -c "from tinygrad import Device; print(list(Device.get_available_devices()))"', shell=True, cwd=env.Dir('#').abspath)
-  if b"AMD" in devs:
-    print("USB GPU detected... building")
-    flags = "DEV=AMD AMD_IFACE=USB AMD_LLVM=1 NOLOCALS=0 IMAGE=0"
-    bp = tg_compile(flags, "big_driving_policy")
-    bv = tg_compile(flags, "big_driving_vision")
-    lenv.SideEffect('lock', [bp, bv])  # tg doesn't support multi-process so build serially
-  else:
-    print("USB GPU not detected... skipping")

selfdrive/modeld/modeld.py

@@ -34,11 +34,13 @@
 PROCESS_NAME = "selfdrive.modeld.modeld"
 SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
 
-VISION_PKL_PATH = Path(__file__).parent / 'models/driving_vision_tinygrad.pkl'
-POLICY_PKL_PATH = Path(__file__).parent / 'models/driving_policy_tinygrad.pkl'
-VISION_METADATA_PATH = Path(__file__).parent / 'models/driving_vision_metadata.pkl'
-POLICY_METADATA_PATH = Path(__file__).parent / 'models/driving_policy_metadata.pkl'
 MODELS_DIR = Path(__file__).parent / 'models'
+VISION_PKL_PATH = MODELS_DIR / 'driving_vision_tinygrad.pkl'
+VISION_METADATA_PATH = MODELS_DIR / 'driving_vision_metadata.pkl'
+ON_POLICY_PKL_PATH = MODELS_DIR / 'driving_on_policy_tinygrad.pkl'
+ON_POLICY_METADATA_PATH = MODELS_DIR / 'driving_on_policy_metadata.pkl'
+OFF_POLICY_PKL_PATH = MODELS_DIR / 'driving_off_policy_tinygrad.pkl'
+OFF_POLICY_METADATA_PATH = MODELS_DIR / 'driving_off_policy_metadata.pkl'
 
 LAT_SMOOTH_SECONDS = 0.0
 LONG_SMOOTH_SECONDS = 0.3
@@ -151,7 +153,13 @@ def __init__(self):
       self.vision_output_slices = vision_metadata['output_slices']
       vision_output_size = vision_metadata['output_shapes']['outputs'][1]
 
-    with open(POLICY_METADATA_PATH, 'rb') as f:
+    with open(OFF_POLICY_METADATA_PATH, 'rb') as f:
+      off_policy_metadata = pickle.load(f)
+      self.off_policy_input_shapes =  off_policy_metadata['input_shapes']
+      self.off_policy_output_slices = off_policy_metadata['output_slices']
+      off_policy_output_size = off_policy_metadata['output_shapes']['outputs'][1]
+
+    with open(ON_POLICY_METADATA_PATH, 'rb') as f:
       policy_metadata = pickle.load(f)
       self.policy_input_shapes =  policy_metadata['input_shapes']
       self.policy_output_slices = policy_metadata['output_slices']
@@ -175,11 +183,13 @@ def __init__(self):
     self.vision_output = np.zeros(vision_output_size, dtype=np.float32)
     self.policy_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()}
     self.policy_output = np.zeros(policy_output_size, dtype=np.float32)
+    self.off_policy_output = np.zeros(off_policy_output_size, dtype=np.float32)
     self.parser = Parser()
     self.frame_buf_params : dict[str, tuple[int, int, int, int]] = {}
     self.update_imgs = None
     self.vision_run = pickle.loads(read_file_chunked(str(VISION_PKL_PATH)))
-    self.policy_run = pickle.loads(read_file_chunked(str(POLICY_PKL_PATH)))
+    self.policy_run = pickle.loads(read_file_chunked(str(ON_POLICY_PKL_PATH)))
+    self.off_policy_run = pickle.loads(read_file_chunked(str(OFF_POLICY_PKL_PATH)))
 
   def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]:
     parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()}
@@ -228,9 +238,17 @@ def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray],
 
     self.policy_output = self.policy_run(**self.policy_inputs).contiguous().realize().uop.base.buffer.numpy().flatten()
     policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(self.policy_output, self.policy_output_slices))
-    combined_outputs_dict = {**vision_outputs_dict, **policy_outputs_dict}
+
+    self.off_policy_output = self.off_policy_run(**self.policy_inputs).contiguous().realize().uop.base.buffer.numpy()
+    off_policy_outputs_dict = self.parser.parse_off_policy_outputs(self.slice_outputs(self.off_policy_output, self.off_policy_output_slices))
+    off_policy_outputs_dict.pop('plan')
+
+
+    combined_outputs_dict = {**vision_outputs_dict, **off_policy_outputs_dict, **policy_outputs_dict}
+    if 'planplus' in combined_outputs_dict and 'plan' in combined_outputs_dict:
+      combined_outputs_dict['plan'] = combined_outputs_dict['plan'] + combined_outputs_dict['planplus']
     if SEND_RAW_PRED:
-      combined_outputs_dict['raw_pred'] = np.concatenate([self.vision_output.copy(), self.policy_output.copy()])
+      combined_outputs_dict['raw_pred'] = np.concatenate([self.vision_output.copy(), self.policy_output.copy(), self.off_policy_output.copy()])
 
     return combined_outputs_dict
 

selfdrive/modeld/models/big_driving_policy.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eb6992bd60bada6162fea298e1a414b6b3d6a326db4eda46b9de62bcd8554754
+size 13393859

selfdrive/modeld/models/big_driving_vision.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86680a657bbb34f997034d1930bb2cb65c38b9222cea199732f72bd45791cfad
+size 13022803

selfdrive/modeld/models/driving_off_policy.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:940e9006a25f27f0b6e85da798e6a8fd1f6dd492dd7d0b9ff1a9436460f46129
-size 46887794
+oid sha256:7af05e03fd170653ff5771baf373a2c57b363da12c4c411cd416dee067b4cf58
+size 23266366

selfdrive/modeld/models/driving_on_policy.onnx

@@ -96,11 +96,17 @@ def parse_vision_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndar
     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)
+    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_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)
     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)
@@ -120,5 +126,6 @@ def parse_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndar
 
   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