modeld_v2: support planplus outputs (#1532)

* conditional concatenation

* v13

* modifucatuins

sunnypilot/modeld_v2/fill_model_msg.py

@@ -11,13 +11,12 @@ SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
 ConfidenceClass = log.ModelDataV2.ConfidenceClass
 
 
-def get_curvature_from_output(output, vego, lat_action_t, mlsim):
+def get_curvature_from_output(output, plan, vego, lat_action_t, mlsim):
   if not mlsim:
     if desired_curv := output.get('desired_curvature'):  # If the model outputs the desired curvature, use that directly
       return float(desired_curv[0, 0])
 
-  plan_output = output['plan'][0]
-  return float(get_curvature_from_plan(plan_output[:, Plan.T_FROM_CURRENT_EULER][:, 2], plan_output[:, Plan.ORIENTATION_RATE][:, 2],
+  return float(get_curvature_from_plan(plan[:, Plan.T_FROM_CURRENT_EULER][:, 2], plan[:, Plan.ORIENTATION_RATE][:, 2],
                                        ModelConstants.T_IDXS, vego, lat_action_t))
 
 

sunnypilot/modeld_v2/modeld.py

@@ -28,6 +28,7 @@ from openpilot.sunnypilot.models.helpers import get_active_bundle
 from openpilot.sunnypilot.models.runners.helpers import get_model_runner
 
 PROCESS_NAME = "selfdrive.modeld.modeld_tinygrad"
+RECOVERY_POWER = 1.0 # The higher this number the more aggressively the model will recover to lanecenter, too high and it will ping-pong
 
 
 class FrameMeta:
@@ -156,11 +157,13 @@ class ModelState(ModelStateBase):
   def get_action_from_model(self, model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action,
                             lat_action_t: float, long_action_t: float, v_ego: float) -> log.ModelDataV2.Action:
     plan = model_output['plan'][0]
+    if 'planplus' in model_output:
+      plan = plan + RECOVERY_POWER*model_output['planplus'][0]
     desired_accel, should_stop = get_accel_from_plan(plan[:, Plan.VELOCITY][:, 0], plan[:, Plan.ACCELERATION][:, 0], self.constants.T_IDXS,
                                                      action_t=long_action_t)
     desired_accel = smooth_value(desired_accel, prev_action.desiredAcceleration, self.LONG_SMOOTH_SECONDS)
 
-    desired_curvature = get_curvature_from_output(model_output, v_ego, lat_action_t, self.mlsim)
+    desired_curvature = get_curvature_from_output(model_output, plan, v_ego, lat_action_t, self.mlsim)
     if self.generation is not None and self.generation >= 10: # smooth curvature for post FOF models
       if v_ego > self.MIN_LAT_CONTROL_SPEED:
         desired_curvature = smooth_value(desired_curvature, prev_action.desiredCurvature, self.LAT_SMOOTH_SECONDS)

sunnypilot/modeld_v2/parse_model_outputs_split.py

@@ -108,6 +108,8 @@ class Parser:
       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:
+        self.parse_mdn('planplus', outs, in_N=plan_in_N, out_N=plan_out_N, out_shape=(SplitModelConstants.IDX_N, SplitModelConstants.PLAN_WIDTH))
 
   def split_outputs(self, outs: dict[str, np.ndarray]) -> None:
     if 'desired_curvature' in outs:

sunnypilot/models/helpers.py

@@ -19,7 +19,7 @@ from openpilot.system.hardware.hw import Paths
 from pathlib import Path
 
 # see the README.md for more details on the model selector versioning
-CURRENT_SELECTOR_VERSION = 12
+CURRENT_SELECTOR_VERSION = 13
 REQUIRED_MIN_SELECTOR_VERSION = 12
 
 USE_ONNX = os.getenv('USE_ONNX', PC)