From eff08cb17f2b0d00d24cd13729b8d16ec394c82c Mon Sep 17 00:00:00 2001
From: YassineYousfi <yyousfi1@binghamton.edu>
Date: Thu, 12 May 2022 19:47:33 -0700
Subject: [PATCH] Joint Model (#24213)

* update models

* wip lanelines mhp parsing

* typos

* newer models

* 95995a49-db0c-4261-8776-b90780dc2a8c/600

* Use laneful policy to prevent planner aggression

* c8c3ab69-bd2c-4d1c-a7f8-9e8457e4827b/950

* ede8f75a-df3e-41fb-bc56-63e2f33858e5/950

* 7205c5af-7532-4215-ad9b-678789054db7/995

* bfb34fa3-1d72-4987-854f-43d42ee01015/950

* change cost

* 8ae2477a-b452-4601-b04c-d3af66ea5e98/950

* Update heading cost

* 9b991314-94f5-4f0b-872e-58d95266e4d0/333

* 9b991314-94f5-4f0b-872e-58d95266e4d0/950

* ede8f75a-df3e-41fb-bc56-63e2f33858e5/950

* undo lateral planner changes

* rm whitespace

* update ref

Co-authored-by: Harald Schafer <harald.the.engineer@gmail.com>
---
 models/supercombo.dlc                         |  4 +--
 models/supercombo.onnx                        |  4 +--
 selfdrive/modeld/models/driving.cc            | 29 ++++++++++-------
 selfdrive/modeld/models/driving.h             | 32 ++++++++++++-------
 .../process_replay/model_replay_ref_commit    |  2 +-
 5 files changed, 43 insertions(+), 28 deletions(-)

diff --git a/models/supercombo.dlc b/models/supercombo.dlc
index a27e3d1180..e46a5ef679 100644
--- a/models/supercombo.dlc
+++ b/models/supercombo.dlc
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ba3fe3e61853cc1434e3e220f40c8e9d1f1b9bab8458196ba3bea6a10b82c6ed
-size 72718099
+oid sha256:d7e58be246bb9fb27e1f7a60ab2ed1ec4ee7bdc99cf82d92af694f77c83e43aa
+size 108294663
diff --git a/models/supercombo.onnx b/models/supercombo.onnx
index 3039035fbc..9efe640760 100644
--- a/models/supercombo.onnx
+++ b/models/supercombo.onnx
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bda57c1a66944f5a633ecd739a24d62702c717a234f2fdcc499dfa1d61c3c19e
-size 73147489
+oid sha256:9b0c23ca7198d945303f407cf374c9e5cffc13e0ca2fde3420748e7fbb3eddf8
+size 108687135
diff --git a/selfdrive/modeld/models/driving.cc b/selfdrive/modeld/models/driving.cc
index e14411df78..696e36078f 100644
--- a/selfdrive/modeld/models/driving.cc
+++ b/selfdrive/modeld/models/driving.cc
@@ -228,19 +228,24 @@ void fill_plan(cereal::ModelDataV2::Builder &framed, const ModelOutputPlanPredic
 
 void fill_lane_lines(cereal::ModelDataV2::Builder &framed, const std::array<float, TRAJECTORY_SIZE> &plan_t,
                      const ModelOutputLaneLines &lanes) {
+
+  const auto &left_far = lanes.get_lane_idx(0);
+  const auto &left_near = lanes.get_lane_idx(1);
+  const auto &right_near = lanes.get_lane_idx(2);
+  const auto &right_far = lanes.get_lane_idx(3);
   std::array<float, TRAJECTORY_SIZE> left_far_y, left_far_z;
   std::array<float, TRAJECTORY_SIZE> left_near_y, left_near_z;
   std::array<float, TRAJECTORY_SIZE> right_near_y, right_near_z;
   std::array<float, TRAJECTORY_SIZE> right_far_y, right_far_z;
   for (int j=0; j<TRAJECTORY_SIZE; j++) {
-    left_far_y[j] = lanes.mean.left_far[j].y;
-    left_far_z[j] = lanes.mean.left_far[j].z;
-    left_near_y[j] = lanes.mean.left_near[j].y;
-    left_near_z[j] = lanes.mean.left_near[j].z;
-    right_near_y[j] = lanes.mean.right_near[j].y;
-    right_near_z[j] = lanes.mean.right_near[j].z;
-    right_far_y[j] = lanes.mean.right_far[j].y;
-    right_far_z[j] = lanes.mean.right_far[j].z;
+    left_far_y[j] = left_far.mean[j].y;
+    left_far_z[j] = left_far.mean[j].z;
+    left_near_y[j] = left_near.mean[j].y;
+    left_near_z[j] = left_near.mean[j].z;
+    right_near_y[j] = right_near.mean[j].y;
+    right_near_z[j] = right_near.mean[j].z;
+    right_far_y[j] = right_far.mean[j].y;
+    right_far_z[j] = right_far.mean[j].z;
   }
 
   auto lane_lines = framed.initLaneLines(4);
@@ -250,10 +255,10 @@ void fill_lane_lines(cereal::ModelDataV2::Builder &framed, const std::array<floa
   fill_xyzt(lane_lines[3], plan_t, X_IDXS_FLOAT, right_far_y, right_far_z);
 
   framed.setLaneLineStds({
-    exp(lanes.std.left_far[0].y),
-    exp(lanes.std.left_near[0].y),
-    exp(lanes.std.right_near[0].y),
-    exp(lanes.std.right_far[0].y),
+    exp(left_far.std[0].y),
+    exp(left_near.std[0].y),
+    exp(right_near.std[0].y),
+    exp(right_far.std[0].y),
   });
 
   framed.setLaneLineProbs({
diff --git a/selfdrive/modeld/models/driving.h b/selfdrive/modeld/models/driving.h
index 47dec0d9d7..b00542dd7b 100644
--- a/selfdrive/modeld/models/driving.h
+++ b/selfdrive/modeld/models/driving.h
@@ -28,6 +28,8 @@ constexpr int META_STRIDE = 7;
 constexpr int PLAN_MHP_N = 5;
 constexpr int STOP_LINE_MHP_N = 3;
 
+constexpr int LANELINES_MHP_N = 4;
+
 constexpr int LEAD_MHP_N = 2;
 constexpr int LEAD_TRAJ_LEN = 6;
 constexpr int LEAD_PRED_DIM = 4;
@@ -77,14 +79,6 @@ struct ModelOutputPlans {
 };
 static_assert(sizeof(ModelOutputPlans) == sizeof(ModelOutputPlanPrediction)*PLAN_MHP_N);
 
-struct ModelOutputLinesXY {
-  std::array<ModelOutputYZ, TRAJECTORY_SIZE> left_far;
-  std::array<ModelOutputYZ, TRAJECTORY_SIZE> left_near;
-  std::array<ModelOutputYZ, TRAJECTORY_SIZE> right_near;
-  std::array<ModelOutputYZ, TRAJECTORY_SIZE> right_far;
-};
-static_assert(sizeof(ModelOutputLinesXY) == sizeof(ModelOutputYZ)*TRAJECTORY_SIZE*4);
-
 struct ModelOutputLineProbVal {
   float val_deprecated;
   float val;
@@ -99,12 +93,28 @@ struct ModelOutputLinesProb {
 };
 static_assert(sizeof(ModelOutputLinesProb) == sizeof(ModelOutputLineProbVal)*4);
 
+struct ModelOutputLaneLinesElement {
+  std::array<ModelOutputYZ, TRAJECTORY_SIZE> mean;
+  std::array<ModelOutputYZ, TRAJECTORY_SIZE>  std;
+  std::array<float, LANELINES_MHP_N> prob;
+};
+static_assert(sizeof(ModelOutputLaneLinesElement) == (sizeof(ModelOutputYZ)*TRAJECTORY_SIZE*2) + sizeof(float)*LANELINES_MHP_N);
+
 struct ModelOutputLaneLines {
-  ModelOutputLinesXY mean;
-  ModelOutputLinesXY std;
+  std::array<ModelOutputLaneLinesElement, LANELINES_MHP_N> prediction;
   ModelOutputLinesProb prob;
+
+  constexpr const ModelOutputLaneLinesElement &get_lane_idx(int lane_idx) const {
+    int max_idx = 0;
+    for (int i = 1; i < prediction.size(); i++) {
+      if (prediction[i].prob[lane_idx] > prediction[max_idx].prob[lane_idx]) {
+        max_idx = i;
+      }
+    }
+    return prediction[max_idx];
+  }
 };
-static_assert(sizeof(ModelOutputLaneLines) == (sizeof(ModelOutputLinesXY)*2) + sizeof(ModelOutputLinesProb));
+static_assert(sizeof(ModelOutputLaneLines) == (sizeof(ModelOutputLaneLinesElement)*LANELINES_MHP_N) + (sizeof(ModelOutputLinesProb)));
 
 struct ModelOutputEdgessXY {
   std::array<ModelOutputYZ, TRAJECTORY_SIZE> left;
diff --git a/selfdrive/test/process_replay/model_replay_ref_commit b/selfdrive/test/process_replay/model_replay_ref_commit
index bf4e1eadb8..1f652422f7 100644
--- a/selfdrive/test/process_replay/model_replay_ref_commit
+++ b/selfdrive/test/process_replay/model_replay_ref_commit
@@ -1 +1 @@
-5fb5cd71a6e878cf45593ecea22c93e932f70c31
\ No newline at end of file
+885bbd4bc01d474795ae81da734be78b75fa4658